💬 NLP

GeoNorm: Unify Pre-Norm and Post-Norm with Geodesic Optimization

本文介绍了 GeoNorm，一种通过将层输出解释为具有层级衰减的流形上的测地更新，从而统一了 Pre-Norm 和 Post-Norm 方法的新型归一化方法，在实现 Transformer 模型性能一致提升的同时，其计算开销微乎其微。

原作者： Chuanyang Zheng, Jiankai Sun, Yihang Gao, Chi Wang, Yuehao Wang, Jing Xiong, Liliang Ren, Bo Peng, Qingmei Wang, Xiaoran Shang, Mac Schwager, Anderson Schneider, Yuriy Nevmyvaka, Xiaodong Liu

发布于 2026-01-30

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Chuanyang Zheng, Jiankai Sun, Yihang Gao, Chi Wang, Yuehao Wang, Jing Xiong, Liliang Ren, Bo Peng, Qingmei Wang, Xiaoran Shang, Mac Schwager, Anderson Schneider, Yuriy Nevmyvaka, Xiaodong Liu

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

想象一下，你正试图教一个机器人写故事。为了做到这一点，机器人使用了一种名为 Transformer 的复杂机器。在这个机器内部，有很多层“工人”（比如注意力工人（attention workers）和前馈工人（feed-forward workers）），他们沿着一条链条传递信息。

每当一个工人向下一个工人传递消息时，消息都需要被“归一化”（normalized）。把归一化想象成一个 音量旋钮 或一把尺子。它确保消息在进入下一步之前，既不会太大声（过大），也不会太小声（过小）。

长期以来，工程师有两种调节这个音量旋钮的主要方法：

Post-Norm（后归一化）： 工人完成工作，将他们的消息添加到链条中，然后有人检查音量并进行调整。
Pre-Norm（前归一化）： 在工人开始工作之前，有人先检查音量，这样工人每次开始时都能得到一个大小完美的信号。

两种方法都有其优缺点。Post-Norm 有时会导致音量剧烈飙升（就像回授产生的尖叫声）；而 Pre-Norm 则可能导致早期的工人比起后期的工人喊得太响。

新的想法：GeoNorm

这篇论文的作者们说：“为什么我们只是用尺子来检查音量呢？让我们思考一下消息行进的路径。”

他们意识到，这些消息移动的方式就像是在一个 地球仪（球体）的表面上行走，而不是在平坦的地面上行走。

旧的方法（投影/Projection）： 想象你正在地球仪上行走。你沿着直线迈出一步（就像一束激光）。如果你一直走直线，你最终会从地球仪掉进太空。为了修复这个问题，旧方法会说：“噢不，你掉了！让我们直接把你拉回到表面上来。”这是一种笨拙且突兀的修正，会扭曲你的方向。
新的方法（测地线/Geodesic）： 作者们建议，与其走直线然后掉下去，不如沿着球体的 弯曲表面 行走。在数学中，曲面上最短的路径被称为 测地线（Ge geodesic，就像飞机为了从纽约飞往伦敦而飞行弧线，而不是直线一样）。

GeoNorm 用这种笨拙的“拉回地面”法，替换成了平滑的“沿曲线行走”法。它使用了一个叫做 指数映射（exponential map）的数学工具，以确保消息自然地留在“地球仪”上，沿着曲线移动，而不会掉落或需要剧烈的修正。

“衰减”技巧

论文还引入了一种处理“步长”（即机器人迈出的步子有多大）的聪明方法。

想象你在爬山。在山脚下，你可以迈出大步且自信；但随着你爬得越高，地形变得越复杂，你应该迈出更小、更谨慎的步伐。
GeoNorm 应用了这种逻辑。它在初始阶段使用较大的更新，并随着消息在网络更深层中移动，逐渐减小步长。这可以防止机器人在链条接近末端时跌跌撞撞或感到困惑。

他们发现了什么？

作者在不同的“徒步路线”（如 Arxiv 和 Books3 等数据集）以及不同规模的“机器人”（从小型到巨大的模型规模）上测试了这种新方法。

性能更好： 使用 GeoNorm 的机器人始终学习得更快，并且最终获得了更好的“分数”（更低的错误率），表现优于使用旧方法的机器人。
稳定性： 旧方法有时会出现“损失峰值”（loss spikes，即性能突然崩溃），尤其是在机器人非常深或者训练时间很长的情况下。GeoNorm 要平滑得多，且不会崩溃。
没有额外成本： 最棒的部分是？这种在地球仪上行走的新方式并不会让机器人背负更重的背包。它不会增加额外的内存消耗，也不会让计算机变慢；它只是改变了机器人 移动的方式。

总结

这篇论文认为，与其强迫 AI 的消息去适应一个盒子（旧的归一化方法），我们应该让它们沿着它们本该旅行的弯曲表面自然流动。通过这样做，AI 会变得更稳定、学习得更好，而且不需要额外的计算能力就能实现这一切。

技术摘要：GeoNorm：通过测地线优化统一 Pre-Norm 与 Post-Norm

