Full Softmax

Full softmax (also called the standard softmax or exact softmax) is the softmax computation that calculates a probability for every possible output class in a classification model, producing an exact probability distribution in which all class probabilities sum to exactly 1. Google's Machine Learning Crash Course defines it directly: "Full softmax is the softmax we've been discussing; that is, softmax calculates a probability for every possible class." ^[15] It is the complete computation of the softmax function over the entire output vocabulary, in contrast to candidate sampling methods (such as sampled softmax, noise contrastive estimation, or negative sampling) that calculate probabilities for the positive labels plus only a random sample of the negative labels to save computation. ^[15][16]

Full softmax guarantees mathematically precise probabilities, but its cost grows linearly with the number of classes K, so it becomes a computational bottleneck when that number is very large, as in language models with vocabularies of tens or hundreds of thousands of tokens. As Google's Crash Course puts it, full softmax is "fairly cheap when the number of classes is small but becomes prohibitively expensive when the number of classes climbs." ^[15] Because of this, candidate sampling is used as a training-time optimization, while full softmax is typically computed at inference time to obtain a properly normalized distribution. ^[16][11]

ELI5 (Explain like I'm 5)

Imagine you have a huge bag with thousands of different colored marbles. You want to figure out the chance of picking each color. Full softmax is like carefully weighing every single marble in the bag, one by one, and then calculating the exact chance for each color. This gives you a perfect answer, but if the bag has a million marbles, it takes a really long time. Some faster methods skip most of the marbles and only look at a few of them to get a "good enough" guess. Full softmax never takes shortcuts; it always checks every marble.

What is full softmax?

Full softmax is the standard form of the softmax function applied across all K output classes. The word "full" distinguishes it from the family of approximate, sampling-based, and tree-based methods that compute the softmax over only part of the output space. In a full softmax layer, the model evaluates the score (logit) of every class, exponentiates each one, and normalizes by the sum over all classes, so every class receives a strictly positive probability and the probabilities sum to exactly 1. ^[13]

Google's Machine Learning glossary frames the same idea by contrast: full softmax "calculates a probability for all the positive labels" and all negative labels, whereas candidate sampling "calculates a probability for all the positive labels, using, for example, softmax, but only for a random sample of negative labels." ^[15][16] In multi-class problems with a small number of classes (for example, ten classes in digit recognition), this distinction barely matters because computing the full softmax is cheap. The distinction becomes important only when the number of classes climbs into the tens of thousands or more.

Mathematical definition

The softmax function was first used in statistical mechanics as the Boltzmann distribution, appearing in foundational work by Ludwig Boltzmann in 1868 and formalized by Josiah Willard Gibbs in 1902. In decision theory, R. Duncan Luce applied the axiom of independence of irrelevant alternatives to formulate what is now known as Luce's choice axiom (1959). ^[14] The modern term "softmax" was coined by John S. Bridle in two 1989 conference papers, where he described it as a "differentiable generalisation of the winner-take-all operation" that "preserves the rank order of its input values." ^[1]

Given an input vector $\mathbf{z} = (z_1, z_2, \ldots, z_K)$ of raw scores (called logits) from the final layer of a neural network, the full softmax function computes the probability for each class i as:

$\sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}}$

where K is the total number of classes, e is Euler's number (approximately 2.71828), and the denominator sums the exponentials of all logit values. The output is a vector of K probabilities that are all positive and sum to exactly 1. ^[13]

For example, given an input vector (1, 2, 3, 4, 1, 2, 3), the full softmax yields approximately (0.024, 0.064, 0.175, 0.475, 0.024, 0.064, 0.175). The largest input value (4) receives the highest probability (0.475), while the relative ordering of all inputs is preserved.

Temperature scaling

The softmax function can be parameterized with a temperature value T (or equivalently, an inverse temperature parameter $\beta = 1/T$):

$\sigma(\mathbf{z})_i = \frac{e^{z_i / T}}{\sum_{j=1}^{K} e^{z_j / T}}$

The temperature controls the "sharpness" of the output distribution:

Scroll sideways for more →

Temperature scaling is widely used in large language models such as GPT and LLaMA to control the randomness of text generation. Lower temperatures produce more deterministic and repetitive text, while higher temperatures increase diversity and creativity at the cost of coherence. Temperature scaling is also used for model calibration, where a single scalar T is learned on a validation set to make predicted probabilities better reflect true likelihoods.

