PaperlayerView exact revision on arXiv

Source-linked AI summary

When Does Muon Actually Win? A Paired Data-, Time-, and Tuning-Budget Benchmark

Mushir Akhtar, Hazem N. Nounou, Majdi Mansouri, Mohamed N. Nounou

arXiv:2610.21555v1cs.LGcs.AI

TL;DR

Optimizer rankings can differ across data, runtime, and tuning budgets, so this paper compares Muon and AdamW under matched data- and time-budget evaluations. Muon improves equal-data validation performance and lowers test loss across paired comparisons, while equal-time winners differ by workload.

  • Problem

    Optimizer efficiency can mean fewer training tokens, shorter runtime, or less tuning, and rankings can differ when update costs and tuning sensitivity vary.

    Source: §1 · p. 1

  • Method

    The benchmark compares data- and time-budget validation trajectories, attributes optimizer-step timing, and evaluates paired seeds across fixed workloads.

    Source: §2.2 · p. 3

  • Results

    Muon reduces equal-data validation-loss AUC by 3.25%–22.44% across workloads and lowers test loss in all 20 paired seeds, while equal-time winners split by task.

    Source: §3.1 · p. 3, §3.3 · p. 4

  • Takeaways & Limitations

    Across the four tested workloads, Muon improves data-budget performance, while equal-time rankings remain conditional on model size, parameter routing, token batch, and implementation.

    Source: §5 · p. 7

    Measured orthogonalization cost depends on matrix dimensions, token batch, and kernel implementation; the measurements do not isolate GPU underutilization or kernel-launch overhead.

    Source: §5 · p. 6

Abstract

from arXiv · show

Muon's gains in data efficiency motivate evaluating its performance under explicit data, time, and tuning budgets. On a single A100 GPU, we compare a reference single-device Muon implementation with AdamW across 17.7M- and 33.6M-parameter causal language models and 2.70M-parameter vision transformers. A common 23-candidate development budget selects configurations evaluated with five matched seeds, full validation trajectories, optimizer-step timing, and a frozen one-shot test. Muon reduces validation-loss area under the curve by 3.25--22.44\% at equal data and lowers test loss in all 20 primary pairs. In this configuration, Muon improves time-normalized loss area} by 19.68--22.31\% on vision, while AdamW is better by 3.38--6.96\% on language. The full Muon optimizer call, including auxiliary AdamW updates, consumes 56.1--69.6\% of training time, compared with 6.1--9.1\% for the AdamW baseline. Muon also has higher expected regret under restricted access to these candidate pools. These results characterize small models with reference kernels; their wall-clock ordering does not establish a ranking for larger models or optimized implementations.

1. Introduction

Muon and AdamW can rank differently depending on whether efficiency is measured by data, runtime, or tuning, motivating a budget-aware comparison. The benchmark finds Muon ahead on equal-data comparisons, while measured equal-time results differ by workload.

1. Introduction

  • Efficiency can mean fewer training tokens, shorter runtime, or less tuning, and these rankings can differ when update costs and tuning sensitivity vary.Prior Muon work distinguishes sample efficiency from wall-clock speed, while distributed work targets orthogonalization and communication costs.Source: §1 · p. 1
  • The benchmark pairs matched initializations across two FineWeb-Edu language models and two vision transformers, using 23 development candidates and five-seed full-horizon runs.It reports equal-data and equal-time validation trajectories, GPU event timing, and a frozen one-shot test evaluation; its scope is small dense models and reference kernels on one accelerator.Source: §1 · p. 1
  • Muon wins every primary equal-data comparison and all 20 paired test-loss comparisons, but the measured clock favors vision Muon and language AdamW.The reference Muon optimizer call accounts for over half of training time; expected validation-AUC regret under one randomly available candidate is 13.66% for Muon and 7.52% for AdamW.Source: §1 · p. 2
  • Optimizer rankings depend on search spaces, workloads, and measurement rules, motivating fixed-hardware benchmarks with explicit tuning protocols.This comparison adds paired crossover curves, optimizer-time attribution, and postselection test evaluation to those benchmark principles.Source: §1 · p. 2

2. A budget-aware comparison

