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

Anand Murugan

Source-linked AI summary

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

Anand Murugan

arXiv:2608.16671v1Aug 17, 2026cs.CL

TL;DR

The paper asks whether the language-model head’s real geometric projection is also a harmful optimization bottleneck. Using backward-only rank reduction alongside a matched low-rank forward head, it finds that removing existing feedback directions hurts, but forward rank reduction hurts substantially more, leaving the causal claim unproven.

Problem
The paper asks whether the language-model head’s real geometric projection is also a harmful optimization bottleneck, separating an established geometric fact from an unresolved causal question.
Method
It uses a backward-only low-rank intervention and compares it with an equally ranked low-rank forward head to isolate causal effects.
Results
Removing directions that reach the hidden state is harmful, but reducing forward-decoder rank is more harmful still.
Takeaways & Limitations
The experiments confirm severe projection but leave the claim that the natural LM-head projection is harmful unproven.
Takeaways & Limitations
The experiments do not establish that the natural projection through the language-model head is a harmful bottleneck.

Abstract

from arXiv · show

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 experiments confirm severe geometric compression through the LM head but do not establish that it is a harmful optimization bottleneck. Removing directions from the hidden-state gradient is harmful, while changing the forward decoder is more harmful still, and several proposed supporting signals have simpler explanations or fail to improve tuned backpropagation.

Geometric compression: Most of the logit-gradient norm can lie outside the LM head’s current feedback directions when vocabulary size greatly exceeds hidden-state width.This confirms the geometric projection but does not by itself show that useful learning information was destroyed.
Residual update: The projection residual still updates the LM-head matrix, changing future predictions and the feedback directions available on later steps.Thus, being invisible to the current hidden state does not mean the component is discarded by the whole training process.
Causal backward-rank test: Deleting directions that already reach the hidden state hurts in both the byte-level and larger subword language models.This backward-only intervention preserves ordinary logits and the LM-head parameter update, isolating the effect of restricting the returned hidden-state gradient.
Forward-versus-backward intervention: Changing the forward decoder hurts much more than changing only the backward path, so low-rank forward-head experiments do not isolate a backward bottleneck.Lowering forward-head rank simultaneously changes token-score patterns, parameterization, decoder conditioning, and backward-signal rank.
Synthetic repeated-token evidence: SpamLang confounds token count with independent data coverage, because repeating one sampled symbol supplies many loss terms but only one independently sampled symbol sequence.When independent sequences per symbol were held fixed, increasing repetition did not change learning.
Additional evidence: Geometric measurements did not reliably predict learning progress, and tested auxiliary feedback routes failed to beat ordinary, well-tuned backpropagation.Together, these results leave the proposed causal problem unproven.

Relation to the Original Experiment

The original experiment showed that lower-rank forward heads train worse at large scale, but did not isolate backward gradients from forward parameterization. This study makes that distinction at small scale: backward rank reduction is harmful, yet forward factorization is substantially worse, without establishing billion-parameter effect sizes.

Original experiment: The original comparison used a six-layer, width-4096 Transformer with about two billion parameters, 11 billion tokens, and a 49,152-token vocabulary, varying forward-head rank 32–4096.The Transformer width remained fixed at 4096 while the factorized head rank varied.
Byte-level experiment: +0.0808 at rank 16, +0.1030 at rank 8, and +0.0913 at rank 4 were the forward-minus-backward-only paired differences in the byte-level experiment.All 15 paired comparisons had the predicted sign, and all three 95% intervals excluded zero.
LM-head update: Removing the vocabulary-space component from the LM-head update was most damaging, showing that ordinary training uses it without proving it should reach the Transformer.The component changes head token vectors and rotates feedback directions available to later examples.

SpamLang: Repeated Tokens Are Not Independent Mapping Examples

SpamLang’s repeated positions count as supervised tokens but not independent samples of the symbol-mapping relation. Controlled comparisons show that failure tracks independently sampled coverage, so it cannot be attributed to the LM-head gradient path without controlling coverage.

Statistical exposure: A length-64 sequence supplies 64 cross-entropy terms for one repeated identity relation, not 64 independently sampled symbols.Repeated terms produce copies of the same gradient contribution apart from loss scaling.
Controlled exposure test: With four independent examples per symbol, repetition increased supervised positions from 4 to 256 without adding mapping examples, and both five-run results matched.Within each matching seed, training paths agreed to numerical precision.
Controlled exposure test: With 131,072 supervised positions, sequence length 1 used 128 independent sequences per symbol and solved the task, whereas length 64 used two sequences and did not.Both conditions had the same 128 supervised positions per symbol.
Interpretation: At 41 million positions, the largest setting provided only about 4.9 independent sequences per symbol, making sparse coverage a direct alternative explanation for failure.The roughly 300 reported occurrences per symbol count repeated positions rather than independent mapping examples.

