跳转至

Progressive Residual Warmup for Language Model Pretraining

会议: NeurIPS 2026
arXiv: 2603.05369
OpenReview: https://openreview.net/forum?id=aoiSdieuFb
代码: https://github.com/dandingsky/ProRes
领域: 预训练
关键词: 语言模型预训练, 残差连接, 渐进式预热, 训练稳定性, 深度扩展性

一句话总结

针对深度 Transformer 预训练初期深层残差过早介入导致表征动荡与梯度混乱的瓶颈,提出渐进式残差预热(ProRes),通过随层深线性递增残差预热步数,使深层等待浅层稳定后再逐步参与学习,全面提升模型收敛速度、深度扩展能力与下游泛化表现。

研究背景与动机

现代大语言模型几乎普遍采用顺序堆叠的 Transformer 架构,其参数规模与深度的扩展高度依赖残差连接和层归一化机制。为了解决深层网络优化中的梯度爆炸、消失以及激活值异常放大等不稳定现象,业界提出了 Pre-LN、Sandwich-LN、DeepNorm 以及按深度/宽度缩放初始权重的各类方案。然而,这些经典稳定策略存在一个根本性局限:它们大多是“训练阶段无关”的静态设计,通常仅根据初始化时刻的理论性质对权重或残差流施加刚性约束,随后便完全交由优化器自行调整。然而,实际的 Transformer 训练具有高度异质的时序阶段性,从初期充满剧烈扰动的预热阶段,到平稳收敛期,再到后期的衰减冷却阶段,模型更新的平滑度与尺度存在显著差异。

更关键的是,顺序堆叠结构内在地决定了前后层之间的非对称依赖逻辑:深层的输入严格依赖于浅层所提炼的中间表征,而浅层反向传播的梯度又直接受制于深层的权重状态。经验观察表明,浅层网络表征往往比深层更早收敛。然而,标准的残差连接从第一步训练开始就赋予所有层同等的表征修改权限。在训练初期浅层特征尚未成形、充满随机噪声的情况下,深层残差分支过早执行激进的特征修改,不仅自身在不断适应动态漂移的输入,还会向后传递混乱的梯度信号,反向干扰浅层的初始稳定,甚至在 Pre-LN 结构中引发激活值随深度呈指数级膨胀的病态现象。

这种“全局同步更新”与“层序单向依赖”之间的结构性矛盾,使得深度扩展时的优化效率大打折扣。如果能在时间维度上显式协调各层的介入时序,让浅层先建立稳健的低级表征基石,深层再阶梯式释放其拟合潜能,便能从根本上规避早期的层间干扰,同时避免静态缩放方法在训练中后期对模型表达容量的过度压制。核心 idea:提出渐进式残差预热(ProRes),为各层残差分支引入随深度线性延长预热周期的动态缩放因子 \(\alpha(l,t)\),严格践行“浅层优先学习”理念,通过从精确恒等映射平滑过渡到全容量表达,实现深度与时序维度的双重平稳更新。

方法详解

整体框架

ProRes 的核心机制极其精炼且具备极高的模块通用性。在标准的 Transformer 层中,残差连接将输入状态与子层变换(注意力模块或前馈网络模块)的输出直接相加。ProRes 则在残差增量前引入一个预定义的、依赖于层索引 \(l\) 与当前训练步数 \(t\) 的标量系数 \(\alpha(l,t)\)。在训练伊始(\(t=0\)),所有层的 \(\alpha(l,0)\) 均为 0,整个网络在数学上严格退化为恒等映射(Identity Mapping);随着训练推进,各层的 \(\alpha(l,t)\) 从 0 线性爬升至 1,且层数越深,完成预热所需的步数越多。浅层率先完全释放其表征学习能力,深层则保持低幅扰动以等待上游特征收敛,最终全网所有层无损恢复至标准架构的全量表达能力。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入表示 $x_l$"] --> Fork["分支计算与恒等映射"]
    Fork --> Res["残差变换 $\mathcal{F}(\mathrm{Norm}(x_l))$"]
    Res --> Scale["层序递进预热调度<br/>$\alpha(l,t)=\min(\frac{t}{T \cdot l}, 1)$"]
    Scale --> Bound["时序与深度联合动态界限<br/>动态松弛:初态恒等映射 $\to$ 终态全容量表达"]
    Fork --> Add["残差累加结合点"]
    Bound --> Add
    Add --> Out["架构无关的即插即用泛化适配<br/>Pre-LN / Post-LN / DeepNorm 输出 $x_{l+1}$"]

