前向传播领域适应（无跨层反向传播）

# 仅前向传播的领域自适应（无跨层反向传播）
来源：https://arxiv.org/html/2608.14563
Rivaan Patil i14 加州大学圣克鲁兹分校 & Simon Dennis i14 墨尔本大学 & 郭浩 i14 & Kevin Shabahang i14

###### 摘要

仅前向传播的MLP训练（FPO）在不需要通过模型主体进行反向传播的情况下，即可对大型语言模型进行自适应。与标准微调相比，其吞吐量提升2.7–3.2倍，峰值训练内存减少约40%，同时使域外基准测试结果保持在基线种子噪声范围内——这是全网络微调无法可靠复现的特性。FPO基于一个经验观察：在Transformer模型的后层，输出层预测误差在六个公开模型中的余弦相似度达到0.47–0.59，近似于真实梯度。我们引入了一个两分钟的诊断工具，可量化任意模型各层的这种近似度，识别后期层自适应可行的区域。基于该诊断，FPO在输出层计算单个误差信号，并将其应用于每个目标层。信号不在层间传播，也不构建自动微分图。我们在三个模型系列（OLMo-2-7B、Qwen3-8B、Falcon3-7B）上评估了FPO。在所有模型中，FPO提升了域内困惑度，同时使MMLU、ARC-Challenge、HellaSwag和Winogrande等基准测试保持在基线种子噪声范围内。将监督微调局部化到FPO的目标层也能进入该状态，但计算耗时是FPO的2.2倍。

## 1引言

在领域数据上微调大型语言模型通常需要推理时两到三倍的内存占用，因为梯度图需要在反向传播过程中存储每一层的激活值。这一内存开销是在消费级GPU、边缘设备或大规模个性化所需高吞吐量场景中进行微调的主要障碍。LoRA（Huet al., 2022）等参数高效方法减少了优化器状态，但仍需通过网络进行完整的反向传播，因此仍需计算其激活值。

在大型Transformer的后层中，与MLP权重相关的梯度与一个无需通过层栈反向传播的量对齐。输出层预测误差通过反嵌入矩阵的转置和最终归一化的雅可比矩阵处理后，生成一个逐层信号，该信号与六个公开Transformer后25%层中真实MLP梯度的余弦相似度为0.47–0.59（特定于下投影层：0.35–0.59）。这一观察既启发了一个识别对齐度足够强、可利用的区域的筛选工具，也催生了基于此的构建方法。

我们的贡献包括：

1. 1. 一种层可行性诊断（第2节 https://arxiv.org/html/2608.14563#S2），它量化了输出投影伪梯度与每个候选层真实梯度之间的余弦相似度。该诊断在单GPU上约两分钟即可完成，可识别任何自回归Transformer的可行目标层。我们将其应用于六个参数从3B到8B的公开模型。
2. 2. 仅前向传播的MLP训练（第3节 https://arxiv.org/html/2608.14563#S3），一种无需构建自动微分图或通过模型主体运行反向传播即可自适应语言模型的算法。FPO比全网络监督微调快2.7–3.2倍，比部分监督微调（SFT-partial，唯一可行的仅后期层自适应的替代实现）快2.2倍，并在固定批量大小下将峰值训练内存降低约40%。
3. 3. 对仅后期层自适应作为基准安全范式的经验量化（第5节 https://arxiv.org/html/2608.14563#S5）。在三个模型系列（OLMo-2-7B、Qwen3-8B、Falcon3-7B）中，FPO使域外基准测试保持在基线种子噪声范围内，而全网络监督微调和LoRA无法复现此行为。同时运行了部分监督微调以确认该范式的属性，也验证了FPO是实现该范式的实用方法。

## 2后期层梯度冗余与余弦诊断

### 2.1 分解

考虑一个参数为θ的自回归Transformer，使用交叉熵损失L在序列上训练。对于第ℓ层的MLP权重W⁽ˡ⁾，真实梯度可加性地分解为可从输出计算的分量和残差：
∇W⁽ˡ⁾L = g̃W⁽ˡ⁾ + RW⁽ˡ⁾，(1)
其中g̃W⁽ˡ⁾是输出投影伪梯度：如果从层ℓ到输出的每一层都作为恒等映射，则在层ℓ处将观察到的梯度。具体来说，g̃W⁽ˡ⁾是（i）输出残差流梯度——通过仅反向传播输出层误差∂L/∂logits经过反嵌入和最终RMSNorm的封闭形式雅可比矩阵得到——与（ii）缓存的输入到W⁽ˡ⁾的前馈激活值的外积。残差RW⁽ˡ⁾包含了恒等近似所丢弃的一切：来自中间层非线性和注意力跨词元路由通过后续层的贡献。

