PaperlayerView exact revision on arXiv

Source-linked AI summary

Spectral Condition for $μ$P under Width-Depth Scaling

Chenyu Zheng, Rongzhen Wang, Xinyu Zhang, Chongxuan Li

arXiv:2603.00541v2cs.LGstat.ML

TL;DR

Large models that scale in both width and depth face unstable feature learning, costly tuning, and unreliable hyperparameter transfer, while existing joint-scaling μP results are fragmented. The paper develops a unified spectral framework and optimizer-level recipe, finding that k ≥2 scaling performs robustly in GPT-2 style experiments whereas standard parameterization and k = 1 μP often do not.

  • Problem

    Joint width-depth scaling complicates stable feature learning and reliable hyperparameter transfer because existing μP extensions remain architecture-, optimizer-, and theory-specific.

  • Method

    The paper extends μP through spectral constraints on residual-network weights and updates, then maps those constraints to concrete hyperparameterizations across many optimizers.

  • Results

    GPT-2 style experiments show that the k ≥2 formulation achieves stable feature learning and robust hyperparameter transfer, whereas standard parameterization and k = 1 μP often fail.

  • Takeaways & Limitations

    The k ≥2, CompleteP-style formulation is better suited to multi-transformation residual branches such as Transformers under width-depth scaling.

  • Takeaways & Limitations

    The spectral results are derived from a simplified linear residual MLP with a one-step update, although the paper reports extensions under additional assumptions.

Abstract

from arXiv · show

Generative foundation models are increasingly scaled in both width and depth, posing significant challenges for stable feature learning and reliable hyperparameter (HP) transfer across model sizes. While maximal update parameterization ($μ$P) has provided a principled solution to both problems for width scaling, existing extensions to the joint width-depth scaling regime remain fragmented, architecture- and optimizer-specific, and often rely on technically involved theories. In this work, we develop a simple and unified spectral framework for $μ$P under joint width-depth scaling. For deep residual networks whose residual blocks contain $k$ transformations, the framework specifies how the norms of weights and their per-step updates should scale with width and depth. It reveals a fundamental transition from $k=1$ to $k\geq 2$, unifying previously disparate $μ$P formulations and identifying the $k\geq 2$ case as more appropriate for practical architectures with multi-transformation branches such as Transformers. Building on this framework, we derive a general recipe for implementing $μ$P across a broad class of optimizers by mapping spectral constraints to concrete HP parameterizations, recovering existing results and extending them to additional optimizers. Finally, experiments on GPT-2 style language models show that the $μ$P formulation derived from the $k\geq 2$ case achieves stable feature learning and robust HP transfer under width-depth scaling, whereas standard parameterization and $μ$P in the $k=1$ case often fail to do so. These results support the practical effectiveness of the proposed spectral framework.

1 Introduction

Joint width-depth scaling makes stable feature learning and reliable hyperparameter transfer difficult, while existing μP extensions are fragmented across architectures and optimizers. The paper proposes a unified spectral framework, optimizer recipe, and empirical validation for practical models.

  • Scaling width and depth can destabilize feature learning and make hyperparameter tuning prohibitively expensive for large models.
  • Existing joint width-depth μP formulations depend on specific residual-block depths, optimizers, and technically involved theoretical tools.
  • The proposed spectral framework specifies weight and per-step update norm scaling for deep residual networks under joint width-depth scaling.
  • The framework distinguishes k = 1 from k ≥2 through higher-order update terms, recovering Depth-μP-style and CompleteP-style formulations respectively.
  • A unified optimizer recipe derives concrete μP hyperparameterizations for multiple optimizers from spectral constraints rather than ad hoc tuning.
  • Controlled GPT-2 style experiments find that the k ≥2 formulation supports stable feature learning and robust hyperparameter transfer, unlike standard parameterization and often k = 1 μP.

2 Preliminaries

The preliminaries define the notation and width-scaling setting underlying μP. They explain μP through scale-invariant feature learning, a spectral condition, and its limitation when depth also grows.

  • The paper uses standard asymptotic notation for functions of jointly increasing width and depth, with fixed quantities absorbed into constants.
  • RMS vector and operator norms are defined alongside spectral norms, with spectral conditions expressed using RMS operator norms for simplicity.
  • The width-scaling setup studies a deep linear MLP with fixed depth and widths tending to infinity, capturing core width-scaling feature-learning behavior.
  • μP targets scale-invariant feature learning while maximizing feature change induced by parameter updates.
  • A spectral condition on weights and per-step updates provides a unified route to deriving hyperparameterizations for broad optimizer classes.
  • The width-only spectral condition applies only when depth is fixed, motivating a simpler unified perspective for joint width-depth scaling.