Increasing the Output Vocabulary Without Changing the Text

Adding never-correct output classes increased the softmax size without changing the byte-level text, targets, inputs, Transformer, batch order, or training-token count. Increasing V/D sixteenfold did not worsen the tuned model, although this control does not show that vocabulary size never matters.

Increasing the Output Vocabulary Without Changing the Text: Increasing V/D sixteenfold did not make the tuned model worse.The control preserved byte-level inputs and unchanged training conditions while adding output classes that were never correct targets.
Increasing the Output Vocabulary Without Changing the Text: Loss was 2.2336 with 256 output classes and 2.2322 with 4,096 output classes for one hidden-width-64 seed.Never-correct output classes provide a control for output-choice count, not a natural larger tokenizer.
Increasing the Output Vocabulary Without Changing the Text: The number of output choices and their competition did not create the predicted optimization problem in this experiment.The result concerns this controlled setup and does not establish that vocabulary size never matters.

Do Projection Measurements Predict Learning?

Across 37 experimental conditions, projection measurements did not reliably predict short-term learning after basic experimental differences were accounted for. Retained gradient length showed nominal predictive improvements, but the uncertainty intervals included worsening as well as improvement.

Results: The other projection measurements were similarly uncertain or unhelpful for predicting learning.Measurements included gradient-norm retention, signal angle, target and competitor preservation, dominant gradient directions, and LM-head condition number.
Evaluation design: Prediction models first used basic training and model variables, then added one projection measurement and evaluated it on a held-out condition.Basic variables included current loss, training progress, learning rate, model width, dataset, and experimental method.
Conclusion: In these runs, none of the measurements reliably predicted short-term learning once basic experimental differences were taken into account.The results do not prove that the measurements contain no useful information.

Secondary Evidence: Tested Auxiliary Feedback Routes

Tested auxiliary feedback routes did not reliably improve optimization over ordinary backpropagation after learning rates were tuned and results were repeated. The adaptive route was mathematically motivated but did not turn dominant logit-space modes into better descent directions, while the negative result does not exclude other designs.

Secondary Evidence: Tested Auxiliary Feedback Routes: None of the tested feedback routes reliably improved ordinary backpropagation after learning-rate tuning and repeated trials.The methods included random, adaptive, residual, depth-specific, intermediate-loss, and one-step meta-learned routes.
Secondary Evidence: Tested Auxiliary Feedback Routes: A 10% increase in ordinary SGD learning rate produced a larger gain than the meta-learned route’s small one-step gain.The meta-learned route therefore did not isolate a useful new source of credit.
Secondary Evidence: Tested Auxiliary Feedback Routes: Tracking dominant squared-gradient modes made mathematical sense but did not produce better optimization.These modes best reconstruct the logit gradient under Euclidean norm, yet high logit-space variance need not yield a useful descent direction in shared parameters.
Secondary Evidence: Tested Auxiliary Feedback Routes: The negative result applies to the tested methods, not to every possible feedback design.A different architecture or learning rule could still help.

What the Experiments Establish

The experiments establish strong geometric compression in logit gradients, while showing that deleting already-received directions worsens training and forward restriction is more damaging than backward-only restriction. They do not establish that the projection is harmful, that the full remainder is useful to the Transformer, or that wider feedback improves ordinary language-model training.

Established: Deleting directions that already reach the current hidden state makes training worse.The result indicates that the retained projected component contributes causally to training.
Established: Restricting the forward decoder is more damaging than applying an equal restriction only during the backward pass.This separates the existence of projection from the claim that backward projection alone is the main harmful bottleneck.
Established: The remaining logit error strongly trains the LM head itself, while repeated positions are not a substitute for independent symbol examples.These findings qualify interpretations based on Transformer feedback and repeated-token failures.
Not established: The experiments do not establish that the full logit-gradient remainder benefits the Transformer, preserves useful cross-text sharing, causes the synthetic failure, or that retained gradient length measures retained learning value.They also do not show that a wider or separate feedback path improves ordinary language-model training.
Conclusion: The projection exists, but its harm has not been established; the experiments distinguish geometry from causation.This is the section’s central interpretive conclusion.

Relation to Earlier Work

