🤖 AI

Stability of Transformers under Layer Normalization

本文对不同层归一化（layer normalization）放置方式如何影响 Transformer 的前向与反向稳定性进行了原则性的理论与数值分析，推导出了关于隐藏状态增长与梯度传播的显式界限，旨在为改进训练动态的架构设计与残差缩放提供指导。

原作者： Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, Markos A. Katsoulakis

发布于 2026-08-10

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Kelvin Kan, Xingjian Li, Benjamin J. Zhang, Tuhin Sahai, Stanley Osher, Krishna Kumar, Markos A. Katsoulakis

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个机器人写故事。你给了它一个由数千个微小层级组成的大脑，就像一座拥有数百层楼高的摩天大楼。每当机器人读到一个词，它都会将信息通过电梯向上传递，每经过一层，都会改变一点点它的理解。这就是现代人工智能（被称为“Transformer”）的工作方式。它们是当今最智能的聊天机器人和图像生成器背后的引擎。但问题在于：随着这些建筑变得越来越高，它们会变得不稳定。有时，信息在上升的过程中会发生如此严重的扭曲，以至于机器人开始发出乱码般的数字尖叫，或者用于教导它的数学逻辑变得异常狂野，导致整个系统崩溃。为了阻止这种情况，工程师们使用了一种叫做“层归一化”（Layer Normalization）的安全装置。把它想象成减速带或汽车发动机上的限速器；它能防止信息失控加速。多年来，人们只是在猜测应该把这些减速带放在哪里——有时在引擎之前，有时在引擎之后——然后听天由命。但没有人真正知道为什么一个位置比另一个位置更好，或者这个机器人是否正在秘密建造一座会在自身重量下坍塌的纸牌屋。

这篇论文就像一群侦探，他们决定不再靠猜，而是开始利用物理定律来解开这个谜团。作者们将 AI 的学习过程视为一场时间之旅，使用了被称为“最优控制理论”的分支数学（这通常用于确定如何最好地驾驶火箭或无人机）。他们问道：“如果我们想让这个机器人完美地学习，信息在向塔顶传输时会发生什么？”他们发现，旧的放置安全装置的方法（被称为“Pre-LN”）实际上是一个陷阱。他们的数学证明了，即使机器人竭尽全力去学习，信息最终也会变得巨大无比，从而导致混乱的局面。这就像试图用消防水柱去填满一个水桶；无论你如何努力瞄准，水都会溢出来。

然而，他们发现了一种更好的方法。通过将安全装置同时移动到每一层的入口和出口处（一种他们称为“Peri-LN”的新方法），信息就能保持冷静且受控。他们的方程显示，数据不再是爆炸式增长，而是呈现出一种可预测的、温和的线性增长，就像一株表现良好的植物。他们还发现了一个让这一切变得更好的简单技巧：放慢机器人在每一层所走的步伐。想象一下，如果机器人迈出的是细小、谨慎的步子，而不是巨大的跨步，这会让整座建筑保持稳固。当他们在真实的 AI 模型上进行测试时，结果与他们的数学推导完美契合。“Peri-LN”模型没有崩溃，而且学习效果与旧模型一样好，甚至更好。这篇论文不仅仅是在说“这感觉是对的”；它提供了一个数学保证，证明这种新设计能让 AI 保持稳定，为构建下一代超级智能机器提供了一份清晰的蓝图，使其不会分崩离析。

技术摘要：层归一化（Layer Normalization）下 Transformer 的稳定性

问题陈述

尽管 Transformer 在现代深度学习中占据主导地位，但训练深层 Transformer 仍然容易出现不稳定性，尤其是在模型深度增加时。这些不稳定性体现在两个关键领域：前向不稳定性（隐藏状态不受控制地增长导致数值问题）和后向不稳定性（反向传播过程中的梯度爆炸）。

缓解这些问题的一个主要架构组件是层归一化（Layer Normalization, LN）。然而，LN 在 Transformer 模块中的位置演变很大程度上是基于经验启发式方法，而非严谨的理论证明。本文识别了三种主要的配置方式：

Post-LN：在残差连接之后应用 LN。通常需要精细的优化调度，并表现出次优的性能。
Pre-LN：在注意力机制和前馈模块之前应用 LN。虽然作为 Post-LN 的改进方案，它在提高训练稳定性方面被视为标准做法，但经验观察表明，它会产生过大的隐藏状态，从而导致深层网络中潜在的数值不稳定和梯度爆炸。
Peri-LN：一种较新的配置，对模块的输入和输出同时进行 LN。经验研究表明，与 Pre-LN 相比，它能提供更优越的稳定性和更规则的隐藏状态，但其理论特性仍未得到充分理解。

