2025-11-12T01:28:29.133817

Stability of Transformers under Layer Normalization

Kan, Li, Zhang et al.
Despite their widespread use, training deep Transformers can be unstable. Layer normalization, a standard component, improves training stability, but its placement has often been ad-hoc. In this paper, we conduct a principled study on the forward (hidden states) and backward (gradient) stability of Transformers under different layer normalization placements. Our theory provides key insights into the training dynamics: whether training drives Transformers toward regular solutions or pathological behaviors. For forward stability, we derive explicit bounds on the growth of hidden states in trained Transformers. For backward stability, we analyze how layer normalization affects the backpropagation of gradients, thereby explaining the training dynamics of each layer normalization placement. Our analysis also guides the scaling of residual steps in Transformer blocks, where appropriate choices can further improve stability and performance. Our numerical results corroborate our theoretical findings. Beyond these results, our framework provides a principled way to sanity-check the stability of Transformers under new architectural modifications, offering guidance for future designs.
academic

Stability of Transformers under Layer Normalization

基本信息

  • 论文ID: 2510.09904
  • 标题: Stability of Transformers under Layer Normalization
  • 作者: Kelvin Kan (UCLA), Xingjian Li (UT Austin), Benjamin J. Zhang (UNC Chapel Hill), Tuhin Sahai (SRI International), Stanley Osher (UCLA), Krishna Kumar (UT Austin), Markos A. Katsoulakis (UMass Amherst)
  • 分类: cs.LG, cs.AI, math.OC
  • 发表时间: 2025年10月10日
  • 论文链接: https://arxiv.org/abs/2510.09904

摘要

尽管Transformer被广泛使用，但训练深度Transformer可能不稳定。层归一化(Layer Normalization)作为标准组件能改善训练稳定性，但其位置选择往往是临时性的。本文对不同层归一化位置下Transformer的前向稳定性(隐藏状态)和后向稳定性(梯度)进行了原理性研究。理论分析揭示了训练动态的关键洞察：训练是否驱动Transformer走向规律解或病态行为。在前向稳定性方面，推导了训练后Transformer中隐藏状态增长的显式界限。在后向稳定性方面，分析了层归一化如何影响梯度反向传播，从而解释了每种层归一化位置的训练动态。分析还指导了Transformer块中残差步长的缩放，适当选择可进一步改善稳定性和性能。

研究背景与动机

问题定义

本研究要解决的核心问题是：不同层归一化位置对Transformer训练稳定性的影响机制。具体包括：

  1. 前向稳定性问题：隐藏状态在深层网络中的增长控制
  2. 后向稳定性问题：梯度在反向传播过程中的稳定性
  3. 架构设计指导：如何为新的Transformer变体提供理论指导

重要性分析

  1. 实用价值：Transformer是现代深度学习的基础，其训练稳定性直接影响模型性能和训练效率
  2. 理论缺失：现有层归一化位置选择主要基于经验，缺乏理论依据
  3. 工业需求：随着模型规模不断增大，训练稳定性问题愈发突出

现有方法局限性

  1. Post-LN：需要精细的优化调度，性能往往次优
  2. Pre-LN：虽然改善了早期训练稳定性，但会产生过大的隐藏状态，导致数值不稳定
  3. Peri-LN：虽然在实践中表现良好，但理论性质理解不足

研究动机

作者采用连续时间动态和最优控制理论的新视角，将Transformer训练问题建模为均场控制问题，从而能够分析训练收敛后的模型性质，而非仅关注初始化时的行为。

核心贡献

  1. 理论框架创新：提出基于最优控制理论的新框架，系统分析不同层归一化位置下的Transformer稳定性
  2. 前向稳定性分析：推导了隐藏状态增长的显式界限，证明Pre-LN导致无界增长而Peri-LN保持受控增长
  3. 后向稳定性分析：揭示了层归一化对梯度反向传播的影响机制
  4. 残差步长缩放：提出改进稳定性和性能的残差步长缩放方法
  5. 实验验证：在GPT-2系列模型上验证了理论发现

