🤖 machine learning

Training nGPT

This paper presents a practical training recipe for the normalized Transformer (nGPT) that incorporates techniques like Logit Gradient Preconditioning and GatedAdamW, enabling 14B-parameter hybrid Mamba-Transformer MoE models to achieve the same validation loss as their unnormalized counterparts using approximately half the training tokens.

Original authors: Ilya Loshchilov, Boris Ginsburg

Published 2026-08-04

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Ilya Loshchilov, Boris Ginsburg

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine the world of artificial intelligence as a giant, bustling library where computers are trying to learn to write stories, solve math problems, and chat like humans. To do this, they use special mathematical structures called "neural networks," which are like massive, interconnected webs of neurons. Inside these webs, information travels as numbers, and the network learns by adjusting the strength of the connections between these neurons. For a long time, scientists have noticed that these numbers can get messy—sometimes they grow too huge, sometimes they shrink to nothing, and the whole system can get confused about which direction to move to get smarter. This is a bit like trying to walk up a steep, foggy mountain while carrying a backpack that keeps changing weight; you might take a step forward, only to slide right back down. The goal of this research is to find a better way to carry that backpack so the computer can climb the mountain of learning much faster and more efficiently.

The paper you are about to read comes from the team at NVIDIA, led by Ilya Loshchilov and Boris Ginsburg. They are tackling a specific problem: how to make these AI models learn better by forcing all their internal numbers to stay on a perfect, invisible sphere. Think of it like a rule that says, "No matter how far you walk, you must always stay exactly one mile away from the center of the city." This rule, called "normalization," keeps the numbers from getting too big or too small, which helps the computer stay focused. The authors built upon a previous idea called "nGPT" (normalized GPT) and tested it on a very modern, powerful type of AI that mixes two different technologies: "Mamba-2" (which is great at remembering long sequences) and "Transformers" (which are excellent at understanding context). They wanted to see if their new set of training rules could make these hybrid models learn faster than the standard way.

The Big Idea: A New Training Recipe

The authors didn't just tweak one thing; they cooked up a whole new "training recipe" for these AI models. Imagine you are teaching a robot to walk. The old way (using a standard method called AdamW) is like letting the robot stumble around, sometimes taking giant, clumsy steps and other times tiny, hesitant ones, while hoping it eventually finds the path. The new recipe, which they call the nGPT training recipe, is like putting the robot on a treadmill with a very specific set of rules to keep its steps perfect.

Here is how they did it, broken down into simple, everyday concepts:

1. The "Logit Gradient Preconditioning" (The Volume Knob)
When the AI tries to guess the next word in a sentence, it produces a list of scores called "logits." In the old system, the volume of these scores could get weirdly loud or quiet as the training went on, confusing the robot about how hard to push. The authors added a special "volume knob" (a learnable scale vector) that they could adjust. But here's the clever part: they realized that if they just turned the knob, it would mess up the robot's memory of how to walk. So, they invented a trick called Logit Gradient Preconditioning. It's like having a volume knob that changes the sound you hear, but when the robot tries to learn from a mistake, the system automatically turns the volume back down to normal so the robot doesn't get confused. This keeps the learning process steady, even if the "loudness" of the predictions changes.

2. The "Logarithmic Learning Rate Decay" (The Speed Racer)
Usually, when training an AI, you start with a fast learning speed and slowly slow it down, like a race car driver easing off the gas pedal. The authors found that the standard way of slowing down (like a smooth curve) wasn't the best for their new "sphere" system. Instead, they used a "Logarithmic Learning Rate Decay." Imagine a race car that zooms off the line, slows down very quickly at first to get its bearings, and then cruises along a long, steady tail for the rest of the race. This "front-loaded" schedule lets the model learn the basics quickly and then refine its skills over a long period, which turned out to be much more efficient.

3. "GatedAdamW" (The Smart Gatekeeper)
The standard way to update the robot's brain is called AdamW. The authors upgraded this to GatedAdamW. Think of the standard method as a gatekeeper who lets every update through, even if the update is tiny and barely noticeable. Sometimes, these tiny updates are just noise. The new GatedAdamW has a "smart gate" that looks at each update. If an update is too small (like a whisper), the gatekeeper gently says, "Not today," and blocks it. If the update is a shout (a significant change), the gate swings wide open. This prevents the robot from wasting energy on tiny, useless adjustments and helps it focus on the big, important changes.

4. The "Angular Step Cap" (The Safety Leash)
Since the robot is walking on a sphere, taking a step that is too big could make it spin out of control or jump to the wrong side of the world. The authors added an "Angular Step Cap," which is like a safety leash. If the robot tries to take a step that is too large, the leash gently pulls it back to a safe, smaller angle. This ensures the robot never takes a wild, dangerous leap, keeping its journey smooth and stable.

The Results: Learning More Efficiently

The team tested this new recipe on a series of AI models, ranging from 1 billion to 14 billion parameters (the "brain size" of the robot). They compared their new nGPT models against the standard GPT models trained with the old methods.

The results were impressive. The new nGPT models consistently achieved lower error rates (measured by "validation loss") than the standard models. Specifically, the 14-billion-parameter nGPT model reached the same level of skill as the standard 14-billion model but required approximately half as many training tokens (the amount of text the model reads) to get there. In other words, to reach the same destination, the nGPT model only needed to read about half the amount of text compared to the standard model.

The paper suggests that this improvement comes from the combination of all these new rules working together, not just one single trick. They also tested a specific setting for their "smart gate" (called the sharpness parameter $a$ ) and found that a softer gate ( $a=0.5$ ) worked slightly better than a strict one, but the biggest win came from the whole new training system, not just the gate itself.

What This Means

