📊 statistics

A Proof of Learning Rate Transfer under $μ$ P

This paper provides the first theoretical proof that $\mu$ P enables learning rate transfer to infinite-width linear MLPs by ensuring the optimal learning rate converges to a non-zero constant, a property that fails under standard and neural tangent parametrizations.

Original authors: Soufiane Hayou

Published 2026-02-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Soufiane Hayou

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 you are a chef trying to bake the perfect cake. You have a recipe (the neural network) and a specific ingredient you need to get just right: the amount of sugar (the learning rate).

In the world of Artificial Intelligence, making bigger cakes (larger neural networks) usually means you have to completely rewrite your recipe. If you double the size of the cake, you might have to cut the sugar in half, or change the baking temperature entirely. This is a huge problem because it means you have to spend weeks testing different amounts of sugar on small cakes just to guess what works for the giant ones.

This paper, written by Soufiane Hayou, introduces a special "magic recipe" called µP (Maximal Update Parametrization). The author proves mathematically that if you use this magic recipe, you can bake a tiny cake, find the perfect amount of sugar, and then use that exact same amount for a cake that is a million times bigger. You don't need to re-test anything.

Here is the breakdown of how this works, using simple analogies:

1. The Problem: The "Goldilocks" Problem

In standard AI training (called Standard Parametrization or SP), the "perfect" amount of sugar (learning rate) changes depending on how big your cake is.

Small Cake: You need a lot of sugar.
Medium Cake: You need a medium amount.
Giant Cake: If you use the same amount as the small cake, the cake burns. You have to use almost no sugar at all.

This is like trying to drive a car where the gas pedal sensitivity changes every time you switch from a Mini Cooper to a semi-truck. You'd have to relearn how to drive every time you changed vehicles.

2. The Solution: The "Magic Recipe" (µP)

The paper focuses on a specific way of setting up the neural network called µP. Think of µP as a special scaling rule that balances the ingredients perfectly.

The author proves that under µP, the "perfect sugar amount" (optimal learning rate) stays the same regardless of the cake size.

The Proof: The author uses math to show that as the network gets infinitely wide (like a cake that keeps getting wider and wider), the optimal learning rate doesn't shrink to zero or explode; it settles on a stable, non-zero number.
The Result: You can tune your hyperparameters on a small model (say, 128 neurons wide), and then apply that exact same setting to a massive model (8,000+ neurons wide) with zero extra effort.

3. The Analogy: The "Volume Knob"

Imagine the learning rate is a volume knob on a radio.

Standard Parametrization (SP): As you add more speakers to your sound system (increase the width), the volume knob gets stuck. To keep the music at a good level, you have to turn the knob down closer and closer to zero. Eventually, the music is silent. You have to constantly adjust the knob just to hear anything.
µP (The Magic Recipe): No matter how many speakers you add, the volume knob stays in the exact same "sweet spot." The music sounds perfect whether you have 2 speakers or 2,000. The system is "self-balancing."

4. Why is this a big deal?

Currently, training massive AI models (like the ones that write this text) is incredibly expensive and slow. Researchers have to run thousands of experiments on small models to guess the right settings for the big ones.

This paper provides the mathematical proof that µP solves this guessing game. It confirms that:

It works: The optimal settings really do stabilize as the model grows.
It's unique: Other methods (like the standard way or the "Neural Tangent Kernel" way) fail at this; they require constant re-tuning.
It saves money: You can tune your settings on a cheap, small computer and then confidently deploy them on a supercomputer.

Summary

Think of this paper as the "User Manual" for building AI at scale. It proves that if you follow the µP instructions, you don't need to be a wizard to figure out how to train giant models. You just find the right setting on a small model, and the math guarantees it will work perfectly on the big one. It turns a chaotic, trial-and-error process into a predictable, stable science.

1. Problem Statement

In deep learning, scaling neural networks (increasing width $n$ ) often requires re-tuning hyperparameters (HPs), particularly the learning rate ( $\eta$ ), to maintain training stability and performance.

The Phenomenon: Empirical observations suggest that under Maximal Update Parametrization (µP), the optimal learning rate converges to a constant as the network width increases. This allows practitioners to tune HPs on small models and transfer them to larger models without re-tuning (Zero-shot HP transfer).
The Gap: While this phenomenon was conjectured by Yang and Hu (2021) to be a result of µP achieving "maximal feature learning" in the infinite-width limit, there was no rigorous mathematical proof. Conversely, under Standard Parametrization (SP) and Neural Tangent Parametrization (NTP), the optimal learning rate was observed to shift (typically vanishing to zero or diverging) as width increases.
Objective: The paper aims to provide the first rigorous proof that learning rate transfer occurs under µP for deep linear Multi-Layer Perceptrons (MLPs) and to demonstrate why it fails under alternative parametrizations.

2. Methodology

The author employs a theoretical analysis of deep linear networks trained with Gradient Descent (GD) in the infinite-width limit ( $n \to \infty$ ).

Model Setup:
- A linear MLP with $L$ layers: $f(x) = V^\top W_L \dots W_1 W_0 x$ .
- Training minimizes quadratic loss using GD.
- Neural Parametrization: The paper formalizes how initialization variances and learning rates scale with width $n$ $n$ via exponents $\alpha_\ell$ $α_{ℓ}$ and $c$ $c$ .
  - µP: $W_\ell \sim \mathcal{N}(0, n^{-1})$ , $V \sim \mathcal{N}(0, n^{-2})$ , $\eta \sim \mathcal{O}(1)$ .
  - SP: $W_\ell \sim \mathcal{N}(0, n^{-1})$ , $V \sim \mathcal{N}(0, n^{-1})$ , $\eta \sim \mathcal{O}(1)$ .