3 Spectral Condition for µP under Width-Depth Scaling

The paper extends μP’s spectral analysis to residual networks whose width and depth grow together, focusing on two-transformation blocks. Higher-order update terms create stricter scaling for k ≥2, while the resulting formulation extends across fixed-depth branches and practical training settings.

  • 3.1 Problem Setup: The setup studies residual networks with k = Θ(1) transformations per branch, focusing on k = 2 as the minimal multi-transformation case.
  • 3.2 Spectral Condition: The spectral condition separately controls initialization for stable forward features and updates for stable, maximally changing one-step features.
  • 3.2 Spectral Condition: Residual contributions and relevant weight or update norm products must scale as Θ(1/L) to control accumulation across L blocks.
  • 3.3 Update Condition: For two-layer residual blocks, expanding one-step updates produces first-order terms when one branch weight changes and second-order terms when both change.
  • 3.4 Implications: For k = 1, the absent second-order term yields looser Depth-μP-style scaling, whereas k = 2 yields stricter CompleteP-style scaling.
  • 3.4 Implications: For any fixed k ≥2, all first- through k-th-order update terms are constrained at Θ(1/L), without changing the resulting μP formulation from the k = 2 case.
  • Scope: Under stated assumptions, the spectral results generalize beyond the linear one-step setup to multiple steps, nonlinearities, and multiple training examples.
  • Derivation: The analysis uses norm estimates based on subadditivity and submultiplicativity, with tightness relying on standard non-cancellation and alignment behavior.

4 Implementation of Spectral Condition

The paper maps the spectral condition to optimizer-specific hyperparameter parameterizations, including a simplified rule for modern optimizers and a ratio-based implementation for practical transfer.

  • Initialization and residual scaling: The chosen variance parameterization aligns initial weight variance with standard width-scaling µP before determining α_l from the spectral conditions.The construction uses RMS operator norms of weight matrices at initialization and treats input, output, and hidden layers separately.
  • Optimizer-specific parameterization: The implementation derives µP parameterizations by mapping spectral constraints to optimizer-specific learning rates and other hyperparameters.The recipe covers Muon-Kimi, Muon, Shampoo, SOAP, AdamW, Sophia, Lion, SGD, and SSO, recovering several existing formulations.
  • Simplified implementation rule: For modern optimizers except SGD, Condition 3.1 reduces to width-scaling µP plus a hidden residual multiplier α_l = Θ(1/L).Normalization or preconditioning removes the depth factor from hidden-layer update norms, whereas SGD requires an additional α_l-dependent learning-rate rescaling.
  • Ratio-based implementation: For Muon-Kimi, the hidden learning rate uses η_l = η_base/√r_n, equivalent to the theoretical scaling η_base/√n.Here r_n = n/n_base is the width ratio, and Table 1 summarizes the resulting width-depth parameterization.
  • Hyperparameter transfer: µP transfers optimal base hyperparameters from a small model to larger targets by applying the prescribed width- and depth-dependent scaling.This limits the expensive hyperparameter search to base values such as η_base.
  • Practical scope: The parameterization is derived from a simplified setup, then applied to standard language-model pretraining to assess practical utility.The paper notes that multiple equivalent parameterization solutions exist and selects one aligned with standard width-scaling µP.

5 Experiments

Experiments on GPT-2-style Transformers evaluate feature stability and hyperparameter transfer under joint width-depth scaling. Condition 3.1 generally outperforms SP and the k = 1 formulation, especially for multi-transformation residual branches.

  • Feature learning: Under SP, feature scale grows rapidly with width and depth, whereas Condition 3.1 maintains stable, scale-invariant final-block RMS feature norms.Models are trained for 10 steps while scaling either width or depth.
  • Hyperparameter transfer: Under width scaling, SP shifts the optimal learning rate substantially, whereas µP keeps the optimal base learning rate nearly invariant.Under depth scaling, µP also preserves transferability and achieves lower loss than SP as depth increases.
  • Optimizer robustness: Similar advantages over SP appear for Muon-AdamW, Shampoo-AdamW, and Sophia in additional experiments.These results extend the observed pattern beyond the primary Muon-Kimi-AdamW example.
  • Residual-block depth: Condition 3.1 yields stable learning-rate transfer for Muon-Kimi-AdamW and Muon-AdamW, while Condition B.1 shifts the optimum as depth increases.This supports stricter scaling when residual branches contain multiple transformations.
  • Additional diagnostics: Without LayerNorm, SP becomes unstable and loses depth-wise hyperparameter transfer, while µP remains stable and transfers reliably up to L = 256.Normalization and moderate tested depths can partially mask SP’s scaling pathology in the primary experiments.
  • Main experimental result: Condition 3.1 (k ≥2) maintains stable feature learning and robust hyperparameter transfer across width-depth scaling, while SP and Condition B.1 often fail.The experiments use GPT-2-style language models and compare the proposed formulation with standard parameterization and k = 1 µP.

