Lost in Backpropagation: The LM Head is a Gradient Bottleneck

Nathan Godey; Yoav Artzi

Source-linked AI summary

Lost in Backpropagation: The LM Head is a Gradient Bottleneck

Nathan Godey, Yoav Artzi

arXiv:2603.10145v2Mar 10, 2026cs.CL

TL;DR

Language models commonly map D-dimensional features to V-dimensional vocabulary logits with D much smaller than V, creating a known softmax bottleneck. The paper theoretically analyzes the resulting gradient compression and empirically finds severe suppression of the training signal, with consequences for learning and training efficiency.

Problem
When D ≪ V, the LM head creates a known expressivity bottleneck, but the optimization effects of backpropagating high-dimensional logit gradients through this low-rank layer have been overlooked.
Method
The paper combines a theoretical analysis of rank-constrained logit updates with controlled synthetic and 2B-parameter language-model experiments and empirical gradient-spectrum measurements.
Results
95–99% of the gradient norm is suppressed through the LM head, while controlled experiments show trivial patterns can become difficult to learn and realistic pretraining convergence slows.
Takeaways & Limitations
The findings identify the LM head as a training-efficiency concern independent of the underlying model architecture and motivate new output-layer designs.
Takeaways & Limitations
The theoretical bound proves inevitable compression but not its severity, since tail singular values could be negligible; model comparisons also differ slightly in total parameter count.

Abstract

from arXiv · show

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.

1 Introduction

The paper reframes the softmax bottleneck as an optimization problem in addition to an expressivity problem. It argues that low-rank output layers compress backpropagated gradients and impair language-model training.

Most autoregressive LMs use a single linear LM head followed by softmax, mapping hidden features to token classes.
When hidden dimension D is much smaller than vocabulary size V, the LM head creates a low-rank mismatch between features and output classes.
The paper identifies gradient compression during backpropagation as an overlooked training-dynamics consequence of the softmax bottleneck.
>95% of the gradient norm is destroyed by the softmax bottleneck, according to the paper’s theoretical and empirical analysis.
Controlled experiments report that low-rank LM heads can make trivial patterns unlearnable and reduce LLM training efficiency by up to ×16 for the same backbone.
The work combines theoretical analysis, controlled 2B-parameter experiments, synthetic-language tests, and empirical measurements of gradient norm and noise transfer.

2 Theoretical Overview

The theoretical overview separates the classical expressivity constraint from a distinct optimization bottleneck. It models logit updates, proves that the gradient is structurally high-rank, and compares this inevitability with empirical compression and alternative head designs.

The paper argues that the classical low-rank expressivity framing is incomplete because the same bottleneck distorts gradients during training.
Problem Setup: The model forms logits as Lθ = HθWθ^T, so the logit matrix has rank at most D.
Optimization: Logit Updates Are Rank-Constrained: The optimal first-order logit update is −∇LL, but updates through the hidden states and LM head produce a logit update with rank at most 2D.
Optimization: Logit Updates Are Rank-Constrained: Under mild data-distribution conditions, ∇LL has full rank V, making the feasible update direction provably suboptimal beyond rank 2D.
Stochastic Gradient Descent: Near convergence, mini-batch prediction error has rank close to V −1, so stochastic training does not remove the bottleneck.
Alternative LM Head Designs: Prior mixture-based and nonlinear heads target expressivity, while the paper says optimization-compliant head designs remain an open direction.
Empirical Validation & Discussion: The theoretical residual bound establishes inevitable compression but not its severity, because tail singular values could in principle be negligible.
Empirical Validation & Discussion: A rank-D projection could theoretically recover 50–70% of the signal, but the learned LM head instead destroys 95–99% of gradient norm with cosine similarity only 0.1–0.2.

3 Experiments

Controlled experiments show that the LM-head gradient bottleneck harms optimization even when the Transformer backbone is fixed, compressing gradients and slowing learning as D/V decreases. This effect appears in large-model training and in a synthetic task where expressivity is sufficient.