Core Analytical Approach:
- Polynomial Expansion: The key insight is that for a linear network, the model output $f^{(t)}(x)$ after $t$ training steps can be expressed as a polynomial function of the learning rate $\eta$ .
- Asymptotic Analysis: The coefficients of this polynomial are random variables dependent on the initialization. The author analyzes the convergence of these coefficients as $n \to \infty$ using the Strong Law of Large Numbers (SLLN) and properties of random matrix products (Tensor Programs).
- Loss Landscape: Since the loss is the square of the output error, the loss function $L_n(\eta)$ becomes a polynomial in $\eta$ . The optimal learning rate is the minimizer of this polynomial.
- Convergence of Minima: The proof relies on showing that the sequence of random polynomials converges almost surely to a deterministic limiting polynomial. The minimizer of the random polynomial then converges to the minimizer of the limiting polynomial.

3. Key Contributions

A. Theoretical Proof of LR Transfer under µP

The paper proves that under µP, the optimal learning rate $\eta_n^*$ converges to a non-zero deterministic constant $\eta_\infty^*$ as $n \to \infty$ .

Step 1 ( $t=1$ ): The output $f^{(1)}(x)$ $f^{(1)} (x)$ is expanded as $f^{(0)}(x) + \sum \phi_\ell \eta^\ell$ $f^{(0)} (x) + \sum ϕ_{ℓ} η^{ℓ}$ .
- Under µP, the coefficient of the linear term ( $\phi_1$ ) converges to a non-zero constant.
- Coefficients of higher-order terms ( $\phi_{\ell}, \ell \ge 2$ ) vanish as $O(n^{-(\ell-1)/2})$ .
- Consequently, the loss function becomes asymptotically quadratic in $\eta$ , ensuring a unique, non-zero minimizer.
General Step ( $t \ge 2$ ): The author extends this to arbitrary training steps. While the polynomial degree increases (becoming non-linear in $\eta$ ), the limiting polynomial still possesses a unique minimizer bounded away from zero, provided the limiting loss has a unique minimizer.
Convergence Rate: The optimal learning rate converges at a rate of $O(n^{-1/2})$ .

B. Proof of Failure under SP and NTP

The paper demonstrates that LR transfer fails under Standard Parametrization (SP) and Neural Tangent Parametrization (NTP).

Under SP: The initialization of the output layer $V$ has variance $n^{-1}$ (instead of $n^{-2}$ in µP). This causes all polynomial coefficients to scale with $\sqrt{n}$ . To minimize the loss, the optimal learning rate must scale as $1/\sqrt{n}$ (or faster), causing $\eta_n^* \to 0$ as $n \to \infty$ .
Under NTP: The scaling leads to a "lazy training" regime where features do not learn. The optimal learning rate tends to diverge or shift significantly to compensate for the kernel regime dynamics.

C. Empirical Validation

Extensive experiments on synthetic linear and non-linear (ReLU) datasets confirm the theory:

µP: Optimal learning rates for widths $n \in [128, 8192]$ cluster around a constant value.
SP: Optimal learning rates shift toward zero as width increases.
Robustness: The phenomenon holds for varying depths ( $L=3, 9, 27$ ), optimizers (Adam), and training steps (near convergence).

4. Key Results

Theorem 1 (LR Transfer at $t=1$ ): For µP, $\eta_n^{(1)} \xrightarrow{P} \eta_\infty^{(1)} = \frac{m}{L} \frac{y^\top K y}{\|Ky\|^2}$ , where $K$ is the normalized input Gram matrix. The convergence rate is $O(n^{-1/2})$ .
Theorem 2 (No LR Transfer under SP): Under SP, $\eta_n^{(1)} \xrightarrow{P} 0$ as $n \to \infty$ .
Theorem 3 (LR Transfer at Step $t$ ): Under mild conditions (unique minimizer of the limiting loss), LR transfer holds for any training step $t$ under µP. The optimal learning rate converges almost surely to a non-zero constant.
Empirical Convergence: Simulations show that for practical widths (e.g., $n=1024$ ), the difference between the optimal LR and the limit is negligible, validating the utility of µP for hyperparameter transfer.

5. Significance and Impact

Theoretical Foundation: This work provides the first rigorous mathematical justification for the "Zero-shot Hyperparameter Transfer" property of µP, moving beyond heuristic arguments.
Practical Utility: It validates the strategy of tuning learning rates on small models and applying them to massive models (e.g., LLMs), significantly reducing the computational cost of hyperparameter search.
Clarification of Parametrizations: It mathematically distinguishes µP from SP and NTP, explaining why µP is superior for feature learning and stable scaling. It highlights that the specific initialization of the output layer ( $V \sim n^{-2}$ ) is the critical factor enabling stable feature updates and constant optimal learning rates.
Limitations & Future Work: The current proof is restricted to linear networks and Gradient Descent. The author notes that extending this to non-linear networks and complex optimizers (like Adam with momentum) requires new proof machinery, though empirical results suggest the phenomenon persists.

In summary, the paper establishes that µP is the unique parametrization (among standard choices) that ensures the optimal learning rate remains stable as network width scales to infinity, providing a solid theoretical basis for the efficiency of µP in modern large-scale AI training.