Key properties

The full softmax function has several mathematical properties that make it useful in machine learning:

Scroll sideways for more →

Gradient and backpropagation

The partial derivative of the softmax output with respect to input logits is:

$\frac{\partial \sigma_i}{\partial z_j} = \sigma_i (\delta_{ij} - \sigma_j)$

where $\delta_{ij}$ is the Kronecker delta (1 when $i = j$, 0 otherwise). When paired with the cross-entropy loss function, the gradient simplifies to:

$\frac{\partial L}{\partial z_i} = \sigma_i - y_i$

where $y_i$ is the one-hot encoded true label. This clean gradient expression (predicted probability minus true label) is one of the reasons the softmax plus cross-entropy combination is so widely used in classification tasks. ^[13] The simplicity of this gradient also makes the backward pass computationally efficient.

Why is full softmax expensive?

The "full" in full softmax refers to the fact that the denominator requires summing over all K output classes. This computation has time complexity $O(K)$ for each input example, which is acceptable when K is small (e.g., 10 classes in digit recognition or 1,000 classes in ImageNet). However, in many real-world applications, K can be extremely large: ^[11]

Scroll sideways for more →

For a language model with a vocabulary of 100,000 tokens and a hidden dimension of 2,048, the output projection matrix alone contains approximately 200 million parameters. At each training step, computing the full softmax requires a matrix-vector multiplication of size $(K \times d)$, followed by exponentiation and normalization across all K values. When K is in the hundreds of thousands, this computation dominates the training cost and can account for a majority of the total floating-point operations per step. This is exactly the regime in which Google's Crash Course warns that full softmax "becomes prohibitively expensive when the number of classes climbs." ^[15]

The memory cost is also significant. The output embedding matrix of size $K \times d$ must be stored in GPU memory, and the full vector of K softmax probabilities must be materialized for each example in the batch during both the forward and backward passes.

What is candidate sampling?

Candidate sampling is the standard alternative to full softmax during training when the number of output classes is large. Google's Machine Learning glossary defines it as "a training-time optimization that calculates a probability for all the positive labels, using, for example, softmax, but only for a random sample of negative labels." ^[16] Where full softmax evaluates every class, candidate sampling evaluates only the true (positive) classes plus a small randomly chosen subset of the remaining (negative) classes. As the glossary illustrates, for an example labeled both beagle and dog, candidate sampling computes loss terms for the beagle and dog outputs together with only a random subset of the other classes (such as cat, lollipop, and fence) rather than all of them. ^[16]

The motivation is efficiency. TensorFlow's documentation describes candidate sampling as a class of algorithms that "speed up training" of models with thousands or millions of output classes by considering only a small randomly chosen subset of contrastive candidates for each batch of training examples. Common candidate-sampling algorithms include sampled softmax, noise contrastive estimation (NCE), negative sampling, and sampled logistic regression. ^[16][11] These methods change the per-step cost from $O(K)$ to roughly $O(k)$, where k is the number of sampled candidates (often on the order of tens to a few thousand) rather than the full vocabulary.

A crucial caveat is that these sampling-based shortcuts are training-time only. Sebastian Ruder's survey notes that "sampling-based approaches are only useful at training time, while during inference the full softmax still needs to be computed to obtain a normalised probability." ^[11] In other words, the common pattern is to train with an approximation (candidate sampling) and then evaluate with the exact full softmax so that the final probabilities are correctly normalized. ^[11][16]

Numerical stability

A naive implementation of the softmax function is prone to numerical overflow and underflow. When logit values are large and positive, computing e^(z_i) can exceed the range of floating-point representation, producing infinity. When logit values are large and negative, e^(z_i) underflows to zero, which can cause division-by-zero errors. ^[10]

The standard solution is the log-sum-exp trick (also called the "exp-normalize trick" or "safe softmax"). Before computing the exponentials, the maximum value m = max(z_1, ..., z_K) is subtracted from every logit:

$\sigma(\mathbf{z})_i = \frac{e^{z_i - m}}{\sum_{j=1}^{K} e^{z_j - m}}$

This transformation is valid because the softmax is translation-invariant. After the shift, the largest exponent is e^0 = 1, which prevents overflow. The corresponding log-softmax can be computed as:

$\log \sigma(\mathbf{z})_i = (z_i - m) - \log \sum_{j=1}^{K} e^{z_j - m}$

