Search papers, labs, and topics across Lattice.
The paper identifies and analyzes a gradient bottleneck in the LM head, where the projection from a low-dimensional feature space to a high-dimensional vocabulary space ($D \ll V$) causes significant gradient suppression during backpropagation. Theoretical analysis and empirical measurements reveal that 95-99% of the gradient norm is suppressed by the output layer, leading to suboptimal update directions. Controlled pretraining experiments demonstrate that this bottleneck hinders the learning of even simple patterns and significantly impacts LLM training dynamics.
LLMs suffer from a severe gradient bottleneck in the output layer, suppressing 95-99% of the gradient norm and crippling training.
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.
