PaperlayerView exact revision on arXiv

Source-linked AI summary

Practical Efficiency of Muon for Pretraining

Essential AI, :, Ishaan Shah, Anthony M. Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh J Shah, Khoi Nguyen, Kurt Smith, Michael Callahan, Michael Pust, Mohit Parmar, Peter Rushton, Platon Mazarakis, Ritvik Kapila, Saurabh Srivastava, Somanshu Singla, Tim Romanski, Yash Vanjani, Ashish Vaswani

arXiv:2505.02222v4cs.LGstat.ML

TL;DR

Pretraining needs an optimizer that balances data efficiency, compute, and training time, while large-model hyperparameter tuning remains costly and imperfectly transferable. The paper evaluates Muon against AdamW through compute-time and token-efficiency tradeoffs, combines it with muP, and introduces telescoping hyperparameter search. Muon expands AdamW’s compute-time Pareto frontier, with its data-efficiency advantage persisting at large batch sizes, while muP transfers cleanly to large models and telescoping controls tuning overhead.

  • Problem

    The paper asks how to compare optimizers across compute, time, batch size, and data efficiency, and how to tune them efficiently at large model scales.

  • Method

    The paper compares Muon and AdamW using iso-loss compute-time frontiers, studies token ratios at large batch sizes, and combines Muon with muP and telescoping hyperparameter sweeps.

  • Results

    Muon expands AdamW’s compute-time Pareto frontier, while its relative data-efficiency advantage persists or grows with batch size and muP transfers successfully up to 3.7B-parameter models.

  • Takeaways & Limitations

    Muon with muP and telescoping transfer forms a practical recipe for more flexible resource allocation and compute-efficient hyperparameter search in large-scale pretraining.

  • Takeaways & Limitations

    muP transfer remains subject to finite-width bias and sampling error, so transfer is only approximately valid at finite model widths.

Abstract

from arXiv · show

We demonstrate that Muon, the simplest instantiation of a second-order optimizer, explicitly expands the Pareto frontier over AdamW on the compute-time tradeoff. We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient, thus enabling more economical training. We study the combination of Muon and the maximal update parameterization (muP) for efficient hyperparameter transfer and present a simple telescoping algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources. We validate our findings through extensive experiments with model sizes up to four billion parameters and ablations on the data distribution and architecture.

1 Introduction

The paper evaluates Muon as a practical second-order optimizer by measuring compute-time tradeoffs directly and studying data efficiency at large batch sizes. It also combines Muon with muP and proposes telescoping hyperparameter transfer to reduce large-scale tuning costs.

  • Compute-time tradeoff: Muon is evaluated by plotting compute and time required to reach identical target losses, rather than abstracting efficiency solely through wall-clock time or FLOPs.This represents optimizers as iso-loss frontiers while marginalizing other variables affecting the tradeoff.
  • Compute-time tradeoff: Muon expands AdamW’s Pareto frontier across variable batch sizes, increasing flexibility in allocating compute resources and training time.The analysis focuses on Muon as the simplest second-order optimizer and studies its relative data efficiency in the large-batch regime.
  • Hyperparameter transfer: Muon and muP are combined to transfer hyperparameters from small proxy models to larger models, reducing the need for large-scale hyperparameter search.The paper addresses both optimizer selection and practical hyperparameter tuning for pretraining.
  • Hyperparameter transfer: The telescoping algorithm adjusts the hyperparameter search grid across model scales while accounting for finite-width and sampling errors, adding O(C log (N)) compute cost.Here C is the cost of training the final model and N its width.
  • Experimental scope: Experiments vary model size up to 4 billion parameters, batch size up to 16 million tokens, data distribution, and architecture.The study is presented as a practitioner-oriented evaluation of optimization for pretraining.

2 Muon Improves the Compute-Time Tradeoff

