[v1.0] Spectral Condition for µP: The Unified Theory of Width-Depth Scaling

TL;DR

Scaling large models usually requires a painful, expensive re-tuning of hyperparameters. While Maximal Update Parameterization (µP) solved this for width scaling, depth scaling remained the "Wild West." This paper introduces a Spectral Condition for joint width-depth scaling. It provides a simple rule: if your residual block has 2+ layers, your weights and updates must shrink with depth at a rate of $\Theta (1/L)$. This allows you to tune a tiny model and "zero-shot" transfer those HPs to a massive, deep foundation model.

Background: Why Depth scaling Kills Training

In Standard Parameterization (SP), as you add layers, the feature magnitudes typically explode or vanish. Modern tricks like LayerNorm mask this pathology, but the underlying optimization landscape still shifts. When you scale from 12 layers to 128 layers, the "optimal" learning rate you found for the small model no longer works for the big one.

µP's goal is Scale-Invariance: making the training dynamics of a 100M parameter model look identical to a 100B parameter model.

The Core Insight: Deeper Blocks, Stricter Rules

The authors utilize a Spectral Perspective, looking at the Root Mean Square (RMS) operator norms of weight matrices ($W$) and their updates ($\Delta W$).

They found a fundamental difference between blocks with $k=1$ layer vs $k\ge 2$ layers:

• Single-layer blocks: Only need a loose constraint ($\Theta (1/\sqrt{L})$).
• Multi-layer blocks (Standard Transformers): Require a much stricter constraint ($\Theta (1/L)$) because weight updates interact across sub-layers, creating "second-order" feature changes that can easily blow up.

Condition 3.1: The Universal Scaling Law

For a residual network to satisfy the µP principle, the weights must satisfy:

1. Initialization: ${\alpha }_l\Vert W_l^{(2)}{\Vert }_R\Vert W_l^{(1)}{\Vert }_R=\Theta (1/L)$
2. Updates: ${\alpha }_l\Vert \Delta W_l^{(2)}{\Vert }_R\Vert W_l^{(1)}{\Vert }_R=\Theta (1/L)$

Methodology: Mapping Math to Optimizers

The beauty of this framework is its portability. The authors provide a recipe to transform these spectral bounds into actual code for any optimizer.

For Muon-Kimi (a popular optimizer for training 1T+ models like Kimi K2), the recipe is:

• Hidden LR: Scale by $1/\sqrt{r_n}$ (where $r_n$ is the width ratio).
• Block Multipliers: Scale by $1/r_L$ (where $r_L$ is the depth ratio).

Experiments: Proving the Theory

The researchers tested this on GPT-2 style Transformers using the OpenWebText dataset. They compared SP vs. the new µP across scaling dimensions.

1. Feature Stability

As shown in the "coordinate checks," under SP, the feature norms explode as depth increases. Under µP, the line remains perfectly flat.

2. Zero-shot HP Transfer

This is the "Holy Grail." In Figure 1(c, d), the optimal learning rate (the "bottom" of the loss curve) for a small model (blue) stays virtually identical for a deep model (red) when using µP.

Figure 1: Comparison of SP and µP. Note how the loss minima align perfectly in (c) and (d) for µP.

Critical Analysis & Takeaways

The paper effectively simplifies a highly "mathematicized" subfield of AI into a set of engineering principles.

• Impact: It provides immediate value for teams training large-scale models with diverse optimizers (Lion, Sophia, AdamW).
• Limitations: The derivation assumes a linear model, though empirical tests on non-linear Transformers suggest it holds. The "Spectral Condition" proves that architecture (residual depth) is just as important as the optimizer in determining scaling laws.

Future Work: Extending this to non-residual architectures (like pure SSMs or newer sparse models) will be the next frontier for the spectral framework.