Softmax: Function & Applications

• Softmax functions are function that normalize a vector into a probability distribution. The function is used in logistic regression and neural networks to model the probabilities of events.
• It is typically used in logistic regression and neural network models, particularly in the final layer of classification tasks, such as image and language classification.
• The softmax function is characterized by properties such as normalization, temperature sensitivity, and robustness, and is evaluated through the use of optimized algorithms and approximations.

Softmax is the normalized exponential map that converts a finite vector of real-valued scores, or logits, into a probability vector. For $z=(z_1,\ldots,z_K)$, its standard form is

$$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$

Every output is positive and the outputs sum to one. Softmax is invariant under addition of a common constant, $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$, and is therefore determined by relative logit differences. It is used in multiclass classification, language modeling, Transformer attention, reinforcement learning, generative adversarial networks, optimization, robustness verification, and low-bit inference. Its behavior is governed by normalization, temperature, numerical precision, output dimensionality, and the optimization objective in which it is embedded.

1. Definition, geometry, and differential properties

Given logits $z\in\mathbb R^K$, the softmax output $p=\operatorname{softmax}(z)$ lies in the interior of the probability simplex,

$p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$

The predicted class is ordinarily $\arg\max_i z_i=\arg\max_i p_i$. A common training objective is cross-entropy with a target distribution $y$,

$L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$

For a one-hot target $t$, this reduces to $$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$0. Its logit derivative is

$$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$1

The derivative is local to every class because the normalization denominator couples all logits.

The Jacobian of softmax is

$$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$2

Its diagonal entries are $$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$3 and its off-diagonal entries are $$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$4. The matrix is the covariance matrix of a categorical random variable with probabilities $$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$5. Consequently, increasing one logit increases its own probability while decreasing the relative probabilities of competing classes.

Softmax is smooth and order-preserving, but it is not globally convex or concave as a vector-valued map. Its translation invariance permits numerically stable evaluation by subtracting the largest logit:

$$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$6

This transformation leaves the exact mathematical distribution unchanged while preventing exponential overflow. Direct evaluation can overflow for sufficiently large positive logits and underflow for sufficiently negative logits; in single precision, the cited analysis reports approximate overflow beyond $$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$7 and underflow below $$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$8 (Dukhan et al., 2020).

A recent norm-uniform analysis establishes the sharper global Lipschitz inequality

$$\operatorname{softmax}(z)_i=\frac{\exp(z_i)}{\sum_{j=1}^{K}\exp(z_j)}.$$9

for every $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$0, where $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$1 is the inverse-temperature parameter. The constant $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$2 is attained for $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$3 and $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$4 norms and is the supremum for intermediate $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$5 norms; for $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$6 and dimension greater than two, it is approached by distributions tending toward $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$7 (Nair, 27 Oct 2025). This improves analyses that use the looser constant $$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$8 and directly affects robustness, smoothness, attention sensitivity, stochastic policies, and entropy-regularized optimization.

2. Temperature, probabilities, and learning dynamics

Temperature rescales the logits:

$$\operatorname{softmax}(z+c\mathbf 1)=\operatorname{softmax}(z)$$9

Small $z\in\mathbb R^K$0 sharpens the distribution, while large $z\in\mathbb R^K$1 flattens it. As $z\in\mathbb R^K$2, the output approaches a one-hot vector when the maximum is unique; as $z\in\mathbb R^K$3, it approaches the uniform distribution. Temperature is equivalent to inverse scaling of the logits, so the effective sharpness is controlled approximately by $z\in\mathbb R^K$4.

The temperature-scaled Jacobian and cross-entropy gradient are

$z\in\mathbb R^K$5

Temperature therefore modifies both the probabilities and the direct gradient scale. High temperature initially produces nearly uniform predictions and small explicit gradients, while low temperature produces concentrated predictions and can still lead to vanishing gradients when the output saturates.

The interaction between temperature and representation geometry has been associated with rank-deficit bias: softmax-trained networks may converge to representations whose effective rank is substantially below the number of classes. High temperature or low logit norm can induce singular-vector alignment, dominance of a few singular directions, and low-rank internal representations. This can reduce effective depth and harm final-layer out-of-distribution generalization while improving out-of-distribution detection. In experiments across MLPs, VGG, ResNets, and ViT models on CIFAR and ImageNet datasets, increasing temperature generally decreased solution rank and worsened final-layer out-of-distribution generalization, while improving some out-of-distribution detection measures (Masarczyk et al., 2 Jun 2025).