关键设计

1. 层序递进预热调度:按深度延迟残差介入以尊重层间逻辑依赖 传统 Transformer 架构中所有残差分支从第 0 步开始全功率运行,导致深层在随机初始化的上游表征上胡乱更新,产生大量无效甚至对抗性的梯度更新。针对这一层间干扰痛点,ProRes 设计了显式遵循层序依赖的预热调度函数。设第一层的基准预热步数为 \(T\),总层数为 \(L\),则第 \(l\) 层的残差权重系数定义为: $\(\alpha(l,t) = \min\left(\frac{t}{T \times l}, 1\right), \quad l \in \{1, 2, \dots, L\}\)$ 在这一调度下,第 \(l\) 层需要经历 \(T \times l\) 步才能使残差系数完全达到 1.0,全网完成预热的总步数为 \(T \times L\)。浅层残差迅速激活并主导早期的表征收敛,深层残差在初期被严重抑制,仅传递微弱信号,直至浅层特征分布逐步稳定后,深层才逐步提升参与度。消融实验证实,这种“浅层到深层”的激活顺序对于稳定优化至关重要:若将各层调度反转(即深层先预热的 reverse 模式)或让各层同步预热(equal 模式),深层会过早主导更新并导致浅层梯度消失,极易引发训练发散。

2. 时序与深度联合动态界限:从严格恒等映射平滑过渡到全容量表达 以往控制深层 Transformer 更新幅度的方案(如 DeepNet、Fixup 或 LayerNorm Scaling)普遍依赖静态数学推导,将残差分支永久或固定地除以随深度增长的因子(如 \(1/\sqrt{l}\) 或 \(1/\sqrt{2L}\))。这种静态约束虽能抑制初始化阶段的梯度爆炸,但具有“过于保守”的副作用,在模型进入稳定训练阶段后人为切除甚至永久压制了深层的表达容量。ProRes 将模型更新有界性从静态时空拓展为时序动态平滑演进:在 \(t=0\) 时,\(\alpha(l,0)=0\) 使网络呈现绝对严格的恒等映射,从根源上消除了随机权重引起的初期激活值异常爆炸;随着训练进入稳定期,各层约束被动态解除,\(\alpha(l,t)\) 均完全归 1。动态过渡既赋予了网络在敏感初期最完备的优化保护,又在平稳期彻底释放深层参数容量,使得超深网络(如 120 层)在扩展时无需承受表达力折损。

3. 架构无关的即插即用泛化适配:统一修正各类归一化范式的层间更新偏差 不同层归一化范式具有各异的层间偏差特征。例如经典的 Post-LN 天生存在向末尾层倾斜的病态趋势,深层梯度远大于浅层,致使浅层更新陷入停滞;而 Pre-LN 则容易在深层积累极大的激活范数。ProRes 作为一种即插即用的残差重加权策略,能够自然地与多种主流规范化与初始化策略解耦集成: - Pre-LN: \(x_{l+1} = x_l + \alpha(l,t) \cdot \mathcal{F}(\mathrm{Norm}(x_l))\) - Post-LN: \(x_{l+1} = \mathrm{Norm}(x_l + \alpha(l,t) \cdot \mathcal{F}(x_l))\) - Sandwich-LN: \(x_{l+1} = x_l + \alpha(l,t) \cdot \mathrm{Norm}(\mathcal{F}(\mathrm{Norm}(x_l)))\) - DeepNorm: \(x_{l+1} = \mathrm{Norm}(\alpha \cdot x_l + \alpha(l,t) \cdot \mathcal{F}_\beta(x_l))\) - LayerNorm Scaling (LNS): \(x_{l+1} = x_l + \alpha(l,t) \cdot \mathcal{F}(\mathrm{Norm}(x_l)/\sqrt{l})\) 特别是在 Post-LN 上,ProRes 的层序预热恰好强力抵消了其“深层主导、浅层冻结”的天然缺陷,使原本极难收敛的 Post-LN 获得了跨越式的性能飞跃。

