The Softmax Bottleneck

A standard language model's output head takes a d-dimensional final hidden state and computes logits as its dot product against a vocab_size x d embedding matrix, then applies softmax. It looks like an unremarkable last step. Yang et al., 2017 (arXiv:1711.03953) showed it is actually a hard structural constraint on what the model can ever represent, no matter how it is trained, how much data it sees, or how long it runs.

Where the rank limit comes from

Imagine stacking, for every possible context c a language model might see, the true conditional log-probability distribution log P(x | c) over the whole vocabulary, into one giant matrix of shape (number_of_contexts, vocab_size). Yang et al. argue, and demonstrate empirically, that the true matrix for natural language is high rank: it takes many independent directions to describe how wildly different contexts favour different words, because "the sky is " and "the mood was " want completely different, non-overlapping shapes of distribution.

Now look at what a single softmax head actually computes. Logits are the matrix product of a d-dimensional context vector and a d-dimensional word embedding matrix. Up to the per-context normalising constant that softmax adds, the resulting log-probability matrix can have rank at most d, the hidden dimension. If the true conditional-distribution matrix needs more independent directions than d to describe, a single softmax cannot represent it exactly, as a matter of linear algebra, regardless of how well the model is trained. That is the bottleneck: an expressivity ceiling set by architecture, not by optimisation.

The fix: Mixture of Softmaxes

Yang et al.'s proposed fix, Mixture of Softmaxes (MoS), replaces the single softmax with several. From the same hidden state, compute K different d-dimensional context vectors (via K separate linear projections), run softmax on each independently to get K full probability distributions over the vocabulary, then combine them as a weighted average, with mixture weights themselves computed from the hidden state through another small softmax. Because the final distribution is a mixture of several rank-d distributions rather than the output of one, its effective rank is no longer capped at d alone. MoS measurably closed the gap to the true high-rank target and improved perplexity over a single-softmax RNN language model on standard benchmarks, at the cost of K separate forward passes through the (typically enormous) output projection.

Why this still matters for modern LLMs

Most production LLMs still ship a single softmax output head, frequently with input and output embeddings tied together (see embeddings-semantic-search) to save parameters. The softmax bottleneck is one real argument, among several, for why hidden dimension matters beyond a generic notion of "more capacity": it directly bounds the rank of the achievable output distribution. It also frames vocabulary size versus hidden dimension as a genuine expressivity tradeoff, not just a memory and compute cost. Models with very large vocabularies (100k or more tokens) paired with comparatively modest hidden dimensions sit closer to the regime this paper describes than models with smaller vocabularies and large hidden states.

When it falls down

• It is a worst-case expressivity bound, not proof of an actual deficiency. The analysis shows a single softmax cannot exactly represent an arbitrarily high-rank target; it does not prove any specific trained model is meaningfully bottlenecked in practice. Later work found the effect measured smaller in large attention-based transformers than in the small RNNs the original paper tested, plausibly because attention lets richer context information reach the final hidden state than a single recurrent state could.
• Mixture of Softmaxes is expensive exactly where it hurts most. It multiplies the cost of the output layer by K, and for a 100k-plus token vocabulary the output projection is already one of the most expensive parts of a forward pass. This is the main reason MoS is rarely used in large-scale production LLMs: the measured perplexity gain has not been shown to be worth the extra compute at that scale.
• Weight tying interacts with the bottleneck in ways the original analysis does not directly address. Forcing the same d-dimensional space to serve as both input embedding and output projection adds a further constraint on top of the rank limit, and the two effects are not simply additive.

Further reading

• Yang et al., 2017, Breaking the Softmax Bottleneck: A High-Rank RNN Language Model, arXiv:1711.03953 - the rank argument and the Mixture of Softmaxes fix.
• Press and Wolf, 2016, Using the Output Embedding to Improve Language Models, arXiv:1608.05859 - the weight-tying practice that interacts with the same output-layer rank constraint.