Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales

AI summary

Overview

Research area: Deep learning optimization and scaling laws, specifically the design and scaling behavior of matrix-preconditioned optimizers (Shampoo, SOAP, Muon, AdaMuon) compared with AdamW for training language models.

Technical level: Advanced. The paper assumes familiarity with the Maximal Update Parameterization (μP), spectral norms, Kronecker-factored preconditioning, and the practice of tuning hyperparameters at scale.

Scope (one sentence): The paper derives and empirically validates width- and depth-scaling rules for the learning rate and weight decay of matrix-preconditioned optimizers, showing that those rules determine whether these optimizers actually retain their speed advantage over AdamW as models grow from 190M to 1.4B parameters.

What This Paper Is About

Several recently proposed optimizers that use matrix-level (rather than elementwise) preconditioning — Shampoo, SOAP, and Muon — have shown promising speedups over AdamW in small-scale experiments, but attempts to replicate these gains at larger scale have produced contradictory results, including reported speedups ranging from a consistent 2× down to roughly 1.1× that decays with model size. The authors argue the main culprit is not the optimizers themselves but the lack of a principled way to choose hyperparameters (especially the learning rate) when models become too large to tune directly, and they set out to derive such scaling rules under μP and test their effect on compute efficiency.

Key Contributions

1. A general procedure for deriving μP scaling rules. The authors show how to derive width-scaling rules for the per-layer learning rate and regularization parameters for a wide range of matrix-preconditioned optimizers, including Shampoo, SOAP, Muon, and AdaMuon, while accounting for commonly used techniques like blocking and grafting (Table 1). The rules are derived from a simplified μP condition: that each layer's output changes by Θ(1) per gradient step.

2. Extension of the scaling rules to model depth. Building on CompleteP-style reasoning, the authors generalize to residual networks by scaling the residual branch output by 1/L (with L the depth) and adjusting the learning rate so that each layer receives Θ(1) feature learning.

3. Diagnosis of where μP transfer breaks down at finite width, and two fixes. They identify failures of μP's infinite-width asymptotics for realistic finite widths — particularly for optimizers with RMS-based update normalization and for SOAP's eigenbasis-projected gradient — and show that blocking with fixed block size and explicit spectral normalization largely repair these failures.

4. A concrete demonstration that hyperparameter transfer determines whether speedups survive scaling. Applying the derived rules, they report that Muon, SOAP, and Shampoo consistently achieve near 1.4× speedup over well-tuned AdamW from 190M to 1.4B parameters, whereas the same optimizers show rapidly vanishing speedups under the Standard Parameterization.

Main Findings

• μP improves learning rate transfer but is not perfect. Across widths D = 128 to D = 4096 on transformers trained on OpenWebText for 100M tokens, μP produced a more stable loss-versus-learning-rate landscape than the Standard Parameterization for all tested optimizers, and gave lower loss when zero-shot transferring the optimal learning rate found at the base width D = 128. However, the optimal learning rate still drifted with width for some optimizers, showing that μP's infinite-width limit is not exactly reached.

• RMS-based update normalization causes late-training inconsistency. For AdaMuon, Shampoo, and Shampoo with Adam-grafting, the optimal learning rate increased with width under μP while feature updates shrank in late training (step 1000). The authors attribute this to update normalizations that force the update RMS to Θ(1), leaving the update's spectral norm scaling as Θ(√(d_in d_out / srank(U))), where the stable rank is bounded by min(D, tB) and is bottlenecked by D at realistic finite widths.

• μP fails the coordinate check for SOAP even at initialization. The authors trace this to sparse eigenbasis-projected gradients in the infinite-width limit versus dense gradients at realistic finite widths and large batch sizes.

• Blocking and spectral normalization fix the finite-width deviations. With a fixed block size of 128, μP achieved good learning rate transfer for all optimizers including SOAP and grafted Shampoo. Explicit spectral normalization to a √(d_out/d_in) spectral norm matched or exceeded μP transfer and produced more consistent feature updates in late training, at a cost of two matrix-vector multiplies per layer per step via online power iteration.

• Depth scaling rules transfer well. Applying a 1/L residual branch multiplier (α = 1) with adjusted learning rates, the authors observed consistently lower loss and more stable early-time (step 10) feature updates than both SP (α = 0) and the depth scaling when transferring from 3 to 192-layer transformers on OpenWebText.

• Compute-optimal learning rate transfer approximately works. On FineWeb with 20 tokens per parameter and block size 512, μP approximately stabilized the optimal learning rate, and transferring the optimum from D = 256 to D = 2048 gave near-optimal performance. Spectral normalization reduced learning rate sensitivity further. The authors caution, citing Everett et al., that μP can overshoot the learning rate under Adam in larger sweeps, so fits may be needed if width varies by more than a factor of 10.

• Optimal independent weight decay scales as 1/width. Contrary to μP's prescription of Θ(1) independent weight decay, the authors find λ = Θ(1/D) is near-optimal across all three optimizers at this scale, consistent with prior findings for AdamW by Xiao.

• Compute-optimal model size differs between optimizer families. The paper reports that the compute-optimal model size is larger for matrix-preconditioned optimizers than for Adam.

• Consistent ~1.4× speedup at scale with correct scaling. Using μP plus Θ(1/D) weight decay, Muon, SOAP, and Shampoo achieved near 1.4× speedups over AdamW for Llama-architecture language models from 190M to 1.4B parameters on FineWeb, with speedups vanishing rapidly under incorrect scaling.

• Hyperparameter transfer matters as much as the optimizer. Because the loss-versus-compute exponent in language modeling has been estimated as low as 0.05, a 2% loss change corresponds to a 40% change in estimated compute — meaning poor hyperparameters can fully mask or manufacture apparent speedups.

