🤖 machine learning

Learning Rate Transfer in Normalized Transformers

This paper introduces $\nu$ GPT, a novel parameterization of the Normalized Transformer that leverages alignment exponents to achieve learning rate transfer across model width, depth, and token horizon, addressing a key limitation of the original nGPT.

Original authors: Boris Shigida, Boris Hanin, Andrey Gromov

Published 2026-05-01

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Boris Shigida, Boris Hanin, Andrey Gromov

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

The Big Picture: The "Goldilocks" Problem

Imagine you are baking a giant cake (a massive AI model). You have a recipe for a small cupcake (a tiny model). You know exactly how much sugar and heat (hyperparameters like the learning rate) work for the cupcake.

The problem? If you just blindly double the ingredients for a bigger cake, it might burn or stay raw. In AI, when researchers make models bigger (wider or deeper), the "perfect" settings for the small model usually stop working for the big one. You have to start from scratch and guess the new settings, which is slow and expensive.

This paper introduces a new way to bake these "cakes" (called nGPT models) so that the settings you find for a small model work perfectly for a giant one without any extra guessing. They call this new recipe νGPT (pronounced "nu-GPT").

The Ingredients: What is nGPT?

First, the authors are working with a specific type of AI called nGPT (Normalized Transformer).

The Old Way: Traditional AI models are like a car with a very sensitive engine. If you push the gas too hard (high learning rate), the engine blows up. If you push too soft, it goes nowhere. To keep it safe, you need a "brake" called weight decay and you need to "warm up" the engine slowly before driving fast.
The nGPT Way: The nGPT model is like a car with a self-stabilizing suspension. It naturally keeps its speed and balance in check. Because of this, it doesn't need the "brakes" (weight decay) or the "warm-up." It trains faster and is more efficient.

The Catch: Even though nGPT is great, the authors found that its "speed settings" (learning rates) still didn't transfer well. If you tuned the speed for a small nGPT, and then tried to use those same settings for a huge nGPT, the big one would perform poorly.

The Solution: The νGPT Recipe

The authors created νGPT. Think of this as a new set of instructions on how to scale your ingredients based on the size of the cake.

They discovered three specific rules to make the settings "transfer" perfectly:

1. The Token Horizon Rule (The "Distance" Rule)

The Concept: Imagine you are teaching a dog. If you only walk it for 10 minutes, you can run fast. If you plan to walk it for 10 hours, you need to start slower so it doesn't get exhausted.
The Finding: The paper found that as the AI "walks" longer (processes more data tokens), the learning rate shouldn't drop as fast as people thought. Instead of dropping like a heavy stone, it should drop like a feather floating down.
The Math: They found the speed should decrease by the cube root of the time walked (specifically, $1/3$ power), not the square root ( $1/2$ ) that other theories suggested.

2. The Width Rule (The "Team Size" Rule)

The Concept: Imagine a choir. If you double the number of singers (width), the sound gets louder. If you don't adjust the conductor's volume (learning rate), the choir might get too loud and distort, or too quiet to hear.
The Finding: The authors realized that in these specific models, the "voices" (activations) and the "singers" (weights) don't align perfectly. They are partially aligned.
The Fix: They adjusted the learning rate for the hidden parts of the model to drop by a specific amount ( $3/4$ power) as the model gets wider. This is different from previous theories (like $\mu$ P) which assumed the voices and singers were perfectly aligned. By assuming they are only partially aligned, they found a "sweet spot" that works better.

3. The Depth Rule (The "Layer" Rule)

The Concept: Imagine a relay race. If you add more runners (layers) to the team, does the baton pass faster or slower?
The Finding: Surprisingly, for this specific type of model, the learning rate for the hidden layers doesn't need to change much as you add more layers. The model is naturally stable. However, they did find that the "starter" settings (initialization) for the very first and last layers need a tiny tweak to keep the race smooth.

The Results: Why It Matters

The authors tested this new νGPT recipe against the old one.

The Old Way (nGPT): When they made the model bigger, the performance curve was messy. The "best" learning rate for a small model was the "worst" for a big model.
The New Way (νGPT): When they made the model bigger, the performance curve was perfect. The learning rate that worked for the small model worked for the big model, and the bigger model performed just as well (or slightly better) than if they had tuned it from scratch.

Summary Analogy

Think of the Learning Rate as the volume knob on a radio.

Old Theory: "If you buy a bigger speaker (wider model), you must turn the volume knob down by half."
The Paper's Discovery: "Actually, because of how this specific speaker works, you only need to turn the volume knob down by a specific, calculated amount (the $3/4$ rule) to get the perfect sound. If you follow this rule, you can buy a speaker 100 times bigger, and the same volume setting will sound perfect."

What They Did NOT Claim

They did not say this makes AI smarter or able to do new things (like curing diseases).
They did not say this works for every type of AI model (it is specific to Normalized Transformers/nGPT).
They did not claim this eliminates the need for tuning entirely; it just means you can tune a small model and trust it will work for a big one.

In short: The paper provides a mathematical "translation guide" that lets engineers take the settings from a small, fast AI and apply them to a massive AI with confidence, saving time and computing power.

1. Problem Statement

The paper addresses the challenge of hyperparameter (HP) transfer in Normalized Transformers (nGPT), a variant of the Transformer architecture introduced by Loshchilov et al. (2025).

