Breaking the softmax bottleneck: your output layer is a rank-d wall

Breaking the Softmax Bottleneck (Yang, Dai, Salakhutdinov, Cohen — CMU, ICLR 2018 oral) is a paper about a wall you cannot see from inside the model.

Deep learning does not have many negative results that matter. Most limitations turn out to be engineering — not enough data, not enough compute, the wrong optimizer. This one is linear algebra, it takes two lines to state, and it is still true of every model you use.

The argument

Take the standard output layer. A network turns the context $c$ into a hidden state $\mathbf{h}_c$, you dot it with every word embedding $\mathbf{w}_x$, and softmax the result:

Now stack every context as a row of $\mathbf{H}_\theta \in \mathbb{R}^{N \times d}$, every embedding as a row of $\mathbf{W}_\theta \in \mathbb{R}^{M \times d}$, and the true log-probabilities as $\mathbf{A} \in \mathbb{R}^{N \times M}$. Your model's logits are $\mathbf{H}_\theta \mathbf{W}_\theta^\top$, and language modelling is now the question of whether that product can equal $\mathbf{A}$.

It cannot, if $d$ is small. The rank of a product is bounded by the shared inner dimension. Whatever the network does, its logits live in a $d$-dimensional subspace.

The move that makes this a paper rather than an observation is handling the obvious objection. Softmax is invariant to adding a constant to a row, so the model doesn't have to hit $\mathbf{A}$ — it can hit any member of the family $F(\mathbf{A}) = \{\mathbf{A} + \mathbf{\Lambda}\mathbf{J}\}$ of row-shifted variants, an infinite set. Surely somewhere in an infinite set there's a low-rank one?

No. Their Property 2: any two matrices in $F(\mathbf{A})$ have ranks differing by at most 1. The entire row-shift freedom is worth one rank. So the corollary is clean:

Corollary 1 (Softmax Bottleneck). If $d < \text{rank}(\mathbf{A}) - 1$, then for any function family $\mathcal{U}$ and any parameter $\theta$, there exists a context $c$ such that $P_\theta(X \mid c) \neq P^*(X \mid c)$.

Read the quantifier: for any function family. Universal approximation doesn't help. You can make the network computing $\mathbf{h}_c$ arbitrarily deep and arbitrarily wide and it changes nothing, because the constraint is on the shape of the factorization, not on the expressiveness of the thing being factorized. All that effort is spent producing a vector that then has to squeeze through a $d$-wide waist.

The part that is not proved

The bound only bites if $\text{rank}(\mathbf{A})$ is actually large, and the paper is upfront that this is a hypothesis:

It is difficult (if possible) to rigorously prove this hypothesis since we do not have access to the true data distribution of a natural language.

The supporting intuitions are decent but soft. Language is context-dependent — "north" is followed by "korea" in a politics article and not in a U.S. history textbook. And if $\mathbf{A}$ were low rank, that would mean a few hundred basis distributions span every meaning humans express, and no one has ever found such a basis.

Neither of those is evidence. The evidence comes later, and it's better than the arguments.

Two easy fixes, priced

Once you see the bound, two fixes suggest themselves, and the paper prices both out before proposing anything.

Use an n-gram model. Non-parametric, no rank constraint, universally approximates any language. Costs $N \times M$ parameters, where $N$ — the number of contexts — is unbounded. Generalizes badly, which is why the field left it.

Raise $d$ until the bound stops binding. To express a full-rank $\mathbf{A}$ you need $d \approx M$, so the embedding matrix costs $M \times M$. Slide the widget above to a modern vocabulary and watch that number: at $|V| = 151{,}936$ it's a 23-billion-parameter output layer, for a model whose whole point was to be small. And empirically it doesn't even work — the paper notes, and everyone else had found, that pushing $d$ past a few hundred stopped helping on these benchmarks.

That is the real tension, and it is why the paper is interesting: expressiveness and generalization are in conflict at the output layer, and the naive ways to buy one spend the other.

Mixture of softmaxes

The fix is small enough to quote in full. Compute $K$ context vectors instead of one, run each through the same embedding matrix, softmax each, and average the resulting probabilities with context-dependent weights:

Here is the whole thing in the reference implementation, essentially unedited:

# one linear layer produces all K context vectors at once
self.latent = nn.Sequential(nn.Linear(nhidlast, n_experts * ninp), nn.Tanh())
self.prior  = nn.Linear(nhidlast, n_experts, bias=False)
 
latent = self.latent(output)                                  # (T*B, K*d)
logit  = self.decoder(latent.view(-1, self.ninp))             # shared W, K times
 
prior = F.softmax(self.prior(output).view(-1, self.n_experts), -1)
prob  = F.softmax(logit.view(-1, self.ntoken), -1)
prob  = prob.view(-1, self.n_experts, self.ntoken)
prob  = (prob * prior.unsqueeze(2).expand_as(prob)).sum(1)    # mix probabilities
log_prob = torch.log(prob.add_(1e-8))

