Pre-LN Transformer中部分残差消融的可重复性研究

# 部分残差消融的可重复性研究：移除注意力跳跃导致坍缩；移除前馈跳跃在10M规模可恢复但在124M规模仍未解决
来源：https://arxiv.org/html/2608.14689
###### 摘要

我研究了从预层归一化Transformer模块中选择性移除注意力跳跃连接和前馈网络跳跃连接时发生的情况，在10M参数（TinyShakespeare）和124M参数（OpenWebText）规模上训练GPT风格模型。主要发现是一种不对称性：移除注意力跳跃（FFNOnly配置）会导致确定性坍缩至无残差基线水平，在所有测试的种子和环境中均如此（10M规模下，3个种子的平均值为3.350±0.002）。移除前馈跳跃（AttnOnly）在10M规模下显示出确认的恢复效应：在A100上进行强制确定性条件下的受控8种子扫描，平均值为1.580±0.003，所有种子的验证损失均至少比坍缩基线低1.76个单位。在不同硬件上进行的中期3种子复现得到2.591±0.544——有一个种子达到了基线——这仍然是一个跨环境的不一致问题。在124M规模下，三个AttnOnly种子的平均值为6.151，标准差为1.105；这一结果具有启发性，但尚未通过受控扫描确认。在实验中期，我识别并修正了一个测量混淆因素：作为运行时乘数应用的增益会在约200步内被AdamW吸收，使得所有增益值等效。我提出了一种跨位置路由假设来解释这种不对称性，并发布了所有代码、检查点和数据——包括未复现的运行——以供独立调查。

## 1引言

残差连接由He等人[1]针对图像分类提出，现已成为Transformer语言模型[2,3]的一个定义性结构特征。在预层归一化公式中，每个模块应用两个带有恒等快捷连接的子层：

xk+1/2 = xk + Attn(LN1(xk)) (attn_res = True)
=x^k + \mathrm{Attn}(\mathrm{LN}_{1}(x^{k}))\quad(\texttt{attn\_res = True}) (1)

xk+1 = xk+1/2 + FFN(LN2(xk+1/2)) (ffn_res = True)
=x^{k+1/2} + \mathrm{FFN}(\mathrm{LN}_{2}(x^{k+1/2}))\quad(\texttt{ffn\_res = True}) (2)

跳跃连接有两个目的：它们提供梯度高速公路，防止深层网络中的梯度消失，并允许每一层学习残差变换而非完整映射[1]。带有残差连接的模块雅可比矩阵包含一个恒等项：

∂xl+1/∂xl = I + ∂f(xl)/∂xl. (3)

无论子层行为如何，这个恒等项都能保证梯度流。没有它，梯度必须穿越完整的非线性变换组合，通常会消失。将任一残差标志设置为False会移除相应的恒等路径，产生四种配置：全残差（两个跳跃都激活）、仅注意力（仅保留注意力跳跃）、仅前馈（仅保留前馈跳跃）和无残差（两者都移除）。

我使用nanoGPT风格实现进行的初步实验揭示了一个看似清晰的结果：所有三种部分配置都发生了相同的坍缩，在10M规模下验证损失达到≈3.35，在124M规模下达到≈7.4——与无残差配置无法区分。我将此作为起始观察而非结论，并探究坍缩是架构性的还是特定实现和评估方法导致的伪影。

#### 贡献.
(1) 一个在nanoGPT下建立联合坍缩观察及其梯度饥饿特征的双规模受控消融（10M和124M）。(2) 一个更清洁的重新实现（ResidualGPT），其固定验证批次揭示了在10M规模下存在2.09倍的不对称性，经受控8种子确定性扫描确认：AttnOnly平均值为1.580±0.003，FFNOnly平均值为3.350±0.002，差距为1.770个验证损失单位，没有种子接近坍缩基线。在不同硬件上进行的中期3种子复现得到2.591±0.544，作为未解决的跨环境不一致问题保留而非省略。(3) 在124M规模下的观察（种子1337时FFNOnly/AttnOnly = 1.56倍），另外两个AttnOnly种子显示出较大方差，这促使下一步进行受控的多种子124M扫描。(4) 对测量混淆因素（运行时增益缩放与权重级别初始化）的识别和诚实报告，该问题在实验中期被发现和修正。(5) 一个包含可证伪预测的机制假设；证伪测试在异常的中间环境下运行，在干净扫描条件下重新运行前结论不确定。(6) 完整发布代码、检查点和所有实验数据，包括中期未复现的运行，以便完整调查轨迹可以独立审查。

