← Nestor G Pestelos Jr · Reference
Mathematics & Machine Learning
Softmax Function
August 25, 2026
Softmax is a normalized exponential mathematical function that maps a \(K\)-dimensional vector of arbitrary real-valued numbers (logits) into a \(K\)-dimensional vector of real values in the range \((0, 1)\) that sum to exactly \(1\), forming a valid discrete probability distribution. Originally introduced by John S. Bridle in 1990 for feed-forward neural networks [1], Softmax serves as the foundational probability layer for multi-class classification, autoregressive next-token prediction, self-attention routing in transformer architectures, and sparse expert gating in Mixture-of-Experts systems.
1. Mathematical Formulation
1.1 Standard Definition
Given an input vector \(\mathbf{z} = (z_1, z_2, \dots, z_K) \in \mathbb{R}^K\), the Softmax function \(\sigma: \mathbb{R}^K o (0, 1)^K\) is defined for each component \(i \in \{1, \dots, K\}\) as:
\[\sigma(\mathbf{z})_i = rac{e^{z_i}}{\sum_{j=1}^K e^{z_j}}\]The denominator \(\sum_{j=1}^K e^{z_j}\) acts as a normalizing partition function (analogous to the Boltzmann distribution in statistical mechanics), ensuring that the output elements satisfy the Kolmogorov probability axioms.
1.2 Core Mathematical Properties
- Probability Constraints: For all \(i\), \(\sigma(\mathbf{z})_i > 0\) and \(\sum_{i=1}^K \sigma(\mathbf{z})_i = 1\).
- Strict Monotonicity: If \(z_a > z_b\), then \(\sigma(\mathbf{z})_a > \sigma(\mathbf{z})_b\), preserving the relative ranking of the input logits.
- Shift Invariance: Adding an arbitrary constant \(c \in \mathbb{R}\) to every element of \(\mathbf{z}\) does not alter the output probability distribution: \[\sigma(\mathbf{z} + c \mathbf{1})_i = rac{e^{z_i + c}}{\sum_{j=1}^K e^{z_j + c}} = rac{e^c e^{z_i}}{e^c \sum_{j=1}^K e^{z_j}} = \sigma(\mathbf{z})_i\]
- Differentiability: The function is smooth and continuously differentiable everywhere on \(\mathbb{R}^K\), making it ideal for gradient-based optimization via backpropagation [2].
2. Numerical Stability
2.1 Floating-Point Hazards
Direct evaluation of \(\sigma(\mathbf{z})_i\) using standard IEEE 754 floating-point representations (such as FP32, FP16, or BF16) is numerically unstable:
- Overflow: For float32 values, \(e^z\) overflows to positive infinity whenever \(z \ge 88.72\), resulting in
NaNoutputs upon division. - Underflow: When all logits are large negative numbers, all exponentials underflow to zero, causing division by zero (\( rac{0}{0} o ext{NaN}\)).
2.2 The Subtract-Max Invariance Trick
Using the shift-invariance property, production implementations subtract the maximum element \(m = \max_j(z_j)\) from the logit vector before computing the exponential:
\[\sigma(\mathbf{z})_i = rac{e^{z_i - \max_j(z_j)}}{\sum_{k=1}^K e^{z_k - \max_j(z_j)}}\]Because \(z_i - \max_j(z_j) \le 0\) for all \(i\), every exponent is bounded in \((-\infty, 0]\), ensuring that \(e^{z_i - m} \in (0, 1]\). Overflow is mathematically eliminated, and at least one term in the denominator equals \(e^0 = 1\), preventing division by zero [3].
2.3 Log-Softmax Formulation
When computing loss functions, calculating \(\log(\sigma(\mathbf{z})_i)\) directly from Softmax outputs introduces precision loss. Deep learning frameworks implement the LogSoftmax primitive via the Log-Sum-Exp operator:
3. Temperature Scaling
Temperature scaling introduces a strictly positive hyperparameter \(T > 0\) into the logit arguments:
\[\sigma(\mathbf{z}, T)_i = rac{e^{z_i / T}}{\sum_{j=1}^K e^{z_j / T}}\]Temperature modulates the entropy (flatness vs. sharpness) of the resulting probability distribution:
- Standard Temperature (\(T = 1\)): The standard Softmax distribution.
- High Temperature (\(T o \infty\)): Logit differences are compressed (\(z_i / T o 0\)), causing the distribution to approach a uniform distribution: \[\lim_{T o \infty} \sigma(\mathbf{z}, T)_i = rac{1}{K}\]
- Low Temperature (\(T o 0^+\)): Logit differences are magnified, causing the distribution to collapse to a deterministic one-hot vector centered on the argmax logit: \[\lim_{T o 0^+} \sigma(\mathbf{z}, T)_i = egin{cases} 1 & ext{if } z_i = \max_j(z_j) \ 0 & ext{otherwise} \end{cases}\]
In post-processing, temperature scaling is also used to calibrate confidence probabilities in neural network classifiers without altering class rankings [4].
4. Derivatives and Gradient Backpropagation
4.1 Jacobian Matrix
The derivative of the \(i\)-th Softmax output \(\sigma_i\) with respect to the \(j\)-th input logit \(z_j\) is given by the quotient rule:
\[ rac{\partial \sigma_i}{\partial z_j} = \sigma_i (\delta_{ij} - \sigma_j) = egin{cases} \sigma_i (1 - \sigma_i) & ext{if } i = j \ -\sigma_i \sigma_j & ext{if } i eq j \end{cases}\]where \(\delta_{ij}\) is the Kronecker delta. In vector form, the Jacobian matrix \(J \in \mathbb{R}^{K imes K}\) is:
\[J = ext{diag}(\sigma) - \sigma \sigma^T\]4.2 Combined Softmax and Cross-Entropy Loss
In classification and autoregressive language modeling, Softmax is paired with categorical Cross-Entropy loss \(\mathcal{L} = -\sum_{k=1}^K y_k \log(\hat{y}_k)\), where \(\mathbf{y}\) is the ground-truth one-hot target vector and \(\hat{\mathbf{y}} = \sigma(\mathbf{z})\).
Applying the chain rule produces an elegant, numerically stable gradient with respect to the pre-activation logit \(z_i\):
\[ rac{\partial \mathcal{L}}{\partial z_i} = \sum_{k=1}^K rac{\partial \mathcal{L}}{\partial \hat{y}_k} rac{\partial \hat{y}_k}{\partial z_i} = \hat{y}_i - y_i\]The gradient is simply the signed residual error between the predicted probability distribution and the true target [2].
5. Applications in Modern AI Architectures
5.1 Scaled Dot-Product Attention in Transformers
In the transformer architecture [5], Softmax computes dynamic attention weights between query matrix \(Q\) and key matrix \(K\):
\[ ext{Attention}(Q, K, V) = ext{softmax}\left( rac{QK^T}{\sqrt{d_k}} ight)V\]The scaling factor \( rac{1}{\sqrt{d_k}}\) prevents dot products from growing excessively large in high dimensions, which would otherwise push the Softmax function into regions with near-zero gradients (the vanishing gradient regime).
5.2 Sparse Mixture-of-Experts Gating
In Mixture-of-Experts (MoE) architectures (such as DeepSeek-V4-Flash and Llama 4), a Softmax gating router determines which feed-forward expert networks process each token [6]:
\[G(x) = ext{softmax}( ext{TopK}(H(x), k))\]5.3 Autoregressive Decoding and Sampling
Autoregressive language models project the final hidden state \(h_t\) across the vocabulary matrix \(W_v\) to produce vocabulary logits \(\mathbf{z}_t \in \mathbb{R}^{|V|}\). Softmax converts these logits into the next-token probability distribution \(P(w_{t+1} \mid w_{1:t}) = ext{softmax}(\mathbf{z}_t / T)\), which is subsequently filtered via Top-\(p\) (nucleus) or Top-\(k\) truncation prior to stochastic sampling [7].
6. Variants and Continuous Relaxations
- Gumbel-Softmax (Concrete Distribution): A continuous, differentiable relaxation of discrete categorical sampling. Enables backpropagation through stochastic discrete choices via the reparameterization trick: \[y_i = rac{e^{(z_i + g_i) / au}}{\sum_{j=1}^K e^{(z_j + g_j) / au}}\] where \(g_i \sim ext{Gumbel}(0, 1)\) and \( au\) is a temperature parameter [8].
- Sparsemax: An alternative activation function based on Euclidean projection onto the probability simplex, yielding exact zero probabilities for low-ranked classes to produce truly sparse attention maps [9].
- Hierarchical and Adaptive Softmax: Tree-structured and cluster-based factorizations designed to reduce computational complexity from \(O(|V|)\) to \(O(\log |V|)\) or \(O(\sqrt{|V|})\) across massive vocabularies [10].
See also
- Large Language Models — Transformer architecture, vocabulary projection, and autoregressive generation.
- Autoregressive Models — Sequential probability factorization and next-token prediction.
- Frontier Models, Architectures, and Tokens — Foundation models, sparse MoE gating, and token economics.
- Deep Neural Networks — Multi-layer perceptrons, backpropagation, and loss functions.
- Reinforcement Learning — Policy gradient methods and stochastic exploration.
References
- [1] J. S. Bridle, "Probabilistic Interpretation of Feedforward Classification Network Outputs, with Relationships to Statistical Pattern Recognition," in Neurocomputing: Algorithms, Architectures and Applications, Springer, 1990, pp. 227–236.
- [2] I. Goodfellow, Y. Bengio, and A. Courville, Deep Learning, MIT Press, 2016. https://www.deeplearningbook.org
- [3] C. M. Bishop, Pattern Recognition and Machine Learning, Springer, 2006.
- [4] C. Guo, G. Pleiss, Y. Sun, and K. Q. Weinberger, "On Calibration of Modern Neural Networks," in International Conference on Machine Learning (ICML), 2017. https://arxiv.org/abs/1706.04599
- [5] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, "Attention Is All You Need," in Advances in Neural Information Processing Systems (NeurIPS), vol. 30, 2017. https://arxiv.org/abs/1706.03762
- [6] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean, "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer," in ICLR, 2017. https://arxiv.org/abs/1701.06538
- [7] A. Holtzman, J. Buys, L. Du, M. Forbes, and Y. Choi, "The Curious Case of Neural Text De-Generation," in ICLR, 2020. https://arxiv.org/abs/1904.09751
- [8] E. Jang, S. Gu, and B. Poole, "Categorical Reparameterization with Gumbel-Softmax," in ICLR, 2017. https://arxiv.org/abs/1611.01144
- [9] A. F. Martins and R. F. Astudillo, "From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification," in ICML, 2016. https://arxiv.org/abs/1602.02068
- [10] E. Grave, A. Joulin, M. Cissé, D. Grangier, and H. Jégou, "Efficient Softmax Approximation for GPUs," in ICML, 2017. https://arxiv.org/abs/1609.04309