Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Test

Abstract: The language-model head maps a hidden state of width D to a vocabulary of size V, so its transpose can return at most D independent directions to the Transformer. Godey and Artzi argue that this severe projection is a harmful optimization bottleneck. We separate the geometry from the causal claim. Our backward-only intervention keeps the ordinary logits and the exact LM-head parameter update while reducing only the rank of the gradient sent into the Transformer. Across five paired seeds on byte-level and BPE-8192 WikiText-2 models, reducing backward rank increases validation loss. An equally ranked factorized forward head, however, increases loss substantially more. At half rank in the larger model, the backward-only loss increase is 0.0586 (95% CI [0.0167, 0.1005]), while the factorized forward head increases loss by 0.1795 ([0.1547, 0.2042]). The vocabulary-space residual also contributes to the ordinary LM-head update, and removing that contribution is harmful. Additional controls show that repeated-token failures are confounded by the number of independently sampled symbols, that adding never-target output classes does not impair learning, and that projection diagnostics do not reliably predict progress in our runs. Tested auxiliary feedback routes do not beat tuned backpropagation. These results confirm strong geometric compression but do not establish that it is a harmful optimization bottleneck.

• The paper uses backward-only rank interventions across byte-level and BPE models, finding that reduced feedback rank worsens validation loss but less than an equally ranked forward decoder, with a BPE loss gap of +0.1209.
• The study shows that the LM-head gradient’s out-of-subspace component actively updates token vectors and future feedback directions, as removing it increased validation loss by +0.6252 in the compact model.
• The results challenge the harmful-bottleneck claim by finding no reliable benefit from auxiliary feedback routes, no loss increase from artificial output-class expansion, and major confounds in repeated-token evidence, while cautioning that the experiments are small-scale.

The claim that the language-model (LM) head acts as a harmful gradient bottleneck rests on a striking geometric observation: for hidden width $D$ and vocabulary size $V$ with $V \gg D$, the transpose $W^T$ can return at most $D$ independent directions to the Transformer, so most of the logit-gradient norm lies outside the feedback subspace. Godey and Artzi report that only 1–5% of the logit-gradient norm survives this projection and argue that the projection is an optimization bottleneck (Godey et al., 10 Mar 2026). This paper separates the geometry from the causal claim. Its central instrument is a backward-only intervention that preserves the ordinary logits and the exact LM-head parameter update while reducing only the rank of the gradient sent into the Transformer. Across five paired seeds on byte-level and BPE-8192 WikiText-2 models, the results confirm severe geometric compression but do not establish that it is harmful.

The decomposition argument

Writing $P_W$ for the orthogonal projector onto the column space of $W$, the logit gradient decomposes as $g_z = g_{\parallel} + g_{\perp}$, where $W^T g_{\perp} = 0$. The component $g_{\perp}$ is invisible to the current hidden state. It is not, however, discarded by training: the LM-head gradient $V$0 includes the term $V$1, which updates $V$2, changes future predictions, and rotates the feedback subspace available to later examples. A local projection is therefore not equivalent to permanent information destruction, and any causal test must distinguish harm from removing directions that already reach the hidden state from hypothetical value in the residual itself.

The backward-only causal control

The original paper's evidence comes from lowering the rank of a factorized forward head, which simultaneously changes the decoder's expressivity, its conditioning, its parameter count, and the backward rank. This confound motivates a three-arm factorial design: ordinary full-rank training; full-rank forward computation with a lower-rank backward signal; and a lower-rank forward head with its exact gradient. The backward-only arm uses the best rank-$V$3 approximation of the current head, matches root-mean-square gradient magnitude to prevent scale artifacts, gives each arm its own learning-rate search on seed 1, and registers seeds 4–5 before inspection. The rank-32 control reproduced ordinary backpropagation exactly, validating the custom backward operation.

In the byte-level model (four layers, width 32), restricting backward rank increased validation loss monotonically with severity: $V$4 at rank 16, $V$5 at rank 8, and $V$6 at rank 4. But the equally ranked factorized forward head was consistently worse, by $V$7, $V$8, and $V$9 respectively, with all 15 paired comparisons showing the predicted sign and all 95% intervals excluding zero. In the larger BPE-8192 model (six layers, width 96, 2.26M parameters), halving backward rank cost $V \gg D$0 (95% CI $V \gg D$1, $V \gg D$2) while the rank-48 forward head cost $V \gg D$3 ($V \gg D$4, $V \gg D$5); the forward-minus-backward difference was $V \gg D$6 ($V \gg D$7). A four-seed sensitivity analysis excluding the screening seed preserved every sign and ordering, though the standalone half-rank backward penalty there had an interval crossing zero ($V \gg D$8), making the forward-versus-backward contrast the more robust result.