## 2背景

### 2.1残差网络与梯度流

He等人[1]表明恒等快捷连接通过确保层l的梯度幅度至少与层l+1一样大，从而允许训练非常深的网络。Elhage等人[4]的残差流形式化将其扩展到Transformer：每个子层从和向累积跨模块信息的共享向量读取和写入。移除残差连接会破坏这种组合结构，而不仅仅是优化景观。

### 2.2初始化与残差耦合

标准初始化（标准差=0.02）隐含地假设恒等路径在训练早期占主导地位。T-Fixup[5]和GPT-2输出缩放[3]表明初始化规模与残差连接紧密耦合。当残差被移除时，这种假设失效：子层输出不再在第0步由恒等路径稳定。我在证伪实验中利用了这种耦合，测试调整初始化是否能挽救部分残差配置。

## 3实验设置

### 3.1两种实现

我报告了两种实现的结果。nanoGPT实现（无权重绑定，每次评估随机验证批次，手动点积注意力）用于在两个规模上的初始联合坍缩观察。ResidualGPT实现使用了令牌嵌入和语言模型头之间的权重绑定、F.scaled_dot_product_attention（Flash Attention）、固定验证批次（使用种子99,991创建一次——在所有配置中相同）以及用于诊断的每层激活存储。所有不对称性声明均使用ResidualGPT作为主要实现。当实现差异影响结果时，会明确报告。

### 3.2超参数

表1总结了所有超参数。

表1：所有实验的超参数。参数10M124M共享n_layer612学习率：3×10^-4（nanoGPT：常数；ResidualGPT 10M：余弦，100步预热，下限3×10^-5）/ 6×10^-4→6×10^-5 余弦，2000步预热（ResidualGPT 124M）n_embd384768β1=0.9, β2=0.95, 权重衰减=0.1n_head612梯度裁剪=1.0, dropout=0.2 (10M) / 0.1 (124M)block_size2561024优化器：AdamW批大小6416梯度累积：1 (10M) / 8 (124M)训练步数3,00020,000参数≈10.7M / 124.4M验证间隔100500数据集：TinyShakespeare / OpenWebText

### 3.3数据集

10M实验使用TinyShakespeare，采用字符级分词（词表大小=65，训练集=1.0M令牌，验证集=111K）。124M实验使用OpenWebText[6]，通过tiktoken进行BPE分词（词表大小=50,257，约1.13M令牌，来自100,000篇文档）。所有ResidualGPT运行中使用固定验证批次（种子99,991），确保所有配置使用相同的评估数据。

### 3.4配置

所有实验测试四种配置：全残差（两个跳跃都激活）、无残差（两者都不激活）、仅注意力（保留注意力跳跃，移除前馈跳跃）和仅前馈（保留前馈跳跃，移除注意力跳跃）。所有超参数在配置间保持不变。唯一的变量是哪些残差连接是激活的。

## 4初始观察：联合坍缩

表2显示了nanoGPT在10M规模下的结果。三种部分配置对称坍缩至验证损失≈3.35，相比全残差退化了2.27倍，与无残差配置无法区分。这就是联合坍缩观察。

表2：nanoGPT实现，10M参数，TinyShakespeare，种子1337。配置注意力前馈最佳验证损失困惑度最终验证损失比率全残差✓✓1.47514.371.4751.00×仅注意力✓×3.353428.603.35342.27×仅前馈×✓3.347528.433.34752.27×无残差××3.352328.573.35232.27×

图1显示了10M和124M nanoGPT规模下的验证损失曲线。模式完全相同：全残差下降到低损失，而所有部分配置在初始下降后立即进入平台期。在124M规模下，退化比率（2.13倍至2.18倍）与10M（2.27倍）一致，证实联合坍缩观察在不同规模、数据集和分词方案下是稳健的。在表2中，对于所有四个10M配置，最佳验证损失和最终验证损失重合：在每个配置的最后一个记录检查点（步骤2,999）时，验证损失仍在单调改进，因此最终测量也是训练期间观察到的最佳值。

参见标题图1：使用nanoGPT实现在10M（左，TinyShakespeare）和124M（右，OpenWebText）规模下的验证损失曲线。所有部分配置都进入相同的平台期，在两个规模上都与无残差配置无法区分。

### 4.1联合坍缩的机制特征

