NGPT：在超球面上进行表示学习的归一化 Transformer

在人工智能领域，神经网络架构的创新不断推动着技术的进步。最近，一篇名为 “NGPT: NORMALIZED TRANSFORMER WITH REPRESENTATION LEARNING ON THE HYPERSPHERE” 的研究论文引起了广泛关注。作者是 Ilya Loshchilov、Cheng - Ping Hsieh、Simeng Sun 和 Boris Ginsburg，他们来自 NVIDIA。这篇论文提出了一种新颖的神经网络架构 —— 归一化 Transformer（nGPT），它在超球面上进行表示学习，展现出了令人瞩目的性能优势。

1. 研究背景

Transformer 架构是现代语言模型的基础，为了提高其训练稳定性、推理成本、上下文长度和鲁棒性等，研究人员提出了大量的修改方案。其中，应用各种归一化技术被认为是有益的，例如添加 LayerNorm 和 RMSNorm 等归一化层，以及通过权重衰减控制权重的范数。同时，也有研究表明在超球面上进行表示学习与更稳定的训练、更大的嵌入空间可分性以及在下游任务中更好的性能相关。在此基础上，本文作者提出了归一化 Transformer，旨在统一该领域的各种发现和观察结果。

2. nGPT 的核心贡献

超球面上的网络参数优化

作者提出将构成网络矩阵嵌入维度的所有向量归一化，使其位于单位范数超球面上。这样，矩阵 - 向量乘法就可以看作是表示在 [-1,1] 范围内的余弦相似度的点积，从而使权重衰减变得不必要。

作为超球面上的变度量优化器

归一化 Transformer 本身在超球面上执行多步优化（每层两步），其中注意力和 MLP 更新的每一步都由特征学习率（可学习的变度量矩阵的对角元素）控制。对于输入序列中的每个标记，归一化 Transformer 的优化路径从超球面上对应其输入嵌入向量的点开始，并移动到超球面上最能预测下一个标记嵌入向量的点。

更快的收敛速度

实验表明，归一化 Transformer 将达到相同精度所需的训练步骤数减少了 4 到 20 倍（取决于序列长度）。

3. 从 GPT 到 nGPT 的演变

标记嵌入和输出逻辑

在原始的仅解码器 Transformer 中，标记嵌入向量的范数不受约束，这可能导致不准确的相似度估计。在 nGPT 中，作者提出在训练算法的每一步之后，对存储在和中的嵌入向量进行归一化。同时，由于所有 nGPT 嵌入都是归一化的，原始公式中的逻辑值代表在 [-1,1] 范围内的点积，这限制了 softmax 生成的概率分布的置信度（温度）。因此，作者引入了一个可训练的缩放参数来调整。

层和块

• 基线 Transformer：对隐藏状态应用层变换，包括交替的自注意力（ATTN）和多层感知器（MLP）块，并使用 RMSNorm 进行归一化。
• 归一化 Transformer：对于超球面上的任意两点和，可以使用 SLERP 或其近似的 LERP 来计算沿着测地线的插值。作者将其改写为 nGPT 中的更新方程，其中涉及到注意力和 MLP 块的更新方程，通过可学习的参数和以及归一化函数 Norm 来控制更新过程。与基线 Transformer 不同，nGPT 在最后一层之后不需要额外的归一化。

自注意力块

• 基线 Transformer：注意力机制是 Transformer 的关键组件，它允许每个标记关注序列中的其他标记。在基线 Transformer 中，首先使用 RMSNorm 对输入隐藏状态进行归一化，然后将其投影为查询、键和值，并应用旋转位置嵌入（RoPE）。通过计算查询和键向量的点积，缩放后应用 softmax 函数得到注意力权重，最后计算值向量的加权和。
• 归一化 Transformer：作者提出对、、和沿着其嵌入维度进行归一化，使得与计算的点积可以解释为单位范数向量之间的余弦相似度。此外，还对和进行额外的归一化，以确保每个查询和键的点积在控制范围内。同时，调整了 softmax 缩放因子。

