反向传播中的迷失：语言模型头是梯度瓶颈 Lost in Backpropagation: The LM Head is a Gradient Bottleneck
语言模型的输出层在反向传播时压缩95-99%梯度，成为训练效率的根本瓶颈
前置知识
LM Head（语言模型头）
语言模型的最后一层，通常是一个线性变换 W∈R^{V×D}，将隐藏状态 h∈R^D 映射到词表大小维度的logits l∈R^V，然后通过softmax得到下一个token的概率分布。在Transformer架构中，这个层将模型内部的特征表示转换为对每个词表token的预测分数。由于通常 D≪V（例如D=4096，V=50000），这个层本质上是一个从低维空间到高维空间的线性投影。
本文的核心论点是这个看似简单的输出层不仅是表达能力的瓶颈，更是梯度流动的瓶颈，理解它的结构是理解全文的关键。
Softmax瓶颈
Yang et al. (2018) 首次提出的概念：当隐藏维度D小于词表大小V时，输出log-probability矩阵的秩最多为D+1，这意味着模型无法表示任意的概率分布。具体来说，对于上下文矩阵H和权重矩阵W，rank(HW^⊤)≤D，而 rank(log σ(HW^⊤))≤D+1。这意味着即使有完美的上下文表示，模型也无法实现某些目标概率分布。
之前的研究将softmax瓶颈视为表达能力限制，本文揭示它同时也是优化效率的限制，这是理解本文创新点的基础。
梯度压缩/投影
在反向传播中，损失函数对logits的梯度 ∇_L L 是V维向量，但通过LM Head的雅可比矩阵 J_f(H_θ)=W_θ 反向传播时，需要投影到D维空间。这个投影过程会将梯度投影到 W_θ^⊤ 的零空间（null space），导致大部分梯度信息丢失。数学上，被破坏的梯度范数比例为 ||P_{ker(W_θ^⊤)}(∇_L L)||_F / ||∇_L L||_F。
这是本文发现的核心现象——95-99%的梯度范数在反向传播中被破坏，直接解释了为什么语言模型训练效率低下。
秩约束下的优化动态
当通过梯度下降更新LM Head权重时，logits的更新方向 Δ_θ^τ 的秩最多为2D（因为它是两个秩至多为D的矩阵之和）。而理想情况下，直接在logit空间进行梯度下降的最优更新方向是 ∇_L L(θ^τ,X)=diag(ω)(P_θ^τ−Ñ)，这个矩阵通常具有接近V的高秩。当 D_u(N)>D 时，实际更新方向必然偏离最优方向。
这个理论分析解释了为什么低秩LM Head会导致训练效率下降——模型被迫在一个受限的子空间中寻找更新方向，无法充分利用梯度信息。
研究动机
当前语言模型架构研究主要集中在改进隐藏层（如Mamba、Differential Transformer等），但所有自回归语言模型都共享相同的输出层设计：一个线性投影加softmax。这个LM Head存在一个根本性的维度不匹配问题——隐藏维度D（如4096）远小于词表大小V（如50000）。已有研究（Yang et al., 2018; Ganea et al., 2019; Godey et al., 2024）关注这种不匹配导致的表达能力限制，即softmax瓶颈，认为它会导致表示退化和小模型性能饱和。然而，一个被忽视的关键方面是：在反向传播过程中，V维的损失梯度需要通过这个秩为D的线性层进行反向传播，这会导致梯度信息的不可避免压缩。这种梯度压缩现象从未被系统研究过，其对训练动态的影响也不清楚。
本文的目标是本文的目标是揭示softmax瓶颈不仅是表达能力问题，更是优化效率问题。作者希望通过理论分析和实验证明：(1) LM Head在反向传播时会造成有损压缩，破坏绝大部分梯度范数；(2) 这种梯度瓶颈会使某些简单模式变得不可学习；(3) 在实际的大规模预训练中，梯度瓶颈会显著降低收敛速度。具体来说，作者希望量化梯度压缩的程度，并证明即使是2B参数的模型，不同隐藏维度（D=32到D=4096）之间也会出现高达16倍的收敛速度差异。
与已有工作不同的是，本文的独特视角在于将softmax瓶颈从表达能力框架重新定位为优化动态框架。之前的Yang et al. (2018)等工作关注的是模型能否表示目标分布（表达能力），而本文关注的是模型能否高效地学习这些分布（优化效率）。这个视角转换揭示了一个关键洞察：即使当D≥2时模型理论上可以正确预测top-1 token（Proposition 2.4），梯度瓶颈仍然会严重阻碍训练过程。换句话说，表达能力的限制不足以解释观察到的训练困难，必须从梯度流动的角度来理解这个问题。这种视角为理解大规模语言模型训练效率低下提供了新的理论基础。
核心方法
本文的方法可以类比为一个水管系统：想象有一个粗水管（V维logit空间）需要通过一个细水管（D维隐藏空间）来输送水流（梯度信号）。当水从粗管流入细管时，大部分水流会被阻挡或溢出，只有少量能通过。作者通过理论分析和实验来量化这种水流损失。技术路线分为三个层次：首先，建立数学框架，将语言建模目标写成矩阵形式 L = -(1/T)⟨N, log σ(HW^⊤)⟩_F；其次，分析梯度下降过程中logits的演化，证明实际更新方向 Δ_θ^τ 的秩最多为2D，而理想更新方向通常具有接近V的高秩；最后，通过控制变量实验验证理论预测，包括合成语言实验和实际预训练实验。
本文的核心创新是揭示了一个被忽视的现象：softmax瓶颈在反向传播时造成有损压缩，破坏95-99%的梯度范数。与已有方法最本质的区别在于，已有工作（如Yang et al., 2018; Ganea et al., 2019）试图通过设计更复杂的softmax变体来提高表达能力（打破秩约束），但本文证明这些方法无法解决优化问题。具体来说，即使使用替代的输出函数 f_θ(H_θ)=L_θ，只要其雅可比矩阵 J_f(H_θ) 的秩为D，梯度压缩问题就依然存在（Equation 9）。这意味着问题的根源不在于softmax的具体形式，而在于任何秩为D的输出映射都会在反向传播时造成信息损失。这个洞察将研究焦点从设计更好的softmax转移到设计能更好保持梯度流的输出层。
方法步骤详情
方法包含以下步骤：(1) 建立矩阵形式的语言建模目标，定义上下文计数矩阵N、行归一化矩阵Ñ、上下文表示矩阵H，将损失写为 L=-(1/T)⟨N, log σ(HW^⊤)⟩_F；(2) 分析表达能力限制，证明当D≥2时，模型可以正确预测top-1 token（Proposition 2.4），说明表达能力不是主要限制；(3) 分析梯度流动，推导LM Head权重更新公式 W_θ^{τ+1}=W_θ^τ - η·(P_θ^τ-Ñ)^⊤ diag(ω) H_θ^τ，证明logit更新 Δ_θ^τ 的秩最多为2D；(4) 证明预测误差矩阵(P-Ñ)的秩下界为 D_u(N)=min(|C_{uu}|, V-1)，在自然语言中通常接近V；(5) 量化梯度压缩程度，测量 ||P_{ker(W_θ^⊤)}(∇_L L)||_F / ||∇_L L||_F；(6) 进行控制变量实验，使用秩受限LM Head（W_θ=A_θ B_θ，A_θ∈R^{V×D}，B_θ∈R^{D×d_m}）来隔离梯度瓶颈的影响。
技术新颖性
本文的技术新颖性体现在三个方面：第一，视角创新——将softmax瓶颈从表达能力问题重新定义为优化效率问题，这在已有文献中是首次系统阐述；第二，理论贡献——Proposition 2.6给出了实际更新方向与理想更新方向之间差距的下界 ||Δ_θ^τ-(P_θ^τ-Ñ)||_F > √(Σ_{i=2D+1}^{D_u(N)} ς_i²)，这个下界明确依赖于D和预测误差矩阵的奇异值；第三，实验设计——使用秩受限LM Head（W_θ=A_θ B_θ）来控制梯度瓶颈强度，同时保持Transformer骨干网络不变，这使得作者能够干净地隔离梯度瓶颈的影响，避免了模型大小变化带来的混淆因素。与已有工作相比，Yang et al. (2018)提出高秩RNN LM来打破表达能力限制，Ganea et al. (2019)提出可学习的单调点非线性，但这些方法都未考虑梯度流动问题。
实验结果
本文的核心发现通过三类实验得到验证：第一，梯度压缩的实证测量：对GPT2、Pythia、Llama3、Qwen3-Base等模型家族，使用FineWeb数据集的10000个文档，测量到95-99%的梯度范数被LM Head的零空间投影破坏（Figure 6）。具体来说，||P_{ker(W_θ^⊤)}(∇_L L)||_F / ||∇_L L||_F 的值在0.95-0.99之间，且投影后的梯度与原始梯度的余弦相似度仅为0.1-0.2。第二，合成语言实验：在SpamLang（由V个符号组成的重复序列语言）上，使用106M参数、D=576的Transformer，词汇表从1024扩大到131072时，即使是学习这种trivial的恒等函数也变得越来越困难（Figure 4、5）。当V=131072时，模型对学习率变得极度敏感，在测试范围内没有合适的学习率能保证收敛。第三，实际预训练实验：训练8个2B参数模型（基于Llama3架构，6层，d_m=4096），D从32到4096，在11B tokens的Fineweb-Edu上训练。结果显示D=4096在700M tokens时就达到了D=32在11B tokens时的验证损失水平，显示16倍的收敛速度差异（Figure 1）。下游任务评估（ARC、HellaSwag、PIQA、SciQ等）显示，D=2048和D=4096的最终平均分数差距为+0.55（加权平均），即使在大隐藏维度下也有明显性能差距。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 语言建模（验证损失） | Validation Loss | D=4096: ~3.0 (700M tokens) | D=32: ~3.0 (11B tokens) | 收敛速度提升16倍 |
| 下游任务平均（加权） | Avg. (weighted) | D=4096: 49.41 (11B tokens) | D=32: 41.37 (11B tokens) | +8.04分 |
| ARC-Challenge | acc_norm | D=4096: 31.06 (11B tokens) | D=32: 25.94 (11B tokens) | +5.12分 |
| HellaSwag | acc_norm | D=4096: 42.64 (11B tokens) | D=32: 35.70 (11B tokens) | +6.94分 |
| Lambada | ppl ↓ | D=4096: 37.27 (11B tokens) | D=32: 406.56 (11B tokens) | 困惑度降低91% |
局限与改进
本文存在以下局限性：第一，实验规模受限——作者训练的最大模型为2B参数，而当前主流LLM为7B-70B级别，梯度瓶颈在更大模型上的影响尚需验证。第二，作者使用秩受限LM Head（W_θ=A_θ B_θ）来模拟不同隐藏维度，这种方法虽然能控制变量，但与实际使用不同隐藏维度训练的模型可能存在差异，因为较小的D通常对应较小的模型，而作者尝试通过保持骨干网络不变来缓解这个问题。第三，SpamLang合成语言虽然能干净地展示问题，但其结构（单一符号重复）与自然语言差异较大，可能高估了梯度瓶颈在实际场景中的影响。第四，作者提到计算奇异值分解来估计Proposition 2.8的界是昂贵的，对于实际batch大小（1-4M tokens）的矩阵甚至无法存储，这限制了理论分析的实证验证。第五，论文未探索任何解决方案，虽然提到了预条件化、优化技术或softmax替代方案作为未来方向，但没有给出具体的改进方法。
独立分析的弱点
本文存在以下弱点：第一，理论分析中的完全表达表示假设（Assumption 2.1）可能过于理想化——假设 φ_θ 是确定性的且足够表达，使得每个上下文可以独立赋值，这忽略了上下文表示之间的实际约束和依赖关系。改进方向是分析当上下文表示共享参数时梯度瓶颈的影响。第二，Proposition 2.7的连通性条件（类似于邻接矩阵的连通图）在实践中可能不总是满足，特别是对于短序列或稀疏数据。改进方向是研究数据分布特性如何影响梯度瓶颈的严重程度。第三，作者测量梯度压缩时使用的是单个模型的权重，但训练过程中权重会变化，梯度瓶颈的影响可能随训练阶段不同而变化。改进方向是追踪训练过程中梯度压缩的动态变化。第四，Figure 8的效率分析虽然直观，但只考虑了单步更新，未考虑累积效应——长期来看，持续的次优更新可能导致模型收敛到不同的局部最优。改进方向是进行长期训练实验，比较不同D值的最终收敛点。
未来方向
作者提出了几个重要的未来研究方向：第一，探索能更好保持梯度流的LM Head设计，包括预条件化（pre-conditioning）方法，通过变换梯度使其更适合低秩投影；优化技术，如在logit空间直接进行部分更新；以及更合适的softmax替代方案，这些替代方案虽然不能提高最大秩，但可能通过雅可比矩阵 J_f 提供更少损失的压缩。第二，将隐藏维度纳入scaling law的考虑，作者推测Hoffmann et al. (2022)的Chinchilla scaling laws可能需要修正，因为当前的scaling laws没有考虑梯度瓶颈对训练效率的影响。第三，基于本文发现的乐观视角是，当前语言模型可能具有比预期更强的潜力，通过更好地引导监督信号通过更适合的logit预测模块，可能获得显著的收敛速度和/或性能提升。第四，研究替代架构（如Mamba、Differential Transformer）是否也存在类似的梯度瓶颈问题，以及它们的设计是否无意中缓解了这个问题。
复现评估
本文的复现性较好：作者承诺将发布代码、数据和检查点（Code, data and checkpoints will be released shortly）。实验设置相对简单——SpamLang是合成数据，易于生成；预训练实验使用公开的Fineweb-Edu数据集和SmolLM2 tokenizer。算力需求方面，SpamLang实验在单个A6000 GPU上运行，预训练实验在B200 GPU上总计约760小时，对于研究机构来说是可承受的。关键的超参数（Table 3、4）都已详细列出，包括层数、注意力头数、隐藏维度、学习率、batch大小等。使用WSD学习率调度和标准的AdamW优化器，没有特殊的训练技巧。复现的主要挑战可能在于：需要训练多个不同D值的2B参数模型，这需要相当的计算资源；以及梯度分析需要对大模型进行QR分解，可能需要较大的GPU内存。
论文图表
