一句话总结

NVIDIA 研究人员提出了一种针对归一化 Transformer (nGPT) 的训练配方，该配方整合了 Logit Gradient Preconditioning、对数学习率衰减、GatedAdamW、角度更新控制以及可选的探索机制，使得一个 14B 参数的混合 Mamba-2–Transformer 混合专家模型能够用大约一半的训练 token 达到与未归一化的 AdamW 训练模型相同的验证损失，并且该配方可扩展到总参数量高达 14B 的模型。

核心贡献

• 本文将归一化 Transformer (nGPT) 扩展到现代混合 Mamba-2–Transformer 混合专家架构，为这些模型建立了一套训练配方。
• 该配方引入了 Logit Gradient Preconditioning、对数学习率衰减、GatedAdamW 以及角度更新控制机制，以实现稳定的超球面表示学习。
• 在 14B 参数模型上，nGPT 配方用大约一半的训练 token 就达到了与未归一化的 AdamW 基线相同的验证损失，并且该方法可扩展到总参数量高达 14B 的模型。

引言

归一化 Transformer (nGPT) 引入了超球面表示，其中所有激活向量和参数向量都被约束为单位范数，将下一个 token 预测重新定义为超球面上的优化过程。尽管 nGPT 在密集 Transformer 训练中展现了优势，但它尚未适应现代混合架构，例如 Mamba-2–Transformer 混合专家 (MoE) 模型，而这些模型对于高效扩展大型语言模型至关重要。作者使用 Nemotron-3 架构和工业级数据混合将 nGPT 扩展到混合 MoE 设置，并开发了一套全面的训练配方，其中包括 Logit Gradient Preconditioning 和 GatedAdamW 等组件。该配方提高了训练归一化 MoE 模型的数据效率，不过研究侧重于有限计算预算下的实际结果。

方法

作者引入了几种优化要素和架构修改，以有效训练归一化 GPT (nGPT) 模型，特别是针对现代混合专家 (MoE) 和 Mamba-2 混合架构。

首先，他们处理 logit 尺度向量的行为。在 nGPT 中，输出嵌入向量和隐藏状态被归一化，因此未缩放的 logit 基于有界的点积。一个可学习的词汇级尺度向量 $\pmb{s}_z \in \mathbb{R}^V$sz​∈RV 控制输出分布的锐度。前向传播计算缩放后的 logit 为 $\pmb{z} = \pmb{s}_z \odot \pmb{u}$z=sz​⊙u，其中 $\pmb{u}$u 是未缩放的 logit。作者观察到，在长时间预训练运行中，$\pmb{s}_z$sz​ 的均值呈对数增长。

这种增长充当了一个隐藏的、时变的全局乘子，作用于通过输出层传播的梯度，使优化器设计复杂化，并导致梯度范数的增长。为了控制这种耦合，作者提出了 Logit Gradient Preconditioning (LGP)。虽然 $\pmb{s}_z$sz​ 在前向传播中使用，但其显式的反向乘子被替换为一个归一化尺度：

$$\frac{\partial \mathcal{L}}{\partial \pmb{u}} \leftarrow \left(\frac{\pmb{s}_z}{\mathrm{mean}(\pmb{s}_z)}\right)^q \odot \frac{\partial \mathcal{L}}{\partial \pmb{z}}$$∂u∂L​←(mean(sz​)sz​​)q⊙∂z∂L​

其中 $q$q 控制预条件强度。设置 $q=1$q=1 会移除显式的时变全局乘子，同时保留 $\pmb{s}_z$sz​ 中相对的词汇级变化。

接下来，作者提出了一种对数学习率衰减计划。由于 nGPT 约束了参数向量的范数，并且不对这些向量使用权重衰减，因此有效步长更直接地由学习率计划控制。在可选的预热之后，无量纲的计划乘子在 $\eta_{\max}$ηmax​ 和 $\eta_{\min}$ηmin​ 之间进行对数插值：

$$\eta(t) = \eta_{\min} + (\eta_{\max} - \eta_{\min}) \left[ 1 - \frac{\log\left(1 + \frac{r}{\rho}\right)}{\log\left(1 + \frac{1}{\rho}\right)} \right]$$η(t)=ηmin​+(ηmax​−ηmin​)​1−log(1+ρ1​)log(1+ρr​)​​

其中 $r$r 是归一化的衰减进度，$\rho > 0$ρ>0 控制形状。