Muon is implemented as a matrix-structured optimizer and evaluated against AdamW through compute-time and token-efficiency measurements. Across large-batch experiments, its relative data-efficiency advantage persists or grows, supporting a broader compute-time frontier despite higher optimizer FLOPs.

  • Muon formulation: Muon applies matrix-structured steepest descent with spectral norm regularization to transformer weight updates.The update transforms each gradient matrix, with the optimizer derived from the singular-value structure of the gradient.
  • Muon formulation: Muon avoids explicit SVD computation by using Newton-Schulz iteration and maintains only a first-moment matrix as optimizer state.The implementation also combines momentum, learning-rate scaling, and weight decay.
  • Experimental setup: The experiments use Gemma-based decoder-only transformers, text and Python-code data, TPU v5p hardware, and model sizes through 4B parameters.Training uses sequence length 8192, converged token budgets, and approximately 50% model FLOP utilization.
  • Initial results: Muon consistently reaches lower training loss at a given number of steps and reaches target loss faster than AdamW in wall time.This initial verification motivates the broader iso-loss compute-time comparison.
  • Compute-time tradeoff: Iso-loss curves compare total training time against compute resources as batch size and device count vary, exposing economically feasible options between optimizer frontiers.The gap between curves corresponds to alternatives such as shorter runs or fewer devices.
  • Relative data efficiency: The token ratio remains above 1 and generally increases at higher losses or stays constant at lower losses as batch size grows.This indicates that Muon’s relative data-efficiency advantage persists beyond fixed large batch sizes.
  • Relative data efficiency: At a target loss of 3.2 nats, 1B, 2B, and 4B models each require around 2B extra tokens with AdamW compared with Muon.The additional-token difference is reported as surprisingly agnostic to model size.

3 Choosing Hyperparameters for Muon

Muon-compatible hyperparameter transfer uses muP to calibrate large models from smaller proxies, while accounting for finite-width and sweep-resolution errors. The telescoping protocol progressively refines hyperparameter searches across widths and substantially reduces brute-force search cost.

  • Motivation: Muon is compatible with muP for calibrating large-language-model hyperparameters, including coupled weight decay.This addresses an open question about transferring hyperparameters when Muon is the optimizer.
  • Motivation: Brute-force searches over k hyperparameters on m-point grids require training m^k models, making large-model tuning prohibitively expensive.A modest 8 × 8 search for a 3.7B-parameter model would require nearly 820,000 device hours.
  • muP parameterization: muP transfers hyperparameters by using width-dependent weight rescaling, initialization variance, and learning-rate scaling.Weights use Wℓ = a(n)wℓ, entries of wℓ are initialized from N(0, b(n)), and the learning rate is c(n)η0.
  • Sources of error: Finite-width effects shift optimal hyperparameters by order 1/n, so transfer from a proxy is only approximately exact and improves as model width increases.The shift includes a network- and loss-dependent constant α, creating an unavoidable finite-width bias unless the proxy is sufficiently large.
  • Sources of error: Grid sweeps add sampling error ε, which can become substantial with coarse meshes, limited budgets, or poorly chosen search ranges.The telescoping algorithm controls finite-width drift and mesh-resolution error separately.
  • Telescoping protocol: The telescoping algorithm begins with a full sweep at the smallest width, then doubles width while reducing sweep points by 4^-1/k and refining the mesh.Under a smooth order-1/n optimum shift, each successive stage has roughly constant cost because doubling width increases FLOPs by about fourfold.
  • Efficiency: Compute savings versus brute-force grid search are typically higher than 50%, while final-model training accounts for 20%–99% of total training and tuning cost depending on confidence.The protocol allocates progressively narrower searches to larger, more expensive models.
  • Experiments: On 3.7B-parameter transformers, the final selected settings achieved training loss 1.61 nats at λ = 0.027 and η0 = 0.0612.Each telescoping iteration reduced variance, and the resulting scaling fit had R2 ≈ 1 with parameter exponent 0.31.

4 Related Work

Prior work established AdamW’s dominance and promising alternatives, but had not conclusively compared optimizers on the compute-time tradeoff. This paper addresses that gap while extending practical evidence for Muon and muP.

  • Dethroning AdamW: Muon expands AdamW’s compute-time Pareto frontier, addressing prior comparisons that fixed compute resources instead of evaluating resource-time tradeoffs.The paper uses iso-loss frontiers to compare the resources required to reach identical losses.
  • Impact of the batch size: Critical batch size is an incomplete summary because it is noise-sensitive and does not describe post-critical-batch behavior.The paper instead studies how optimizer efficiency changes across the broader batch-size regime.
  • Practical aspects of muP: muP relates infinite-width behavior to finite-width networks and aims to transfer hyperparameters from small models to larger ones.Existing studies had shown theoretical correctness and empirical effectiveness, while precise transfer across other scale dimensions remained approximate.
  • Practical aspects of muP: Muon’s compatibility with muP fills an empirical gap left by prior work on second-order optimizers such as Shampoo and KFAC.The same muP scaling used for AdamW works for Muon in the paper’s experiments.

5 Conclusion