3.1 Consequences for LLM Training: ×16 convergence speedup occurs between D=32 and D=4096 when training otherwise matched 2B models.Validation loss and downstream performance both improve consistently with D.
3.1 Consequences for LLM Training: Models with larger controlled LM-head dimension D converge faster and achieve better validation and downstream performance.The experiments keep the Transformer stack fixed while varying the low-rank output head.
3.1 Consequences for LLM Training: A slight parameter-count difference from 1.8B to 2.0B is a potential confound, although the reported gaps exceed what this discrepancy would predict.The controlled models share the same Transformer backbone and differ primarily in output-head rank.
3.2 Disentangling Expressivity and Optimization Issues: As vocabulary size V increases, SpamLang becomes increasingly difficult to learn despite Transformers being expressive enough for the task.The experiment uses a fixed hidden dimension and shows that increasing logit-gradient dimensionality makes supervision more compressed.
3.3 Analysis of the Gradient Compression: 95–99% of the logit-gradient norm is projected into ker(Wθ⊤) across GPT-2, Pythia, Llama 3, OLMo 2, and Qwen 3.The surviving projection is usually only mildly aligned with the original gradient, with cosine similarities mostly between 0.1 and 0.3.
3.4 Efficiency of the Update Direction: Updating hidden states along the backpropagated gradient is orders of magnitude less efficient than directly updating logits along the full gradient.Update efficiency improves as D increases, and the effect is independent of model size when hidden dimensions match.

4 Related Works & Discussion

Prior work primarily treats the softmax bottleneck as an expressivity and representation problem, while this paper emphasizes destructive gradient compression during optimization. Preliminary mitigation attempts were slower, motivating architectural approaches that better preserve gradient flow.

Softmax Bottleneck and Alternatives: Earlier studies show that D<V can limit the rank and attainable structure of next-token distributions.The literature includes low-rank expressivity limits, multi-mode distribution constraints, and unreachable probability orderings.
Representation Degeneration and Gradient Flow: Representation degeneration has been linked to the softmax bottleneck and may coincide with performance drops in small language models.Related work observes output embeddings concentrating in a narrow high-dimensional cone.
Logit Gradient & Equilibrium Analysis: The paper connects its gradient-bottleneck analysis to prior studies of logits-gradient dynamics and sign-pattern effects.It hypothesizes that sign conservation under low-rank projection may help connect these findings.
Toward Mitigations: Regularization, auxiliary alignment losses, and feedback alignment all produced slower convergence in preliminary mitigation experiments.These results suggest that the bottleneck is not easily patched without redesigning the LM head.
Discussion & Implications: Severe output-layer gradient compression is presented as an inherent limitation on LLM training efficiency and a motivation for better logits-prediction modules.The discussion identifies preconditioning, optimization methods, and softmax alternatives as possible directions.

5 Conclusion

The paper concludes that the softmax bottleneck is a fundamental optimization bottleneck, not merely an expressivity limitation: low-rank LM heads compress gradients and impair training. Controlled experiments show that this compression makes trivial patterns difficult to learn and slows convergence in realistic 2B-parameter pretraining runs.

95–99% of the supervision signal is lost during backpropagation through the output layer, transferring informative components into random-noise-like tail components.The paper characterizes this as lossy gradient compression through the LM head.
Controlled experiments show that gradient compression makes trivial patterns difficult to learn as vocabulary size grows and significantly slows convergence in realistic 2B-parameter pretraining runs.The paper reports reduced training efficiency and argues that current LMs train less efficiently than they could.
When D ≪ V, the LM head imposes a strict low-rank constraint on realizable log-probabilities, preventing arbitrary vocabulary distributions even with perfect context representations.The empirical conditional distributions are the unconstrained optimum, but rank constraints typically make that optimum unattainable.
For D ≥ 2, the model can still identify the most likely next token and its probability to arbitrary precision, so expressivity alone does not explain the training difficulty.The reachable top-1 probabilities are theoretically unconstrained for greedy decoding.
The actual global logit update has rank at most 2D, whereas direct optimization would follow the full logit gradient, which can have rank up to V.The resulting update direction cannot align with the logit gradient when the gradient has intrinsic rank exceeding 2D.
Under mild data conditions, the logit gradient is full-rank or structurally high-rank, forcing a residual bounded by tail singular values beyond rank 2D.This establishes a provably suboptimal update direction whenever the gradient rank exceeds the LM head’s effective update rank.

A.3.1 Stochastic gradient descent

Mini-batch stochasticity does not remove the gradient bottleneck: near convergence, logit gradients remain high-rank, so a rank-D LM head still produces suboptimal updates.

High-rank logit-gradient structure persists in SGD, particularly as the model approaches convergence.
The theoretical result applies to mini-batches whose contexts have distinct next tokens and satisfy a connectivity condition.
As predictions approach the empirical distribution, the first-order optimal batch logit update becomes high-rank and the actual update is suboptimal.
Replacing softmax with another output mapping does not remove the first-order limitation when its Jacobian has rank at most D.
Empirical ranks of Pythia logit gradients grow with batch size and approach the vocabulary size.
Almost exact reconstruction of the logit gradient requires up to 30,000 SVD components at the largest tested batch size.