Two things worth noticing in that code. The embedding matrix self.decoder is used $K$ times — MoS does not buy $K$ embedding tables, it buys $K$ readings of one table, so the parameter cost is the latent projection and nothing else. And the mix happens after the softmax, on probabilities, which is the only reason any of this works.

Because the resulting log-probability matrix is

and $\log \sum \exp$ is nonlinear, $\hat{\mathbf{A}}_{\text{MoS}}$ has no rank ceiling at all. It is a nonlinear function of $K$ rank-$d$ matrices, and nonlinear functions of low-rank matrices are generically full rank.

The trap, which is the best part

Now the near-miss. Suppose you mix the context vectors instead of the probabilities — average $\mathbf{h}_{c,k}$ with the same weights, then take one softmax. Call it mixture of contexts. It looks like the same idea, it has the same parameter count, and it is completely useless:

which is the original softmax with a different $\mathbf{h}$. Rank still bounded by $d$. Mixing in feature space makes the function family richer and leaves the ceiling exactly where it was.

MoC exists in the paper as a control, and it is the sharpest instrument in it: same parameters, same layers, same hyperparameters, one design decision moved, and the theory says one should work and the other shouldn't.

There is also a footnote in the related-work section that has aged into the most consequential sentence in the paper:

Although Shazeer et al. (2017) name their architecture as MoE, it is not a standard MoE and should be classified as MoC under our terminology.

Shazeer et al. 2017 is the sparsely-gated mixture-of-experts layer — the direct ancestor of every MoE LLM shipping today, from Switch Transformer to DeepSeek-V3 to Kimi. Under this paper's taxonomy all of them are mixture-of-contexts. They mix in feature space. They make the function family enormously richer and they do not touch the rank of the output layer by one.

That is not a criticism of MoE — sparse experts are solving conditional computation, not expressiveness — but it does mean the thing people reach for when they want "more capacity" is provably not the thing that lifts this particular ceiling.

The evidence

Two measurements, and the first one is the kind that could have embarrassed everybody.

They compute the empirical log-probability matrix on PTB and estimate its rank. Softmax with $d = 400$ measures 400. MoC with $d = 280$ measures 280. Not near the bound — the bound, to the digit. MoS with the same 280 dimensions measures 9,981 out of a possible 10,000.

Then the dose-response. Sweep $K$ from 3 to 20 and rank climbs — 6,467, 8,930, 9,973 — with perplexity falling alongside it. At $K = 15$ rank has saturated at 9,981 and perplexity is at its best. At $K = 20$ rank does not move, because there is nothing left to buy, and perplexity gets worse.

That reversal is what makes the sweep an argument. A "more parameters help" story predicts a monotone curve. A "mixtures buy rank until rank runs out, and then you're just overfitting" story predicts a curve that turns exactly where rank saturates. It turns exactly where rank saturates.

expand

Counting non-zero singular values is a roundoff-sensitive way to measure rank, so they plot the spectrum instead and the picture is unambiguous. Softmax and MoC dump ~96% of their normalized singular values below $10^{-9}$; MoS's are spread from $10^{-5}$ upward. Same conclusion, no thresholding decision required.

A third check, in the appendix: expected pairwise KL divergence between next-token distributions at different contexts — how much the model's prediction actually changes when the context changes. Softmax 4.763, MoC 4.864, MoS 5.284 on PTB test.

Three controls that turn it into a mechanism

Any of the above is consistent with "MoS is a good regularizer and the rank story is decoration." The paper runs the experiments that separate those.

Ablation. MoC with matched everything is worse than MoS on both datasets — and on WikiText-2 it is worse than the plain AWD-LSTM baseline it was built from (65.98 against 65.40). So mixing per se isn't the win. Separately, training the baseline with MoS's hyperparameters is a disaster (74.86 against 58.95 on PTB), which rules out "they just found better hyperparameters."

Regularization control. On the 1B Word dataset, where overfitting is unlikely and no dropout is used at all: Softmax reaches 41.47 train / 42.77 test; MoS reaches 36.39 train / 37.10 test. MoS's training perplexity is 5.08 points lower. If the gain were regularization, training perplexity would have gone up, not down. (Their word for the generalization gaps is "similar"; strictly, MoS's is a bit narrower — 0.71 against 1.30 points, or ratios of 1.020 and 1.031 — which if anything strengthens the reading.)

The inverse experiment, which is the one I'd point at. If the mechanism really is the rank bound, then in a setting where the bound cannot bind, MoS should do nothing. Character-level language modelling is exactly that setting: $\text{rank}(\mathbf{A}) \leq |V| \approx 27$, and $d$ is in the hundreds, so there is no bottleneck to break. On text8, at matched parameter counts:

Identical. A method that improves everything improves this too; a method that breaks a specific bound does nothing when the bound is absent. Papers that predict their own null results are rare, and this one went and measured it.

What it won