核心问题在于缺乏一个原则性的理论框架来解释为什么这些不同的放置方式会导致截然不同的训练动态和稳定性特征，以及训练究竟是驱动模型走向规则解还是病态行为。

方法论

作者提出了一个基于离散时间最优控制理论的新型理论框架，用于分析 Transformer 的训练动态。

建模为最优控制：将 Transformer 训练问题构建为一个离散时间平均场控制问题。层索引 $i$ 代表离散时间，隐藏状态 $X_i$ 代表系统状态。残差连接充当动力学的速度。
Hamilton-Jacobi-Bellman (HJB) 分析：作者使用 HJB 偏微分方程（PDE）结合质量传输方程来分析训练问题的最优性条件。
- 他们研究了哈密顿量（Hamiltonian），该量作为一个反馈控制机制，定义了隐藏状态的最优速度场。
- 对于 Pre-LN，他们证明了哈密顿量是无界的，导致最优性条件病态，使得隐藏状态可以无限增大。
- 对于 Peri-LN，输出归一化将速度场约束在一个椭球体内（基于引理 1），确保了哈密顿量是有界的且问题是适定的。
梯度分析：通过研究跨层的局部敏感度（雅可比矩阵）乘积，分析了梯度的后向稳定性。
残差缩放：该框架通过引入时间增量参数 $\Delta t < 1$ 进行了扩展，有效地缩放残差步的幅度以控制信号传播。

核心贡献

1. 已训练模型的理论框架

不同于以往通常关注模型初始化或依赖经验证据的工作，该框架分析的是收敛时的已训练模型。它提供了对训练动态是导致规则解还是病态行为的系统性评估。

2. 对 LN 放置方式的严谨分析

Pre-LN 不稳定性：论文证明，对于 Pre-LN，训练问题的最优解在量级上是无界的。即使使用权重衰减，隐藏状态也会呈指数级增长（ $O(e^D)$ ），且梯度随激活值同步增长，导致训练不稳定。
Peri-LN 稳定性：作者推导出了显式界限，证明 Peri-LN 保证了隐藏状态绝对值的线性增长（ $O(D)$ ）以及方差的二次增长（ $O(D^2)$ ）。至关重要的是，研究表明 Peri-LN 的局部敏感度对激活值的大小具有不变性，从而防止了梯度爆炸。

3. 不确定性量化

分析表明，Peri-LN 能够实现从输入到最终隐藏状态的可控不确定性传播（通过 Wasserstein 距离衡量），而 Pre-LN 则可能将输入不确定性放大到无界水平。

4. 残差步缩放

在稳定性分析的指导下，作者引入了残差步缩放机制（ $\Delta t < 1$ ）。他们在理论上证明了这种缩放：

强化了隐藏状态增长和输出不确定性的界限。
通过缩放局部敏感度来缓解梯度爆炸。
在不增加额外计算成本的情况下，同时提高了前向和后向稳定性。

实验结果

作者使用在 OpenWebText 数据集上训练的 GPT-2 变体（124M 至 1.5B 参数）验证了其理论发现。

训练稳定性：在 5 次试验中，当启用权重衰减时，Pre-LN 模型在 1 次试验中发散；当禁用权重衰减时，在 3 次试验中发散。相比之下，Peri-LN 模型在两种权重衰减设置下的 5 次试验中均保持稳定。
性能：即使在针对 Pre-LN 进行超参数调优的情况下，Peri-LN 仍取得了与 Pre-LN 相当的性能。
缩放的影响：将 Peri-LN 与残差缩放（ $\Delta t = 0.1$ ）结合使用，在所有指标（验证损失、困惑度、ROUGE、BERT 分数）上始终取得了最佳性能，并有效控制了跨层隐藏状态矩（绝对值均值和方差）的增长，这与理论预测一致。

重要性与主张

本文声称提供了对 Transformer 中层归一化放置现象的首次严谨理论证明。

解释不稳定性：它解释了为什么 Pre-LN 会通过底层 HJB PDE 的病态性导致无界隐藏状态和梯度爆炸。
证明 Peri-LN 的合理性：它为 Peri-LN 的稳定性提供了理论基础，表明其维持了受控的增长和不变的敏感度。
设计指导：该框架提供了一种原则性的方法，可以在进行昂贵的经验训练之前对新的架构修改进行“合理性检查”。
实际改进：残差步缩放被呈现为一种简单的、无成本的修改，它显著增强了稳定性和性能，并通过实验验证了带有缩放的 Peri-LN 优于标准配置。

作者强调，其工作弥合了经验启发式方法与理论理解之间的鸿沟，提供了一个连续时间视角，从而产生了既有理论洞察力又有实际架构改进的研究成果。