Breaking the Softmax Bottleneck: A High-Rank RNN Language Model

ListenThe Softmax Bottleneck: When Your Output Layer Runs Out of Rank10:25

Show notes

A deep dive into Yang, Dai, Salakhutdinov & Cohen's 2017 paper reframing language modeling as matrix factorization — and diagnosing the "Softmax bottleneck," where an output layer with embedding dimension d can never represent a high-rank true log-probability matrix. We walk through the Mixture of Softmaxes fix, the state-of-the-art perplexity numbers on Penn Treebank, WikiText-2, and 1B Word, and the ablations that show it's the structure, not the tuning, doing the work. We close with why image and video modelers should care about a structural d-versus-vocabulary mismatch. **Sources** - Breaking the Softmax Bottleneck: A High-Rank RNN Language Model

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

Yang, Dai, Salakhutdinov & Cohen formalize language modeling as a matrix factorization problem: a Softmax over hidden state $\cdot$ word embedding produces a log-probability matrix whose rank is upper-bounded by the embedding dimension $d$ . They argue natural language’s true log-probability matrix is high-rank (context depends on much more than $d$ latent factors), so Softmax with $d \ll |V|$ is expressively insufficient — the “Softmax bottleneck.” Their fix, Mixture of Softmaxes (MoS), computes $K$ separate Softmax distributions and weighted-averages them; because the log-sum-exp of Softmaxes is nonlinear in the context, MoS can produce an arbitrarily high-rank log-probability matrix at similar parameter count. MoS advanced SOTA perplexity on Penn Treebank (47.69), WikiText-2 (40.68), and beat a matched-parameter Softmax baseline by 5.6 perplexity on 1B Word.

Key claims

Learning a Softmax-based LM is equivalent to factoring the log-probability matrix $A$ as $H W^\top$ ; since $H \in \mathbb{R}^{N \times d}$ and $W \in \mathbb{R}^{V \times d}$ , the achievable rank is at most $d$ [§2.1, Proposition 1].
The set $F(A)$ of log-probability matrices consistent with the true data distribution has rank differing by at most 1 across members, so a Softmax LM with $d < \mathrm{rank}(A) - 1$ cannot represent the true distribution regardless of universal-approximator power in $H$ [§2.1, Corollary 1].
Natural language’s log-probability matrix is hypothesized to be high-rank (rank possibly at the scale of $V$ ), motivated by contextual polysemy and the difficulty of expressing all semantic bases in a few hundred dimensions [§2.2].
Mixture of Softmaxes computes $P(x|c) = \sum_{k=1}^K \pi_{c,k} \mathrm{softmax}(h_{c,k} W^\top)$ ; because $\log \sum_k \pi_k \exp(\cdot)$ is nonlinear, the resulting log-probability matrix can be arbitrarily high-rank at fixed embedding dim $d$ [§2.4].
Mixture of Contexts (mixing hidden states before Softmax) is a low-rank baseline: it reduces algebraically to a single Softmax with an averaged context and inherits the rank- $d$ ceiling [§2.5].
On Penn Treebank, AWD-LSTM-MoS + dynamic evaluation reaches test perplexity 47.69 (prior SOTA 51.1) at 22M parameters [Table 1].
On WikiText-2, AWD-LSTM-MoS + dynamic evaluation reaches test perplexity 40.68 (prior SOTA 44.3) at 35M parameters [Table 2].
On 1B Word, MoS at 113M parameters beats a matched 119M Softmax baseline by 5.6 test perplexity (37.10 vs 42.77) with no dropout tuning [Table 3].
Empirical rank measurement: on PTB validation with $|V|=10{,}000$ , the empirical rank of the log-probability matrix is 400 for Softmax ( $d=400$ ), 280 for MoC ( $d=280$ ), and 9981 (near-full) for MoS ( $d=280$ , $K=15$ ) [Table 6].
Increasing $K$ from 3 to 15 monotonically increases both empirical rank (6467 → 9981) and reduces perplexity (58.62 → 55.97); beyond 15 the rank saturates and further $K$ overfits [Table 7].
Ablation isolates MoS’s contribution from hyperparameters and extra parameters: matched-parameter MoC underperforms MoS by 1.6–2.7 perplexity on PTB/WT2, and vanilla AWD-LSTM run with MoS’s hyperparameters performs worse than its own tuned baseline [Table 5].
Character-level LM (small $|V|$ , no rank bottleneck) shows no MoS gain — an inverse experiment supporting the rank-limitation diagnosis [App. C.2].