State of the art at the time, at comparable model size:

22M parameters on PTB against 24M baselines, and 35M on WT2 against 33M — so slightly under on one and slightly over on the other, which is the honest way to read "comparable." The 1B Word row is the one that ages best: 5.67 points on a dataset large enough that regularization tricks aren't doing the work, against a plain 2-layer LSTM softmax at 119M parameters, with hyperparameters they admit they never tuned.

They also bolt MoS onto a Seq2Seq decoder for dialogue on Switchboard and it wins on perplexity and on every BLEU precision and recall figure, which is a reasonable check that this is about context-dependent distributions in general rather than about language-modelling benchmarks in particular.

So why isn't it in your model

Cost. $K$ softmaxes means $K$ passes over the vocabulary. Measured at matched batch size it's 1.9× on PTB, 2.5× on WikiText-2, 3.8× on 1B Word; at the settings where each model does its best, 2.8× and 6.4× on one GPU. Sub-linear in $K$ thanks to GPU matmul efficiency, but "sub-linear" still means two to three times the training cost, and that is before you consider that the output layer is now $K$ times the memory.

Then scale the setting. PTB has a 10,000-token vocabulary. A modern model has 150,000–200,000, and the vocabulary projection is already one of the most expensive tensors in the network — it's why chunked-and-fused cross-entropy kernels exist at all. Fifteen softmaxes over 200,000 logits per position is not a rounding error, it is the model.

So the field made a choice, and the choice was not obviously wrong: buy quality with data and depth, where the cost curve is friendlier, and leave the output layer alone.

What happened to the idea

It didn't disappear so much as fragment into a small literature that no one reads together.

Sigsoftmax (Kanai et al., 2018) re-derives the bottleneck and argues the culprit is specifically the exponential in softmax, proposing a cheaper output nonlinearity that also escapes the rank limit — the same diagnosis, a one-softmax fix.

Stolen Probability (Demeter, Kimmel, Downey, 2020) finds a different consequence of the same geometry: embeddings in the interior of the convex hull of the embedding cloud can never be the argmax, no matter the context, so certain words are structurally unpredictable.

And the honest counterweight, Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practice (Grivas, Bogoychev, Lopez, 2022), goes looking for that failure in real systems: 13 of 150 public models have unargmaxable tokens, and they are rare enough not to matter. Which is a useful correction — a bound being real is not the same as a bound costing you anything — though note it tests one specific symptom, the argmax-unreachable token, not the broader claim that the expressible distribution is lower-rank than the one you want.

The bottleneck's authors moved on to Transformer-XL and XLNet, and Zhilin Yang went on to found Moonshot AI, whose Kimi K3 ships a 7,168-dimensional hidden state against a 163,840-token vocabulary — a ratio of 23, in a model from the person who wrote the paper about the ratio. (Kimi K2 sits in the chart below at the same two numbers.)

The ratio, today

I pulled hidden_size and vocab_size from published configs to see whether nine years of scaling relaxed the geometry. It didn't. The paper's own bottlenecked setup — the one where breaking the bound was worth 3.6 perplexity — had $|V|/d = 25$. Almost every open model checked is above that, several by a factor of six.

Anchoring on the paper's own numbers: from PTB's 10,000 tokens to a modern 151,936, vocabulary grew about 15×, while $d$ went from 400 to roughly 4,096 — about 10×. And the mismatch lands hardest on small models, which inherit a large tokenizer from their big siblings and get a fraction of the width to read it with. Qwen3-0.6B carries the same 151,936-token vocabulary as Qwen3-32B through 1,024 dimensions instead of 5,120.

I want to be careful about what that does and doesn't show. $|V|/d$ is geometry, not severity. Nobody knows $\text{rank}(\mathbf{A})$ for natural language; a 2026 model has representations no 2017 LSTM had; and the bottleneck may be so far from binding at this scale that it costs nothing measurable. The claim is narrower and, I think, harder to argue with: the constraint everyone stopped worrying about is tighter now than when they stopped, and the experiment that would tell us what it costs — an MoS-style rank measurement on a modern LLM's output layer — appears not to have been run.

Why I keep coming back to it

The full-bandwidth transformer paper argues that the feedback path between decoding steps is one token wide — $\log_2|V|$ bits — while the hidden state that produced that token gets thrown away, and that chain-of-thought is partly a workaround for the narrow pipe.

That is the same argument, at the other end of the model. Both say: the network is not the limiting factor, the interface is. One narrow shape sits between a rich internal state and the thing you actually want, and everything upstream is spending its capacity on getting through it.

The best thing about the softmax bottleneck paper isn't the mixture of softmaxes, which nobody uses. It's that it demonstrated the move: take the part of the architecture that's so standard nobody writes it down, ask what it structurally cannot do, and then — this is the rare part — go and measure whether it costs anything. Softmax 400. MoC 280. Character-level, no change.

Most papers proposing an architecture would have stopped after the perplexity table.