Temperature is not equivalent to calibration in all contexts. In active learning, softmax confidence is often used through least confidence, entropy, or max-margin acquisition. However, softmax values need not be calibrated probabilities: neural networks can produce extreme confidence for unfamiliar or incorrect inputs, and the simplex constraint forces every input to distribute all probability mass among known classes. Experiments with BERT and RoBERTa found that vanilla softmax remained competitive with eight alternatives, while selecting the most uncertain outliers exclusively could degrade active-learning performance. Uncertainty Clipping, which ignores the top 5% of uncertainty scores before selection, often improved acquisition performance (Gonsior et al., 2022).

3. Computational complexity and scalable normalization

For a classifier with hidden representation $z\in\mathbb R^K$6 and $z\in\mathbb R^K$7 classes, logits of the form $z\in\mathbb R^K$8 require $z\in\mathbb R^K$9 computation. The normalization and gradient also involve all classes. This becomes a bottleneck in language modeling, recommendation, extreme classification, and large-vocabulary prediction.

Full softmax can be replaced or reorganized in several ways. Sampled softmax evaluates the target class and a subset of negative classes. If $p=\operatorname{softmax}(z)$0 negatives are sampled from proposal $p=\operatorname{softmax}(z)$1, the negative logit is adjusted by $p=\operatorname{softmax}(z)$2. The resulting cost is approximately $p=\operatorname{softmax}(z)$3, but the estimator is generally biased because the sampled denominator appears inside a nonlinear reciprocal or logarithm. Exact unbiasedness under the cited formulation requires sampling from the model’s own softmax distribution, which is itself expensive to compute (Blanc et al., 2017).

The quality of sampled softmax is therefore controlled by the proposal distribution. Uniform, unigram, log-uniform, and other fixed proposals can differ substantially from the input-dependent model distribution. Kernel-based sampling constructs

$p=\operatorname{softmax}(z)$4

using a nonnegative kernel with a finite-dimensional feature map. Precomputed feature summaries in a binary tree allow exact sampling from the kernel distribution in $p=\operatorname{softmax}(z)$5 time per negative sample, where $p=\operatorname{softmax}(z)$6 is the feature dimension. The quadratic kernel $p=\operatorname{softmax}(z)$7 has $p=\operatorname{softmax}(z)$8; Random Fourier Softmax instead approximates the exponential dot-product kernel with random Fourier features and achieves sampling cost $p=\operatorname{softmax}(z)$9 (Blanc et al., 2017, Rawat et al., 2019).

Adaptive softmax uses a hierarchical classifier rather than sampling negatives. The vocabulary is divided into a frequent-word head and increasingly large rare-word clusters. The root classifier predicts either a head word or a cluster symbol; a tail classifier then predicts the word within the selected cluster. For a head word $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$0,

$p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$1

whereas for $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$2 in cluster $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$3,

$p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$4

The model is exactly normalized for its hierarchical distribution, but it is not generally the same as the original flat $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$5-way softmax. Cluster sizes are chosen from word frequencies and measured GPU multiplication costs. Reported GPU speedups relative to full softmax range approximately from $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$6 to $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$7800{,}000$ words (Grave et al., 2016).

A separate optimization concerns the softmax kernel itself. CPU implementations commonly require three logical passes: maximum reduction, normalization reduction, and output generation. The Two-Pass Softmax algorithm represents each exponential as a mantissa–exponent pair, accumulates the denominator in the same representation, and avoids a separate maximum-finding pass. Because optimized softmax can be memory-bandwidth bound, reducing memory transfers improves performance. Reported out-of-cache gains reach up to 28% with AVX512 and 18% with AVX2 (Dukhan et al., 2020). For Intel Xeon inference, replacing general Eigen operations with MKL routines, SIMD reductions, and selective JIT assembly produced approximately a twofold softmax speedup and 15–20% faster end-to-end DAM inference (Czaja et al., 2019).

4. Softmax variants and alternative objectives

Softmax is not a single training principle: changing its normalization, logits, gradients, or active classes produces distinct objectives.

Margin-based softmax modifies the target logit. With margin $p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$8,

$p_i>0,\qquad \sum_{i=1}^{K}p_i=1.$9

