归一化Transformer（nGPT）：（LLM）的训练速度提高多达20倍，同时保持原有的精度水平！

英伟达团队最近提出了一种新的神经网络架构—归一化Transformer（nGPT），这一创新在人工智能领域引起了巨大轰动。

论文出自NVIDIA核心研究团队，作者包括 ：

• Ilya Loshchilov：深度学习优化领域专家，AdamW优化器的提出者
• Cheng-Ping Hsieh：专注于模型架构设计
• Simeng Sun：深度学习系统优化专家
• Boris Ginsburg：NVIDIA AI研究主管

nGPT基于超球面进行表示学习，相较于传统的Transformer架构，它能够将大型语言模型 （LLM）的训练速度提高多达20倍，同时保持原有的精度水平。

这项技术的核心在于将所有向量，包括嵌入、MLP、注意力矩阵和隐藏状态，归一化为单位范数。这意味着在输入后的token在超球面表面上移动，每一层通过“位移“来贡献最终的输出预测。

实验结果表明，nGPT在达到相同精度所需的训练步骤上减少了4到20倍，这一加速效果随着上下文长度的增加而变得更加显著。在1k上下文长度下，nGPT的训练速度提高了4倍；在4k 上下文长度下，提高了10倍；而在8k上下文长度下，训练速度暴增20倍。 Reddit上的用户对nGPT的潜力表示出了极大的兴趣，他们好奇这一技术能扩展到什么程度，并认为它可能为像o1这样的模型带来显著的训练速度优势 。

nGPT的提出，不仅仅是对现有Transformer架构的一次重大改进，更是对整个Ali练过程的一次革命。 它通过在超球面上优化网络参数，将归一化Transformer作为超球面上的可变度量优化器 ，实现了更快的收敛速度。此外，nGPT还引入了可学习的缩放参数，使得模型在训练过程中可以更灵活地调整置信度，从而做出更准确的预测。

Normalized Transformer：基于超球面表示学习的模型架构深度

1. 论文基本信息详解

论文已在arXiv上发布（编号2410.01131v1），其重要性体现在三个方面：

1. 从理论上统一了多个看似无关的优化技术
2. 提供了一个全新的视角理解Transformer的工作机制
3. 显著提升了训练效率，具有重要的工程价值

2. 研究背景深入分析

深入理解本文的价值，需要从Transformer架构的发展历程说起。自2017年提出以来，Transformer已成为现代语言模型的基石。研究界对其持续改进主要集中在以下方向：

归一化技术的演进 从最初的LayerNorm到后来的RMSNorm，归一化层的位置和形式不断优化。研究发现，合适的归一化不仅能提升性能，还能增强训练稳定性。但归一化层的引入也带来了计算开销和表达能力的潜在影响。

权重衰减的深入认识 权重衰减最初被视为一种正则化技术，但近期研究发现其作用远不止于此。通过控制权重范数，权重衰减实际上在隐式地进行参数空间的归一化。这种认识启发了对参数空间几何特性的思考。

超球面表示学习的兴起 在表示学习领域，研究者发现将特征约束在超球面上有诸多优势：

• 更好的嵌入空间可分性
• 更稳定的训练过程
• 更优的下游任务性能

Transformer作为元优化器的新理解 最近的研究揭示了一个惊人的发现：Transformer架构本身可能在执行类似梯度下降的优化过程。这为理解和改进Transformer提供了全新视角。

面对这些研究进展，现有Transformer架构仍存在几个基本问题：

嵌入向量范数问题 传统Transformer中，嵌入向量的范数完全不受约束。这导致：

• 相似度计算不准确
• 数值稳定性差
• 训练过程不稳定

隐藏状态更新机制的局限 目前的设计中，注意力和MLP块的输出简单相加到隐藏状态上，这种处理方式存在明显缺陷：

• 范数剧烈波动
• 更新步长难以控制
• 优化过程效率低下

归一化层的累赘 为了缓解上述问题，现有方法不得不引入多个归一化层，这带来了：

