# Lost in Backpropagation: The LM Head is a Gradient Bottleneck

arXiv: https://arxiv.org/abs/2603.10145 · HTML: https://mlsanity.com/paper/2603.10145/

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 31 references that are in the map.
- First posted on arXiv 10 Mar 2026.

## Builds on (most cited first, 17 of 17)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [SmolLM2: When Smol Goes Big -- Data-Centric Training of a Small Language Model](https://mlsanity.com/paper/2502.02737/) (2025)
- [Representation Degeneration Problem in Training Natural Language Generation Models](https://mlsanity.com/paper/1907.12009/) (2019)
- [Breaking the Softmax Bottleneck: A High-Rank RNN Language Model](https://mlsanity.com/paper/1711.03953/) (2017)
- [The Impact of Depth on Compositional Generalization in Transformer Language Models](https://mlsanity.com/paper/2310.19956/) (2023)
- [Closing the Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/2310.01693/) (2023)
- [Why do small language models underperform? Studying Language Model Saturation via the Softmax Bottleneck](https://mlsanity.com/paper/2404.07647/) (2024)
- [Sigsoftmax: Reanalysis of the Softmax Bottleneck](https://mlsanity.com/paper/1805.10829/) (2018)
- [Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practice](https://mlsanity.com/paper/2203.06462/) (2022)
- [Breaking the Softmax Bottleneck via Learnable Monotonic Pointwise Non-linearities](https://mlsanity.com/paper/1902.08077/) (2019)

## Built on by (most cited first, 9 of 9)

- [Channel Location Constrains the Auditability of Subliminal Learning](https://mlsanity.com/paper/2606.22019/) (2026)
- [Echo: KV-Cache-Free Associative Recall with Spectral Koopman Operators](https://mlsanity.com/paper/2605.06997/) (2026)
- [The Expressivity Boundary of Probabilistic Circuits: A Comparison with Large Language Models](https://mlsanity.com/paper/2605.12940/) (2026)
- [GradSentry: Gradient Spectral Entropy for Backdoor Sample Filtering in Large Language Model Fine-Tuning](https://mlsanity.com/paper/2605.26574/) (2026)
- [Toward a First-Principles Update Geometry for the Language-Model Head](https://mlsanity.com/paper/2608.22253/) (2026)
- [HeadEdit: Calibrating Language Model Behavior Through the Frozen Unembedding Matrix](https://mlsanity.com/paper/2610.01170/) (2026)
- [Cross-Lingual Alignment for Decoder-Only Models using MoE Routers](https://mlsanity.com/paper/2610.01921/) (2026)
- [Why Does Adaptive Batching Help LLM Pretraining? A Perspective from Unbounded Variance](https://mlsanity.com/paper/2610.02355/) (2026)
- [Copy-Paste Language Models](https://mlsanity.com/paper/2610.19247/) (2026)

## Abstract (arXiv)

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.