💬 NLP

Gated Normalization Removal and Scale Anchoring in Pre-Norm Transformers

本文介绍了 TaperNorm，这是一种针对预归一化 Transformer 的门控方法，通过逐步移除内部归一化层来提升推理吞吐量，并揭示了最终归一化层的主要作用在于锚定预对数 logits 表示的尺度，而这一功能可通过固定目标缩放来替代，从而实现完全无归一化的模型。

原作者： Andrei Kanavalau, Carmen Amo Alonso, Sanjay Lall

发布于 2026-05-21

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Andrei Kanavalau, Carmen Amo Alonso, Sanjay Lall

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

想象一个 Transformer 模型（现代 AI 聊天机器人背后的“大脑”）是一座巨大的多层工厂。每当一条数据（一个词）穿过工厂时，它都要经过几个房间。在每一个房间里，都有一位“质量检查员”（称为归一化层），它在数据进入下一个房间之前，检查数据、调整其规模，并确保它不会变得过于失控或过于微小。

多年来，工程师们认为这些检查员在所有步骤中、任何时候都是绝对必要的。但这篇论文提出了一个简单的问题：在工厂完成训练之后，我们真的还需要这些检查员保持活跃吗？

以下是他们发现的分解，使用了简单的类比：

1. “渐隐”技巧（将检查员转变为门道）

作者创造了一种名为TaperNorm的新方法。把它想象成质量检查员的调光开关。

训练期间：在开始时，检查员努力工作，像往常一样检查每一条数据。
过渡阶段：随着训练结束，作者逐渐“调暗”检查员。他们并没有直接解雇他们，而是教导他们停止检查数据，而是让数据通过一个简单的固定规则（例如“乘以 1.5"）直接通过。
结果：到了最后，检查员消失了。他们被一个简单的静态门道所取代。因为这个门道只是一个固定规则，工厂可以将其“折叠”到下一个房间的机器中。这意味着计算机不再需要执行任何额外的数学运算来检查数据；它只需将数据传递下去。

这为什么重要？
该论文在小型模型（TinyStories）和一个著名模型（GPT-2）上测试了这一点。他们发现，移除这些内部检查员并没有显著损害模型的智能（性能仅出现微小下降）。然而，由于移除了检查员，工厂变得更快了。

速度提升：当他们测试模型生成文本的速度时，速度提高了1.18 倍。这就像一辆汽车仅仅通过移除几个不必要的减速带，就从每小时 60 英里提速到了每小时 70 英里。

2. “锚点”问题（为什么最后一个检查员必须保留）

这是最有趣的发现。虽然他们可以移除工厂中间的检查员，但他们发现最后一个检查员（就在 AI 给出最终答案之前）扮演着特殊且独特的角色。

锚点的类比：
想象 AI 的最终想法是一个在风中飘浮的气球。

有最后一个检查员时：检查员就像一个锚。它将气球固定在特定高度。无论风刮得多大（数学试图让答案更加自信），气球都保持在稳定的水平。这保持了 AI 的稳定性。
没有最后一个检查员时：如果你移除了那个最后的锚，气球就会自由漂移。AI 的数学运算自然会试图让气球飞得越来越高（使答案更加“自信”或极端），仅仅是为了降低其误差分数。这被称为**“Logit 追逐”**。AI 变得不稳定，因为它不断无限地膨胀自己的自信。

解决方案：
如果你必须移除最后一个检查员（为了让模型更快），你就必须用别的东西来替代这个锚。作者使用了一种**“固定目标尺度损失”**。

隐喻：想象一根绳索，如果气球试图飘得太高或沉得太低，它会轻轻地将气球拉回特定高度。这根绳索充当了“虚拟锚”，阻止 AI 发疯，从而允许他们安全地移除最后一个检查员。

3. 底线

该论文得出了两个主要结论：

内部检查员是可选的：你可以用检查员训练 AI，然后在最后将它们转变为简单的门道。这使 AI 更快（在他们的测试中快达 18%），且几乎不会损失智能。
最后一个检查员很特殊：最后的检查至关重要，因为它阻止 AI 变得“过度自信”和不稳定。如果你移除了它，你必须添加一根“绳索”（一个特定的数学规则）来约束 AI 的自信。

简而言之：你可以剔除中间环节以加快 AI 的运行速度，但你必须小心最后一步，否则 AI 可能会因自身的自信而失控。

技术摘要：Pre-Norm Transformer 中的门控归一化移除与尺度锚定

问题陈述

归一化层（如 LayerNorm、RMSNorm）是 Transformer 架构中的标准组件，能够稳定激活值并支持深度网络的训练。然而，这些层引入了依赖于样本的非线性计算，在推理过程中造成了显著开销。目前尚不清楚这些逐 token 的统计量在整个推理阶段是否严格必要，特别是在 Pre-Norm Transformer 中。尽管先前的工作试图通过初始化技巧或顺序微调来移除归一化，但人们对内部归一化层与最终归一化层在功能上的区别缺乏理解，特别是关于它们如何控制预 logits 表示的尺度。

方法论

作者提出了TaperNorm，这是一种专为 Pre-Norm Transformer 设计的门控归一化移除方法。该方法允许模型在训练期间从标准归一化过渡到学习的、与样本无关的线性或仿射映射，从而在推理时移除归一化操作。

核心机制：TaperNorm