尽管验证损失相同，四种配置显示出质上不同的内部行为（图2）。对于所有部分配置，第0层梯度范数从步骤300起恰好坍缩至0.000，而全残差配置保持≈0.114。这就是梯度高速公路效应：方程1中的恒等项提供了一条直接的梯度路径通向早期层，无论子层行为如何。隐状态范数增长比率显著不同：全残差增长1.38倍，而仅注意力增长14.03倍，仅前馈增长5.37倍。这种分歧——尽管损失相同——表明两种部分配置以不同方式失败。

## 5挑战观察

### 5.1测量混淆：运行时缩放与权重初始化

我首次尝试测试初始化敏感性时，将增益作为运行时乘数应用：在每个前向传播中输出 = g×子层(x)，其中g∈{0.25, 0.5, 2.0, 4.0}。在g=2.0时，仅注意力似乎恢复到最佳验证损失1.54。在报告结果之前，我识别出这是一个混淆因素：AdamW的每参数自适应学习率在约200步内补偿了恒定的乘法重参数化，使得所有增益值在功能上等效。所有增益值的损失曲线完全重叠（任何步骤的总展宽<<0.002）。这种表面上的恢复是优化器伪影，而非真实的初始化效应。我修正了实现，使用在优化器构建之前应用的权重级别重新初始化，并重建了实验。这里完整报告这个混淆因素作为方法论贡献：它很容易被忽视，产生看起来令人信服的结果，并且需要仔细的实验设计来避免。

参见标题图2：来自10M nanoGPT的机制诊断。(a) 梯度饥饿：对于所有部分配置，第0层范数在步骤300坍缩至0.000；全残差保持≈0.114。(b) 表征漂移：仅注意力14.0倍，仅前馈5.4倍，全残差1.4倍。(c) 训练结果（ResidualGPT，非nanoGPT）：仅注意力原始记录值（1.600）以及后续未恢复的复现值（2.591）；仅前馈（3.349）复现接近，显示为无斜线。

### 5.2实现差异：固定验证批次

nanoGPT实现在每个检查点评估随机抽样的验证批次。对于像仅注意力这样学习缓慢的模型，这引入了高评估方差：一个真正改进的模型在有噪声的评估下可能看起来已进入平台期。ResidualGPT在启动时创建一次固定验证批次（种子99,991），并在所有四个配置的每次评估中重复使用相同的样本。这消除了评估噪声作为混淆因素，这是仅注意力在nanoGPT中显示3.35而在ResidualGPT中显示1.58的最可能解释。我验证了set_global_seed()不会重新播种批处理生成器，确认固定验证批次在所有四个配置中是相同的。这个解释没有从两种实现之间的其他差异中正式分离出来：ResidualGPT还与nanoGPT在令牌嵌入和输出头之间的权重绑定、使用Flash Attention（F.scaled_dot_product_attention）而非手动实现、以及次要的优化器/初始化细节上有所不同。我没有运行除了验证批处理外固定所有这些的消融，因此固定验证批次解释应被读作给定现有证据下最合理的假设，而非已证实的因果机制；这些捆绑差异中的任何一个，单独或组合，都可能是原因。

## 6核心观察：确认的10M不对称性和未解决的124M缩放行为

### 6.110M规模 — ResidualGPT

表3显示了ResidualGPT在10M规模下的主要结果，如最初记录：仅注意力的测量值低于坍缩基线，而仅前馈发生坍缩。后来在不同硬件上的中期复现未能复现这一统计量，但随后进行的受控8种子确定性扫描确认了10M效应：8个种子的仅注意力平均值为1.580±0.003，仅前馈平均值为3.350±0.002，差距为1.770个验证损失单位，没有种子接近基线。完整的三阶段说明见第8.3节。

表3：ResidualGPT实现，10M参数，TinyShakespeare。当可用时，3个种子的平均值±标准差。仅注意力值为最初记录值；后续未恢复该精确统计量的复现尝试见第8.3节。配置最佳验证损失±标准差困惑度种子数相对全残差全残差1.5017 –4.491337–（基线）仅注意力1.5999 0.01494.951337, 42, 123+6.5% [原始记录低于基线]仅前馈3.3493 0.000928.481337, 42, 123+123.0% [坍缩]无残差3.3564 –28.691337+123.5% [坍缩]

图3可视化了表面上的不对称性，以及后续未恢复原始统计量的逐种子复现尝试。