cs.CLAug 17, 2026

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

Authors: Anand Murugan

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

Explore similar work

CardsList
  1. Lost in Backpropagation: The LM Head is a Gradient Bottleneck

    Mar 10, 2026Nathan Godey, Yoav ArtziLLM TrainingEfficient Language Model Training

  2. Toward a First-Principles Update Geometry for the Language-Model Head

    Aug 23, 2026Aditya Somasundaram, Charles Guille-Escuret, Alexander Moreno +2

  3. ARCHead: Activation-Metric Residual Correction for Large Language Model Output Heads

    Aug 3, 2026Şuayp Talha Kocabay, Talha Rüzgar Akkuş, Kamer Ali YukselLLM QuantizationLLM Compression