Lost in Backpropagation: The LM Head is a Gradient Bottleneck

TL;DR

The softmax bottleneck in LM heads (hidden dim $D \ll$ vocabulary size $V$ ) is not just an expressivity limitation; it is an optimization bottleneck. Backpropagating $V$ -dimensional logit gradients through a rank- $D$ linear layer destroys 95–99% of the gradient norm before it reaches the backbone, transferring energy from the informative top components to the tail as random noise. The authors prove this with an Eckart–Young–Mirsky bound on the update direction, validate it empirically on GPT-2, Pythia, Llama-3, and Qwen3-Base, and show in a controlled 2B-parameter pretraining sweep (same Transformer stack, varying low-rank LM-head structure) that larger effective $D$ converges 1.75× faster and gains +0.55 average downstream score at the final cooldown. Frames a new design axis for LMs at scale: better-suited logit prediction modules that preserve gradient flow.

Key claims

The softmax-bottleneck rank constraint does not limit top-1 token probabilities — for any context and target distribution, there exist representations achieving arbitrarily-close top-1 mass even when $D \ll V$ — so the bottleneck must be analyzed as an optimization phenomenon, not an expressivity one [§2.2, Prop. 2.4].
Under full-batch GD, the actual logit update $\widetilde{\Delta Z}$ from one step through the LM head has rank at most $D$ , while the logit gradient $G$ has rank close to $V$ in practical setups; by Eckart–Young–Mirsky the residual is bounded below by $\sum_{i=D+1}^{r} \sigma_i^2(G)$ , i.e. all tail singular values of $G$ are inevitably lost [§2.3, Eq. 6, Prop. 2.6].
The high-rank property of $G$ persists under SGD: when the model approaches the empirical distribution, the per-batch logit gradient has rank close to the number of contexts in the batch with unique observed continuations, so mini-batching does not relax the bottleneck [§2.3.1, Prop. 2.7, Cor. 2.8].
Expressivity-improving softmax alternatives (Mixture of Softmaxes, Ganea et al.) do not resolve the optimization bottleneck — for any output function $f(Z)$ the chain rule injects a rank- $D$ Jacobian, so the same gradient compression occurs regardless of $f$ ‘s form [§2.3, Eq. 9].
Empirically across GPT-2, Pythia, Llama-3, Qwen3-Base, 95–99% of the logit gradient norm projects onto the null space of $W$ (the LM head); the projected and original gradients are only mildly aligned (cosine 0.1–0.2), and the fraction of norm lost diminishes with hidden dimension [§3.3, Fig. 6, App. D].
The projection redistributes gradient energy from the top 5–10 informative coefficients to the tail (low-probability tokens) in the form of increased variance — i.e. compressed gradients arrive at the backbone as noisier feedback, not just lower-magnitude feedback [§3.3, Fig. 7].
In a 1.8B-Transformer-backbone + low-rank LM-head sweep on 11B FineWeb-Edu tokens (controlling for backbone capacity), validation loss curves and downstream zero-shot scores cleanly separate by $D \in \{64, 128, \dots, 4096\}$ ; $D=4096$ reaches $D=512$ ‘s final loss in 700M tokens — 1.75× convergence speedup — and final downstream average gap between $D=512$ and $D=4096$ is +0.55 [§3.1, Figs. 1, 3, App. F].
A trivial synthetic language “SpamLang” (a single uniformly-drawn token repeated for the whole sequence, for which Transformers are provably expressively-sufficient) becomes harder and eventually impossible to learn as $V$ increases at fixed $D$ , with the failure being optimization-driven rather than expressivity-driven; LR sweeps cannot recover convergence at $V=131{,}072$ [§3.2, Fig. 4, Table 1].
Updating logits along the true logit gradient $G$ vs. along the realizable direction $JG$ (gradient backpropagated through $W$ , then forward through $W^\top$ ) yields orders-of-magnitude different per-budget loss reduction, with the gap shrinking as $D$ grows; the gap is independent of overall model size (Pythia 1B and 1.4B collapse onto each other at fixed $D$ ) — implicating hidden dim, not parameter count, as the relevant axis [§3.4, Fig. 8].
The paper hypothesizes that hidden dimension should be factored into scaling laws (Hoffmann/Chinchilla) as a first-class variable, since extrapolations that fix $D$ will systematically misestimate compute-optimal allocations once the gradient bottleneck dominates [§4, Discussion].

Method