The comparison spans language and vision workloads and routes eligible matrices to Muon while retaining AdamW for other parameters. Configurations are selected under a fixed development budget, then assessed with paired trajectories and a frozen test protocol.

2.1. Workloads and optimizer routing

  • The language workloads use 17.7M- and 33.6M-parameter transformers, while both vision workloads use a 2.70M-parameter ViT on balanced or 100:1 long-tailed CIFAR-100.Language runs use 5,000 steps; vision runs use 100 epochs, with held-out validation data for development.Source: §2.1 · p. 2
  • Muon updates hidden attention and feed-forward matrices, including fused QKV, while AdamW updates embeddings, normalization parameters, biases, and output heads.Muon covers 26.7% and 42.2% of language-micro and language-small parameters; both methods share the stated training controls and data order within each seed.Source: §2.1 · p. 2

2.2. Selection, estimands, and test isolation

  • One-seed development evaluates 23 candidates per workload and strategy, selecting configurations by data-normalized validation-loss AUC with tie-breaks by final loss and run identifier.Full-horizon runs use seeds 101–105 and estimate performance conditional on the selected configurations; selection stability across development seeds remains unmeasured.Source: §2.2 · pp. 2–3
  • The evaluation compares interpolated validation loss over paired common horizons using data or measured training seconds, with the seed as the replication unit.Training time excludes validation and checkpoint I/O; CUDA events time optimizer steps, and the official test evaluation is frozen before access and performed once per selected checkpoint.Source: §2.2 · p. 3

3. The winner depends on the budget

Muon wins every equal-data validation comparison and every paired test-loss comparison, but measured time-budget rankings differ by workload: Muon leads on vision and AdamW on language. The reference Muon optimizer call also takes a substantially larger share of training time than AdamW’s.

3.1. Data efficiency is consistent; time efficiency is not

  • Muon lowers equal-data validation-loss AUC by 3.25% and 4.46% on language and 16.98% and 22.44% on vision, while time-normalized AUC favors Muon on vision and AdamW on language.Time-normalized AUC reductions are 19.68% and 22.31% for Muon on vision, and AdamW leads by 6.96% and 3.38% on language.Source: §3.1 · p. 3, Figure 1 · p. 4

3.2. Optimizer-step cost in the reference implementation

  • Muon’s optimizer steps consume 56.1–69.6% of training time, versus 6.1–9.1% for AdamW, with Muon throughput at 34–47% of AdamW’s.The timed Muon call includes matrix updates and auxiliary AdamW; separate Newton–Schulz timing and kernel-utilization measurements are unavailable.Source: §3.2 · pp. 3–4, Table 1 · p. 4

3.3. The data-budget advantage transfers to the test set

  • Muon lowers test loss in all 20 paired comparisons, with mean reductions of 4.77% and 7.01% on language and 11.52% and 26.03% on vision.Across the four equally weighted workloads, mean relative test-loss reduction is 12.33% (stratified bootstrap 95% interval [11.84, 12.83]%), conditional on the chosen workloads.Source: §3.3 · p. 4, Figure 2 · p. 5

4. What changes the conclusion?

Restricted access to the frozen candidate pools yields higher expected validation-AUC regret for Muon than AdamW. Exploratory variants and approximation audits show trade-offs, but are descriptive and limited in replication.

Tuning allowance

  • With one randomly available candidate, expected relative AUC regret is 7.52% for AdamW and 13.66% for Muon; with five, it is 0.42% and 1.69%.With ten candidates, regret is 0.09% for AdamW and 0.67% for Muon; estimates condition on candidate ranges, pool construction, development seed, and uniform sampling.Source: §4 · p. 5

Variants and approximation

  • Full-horizon Muon achieves lower test loss than switching from AdamW after 25% of training on all four workloads; three-seed Variance-Muon raises long-tailed test loss by 11.71%.In one-seed language-small audits, three rather than five Newton–Schulz steps save 21.87% of training time with a 0.74% AUC penalty; restricted matrix routing saves 28–45% with penalties of 1.13–2.76%.Source: §4 · p. 5

5. Discussion and limitations