伪梯度仅使用在前向传播期间计算或从logits解析获得的量。该构造中组合的封闭形式RMSNorm雅可比矩阵是针对单个归一化操作的恒定成本VJP（向量-雅可比积）——而非遍历层栈的图遍历——且不需要自动微分工具。

### 2.2 诊断

对于一个小的校准批次（我们使用16个长度为512的序列，在单张A100上约需两分钟运行），我们测量伪梯度g̃W⁽ˡ⁾与真实梯度g̃W⁽ˡ⁾+RW⁽ˡ⁾之间的方向对齐度：
cos⁽ˡ⁾ = ⟨g̃W⁽ˡ⁾, g̃W⁽ˡ⁾+RW⁽ˡ⁾⟩ / (‖g̃W⁽ˡ⁾‖ · ‖g̃W⁽ˡ⁾+RW⁽ˡ⁾‖)。(2)
高余弦值意味着残差RW⁽ˡ⁾在垂直于g̃W⁽ˡ⁾的方向上很小，因此输出投影分量捕获了真实梯度的主要方向。接近零的余弦值表示伪梯度不具信息量。

该诊断特别测量了下投影（W₃）更新的保真度。下投影位于SwiGLU MLP块的末端，是唯一一个其梯度可直接表示为（i）层残差流输出梯度与（ii）缓存的MLP中间激活值乘积的权重矩阵——这两个量在此阶段的前馈传播中都可用。门控（W₁）和上投影（W₂）位于MLP图的较早位置。它们的梯度需要通过W₃ᵀ进行额外的反向步骤才能到达。FPO使用缓存的W₃ⁱⁿⁱᵗ作为该反馈（第3.1节 https://arxiv.org/html/2608.14563#S3.SS1），当W₃变化缓慢时，这是一个有效的近似。因此，该诊断测量的是最清晰的情况：其理由仅依赖于输出投影近似的情况。对W₁和W₂的更新也基于缓存的雅可比近似，并通过下游指标进行了经验验证。

### 2.3 六个模型的调查

在涵盖3B到8B参数的六个公开Transformer中，后期层余弦值落在0.47–0.59范围内，其中特定于W₃的余弦值范围为0.35–0.59（表1 https://arxiv.org/html/2608.14563#S2.T1）。逐层余弦曲线具有特征性的双区间形状：在网络大部分区域低且嘈杂，然后在最后几层急剧上升并趋于平稳。上升的位置因架构而异，但在每个调查模型中都稳定地位于最后四分之一内，我们报告最后25%层的平均余弦作为模型级别的摘要。FPO目标层根据每个模型从高平台区域选择（逐层曲线见附录C https://arxiv.org/html/2608.14563#A3）。跨模型的余弦排序与域内增益和保留优势的实验排序相符（第5节 https://arxiv.org/html/2608.14563#S5），因此该诊断可对FPO最有效的区域进行排序。

表 1：六个模型的余弦诊断。#### 架构模式。
具有28层的模型往往比32+层的模型表现出更高的余弦值，这符合更深的模型将信息分布在更多层的直觉。Qwen3-8B是调查中余弦值最低的模型：输出余弦0.47，W₃余弦0.35。我们包括Qwen3作为对FPO有效范围的压力测试，预期其FPO优势最小。我们的实验证实了这一方向性预测：Qwen3显示出比全网络方法最窄的保留优势（第5.2节 https://arxiv.org/html/2608.14563#S5.SS2），并在我们三个测试模型中域内增益最小（第5.4节 https://arxiv.org/html/2608.14563#S5.SS4）。尽管如此，FPO在此范围内仍能产生有用的自适应（域困惑度下降，且基准测试在50M token时保持在基线0.14分以内），表明FPO的有效范围显著超过了我们调查的保守端点。

### 2.4 与更新幅度的关系

余弦量化了方向保真度，而更新幅度由隐藏状态和权重范数单独设定。逐层比率‖g̃⊗h‖/‖W‖，加上非标准归一化所需的逐层校正（第3.2节 https://arxiv.org/html/2608.14563#S3.SS2），确定了一个起始学习率范围。我们观察到的OLMo（最优约在5×10⁻⁴）和Qwen3（最优约在10⁻²）之间20倍的LR差距与这些逐层比率的顺序一致，但我们的六模型调查数据太稀疏，无法拟合出一个封闭形式的缩放规则。完整逐层测量见附录C（https://arxiv.org/html/2608.14563#A3）。我们建议在括号内的起始点周围进行5倍LR扫描。