The paper answers two pretraining questions: how to trade compute against time, and how to tune the optimizer economically. It reports that Muon improves the resource tradeoff over AdamW and that muP with telescoping sweeps enables scalable hyperparameter transfer.

  • Optimizer choice: Muon requires 10–15 % fewer tokens than AdamW to reach an identical loss, with the advantage remaining constant or growing as batch size increases.The result is validated across five model sizes, two data modalities, and several decades of global batch sizes.
  • Optimizer choice: Muon expands AdamW’s compute-time Pareto frontier, increasing flexibility in allocating devices, training time, and data.The paper characterizes this advantage using a token ratio measuring relative data efficiency at large batch sizes.
  • Hyperparameter transfer: muP transfers learning-rate and coupled-weight-decay hyperparameters cleanly to models up to 3.7B parameters at sequence length 8192.This provides an affirmative answer to whether muP remains valid with Muon.
  • Hyperparameter transfer: The telescoping sweep contracts the search grid at each width doubling and bounds tuning overhead by O(C log N).The procedure accounts for known transfer-error sources and leaves more than 20% of total compute for the final full-scale run.
  • Unified recipe: The resulting recipe combines Muon optimization, muP scaling, and telescoping hyperparameter transfer for large-scale pretraining.The paper presents this combination as a practical replacement for the AdamW baseline.

B Reduction of Shampoo and Soap to Muon

Under simplifying point-estimate assumptions, Shampoo and Soap reduce to Muon. This connects Muon to second-order preconditioning through singular-vector transformations and inverse-square-root approximations.

  • Shampoo: Shampoo without momentum is equivalent to Muon, establishing a direct connection between the two second-order optimizers.The reduction is considered under simplifying assumptions described in the appendix.
  • Shampoo: Muon’s gradient transformation can be obtained from Shampoo’s point-estimate formulation under the stated matrix assumptions.The derivation uses the gradient’s singular-value decomposition and corresponding left and right singular vectors.
  • Soap: Soap performs Adam in a second-order basis and likewise reduces to Muon under the simplifying assumptions.The reduction follows by identifying the relevant eigendecomposition and singular-vector factors.
  • Interpretation: The Shampoo and Soap reductions transfer their inverse-square-root empirical-Fisher motivation to Muon.The empirical Fisher is described as an approximation of the inverse Hessian in this motivation.

C Initial Verification of Muon’s Performance

Initial experiments show that Muon converges faster and reaches lower training loss than AdamW across model sizes and datasets. In a 1B-parameter DCLM example, Muon remains below AdamW in both steps and wall-clock time through the end of training.

  • Initial verification: Muon consistently achieves lower training loss and faster convergence across different model sizes and datasets.The paper illustrates this pattern with best-run loss curves for a 1B model trained on DCLM.
  • Initial verification: Muon strictly lower-bounds AdamW’s loss curves in the 1B DCLM example, with no crossover even beyond the Chinchilla-optimal training budget.The comparison is shown as training loss versus both steps and hours.
  • Initial verification: The wall-time advantage appears with a minimal Muon implementation and may improve with more sophisticated implementation choices.The passage specifically mentions lower-precision first-moment computation as a possible improvement.

D Deeper Analysis of the Critical Batch Size

Critical batch size is better analyzed through steps-to-loss curves and their slopes than by visually locating kinks alone. The paper notes that these curves can lack clear transitions and that slope values determine whether larger batches remain beneficial.

  • Motivation: Increasing batch size can accelerate training on data-parallel hardware, but may reduce loss reduction per token.The tradeoff motivates evaluating both training speed and data efficiency.
  • Definition: B⋆ is defined from SL(B), the steps required to reach loss L, with TL(B) = B × SL(B) measuring consumed tokens.The critical batch size is identified through the token-consumption behavior around the transition.
  • Limitations: Visual kinks in log-log steps-to-loss plots can misidentify the critical batch size when transitions are unclear or slopes differ from −1.A slope above −1 indicates diminishing returns already apply, while a slope below −1 means observed kinks are not critical-batch transitions.
  • Assumptions: The post-critical token function incL(B) is only assumed to be increasing, without a prescribed specific functional form.Appendix F supplies a simple parametric model for completeness.
  • Empirical analysis: Figure 8 plots steps-to-loss against batch size for 500M-parameter models at target losses near convergence on a log-log scale.These plots provide the empirical basis for examining batch-size behavior.

D.1 Token-optimal batch size

