Spectral Condition for μP under Width-Depth Scaling

TL;DR

Extends the spectral-norm view of µP from width-only scaling to joint width–depth scaling of residual networks, using only elementary linear algebra and probability. Introduces a single spectral condition (Condition 3.1) on the RMS operator norms of weights and per-step updates that unifies previously disparate µP formulations (Yang et al., Bordelon et al., Dey et al., Qiu et al.) as special cases of varying residual block depth. From this condition the paper derives concrete HP recipes — initial variance, block multipliers, and learning rates — for SGD, AdamW, Muon, Muon-Kimi, Shampoo, SOAP, SSO, Lion, and Sophia. GPT-2 experiments confirm scale-invariant feature learning and zero-shot LR transfer across width and depth.

Key claims

A single spectral scaling condition (RMS operator norms of weights and updates) characterizes µP under joint width–depth scaling for two-layer residual blocks; deeper residual blocks impose strictly stricter conditions, and prior µP formulations arise as special cases at different block depths [§3.2, Condition 3.1].
Hidden weights’ RMS operator norms (and updates) must shrink with depth as $L_b^{-1/2}$ to prevent feature explosion across the residual stack — the width-only condition (Yang et al. 2023) is recovered when this depth factor is dropped [§3.2].
For a residual block of fixed depth $b$ , the initialization requires the product of $b$ norms of hidden weights to scale as $L_b^{-1/2}$ , and the update condition constrains all first- through $b$ -th order update terms to scale as $L_b^{-1/2}$ [§3.2, Appendix B.2].
Once the first-order update condition (C2.2) and initial condition (C1.2) are imposed, the second-order update condition (C2.3) on hidden weights is automatically satisfied — they are equivalent under the framework [§3.3.3].
The derivation uses only subadditivity, submultiplicativity, and the spectral norm of i.i.d. Gaussian matrices, in contrast to prior derivations via Tensor Programs or dynamical mean-field theory [§2.1, §3.3].
For Muon-Kimi, hidden-layer LR scales as $\eta_h = \eta_b \cdot m_w^{-1} \cdot m_d^{-1/2}$ where $m_w, m_d$ are width and depth ratios versus a small base model [§4.3, Table 1, Eq. 15].
The framework recovers existing µP formulations for SGD (Bordelon et al. 2024), AdamW (Dey et al.), and matrix-preconditioned optimizers (Qiu et al.) under width-depth scaling, and extends µP to Muon, Muon-Kimi, SSO, Sophia, and Lion [§4, Appendix C].
On GPT-2 style models, SP exhibits substantial drift in optimal LR with width scaling, while µP preserves a nearly invariant optimal LR across both width and depth scaling [§5.2, Fig. 1(c,d)].
Removing LayerNorm exposes SP’s underlying instability: SP training becomes unstable and depth-wise HP transfer breaks down, while spectral-µP remains stable up to depth 60 with robust HP transfer — i.e. modern normalization was masking the SP pathology that µP was originally designed to fix [§5.2 Discussion, Appendix D.3, Fig. 2].

Method

The setup is a two-layer linear residual MLP: $x_{\ell+1} = x_\ell + \alpha_h \cdot W^{(2)}_\ell W^{(1)}_\ell x_\ell$ , with layerwise learning rates and block multipliers, scaled to infinite width and depth with $n_\text{in}, n_\text{out}$ fixed. The analysis proceeds in three steps: (i) a preliminary initialization condition $\|W\|_\text{RMS}^2 = O(L_b^{-1/2}/n_\text{in})$ from $\Theta(1)$ feature propagation through the residual recursion, using that the RMS norm of a sum of independent zero-mean vectors scales with the square root of the variance sum; (ii) an update condition expanding $x_\ell$ after one gradient step into zero-, first-, and second-order terms in $\Delta W$ , requiring each term to be $\Theta(L_b^{-1/2})$ to enforce both P1 (stable features) and P2 (maximally updated); (iii) refining the initial condition by combining the two first-order update inequalities and the second-order update condition. Mapping the resulting spectral constraints onto concrete HP parameterizations (initial variance, block multiplier, learning rate) yields the per-optimizer recipe.

For Muon-Kimi specifically — where the update is the RMS-normalized SVD-based orthogonalization $\eta \cdot UV^T \cdot \sqrt{n_\text{out}/n_\text{in}}$ — Table 1 gives the full HP set: input/output weights use width-only LR scaling, hidden weights use the additional $L_b^{-1/2}$ factor, and block multipliers carry $\Theta(L_b^{-1/2})$ for the residual term.

Results

Coordinate-check on the final Transformer block’s RMS norm: SP grows rapidly with both width and depth; spectral-µP stays flat [Fig. 1(a,b)].
HP transfer on GPT-2 trained for 300M tokens (batch 240, warmup + cosine decay): SP’s optimal LR drifts under width scaling; spectral-µP’s optimal LR is invariant across width and depth, with monotonically lower loss as size increases [Fig. 1(c,d)].
Depth transfer under SP appears to work in the small-depth regime only because of LayerNorm + QKNorm masking the instability; removing LayerNorm breaks SP depth transfer at depths up to 60 while spectral-µP remains stable [Appendix D.3, Fig. 2].
Optimizers covered with explicit recipes (Appendix C): SGD, AdamW, Muon, Muon-Kimi, Shampoo, SOAP, SSO, Lion, Sophia. Several recover prior published µP results as special cases.

Why it’s interesting

This is the cleanest derivation in the wiki of a unified µP recipe across both scaling axes — and it directly resolves the framing tension between the two papers on the existing Hyperparameter scaling laws page. It substantiates the “Enforce” position from Controlled LLM Training on Spectral Sphere (spectral conditions are the right primitive) and provides the theoretical scaffolding SSO was operating under; conversely it puts more pressure on the “Fitting” position from How to Set the Learning Rate for Large-Scale Pre-training? by showing that, once LayerNorm is ablated, SP’s depth transfer collapses while spectral-µP holds — i.e. modern normalization may be masking exactly what spectral-µP was designed to address. The Muon-Kimi recipe is also load-bearing for the broader Training stability at scale cluster, since several recently-filed reports (Kimi, Step 3.5 Flash) use Muon-Kimi at trillion-parameter scale.