Methodology in Plain English

The authors take a two-part approach. First, they do pen-and-paper analysis. They consider a deep network and look only at the very first gradient step on a single training example, a simplification that prior μP work showed is sufficient to determine the correct scaling relationships. In this setting the gradient is rank-1 and momentum can be ignored, so the preconditioner matrices of Shampoo and related methods reduce to simple rank-1 objects with known eigenvalues. They then require that each layer's output changes by a fixed amount (Θ(1)) per step, and solve for how the learning rate and regularization constants must depend on the layer's input and output dimensions, block size, and depth. They repeat this for blocking, for grafting (where the update is rescaled to another optimizer's norm, making the learning rate follow whichever optimizer supplies the norm), and for residual networks where the branch output is scaled by 1/L.

Second, they run experiments to test the predictions. They train small transformers on OpenWebText (100M tokens, vocabulary size 96 so that full preconditioners fit on the embedding and readout layers) across widths from 128 to 4096 and depths from 3 to 192, comparing μP against the Standard Parameterization. They check the "coordinate check" — whether the RMS of the one-step feature update is invariant to width. They then run compute-optimal studies on FineWeb with the GPT-2 tokenizer at 20 tokens per parameter with block size 512, sweeping learning rate and weight decay. Finally, they train Llama-architecture models from 190M to 1.4B parameters following the setup of Wen et al., and compare the throughput-adjusted compute efficiency of Muon, SOAP, and Shampoo against well-tuned AdamW.

Why This Matters

The paper reframes a disputed empirical question — do matrix-preconditioned optimizers really beat AdamW at scale? — as a hyperparameter-transfer question, and shows that the answer changes qualitatively depending on whether the scaling rules are correct. It also provides a general recipe that can be reused for new optimizers rather than re-deriving μP from scratch each time.

Real-world applications:

• Foundation model pretraining: Choosing an optimizer for large language model runs where a 1.4× compute efficiency difference translates directly into cost and time savings.
• Compute budget planning: Using the reported 1/D weight decay scaling and transfer rules to decide model sizes and training horizons without re-tuning at every scale.
• Optimizer research and evaluation: Providing a fair comparison protocol that controls for tuning budget, addressing the mixed results in prior replication studies.
• Training infrastructure engineering: Informing which practical knobs (blocking, grafting, spectral normalization) to implement, given that blocking and normalization both improve transfer and reduce optimizer overhead.

Industry relevance: The paper's central claim — that the practical value of a new optimizer is only realized with correct hyperparameter transfer — directly affects how teams with fixed compute budgets should evaluate and adopt optimizers. The finding that Θ(1/D) weight decay is near-optimal across optimizers also provides an immediately actionable default.

Future Directions

• Finer modeling of compute-optimal learning rate scaling. The authors note their results align with μP working well in the compute-optimal regime at their tested scale, but cite larger-scale Adam sweeps where μP overshoots, and suggest fitting a power-law correction on top of μP when width varies by more than a factor of 10.
• Wider-scale validation. Their compute-efficiency comparisons cover 190M to 1.4B parameters; whether the near-1.4× speedup holds at trillion-parameter scale, where Muon has already been used, is left open.
• Comparison of 1/D weight decay against constant-EMA-timescale approaches. The authors state that at their scale the two are similar and leave a finer comparison to future work.
• Extension of the derivation to architectures and optimizers beyond those tested. The paper notes the general procedure should extend straightforwardly to other optimizers, and their batch-size formulation is sketched for batch sizes greater than one, but broader empirical coverage is not reported.

Target Audience

Optimizer researchers and engineers working on large-scale language model training; practitioners responsible for hyperparameter tuning pipelines or compute-efficiency benchmarking; and theoretically inclined readers interested in infinite-width limits, μP, and the spectral-norm characterization of stable updates. Readers without background in scaling parameterizations will find Sections 3.2 through 3.4 the most practically useful, while the derivations in the appendices (referenced but not fully included in the provided content) are aimed at those who want to re-derive the rules for a new optimizer.

Note on completeness: the provided paper content is truncated partway through Section 4, so the detailed model configuration table (Table 3), the full speedup breakdown per model size, and the remaining ablations are not available for summary here. All figures and results above are drawn only from the content that was provided.

Authors’ abstract

Several recently introduced deep learning optimizers utilizing matrix-level preconditioning have shown promising speedups relative to the current dominant optimizer AdamW, particularly in relatively small-scale experiments. However, efforts to validate and replicate their successes have reported mixed results. To better understand the effectiveness of these optimizers at scale, in this work we investigate how to scale preconditioned optimizers via hyperparameter transfer, building on prior works such as $μ$P. We study how the optimal learning rate and weight decay should scale with model width and depth for a wide range of optimizers, including Shampoo, SOAP, and Muon, accounting for the impact of commonly used techniques such as blocking and grafting. We find that scaling the learning rate according to $μ$P improves transfer, but can still suffer from significant finite-width deviations that cause drifting optimal learning rates, which we show can be mitigated by blocking and explicit spectral normalization. For compute-optimal scaling, we find scaling independent weight decay as $1/\mathrm{width}$ is nearly optimal across optimizers. Applying these scaling rules, we show Muon, SOAP and Shampoo consistently achieve near $1.4\times$ speedup over AdamW for training Llama-architecture language models of sizes ranging from $190$M to $1.4$B, whereas the speedup vanishes rapidly with scale under incorrect scaling. Based on these results and further ablations, we argue that studying optimal hyperparameter transfer is essential for reliably comparing optimizers at scale given a realistic tuning budget.