The Softmax Bottleneck and What It Limits

The last operation in a language model is a matrix multiply followed by a softmax, and that shape imposes a ceiling on which probability distributions the model can produce at all — a limit no amount of extra depth or data can lift. The argument is short and entirely linear-algebraic.

The final layer as a matrix product

Collect N different contexts — every prefix the model might ever be asked to continue. Each produces a hidden vector of width d, so stack them into a matrix H of shape N × d. The output embedding matrix W has shape V × d, where V is the vocabulary size. The logits for all contexts at once are:

  logits = H W^T        shape (N x V)

For a published 7B-class model:  d = 4096, V = 32000, N unbounded.

Softmax turns each row into a distribution. Because log-softmax is the logits minus a per-row constant (the log-sum-exp), the matrix of log-probabilities the model produces is:

  A_model = H W^T - c 1^T

where c is a column of per-row constants and 1 is a column of ones.

The rank argument

The rank of a product is at most the smaller of the two ranks, so H W^T has rank at most d — here, at most 4096. Subtracting the rank-1 term c1^T can add at most one, so:

  rank(A_model) <= d + 1 = 4097

while the matrix of TRUE conditional log-probabilities A* has, in
principle, rank up to min(N, V) = 32000.

This is Yang et al.’s argument from “Breaking the Softmax Bottleneck: A High-Rank RNN Language Model” (2018). Their claim is empirical as well as algebraic: natural language’s true conditional log-probability matrix appears to be high-rank, higher than the hidden widths in common use. If that is right, then no setting of the weights makes the model match every conditional distribution exactly. The gap is not a training failure; it is outside the family of functions the architecture can express.

Notice what the bound depends on. Not depth — stacking more blocks changes H but not its width. Not data. Not parameter count as such. Only d, the width of the residual stream, and the parameterisation of the output layer.

There is a tension here with a cost pressure pulling the other way. The output matrix has V × d entries, so for a 32000-token vocabulary at d = 4096 it holds 131 million parameters, and vocabularies have been growing — larger vocabularies compress text into fewer tokens, which reduces the number of forward passes an answer needs. Growing V while holding d fixed widens the gap between the rank the model can express and the rank the vocabulary could in principle demand. It is a small effect at current sizes, and it is the direction the bound cares about.

What it actually limits

The clean statement is: the log-probability vectors the model can produce, across all contexts, live in a subspace of dimension at most d + 1 out of V. So there are combinations of relative token probabilities that are simply unreachable, and pushing one context toward a target distribution drags others along.

• Weight tying makes it sharper. Many models tie the input and output embedding matrices (Press and Wolf, 2017), which saves V × d parameters and forces one geometry to serve both jobs. The rank bound is unchanged; the flexibility within it is reduced.
• It bounds logit-level steering. Anything operating on the output distribution — constrained decoding, logit bias, classifier heads reading logprobs — works inside this subspace. You can renormalise and mask what the model produced; you cannot conjure a distribution it cannot represent.
• It is a reason width matters. When a model family grows, d grows along with depth. Part of what widening buys is headroom on exactly this bound.

What can be done about it

Yang et al. proposed Mixture of Softmaxes: compute K different hidden vectors per context, softmax each against the same output matrix, and take a weighted mixture of the K distributions. Because the mixture is formed in probability space and then logged, the resulting log-probability matrix is not constrained to rank d + 1 — a log-of-sum-of-exponentials is not a linear function of anything. The cost is K times the output computation, on the layer that already involves the largest matrix in the model, which is why it did not become standard.

The path the field actually took was cruder and cheaper: make d large. A 4096-wide stream with a 32000-token vocabulary is a far less binding constraint than the 300-to-1000-wide models the original paper studied. Sigmoid-based and other output parameterisations have been proposed since; none has displaced the plain softmax in production models.

Whether it matters to you

Honestly: not directly, and a page that oversold it would be doing you a disservice. Nothing you will diagnose in an application traces back to the rank of the output layer. Current models are far from being limited by this on any task where they are also limited by knowledge, reasoning or context.

It earns its place for a different reason. It is the cleanest example in the whole architecture of a hard ceiling — a thing the model cannot do, provable in three lines, unaffected by scale along every axis except one. Most claims about what language models “cannot do” are claims about training or about current performance. This one is a claim about the function class, and knowing the difference is the useful part.