Token-optimal batch size uses tokens-to-loss rather than steps-to-loss to obtain a uniquely defined measure of data-efficient batching. The measure increases during training and is generally larger for Muon, but it can miss benefits from very large batches.

  • Definition: Tokens-to-loss TL(B) is flat under perfect scaling and increases in the post-critical regime, making it useful for measuring batch efficiency.The corresponding token-optimal batch size is the largest data-optimal batch size.
  • Definition: Unlike steps-to-loss kinks, Btok,L is uniquely defined and coincides with B⋆ when the critical-batch assumption holds.This makes token-optimal batch size convenient for empirical measurement.
  • Results: Muon generally has a larger token-optimal batch size than AdamW across architecture, dataset, and model-size ablations.For both optimizers, the token-optimal batch size increases during training.
  • Limitations: At DCLM loss 3.25, AdamW and Muon share a 3.5M token-optimal batch size even though Muon’s token consumption then grows more slowly.Thus the single token-optimal point does not capture all benefits in the over-large-batch regime.

E Additional Compute-Time Tradeoff Plots

Compute-time tradeoff plots compare optimizers across target-loss thresholds using a piecewise-linear log-log model for steps and tokens. The curves shift upward for lower target losses, and the token-ratio condition characterizes persistent post-critical data-efficiency differences.

  • Threshold analysis: As the target loss decreases, the compute-time tradeoff curve moves upward because achieving lower losses requires more time and compute.Figure 12 overlays curves for different loss thresholds.
  • Modeling assumption: The analysis assumes SL(B) is piecewise linear in log-log scale, motivated by empirical observations.The assumption is explicitly described as strong.
  • Ablations: Figures 9–11 vary architecture, dataset, and model size, covering Gemma 3 versus Gemma 2, Python versus DCLM, and 1B, 2B, and 4B models.These plots extend the batch-size analysis across the stated ablations.
  • Modeling assumption: The paper derives parametric forms for steps, tokens, and the token ratio RL(B) from the log-log batch-size model.The construction uses continuous log SL(B) with intercepts b1 and b2.
  • Results: RL(B) does not vanish in the post-critical regime if and only if mA ≥ mB.This condition determines whether the relative token overhead remains nonvanishing at large batch sizes.

G Muon Implementation

The implementation adapts Muon’s Newton–Schulz orthogonalization to transformer parameter shapes and applies it through an Optax gradient transformation. It routes parameter groups by path and combines Muon with Adam-style updates and weight decay.

  • Core implementation: The Muon implementation adapts Newton–Schulz application for differently shaped transformer weights under jax.lax.scan.It retains Jordan et al.’s default Newton–Schulz hyperparameters.
  • Core implementation: orthogonalize_matrix accepts batched matrices shaped (batch_size, d, D), performs Newton–Schulz iterations, and returns the orthogonalized result.The routine uses coefficients a, b, and c and handles transposition based on matrix dimensions.
  • Tensor handling: Attention and MLP tensors are transposed, reshaped across layers, orthogonalized, and reshaped back to their original layouts.Separate routines handle attention QKV, attention output, and MLP weights.
  • Tensor handling: Path-based logic selects the MLP, attention-QKV, or attention-output orthogonalization routine and raises an error for unidentifiable paths.The result is scaled by base_scale × sqrt(max(d1, d2)) when base_scale is positive.
  • Optimizer composition: The combined optimizer labels parameters for Muon or Adam, adds decayed weights, and scales updates by the learning rate.Parameters containing norm, logits, or embedding in their path are assigned the Muon label.

H Parameterization of Telescoping and Distribution of Losses

The paper frames muP through spectral and dynamical scaling conditions, then analyzes finite-width corrections that cause hyperparameter optima to drift. This drift is quantified as O(1/n) and identified as a key error source in muP-based transfer.

  • The Maximal Update Parameterization: muP maintains scale-invariant propagation by requiring layer spectral norms to map typical inputs to outputs with width-matched norms.The spectral condition controls the largest singular value of each layer and preserves typical activation norms across widths.
  • The Maximal Update Parameterization: Stable forward, backward, and update dynamics are enforced through compatible spectral and learning-rate scalings.The paper connects these conditions to stable activation and gradient propagation and to order-one update effects.
  • Hyperparameter Drift under MuP Scaling: Finite-width effects make optimal hyperparameters drift from their infinite-width values as network width changes.The analysis models a width-dependent output f(x, n), its infinite-width limit f0(x), and a width-dependent optimum x⋆(n).
  • Hyperparameter Drift under MuP Scaling: At extrema with a vanishing denominator term, the general drift expression reduces to simplified drift laws.These cases follow from imposing optimality after expanding the loss around the infinite-width limit.
  • Hyperparameter Drift under MuP Scaling: O(1/n) hyperparameter drift is governed by the structure of f0 and f1 and constitutes a key error source in muP-based transfer.The result formalizes why minima shift during width scaling and quantifies both the scale and direction of that shift.
Loading 2505.02222v4…