How to Transfer Your Learning Rate to Larger Models

The Problem: The Width Scaling Trap

No one runs comprehensive hyperparameter sweeps on a trillion-parameter model — instead, researchers tune on smaller models and hope those insights transfer. But smaller models behave differently from larger ones. So how do you actually transfer what you learned? That's what this post is about: a technique called μP (Maximal Update Parameterization) that makes hyperparameters transfer reliably across model scales.

This isn't just an inconvenience — it's a massive waste of compute. Every time you scale up, you need expensive hyperparameter sweeps that can cost millions of dollars at frontier scale. So how can we optimize it?

Why Gradients Grow with Model Width

To understand the problem, let's trace what happens inside a neural network. Consider a hidden layer h = W @ x, where W is the weight matrix and x is the input.

Step 1: Activation Norm Grows with Width. Each neuron in the hidden layer has some output value. When we look at the norm of the entire hidden layer, it grows with width:

Step 2: Gradient Norm Inherits This Growth. For the output layer y = W_out @ h, the gradient depends on the hidden activations h. Since ||h|| grows with √width, so does the gradient:

Step 3: The Consequence. This means wider models experience larger weight updates for the same nominal learning rate. The effective learning rate becomes width-dependent:

Effective LR ∝ actual_LR × √width

The Experiment: Standard Parameterization Under the Microscope

To prove this theory experimentally, we ran a controlled study. We trained MLPs of varying widths (64 to 2048) and depths (2 to 16) on a spiral classification task, using a fixed learning rate of 0.01 for all configurations.

If standard parameterization (SP) is width-invariant, all widths should train similarly. If our theory is correct, wider models should show larger gradient/weight ratios and potentially unstable training.

The Results: Width Breaks Standard Training

Look at the gradient/weight ratio (||∇W|| / ||W||) across all width-depth combinations. This ratio determines the effective learning rate — it measures how much weights change relative to their current magnitude:

The pattern is unmistakable. At depth=2, the ratio goes from 0.6 (width=64) to 1.5 (width=2048) — a 8× increase on a log scale. The wider your model, the larger the relative gradient, and the more aggressive your effective learning rate becomes.

Training Dynamics: Loss Curves by Width

This width-dependent gradient scaling has real consequences for training. Look at how the loss evolves for different widths with a fixed learning rate:

Observation 1: Gradient Norm Scales with √width

The data confirms our theoretical prediction. Gradient norms follow the √width line almost exactly, regardless of depth. This isn't a bug in specific architectures — it's a fundamental property of standard parameterization.

Observation 2: The Ratio Keeps Growing

This plot shows the key comparison. The red SP lines climb steadily as width increases — wider models get progressively larger updates. Meanwhile, the μP lines (blue and green) stay remarkably flat, achieving width invariance.

More Results: Gradient Dynamics During Training

How do gradient norms evolve during training? This visualization shows that the width-dependent gradient scaling persists throughout training, not just at initialization:

Key insight: The √width scaling isn't just an initialization effect — it persists throughout training. This means the relative update imbalance compounds over thousands of training steps, making hyperparameter transfer fundamentally impossible with standard parameterization.

The Solution: Maximal Update Parameterization (μP)

In 2022, Greg Yang and colleagues at Microsoft Research proposed an elegant fix called Maximal Update Parameterization (μP) [1]. The insight is simple: if gradients scale with √width, we should scale the output of each layer by 1/√width to compensate.

class StandardLinear(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.W = nn.Parameter(
            torch.randn(out_dim, in_dim)
            * (2/in_dim)**0.5
        )

    def forward(self, x):
        return x @ self.W.T

class MuPLinear(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.W = nn.Parameter(
            torch.randn(out_dim, in_dim)
            * (2/in_dim)**0.5
        )
        self.scale = 1 / in_dim**0.5  # μP fix
    def forward(self, x):
        return self.scale * (x @ self.W.T)

The Proof: μP Achieves Width Invariance

The side-by-side comparison below makes the difference clear. While SP (left) shows a dramatic color gradient from left to right as width increases, μP (right) displays horizontal uniformity — the gradient/weight ratio remains nearly constant across all widths. At depth=2, the ratio is ~0.7 whether width is 64 or 2048. The model behaves the same at any scale.

Full Stability Analysis

The top-left panel shows SP's gradient ratio growing from ~5 to ~25 as width increases; the top-middle shows μP staying flat between 0.4-1.0. The bottom-left bar chart quantifies this: μP achieves a 42.8% reduction in coefficient of variation across widths. The bottom-center confirms gradient norms in SP scale with √width (dotted line), while μP (dashed) remains constant.

μP in Frontier LLM Research

This tutorial covered only the basics of μP — specifically, learning rate transfer across widths. We haven't touched initialization scaling, which uses the same principles. At frontier AI labs, μP is standard practice for scaling up training runs efficiently.

Active research areas include GQA-μP [3] (scaling rules for Grouped Query Attention in Llama-style models) and token horizon scaling [2] — while μP handles width, optimal LR still depends on total training tokens, leading to 2D scaling laws.

Production implementations must also address: (1) non-parametric norms (learnable LayerNorm gains can break μP limits), (2) weight decay scaling [4] (AdamW isn't naturally μP-invariant), and (3) embedding multipliers to prevent gradient vanishing into wide layers.

Experiments and post generated with Orchestra.