Method

Language modeling factors the joint probability autoregressively; each conditional $P(x \mid c)$ is standardly parameterized as $\mathrm{softmax}(h_c W^\top)$ with $h_c \in \mathbb{R}^d$ and word embeddings $W \in \mathbb{R}^{V \times d}$ . Stacking log-probabilities over all contexts gives $A \in \mathbb{R}^{N \times V}$ with $A = HW^\top + \text{row shifts}$ ; matrix-rank arithmetic caps the achievable rank at $d+1$ . The paper’s contribution is (a) diagnosing this cap as an expressiveness problem when $\mathrm{rank}(A^\text{true}) \gg d$ , and (b) proposing MoS: an RNN produces $K$ context vectors $h_{c,1}, \ldots, h_{c,K}$ per position plus a prior $\pi_c \in \Delta^{K-1}$ , and the model outputs $P(x \mid c) = \sum_k \pi_{c,k} \mathrm{softmax}(h_{c,k} W^\top)$ . Since $\log\sum_k \pi_k \exp(z_k)$ is nonlinear in $z$ , the effective log-probability matrix $\log(\Pi \odot \mathrm{softmax}(HW^\top))$ is not constrained to rank $d$ . To hold parameter count constant against Softmax, MoS reduces $d$ from 400 to 280 and uses $K=15$ mixtures. All other regularizers follow Merity et al.’s AWD-LSTM recipe. The MoC baseline mixes hidden states rather than probabilities and is shown to remain rank- $d$ because $\sum_k \pi_k h_k$ collapses back into a single-context Softmax.

Results

Penn Treebank: 47.69 test perplexity (SOTA at time; -3.4 over dynamic-eval AWD-LSTM baseline) with 22M params [Table 1].
WikiText-2: 40.68 test perplexity (SOTA; -3.6 over baseline) with 35M params [Table 2].
1B Word: 37.10 test perplexity vs Softmax baseline’s 42.77 at matched params, no dropout [Table 3].
Switchboard dialog Seq2Seq: MoS beats Softmax and MoC on perplexity (32.7 vs 34.7 / 33.3) and BLEU-1–4 precision/recall [Table 4].
Rank vs perplexity is monotonic up to saturation: rank 6467 → 9981 as $K$ grows from 3 to 15, perplexity 58.62 → 55.97 [Table 7].
Ablations: matched-param MoC gets 57.55 / 65.98 PTB / WT2 (vs MoS 55.97 / 63.33); MoS hyperparameters applied to vanilla AWD-LSTM catastrophically hurt it (74.86 / 69.18) — MoS structure, not tuning, is doing the work [Table 5].
MoS wall-clock: sub-linear in $K$ ; typical 2–3× slowdown at $K \le 15$ [App. C.3].

Why it’s interesting

This is the seminal work naming the “Softmax bottleneck” — the load-bearing term for the argument that both Lost in Backpropagation: The LM Head is a Gradient Bottleneck and Universal One-third Time Scaling in Learning Peaked Distributions build on eight years later. Yang et al.’s framing is expressivity: the achievable log-probability matrix has rank $\le d$ , so if the true matrix is high-rank, Softmax cannot represent it. Godey & Artzi in 2603.10145 explicitly cite MoS as a case where the expressivity fix does not resolve the optimization bottleneck (the chain rule still injects a rank- $D$ Jacobian regardless of the output function shape) — reframing the same $d$ -vs- $V$ mismatch as a gradient-compression problem rather than a representability problem. Liu et al. in 2602.03685 go further: even with sufficient expressivity, the softmax-CE pairing on peaked targets produces a universal $t^{-1/3}$ loss decay whose slope no optimizer or data intervention moves. Together the three papers place the LM head at the center of three orthogonal scaling-law pathologies: expressivity (Yang 2017), optimization / gradient rank (Godey 2026), and asymptotic slope (Liu 2026). This 2017 paper is the historical anchor — reading it makes the ” $d \ll V$ is a structural primitive” framing on Training stability at scale concrete rather than abstract.