Lost in Backpropagation: The LM Head is a Gradient Bottleneck

神经语言模型（LM）的最后一层将维度为 $D$ 的输出特征映射为词汇表大小 $V$ 维的 logits，通常满足 $D \ll V$。这种维度不匹配已知会限制神经语言模型的表达能力，从而引发所谓“softmax 瓶颈”。我们指出，该 softmax 瓶颈不仅是一种表达能力瓶颈，更是一种优化瓶颈。当 $V$ 维梯度经由一个秩为 $D$ 的线性层反向传播时，必然发生信息压缩——这一过程不可避免地扭曲了传递给绝大多数模型参数的训练反馈信号。我们对该现象进行了理论分析，并通过实验测量发现：输出层压制了 95%–99% 的梯度模长，致使参数更新方向严重偏离最优。我们开展了一系列受控的预训练实验，结果表明：该梯度瓶颈使得一些本应极易学习的简单模式完全无法被模型掌握，并显著干扰大语言模型（LLM）的训练动力学行为。我们认为，这一固有缺陷独立于具体模型架构，是导致大规模训练效率低下的根本原因之一，因而亟需设计新型的语言模型输出头（LM head）。

论文指出神经语言模型（LMs）中经典的softmax输出层存在双重瓶颈：不仅是已知的表达能力瓶颈（expressivity bottleneck），更关键的是一个被忽视的优化瓶颈（optimization bottleneck）——由于最后线性层秩受限（rank-D，D≪V），反向传播时V维梯度被迫压缩至D维，导致95–99%梯度范数被抑制，严重扭曲参数更新方向，使训练效率大幅下降。该问题在大模型预训练中普遍存在且未被系统量化，属于对经典架构缺陷的新诊断。

首次从优化动力学角度形式化分析softmax瓶颈的梯度压缩本质，证明其导致不可逆的信息损失和次优更新方向；提出该瓶颈是架构固有缺陷，与模型主体无关，因而呼吁重构LM输出头（LM head）设计，而非仅改进主干网络。

理论推导了梯度压缩率下界，并在多个标准预训练设置（如WikiText、OpenWebText子集）上实证测量到95–99%梯度范数被抑制；通过控制实验（如学习人工可解的token重复模式）证明该瓶颈可使模型完全无法收敛；未开源代码，但实验设计高度可控（固定种子、模型规模、数据分布），结果可复现；值得深入的方向包括低秩自适应输出层、隐式词汇表建模、以及梯度感知的head初始化策略。

‘Breaking the Softmax Bottleneck: A High-Rank RNN Language Model’ (Yang et al., ICLR 2018)；‘Tied Transformers: Tying Word Embeddings and Output Projections in Transformer Language Models’ (Press & Wolf, ACL 2017)；‘LLaMA: Open and Efficient Foundation Language Models’ (Touvron et al., arXiv 2023) 中隐含了输出层瓶颈的实践缓解；‘Gradient Starvation: A Learning Dynamic that Provokes Over-Simplification of Neural Networks’ (Saxe et al., NeurIPS 2021) 提供了相关优化动态视角。