• 额外的计算开销
• 模型表达能力的潜在损失
• 训练参数的增加

理论框架的缺失 现有的优化方法大多是经验性的，缺乏统一的理论指导。这导致：

• 改进方法的碎片化
• 优化效果难以预测
• 进一步改进方向不明确

基于对这些问题的深入思考，本文提出了nGPT框架。这个框架最大的特点是将所有改进统一到超球面表示学习的理论框架下，不仅解决了上述问题，还提供了理解和改进Transformer的新思路。

3. 方法论详解

nGPT的设计理念建立在一个核心观察之上：如果将所有运算约束在单位超球面上，可以从根本上解决Transformer中的数值稳定性问题。下面详细展开这一框架的四个关键组成部分。

3.1 超球面参数优化机制

矩阵归一化 本文首次提出将所有网络矩阵的嵌入维度归一化到单位超球面上。这包括：

• 输入嵌入矩阵
• 输出嵌入矩阵
• 注意力机制相关矩阵 , , ,
• MLP层相关矩阵 , ,

归一化操作在每个训练步骤后执行，确保所有向量严格位于单位超球面上。这种设计带来三个直接好处：

1. 数值稳定性大幅提升
2. 计算效率显著提高
3. 模型表达能力得到保持

余弦相似度解释 在这个框架下，传统的矩阵-向量乘法获得了新的几何解释：它等价于计算向量间的余弦相似度。由于所有向量都位于单位超球面上，点积自然被限制在[-1,1]范围内。这种重新解释不只是理论上的优雅，更带来实际的计算优势：

• 避免了数值溢出风险
• 提供了自然的相似度度量
• 简化了后续的缩放操作

权重衰减的消除 一个意料之外的发现是：完全不需要权重衰减。这是因为：

• 所有参数都被约束在单位超球面上
• 范数通过直接归一化得到控制
• 优化过程变得更加简洁

3.2 变量度量优化器的精确设计

nGPT将Transformer重新解释为在超球面上进行优化的变量度量优化器。这体现在两个核心更新公式：

这里的设计蕴含着深刻的数学思想：

可学习特征学习率 和 是可学习的向量，其每个分量对应一个特征维度的学习率。这种设计允许：

• 不同特征有不同的更新步长
• 模型自适应地调整优化过程
• 在训练过程中动态调整更新策略

几何意义 更新公式可以从Riemannian优化的角度理解：

• 和 代表切空间中的搜索方向
• 和 控制在各个方向上的步长
• Norm操作将更新后的点投影回流形表面

实验发现 研究发现模型倾向于选择适度的更新步长（20%-30%），这个发现极具启发性：

• 表明模型在自动平衡探索与利用
• 证实了渐进式学习的重要性
• 为学习率设计提供了理论指导

3.3 注意力机制的深度改进

注意力机制是Transformer的核心组件，nGPT对其进行了根本性的重新设计。改进从查询和键向量的计算开始：

这个设计包含了多层深意。首先，通过Norm操作确保查询和键向量位于单位超球面上。这种归一化直接改变了注意力分数的计算性质：不再依赖向量的模长，而是纯粹反映方向的相似性。通过引入可学习的缩放因子 ，模型获得了调节注意力分布锐度的能力。

更深层的改进体现在softmax缩放因子的调整上。传统Transformer使用 作为缩放因子，这源于对未归一化向量点积方差的考虑。在nGPT中，由于向量已经归一化，点积的期望方差变为 。因此，为了恢复单位方差，缩放因子需要调整为 。这个看似简单的修改实际上保证了注意力计算的数值稳定性。

注意力权重的计算也获得了新的几何解释。在超球面上，注意力权重实际上表示了当前查询向量与各个键向量之间的测地线距离。这种解释不仅在数学上更优雅，也为理解注意力机制的工作原理提供了新视角。

3.4 MLP块的系统优化

MLP块的改进同样深入而系统。核心公式为：

这里的设计考虑了三个关键因素：

首先，引入可学习的缩放因子 和