## 3仅前向传播的MLP训练（FPO）

FPO从单个前向传播中获得所有权重更新：输出层误差从logits解析计算得出，与所遇到的归一化封闭形式雅可比矩阵组合，并直接应用于每个目标层的权重。在训练期间的任何时刻都没有梯度信号遍历层栈。

我们称之为*仅前向传播*，因为前向传播产生了FPO用于计算更新的所有信息——不仅是误差信号，还有逐层权重更新方向本身。在标准的基于反向传播的训练步骤中，前向传播仅产生损失；一个单独的阶段以反向方式遍历模型，将损失责任分配到网络的层和权重。FPO没有这样的阶段。它使用的封闭形式雅可比矩阵应用（反嵌入的转置、RMSNorm导数）是应用于前向传播量的固定架构函数——它们不在层间分配责任，也不会随着模型训练而改变。每个权重更新都在前向传播之后，其间没有模型遍历步骤。

### 3.1 方法

给定一个预训练模型和由诊断确定的目标层集ℒ，每个FPO训练步骤如下进行：

1. 1. 执行标准前向传播，在每个目标层ℓ∈ℒ处缓存MLP中间激活值hₘᵢᵈ⁽ˡ⁾（门控后、上投影后、下投影前），以及我们需要校正的每个归一化的输入。
2. 2. 从输出logits z和目标标签𝟏ᵧ，形成输出层误差e = softmax(z) - 𝟏ᵧ ∈ ℝᴮˣᴸˣⱽ。
3. 3. 通过语言模型头的转置将e投影到残差流空间：g = e Wₗₘ ∈ ℝᴮˣᴸˣᴰ。(3)
   这是交叉熵损失关于反嵌入输入的精确梯度。
4. 4. 将g与最终RMSNorm的封闭形式雅可比矩阵组合。对于具有逐层前馈后RMSNorm的架构（OLMo-2、Gemma 3），进一步与该层的雅可比矩阵组合，以产生层特定信号g̃⁽ˡ⁾。这些雅可比矩阵应用是简单的封闭形式计算，不构建自动微分图。
5. 5. 在每个目标层应用权重更新。下投影更新在输出投影近似下是精确的：ΔW₃⁽ˡ⁾ ∝ -g̃⁽ˡ⁾,⊤ hₘᵢᵈ⁽ˡ⁾。(4)
   门控（W₁）和上投影（W₂）更新需要在MLP隐藏空间中的梯度信号——即量W₃ᵀ g̃⁽ˡ⁾，这将通过自动微分对MLP的反向传播产生。如第2.2节（https://arxiv.org/html/2608.14563#S2.SS2）所预期的，FPO用一个缓存版本替代：在训练开始时，我们记录W₃ⁱⁿⁱᵗ⁽ˡ⁾并在整个自适应过程中将其用作固定反馈矩阵：errorₕ⁽ˡ⁾ = g̃⁽ˡ⁾ W₃ⁱⁿⁱᵗ⁽ˡ⁾。(5)
   在第0步，此信号是精确的；随着W₃漂移，它成为缓存近似。在我们的自适应机制下（后期层，相对梯度裁剪在1%），W₃的累积漂移在操作范围内是有界的（附录D https://arxiv.org/html/2608.14563#A4），使W₃ⁱⁿⁱᵗ接近实时的W₃ᵀ。对W₁和W₂的更新由errorₕ⁽ˡ⁾与适当的激活导数和缓存输入形成。

我们将这些更新公式与AdamW（权重衰减0.01）或带动量的SGD（β=0.9）配对。两者应用于相同的梯度信号。

由于在任何点都不构建自动微分图，因此相对于推理的峰值内存开销是（i）目标层的缓存前向激活，（ii）这些层的优化器状态（仅目标权重，约10%的参数），以及（iii）这些权重的float32主副本（附录A https://arxiv.org/html/2608.14563#A1）。这三者与监督微调步骤所需的梯度图相比都很小。测量见第5.1节（https://arxiv.org/html/2608.14563#S5.SS1）。

### 3.2 特定架构的雅可比矩阵

#### 最终RMSNorm（所有模型）。

在所有情况下都应用最终RMSNorm的雅可比矩阵。对于形式为 y = γ x / √(mean(x²) + ε) 的归一化，雅可比矩阵是关于x和γ的封闭形式表达式。

#### 前馈后RMSNorm（OLMo-2、Gemma 3）。

这些架构在MLP输出之后、残差相加之前应用额外的RMSNorm。由于逐层γ权重较小，且MLP输出的逐token RMS在离群token上高度可变，雅可比矩阵的1/rms因子会放大信号。