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.