The authors are careful to note that they didn't try every possible variation of these rules (they didn't do a "complete ablation study"), so there might be even better settings they haven't found yet. However, the results they did find are solid: by forcing the AI to walk on a perfect sphere and using a smarter, more controlled way to take steps, they can teach these massive models to be smarter using significantly less data and time. It's a practical recipe that suggests we can build better AI without needing to feed them the entire internet, just a much more efficient portion of it.

Technical Summary: Training nGPT on Modern Hybrid MoE Models

Problem Statement

The normalized Transformer (nGPT) framework proposes representing model parameters and activations on a unit hypersphere to facilitate hyperspherical representation learning. While the original nGPT formulation (Loshchilov et al., 2024) demonstrated the viability of this approach in dense Transformer settings, applying it to modern, large-scale hybrid architectures—specifically those combining Mamba-2 state-space models with Transformer Mixture-of-Experts (MoE)—requires a robust training recipe. Standard optimizers like AdamW, designed for unconstrained Euclidean spaces, may not interact optimally with the geometric constraints of nGPT, particularly regarding gradient scaling, learning rate schedules, and the handling of small second-moment statistics.

Methodology

The authors propose a comprehensive "nGPT training recipe" tailored for hybrid Mamba-2–Transformer MoE models (specifically the Nemotron-3 architecture). This recipe integrates normalization constraints with specific optimization modifications to stabilize training and improve convergence.

1. Normalization and Scaling

Parameter Constraints: All parameter vectors forming matrices are constrained to unit norm along the embedding dimension.
Activation Normalization: Hidden states are re-normalized after recombination with attention/SSM and MLP/MoE outputs.
Learnable Scaling Vectors: To prevent over-constraint in nonlinear layers, the model employs learnable scaling vectors:
- $s_{qk}$ for rescaling normalized queries and keys.
- $s_u$ for rescaling intermediate MLP/MoE activations.
- $s_z$ for rescaling output logits to control distribution sharpness.
Mamba-2 and MoE Specifics: Trainable scalars ( $s_{mamba}$ , $s_{moe}$ ) are introduced to rescale inputs to SiLU and sigmoid router functions, respectively, replacing standard normalization layers like RMSNorm.

2. Optimization Ingredients

The recipe introduces several novel components to the optimization process:

Logit Gradient Preconditioning (LGP): In nGPT, the learned logit scale vector $s_z$ grows logarithmically during training, acting as a hidden, time-varying multiplier on gradients. LGP decouples this effect by normalizing the backward gradient scale: $\partial L / \partial u \leftarrow (s_z / \text{mean}(s_z))^q \odot \partial L / \partial z$ . With $q=1$ , this removes the global mean scale while preserving relative vocabulary-wise variations, stabilizing gradient flow without altering the forward pass.
Logarithmic Learning Rate Decay: Instead of standard cosine decay, the authors employ a logarithmic decay schedule. This front-loads the learning rate budget, allowing for a sharper initial decrease followed by a long tail, which aligns better with the constrained parameter dynamics of nGPT.
GatedAdamW: This optimizer extends AdamW by separating numerical stability from update suppression.
- It uses a numerical constant $\epsilon_{num}$ strictly for stability.
- It introduces a soft gate $\gamma_t = \sigma(a \log(d_t / \epsilon_{gate}))$ based on the second-moment scale $d_t$ .
- This allows explicit control over which coordinates receive updates when second-moment statistics are small, preventing the suppression of updates that might occur with a single $\epsilon$ value in standard AdamW.
Geometric Controls:
- Angular Step Cap (ASC): Limits the angular displacement of normalized non-embedding parameter vectors to ensure geometric stability.
- Second-Moment Growth Clipping (SMGC): A safeguard against extreme gradient spikes by clipping the growth of second moments.
- Post-Moment Exploration Noise (PMEN): Adds noise to encourage exploration in low-variance regions.

Key Results

The authors evaluated the recipe on hybrid Mamba-2–Transformer MoE models ranging from 1B to 14B total parameters, using an industry-grade data blend.

Data Efficiency: The primary finding is a significant improvement in data efficiency. Specifically, the 14B-parameter nGPT model reaches the same validation loss as an unnormalized GPT baseline (trained with AdamW) using approximately half the number of training tokens. This result is derived from interpolation of the loss–token curves in Figure 7, rather than a direct table comparison across all model sizes.
Consistent Scaling: The performance gains are consistent across model sizes (1B, 2B, 4B, 7B, and 14B). The nGPT variants consistently achieve validation losses approximately 2.5% lower than their unnormalized counterparts trained with AdamW.
Convergence Dynamics: nGPT exhibits a slower initial loss decrease followed by faster convergence later in training, a behavior attributed to the learning rate schedule more directly controlling the effective step size in the absence of weight decay interactions.
Component Analysis: Ablation studies on the gating sharpness parameter $a$ in GatedAdamW suggest that while $a=0.5$ offers a modest advantage over $a=1$ (which recovers standard Adam behavior), the bulk of the performance gain stems from the normalized parameterization and the complete training recipe rather than the gating mechanism alone.

Significance and Claims

The paper claims to provide a practical training recipe that successfully extends nGPT from dense Transformers to modern, large-scale hybrid MoE architectures. The authors emphasize that the core improvements in data efficiency are not specific to MoE models but are a result of the normalization and optimization synergy.

The work is presented with modesty regarding its scope:

The experiments were conducted with a limited compute budget, focusing on a recipe that "worked in practice" rather than an exhaustive ablation study of every hyperparameter.
The authors explicitly state that comprehensive component-wise ablations and hyperparameter scaling laws are left for future work.
The significance lies in demonstrating that hyperspherical representation learning can be scaled to 14B-parameter models with substantial gains in training efficiency, offering a viable alternative to standard Transformer training pipelines.