Context: nGPT achieves significant training speedups and eliminates the need for weight decay and learning rate warmup by extensively normalizing weights and activations (using $\ell_2$ -norms instead of LayerNorm/RMSNorm).
The Issue: While nGPT has hyperparameters that scale with model size, the authors observe that it fails to exhibit learning rate transfer across model width (dimension), depth (number of layers), and token horizon (training iterations).
Limitations of Existing Methods: Standard transfer techniques like $\mu$ P (Maximal Update Parametrization) rely on specific theoretical assumptions about the alignment of weights and activations (specifically, that they are fully aligned or unaligned in specific ways). The authors find these assumptions do not hold for nGPT, leading to suboptimal performance when scaling up.

2. Methodology

The authors propose a novel parameterization called $\nu$ GPT (nu-GPT). Their approach combines numerical experiments with a principled theoretical framework based on alignment exponents (Everett et al., 2024).

A. Theoretical Framework: Alignment Exponents

Instead of assuming the "full alignment" ( $\alpha=1$ ) or "no alignment" ( $\alpha=1/2$ ) regimes used in $\mu$ P, the authors empirically measure the alignment exponents ( $\alpha, \omega, \nu$ ) which describe how weight updates and activations scale relative to the model dimension $d_{model}$ .

They define exponents for how the change in a matrix-vector product scales: $\|\Delta W h\| \sim d^{\alpha} \|\Delta W\| \|h\|$ .
Key Finding: In nGPT, the "mid-alignment" regime is observed, where $\max\{\alpha, \nu\} \approx 3/4$ . This contradicts the $\mu$ P assumption of $\max\{\alpha, \nu\} = 1$ (full alignment) or $1/2$ (no alignment).

B. The $\nu$ GPT Parameterization

Based on the "mid-alignment" hypothesis ( $\max\{\alpha, \nu\} = 3/4$ ) and empirical scaling laws, the authors derive specific scaling rules for learning rates ( $\eta$ ) and initialization parameters. Let $m_{width}$ , $m_{depth}$ , and $m_{data}$ be the multipliers for width, depth, and token count relative to a base model.

Key Scaling Rules in $\nu$ GPT:

Global Learning Rate: Scales as $m_{data}^{-1/3}$ . This corrects for the token horizon, differing from the $m_{data}^{-1/2}$ suggested by some prior SDE-based theories.
Embedding Weights ( $\eta_{input}$ ): Scales as $m_{width}^{-1/2}$ .
Hidden/Unembedding Weights ( $\eta_{hidden}, \eta_{output}$ ): Scales as $m_{width}^{-3/4}$ . This is the critical deviation from $\mu$ P (which uses $m_{width}^{-1}$ or $m_{width}^{-1/2}$ depending on the variant).
Depth Corrections:
- The trainable rescaling vectors $\alpha_{A, init}$ and $\alpha_{M, init}$ are scaled by $m_{depth}^{-1}$ .
- Unlike "Depth- $\mu$ P" (which scales hidden learning rates by $m_{depth}^{-1/2}$ ), $\nu$ GPT finds that hidden learning rates should not be scaled with depth in nGPT, as the trainable rescalers naturally handle depth stability.
Initialization Constants: Fixed constants (e.g., 0.03) are used for scale parameters ( $\alpha_{scale}$ ) rather than scaling with $d_{model}^{-1/2}$ .

3. Key Contributions

Identification of $\mu$ P Failure in nGPT: Demonstrated that standard $\mu$ P assumptions regarding weight-activation alignment are invalid for normalized architectures, leading to poor learning rate transfer.
Proposal of $\nu$ GPT: Introduced a new parameterization that explicitly targets the "mid-alignment" regime ( $\approx 3/4$ ) observed in nGPT.
Token Horizon Scaling: Established that the optimal learning rate in nGPT scales with token count as $T^{-1/3}$ , refining previous theoretical predictions of $T^{-1/2}$ .
Depth Transfer Insight: Showed that for normalized models with trainable rescalers, explicit depth scaling of the hidden learning rate is unnecessary and potentially harmful; instead, scaling the initialization of the rescalers is sufficient.

4. Experimental Results

The authors validated $\nu$ GPT on the FineWeb-Edu dataset using the OLMo 2 tokenizer.

Width Transfer:
- Baseline nGPT: Showed no learning rate transfer; performance degraded as width increased.
- $\mu$ P-style corrections: Showed slight drift in optimal learning rates.
- $\nu$ GPT: Achieved perfect learning rate transfer across width (from 0.26B to 3.22B parameters) with no loss in validation loss compared to the tuned baseline.
Depth Transfer:
- Baseline nGPT showed decent depth transfer but became unstable at extreme depths.
- $\nu$ GPT maintained stability and performance across depths (8 to 128 layers), showing lower sensitivity to learning rate choices at large depths.
Token Horizon Transfer:
- Experiments confirmed the $T^{-1/3}$ scaling law for the optimal learning rate as the number of training iterations increased.
Alignment Exponent Verification:
- Empirical measurements confirmed that $\omega_{hidden} = \omega_{output} = 1/2$ and $\max\{\alpha, \nu\} \approx 0.75$ (mid-alignment), validating the theoretical basis of $\nu$ GPT.

5. Significance

Practical Impact: $\nu$ GPT allows researchers to train large-scale Normalized Transformers by simply tuning hyperparameters on small models and transferring them directly to massive scales, saving significant computational resources.
Theoretical Advancement: The paper challenges the universality of $\mu$ P, demonstrating that different architectural choices (specifically normalization schemes) require different alignment assumptions and scaling laws.
Efficiency: By removing the need for weight decay and warmup while enabling robust scaling, $\nu$ GPT offers a more streamlined and efficient training pipeline for next-generation foundation models.

In summary, the paper provides a rigorous, empirically grounded solution to the scaling problem in normalized transformers, replacing the "one-size-fits-all" $\mu$ P approach with a tailored parameterization ( $\nu$ GPT) that respects the unique dynamics of normalized architectures.