语言模型预训练中的渐进残差预热方法 Progressive Residual Warmup for Language Model Pretraining
通过逐层渐进激活残差连接，提升Transformer预训练稳定性和性能
前置知识
残差连接(Residual Connection)
残差连接是深度神经网络中的关键设计，通过跳跃连接将层的输入直接加到输出上，即 $x_{l+1} = x_l + F(x_l)$。这种设计使得梯度能够直接流过网络，有效缓解了深层网络的梯度消失问题。在Transformer中，残差连接用于连接自注意力层和前馈网络层，使得模型能够训练到数百甚至上千层的深度。残差连接的核心思想是让网络学习输入到输出的残差映射，而非直接学习完整映射，这大大简化了优化过程。
本文的核心创新正是对残差连接的改进，理解残差连接的工作原理和重要性是理解ProRes方法的基础。
Pre-LN与Post-LN Transformer
Pre-LN和Post-LN是Transformer中层归一化(Layer Normalization)的两种放置方式。Pre-LN将归一化层放在残差连接之前，即 $x_{l+1} = x_l + F(\text{Norm}(x_l))$，这种设计训练更稳定，已成为现代LLM的标准做法。Post-LN则将归一化放在残差连接之后，即 $x_{l+1} = \text{Norm}(x_l + F(x_l))$，虽然表达能力更强，但训练稳定性较差，需要更复杂的初始化技巧。两者在深度扩展时表现出不同的特性：Pre-LN在深层时可能出现指数级激活增长，而Post-LN则面临深层梯度消失的问题。
本文在多种LN变体上验证ProRes的有效性，理解这两种架构差异对于理解实验结果至关重要。
学习率预热(Learning Rate Warmup)
学习率预热是训练深度神经网络时的常用技巧，在训练初期将学习率从一个很小的值逐渐增加到目标值。这样做的原因是训练初期模型参数是随机初始化的，梯度方向不稳定，如果一开始就使用大学习率容易导致训练发散。预热阶段让模型先适应数据分布，找到相对稳定的优化方向，然后再用正常的学习率进行训练。通常预热步数占总训练步数的1-10%，之后进入稳定训练阶段，最后进行学习率衰减。
ProRes本质上是一种新的预热策略，但预热对象不是学习率而是残差连接的贡献强度，理解传统学习率预热有助于理解ProRes的设计动机。
模型深度扩展(Depth Scaling)
模型深度扩展是指通过增加Transformer层的数量来提升模型容量和性能。然而，简单地增加层数往往会遇到训练不稳定、性能饱和甚至下降的问题。现有的深度扩展方法包括：使用特定的初始化策略（如DS-Init、DeepNorm）、调整层归一化的位置和方式、以及设计特殊的残差连接。这些方法的目标都是在增加深度时保持训练稳定性，但大多是在初始化阶段设计静态约束，没有考虑训练过程中的动态特性。
本文的核心贡献之一是证明ProRes能够显著改善深度扩展性能，使得模型能够扩展到120层并持续获得性能提升。
研究动机
现代大语言模型基于Transformer架构，其预训练稳定性和收敛速度是核心关注点。尽管残差连接和归一化技术已经使得训练万亿参数成为可能，但在深度扩展时仍面临显著挑战。具体来说，现有方法（如Pre-LN、Post-LN、DeepNorm等）大多在初始化阶段施加静态约束，然后让优化器自行适应整个训练过程。然而，Transformer训练表现出明显的阶段性特征：早期训练阶段（如预热期）模型更新幅度大且混乱，而稳定训练阶段的更新相对较小和平缓。更关键的是，不同深度的层具有异质性的收敛动力学——浅层往往比深层更早收敛。当所有层从初始化开始就同时修改表示时，深层可能在上游表示尚未稳定之前就开始贡献，导致低效更新或相互冲突的学习信号。实验数据显示，在1.3B参数模型上，标准Pre-LN的测试困惑度为10.32，而深层模型（120层）的困惑度会显著恶化，表明现有方法在深度扩展时的局限性。
本文的目标是本文的目标是提出一种简单且可扩展的方法来协调Transformer训练过程中的逐层残差学习。具体而言，作者希望设计一种'训练阶段感知'的残差调度策略，使得浅层在训练早期被优先学习，而深层等待上游表示稳定后再逐步参与学习。这种方法应该能够：（1）稳定预训练过程，减少损失和梯度尖峰；（2）改善深度扩展性能，使更深的模型能够持续获得性能提升；（3）加速收敛，引入独特的优化轨迹；（4）与现有各种Transformer变体兼容，包括Pre-LN、Post-LN、Sandwich-LN、DeepNorm和LayerNorm Scaling。最终目标是在71M到7B参数规模范围内验证方法的有效性和通用性。
与已有工作不同的是，本文的独特切入角度是从'时间维度'重新审视残差连接的作用。现有方法主要关注模型在初始化时的行为，设计静态约束来控制激活、权重、梯度的规模。然而，ProRes认识到Transformer优化是一个动态过程，不同训练阶段和不同深度的层需要不同的处理方式。这种'早期层优先学习'的哲学体现了对Transformer训练本质的深刻理解：由于层的顺序堆叠，深层的输入依赖于浅层，而浅层的梯度依赖于深层。通过延迟深层残差分支的贡献，ProRes确保深层在稳定的表示上构建，而不是在早期阶段放大随机初始化的噪声。这种从静态到动态、从初始化到整个训练轨迹的视角转变，是ProRes相对于现有方法的根本创新。
核心方法
ProRes的核心思想是实现'早期层优先学习'的训练策略。直觉上，想象一个12层的Transformer就像一个流水线工厂：浅层是前几道工序，深层是后几道工序。如果所有工序同时开始加工，后工序可能基于不稳定的半成品进行加工，导致质量问题。ProRes的做法是让前几道工序先稳定运行，等它们的产品质量稳定后，后面的工序再逐步加入。技术上，ProRes为每一层的残差分支引入一个预定义的缩放因子 $\alpha(l, t)$，这个因子从0开始，随训练逐步增加到1，且深层比浅层需要更长的预热时间。这样，残差分支从浅到深逐步被激活，浅层在早期训练中被优先优化，深层则在上游表示稳定后再逐步贡献。整个过程通过一个简单的线性调度实现，不引入额外的可学习参数或计算开销。
ProRes的核心创新在于将'训练阶段感知'引入残差连接的设计。与现有方法的本质区别在于：（1）动态性vs静态性——现有方法如DeepNorm、LayerNorm Scaling在初始化时设定固定的缩放因子，而ProRes的缩放因子随训练动态变化；（2）时间协调性——ProRes不仅考虑深度维度的缩放，还考虑时间维度的渐进激活，实现层间贡献的时序协调；（3）尊重序列依赖——ProRes认识到Transformer的层是顺序堆叠的，深层依赖浅层的输出，因此应该让浅层先稳定再让深层参与。这种设计使得ProRes既能在训练初期控制模型更新（类似传统预热），又能在训练后期充分利用所有层的学习能力（不像静态方法那样保守）。实验表明，这种动态策略比静态约束更有效，例如在Post-LN上，ProRes的困惑度改善比DeepNorm更显著。
方法步骤详情
ProRes的实现分为以下步骤：（1）为每一层 $l$（$l = 1, ..., L$）的残差分支引入缩放因子 $\alpha(l, t)$，其中 $t$ 是训练步数；（2）在训练开始时（$t = 0$），将所有层的缩放因子初始化为0，确保网络从恒等映射开始；（3）采用线性调度策略，$\alpha(l, t) = \min\left(\frac{t}{T \times l}, 1\right)$，其中 $T$ 是第一层的预热长度；（4）随着训练进行，浅层（小 $l$）更快达到 $\alpha = 1$，深层（大 $l$）需要更长的训练步数；（5）当所有层的 $\alpha$ 都达到1后，网络恢复为标准的Transformer架构，不引入任何额外计算。对于Pre-LN，修改后的前向传播为 $x_{l+1} = x_l + \alpha(l, t) \cdot F(\text{Norm}(x_l))$。对于其他架构变体，遵循相同原则，在残差分支前乘以 $\alpha(l, t)$。整个过程只需修改一行代码，不增加可学习参数，不影响推理速度。
技术新颖性
ProRes的技术新颖性体现在多个层面。首先，这是首个明确从'时间维度'协调层间残差贡献的方法，将优化视角从静态初始化扩展到整个训练轨迹。其次，ProRes的设计简洁优雅——仅通过一个标量缩放因子就实现了复杂的训练动态控制，这与需要修改架构或引入额外模块的方法形成鲜明对比。第三，ProRes具有广泛的兼容性，能够无缝集成到Pre-LN、Post-LN、Sandwich-LN、DeepNorm、LayerNorm Scaling等多种架构中，这在现有方法中是罕见的。第四，ProRes的'恒等映射起点'设计（$\alpha(l, 0) = 0$）比传统初始化方法更严格地控制了初始行为，为网络提供了明确的优化起点。最后，通过消融实验，作者系统验证了'浅层优先'的贡献顺序优于'同时贡献'或'深层优先'，这为理解Transformer训练动力学提供了新的理论视角。
实验结果
本文通过大量预训练实验验证了ProRes的有效性，涵盖了从130M到7B参数的多种规模。在C4-en测试集上，ProRes在所有配置下都带来了一致性的困惑度改善：对于1.3B参数的Pre-LN模型，困惑度从10.32降至9.86，改善0.46；对于Post-LN模型，从11.62降至10.53，改善1.09，这是所有架构中改善最大的。在零样本推理基准测试中，ProRes在1.3B模型上平均提升1.27%准确率，其中LAMBADA提升2.89%，HellaSwag提升2.67%，ARC-Easy提升1.85%。更重要的是，ProRes显著改善了深度扩展性能：在120层模型上，Pre-LN (ProRes)的困惑度为12.0，而标准Pre-LN为12.5，LayerNorm Scaling为12.3。训练稳定性方面，ProRes使损失尖峰分数保持接近零，而标准Pre-LN在深层时出现显著尖峰。在分布外数据上，ProRes带来更大改善：在LAMBADA测试集上，平均困惑度降低4.86，在WikiText上降低1.58。7B规模实验进一步验证了ProRes的可扩展性，即使在所有残差完成预热后，ProRes仍保持甚至扩大其优势。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| C4-en测试集困惑度（1.3B参数） | Perplexity | Pre-LN (ProRes): 9.86 | Pre-LN: 10.32 | 降低4.5% |
| C4-en测试集困惑度（1.3B参数） | Perplexity | Post-LN (ProRes): 10.53 | Post-LN: 11.62 | 降低9.4% |
| 零样本推理基准（1.3B参数平均） | Accuracy | 平均45.00% | 平均43.42% | +1.58% |
| LAMBADA语言建模（1.3B参数） | Accuracy | 42.71% | 39.30% | +3.41% |
| HellaSwag常识推理（1.3B参数） | Accuracy | 56.39% | 52.45% | +3.94% |
| WikiText困惑度（1.3B参数） | Perplexity | Pre-LN (ProRes): 23.67 | Pre-LN: 25.61 | 降低7.6% |
| LAMBADA困惑度（1.3B参数） | Perplexity | Pre-LN (ProRes): 15.05 | Pre-LN: 20.46 | 降低26.4% |
局限与改进
尽管ProRes取得了显著成效，但仍存在一些局限性。首先，ProRes的最优调度策略是架构依赖的：'线性'调度对Pre-LN和Sandwich-LN最有效，但Post-LN需要'linear-square'或'stagewise-L'调度才能获得最佳性能，这增加了实际应用中的调优复杂度。其次，ProRes与LayerNorm Scaling (LNS)的组合在大模型上效果不佳，因为LNS已经通过 $1/\sqrt{l}$ 对深层残差进行了缩放，ProRes的额外缩放进一步抑制了深层的贡献，导致在固定训练步数内深层无法充分学习。第三，ProRes的预热长度 $T$ 需要根据模型深度进行调整：在120层模型上，$T = 1000$ 会导致性能下降，需要使用 $T = 500$ 或 $T = 12k$（对于equal调度），这表明超参数选择仍需谨慎。第四，论文主要在C4-en数据集上进行实验，虽然附录中在CimbMix上验证了鲁棒性，但在更多样化的数据分布和任务上的验证仍不充分。最后，ProRes增加了训练初期的计算浪费（因为残差分支被抑制），尽管这个开销很小，但在大规模训练中可能需要考虑。
独立分析的弱点
本文存在几个值得深入探讨的弱点。首先，消融实验表明'equal'和'reverse'调度会导致训练发散，这暴露了ProRes对调度策略的敏感性，但论文没有提供理论分析来解释为什么某些调度有效而其他无效。改进方向：可以建立理论框架分析不同调度下的梯度流和优化动力学。其次，ProRes在Post-LN上的改善最大，但在Sandwich-LN和LayerNorm Scaling上的改善相对较小，这可能与这些架构本身已经对残差进行了某种形式的缩放有关。改进方向：可以设计自适应的调度策略，根据架构特性自动调整预热曲线。第三，论文没有探索ProRes与其他训练技巧的组合效果，如梯度累积、混合精度训练、模型并行等。改进方向：在更实际的大规模训练场景中验证ProRes的兼容性。第四，虽然论文声称ProRes不增加推理开销，但没有讨论模型在完成预热后是否可以安全地移除缩放因子，或者是否需要保留以维持性能。改进方向：研究ProRes训练后模型的权重分布特性，探索模型压缩的可能性。
未来方向
基于本文的研究成果，可以延伸出多个有前景的未来方向。首先，作者在讨论中提到可以探索'非单调'调度策略，即允许某些层的残差贡献先增加后减少，这可能更好地模拟人类学习的阶段性特征。其次，ProRes的核心思想可以扩展到其他类型的神经网络，如卷积神经网络、图神经网络等，探索'层优先级学习'在不同架构中的普适性。第三，可以研究ProRes与模型剪枝、知识蒸馏等技术的结合，利用训练过程中获得的层重要性信息来指导模型压缩。第四，论文的消融实验暗示'浅层优先'的学习顺序对Transformer训练至关重要，这启发我们设计更精细的层间协调机制，如基于梯度信息的动态调度。第五，ProRes在7B规模上展示了良好的扩展性，但在更大规模（如70B、175B）上的验证仍然缺失，未来工作可以在这些规模上进行系统研究。最后，可以探索ProRes在多模态模型、长上下文模型等新兴架构中的应用。
复现评估
本文在可复现性方面做得相当好。作者在GitHub上开源了完整代码（https://github.com/dandingsky/ProRes），包括所有实验配置和训练脚本。数据方面，论文使用了公开可用的C4-en数据集，tokenizer使用Llama Tokenizer，这些都可以轻松获取。算力方面，所有实验在单节点8×H800 GPU上完成：130M模型需要7小时，350M需要20小时，1.3B需要50小时，7B需要260小时，这些资源需求对于学术研究机构是可接受的。论文详细列出了所有超参数设置（学习率、batch size、训练步数等），并提供了附录中的学习率调优细节。ProRes方法本身的实现极其简单——只需修改残差连接处的一行代码，引入一个随训练步数和层索引变化的缩放因子，这大大降低了复现难度。唯一需要注意的是，不同架构需要不同的调度策略（线性、线性平方等），但这些都在论文中有明确说明。总体而言，本文的复现难度较低，代码质量高，文档齐全。
论文图表