MLP 块

• 基线 Transformer：MLP 块的输入隐藏状态首先使用 RMSNorm 进行归一化，然后通过两个单独的线性投影产生两个中间向量和，使用 SwiGLU 激活函数进行组合，最后通过一个最终的线性变换得到输出。
• 归一化 Transformer：作者提出对矩阵和沿着嵌入维度进行归一化，使得和向量分别代表与存储在和中的向量之间的余弦相似度。为了控制它们的影响，引入了缩放因子和。

有效学习率在 ADAM 中的应用

在 nGPT 中，对于任何可训练的缩放参数向量，如，使用两个标量和来控制其有效学习率。通过调整，可以在保持全局学习率不变的情况下，控制的有效学习率。

总结

将基线 Transformer 转换为归一化 Transformer 的步骤包括：移除所有归一化层；在每次训练步骤后，对所有矩阵沿着其嵌入维度进行归一化；替换更新方程；改变注意力中的 softmax 缩放因子并对和进行重新缩放和归一化；对 MLP 块的中间状态进行重新缩放；对逻辑值进行重新缩放；移除权重衰减和学习率预热。

4. 实验结果

训练加速

作者在 OpenWebText 数据集上训练了基线 Transformer（GPT）和归一化 Transformer（nGPT），并在一组标准下游任务上对它们进行了评估。实验结果表明，在不同的上下文长度和网络大小下，nGPT 的训练速度比 GPT 快 4 到 20 倍。例如，在 4k 上下文长度下，具有 10 亿参数的 nGPT 在 20k 次迭代后达到了与 GPT 在 200k 次迭代后相同的验证损失，展示了 10 倍的迭代速度提升和使用的标记数量提升。

网络参数检查

• 嵌入的范数分布：nGPT 保持嵌入的固定范数，而 GPT 的嵌入范数有显著变化。GPT 的输入嵌入具有较高的条件数，尤其是在 1B 模型中。嵌入之间的成对点积分布表明，即使在 nGPT 中，嵌入也不是均匀分布在超球面上，而是形成簇，这可能反映了语言数据中的自然模式。
• 注意力和 MLP 矩阵的条件数：GPT 模型在其注意力矩阵中表现出明显更高的条件数，与 nGPT 相比，这些矩阵可能退化为低秩矩阵，潜在地降低了这些块的学习能力。
• 特征学习率和缩放因子：注意力和 MLP 块对隐藏状态的贡献由特征学习率和控制。网络学习在和所建议的方向上只采取适度的步骤。缩放因子、和在各层之间相对稳定，它们似乎补偿了在归一化矩阵和嵌入时丢失的幅度信息。

消融研究

作者进行了大量的消融实验，结果表明，对于、、使用固定（不可学习）值，以及对于使用单个全局可学习值，只会导致准确性的轻微下降。此外，nGPT 可以在不需要对 RoPE 进行任何修改的情况下处理更长的上下文。

5. 相关工作

本文的研究与之前关于超球面表示学习的工作相关。例如，在变分自动编码器的潜在空间和用于面部验证的嵌入中，球形表示与更稳定的训练相关。同时，之前的研究也发现下游任务性能与嵌入在超球面上的对齐（紧密性）和均匀性之间存在强烈的经验相关性。作者还讨论了 nGPT 的更新与 GPT 中应用 RMSNorm 的近似关系，以及 nGPT 中 QK 归一化与之前工作的相似性。

6. 讨论与结论

这项工作建立在该领域的众多关键发现和观察结果之上，主要贡献包括对所有矩阵的嵌入维度进行归一化，以及将特征学习率从网络的其他部分解耦，使其成为可训练的参数。

通过这些创新，nGPT 作为一种变度量优化器，能够利用数据驱动的梯度信息在超球面上搜索输出解决方案。实验结果表明，nGPT 在训练速度上有显著提升，同时也为进一步探索新的算法和架构提供了基础。未来的工作可以探索将 nGPT 扩展到更大的网络规模、真实世界数据集以及更广泛的任务范围。