TaperNorm 使用全局门控 $g \in [0, 1]$ 在归一化分支和固定缩放分支之间进行动态插值。对于 token 向量 $h$ ：
$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$
其中：

$r(h)$ 是 RMSNorm 缩放因子。
$D\gamma$ 和 $D\tilde{\gamma}$ 是可学习的对角缩放矩阵。
$c$ 是每层的标量。
$g$ 是全局门控，初始值为 1（标准归一化），并在训练过程中按余弦衰减至 0。

训练过程：

热身（Warmup）： 门控 $g$ 保持在 1。模型作为标准 Pre-Norm Transformer 进行训练。
校准（Calibration）： 在热身结束时，计算标量 $c$ 以最小化归一化分支与缩放分支之间的最小二乘差异，确保过渡平滑且无分布偏移。
** tapering（衰减）：** 门控 $g$ 衰减至 0。缩放分支的参数（ $\tilde{\gamma}$ ）保持可训练。
推理折叠（Inference Folding）： 一旦 $g=0$ ，TaperNorm 层变为简单的线性/仿射映射（ $c \cdot h D\tilde{\gamma}$ ）。该操作在数学上可以折叠到相邻的线性投影中（例如 QKV 投影或 MLP 权重），从而从推理图中移除显式的归一化操作。

尺度锚定分析

该论文指出了内部归一化层与最终归一化层之间的关键区别：

最终归一化作为锚点： 最终归一化层充当隐式的尺度锚点。从数学上讲，它是 0 次齐次的，意味着它移除了输出处的径向梯度分量。如果没有这个锚点，最小化交叉熵可能会驱使模型无限增加 logits 的幅度（"logit chasing"）以降低损失，从而导致不稳定。
显式替代方案： 当移除最终归一化时（在"All-Taper"设置中），作者引入了固定目标尺度损失（ $L_{aux}$ ）。该辅助损失惩罚 token 尺度与固定目标 $s_{tgt}$ 的偏差，提供了一种显式的径向恢复力，以抵消 logit 追逐倾向。

测试了两种实验配置：

Internal-Taper（内部衰减）： 内部块归一化器被衰减为线性映射，但最终归一化保持激活状态（隐式锚点）。
All-Taper（全部衰减）： 所有归一化器（包括最终归一化器）均被衰减。这需要固定目标尺度损失来维持稳定性（显式锚点）。

主要贡献

门控归一化移除： 一种将 RMSNorm/LayerNorm 模块渐变为固定线性/仿射重缩放的方法，允许通过权重折叠在推理时消除逐 token 统计量。
输出尺度锚定理论： 一项分析表明，最终归一化层的主要独特作用是移除径向梯度。论文显示，如果没有这个锚点，交叉熵会导致 logits 无界增长，而固定目标尺度损失可以有效地替代这一稳定作用。
实证验证： 在 TinyStories 预训练和 GPT-2 微调上的实验表明，内部归一化可以被移除，且验证损失增加极小。在折叠内部缩放后，该方法在 KV 缓存的自回归解码中实现了高达1.18 倍的吞吐量提升。

结果

预训练（TinyStories）：
- Internal-Taper (+aux)： 在不同模型规模（100 万到 3000 万参数）下，验证损失保持在基线的 0.7–1.5% 以内。
- All-Taper (+aux)： 损失保持在基线的 0.8–1.8% 以内，证明了具有显式尺度锚定的完全无归一化模型是可行的。
- 动态特性： 没有尺度锚点的模型表现出 logits 幅度的快速增长（"logit chasing"）。固定目标尺度损失成功抑制了这种行为，使梯度幅度与锚定基线保持一致。
推理吞吐量：
- 在 NVIDIA H100 上，Internal-Taper (fused) 配置相比标准 RMSNorm 基线实现了1.15–1.18 倍的每秒 token 吞吐量提升。Internal-Taper (unfused) 配置（显式缩放操作）实现了1.10–1.14 倍的加速。
微调（GPT-2）：
- 提出的 Internal-TaperLN 和 All-TaperLN 方法在 OpenWebText 和 The Pile 数据集上，实现了与"FakeLN"方法（Baroni et al., 2025）及标准 Vanilla 模型相当的交叉熵损失。
- 与顺序移除方法不同，TaperNorm 使用单个全局门控同时移除所有选定层。

意义与主张

该论文声称，内部和最终归一化层在 Pre-Norm Transformer 中扮演着不同的角色。虽然内部归一化器可以被学习的、与样本无关的线性映射所替代，且性能下降可忽略不计，但最终归一化层作为输出尺度锚点发挥着独特的作用。

这项工作的意义在于：

效率： 它提供了一条从推理图中移除计算昂贵的归一化操作的实用途径，无需从头重新训练即可获得可测量的吞吐量提升（高达 1.18 倍）。
理论洞察： 它阐明了最终归一化层提供的稳定性主要归因于其移除了径向梯度。这一见解使得通过用显式固定目标尺度损失替代隐式锚点来设计“完全无归一化”模型成为可能。
鲁棒性： 该方法在训练早期保持标准的 Pre-Norm 动态，仅在后期过渡到线性映射，这似乎比试图从一开始就在无归一化条件下训练的方法更稳定。

作者在范围上保持谦逊，指出实验仅限于高达 3000 万参数的模型和特定的硬件栈。他们承认，在十亿参数规模下对激活异常值的鲁棒性，以及该方法在交叉熵之外其他目标上的适用性，仍是未来工作的开放性问题。