Lost in Backpropagation: The LM Head is a Gradient Bottleneck

Nathan Godey; Yoav Artzi

cs.CLMar 10, 2026

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

Authors: Nathan Godey, Yoav Artzi

Organizations: Cornell University, New York, USA

Abstract

The last layer of neural language models (LMs) projects output features of dimension $D$ to logits in dimension $V$ , the size of the vocabulary, where usually $D \ll V$ . This mismatch is known to raise risks of limited expressivity in neural LMs, creating a so-called softmax bottleneck. We show the softmax bottleneck is not only an expressivity bottleneck but also an optimization bottleneck. Backpropagating $V$ -dimensional gradients through a rank- $D$ linear layer induces unavoidable compression, which alters the training feedback provided to the vast majority of the parameters. We present a theoretical analysis of this phenomenon and measure empirically that 95-99% of the gradient norm is suppressed by the output layer, resulting in vastly suboptimal update directions. We conduct controlled pretraining experiments showing that the gradient bottleneck makes trivial patterns unlearnable, and drastically affects the training dynamics of LLMs. We argue that this inherent flaw contributes to training inefficiencies at scale independently of the model architecture, and raises the need for new LM head designs.

Figures & tables

Panel of Figure 1 — Figure 1: Two views of the softmax bottleneck. (a) Expressivity view: The LM head $W_{\theta}$ projects hidden states $H_{\theta}$ into a rank- $D$ subspace of the $V$ -dimensional log-probability space, constraining where model log-probabilities ( $\bullet$ ) can lie relative to the true log-probabilities ( $\bullet$ ); (b) Optimization view: During backpropagation, the full logit gradient ( $\rightarrow$ ) lives in the $V$ -dimensional space, but only a projected component ( $\rightarrow$ ) is visible when passing through $W_{\theta}$ back to the hidden state space.

Panel of Figure 2 — Figure 2: Impact of the gradient bottleneck on 2B model training. We control the rank $D$ of the output linear layer to isolate the bottleneck effect without changing the Transformer backbone. Both validation loss and downstream performance improve consistently with $D$ , with a $\times$ 16 convergence speedup between $D{=}32$ and $D{=}4096$ .

Panel of Figure 3 — Figure 3: SpamLang experiments. Despite Transformers being provably expressive enough to learn this synthetic language for any vocabulary size ( Proposition A.4 ), the gradient bottleneck makes learning increasingly difficult as $V$ grows.

V	Generated sequence
16384	A AAAAAAAAAAAAAAAAAAA
32768	A A Ġp A Ġ AA Ġ A ĠiĠ AA ’ĠĠ A Ġ
65536	A -\|lu--\|B-\| A NĠRĠb.Ġ\|
131072	A ç.\|çĠççççyççæmçè×çç

Table 1: Generated samples from our SpamLang models for different vocabulary sizes (LR: 2.5e-4). We use a single token (here A ) as a seed and generate with temperature 0.9. We use the first character of the Qwen3-Base tokens to represent SpamLang tokens. We use colors to label seed , and correct / incorrect predictions.

Panel of Figure 4 — Figure 4: Gradient compression as a function of $D/V$ , for several model families. The logit gradient norm is almost entirely destroyed by backpropagation through $W_{\theta}$ , and the surviving projected gradient is only weakly aligned with the original.

Panel of Figure 5 — Figure 5: Average full and projected logits gradient coefficients, sorted by absolute values of the full gradient. The gradients are taken for Llama-3.1-70B on Fineweb documents. The plot on the left-hand side focuses on the main coefficients, while the plot on the right-hand side zooms in on smaller coefficients. The standard deviation is reported using filled areas around the curves.

Panel of Figure 6 — Figure 6: Effect of first-order updates on logits and hidden states for the language modeling loss for different families of models.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 7 — Figure 7: Empirical and effective rank of the logit gradient $\nabla_{L}\mathcal{L}$ as a function of batch size. (a) Empirical rank across Pythia model sizes, measured on the Pile ( Gao et al., 2020 ) via QR decomposition. (b) Number of SVD components required to reconstruct $\nabla_{L}\mathcal{L}$ up to a given fraction of its total norm (color), for Pythia-1B.