All major deep learning frameworks, including PyTorch, TensorFlow, and JAX, implement this numerically stable version by default. ^[10] In practice, combining the softmax and cross-entropy into a single fused operation (such as PyTorch's CrossEntropyLoss or TensorFlow's tf.nn.softmax_cross_entropy_with_logits) provides even better numerical stability by avoiding explicit computation of the intermediate probabilities.

What are the alternatives to full softmax?

Because full softmax becomes expensive for large output spaces, researchers have developed numerous approximation techniques. These generally fall into two categories: softmax-based approaches that restructure the computation, and sampling-based (candidate-sampling) approaches that avoid computing the full normalization constant. ^[11]

Softmax-based approaches

Hierarchical softmax

Hierarchical softmax (H-Softmax), introduced by Morin and Bengio in 2005, replaces the flat output layer with a binary tree structure in which each word is a leaf node. ^[2] Instead of computing probabilities over all K classes, the model traverses a path from the root to the target leaf, making a binary decision at each internal node. This reduces the computational complexity from $O(K)$ to $O(\log_2 K)$.

Using a Huffman tree (as in Word2Vec) assigns shorter paths to more frequent words, further reducing the average computation. Hierarchical softmax achieves training speedups of 25x to 100x compared to full softmax. ^[11] However, it has drawbacks: the tree structure introduces a dependency on the quality of the word hierarchy, and at test time the full distribution over all words must still be computed when ranking is needed. It also tends to perform poorly on small vocabularies.

Adaptive softmax

Adaptive softmax, proposed by Grave et al. in 2017, exploits Zipf's law (the observation that a small fraction of words account for most of the probability mass in natural language). It partitions the vocabulary into clusters based on word frequency, assigning smaller projection dimensions to infrequent word clusters. ^[6] This approach is specifically optimized for GPU computation, taking advantage of efficient matrix-matrix operations. Experiments on benchmarks such as EuroParl and the One Billion Word dataset showed speedups of 2x to 10x over full softmax with minimal loss in perplexity. ^[6]

Differentiated softmax

Differentiated softmax (D-Softmax) allocates varying embedding dimensions to words based on their frequency. Frequent words receive larger embeddings, while rare words receive smaller ones. Unlike hierarchical softmax, the speedup persists at test time because fewer parameters are needed overall. The trade-off is that rare words may be modeled less accurately due to their smaller representation. ^[11]

Sampling-based approaches

Noise contrastive estimation (NCE)

Noise contrastive estimation reformulates the softmax computation as a binary classification problem. Instead of computing probabilities over all K classes, the model learns to distinguish between true data samples and artificially generated noise samples drawn from a known distribution. For each training example, only k noise samples are drawn (typically k = 25), and a logistic regression classifier is trained on the k + 1 samples. The computational cost becomes $O(k)$ per example rather than $O(K)$. Gutmann and Hyvarinen (2010) proved that as the number of noise samples increases, the NCE gradient converges to the gradient of the full softmax. ^[4] NCE achieves speedups of 8x to 45x during training. ^[11]

Negative sampling

Negative sampling, popularized by Mikolov et al. (2013) in the Word2Vec system, is a simplified variant of NCE. ^[3] It drops the normalization constant entirely and uses a simplified objective function. For each positive (target) word, k negative words are sampled from a noise distribution (often the unigram distribution raised to the 3/4 power). While negative sampling lacks the asymptotic consistency guarantees of NCE (meaning it does not converge to the true softmax distribution), it works well in practice for learning word embeddings. ^[3] It is generally not recommended for language modeling tasks where accurate probability estimates are needed. ^[11]

Importance sampling

Importance sampling approximates the softmax gradient by sampling from a proposal distribution Q rather than summing over all classes. The samples are reweighted by importance weights P(w)/Q(w) to correct for the sampling bias. Adaptive importance sampling, which learns a mixture of bigram and unigram distributions during training, has achieved speedups of up to 100x. ^[11] However, importance sampling can diverge if the proposal distribution is a poor match for the true distribution.

Comparison of softmax methods

Scroll sideways for more →

The softmax bottleneck

Beyond the computational cost, Yang et al. (2018) identified a theoretical limitation called the softmax bottleneck. They formulated language modeling as a matrix factorization problem and showed that when the true data distribution has a higher rank than the hidden dimension of the model, a single softmax layer cannot represent the distribution exactly. ^[8] In other words, the expressiveness of softmax-based language models is fundamentally limited by the rank of the context-to-logit mapping.