6 Conclusion

The paper develops a unified spectral framework for μP under joint width-depth scaling, clarifying how residual-block depth changes the required scaling rules. It shows that the k≥2 formulation generalizes across fixed-depth residual branches and provides a practical basis for optimizer-specific parameterizations.

  • Existing joint width-depth μP extensions are fragmented across architectures and optimizers and often require technically involved analyses.
  • The framework explains a fundamental transition from one-layer residual blocks to multi-layer blocks through higher-order update terms.For k=1, second-order terms are absent; for k≥2, they impose tighter scaling conditions.
  • The two-layer residual block is the minimal representative of all fixed-depth residual branches with k≥2.Increasing fixed internal depth adds higher-order interactions but does not change the resulting μP implementation.
  • The analysis assumes fixed internal block depth k = Θ(1) and does not address regimes where internal block depth scales with width or network depth.
  • Higher-order update contributions for fixed k≥2 are automatically controlled at Θ(1/L), so no additional parameterization is required beyond the two-layer case.

B.2.3 Derivation for Preliminary Initial Condition

The derivation establishes initialization and update constraints for k-layer residual blocks under joint width–depth scaling, including bias extensions. For fixed k≥2, the resulting implementation reduces to the k=2 case.

  • Input layer: Choosing α0∥W0∥R = Θ(1) ensures order-one input features when the input norm is Θ(1).
  • Hidden layers: Hidden feature magnitudes remain stable when each residual contribution has bounded RMS norm under zero-mean Gaussian initialization.The summands are treated as approximately independent, so the squared RMS norm of their sum scales with the sum of squared norms.
  • Update conditions: Maximal first-order feature updates require αl∥∆Wl∥R∥Wl∥R = Θ(1/L) for every hidden layer and transformation.
  • Update conditions: For any subset of j updated transformations, the corresponding j-th-order update term is controlled by imposing the same Θ(1/L) scaled-product condition.
  • Final condition: For every fixed k≥2, the spectral constraints reduce to hidden residual multipliers αl = Θ(1/L) and common hidden update-norm scaling.
  • Bias parameters: Biases can be added with order-one initialization and update scales without modifying the matrix-weight hyperparameter parameterization.

C.4.2 Derivation of Parameterization

This section derives concrete SGD hyperparameter scaling from the k=2 spectral condition, using layer-specific update requirements and gradient norms. It also shows that second-order hidden updates follow from initialization and first-order constraints.

  • Hidden layers: For hidden layers, the first-order spectral condition combines αl = Θ(1/L) with order-one weight and update norms.
  • Hidden layers: Hidden second-order update conditions are automatically satisfied once the initialization and first-order update conditions hold.
  • Input and output layers: SGD parameterization is derived by matching raw gradient scales to the spectral requirement that input and output updates are order one.

C.5.2 Derivation of Parameterization

This section derives k=2 parameterizations for AdamW, Lion, Sophia, Shampoo, and related optimizer families from the same spectral condition. Several optimizers share parameterizations when their effective update rules have matching scales.

  • AdamW: AdamW parameterization is derived from its update rule, with momentum omitted for scaling analysis and stabilization handled separately.
  • Unified implementation: The framework derives optimizer-dependent learning rates, weight decay, and stabilization terms while retaining common initialization and block-multiplier rules.
  • Lion: Lion reuses AdamW parameterizations when both momentum terms are omitted, because both reduce to sign gradient descent.
  • Sophia: Sophia shares AdamW parameterizations because its estimated second-order scale has the same order as the AdamW estimate.
  • Shampoo: Shampoo shares Muon learning-rate and weight-decay parameterizations because its preconditioned update matches Muon’s update rule.

C.9.2 Derivation of Parameterization