The measured equal-time ranking depends on workload and optimizer throughput, while restricting access to the frozen candidate pools raises Muon’s expected regret more than AdamW’s. These conclusions are bounded by the benchmark’s model scale, implementation, and selection design.

A break-even interpretation

  • Muon’s throughput is 0.339–0.468 of AdamW’s; vision gains overcome this gap, but language gains do not, and faster orthogonalization would move the break-even point.The time-budget comparison couples loss trajectories to the throughput that determines how much data each optimizer processes.Source: §5 · pp. 5–6

Tuning-budget implications

  • With restricted access to the frozen pools, Muon has higher expected validation-AUC regret than AdamW, and the audit does not model adaptive search.The figure averages 10,000 random orderings within each 23-candidate workload–strategy pool.Source: §5 · p. 6, Figure 3 · p. 6

Scale and implementation

  • The benchmark covers small dense models and reference kernels on one accelerator; its wall-clock ordering does not establish a ranking for large-scale or distributed systems.The relative orthogonalization cost depends on matrix dimensions, token batch, and kernel implementation; GPU underutilization and kernel-launch overhead were not isolated.Source: §5 · p. 6

Further limits

  • The largest test effect, 26.03%, comes from one 100:1 CIFAR-100 construction and cannot establish a general long-tailed result.Other limits include one-seed selection, short horizons, checkpoint-level crossover resolution, and fixed-pool tuning.Source: §5 · p. 7

Strength of evidence

  • The five-seed primary comparison is stronger than the three-seed variants and one-seed mechanism audits; broader evidence needs more seeds, architectures, imbalance constructions, and optimized multi-accelerator kernels.Equal-time ranking remains conditional on model size, parameter routing, token batch, and implementation.Source: §5 · p. 7

A.1. Data construction

The appendix specifies deterministic language and CIFAR-100 data construction, then describes the workloads and model configurations used for comparison.

A.1. Data construction

  • The language data use a frozen FineWeb-Edu revision, deterministic document assignment, GPT-2 tokenization, and disjoint 100M-token training, 5M-token validation, and 5M-token test splits.The evaluator reads 50 deterministic batches at each checkpoint using its frozen batch size and sequence length.Source: §A.1 · p. 10
  • CIFAR-100 uses 45,000 balanced training examples or 9,754 long-tailed examples, with class counts from 450 to 4 and a 100:1 nominal imbalance ratio.The validation set contains 5,000 examples selected with seed 2026; the official 10,000-example test split is evaluated in full.Source: §A.1 · p. 10

A.2. Models and horizons

  • The two language workloads vary depth, width, and parameter count, while the vision workloads hold the ViT architecture fixed and change the training distribution.The ViT uses 4×4 patches, six pre-normalization blocks, width 192, three attention heads, and MLP ratio 4; evaluations occur every 50 language steps and every vision epoch.Source: §A.2 · p. 10, Table 2 · p. 10

A.3. Selected hyperparameters

The appendix reports frozen primary configurations with separate learning rates and weight decay for Muon’s eligible matrices and auxiliary AdamW parameters. It also documents the 23-candidate development search and its selection rule.

A.3. Selected hyperparameters

  • Muon’s frozen configuration separates matrix learning rate and weight decay from auxiliary AdamW settings, whereas the AdamW baseline applies its values to all parameters.The table reports matrix and auxiliary learning rates separately to make parameter routing explicit.Source: §A.3 · p. 10, Table 3 · p. 11
  • Muon auxiliary weight decay is 0.10 for language and 0.05 for vision; auxiliary AdamW uses β2 = 0.95.Muon-eligible matrix counts are 4.72M, 14.16M, and 2.65M for language-micro, language-small, and each ViT, respectively.Source: §A.3 · p. 11

A.4. Tuning protocol

  • Each of 12 workload–strategy cells contains 23 candidates, selected by normalized validation-loss AUC, then final validation loss, then run identifier.AdamW and Muon select interior learning rates on all four workloads; the variance-adapted control selects audited boundary values, motivating caution in its interpretation.Source: §A.4 · p. 11

A.5. Timing and reproducibility

