语言模型输出层会严重压缩梯度，导致训练效率下降。
Lost in Backpropagation: The LM Head is a Gradient Bottleneck
- 输出层将高维特征压缩到词汇表维度，造成梯度信息丢失。
- 实测95%-99%的梯度范数被输出层抑制，更新方向严重失真。
- 此问题影响大模型训练动态，适用于所有架构的自回归语言模型。
神经语言模型的最后层将维度为D的输出特征映射到维度为V的logits（通常D≪V），这种不匹配不仅带来表达能力受限的软性瓶颈，更构成优化瓶颈。通过一个秩为D的线性层反向传播V维梯度，不可避免地产生压缩，改变了绝大多数参数所接收的训练反馈。我们进行了理论分析并实测发现，95%-99%的梯度范数被输出层抑制，导致更新方向远非最优。在受控预训练实验中，梯度瓶颈使得简单模式无法学习，并显著影响大模型的训练动态。我们认为这一固有缺陷独立于模型架构，在大规模训练中引发效率问题，亟需新的语言模型输出头设计。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。