损失函数 / 训练策略

模型采用标准自回归因果语言建模交叉熵损失函数进行预训练: $\(\mathcal{L}_{\mathrm{LM}} = -\frac{1}{N}\sum_{i=1}^{N} \log P(w_i \mid w_{<i})\)$ 优化器采用 AdamW,超参数设置为 \(\beta = (0.9, 0.95)\),权重衰减为 0.1,\(\epsilon = 10^{-8}\),梯度裁剪阈值设为 1.0。学习率调度采用 Warmup-Stable-Decay (WSD) 策略,总训练步数为 100k 步,全局 Batch Size 为 512,前 2,000 步线性预热至最大学习率,在最后 10% 步数内线性衰减至 0。默认主实验设置第一层预热基准长度 \(T=1000\) 步,无需任何复杂的超参数网格搜索即可在各模型尺寸下稳定运行。

实验关键数据

主实验

在 50B tokens 的 C4-en 预训练语料上对 130M、350M 和 1.3B 规模模型进行全面评测,困惑度(Perplexity, PPL \(\downarrow\))对比结果如下表所示:

方法 ProRes 130M PPL 350M PPL 1.3B PPL
Pre-LN ✘ 14.67 12.36 10.32
Pre-LN ✔ 14.30 11.74 9.86
Pre-LN (DS-Init) ✘ 14.62 12.24 10.32
Pre-LN (DS-Init) ✔ 14.29 11.73 9.85
Pre-LN (Scaled Init) ✘ 14.63 12.29 10.30
Pre-LN (Scaled Init) ✔ 14.28 11.72 9.84
Sandwich-LN ✘ 14.55 11.97 10.16
Sandwich-LN ✔ 14.50 11.78 9.94
LayerNorm Scaling ✘ 14.45 11.74 9.93
LayerNorm Scaling ✔ 14.22 11.62 9.89
Post-LN ✘ 14.84 12.74 11.62
Post-LN ✔ 14.72 11.92 10.53
DeepNorm ✘ 14.57 12.38 10.32
DeepNorm ✔ 14.45 11.97 10.09

在 1.3B 模型上评估 10 个通用推理与问答基准(Zero-shot Accuracy \(\uparrow\)),ProRes 带来平均 +1.27% 的跨架构性能增益,在长程建模基准 LAMBADA 上准确率跃升 +2.89%(PPL 从 20.46 大幅降至 15.05),在 WikiText 域外泛化测试中 PPL 从 25.61 降至 23.67。

消融实验

基于 350M 参数、24 层(\(L=24\))架构,在 6B tokens C4 语料上对多种预热函数及激活次序进行消融分析:

调度策略类型 策略配置 预热完成总步数 Pre-LN PPL Post-LN PPL Sandwich-LN PPL
基线(无 ProRes) – – 15.21 16.83 14.76
线性层序(标准) linear (T=1k) 24k 14.43 15.12 14.46
线性层序(短预热) linear (T=250) 6k 14.44 15.33 14.51
线性非线性变形 linear-square 24k 14.46 14.84 14.47
线性非线性变形 linear-sqrt 24k 14.64 15.44 14.59
全局同步预热 equal (T=1k) 1k 14.79 16.29 14.76
全局同步预热 equal (T=24k) 24k 15.29 发散 (>1000) 15.14
倒序预热(深层先) reverse (T=1k) 24k 15.65 发散 (>1000) 16.78
分段释放 stagewise-0 24k 19.51 发散 (>1000) 18.31
静态深度约束 fix-\(1/\sqrt{l}\) – 14.98 16.21 14.82
静态全网络约束 fix-\(1/L\) – 15.77 20.18 15.47