B Proofs

The proofs establish full-rank or near-full-rank structures in logit-gradient submatrices under connectivity and distinct-token conditions, while experiments support these assumptions for natural language data.

For D > 2, the expressivity construction uses the first two hidden dimensions and sets all remaining entries to zero.
A fully connected graph case yields rank V −1, while strict diagonal dominance yields full rank by Gershgorin’s circle theorem.
Distinct batch next tokens and connectivity produce a Laplacian-like submatrix with rank V −1 when its size equals V.
When the relevant submatrix is smaller than V, strict diagonal dominance gives full rank equal to its size.
Nearly all observed contexts are unique, and most tokens appear in usual Fineweb batches of up to a few million tokens.
Fineweb contexts include a non-negligible number with next-token entropy above 5, supporting high connectivity in the underlying graph.

D Analysis of the Gradient Compression

Gradient compression smooths the dominant logit-gradient components while making the coefficient tail noisier in both Llama-3.1-8B and OLMo2-32B.

In Llama-3.1-8B and OLMo2-32B, compression smooths the main gradient components and makes the tail coefficients noisier after projection.

E Training Dynamics

Compression metrics remain broadly stable during OLMo-2-1B training, with an early worsening and a midtraining improvement whose cause is unclear.

Compression slightly worsens during the first few billion tokens, then reaches a plateau sustained throughout training.
Both compression metrics improve slightly during midtraining, coinciding with a data-distribution change and learning-rate decay to 0.
The cause of the midtraining transition is unclear because the data distribution and learning rate change at the same time.
Further exploration of the minor midtraining effect is left for future work.

F Downstream Evaluation – Detailed Results

Across downstream evaluations, models with larger effective hidden dimensions generally perform better along training, while the reported experiments find little effect from tied versus untied embeddings. Benchmark scores are grouped by training-token budget and effective hidden dimension for detailed comparison.

Downstream benchmark trends: Higher values of D generally produce better downstream benchmark performance along training.Some benchmarks show greater variance, but the overall trend favors larger effective hidden dimensions.
Downstream benchmark trends: Benchmark scores are grouped by 5B, 8.5B, and 11B training tokens and ordered by effective hidden dimension D.Both weighted and unweighted averages are reported.
Embedding and hyperparameter checks: The SpamLang experiments use learning rate 1e-4 for D = 32 and D = 768, obtaining final losses within the [0.2, 0.4] range.Initial experiments found minimal impact from weight decay on training dynamics.
Embedding and hyperparameter checks: Final validation loss shows no strong effect from tied versus untied embeddings across vocabulary sizes and learning rates.The models use 106M non-embedding parameters and a hidden dimension of 576.

G.2 Pretraining Experiments

The pretraining experiments use a staged WSD decay schedule with three intermediate decay points. Decay phases occupy 5-10% of total training and produce considerable performance gains.

Pretraining schedule: Three intermediate WSD decay phases occur after 80,000, 120,000, and 170,000 training steps.The corresponding decay lengths are 5,000, 10,000, and 10,000 steps.
Pretraining schedule: Decay length is adjusted to remain within a 5-10% ratio of total training steps.Each model performs the decay phase on the same data as the stable phase.
Pretraining outcomes: Considerable performance gains occur during the decay phase.The reported observation is consistent with prior findings cited by the authors.
Experimental setup: The pretraining experiments with constrained heads run on B200 GPUs for approximately 760 hours in total.The setup and hyperparameters are summarized in Table 4.

H Gradient Compression Example

Gradient projections through the LM head are poorly aligned with the original gradients and redistribute training feedback across related tokens and unrelated noise. The examples show dilution of precise token-level signals into broader or semantically mixed categories.

Overall gradient distortion: Projected gradients are poorly aligned with original gradients across target tokens and model sizes.The strongest coefficient remains on the ground-truth token, but the rest of the signal is clearly distorted.
Target-token examples: For “Wayne,” projected gradients reinforce related tokens such as “Rooney,” Batman-associated names, and “Detroit” or “Michigan.”Subtokens and script variants of “Wayne” also carry nearly as much gradient as the target token across model sizes.
Target-token examples: For “May,” gradient coefficients spread across typographical variants, valid alternative months, semantically related modal verbs, and lexical neighbors such as “Mayhem” or “Maya.”The example target follows the phrase “John Wayne was born on”.
Interpretation: The transferred gradient weight mixes meanings non-contextually, diluting precise feedback into broader categories or noise.The paper contrasts precise feedback such as “1907” with surrounding years and noise such as “metry”.