方法详解

任务定义

研究不同层归一化位置下Transformer的稳定性，包括：

  • 输入：嵌入和位置编码后的序列 X0∈Rd×nX_0 \in \mathbb{R}^{d \times n}
  • 输出：经过D层Transformer块后的隐藏状态 XDX_D
  • 目标：分析前向和后向传播的稳定性

连续时间建模

Transformer的连续时间表示

将标准Transformer块的跳跃连接结构解释为连续时间动态的欧拉离散化：

f_{attn}(X(t), t; \theta_{attn}(t)), & t \in [t_i, t_i + \Delta t) \\ f_{ffn}(X(t), t; \theta_{ffn}(t)), & t \in [t_i + \Delta t, t_{i+1}) \end{cases}$$ 其中 $\Delta t = \frac{T}{2D}$，$t_i = 2i\Delta t$。 #### 均场控制问题建模 将训练问题表述为连续时间均场控制问题： $$\min_\theta \mathbb{E}_{(X_0,y)} G(X(T), y)$$ $$\text{s.t. } \frac{dX(t)}{dt} = f(X(t), t; \theta(t))$$ 其中$f \in \{f_{Pre}, f_{Peri}\}$对应不同的层归一化位置。 ### 层归一化的几何性质 **关键引理1**：层归一化输出位于椭球面上 $$\mathcal{E} = \{z \in \mathbb{R}^d : (z - \beta)^T\Gamma^{-2}(z - \beta) = d\}$$ 其中 $\Gamma = \text{diag}(\gamma)$。 ### 前向稳定性分析 #### Pre-LN的无界性 **定理2**：Pre-LN训练问题的最优解在幅度上无界。 **证明思路**：通过分析Hamilton-Jacobi-Bellman (HJB) 偏微分方程，证明对应的哈密顿量不存在，导致训练问题退化。 **定理3**：即使使用权重衰减，Pre-LN Transformer的隐藏状态仍呈指数增长： $$MA(X_D) \leq (1 + C(\lambda))^D \frac{\|X_0\|_F}{\sqrt{nd}} = O(e^D)$$ #### Peri-LN的受控增长 **定理4**：Peri-LN Transformer的隐藏状态呈线性增长： $$MA(X_D) \leq \frac{\|X_0\|_F}{\sqrt{nd}} + 2D(\gamma_{max} + \beta_{max}) = O(D)$$ 方差呈二次增长： $$\text{Var}(X_D) \leq \frac{(\|X_0\|_F + 2D\sqrt{nd}(\gamma_{max} + \beta_{max}))^2}{nd - 1} = O(D^2)$$ ### 后向稳定性分析 梯度计算公式： $$\nabla_{\theta_i} G(X_D) = \nabla_{\theta_i} X_{i+1} \cdot J_{i:D} \cdot \nabla_{X_D} G(X_D)$$ 其中雅可比矩阵： $$J_{i:D} = \prod_{j=i+1}^D (I + \nabla_{X_{j-1}} f(X_{j-1}; \theta_{j-1}))$$ **命题7**：Pre-LN下，敏感性 $\nabla_{X_{j-1}} f_{Pre}$ 与激活值成正比增长。 **命题8**：Peri-LN下，敏感性 $\nabla_{X_{j-1}} f_{Peri}$ 对激活值幅度不变。 ## 实验设置 ### 数据集 - **OpenWebText数据集**：约90亿训练token，400万验证token - 使用GPT-2系列架构进行预训练 ### 模型配置 - **GPT-2** (124M参数) - **GPT-2 Large** (774M参数) - **GPT-2 XL** (1.5B参数) ### 评价指标 - **困惑度 (Perplexity)** - **ROUGE分数** (Rouge1, Rouge2, RougeL) - **BERT分数** (BertP, BertR, BertF1) - **训练稳定性**：发散运行次数统计 ### 实现细节 - 使用为Pre-LN调优的超参数，未对Peri-LN单独优化 - 残差步长缩放：$\Delta t \in \{0.1, 1\}$ - 硬件：NVIDIA H200 GPU ## 实验结果 ### 训练稳定性对比 | 层归一化设置 | 权重衰减开启 | 权重衰减关闭 | |-------------|-------------|-------------| | Pre-LN | 1/5发散 | 3/5发散 | | Peri-LN | 0/5发散 | 0/5发散 | | 无LN | 5/5发散 | — | ### 性能对比结果 **GPT-2 (124M)模型结果**： - Pre-LN ($\Delta t=1$): 验证损失5.43，困惑度247.52 - Pre-LN ($\Delta t=0.1$): 验证损失3.13，困惑度24.43 - Peri-LN ($\Delta t=1$): 验证损失3.12，困惑度24.17 - Peri-LN ($\Delta t=0.1$): 验证损失3.10，困惑度23.63 ### 隐藏状态增长分析 实验验证了理论预测： - Pre-LN在较大$\Delta t$时表现出快速增长 - Peri-LN保持更规律的线性增长 - 残差步长缩放有效控制了增长率 ### 残差步长缩放效果 1. **性能提升**：Peri-LN + $\Delta t=0.1$在所有指标上表现最佳 2. **稳定性改善**：Pre-LN在$\Delta t=0.1$时从不稳定变为稳定 3. **增长控制**：有效降低隐藏状态的均值和方差增长率 ## 相关工作 ### 层归一化研究 - **Post-LN**：最早的Transformer设计，需要精细调度 - **Pre-LN**：改善训练稳定性但产生大激活值 - **Peri-LN**：最近在大规模模型中采用，如Gemma2、OLMo2 ### 理论分析方法 - 现有工作多关注初始化时的行为或依赖经验观察 - 本文创新性地分析训练收敛后的模型性质 - 连续时间视角为架构分析提供新工具 ## 结论与讨论 ### 主要结论 1. **Pre-LN理论缺陷**：训练问题本质上是病态的，导致无界解 2. **Peri-LN优势**：提供良好定义的优化问题和受控的隐藏状态增长 3. **残差缩放价值**：简单有效的稳定性改善方法 ### 局限性 1. **简化假设**：理论分析基于连续时间近似 2. **超参数依赖**：实验使用Pre-LN调优的超参数 3. **规模限制**：实验主要在中等规模模型上进行 ### 未来方向 1. **架构筛选框架**：为新架构修改提供理论筛选准则 2. **更大规模验证**：在更大模型上验证理论发现 3. **其他归一化方法**：扩展分析到RMSNorm等变体 ## 深度评价 ### 优点 1. **理论创新性强**：首次用最优控制理论分析层归一化位置问题 2. **数学严谨性**：提供完整的理论推导和证明 3. **实用价值高**：残差步长缩放方法简单有效 4. **实验设计合理**：在多个模型规模上验证理论 ### 不足 1. **理论与实践差距**：连续时间假设与实际离散实现存在差距 2. **实验范围有限**：主要在GPT-2系列上验证，缺乏更多架构验证 3. **超参数公平性**：未对Peri-LN进行专门的超参数优化 ### 影响力评估 1. **学术贡献**：为Transformer稳定性分析提供新的理论框架 2. **实用价值**：指导实际模型设计和训练策略 3. **可复现性**：承诺公开代码和模型 ### 适用场景 1. **深度Transformer训练**：特别适用于大规模深层模型 2. **架构设计指导**：为新架构修改提供理论依据 3. **训练稳定性改善**：通过残差缩放提升训练稳定性 ## 参考文献 论文引用了多个重要工作，包括： - Ba et al. (2016): Layer Normalization原始论文 - Xiong et al. (2020): Pre-LN vs Post-LN比较研究 - Kim et al. (2025): Peri-LN的经验研究 - He et al. (2016): 残差连接的开创性工作 --- **总体评价**：这是一篇理论与实践结合良好的高质量论文，为Transformer稳定性分析提供了新的数学框架，具有重要的学术价值和实用意义。理论分析严谨深入，实验验证充分，为深度学习架构设计提供了有价值的指导。