This section derives optimizer-specific μP parameterizations from the spectral conditions, covering multiple optimizers and accounting for weight decay, stabilization, and residual-block scaling. It also shows that SOAP and Shampoo share Muon’s parameterizations under the stated simplifications.

  • Shampoo: Shampoo reduces to Muon under the stated simplification, so their learning-rate and weight-decay parameterizations coincide.The reduction follows by matching the preconditioned direction to Muon’s update rule.
  • SOAP: SOAP likewise reduces exactly to Muon under the analyzed eigendecomposition and update simplifications, yielding Muon’s parameterizations.The conclusion applies to the parameterizations summarized in the corresponding table.
  • k = 1 implementation: For the k = 1 implementation, hidden residual multipliers use the Depth-μP-style choice α_l = Θ(1/L), while input and output multipliers remain α_0 = Θ(1) and α_L+1 = Θ(1/n_in).The analysis treats weight decay as comparable in scale to the optimizer-driven update.
  • Unified derivation: The derivation maps spectral update constraints to concrete learning-rate, weight-decay, stabilization, and multiplier choices across optimizer families.The section organizes these rules under initialization, update, and weight-decay conditions for both k = 2 and k = 1 formulations.

D.3.1 Derivation of Parameterization

This section derives k = 1 μP parameterizations for SGD, AdamW, Lion, Sophia, and SSO by enforcing spectral update conditions across layers, biases, stabilization terms, and weight decay.

  • General setup: The k = 1 derivation uses α_l = Θ(1/L) for hidden residual multipliers and preserves the k = 2 input and output layer conditions.The resulting optimizer rules are summarized in implementation tables for the covered families.
  • Stabilization: The derivation also parameterizes stabilization terms to match gradient or preconditioner scales rather than dominate the optimizer-driven update.This treatment is applied to Shampoo and AdamW-style stabilization parameters.
  • SGD: SGD parameterization accounts for the raw gradient’s dependence on the residual multiplier when enforcing α_l∥∆W_l∥_R = Θ(1/L).Input and output learning rates are η_0 = Θ(n_out) and η_L+1 = Θ(n_in), with corresponding weight-decay scalings.
  • AdamW, Sophia, and Lion: AdamW, Sophia, and Lion share the same k = 1 parameterization under the sign-gradient simplification used in the derivation.The simplification sets AdamW momentum and stabilization terms to zero, while Sophia and Lion inherit the same scaling rules.
  • Sign-style updates: For sign-style optimizers, the hidden update direction’s norm is dimension-determined, so changing the hidden multiplier is the main difference from the k = 2 implementation.The weight-decay parameterization remains unchanged because its scale matches the optimizer update and does not depend on the residual multiplier.

Appendix E Additional Details and Results of GPT-2 Experiments

The appendix provides additional GPT-2 experiment details and numerical-result organization, including comparisons among SP, k ≥ 2 μP, and k = 1 μP, plus an asset-license inventory.

  • Experimental comparisons: The experiments compare SP with Condition 3.1 μP (k ≥ 2) and Condition B.1 μP (k = 1) in GPT-2-style language models.Condition 3.1 is described as CompleteP-style scaling for residual branches with multiple transformations.
  • Reproducibility: Table 14 lists the datasets, code, and licenses used for the experiments.The appendix explicitly identifies the used assets and their licenses.

E.2 Additional Details of Feature Learning Experiments

Additional feature-learning experiments evaluate final-block RMS norms across width and depth scaling, using repeated GPT-2-style Transformer runs and several optimizer settings. The results show scale growth for SP but approximate scale invariance for k ≥ 2 μP.

  • Evaluation protocol: Feature-scale stability is measured by the final Transformer block’s RMS norm after short training runs, averaged over three seeds.The coordinate-check experiments vary width and depth while using a fixed initialization and training setup.
  • Optimizer settings: The appendix reports coordinate-check results for Muon-Kimi-AdamW, Muon-AdamW, Shampoo-AdamW, and Sophia.Each optimizer configuration uses specified learning rates and optimizer hyperparameters, with results assigned to corresponding figures.
  • Results: SP exhibits feature-scale growth under width or depth scaling, whereas Condition 3.1 μP keeps the final-block feature norm approximately scale-invariant.The same qualitative pattern appears across the optimizer settings reported in the coordinate checks.

E.3 Additional Details of HP Transfer Experiments