较小的 $\rho$ρ 值会将更多的学习率预算分配在早期，在预热后产生急剧下降和长尾，而较大的 $\rho$ρ 则趋近于线性衰减。

作者还引入了 GatedAdamW 来改进自适应梯度更新。在标准 AdamW 中，常数 $\epsilon$ϵ 同时用于数值稳定性和在二阶矩估计较小时抑制更新。GatedAdamW 通过使用 $\epsilon_{\mathrm{num}}$ϵnum​ 保证数值稳定性，并使用一个单独的软门来控制小二阶矩尺度下的更新，从而分离了这些角色。坐标级门定义为：

$$\gamma_t = \sigma\left(a \log \frac{\pmb{d}_t}{\epsilon_{\mathrm{gate}}}\right)$$γt​=σ(alogϵgate​dt​​)

其中 $\pmb{d}_t = \sqrt{\hat{\pmb{v}}_t} + \epsilon_{\mathrm{num}}$dt​=v^t​​+ϵnum​，$\sigma$σ 是 sigmoid 函数，$a$a 控制锐度。参数更新变为：

$$\pmb{\theta}_t = \pmb{\theta}_{t-1} - \eta_t \left(\alpha \pmb{\gamma}_t \odot \frac{\hat{\pmb{m}}_t}{\pmb{d}_t} + \lambda \pmb{\theta}_{t-1}\right)$$θt​=θt−1​−ηt​(αγt​⊙dt​m^t​​+λθt−1​)

较小的 $a$a 值会产生更平滑的过渡，对低于阈值的坐标抑制较弱。此外，作者对归一化的非嵌入参数向量应用了 Angular Step Cap (ASC)。这通过限制候选向量相对于前一个归一化向量的角位移来保护几何步长，确保超球面上的稳定更新。

最后，作者详细说明了训练 nGPT 与混合 Mamba-2 和 MoE 模型所需的具体修改。所有构成矩阵的参数向量通过沿嵌入维度归一化被约束为单位范数，并移除了 RMSNorm 和 LayerNorm 等标准归一化层。当激活与流重新组合时，使用可学习的特征学习率将其归一化为单位范数。为了防止归一化向量的乘积变得过于受限，可训练的缩放向量重新缩放归一化的查询和键，以及中间的 MLP/MoE 激活。特别是对于 Mamba-2 和 MoE 块，引入了可训练的标量来重新缩放输入激活，将非线性函数（如 SiLU 和 sigmoid 路由器）的输入置于合适的尺度。完整的配方包括移除权重衰减、应用不带权重衰减的 GatedAdamW、调整注意力中的 softmax 缩放因子，以及实现前述的 logit gradient preconditioning 和对数学习率衰减。

实验

评估在 Nemotron-3 Nano 模型（1B–14B 参数）的规模阶梯上，使用相同的数据和训练 token 预算，将归一化 GPT (nGPT) 与标准 GPT 进行比较。nGPT 始终实现了约 2.5% 更低的训练和验证损失，并且用大约一半的训练 token 就达到了相当的损失水平，改进主要归功于归一化参数化和训练配方，而非仅靠门控优化器。消融研究表明，GatedAdam 的门控机制额外节省了约 15–20% 的 token，而其他提出的组件对验证损失只有微小的影响。

在从 1B 到 14B 参数的模型规模上，带有 GatedAdam 的归一化 Transformer (nGPT) 始终比使用 AdamW 训练的基线 GPT 获得更低的训练和验证损失，相对降低约 2.5%。14B 模型显示出与缩放趋势的轻微偏离，可能是由于参数分配差异和较短的训练 token 周期。进一步分析表明，nGPT 可以用大约一半的训练 token 达到基线的损失，整体改进主要源于归一化参数化和训练配方，而非仅靠门控机制。在所有模型规模上，nGPT 的训练和验证损失比 GPT 基线低约 2.5%，百分比差异范围在 2.1% 到 3.4% 之间。14B 模型的损失降低与整体缩放趋势略有偏离，可能是因为各组件间的参数分配不同，以及相对于其规模的 token 周期较短。

在从 1B 到 14B 参数的模型规模上，带有 GatedAdam 的归一化 Transformer (nGPT) 始终比使用 AdamW 训练的基线 GPT 获得更低的训练和验证损失，相对降低约 2.5%。改进主要源于归一化参数化和训练配方，而非仅靠门控机制，并且 nGPT 可以用大约一半的训练 token 达到基线的损失。14B 模型显示出与缩放趋势的轻微偏离，可能是由于参数分配差异和较短的训练 token 周期。