while non-target logits remain unchanged. The target must therefore exceed every competitor by a positive score gap. Ensemble soft-margin softmax combines this mechanism with multiple classifier heads whose diversity is regularized using the Hilbert–Schmidt Independence Criterion. The resulting EM-Softmax objective improves feature separation through the margin and classifier strength through diverse heads. On the reported MNIST, CIFAR, and ImageNet32 experiments, EM-Softmax outperformed ordinary softmax and the individual margin-only and ensemble-only ablations (Wang et al., 2018).

Taylor softmax replaces the exponential with an even-order Taylor polynomial,

$\arg\max_i z_i=\arg\max_i p_i$0

Even-order polynomials remain positive for all real inputs. SM-Taylor combines this transformation with a target-logit margin. Experiments on MNIST, CIFAR-10, and CIFAR-100 found that the best Taylor order was dataset-dependent; the best SM-Taylor configuration outperformed standard softmax on all three datasets, although the paper did not establish a universally superior replacement (Banerjee et al., 2020).

Adaptive Sparse Softmax masks non-target classes whose probabilities are sufficiently below the target:

$\arg\max_i z_i=\arg\max_i p_i$1

The remaining classes are renormalized, and when all competitors are masked the loss becomes exactly zero. Thus, unlike ordinary cross-entropy, easy examples can cease contributing gradients at finite logits. The method aims to align training with test-time $\arg\max_i z_i=\arg\max_i p_i$2 classification by retaining strong competitors while discarding weak ones. Its associated AS-Speed procedure accumulates gradients according to the masked-sample ratio, producing a reported approximate training speedup of $\arg\max_i z_i=\arg\max_i p_i$3 (Lv et al., 5 Aug 2025).

$\arg\max_i z_i=\arg\max_i p_i$4-softmax modifies the largest ordinary softmax probability. If $\arg\max_i z_i=\arg\max_i p_i$5 and $\arg\max_i z_i=\arg\max_i p_i$6, then

$\arg\max_i z_i=\arg\max_i p_i$7

The output remains normalized and approaches a one-hot vector as $\arg\max_i z_i=\arg\max_i p_i$8 grows. The distance to the nearest one-hot vector is bounded by

$\arg\max_i z_i=\arg\max_i p_i$9

When used with cross-entropy, the transformation reduces the gradient for already-correct confident examples, producing a “soft early-stopping” effect. Combined with MAE, it yielded improved robustness to synthetic, human-annotated, instance-dependent, and real-world label noise (Wang et al., 4 Aug 2025).

Mellowmax is an alternative aggregation operator for reinforcement learning:

$y$0

It approaches the arithmetic mean as $y$1 and the maximum as $y$2. Unlike the Boltzmann softmax value operator, mellowmax is a non-expansion under the infinity norm. Consequently, a discounted Bellman operator using mellowmax is a contraction for $y$3, yielding a unique fixed point under standard tabular assumptions. Its associated maximum-entropy policy has a Boltzmann form but uses a state-dependent inverse temperature selected so that the policy’s expected action value equals the mellowmax value (Asadi et al., 2016).

Self-Adjust Softmax modifies attention by multiplying softmax by a score-dependent factor:

$y$4

The derivative of the unnormalized form contains an additional diagonal term:

$y$5

This can mitigate the vanishing Jacobian of saturated softmax. Normalized variants use rowwise minimum and maximum values to constrain the multiplier, but the resulting attention weights generally do not sum to one. Experiments with Transformer models up to 2.7 billion parameters reported generally modest improvements, with configuration-dependent exceptions and additional rowwise reduction overhead (Zheng et al., 25 Feb 2025).

5. Applications beyond ordinary classification

In language modeling, softmax converts the final classifier logits over a vocabulary into next-token probabilities. The full output layer can dominate training cost, motivating adaptive, sampled, and kernel-based methods. In attention, softmax is applied row-wise to scaled query–key scores:

$y$6

Each row defines normalized weights over keys. Softmax’s normalization couples all positions and gives attention weights a probabilistic interpretation, but it also introduces exponential evaluation, reductions, and sensitivity to saturation.