Earlier methods show that learning can tolerate approximate feedback, but they do not establish that approximate signals outperform exact gradients when exact gradients are available. The paper distinguishes its approach from synthetic gradients, deep supervision, online gradient reconstruction, and post-update parameter-gradient compression.

Approximate feedback: Feedback alignment and direct feedback alignment show that exact backward symmetry is not always required for learning, not that approximate feedback is superior to exact gradients.Feedback alignment uses fixed random matrices, while direct feedback alignment sends output errors directly to earlier layers.
Alternative training routes: Synthetic gradients predict later gradients to enable asynchronous training, whereas deep supervision adds intermediate prediction losses and changes both the objective and backward path.These methods therefore differ from interventions that alter only the feedback signal while retaining the original training objective.
Online compression: Oja's rule and Frequent Directions compress recent logit gradients for accurate reconstruction, not necessarily for selecting the best future parameter update.They track principal directions in a data stream and are presented as natural tools for compressing recent logit gradients.
Parameter-gradient compression: PowerSGD compresses parameter gradients after an update exists in parameter space, unlike treating an unconstrained vocabulary-space residual as supervision for the shared model.The distinction is between communication-oriented compression of an existing parameter update and feedback from a vocabulary-space residual.

Limits of This Study

The study’s conclusions are limited by small models, short training, experimental design choices, and wide uncertainty. In particular, failing to find a superior feedback method does not rule out one that consistently outperforms exact backpropagation.

Scope: The largest model had about 2.26 million parameters, six layers, hidden width 96, and 1.23 million supervised training tokens.The results therefore do not establish what happens in billion-parameter models or during long pretraining runs.
Experimental design: The larger backward-only experiment recomputed its lower-rank approximation every 20 steps, so updating it every step could produce a different result.
Experimental design: The added-output-class experiment changes output dimension cleanly but is not a natural tokenizer comparison because natural tokenizers also change the data seen by the model.
Statistical uncertainty: Saved training points from related small-model runs are not independent observations, and the uncertainty intervals remain wide.
Interpretation: Failing to find a better feedback method does not prove that none exists.Stronger support for the bottleneck claim would require a prediction-preserving route for unavailable logit error to reliably beat exact backpropagation across learning rates, seeds, tasks, and model sizes, with gains increasing predictably with bottleneck severity.

Reproducing the Results

The public release provides code, frozen configurations, data recipes, tests, figures, and filtered aggregate results, while omitting large generated artifacts. Reproducibility is supported by source manifests, release verification, documented runners, and tests of intervention correctness.

Release contents: The repository includes code, frozen configurations, data recipes, tests, figures, and filtered per-run aggregate results, but omits large generated artifacts.Omitted artifacts include prepared corpora, checkpoints, step-by-step logs, and machine-environment captures.
Release contents: A source manifest records the origin and checksum of every included result.
Reproduction materials: Reproduction materials identify frozen choices, commands, evidence mapping, experiment runners, prediction analysis, and dataset sources with checksums.Files include docs/PROTOCOL_LOCK.md, docs/REPRODUCIBILITY.md, several src/bottleneck runners, and data/README.md.
Verification: The release verifier checks result checksums and scans for accidental local paths and credentials.It is run with uv run python scripts/verify_release.py.
Verification: The test suite verifies unchanged forward computation in the backward-only intervention and agreement between full-rank custom feedback and ordinary backpropagation.

Conclusion

The LM head strongly compresses the error returned to the final hidden state, but the evidence does not show that this limits language-model optimization or performance. Backward-rank reduction hurts less than forward-decoder reduction, while several proposed failure signals and feedback routes do not improve tuned training.

Evidence and interpretation: The LM head can return only a small part of an arbitrary vocabulary-space error, but severe projection alone does not establish a harmful bottleneck.The missing portion still trains the LM head.
Evidence and interpretation: Reducing backward rank hurts, but reducing the forward decoder at the same rank hurts much more.This separates the geometry of gradient compression from the causal claim that it is an optimization bottleneck.
Control results: Repeated-token failures are largely explained by how few independent symbols the model observes, while increasing output dimension alone does not reproduce the claimed failure.These controls weaken interpretations based on repeated-token behavior or vocabulary size alone.
Control results: Retained gradient length does not reliably predict learning progress, and the tested separate feedback routes do not improve tuned training.The tested auxiliary routes therefore did not outperform tuned backpropagation.
Future-work standard: A convincing bottleneck result should hold the forward model fixed, express the missing signal in the model’s parameter space, improve actual training, and replicate across sizes and tasks.The proposed standard requires more than a geometric score.