The setup writes language modeling as $\min_{H, W} \sum_c \omega_c \mathrm{CE}(W h_c, n_c)$ over distinct contexts $c$ (with empirical counts $n_c$ ), then tracks the logit matrix $Z = H W^\top$ under gradient descent. The key object is the gradient of the loss w.r.t. the logits, $G = (P - \widehat{N})$ where $P$ is the model’s softmax and $\widehat{N}$ is the row-normalized count matrix. Under one GD step on $(H, W)$ , the induced rescaled logit update $\widetilde{\Delta Z}$ is a sum of two rank- $D$ outer products, so $\mathrm{rank}(\widetilde{\Delta Z}) \le 2D$ — but $G$ has rank close to the number of contexts with unique observed continuations, which the paper argues (and verifies on Pythia/Pile) is close to $V$ in practice. Eckart–Young–Mirsky then forces a residual lower bound proportional to $G$ ‘s tail singular values.

Empirical validation has three legs: (1) for pretrained Pythia, Llama-3, Qwen3, GPT-2 models, count non-null diagonals in the SVD of $G$ to measure its empirical rank, and measure the projection of $G$ onto $\ker(W^\top)$ to quantify the lost fraction (95–99% across families); (2) a controlled 2B-parameter sweep — same Llama3-style 6-layer 4096-hidden backbone for all models, with only the LM head’s effective rank $D \in \{64, 128, 256, 512, 1024, 2048, 4096\}$ varied via a low-rank factorization $W = W_1 W_2$ ; (3) a synthetic SpamLang experiment that holds expressivity constant and varies $V$ to isolate the optimization side. The SmolLM2 tokenizer (~50k tokens) is used to avoid exaggerating the effect with a large vocabulary; the WSD LR schedule with cooldowns at 5B/8.5B/11B isolates final-loss gaps.

Results

Convergence speedup: in the controlled 2B pretraining sweep, $D=4096$ reaches $D=512$ ‘s final validation loss in 700M tokens out of 11B — a 1.75× speedup at iso-backbone [§3.1].
Downstream gap: final post-cooldown zero-shot average for $D=512$ vs $D=4096$ is +0.55 absolute points, monotonically increasing with $D$ across all 8 tested values [§3.1, App. F].
Lost gradient norm: 95–99% of logit gradient norm is destroyed across GPT-2, Pythia, Llama-3, Qwen3-Base; the cosine between projected and original gradient is only 0.1–0.2 [§3.3, Fig. 6, App. D].
SpamLang impossibility: at $V=131{,}072$ with $D$ fixed, no learning rate in the swept range produces convergence on a trivial repetition task that the architecture provably can represent — making the optimization-not-expressivity diagnosis empirically clean [§3.2, Fig. 5, Table 1].
Update-direction inefficiency: at fixed budget for $\|\Delta Z\|$ , the realizable direction $JG$ is orders of magnitude less loss-reducing than the true logit gradient $G$ direction; the gap is set by $D$ (not by $N$ ), and collapses cleanly across Pythia 1B/1.4B which share $D$ [§3.4, Fig. 8].

Why it’s interesting

This complements Beyond MuP 3: Special Cases, Special Treatment (Embedding, LM Head, RMS Norm) from the opposite direction: Su’s post argues Muon’s spectral-norm steepest-descent derivation breaks down for Embedding and LM Head (and gives the wrong update direction), and Lost-in-Backprop shows there is an information-theoretic bottleneck at exactly the same layer that no choice of optimizer can paper over — only $D$ can. Together they reframe the LM head from a “trivially handled output layer” to a structural primitive on the same footing as the residual stream and attention sinks. It also contrasts with Universal One-third Time Scaling in Learning Peaked Distributions, which derives the universal-1/3 slope of LM scaling laws from softmax-CE saturation — Lost-in-Backprop adds a complementary mechanism (rank- $D$ Jacobian compression) that affects the intercept and convergence speed rather than the asymptotic slope, but both papers locate scaling-law behavior in the same softmax+CE primitive. The “factor hidden dim into Chinchilla” suggestion lines up with How to Set the Learning Rate for Large-Scale Pre-training?‘s finding that $D$ affects optimal-LR scaling more sharply than $N$ , and with the existence of STEM: Scaling Transformers with Embedding Modules‘s STEM, which decouples the FFN-side embedding lookup from hidden width — symmetric on the input side. Sam’s commentary (“kind of an obvious paper ngl”) is a fair reading: the rank- $D$ Jacobian argument is elementary, but the empirical 95-99% number on production LM families and the controlled iso-backbone sweep are the new evidence.