关键发现

  • 层间介入次序是决胜关键:自浅至深的 linear 预热在所有架构上均显著胜出;若改为全层同步预热的 equal 模式,随着预热步数拉长性能显著恶化甚至发散;而一旦采用深层先学的 reverse 调度,优化轨迹完全被深层破坏,Post-LN 彻底崩溃,说明“上游表征成熟前深层必须节制”是深层网络的铁律。
  • 动态松弛超越静态约束:对比 stagewise(动态恢复)与 fix(全程固定缩放),动态策略在所有配置下全面占优。静态方法过度压缩深层学习能力,而 ProRes 仅在敏感期施加缓冲,在平稳期完整保留表达力。
  • 深度扩展稳定性断层领先:在 12 层至 120 层的超深扩展实验中,Pre-LN + ProRes 的 Loss Spike 与梯度范数突刺得分几乎保持在绝对零度,在 72 层以上显著超越静态的 LayerNorm Scaling 和 DeepNorm。

亮点与洞察

  • 极简工程实现换取深远优化收益:仅需在每个残差分支乘以一行标量 \(\alpha(l,t)\),既不增加额外可学习参数,也不带来任何推理开销(推理时 \(\alpha \equiv 1\) 直接等效于原始标准结构)。
  • 天然解决 Pre-LN 激活值指数膨胀困境:特征动力学追踪表明,Pre-LN 早期激活值随深度呈指数爆炸的根源在于无序并发更新,ProRes 通过分步时钟将激活增长由指数型拉平为温和的线性增长,免除了额外复杂 norm 的需求。
  • 可复用的训练方法学启示:不仅适用于预训练,亦可迁移至超深视觉模型、多模态架构或 MoE 稀疏专家的预热路由设计中,凡存在顺序依赖与残差流的深度拓扑均可借鉴其时序分阶段解耦思路。

局限与展望

  • 作者承认的局限:目前主要针对标准前馈顺序堆叠的 Transformer 解码器进行验证;对于包含密集交叉注意力(Cross-Attention)的编码器-解码器架构,或近年来兴起的超大 Cross-layer 跳跃连接与复杂有向无环图(DAG)拓扑,线性层序调度的直接套用可能需要重新设计。
  • 自身发现的局限:主实验对于不同模型深度默认绑定了统一的 \(T\) 尺度,未对超浅层与超深层的异构学习率组合展开联合超参搜索;此外当与已有残差衰减方法(如 LayerNorm Scaling)共存时,双重抑制会在固定训练步数内过度推迟深层的充分训练。
  • 未来改进思路:探索基于每层梯度范数或表征余弦变化率自适应触发各层残差松弛的动态自适应预热(Adaptive ProRes),取代手工预设的线性步数调度。

相关工作与启发

  • vs DeepNet / DeepNorm (Wang et al., 2024): DeepNorm 通过在初始化阶段按公式静态缩减残差权重 \(\beta\) 并放大主干残差系数 \(\alpha\) 来保障千层收敛;而 ProRes 采用时序感知的动态预热,在训练平稳期完全撤除约束,避免了 DeepNorm 对深层拟合容量的永久损害,且无需繁琐的超参网格微调。
  • vs LayerLock (Erdogan et al., 2025): LayerLock 观察到浅层先收敛并采取渐进式“冻结”早期层参数以节省算力;ProRes 则从残差激活角度出发,实行浅层先启动、深层后接入的“解冻/预热”范式,二者在层间时序收敛动力学上形成了优雅的镜像呼应。
  • vs ReZero / Fixup (Bachlechner et al., 2021; Zhang et al., 2019b): ReZero 引入可学习标量参数并将残差初值设为 0,让优化器自主学习权重;但可学习标量往往极快被拉大,无法有效约束长程预热,且容易陷入病态局部极值。ProRes 采用显式确定的层序时间表,提供了更为鲁棒且可控的全局协同轨迹。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从时间与层深双重视角审视残差协同,打破了传统静态初始化与静态归一化的思维定势。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 71M 至 7B 参数、12 至 120 层深度、多达 7 种归一化/初始化组合及数十种调度消融,分析深入透彻。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论动机深刻自洽,数学表述精炼,实验分析与特征动力学可视化相得益彰。
  • 价值: ⭐⭐⭐⭐⭐ 极其简单的单行修改即可带来跨尺度、跨架构的显著收益与数值稳定性,具备极高的工业落地与推广价值。