The experiments evaluate SP and µP from Condition 3.1 across width and depth scaling using GPT-2 style models and report final validation loss. The setup includes multiple optimizers and standardized training procedures.

  • Experimental scope: The study compares SP with µP from Condition 3.1 (k ≥2) under both width and depth scaling, with additional optimizer-specific comparisons where applicable.For each optimizer, the experiments assess feature learning and HP transfer; Muon-Kimi-AdamW also includes weight-decay transfer and no-LayerNorm diagnostics.
  • Reported pattern: Across the reported comparisons, µP from Condition 3.1 supports stable feature learning and HP transfer, while SP becomes less reliable as model size or depth increases.The figures report lower loss or stronger transfer for µP in the tested settings.
  • Model setup: GPT-2 style models use OpenWebText, sequence length 1024, a GPT-2 tokenizer, width 256 and depth 4 for the base model.Width experiments vary n with L = 4, while depth experiments vary L with n = 256.
  • Training setup: All models train for about 300M tokens with batch size 240, 120 warmup iterations, cosine decay, gradient clipping of 1.0, and optimizer-specific settings.The schedule runs for 1221 iterations and decays to a minimum learning rate of 3 × 10^-5.
  • Evaluation: The reported HPs are base HPs transformed by the corresponding SP or µP scaling rules, and the experiments report final validation loss.Condition 3.1 and Condition B.1 implementations are specified in separate parameterization tables.

E.3.2 Additional Details of Muon-Kimi-AdamW

Muon-Kimi-AdamW experiments compare SP and two µP formulations across width and depth, with and without LayerNorm, and evaluate learning-rate and weight-decay transfer. Condition 3.1 generally provides the most robust transfer and stability, especially at larger depths and without LayerNorm.

  • Experimental setup: Muon-Kimi-AdamW uses Muon-Kimi for hidden matrices and AdamW for embeddings, the LM head, and biases.Muon-Kimi uses Nesterov-style momentum 0.95; AdamW uses β1 = 0.9, β2 = 0.95, and ϵbase = 10^-16.
  • Width-wise learning-rate transfer: SP fails to transfer the optimal base learning rate across widths, whereas µP from Condition 3.1 approximately transfers it and achieves lower loss as width increases.The tables highlight the best validation loss for each width.
  • Depth-wise learning-rate transfer: With LayerNorm, SP transfers the optimal base learning rate across tested depths, but this apparent transfer should be interpreted cautiously because LayerNorm and the tested depth range can mask feature-scale instability.Condition 3.1 also transfers the optimal base learning rate across depths and achieves lower loss than SP as depth increases.
  • No-LayerNorm diagnostic: Without LayerNorm, SP fails to preserve stable training, while Condition 3.1 has stable runs and approximately transfers the optimal base learning rate at L ≥32.NaN values indicate training instability where the loss explodes.
  • Weight-decay transfer: SP fails to transfer optimal base weight decay across widths, whereas Condition 3.1 approximately transfers it across widths and depths while achieving lower loss than SP as depth increases.The depth-wise comparison is reported with LayerNorm.
  • Muon-AdamW comparison: The Muon-AdamW comparisons show the same broad pattern: SP shifts or fails to transfer optimal learning rates, while Condition 3.1 approximately transfers them across widths and depths.Condition 3.1 also achieves lower loss than SP as depth increases.

G.2 Experimental Details

The appendix tests assumptions used to extend the spectral analysis beyond its simplified setting. It uses deep residual MLPs on a small binary CIFAR-10 subset and evaluates metrics across depths, training phases, and representative layers.

  • Verification protocol: The simulations are designed to empirically test Assumptions G.1–G.3, with emphasis on depth scaling rather than width scaling.The experimental setup largely follows Yang et al. [45].
  • Dataset: The verification dataset contains 100 airplane and 100 automobile images from CIFAR-10, flattened into vectors in R^3072 with binary labels.The task is binary classification.
  • Architecture and training: The model is a deep residual MLP with ReLU activations, an input layer, L residual blocks, and a final linear output layer.The width is n = 256, each residual block has width n, and the output dimension is 1.
  • Architecture and training: Models minimize binary cross-entropy with full-batch gradient descent for T = 200 steps.The setup uses the width-depth µP parameterization for SGD derived in Table 5 of Appendix C.5.
  • Verification protocol: The depth-scaling analysis trains networks with L ∈ {4, 8, 16, 32, 64, 128, 256} and tracks assumption-related metrics at initialization, halfway through training, and the end.Measurements include the input layer and representative internal layers of the final residual block.
Loading 2603.00541v2…