In low-bit FlashAttention-style inference, EFQ-Softmax directly maps shifted attention scores to block-scaled nonnegative E2M1 operands rather than evaluating high-precision exponentials and subsequently quantizing them. A microscaling block uses a power-of-two scale and an affine residual-domain code rule. The same approximate probability operand is used in both the numerator $y$7 and denominator $y$8, while row-maximum updates, historical rescaling, high-precision accumulation, and final normalization remain unchanged. On Qwen3 and Qwen3-VL, reported task means improved over MXFP4, and the A5 vector-unit probability-generation stage was reduced by approximately 40.33% (Han et al., 9 Sep 2026).

Softmax GAN applies normalization across individual samples in a minibatch rather than independently classifying each sample as real or generated. For a batch $y$9, the discriminator defines

$L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$0

The discriminator target places uniform probability on real samples and zero probability on generated samples. The generator target is uniform over all samples in the batch. The shared normalization couples all examples, and the resulting discriminator can be interpreted as an importance-sampling approximation to energy-based maximum likelihood. The method was reported to stabilize training and reduce some mode-collapse behavior in CelebA experiments (Lin, 2017).

Softmax also appears in convex optimization and verification. In a parameterized softmax regression problem, a matrix $L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$1 and parameter vector $L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$2 produce $L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$3, and the objective fits this entire probability vector to a target $L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$4 using squared Euclidean distance. The Hessian has the structured form

$L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$5

with $L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$6 under the stated assumptions. A sufficiently strong diagonal regularizer makes the objective strongly convex and supports a local approximate-Newton convergence theorem (Deng et al., 2023).

For robustness verification, softmax components are bounded by convex lower bounds and concave upper bounds. The exponential–reciprocal and log-sum-exp decompositions produce tighter relaxations than earlier linear compositions. These bounds improve certified robustness of Transformer attention and certified uncertainty estimates for deep ensembles, although nonlinear formulations are substantially more expensive than affine relaxations (Wei et al., 2023).

6. Limitations, trade-offs, and interpretation

Softmax is an exact normalized model only when the full denominator is evaluated. Sampled softmax trades exactness for reduced computation and is biased unless the proposal coincides with the model distribution. Hierarchical and adaptive softmax define different normalized distributions rather than merely omitting terms from the flat model. Kernel proposals reduce sampling bias but remain approximations whose quality depends on the feature map, sample size, and output-layer structure.

The probabilistic interpretation is also objective-dependent. Ordinary softmax produces a normalized distribution over all classes, but margin-based, sparse, adaptive, and score-multiplied variants may normalize over subsets, cease to sum to one, or lose calibration. This is acceptable in objectives concerned primarily with ranking, representation geometry, robustness, or optimization, but it is not equivalent to modeling a calibrated categorical distribution.

Saturation is a recurring limitation. When one probability approaches one and the others approach zero, the ordinary softmax Jacobian approaches zero. Temperature, self-adjusting factors, adaptive masking, and $L(y,p)=-\sum_{i=1}^{K}y_i\log p_i.$7-softmax address different aspects of this behavior, but none universally removes all optimization or calibration problems. Conversely, sharper or more one-hot outputs can improve noise tolerance, classification margins, or out-of-distribution detection while reducing representation diversity, clean-data fitting ability, or out-of-distribution generalization.

Numerical stability requires maximum subtraction or an equivalent representation. Direct exponentiation can overflow, while low-probability terms can underflow. CPU and accelerator implementations are often limited by memory traffic, transcendental evaluation, or quantization overhead rather than by the abstract arithmetic complexity alone. Optimizations must therefore preserve stable normalization, distinguish training from inference safeguards, and validate end-to-end effects rather than relying solely on kernel-level speedups.

Theoretical guarantees are typically conditional. Mellowmax provides a tabular contraction guarantee but not a general nonlinear-function-approximation guarantee. Approximate-Newton softmax regression requires strong regularization, full column rank, local initialization, and a positive smallest singular value. NTK analyses of softmax networks rely on over-parameterization, training close to initialization, and a positive-definite initial NTK. Robustness and verification bounds involve box domains, curvature-compatible relaxations, and solver-dependent costs. Empirical improvements reported for alternative softmax operators are consequently task-, architecture-, temperature-, margin-, and calibration-dependent.

Softmax remains widely used because it simultaneously provides order-preserving normalization, differentiability, probabilistic outputs, compatibility with cross-entropy, and a natural interface between real-valued neural scores and categorical decisions. Contemporary research treats it not merely as a final activation but as a source of computational cost, optimization geometry, representation bias, uncertainty behavior, robustness structure, and hardware-specific design constraints.