在超球面上进行表示学习的归一化 Transformer
- 1. 研究背景
- 2. nGPT 的核心贡献
- 超球面上的网络参数优化
- 作为超球面上的变度量优化器
- 更快的收敛速度
- 3. 从 GPT 到 nGPT 的演变
- 标记嵌入和输出逻辑
- 层和块
- 自注意力块
- MLP 块
- 有效学习率在 ADAM 中的应用
- 总结
- 4. 实验结果
- 训练加速
- 网络参数检查
- 消融研究
- 5. 相关工作
- 6. 讨论与结论
在人工智能领域,神经网络架构的创新不断推动着技术的进步。最近,一篇名为 “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 扩展到更大的网络规模、真实世界数据集以及更广泛的任务范围。