Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Test
2026-08-17 • Computation and Language
Computation and Language
AI summaryⓘ
The authors investigate how a part of language models called the LM-head maps internal information to words and whether its limited directions cause learning problems. They show that reducing how much information flows back to the model during training modestly hurts performance, but changing the forward pass hurts more. They also find that removing some parts of this mapping is harmful and that adding unused words doesn't affect learning. Overall, their results suggest this geometric compression exists but may not be a serious training bottleneck as previously thought.
language modelLM-headTransformerbackpropagationlogitsrankgradientvocabularyoptimizationWikiText-2
Authors
Anand Murugan
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.