问题陈述

Transformer 架构中归一化层的位置——即在 Pre-Norm 和 Post-Norm 之间的选择——仍然是一个影响训练稳定性、模型性能和梯度流的关键设计决策。虽然 Pre-Norm 因其稳定性已成为现代大语言模型（LLMs）的主流选择，但它会引入不平衡的梯度，导致较低层接收到的更新比较高层更大。相反，最初旨在促进深层网络训练的 Post-Norm 则容易出现训练不稳定和损失峰值（loss spikes）的问题。现有的解决方案（如 DeepNorm 和 SandwichNorm）通过参数初始化或额外的归一化层来提高稳定性，但缺乏用于分析这些机制的统一理论框架。

方法论

作者提出了 GeoNorm，一种基于流形优化的新型归一化方法。其核心理论见解是将标准的 Transformer 更新操作（ $x + \text{FFN}(x)$ 和 $x + \text{Attention}(x)$ ）重新解释为在球面上的优化步骤。在这种视角下：

Token 嵌入（Token embeddings） 充当流形上的初始状态。
Attention 和 FFN 模块 生成更新方向（估计梯度）。
归一化（Normalization） 被视为向约束集（球面）的投影回退。

标准归一化（Post-Norm）执行一个无约束更新，随后进行向球面的欧几里得投影。作者认为，这种投影丢弃了几何信息并可能扭曲更新方向。相反，GeoNorm 使用指数映射（exponential map）通过**测地线更新（geodesic update）**取代了投影步骤。

对于球面上的一个点 $x$ 和一个切向量 $v$ （投影后的更新方向），更新是沿着测地线内在地进行的：
$\text{exp}_x(v) = \cos\left(\frac{\|v\|}{\|x\|}\right)x + \|x\|\sin\left(\frac{\|v\|}{\|x\|}\right)\frac{v}{\|v\|}$

这种公式确保了表示能够沿着流形平滑移动，从而保留了特征空间的内在几何结构。此外，作者还为步长（ $\alpha_k$ ）引入了一种**层级更新衰减（layer-wise update decay）**机制，类似于学习率调度。他们评估了三种衰减策略：

调和衰减（Harmonic）： $\alpha_k = \alpha / k$
平方根衰减（Square Root）： $\alpha_k = \alpha / \sqrt{k}$
线性衰减（Linear）： $\alpha_k = \alpha (T - k) / T$

论文还从理论上证明了 Pre-Norm 可以被解释为 GeoNorm 的一个特例，即具有非自适应、依赖于深度的角度，从而将这两种方法统一在一个单一的几何框架之下。

核心贡献

理论框架： 作者提供了一个统一的理论视角，将 Transformer 层解释为球面流形上的迭代优化步骤，弥合了 Pre-Norm 和 Post-Norm 之间的鸿沟。
GeoNorm 算法： 他们引入了 GeoNorm，该方法通过指数映射使用测地线更新取代了标准的基于投影的归一化。该方法不需要额外的可训练参数，且增加的计算开销微乎其微。
层级衰减： 提出了一种用于注意力（attention）和 FFN 组件的层级更新衰减新机制，以稳定训练动态。
实证验证： 在不同模型规模（125M 到 1.3B）、数据集（Arxiv, Books3, FinWeb-Edu）和训练长度下的广泛实验证明了该方法的有效性。

实验结果

论文报告称，在所有测试配置中，GeoNorm 的表现始终优于现有的归一化方法（Pre-Norm, Post-Norm, DeepNorm, SandwichNorm）：

性能： 在 125M 模型的 Arxiv 和 Books3 数据集上，与所有基准模型相比，GeoNorm 实现了最低的验证损失（Arxiv 为 1.8792；Books3 为 3.4040）。
可扩展性： 随着模型规模的增加，性能优势持续存在并扩大。对于 1.3B 模型，GeoNorm 的损失为 3.1759，优于表现最好的基准模型（Pre-Norm 为 3.2067）。
稳定性： GeoNorm 展示了卓越的训练稳定性，特别是在长序列长度（高达 4096）的长时间训练运行中。虽然 Post-Norm 在较长长度下会出现损失峰值和训练崩溃，但 GeoNorm 保持了稳定的收敛。
衰减机制： 在所有衰减策略中，调和衰减（ $\alpha/k$ ）表现最好，在测试的数据集上建立了新的基准。

重要性与主张

论文声称 GeoNorm 可以无缝集成到标准的 Transformer 架构中，无需修改核心模块（Attention 或 FFN），也不引入额外的可训练参数。作者断言，通过利用测地线优化，GeoNorm 使表示能够在球面流形上更连贯地移动，从而改善了信息传播和训练稳定性。

这项工作旨在通过黎曼几何的角度为理解归一化机制奠定基础。它表明，Pre-Norm 的经验成功可以通过其隐式的步长衰减来解释，而 GeoNorm 使这一机制变得显式且具有自适应性。作者总结道，这种统一的视角为未来 Transformer 架构的设计（特别是将其扩展到更深和更复杂的深度）提供了稳健的基础。