Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Test
Organizations: AI RESEARCH
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.
Figures & tables
| FORWARD COMPUTATION | BACKWARD PATH | WHAT IT TESTS |
|---|---|---|
| ordinary full-rank head | exact gradient | normal training |
| ordinary full-rank head | lower-rank gradient | backward rank only |
| lower-rank head | its exact gradient | forward and backward changes together |
| PROPERTY | ORIGINAL PRETRAINING COMPARISON | OUR LARGER CAUSAL CONTROL |
|---|---|---|
| Transformer layers | 6 | 6 |
| Transformer width | 4,096 | 96 |
| Parameters | about 2B | 2.26M |
| Vocabulary | 49,152 | 8,192 |
| Training tokens | about 11B | 1.23M per run |
| Intervention | factorized forward head | factorized forward and backward-only |
| FORWARD COMPUTATION | BACKWARD RANK | VALIDATION LOSS | INCREASE OVER ORDINARY TRAINING |
|---|---|---|---|
| ordinary head | 32 | — | |
| ordinary head | 16 | ||
| ordinary head | 8 | ||
| ordinary head | 4 | ||
| rank-16 forward head | 16 | ||
| rank-8 forward head | 8 |
| RANK | BACKWARD-ONLY MINUS EXACT, 95% CI | FORWARD MINUS BACKWARD-ONLY, 95% CI | ||
|---|---|---|---|---|
| 16 | ||||
| 8 | ||||
| 4 |
| FORWARD COMPUTATION | RANK | VALIDATION LOSS | INCREASE OVER ORDINARY TRAINING |
|---|---|---|---|
| ordinary head and exact gradient | 96 | — | |
| ordinary head and lower-rank gradient | 48 | ||
| rank-48 forward head and exact gradient | 48 |
| HEAD UPDATE | VALIDATION LOSS | INCREASE OVER ORDINARY TRAINING |
|---|---|---|
| complete logit gradient | — | |
| projected logit gradient only |
| SEQUENCE LENGTH | INDEPENDENT DRAWS PER SYMBOL | COUNTED TOKEN OCCURRENCES | VALIDATION LOSS |
|---|---|---|---|
| 1 | 4 | 4 | |
| 64 | 4 | 256 |
| SEQUENCE LENGTH | INDEPENDENT DRAWS PER SYMBOL | COUNTED TOKEN OCCURRENCES | VALIDATION LOSS | ACCURACY |
|---|---|---|---|---|
| 1 | 128 | 128 | ||
| 64 | 2 | 128 |
| OUTPUT CLASSES | HIDDEN WIDTH | VALIDATION LOSS |
|---|---|---|
| 256 | 32 | |
| 1,024 | 32 | |
| 4,096 | 32 |
Transformer
Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.
See all 540 papers