Lost in Backpropagation: The LM Head is a Gradient Bottleneck

Abstract: The last layer of neural 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.

• The paper shows that a reduced LM head output rank causes significant gradient signal loss (up to 99%), impeding optimal backpropagation during training.
• Empirical results indicate that lower-dimensional LM heads delay convergence and worsen downstream zero-shot performance, even in trivial learning tasks.
• The study suggests that redesigning LM head architectures to preserve more gradient information could enhance training efficiency and reduce compute needs.

The LM Head as a Gradient Bottleneck: Theoretical and Empirical Insights

Introduction

"Lost in Backpropagation: The LM Head is a Gradient Bottleneck" (2603.10145) critically reexamines the role of the LM head in neural LMs, advancing the perspective that the softmax bottleneck is both an expressivity and optimization bottleneck. The standard LM head projects a $D$-dimensional feature space onto the vocabulary space of size $V$ through a linear map, with $D \ll V$ in prevailing architectures. This introduces a well-known expressivity constraint; the novel contribution of this work is a thorough theoretical and empirical demonstration that this architecture also fundamentally restricts the backpropagation of gradients, destroying up to 99% of the supervision signal during training. This effect inhibits both efficient convergence and the learning of even trivial patterns, with significant consequences for large-scale LLM pretraining.

Theoretical Analysis: Softmax Bottleneck Beyond Expressivity

Previous studies [softmax_bottleneck, pmlr-v97-ganea19a] have established that the rank mismatch imposed by the LM head limits the attainable log-probability distributions, but this work provides a more nuanced characterization:

• Softmax Layer Rank Constraint: With $D \ll V$, the model's logit matrix ($HW^\top$) and $\log$-softmax outputs are at most rank $D$ and $D+1$, respectively. Thus, only a very low-dimensional manifold of distributions is reachable, precluding arbitrary adjustment of next-token probabilities.
• Top-1 Prediction Unconstrained: Despite these constraints, top-1 predictions remain attainable with arbitrary precision for $D \ge 2$, demonstrating that the expressivity bottleneck does not prohibit correct token classification in the standard greedy setting.

The work goes beyond classical expressivity limits and analyzes the LM head during backpropagation. For both full-batch and minibatch stochastic gradient descent, the backpropagated logit gradient is projected from $\mathbb{R}^{V}$ to a $D$-dimensional subspace, eliminating most gradient information necessary for efficient learning.

Figure 1: Empirical rank of logit gradients as batch tokens increase; most batches yield nearly full-rank gradients, far greater than the LM head can propagate.

Theoretical results derived show that, in a realistic data regime, the effective logit gradient is high-rank for almost all batches, and the update induced by the LM head is only aligned with the true optimal gradient in a vanishingly small fraction of cases. The model's updates are thus necessarily misaligned with the true direction of fastest loss decrease.

Empirical Evidence: Convergence and Learning Dynamics

The empirical investigation comprises several key controlled experiments:

Training Dynamics with Constrained LM Head

A suite of 2B parameter LMs was trained with identical Transformer backbones ($d_m = 4096$), varying only the rank of the output LM head ($D$). Results indicate that reducing the output rank significantly delays convergence and impairs downstream zero-shot generalization.

Figure 2: Training dynamics under varying LM head dimensionality; reduced $D$ severely impedes convergence, independently of backbone expressiveness.

Figure 3: Average zero-shot performance on downstream tasks increases consistently with $D$.

The models with larger $D$ reach their final loss values up to 16× faster than their low-rank counterparts. Notably, the final zero-shot downstream performance continues to improve even for large $D$, underscoring the optimization bottleneck’s impact beyond mere expressivity saturation.

Learning Trivial Patterns: The SpamLang Experiment

A synthetic language ("SpamLang") where each sample consists of a single symbol repeated (a trivial learnable rule) was used to isolate the optimization effect from expressivity. Neural LMs with $D=576$ were trained on this task across vocabulary sizes up to 131,072.

Figure 4: Training loss on SpamLang indicates that as $V$ increases, models increasingly fail to learn even a trivial pattern, despite sufficient model expressivity.

Figure 5: Final validation loss in SpamLang as a function of $V$ and learning rate; increasing $V$ renders the pattern unlearnable across all learning rates.

This experiment demonstrates that increasing $V$ can render trivial sequence patterns unlearnable, entirely due to the destructive gradient compression in the output layer.

Gradient Compression Analysis

The critical measurement is the fraction of the logit gradient norm eliminated during the backpropagation step at the LM head. This is evaluated across several open-source architectures (GPT2, Pythia, Llama3, Qwen3), showing:

Figure 6: 95–99% of the logit gradient is destroyed by the LM head for all model families surveyed, regardless of architecture.

The in-depth gradient analyses reveal that:

• The remaining update is only mildly aligned with the true logit gradient (cosine similarity ∼0.1–0.2, see Appendix and Figure 7).
• Informative components of the backprop signal are preferentially suppressed, transferring supervision energy to random, low-magnitude directions, effectively injecting noise into the learning signal.

Figure 8: Projected logit gradient coefficients are highly diminished for informative tokens; energy is systematically dissipated into the tail.

Downstream Implications and Broader Impact

The main implication is that the prevailing LM head design needlessly discards optimization signal, slowing convergence and requiring greater compute/data for equivalent performance, especially as vocabulary sizes grow. This bottleneck is independent of the underlying architecture, as illustrated by consistent effects across major Transformer families and across both embedding-tied and untied variants.

Theoretical insights further indicate that proposed softmax-alternative heads that only alter the rank of the output distribution but not the backpropagation Jacobian fail to fix the optimization bottleneck. Only designs that permit higher-rank supervision propagation (i.e., architectures with more expressive Jacobian structure) could preserve more of the logit gradient.

Consequences and speculative future avenues include:

• Revision of LM scaling laws: Incorporating a hidden-to-vocab dimension ratio ($D/V$) term to correctly extrapolate pretraining efficiency and downstream generalization.
• Architectural innovation: Exploration of more expressive heads, specialized output structures, or preconditioning mechanisms specifically designed to improve gradient flow.
• Potential for substantial efficiency gains: If the bottleneck is addressed, existing LLMs could reach higher quality or require significantly less pretraining compute.

Conclusion

This work rigorously establishes that in neural LLMs, the narrow LM head is a fundamental optimization bottleneck—not merely an expressivity limit. The severe lossy gradient compression (destroying up to 99% of the supervision signal) restricts the efficiency and capabilities of large-scale LLM training. These findings urge the research community to re-examine this canonical architectural choice and to seek new solutions that directly address the limitations imposed by low-dimensional output heads on gradient propagation.

For foundational and related literature, see the original citations in (2603.10145) and, in particular, [softmax_bottleneck], [pmlr-v97-ganea19a], [godey2024why], [grivas-etal-2022-low], [chang-mccallum-2022-softmax], and others as referenced throughout the essay.