测试语言模型头的梯度压缩是否导致优化瓶颈
Does the LM Head Create a Harmful Gradient Bottleneck? A Causal Test

- 仅降低反向传播梯度秩,保持原有参数更新和输出不变
- 梯度秩减半时验证损失上升0.0586,远低于前向分解头的0.1795
- 结果表明梯度压缩虽存在但未必有害,适合研究模型优化机制者
语言模型头将宽度为D的隐藏状态映射到大小为V的词表,其转置最多只能向Transformer返回D个独立方向。Godey和Artzi认为这种严重投影是危害性优化瓶颈。本文将几何结构与因果假设分离:采用仅反向干预,保持普通logits和精确的LM头参数更新,仅降低送入Transformer的梯度秩。在byte-level和BPE-8192 WikiText-2模型上,五组配对种子实验显示,降低反向秩会增加验证损失。而同等秩的分解前向头则造成更大损失。在大模型中降至一半秩时,反向仅损失增加0.0586(95% CI [0.0167, 0.1005]),而分解前向头增加0.1795([0.1547, 0.2042])。词表空间残差也参与常规LM头更新,移除该贡献有害。额外控制实验显示,重复标记失败由独立采样符号数混淆;添加未目标输出类不损害学习;投影诊断无法可靠预测进展。测试的辅助反馈路径亦未能超越调优的反向传播。结果确认强几何压缩,但未证明其为有害优化瓶颈。
原文摘要 · Abstract (English)
The language-model head maps a hidden state of width D to a vocabulary of size V, so its transpose can return at most D independent directions to the Transformer. Godey and Artzi argue that this severe projection is a harmful optimization bottleneck. We separate the geometry from the causal claim. Our backward-only intervention keeps the ordinary logits and the exact LM-head parameter update while reducing only the rank of the gradient sent into the Transformer. Across five paired seeds on byte-level and BPE-8192 WikiText-2 models, reducing backward rank increases validation loss. An equally ranked factorized forward head, however, increases loss substantially more. At half rank in the larger model, the backward-only loss increase is 0.0586 (95% CI [0.0167, 0.1005]), while the factorized forward head increases loss by 0.1795 ([0.1547, 0.2042]). The vocabulary-space residual also contributes to the ordinary LM-head update, and removing that contribution is harmful. Additional controls show that repeated-token failures are confounded by the number of independently sampled symbols, that adding never-target output classes does not impair learning, and that projection diagnostics do not reliably predict progress in our runs. Tested auxiliary feedback routes do not beat tuned backpropagation. These results confirm strong geometric compression but do not establish that it is a harmful optimization bottleneck.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。