Two implications follow directly. First, directions already present in the hidden gradient matter—deleting them hurts. Second, and more consequentially for the bottleneck thesis, a low-rank forward-head experiment does not isolate a backward bottleneck, because changing the decoder harms far more than changing only the return path at equal rank. The study is explicit that these are small-scale causal identifications, not reproductions of the original two-billion-parameter run, and that effect sizes may not transfer to billion-parameter models.

Is the residual already used?

A complementary intervention removed $V \gg D$9 from the LM-head update while keeping the exact hidden-state gradient. This was the most damaging intervention in the compact study, increasing validation loss by $W^T$0. The result does not show that the residual should be routed into the Transformer; it shows that ordinary training already exploits it to shape token vectors and future feedback directions. Describing the residual's full length as destroyed is therefore inaccurate.

Re-examining SpamLang

The repeated-token synthetic task used as supporting evidence has a statistical confound. SpamLang samples one symbol per sequence and repeats it at every position, so a length-64 sequence supplies 64 cross-entropy terms but only one independently sampled mapping example. Holding independent draws fixed at four per symbol, sequence lengths 1 and 64 produced identical validation loss ($W^T$1), agreeing to numerical precision within seeds. Conversely, holding supervised positions fixed at 131,072, length-1 sequences (128 independent draws per symbol) reached 100% accuracy while length-64 sequences (two draws per symbol) reached only 26.6%. At the original paper's largest setting, roughly 41 million token positions over a 131,072-symbol vocabulary yield only about 4.9 independently sampled sequences per symbol—not the ~300 occurrences the position count suggests. The large-vocabulary failure thus admits a direct data-coverage explanation that must be controlled before attributing it to gradient geometry. The authors concede that repeated positions can still affect contextual Transformers' activations, so the objection is specifically about independent coverage, not computational equivalence.

Vocabulary-size controls and predictive failure

Because changing vocabulary normally changes tokenization and data statistics, the paper adds never-correct output classes to byte-level inputs, inflating $W^T$2 sixteenfold while holding text, targets, embeddings, batch order, and token count fixed. Validation loss was unchanged across 256, 1,024, and 4,096 output classes ($W^T$3, $W^T$4, $W^T$5). This does not show vocabulary size never matters—never-target classes are not a natural tokenizer—but it shows output-dimension competition alone did not create the predicted optimization problem.

Projection diagnostics also failed as predictors. Across 894 saved training points from 37 conditions, leave-one-condition-out prediction tests found that retained gradient length improved final-loss prediction error by an estimated 8.2%, but with a 95% interval spanning a 23.0% improvement to a 3.4% worsening; other measurements were similarly uncertain. Geometric measurements describe the projection but do not reliably identify which runs will learn faster.

Auxiliary feedback routes

Six families of alternative return paths—fixed random matrices, adaptively tracked dominant logit-gradient directions, residual-injection routes, multi-depth routes, intermediate-layer losses, and one-step meta-learned routes—were tested against tuned backpropagation. None gave reliable improvement. The adaptive method is instructive: tracking high-variance logit-space modes maximizes Euclidean reconstruction without yielding useful descent directions in shared parameter space. The meta-learned route's small one-step gain was exceeded by a 10% increase in the ordinary SGD learning rate. These are negative results for the tested methods, not proofs that no feedback design could help.

Limits of the study

The strongest caveats concern scale and coverage. The largest model has 2.26M parameters and trained on 1.23M tokens, so nothing here establishes behavior during billion-parameter pretraining. The larger experiment recomputed its low-rank backward approximation every 20 steps rather than every step. The added-output-class design is not a natural tokenizer comparison. The predictive analysis treats non-independent points from the same run as observations, and its intervals are wide. Most importantly, failing to find a better feedback route does not prove none exists; the authors specify what would count as stronger support—a route that keeps predictions unchanged yet reliably beats exact backpropagation across learning rates, seeds, tasks, and scales, with gains growing predictably with measured projection severity.

Conclusion

This paper confirms that the LM head strongly compresses the immediate error returned to the final hidden state, and demonstrates causally that deleting existing hidden-gradient directions is harmful—but less harmful than equivalently restricting the forward decoder. It further shows that the out-of-subspace residual actively trains the LM head, that the repeated-token evidence is confounded by independent-example coverage, that output-dimension inflation alone does not impair tuned training, and that neither projection diagnostics nor tested auxiliary routes improve on standard backpropagation. The distinction it draws is between geometry and causation: the projection exists; its harm has not been demonstrated.