To address this, they proposed Mixture of Softmaxes (MoS), which introduces discrete latent variables and computes the output distribution as a weighted average of multiple softmax distributions. MoS achieved state-of-the-art perplexity on Penn Treebank (47.69) and WikiText-2 (40.68) at the time of publication. ^[8] Later work by Chang and McCallum (2022) further showed that the softmax bottleneck makes language models unable to represent multi-mode word distributions. ^[12]

Sparse alternatives to softmax

While full softmax always produces dense probability distributions (every class receives nonzero probability), several alternatives have been proposed that produce sparse outputs:

Sparsemax, introduced by Martins and Astudillo at ICML 2016, projects the input onto the probability simplex in a way that can assign exactly zero probability to some classes. ^[5] This produces more interpretable and selective distributions, which is useful in attention mechanisms where focusing on a small number of inputs is desirable. Sparsemax achieves similar performance to softmax on natural language inference tasks while producing more compact attention patterns. ^[5]

Entmax (alpha-entmax) generalizes both softmax and sparsemax as a continuous family parameterized by alpha. When alpha = 1, entmax reduces to softmax. When alpha = 2, it reduces to sparsemax. Intermediate values allow fine-grained control over the degree of sparsity.

Gumbel-softmax

The Gumbel-softmax trick, introduced independently by Jang et al. (2017) and Maddison et al. (2017), provides a differentiable approximation to sampling from categorical (discrete) distributions. ^[7] Standard softmax outputs can be used to define categorical distributions, but sampling from them is not differentiable, which prevents the use of the reparameterization trick for backpropagation.

Gumbel-softmax addresses this by adding Gumbel noise to the logits before applying temperature-scaled softmax:

$y_i = \frac{\exp((\log \pi_i + g_i) / \tau)}{\sum_{j=1}^{K} \exp((\log \pi_j + g_j) / \tau)}$

where $\pi_i$ are the category probabilities, $g_i$ are independent samples from the Gumbel(0, 1) distribution, and $\tau$ is a temperature parameter. As $\tau \to 0$, the Gumbel-softmax output approaches a one-hot vector (a true categorical sample). As $\tau \to \infty$, it approaches a uniform distribution. ^[7] This technique is widely used in variational autoencoders and other generative models that require differentiable discrete sampling.

Softmax in attention mechanisms

In transformer architectures, the softmax function is applied at every layer of the self-attention mechanism to compute attention weights. Given query Q, key K, and value V matrices, the attention output is:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$

Here the softmax is applied over the sequence dimension (not the vocabulary dimension), so the number of classes K equals the sequence length. For long sequences, this becomes a bottleneck because the attention matrix is of size (sequence_length x sequence_length), leading to quadratic memory and computation costs.

FlashAttention, introduced by Dao et al. in 2022, addresses this by tiling the softmax computation to avoid materializing the full attention matrix in GPU high-bandwidth memory (HBM). It uses an online softmax algorithm that computes the normalization constant incrementally, block by block, achieving 2x to 4x speedups and 10x to 20x memory savings compared to standard attention implementations. ^[9] FlashAttention computes the exact same result as full softmax attention; it is a hardware-aware optimization, not an approximation. ^[9]

Applications of full softmax

Full softmax is used across many areas of machine learning and artificial intelligence:

Multi-class classification. In tasks like image classification (e.g., classifying images into one of 1,000 ImageNet categories), the full softmax is applied at the output layer to produce class probabilities. Combined with cross-entropy loss, it forms the standard training objective for classification networks such as ResNet, VGG, and Vision Transformers.

Language modeling. Large language models use full softmax to compute the probability distribution over the next token at each position. Despite the large vocabulary sizes, modern hardware and optimizations (such as mixed-precision training and tensor parallelism) have made full softmax feasible even for vocabularies exceeding 100,000 tokens.

Machine translation. Encoder-decoder models for machine translation apply full softmax at each decoding step to select the next word in the target language. Beam search and other decoding strategies rely on the exact probability values produced by full softmax.

Reinforcement learning. In policy-based methods, the softmax function converts action-value estimates (Q-values) into a policy (a probability distribution over actions). This is known as Boltzmann exploration or softmax exploration, where the temperature parameter controls the trade-off between exploration and exploitation. Higher temperatures encourage more exploration by making the policy more uniform, while lower temperatures make the agent more likely to select the action with the highest estimated value.

