Lost in Backpropagation: The LM Head is a Gradient Bottleneck

• URL: http://arxiv.org/abs/2603.10145v1
• Date: Tue, 10 Mar 2026 18:27:01 GMT
• Title: Lost in Backpropagation: The LM Head is a Gradient Bottleneck
• Abstract summary: We show the softmax bottleneck is not only an expressivity bottleneck but also an optimization bottleneck.<n>Backpropagating $V$-dimensional gradients through a rank-$D$ linear layer induces unavoidable compression.<n>We argue that this inherent flaw contributes to training inefficiencies at scale independently of the model architecture.
• Score: 14.297816005348539
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• The Implicit Bias of Depth: From Neural Collapse to Softmax Codes [29.85277126753054]
  We study the deep unconstrained feature model (UFM)-equivalent to a deep linear network with inputs-trained without regularization.<n>We show that depth induces an implicit low-rank bias, promoting low-rank alternatives to NC.<n>Some effects act in the opposite direction: for randomly networks, increasing width training toward higher-rank solutions.
  arXiv  Detail & Related papers  (2026-05-21T22:37:25Z)
• Gradient Flow Polarizes Softmax Outputs towards Low-Entropy Solutions [31.988459260329353]
  We analyze flow dynamics of the value-softmax model, defined as $L(mathbfV mathbfa))$, where $mathbfV and $mathbfa$ are a learnable value matrix and attention vector, respectively.<n>We reveal that this gradient flow on the optimization toward solutions characterized by lowentropy outputs is polarizing across various objectives, including logistic square loss.
  arXiv  Detail & Related papers  (2026-03-06T13:08:56Z)
• $\nabla$-Reasoner: LLM Reasoning via Test-Time Gradient Descent in Latent Space [71.23672814629448]
  $nabla$-Reasoner is an iterative generation framework that integrates differentiable optimization over token logits into the decoding loop.<n>$nabla$-Reasoner achieves over 20% accuracy improvement on a challenging mathematical reasoning benchmark.
  arXiv  Detail & Related papers  (2026-03-05T08:42:54Z)
• Curvature-Weighted Capacity Allocation: A Minimum Description Length Framework for Layer-Adaptive Large Language Model Optimization [8.029535985033485]
  Layer-wise capacity in large language models is non-uniform; some layers contribute disproportionately to loss reduction while others are near-redundant.<n>Existing methods for exploiting this non-uniformity, such as influence-function-based layer scoring, produce sensitivity estimates but offer no principled mechanism for translating them into allocation or pruning decisions.<n>We address this gap with a unified, curvature-aware framework grounded in the Minimum Description Length (MDL) principle.
  arXiv  Detail & Related papers  (2026-03-01T04:14:15Z)
• Universal One-third Time Scaling in Learning Peaked Distributions [48.44706450307606]
  Training large language models (LLMs) is computationally expensive, partly because the loss exhibits slow power-law convergence.<n>We show that this behavior can arise intrinsically from the use of softmax and cross-entropy.
  arXiv  Detail & Related papers  (2026-02-03T16:06:18Z)
• LoRIF: Low-Rank Influence Functions for Scalable Training Data Attribution [62.830878652285406]
  Training data attribution identifies which training examples most influenced a model's prediction.<n>LoRIF exploits low-rank structures of gradient to address both bottlenecks.<n>On models from 0.1B to 70B parameters trained on datasets with millions of examples, LoRIF achieves up to 20$times$ storage reduction and query-time speedup.
  arXiv  Detail & Related papers  (2026-01-29T16:18:34Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• The Curse of Depth in Large Language Models [28.37870372690079]
  In large language models, nearly half of the layers are less effective than expected.<n>LayerNorm Scaling (LNS) scales the variance of output of the layer normalization inversely by the square root of its depth.<n>LNS consistently outperforms previous normalization and scaling techniques in enhancing LLM pre-training performance.
  arXiv  Detail & Related papers  (2025-02-09T07:03:36Z)
• Zeroth-Order Fine-Tuning of LLMs in Random Subspaces [63.10833446782114]
  As language models grow in size, memory demands for backpropagation increase.<n>Zeroth-order (ZO) optimization methods offer a memory-efficient alternative.<n>In this paper, we propose Subspace Zero-order optimization to address the challenges posed by posed by high dimensionality perturbations.
  arXiv  Detail & Related papers  (2024-10-11T17:01:43Z)
• Large Scale Private Learning via Low-rank Reparametrization [77.38947817228656]
  We propose a reparametrization scheme to address the challenges of applying differentially private SGD on large neural networks. We are the first able to apply differential privacy on the BERT model and achieve an average accuracy of $83.9%$ on four downstream tasks.
  arXiv  Detail & Related papers  (2021-06-17T10:14:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.