Panel of Figure 8 — Figure 8: Unique tokens/contexts count as batch size increases. We use the SmolLM2 tokenizer on Fineweb documents.

Figure 9: Next-token entropy distribution for Llama-3.1-8B on Fineweb documents.

Panel of Figure 10 — Figure 10: Average full and projected logits gradient coefficients, sorted by absolute values of the full gradient. The gradients are taken for Llama-3.1-8B and OLMo2-32B on Fineweb documents. The plots on the left-hand side focus on the main coefficients, while the plots on the right-hand side zoom in on smaller coefficients. The standard deviation is reported using filled areas around the curves.

Panel of Figure 11 — Figure 11: Gradient compression metrics along training for OLMo2-1B. Standard deviation is reported using filled areas.

$D$	ARC-C	ARC-E	HellaSwag	PIQA	SciQ	OpenBookQA	Lambada $\downarrow$	Avg.	Avg. (weighted)
5B tokens
32	23.63	42.21	33.17	60.88	50.40	30.20	593.70	40.08	39.22
64	25.77	45.50	34.57	62.62	60.50	31.80	186.90	43.46	41.46
128	26.96	51.30	36.31	64.36	66.60	34.00	104.50	46.59	43.96
256	28.24	52.23	37.39	64.58	69.30	33.40	70.99	47.52	44.93
512	27.82	53.41	37.66	62.68	72.90	32.40	62.47	47.81	45.06

Appendix

Table 2: Benchmark scores grouped by training tokens (5B, 8.5B, 11B) and ordered by effective hidden dimension

D

. We report both weighted and unweighted averages.

Panel of Figure 12 — Figure 12: Downstream scores evolution along training for different benchmarks.

Hyperparameter	Value
Num. Layers	30
Head Dimension	64
Num. Attention Heads	9
Num. KV Heads	3
Hidden size	$576$
Intermediate size	$1536$

Appendix

Table 3: Training hyperparameters used for the SpamLang experiment. Experiments were run on single A6000 NVIDIA GPUs.

Panel of Figure 13 — Figure 13: Final validation loss for 106M (non-embedding) parameter Transformer models trained on the SpamLang synthetic language for various vocabulary sizes and learning rates. We report untied and tied embedding setups, and observe no strong effect from this hyper-parameter choice. The hidden dimension is set to 576 for all models.

Hyperparameter	Value
Num. Layers	6
Head Dimension	128
Num. Attention Heads	32
Hidden size	$4096$
Intermediate size	$16true384$
Weight Tying	No

Appendix

Table 4: Training hyperparameters used for the pretraining experiments with constrained heads. Experiments were run on B200 GPUs, totaling

\sim

760 hours.

Panel of Figure 14 — Figure 14: Most important gradient coefficients ranked by post-projection highest absolute values (normalized). The sentence is “ John Wayne was born on May 26th, 1907. ”, where the target token for this specific gradient is bolded.

Panel of Figure 15 — Figure 15: Most important gradient coefficients ranked by post-projection highest absolute values (normalized). The sentence is John Wayne was born on May 26th, 1907. , where the target token for this specific gradient is bolded.

Panel of Figure 16 — Figure 16: Most important gradient coefficients ranked by post-projection highest absolute values (normalized). The sentence is John Wayne was born on May 26th, 1907 . , where the target token for this specific gradient is bolded.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 342 papers

Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Test
Aug 17, 2026Anand MuruganTransformer Causal Interventions in Language Models
AI RESEARCH
Emergent One-Third Scaling Law as Attention Tries to Concentrate
Sep 26, 2026Yizhou Liu, Sara Kangaslahti, Jeff GoreLanguage Model Scaling Laws Scaling Laws
Massachusetts Institute of Technology · Harvard University
Exact Scaling Laws for a Random-Features Bigram Model
Oct 13, 2026Julian Enoch, Elvis DohmatobZipf's Law Random Feature Methods
Concordia University · Mila – Quebec AI Institute

Aug 17, 2026J/K move · Enter open · S save

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

Anand Murugan

AI RESEARCH

Transformer Causal Interventions in Language Models

Open paper AskPDF arXiv

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.

Abstract

LLM Training

Explore similar work

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

Emergent One-Third Scaling Law as Attention Tries to Concentrate

Exact Scaling Laws for a Random-Features Bigram Model