Recommender systems. Softmax is used in recommendation systems to convert item scores into selection probabilities, especially in session-based and sequential recommendation models. Candidate sampling is especially common here because the catalog (the number of output classes) can reach into the millions. ^[16]

Implementations in deep learning frameworks

All major deep learning frameworks provide optimized implementations of the full softmax:

Scroll sideways for more →

TensorFlow additionally exposes candidate-sampling losses, such as tf.nn.sampled_softmax_loss and tf.nn.nce_loss, that approximate the full softmax during training of large-vocabulary models. ^[16] The fused softmax-cross-entropy functions are preferred in practice because they are more numerically stable (they operate in log space internally) and more computationally efficient (they avoid materializing the intermediate probability vector).

Summary

Full softmax remains the default choice for classification and language modeling tasks where exact probability distributions are needed and the number of output classes is manageable. Its mathematical properties (differentiability, normalization, translation invariance) make it well suited for gradient-based optimization. For very large output spaces, candidate-sampling and other approximate methods offer significant speedups at the cost of some accuracy, and a common pattern is to train with such an approximation while computing the full softmax at inference time. ^[11][16] The choice of method depends on the specific application: full softmax for exact probabilities, hierarchical softmax or adaptive softmax for large-vocabulary training, and negative sampling for learning word embeddings.

See also

• Softmax function
• Candidate sampling
• Hierarchical softmax
• Noise contrastive estimation
• Negative sampling
• Cross-entropy
• Logistic regression

References

1. Bridle, J. S. (1990). "Probabilistic Interpretation of Feedforward Classification Network Outputs, with Relationships to Statistical Pattern Recognition." *Neurocomputing: Algorithms, Architectures and Applications*. Springer. ↩
2. Morin, F. and Bengio, Y. (2005). "Hierarchical Probabilistic Neural Network Language Model." *Proceedings of AISTATS*. ↩
3. Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., and Dean, J. (2013). "Distributed Representations of Words and Phrases and their Compositionality." *Advances in Neural Information Processing Systems (NeurIPS)*. ↩
4. Gutmann, M. U. and Hyvarinen, A. (2010). "Noise-contrastive estimation: A new estimation principle for unnormalized statistical models." *Proceedings of AISTATS*. ↩
5. Martins, A. and Astudillo, R. (2016). "From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification." *Proceedings of ICML*. ↩
6. Grave, E., Joulin, A., Cisse, M., Grangier, D., and Jegou, H. (2017). "Efficient softmax approximation for GPUs." *Proceedings of ICML*. ↩
7. Jang, E., Gu, S., and Poole, B. (2017). "Categorical Reparameterization with Gumbel-Softmax." *Proceedings of ICLR*. ↩
8. Yang, Z., Dai, Z., Salakhutdinov, R., and Cohen, W. W. (2018). "Breaking the Softmax Bottleneck: A High-Rank RNN Language Model." *Proceedings of ICLR*. ↩
9. Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Re, C. (2022). "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." *Advances in Neural Information Processing Systems (NeurIPS)*. ↩
10. Blanchard, P. and Higham, D. J. (2021). "Accurately Computing the Log-Sum-Exp and Softmax Functions." *IMA Journal of Numerical Analysis*, 41(4), 2311-2330. ↩
11. Ruder, S. (2016). "Approximating the Softmax." Blog post. https://www.ruder.io/word-embeddings-softmax/ ↩
12. Chang, H. and McCallum, A. (2022). "Softmax Bottleneck Makes Language Models Unable to Represent Multi-mode Word Distributions." *Proceedings of ACL*. ↩
13. Goodfellow, I., Bengio, Y., and Courville, A. (2016). *Deep Learning*. MIT Press. Chapter 6.2.2. ↩
14. Luce, R. D. (1959). *Individual Choice Behavior: A Theoretical Analysis*. John Wiley and Sons. ↩
15. Google. "Neural networks: Multi-class classification" and "Machine Learning Glossary: full softmax." Google Machine Learning Crash Course / Glossary. https://developers.google.com/machine-learning/crash-course/neural-networks/multi-class ↩
16. Google. "Machine Learning Glossary: candidate sampling." Google for Developers. https://developers.google.com/machine-learning/glossary ; TensorFlow, "Candidate Sampling." https://www.tensorflow.org/extras/candidate_sampling.pdf ↩