The appendix defines the measured training clock and optimizer-time attribution, and specifies the confirmation-run hardware, software, and test-evaluation controls.

A.5. Timing and reproducibility

  • The training clock includes forward, backward, clipping, scheduling, and optimizer steps, but pauses during validation, serialization, and logging.CUDA events surround the optimizer call; Muon’s timed call includes auxiliary AdamW, matrix momentum, Newton–Schulz iterations, scaling, and decay.Source: §A.5 · p. 11
  • Confirmation runs use deterministic algorithms and bfloat16 on an NVIDIA A100-SXM4-40GB, and test metrics are excluded from training, selection, and trajectory analysis.Checkpoint integrity is verified before each test evaluation.Source: §A.5 · p. 11

Appendix B. Full primary results

The appendix reports the frozen primary comparisons and audits how optimizer costs and restricted tuning access affect the benchmark. The language-small equal-time endpoint does not show a resolved advantage, and the tuning audit evaluates regret over a frozen candidate pool.

Appendix B. Full primary results

  • −0.73% is the language-small equal-time endpoint gain, with a paired bootstrap 95% interval of [−1.60, 0.22]%; the ordering does not establish an endpoint advantage.Table 4 distinguishes validation-loss AUC gains at equal data and time from the equal-time endpoint loss gain.Source: Appendix B · p. 12, Table 4 · p. 12
  • Table 5 reports frozen final-test endpoints as mean ± sample standard deviation over five paired seeds, with intervals for paired relative loss reduction within each workload.Table 6 reports absolute systems metrics averaged across five seeds; Figure 4 shows optimizer-step time as a fraction of measured training time.Source: Table 5 · p. 12, Table 6 · p. 12, Appendix B · p. 12

Appendix C. Tuning-budget audit

  • The tuning audit models candidate access by randomly permuting 23 configurations per workload–strategy cell and computes regret against the best AUC in the full pool.It averages 10,000 permutations per cell and then averages 12 cell curves by strategy; the completed pools include extensions informed by earlier development outcomes.Source: Appendix C · pp. 12–13
  • Figure 3 shows the budget curve and four prespecified checkpoints; its common 23-candidate endpoint measures regret against the best configuration in the same frozen pool.Figure 4’s bars compare AdamW and Muon optimizer-step time fractions, excluding validation and checkpoint I/O from the training clock.Source: Appendix C · p. 13, Figure 4 · p. 13

Appendix D. Secondary strategies

The appendices describe two secondary optimizer strategies and a one-workload mechanism audit, alongside the statistical conventions used for primary contrasts. Secondary comparisons are exploratory, and the mechanism results are descriptive rather than broad performance claims.

Appendix D. Secondary strategies

  • The long-tailed workload separates the secondary strategies sharply, so the variance-adapted result is treated as a controlled ablation rather than a uniformly stronger optimizer.Figure 5 shows paired uncertainty for the 25% AdamW-to-Muon switch and Variance-Muon; the three-seed comparisons are not power-matched to the five-seed primary study.Source: Appendix D · p. 13, Figure 5 · p. 14, Table 7 · p. 14

Appendix E. Descriptive mechanism audit

  • The mechanism audit uses language-small development data and seed 11 to compare Newton–Schulz depths and eligible-matrix roles under the same normalized validation-AUC selection rule.Because Q, K, and V are stored in one fused matrix, the audit cannot separate their individual effects; Figure 6 presents the quality–time trade-off and Table 8 compares selected three- and five-step configurations.Source: Appendix E · p. 14, Table 8 · p. 14, Figure 6 · p. 15

Appendix F. Statistical scope and reporting conventions

  • Primary contrasts pair workload and seed, use trajectories for interpolation and integration rather than independent replicates, and define equal-budget horizons by the minimum attained paired horizon.Within-workload intervals resample five paired seed effects; the pooled final-test interval is conditional on the fixed workload set, and secondary intervals use three paired seeds.Source: Appendix F · p. 15
  • Positive gains indicate lower loss or AUC for the named alternative, while accuracy effects are percentage-point differences and time metrics use validation curves.Test evaluation occurs only after training.Source: Appendix F · p. 15

Related papers

Loading 2610.21555v1…