One-sentence Summary

NVIDIA researchers propose a training recipe for the normalized Transformer (nGPT) that integrates Logit Gradient Preconditioning, Logarithmic Learning Rate Decay, GatedAdamW, angular update control, and optional exploration mechanisms, enabling a 14B-parameter hybrid Mamba-2–Transformer Mixture-of-Experts model to reach the same validation loss using approximately half the training tokens as an unnormalized AdamW-trained counterpart, with the recipe scaling across models up to 14B total parameters.

Key Contributions

• The paper extends the normalized Transformer (nGPT) to modern hybrid Mamba-2–Transformer Mixture-of-Experts architectures, establishing a training recipe for these models.
• The recipe introduces Logit Gradient Preconditioning, logarithmic learning rate decay, GatedAdamW, and angular update control mechanisms to enable stable hyperspherical representation learning.
• On a 14B-parameter model, the nGPT recipe reaches the same validation loss with approximately half the training tokens compared to an unnormalized AdamW baseline, and the approach scales to models with up to 14B total parameters.

Introduction

The normalized Transformer (nGPT) introduced hyperspherical representations where all activation and parameter vectors are constrained to unit norm, reframing next-token prediction as an optimization process on the hypersphere. While nGPT demonstrated benefits for dense Transformer training, it had not been adapted to modern hybrid architectures such as Mamba-2–Transformer mixture-of-experts (MoE) models, which are critical for scaling large language models efficiently. The authors extend nGPT to a hybrid MoE setting using the Nemotron-3 architecture and an industry-grade data blend, and they develop a comprehensive training recipe that includes components like Logit Gradient Preconditioning and GatedAdamW. This recipe improves data efficiency for training normalized MoE models, though the study focuses on practical results under a limited compute budget.

Method

The authors introduce several optimization ingredients and architectural modifications to train normalized GPT (nGPT) models effectively, particularly for modern Mixture of Experts (MoE) and Mamba-2 hybrid architectures.

First, they address the behavior of the logit scale vector. In nGPT, output embedding vectors and hidden states are normalized, so unscaled logits are based on bounded dot products. A learnable vocabulary-wise scale vector $\pmb{s}_z \in \mathbb{R}^V$sz​∈RV controls the sharpness of the output distribution. The forward pass computes scaled logits as $\pmb{z} = \pmb{s}_z \odot \pmb{u}$z=sz​⊙u, where $\pmb{u}$u are the unscaled logits. The authors observe that during long pretraining runs, the mean value of $\pmb{s}_z$sz​ grows logarithmically.

This growth acts as a hidden, time-varying global multiplier on the gradient propagated through the output layer, complicating optimizer design and contributing to the growth of gradient norms. To control this coupling, the authors propose Logit Gradient Preconditioning (LGP). While $\pmb{s}_z$sz​ is used in the forward pass, its explicit backward multiplier is replaced by a normalized scale:

where $q$q controls the preconditioning strength. Setting $q=1$q=1 removes the explicit time-varying global multiplier while preserving the relative vocabulary-wise variation in $\pmb{s}_z$sz​.

Next, the authors propose a Logarithmic Learning Rate Decay schedule. Because nGPT constrains parameter-vector norms and does not use weight decay for these vectors, the effective step size is more directly controlled by the learning rate schedule. After an optional warmup, the dimensionless schedule multiplier interpolates logarithmically between $\eta_{\max}$ηmax​ and $\eta_{\min}$ηmin​:

where $r$r is the normalized decay progress and $\rho > 0$ρ>0 controls the shape.

Smaller values of $\rho$ρ allocate more of the learning rate budget early, producing a sharp decrease after warmup and a long tail, while large $\rho$ρ approaches linear decay.

The authors also introduce GatedAdamW to refine the adaptive gradient update. In standard AdamW, the constant $\epsilon$ϵ serves both numerical stability and the suppression of updates when the second-moment estimate is small. GatedAdamW separates these roles by using $\epsilon_{\mathrm{num}}$ϵnum​ for numerical stability and a separate soft gate to control updates at small second-moment scales. The coordinate-wise gate is defined as:

where $\pmb{d}_t = \sqrt{\hat{\pmb{v}}_t} + \epsilon_{\mathrm{num}}$dt​=v^t​​+ϵnum​, $\sigma$σ is the sigmoid function, and $a$a controls the sharpness. The parameter update becomes:

Smaller values of $a$a produce a smoother transition, suppressing coordinates below the threshold less strongly. Additionally, the authors apply an Angular Step Cap (ASC) to normalized non-embedding parameter vectors. This safeguards the geometric step by capping the angular displacement of the candidate vector relative to the previous normalized vector, ensuring stable updates on the hypersphere.

Finally, the authors detail the specific modifications required to train nGPT with hybrid Mamba-2 and MoE models. All parameter vectors forming matrices are constrained to unit norm by normalizing them along the embedding dimension, and standard normalization layers like RMSNorm and LayerNorm are removed. Activations are normalized to unit norm when recombined with the stream, using learnable eigen learning rates. To prevent products of normalized vectors from becoming too constrained, trainable scaling vectors rescale normalized queries and keys, as well as intermediate MLP/MoE activations. For Mamba-2 and MoE blocks specifically, trainable scalars are introduced to rescale input activations, placing inputs to nonlinearities like SiLU and the sigmoid router at appropriate scales. The complete recipe involves removing weight decay, applying GatedAdamW without weight decay, adjusting the softmax scaling factor in attention, and implementing the aforementioned logit gradient preconditioning and logarithmic learning rate decay.

Experiment

The evaluation compares a normalized GPT (nGPT) against a standard GPT across a scaling ladder of Nemotron-3 Nano models (1B–14B parameters) using the same data and training token budgets. nGPT consistently achieves approximately 2.5% lower training and validation losses and reaches comparable loss levels with roughly half the training tokens, with the normalized parameterization and training recipe driving the improvement rather than the gated optimizer alone. Ablation studies indicate that GatedAdam's gating contributes a modest additional token savings of 15–20%, while other proposed components have only minor effects on validation loss.

Across model sizes from 1B to 14B parameters, the normalized transformer (nGPT) with GatedAdam consistently achieves lower training and validation losses than the baseline GPT trained with AdamW, with relative reductions around 2.5%. The 14B model shows a slight deviation from the scaling trend, likely due to differences in parameter allocation and a shorter training token horizon. Further analysis reveals that nGPT can match the baseline's loss using roughly half the training tokens, and the overall improvement stems primarily from the normalized parameterization and training recipe rather than the gating mechanism alone. nGPT yields training and validation losses that are approximately 2.5% lower than the GPT baseline across all model sizes, with the percentage difference ranging from 2.1% to 3.4%. The 14B model's loss reduction deviates slightly from the overall scaling trend, potentially because of a different allocation of parameters across components and a shorter token horizon relative to its size.

The normalized transformer (nGPT) with GatedAdam achieves consistently lower training and validation losses than the baseline GPT trained with AdamW across model sizes from 1B to 14B parameters, with relative reductions around 2.5%. The improvement stems primarily from the normalized parameterization and training recipe rather than the gating mechanism alone, and nGPT can match the baseline's loss using roughly half the training tokens. The 14B model shows a slight deviation from the scaling trend, likely due to differences in parameter allocation and a shorter training token horizon.