PaperlayerView exact revision on arXiv

Source-linked AI summary

The Best Optimizer Depends on Batch Size

Xingyu Dang, Kaiyue Wen, Sadhika Malladi

arXiv:2610.08975v1cs.LGstat.ML

TL;DR

Batch-size scaling rules promise to preserve optimizer performance, but their transferability and the stability of optimizer rankings across batch sizes are open questions. The paper derives and tests Muon scaling rules, tunes optimizers across batch sizes, and analyzes noisy quadratic models; it finds that the best rule depends on setting and optimizer rankings can reverse with batch size.

  • Problem

    Scaling rules promise to preserve performance as batch size changes, while existing optimizer benchmarks leave open whether optimizer rankings persist across batch sizes.

    Source: §1 · p. 1

  • Method

    The authors derive theoretical Muon scaling rules, compare them with candidate rules across training settings, independently retune optimizers at multiple batch sizes, and analyze noisy quadratic models.

    Source: §1 · pp. 1–2

  • Results

    The best Muon scaling rule varies across settings and target batches, while SOAP’s lead at smaller batches gives way to Shampoo at the largest batch under independent tuning.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    Optimizer evaluations should tune across several batch sizes and account for how batch size changes the benefits and costs of adaptation.

    Source: §1 · p. 2

    The SDE-derived Muon rule is a candidate rather than an exact invariance guarantee because the implementation ties its Nesterov mixing weight to momentum across batch sizes.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

A plethora of new adaptive optimizers are designed to efficiently estimate and use minibatch gradient statistics to shape parameter updates, but they are typically benchmarked at a single batch size. Hyperparameter scaling rules promise to preserve performance as batch size and gradient noise change, suggesting that the best optimizer at one batch size should remain the best at another. We challenge this approach to developing and evaluating optimizers by showing: (1) no principled scaling rule for Muon works consistently across training settings, and (2) the best optimizer for language model pretraining changes with batch size even after extensive hyperparameter tuning.

1 Introduction

The paper tests whether batch-size scaling rules transfer across training settings and whether optimizer rankings remain stable after independent tuning at each batch size. It finds that neither scaling prescriptions nor optimizer rankings are reliably invariant to batch size or setting.

1 Introduction

  • Neither theoretical Muon scaling prescription consistently performs best across settings, and both fall behind no scaling in language modeling at the largest batches.Among broader candidates, the best rule also varies by task and target batch size.Source: §1 · p. 1
  • SOAP leads at smaller batches, but Shampoo leads at the largest tested batch under both decoupled weight decay and HyperBall.The crossover persists despite independent tuning at each batch size.Source: §1 · p. 2, Figure 1 · p. 2
  • A noisy quadratic model explains how direction-dependent curvature and gradient noise can shift tuned learning-rate scaling from approximately square-root toward linear.The analysis also proves that a bias–variance tradeoff can favor SGD at small batches and Newton’s method at large batches, even with independently optimal learning rate and momentum.Source: §1 · p. 2
  • Retaining small-batch updates in the sharpest 768 of roughly 85M preconditioned-Hessian directions removes 35–59.5% of the loss gap over the next 8% of training.A random subspace of the same size changes the gap by at most 3%; the paper assesses loss shortly after the switch because the quadratic approximation is local.Source: §1 · p. 2
  • The authors recommend tuning optimizers across several batch sizes and accounting for how batch size changes adaptation’s benefits and costs.They argue that a single-batch comparison is insufficient to establish optimizer efficacy.Source: §1 · p. 2

2 Preliminaries

The preliminaries define batch-size scaling as transferring tuned hyperparameters from a reference batch to a larger batch. They distinguish rules derived by preserving training properties from rules derived by minimizing convergence bounds, then specify the MuonW–AdamW optimizer used in the paper.

2.1 Scaling Rules

  • A batch-size scaling rule maps hyperparameters tuned at reference batch B0 to prescribed values at B = κB0, with training budget fixed in tokens or examples.The paper contrasts preserving training invariants with minimizing a convergence bound to derive such rules.Source: §2.1 · pp. 2–3
  • For SGD, preserving η/B in an SDE approximation yields linear learning-rate scaling, while related approaches preserve other quantities such as moment memory or weight-decay timescales.The SDE describes random updates through continuous-time dynamics; the SGD noise covariance scales as Σ(x)/B.Source: §2.1 · p. 3

2.2 MuonW-AdamW Optimizer

  • The paper’s MuonW–AdamW setup uses MuonW and AdamW parameter groups, with Muon momentum updated before orthogonalization by a Newton–Schulz map.The implementation applies decoupled weight decay; AdamW’s powers, roots, and division are elementwise.Source: §2.2 · pp. 3–4

3 Deriving and Testing Batch-Size Scaling Rules

The paper derives two theoretically motivated Muon scaling rules and compares them with a broad sweep of candidate prescriptions. No single rule is best across the tested settings or target batch sizes.

3.1 Two Scaling Rules for Muon

  • The SDE-derived Muon rule scales learning rate and momentum to match modeled sample-time dynamics, but tying the implementation’s Nesterov mixing weight to momentum makes it a candidate rather than an exact invariance guarantee.The derivation holds that mixing weight fixed while matching the SDE coefficients.Source: §3.1 · p. 4
  • The bound-minimization candidate balances an update-count term against a smoothness cost, yielding square-root growth in effective step size as batch size increases.The constrained minimizer follows square-root scaling until the effective step-size cap at 1 becomes active.Source: §3.1 · pp. 4–5

3.2 Sweeping Scaling Rules

  • The authors compare more than 200 candidate MuonW–AdamW scaling rules, evaluating 216 complete rules in language modeling and 648 on CIFAR-5M.The search varies learning rates, weight decay, momentum, and Adam moment coefficients; rules are selected by mean batch-size regret across target batches.Source: §3.2 · p. 5
  • The best scaling prescription differs by setting and target batch: language modeling favors fixed learning rates with linearly scaled matrix weight decay at 1M–2M, while CIFAR-5M favors square-root-scaled learning rates.The empirical LLM rule stays closest to full retuning in language modeling, whereas bound minimization stays within 0.006 nats of full retuning on CIFAR-5M; Figure 2 compares the transfer rules across both settings.Source: §3.2 · pp. 5–6, Figure 2 · p. 5
  • The selected rule within language modeling also changes with target batch: square-root scaling for both learning rates at 256K–512K gives way to fixed rates and linear matrix weight-decay scaling at 1M–2M.Keeping Muon momentum and Adam’s first-moment coefficient fixed is generally preferable in the sweeps.Source: §3.2 · p. 6

4 Batch Size Changes the Best Optimizer

After independently tuning optimizers at each batch size, the authors find that the lowest-loss optimizer changes as batch size increases. The crossover appears under both weight decay and HyperBall, limiting what single-batch optimizer comparisons establish.

4 Batch Size Changes the Best Optimizer

  • The benchmark compares Adam, Lion, Muon, SOAP, and Shampoo across 128K, 512K, 1M, and 2M tokens with an approximately 1.7B-token budget.Each optimizer is extensively tuned at each batch size; weight decay and HyperBall provide two hidden-matrix norm-control mechanisms.Source: §4 · p. 6
  • SOAP is best at smaller batches, while Shampoo is best at the largest tested batch; the optimizer crossover also appears under HyperBall.Figure 3 shows the HyperBall comparison, with KL-SOAP lowest from 128K through 1M and Shampoo lowest at 2M.Source: §4 · p. 6, Figure 3 · p. 6
  • Crossovers occur at different batch sizes and between different optimizer pairs, and a SOAP-to-Shampoo crossover also appears in the nanochat setting.The authors conclude that comparisons at a single batch size provide only a limited understanding of optimizer efficacy.Source: §4 · p. 6

5 Curvature and Noise Shape Batch-Size Scaling

In noisy quadratic models, curvature-to-noise differences make batch-size scaling depend on the setting and can reverse optimizer rankings even after retuning. The analysis connects these effects to direction-dependent bias–variance tradeoffs.

5 Curvature and Noise Shape Batch-Size Scaling

  • The noisy quadratic model isolates curvature and gradient-noise effects; on a 1 × 1 block, Muon reduces to SignSGD with momentum, so the analysis omits weight decay.Source: §5 · pp. 6–7

5.1 Scaling Rules Depend on the Training Setting

  • The better SDE- versus bound-minimization scaling rule depends on both reference and target batch sizes under a fixed sample budget.The comparison is against independently retuned target-batch hyperparameters.Source: §5.1 · p. 7, Figure 4 · p. 7
  • As CNR rises, tuned SignSGD-M learning rates shift approximately from square-root toward linear scaling.Noise-dominated movement motivates square-root scaling, while signal-dominated steps saturate and motivate linear scaling.Source: §5.1 · p. 8

5.2 Batch Size Can Reverse Optimizer Rankings After Retuning

  • Across flat noisy and sharp noiseless directions, Newton reduces initialization error equally but amplifies flat-direction noise by 1/h_i, creating a batch-dependent bias–variance balance.The theorem shows this tradeoff can reverse the optimally tuned ranking.Source: §5.2 · p. 8
  • Optimally tuned SGD can outperform Newton at B = 1 yet lose at B = T, with or without momentum; a three-direction NQM also shows a retuned optimizer crossover.Crossovers also appear on three matrix regression tasks.Source: §5.2 · p. 8, Figure 5 · p. 8

6 Sharp Directions Carry Most of the Local Batch-Size Penalty in Pretraining

The pretraining experiment branches from a 128K-token MuonW run to test larger-batch updates while retaining small-batch updates in selected directions. Holding sharp directions recovers part of the 2M branch penalty, whereas a random subspace has little effect.

6 Sharp Directions Carry Most of the Local Batch-Size Penalty in Pretraining

  • From 128K MuonW checkpoints, 1,024-step branches compare base updates with 2M updates emulated by averaging 16 gradients and square-root-scaled learning rates.Source: §6 · p. 8
  • 35–59.5% of the 2M branch penalty is removed by holding the sharpest 0.001% of directions at 128K for branches started at steps 1,000–9,000; random 768-dimensional holding changes the penalty by at most 3%.For branches started at steps 11,000 and 12,000, no held subspace significantly changes the penalty after learning-rate decay.Source: Figure 6 · p. 9, §6 · p. 9
  • The held subspace stays fixed per branch and selects top eigenvectors of the Muon-preconditioned Gauss–Newton proxy; gradient-noise anisotropy is ignored.A random 768-dimensional subspace serves as a baseline.Source: §6 · p. 9

A Related Work

Prior work develops batch-size scaling rules and optimizer comparisons, while this section distinguishes direction-specific scaling behavior from scalar summaries and sets up two Muon scaling analyses under a fixed data budget.

A Related Work

  • Prior batch-size scaling prescriptions preserve different invariants, and theoretical derivations for Muon can disagree or fail to test whether transfer is optimal.The literature includes SDE-based rules and rules derived by minimizing convergence bounds; some experiments apply square-root scaling without testing its optimality.Source: A Related Work · p. 14
  • Critical-batch-size summaries reduce training to a scalar, while directions with different curvature-to-noise ratios can favor different batch-size scaling rules.Prior work also finds that the critical batch size depends on the model and optimizer, and to a lesser extent the dataset.Source: A Related Work · pp. 14–15
  • KL-SOAP leads at 128K tokens per batch and Shampoo at 2M after five optimizers are retuned, with the shift holding under both weight decay and HyperBall.The comparison fixes the model, data distribution, and token budget; the large-batch shift differs from a prior SOAP–Shampoo comparison using a different SOAP variant and model.Source: A Related Work · p. 15
  • Batch-size-invariant methods modify the training algorithm, whereas this analysis keeps the MuonW–AdamW hybrid fixed and rescales its hyperparameters.The cited Batch Size Invariant Adam approach changes how second moments are updated.Source: A Related Work · p. 15

B Derivations and comparison of batch-size scaling rules

  • The appendix derives two Muon scaling candidates: one preserves a continuous-time model at matched data consumption, and the other minimizes a stationarity bound at fixed data budget.Their different objectives yield different learning-rate rules.Source: B Derivations and comparison of… · p. 16

B.1 Setup and notation

  • Comparisons hold the model and data budget τ fixed as batch size changes, so the number of updates is n = τ/B.The budget uses tokens for language modeling and examples for image classification; learning-rate schedule shape is held fixed against the fraction of budget consumed.Source: §B.1 · p. 16
  • The minibatch gradient is the mean gradient plus B^-1/2-scaled zero-mean noise with conditional covariance Σ(θk); matrix covariance is defined after vectorization.The optimizer notation permits momentum retention µ and Nesterov mixing weight ω to vary independently, although experimental MuonW ties ω = µ.Source: §B.1 · p. 16
  • In the theoretical recursion, Ψ = Orth gives the studied Muon variant, while experimental MuonW restricts the parameters by tying ω = µ.That restriction is not imposed on theoretical batch-size transport.Source: §B.1 · p. 16

B.2 Matching the sample-time Muon diffusion

This section examines conditional batch-size scaling for Muon, including rules that target its sample-time diffusion and separate prescriptions based on other invariants. It also distinguishes transporting a fixed optimizer from selecting one by tuned finite-budget risk, framing batch size in terms of gradient signal-to-noise ratios and available updates.

B.2 Matching the sample-time Muon diffusion

  • For B′ = κB, the stated transformation of η, µ, ω, and λ preserves the sample-time diffusion law under globally Lipschitz, linearly growing coefficients.The result concerns the continuous-time diffusion at matched sample time, not necessarily the discrete trajectory.Source: §B.2 · p. 17
  • Under bounded-noise, smoothness, and momentum-lag assumptions, Proposition B.2 gives a finite-horizon diffusion approximation with a coupling to the discrete updates.Its error constant depends on the stated regularity parameters and the initial state's second moment.Source: §B.2 · pp. 17–18

B.3 Minimizing a stochastic Frank–Wolfe bound

  • Minimizing the stochastic Frank–Wolfe bound at fixed weight decay gives a learning-rate rule proportional to √B when Δ1 > 0, until the effective step-size cap becomes active.The theorem assumes an exact linear minimization oracle and fixed momentum parameters; when Δ1 = 0, the bound is minimized as γ ↓ 0.Source: §B.3 · pp. 19, 22

B.4 Comparison with the published prescriptions

  • Exact sample-memory preservation uses q′ = q^κ, whereas linear refresh scaling uses q′ = 1 − κ(1 − q); the formulas agree only to first order near q = 1.The passage distinguishes these finite-change prescriptions and their sample-time half-life interpretation.Source: §B.4 · p. 22
  • The Adam SDE prescription of Malladi et al. (2022) uses square-root learning-rate scaling and affine moment refresh, while Completed Hyperparameter Transfer also scales weight decay at fixed model size and token budget (Mlodozeniec et al., 2026).The appendix notes that discrete-trajectory use depends on the respective SDE approximation conditions.Source: §B.4 · p. 22
  • Monzio Compagnoni et al. (2025) preserve an infinite-time asymptotic-loss bound under strong-convexity, smoothness, and noise-regime assumptions, a different invariant from trajectory matching.Their stated prescription scales η and λ by √κ and uses affine refresh for βi.Source: §B.4 · p. 22
  • Cesista & Wen (2025, Section 5) derive η′ = κη at fixed token budget from a fixed per-step update-noise variance heuristic, distinct from their reported empirical square-root rule for AdamW and Muon.The fixed-budget conversion assumes the per-step variance is proportional to η2/B with a fixed proportionality factor.Source: §B.4 · p. 22
  • The published timescale prescription preserves ηλ via η′λ′ = κηλ, with Power Lines favoring p = 0, but its nominal timescale is approximate and may drift empirically.The timescale is 1/(ηλ); the preferred timescale also drifts at large batches and with tokens per model parameter. The averaging identity also applies to MuonW with its direction substituted for the update.Source: §B.4 · p. 23

C.1 Finite-budget risks and optimizer comparisons

  • At fixed horizon T, batch size changes both directional minibatch-gradient signal-to-noise ratios and update count T/B, so optimizer choice compares families’ tuned finite-budget risks at the same batch size.The diagonal quadratic model gives each coordinate its own curvature, noise level, and signal-to-noise ratio.Source: §C.1 · p. 24

C.1.1 Setup, noise ratios, and tuned finite-budget risk

  • The diagonal quadratic setup assumes curvature and single-example gradient-noise covariance are simultaneously diagonalizable, with coordinate minibatch-noise variance c_i/B.Increasing batch size raises the initial coordinate signal-to-noise ratio unevenly when the coordinate noise-to-curvature ratios differ.Source: §C.1.1 · pp. 24–25
  • The batch-B comparison fixes horizon T, uses S = T/B gradient queries, and allows each optimizer family’s hyperparameters to be retuned at every batch size.The criterion compares the resulting tuned risks across optimizer families at the same (B, T).Source: §C.1.1 · p. 25

C.1.2 Risk decomposition and crossover proof

This section analyzes how batch size and gradient noise shape optimizer risk and response, deriving an SGD–Newton risk crossover and examining tuned learning-rate behavior, momentum-SignSGD, and RMSProp. It also characterizes RMSProp’s stationary response across noise- and signal-dominated regimes and tracks conditional gradient signal-to-noise during training under fixed CNR.

C.1.2 Risk decomposition and crossover proof

  • At B = 1, a feasible SGD configuration beats Newton’s uniform lower bound; at B = T, Newton’s one-step risk beats SGD’s optimized bias lower bound for sufficiently large T.The Newton small-batch lower bound and large-batch upper bound are both 1/(8T) + o(T^-1).Source: §C.1.2 · pp. 26–27
  • For dyadic budget T, at least one adjacent pair of dyadic batches has a risk difference whose sign changes, allowing an intermediate tie.The proof chooses the first batch where the difference becomes positive; the preceding batch has a nonpositive difference.Source: §C.1.2 · p. 27

C.1.4 Learning-rate scaling under a finite budget

  • Tuned learning rates do not remain on the linear-scaling reference: Newton bends with the horizon T/B, while SignSGD and RMSProp bend with sign reliability and empirical-Fisher feedback.Figure 7 compares tuned rates with linear scaling anchored at B = 2.Source: §C.1.4 · p. 27, Figure 7 · p. 28

C.1.5 Optimizer crossover in the noisy quadratic model

  • Figure 8’s tuned terminal-risk comparison shows optimizer-family rankings change with batch size, even after each family is retuned at every batch size.The compared families are SGD-M, Newton-M, SignSGD-M, and Adam; Adam also tunes β2, and momentum improves small-batch SignSGD behavior.Source: §C.1.5 · p. 28, Figure 8 · p. 29

C.2 Finite-budget SignSGD batch-size transfer

  • A batch-size scaling rule transports hyperparameters along a low-dimensional path, but need not reach the family optimum because coordinates can demand different behavior.The target risk depends on effective coordinate tuples, while SignSGD depends on sign reliability and RMSProp’s empirical-Fisher mean depends on the stationary law of θi.Source: §C.2 · p. 29

C.3.1 Shared movement ansatz

  • Matching frozen-state mean movement per processed sample yields a batch-dependent learning-rate ansatz; momentum-SignSGD’s response interpolates from η ∝ √B under noise dominance to η ∝ B under signal dominance.Momentum changes how much noise is needed to reach the square-root regime, not the two limiting scalings.Source: §C.3.1 · p. 30, §C.3.2 · p. 30

C.3.2 Momentum-SignSGD response

  • 0.591 vs. 0.588: the frozen-state ansatz predicts and measures this exponent at CNR = 0.001; at higher CNR it overpredicts the measured exponents.At µ = 0.9, predicted/measured exponents are 0.911/0.794 for CNR = 0.03 and 1.000/0.904 for CNR = 1; the ansatz freezes w at its initial scale, unlike tuned trajectories near the optimum.Source: §C.3.2 · p. 30

C.3.3 Momentum-SignSGD learning-rate scaling across CNR conditions

  • 27.2×, 67.7×, and 135.8×: selected µ = 0.9 learning rates rise by these factors from B = 1 to B = 256 across increasing CNR conditions.For CNR = 0.001, 0.03, and 1, fitted exponents were 0.588, 0.794, and 0.904 at µ = 0.9, versus 0.505, 0.591, and 0.853 without momentum.Source: §C.3.3 · p. 31

C.3.4 RMSProp response

  • RMSProp’s local logarithmic slope approaches 1/2 when noise dominates and 1 when signal dominates.The slope is 1 − [2(1 + B CNR)]^-1.Source: §C.3.4 · p. 31
  • The dashed curves in Figure 9 use measured stationary responses because E[g/√v] need not equal E[g]/E[v].This avoids assuming the mean-denominator approximation is exact.Source: §C.3.4 · p. 31
  • The conditional gradient signal-to-noise ratio along training is Bw2 CNR, which changes even though CNR remains fixed.Source: §C.3.4 · p. 31

C.3.5 RMSProp learning-rate scaling across CNR conditions

RMSProp’s tuned learning-rate scaling shifts from approximately square-root to linear as the curvature-to-noise ratio (CNR) rises. Adam shows the same CNR dependence, while tuned SGD remains near linear across the tested conditions.

C.3.5 RMSProp learning-rate scaling across CNR conditions

  • RMSProp’s normalization denominator shifts from noise-dominated to signal-dominated as batch size increases, changing which learning-rate scaling compensates for fewer updates.The one-dimensional conditional second moment is h^2w^2 + c/B: the signal term stays fixed at a given parameter value, while the noise term falls with batch size.Source: §C.3.5 · p. 31
  • Adam’s endpoint slopes also rise with CNR, while tuned SGD stays between 0.96 and 0.99 across the three conditions.For Adam, slopes rise from 0.47 to 0.94 at β2 = 0.9 and from 0.48 to 0.95 at β2 = 0.99; Figure 11 shows the corresponding sweep.Source: §C.3.5 · pp. 31–32, Figure 11 · p. 33
  • The RMSProp sweep tunes terminal loss over batches B = 2^j, j = 0,…,8, with a fixed T = 4096-sample budget and independent trajectory-based evaluation.Selected candidates’ held-out loss was at most 1.00838 times the minimum on the validation grid.Source: §C.3.5 · pp. 31–32
  • The tuned RMSProp exponent rises from 0.48 at CNR = 0.01 to 0.75 at CNR = 1 and 1.02 at CNR = 100.These endpoint slopes are 0.477, 0.748, and 1.017, respectively, for batches from B = 1 to B = 256.Source: Figure 9 · p. 31, §C.3.5 · pp. 31–32

D Additional Matrix-Model Results

At higher noise, optimizer crossovers also occur across matrix regression tasks, with Muon best at small batches and EMA-Shampoo at large batches in matrix linear and logistic regression. The paper evaluates scaling rules across language modeling and CIFAR-5M and defines the matrix models used in additional results.

D Additional Matrix-Model Results

  • At higher noise, the best optimizer changes from Muon at small batches to EMA-Shampoo at large batches in matrix linear and logistic regression.The paper also reports optimizer crossovers on a noisy quadratic and presents tuned optimizer families across three matrix tasks at baseline and higher noise levels.Source: D Additional Matrix-Model Results · p. 34, Figure 13 · p. 35

D.1 Matrix-model definitions

  • The matrix tests use rank-one covariates in linear, logistic, and normalized linear regression, exposing feature geometry, state-dependent curvature, and radial-null-direction dynamics.In normalized linear regression, positive rescaling leaves the objective invariant and the optimizer dynamics are angular.Source: §D.1 · p. 34

E Scaling-rule sweeps on CIFAR-5M and language modeling

  • The scaling-rule comparison tests whether MuonW–AdamW hyperparameter choices transfer across batch sizes and training settings, using complete grids in CIFAR-5M and language modeling.The grids contain 648 CIFAR-5M rules and 216 language-model rules; auxiliary weight decay is fixed in the latter.Source: E Scaling-rule sweeps on CIFAR-5M and… · p. 36, §E.1 · p. 36

E.1 Shared design and evaluation

  • Learning rates and weight decays use fixed, square-root, or linear scaling, while retention coefficients either stay fixed or preserve decay in data units or over the run.The coordinates include MuonW and auxiliary learning rates and weight decays, Muon momentum, and Adam’s two moment-retention coefficients.Source: §E.1 · p. 36
  • The sweeps define batch-size regret against the best tested rule at each target batch and select a common rule by its mean regret across target batches.Coordinate panels fix one coordinate choice and average loss over all combinations of the others.Source: §E.1 · p. 36

E.2 CIFAR-5M sweep

On CIFAR-5M, square-root learning rates with the other coordinates fixed minimize mean batch-size regret among the tested scaling rules. The individual winning rules nevertheless change between target batch sizes.

E.2 CIFAR-5M sweep

  • The CIFAR-5M comparison trains a GroupNorm ResNet over one pass through the data, holding processed examples and schedule fractions fixed across target batches.The reference batch is B0 = 256; target batches range from 64 to 4096, and validation loss is averaged over the final ten evaluations.Source: §E.2 · p. 36
  • 0.00492 mean and 0.00913 maximum batch-size regret make square-root learning rates, with other coordinates fixed, the best common rule.At batches 64, 128, 512, 1024, 2048, and 4096, its losses are 0.23370, 0.23136, 0.23373, 0.22778, 0.23290, and 0.23683.Source: §E.2 · p. 36
  • The best individual rule changes with batch size: linear scaling wins at B = 64, while at B = 4096 the Muon learning rate and auxiliary weight decay scale by the square root.At B = 64, both learning rates and both weight decays scale linearly; at B = 4096, β2 is transported and the remaining coordinates stay fixed. Figure 14 summarizes coordinate-wise mean losses across the 648 rules.Source: §E.2 · p. 37, Figure 14 · p. 37
  • The full-retuning comparison tunes all seven hyperparameters independently at each target from 512 to 4096, with repeated-seed acceptance checks for candidate improvements.Retuned-configuration losses span up to 0.0064; final-window standard deviations describe variation within a run, not confidence intervals across seeds.Source: §E.2 · p. 37

E.3 Language-model sweep

The language-model sweep examines how scaling choices and optimizer rankings vary across target batch sizes: lowest-loss scaling choices change with batch size, while a common fixed-learning-rate, linear matrix weight-decay rule performs well across targets, and independently retuned runs shift from KL-SOAP at 128K to Shampoo at 2M under both decoupled weight decay and HyperBall.

E.3 Language-model sweep

  • The lowest-loss rules use square-root scaling of both learning-rate groups at 256K and 512K, but fixed rates and linear matrix weight decay at 1M and 2M; all transport Adam β2 and fix Muon momentum and Adam β1.Table 5 reports the batch-specific winners.Source: §E.3 · p. 38, Table 5 · p. 38
  • With fixed learning rates and linear matrix weight-decay scaling, the selected common rule keeps Muon momentum and Adam β1 fixed, transports only Adam β2, and has mean and maximum BSZ regret 0.00151 and 0.00416.Transporting β2 lowers mean loss by 0.000895 but raises loss by 0.00017 at 256K.Source: §E.3 · pp. 38–39
  • At 1M and 2M tokens, square-root auxiliary-decay scaling raises median loss by 0.001 and 0.005 nats across rules, while linear scaling raises it by 0.012 and 0.033 nats.For the common rule, square-root scaling instead reduces loss by 0.0037 and 0.0027 nats at those targets.Source: §E.3 · p. 40, Table 9 · p. 51

E.4 Comparison across settings

  • The selected common rules differ across settings—fixed rates and linear matrix weight decay for language modeling versus square-root rates and fixed weight decays for CIFAR-5M—and bound-minimizing versus trajectory-matching ranks are 52.1 versus 141.9 and 59.9 versus 347.7, respectively.Both selected rules fix Muon momentum and Adam β1; only the language-model rule transports β2.Source: §E.4 · p. 40

E.5 Complete measurements

  • All 864 original language-model runs reached their terminal step with finite loss; Table 9 combines those runs with 864 auxiliary weight-decay measurements evaluated only at 1M and 2M.The additional measurements include square-root and linear auxiliary-decay variants of each original rule.Source: §E.5 · p. 41, Table 9 · p. 51

F Language-Model Optimizer Crossover Experiments

  • The sweep compares Adam, Lion, Muon, KL-SOAP, and Shampoo at 128K, 512K, 1M, and 2M tokens with an approximately 1.7B-token budget, independently tuning each optimizer at each batch size.Configurations are ranked by terminal validation loss, with lower values better; a crossover means the ordering reverses across batch sizes after retuning.Source: F Language-Model Optimizer Crossover… · p. 60

F.1 Implementations and common protocol

  • All methods share the model, token stream, validation set, precision, distributed code, and schedule shape; weight decay applies decoupled matrix decay, whereas HyperBall uses a norm-preserving matrix projection without matrix decay.Preconditioner frequency, initialization, and schedule family are not tuned.Source: §F.1 · p. 60

F.2 Optimizer definitions

  • HyperBall acts only on hidden matrices, other parameter groups use bias-corrected Adam, and KL-SOAP refreshes its covariance bases every step.Source: §F.2 · pp. 60–61

F.3 Coordinate-descent protocol

  • Coordinate-descent candidates are accepted only when median validation loss improves by more than 0.002 under a multi-seed repeat policy, and final centers are checked against adjacent grid neighbors.Gains from 0.002 to under 0.006 require three matching repeats; larger gains require two, and selected centers have at least two repeats.Source: §F.3 · p. 62

F.4 Final weight-decay hyperparameters

  • At 128K, KL-SOAP-W records terminal validation loss 3.2536 versus ShampooW’s 3.2650; at 2M, ShampooW leads at 3.3622 versus KL-SOAP-W’s 3.3833.The reversal persists despite independent hyperparameter tuning at each batch size.Source: §F.4 · p. 62, Table 11 · p. 63

F.5 Optimizer crossover under HyperBall

  • MuonH beats ShampooH at 128K but loses at each larger batch, while LionH beats AdamH only at 128K; both pairs’ members remain above MuonH, KL-SOAP-H, and ShampooH throughout.The MuonH–ShampooH difference at 512K is about 0.0007.Source: §F.5 · p. 63
  • At 128K, KL-SOAP-H has the lowest loss at 3.2585, while at 2M ShampooH leads at 3.3390; at 1M their losses are 3.2962 and 3.2965, respectively.The selected HyperBall configurations use no matrix or auxiliary weight decay and are reported in Tables 13 and 14.Source: §F.5 · p. 63, Table 13 · p. 64, §F.6 · p. 64

F.7 Complete One-Coordinate Sweep Results

Appendix F.7 presents one-coordinate hyperparameter sensitivity across optimizer–batch settings and compares branch-penalty trajectories across anchors for branch variants against same-anchor controls. These results provide local sensitivity evidence, not a new configuration-selection rule.

F.7 Complete One-Coordinate Sweep Results

  • 643 completed one-coordinate trials cover 40 optimizer–batch settings, with each trial varying one hyperparameter from its accepted final configuration and reporting terminal validation loss.Only trials with a validation result at the configured final step are included; repeated trials use the latest completed result, without choosing by validation loss.Source: §F.7 · p. 65
  • The three auxiliary learning rates are tied during search, so changes to ηOut. and ηScalar track the absolute embedding rate ηEmb. by the same factor.The reported metric is terminal validation loss, not the best loss observed earlier in training.Source: §F.7 · p. 65
  • A lower single-seed arm is not accepted unless it also meets the multi-seed threshold policy in Appendix F.3.The sweeps are a local sensitivity record rather than a configuration-selection rule.Source: §F.7 · p. 65

Optimizer ablations

  • The ablation tables cover MuonH, KL-SOAP-H, ShampooH, AdamH, LionH, MuonW, KL-SOAP-W, ShampooW, AdamW, and LionW across batch-size settings.Each table reports a hyperparameter ablation; the supplied result passages expose terminal-loss values but do not preserve enough column labels to attribute individual values to specific coordinates.Source: Table 15 · p. 65, Table 19 · p. 67, Table 23 · p. 69

G Directional Branching in the Language Model

  • Figure 17 tracks each 2M branch’s held-out-loss difference from its same-anchor 128K control across fully scaled, top-r-held, and random-768-held variants.Trajectories use 1M held-out tokens, evaluated every 32 base steps and smoothed over three evaluations. Table 55 endpoints use 10.5M validation tokens; top-128 was not run at anchors 7,000, 9,000, and 11,000.Source: Figure 17 · p. 89

G.1 A local noisy-quadratic model for Muon

The section constructs a local noisy-quadratic model for Muon by freezing its update Jacobian and Gauss–Newton matrix at an anchor. It establishes a local correspondence with momentum/Nesterov SGD and bounds finite-horizon deviations under local assumptions.

G.1 A local noisy-quadratic model for Muon

  • P0 is positive semidefinite because it is the derivative of the ideal polar direction map at the anchor input.The model describes centered trajectory deviations; it does not replace the absolute update with P0q.Source: §G.1 · pp. 90–91
  • The frozen Muon model transforms exactly into momentum/Nesterov SGD on the range of P0, with quadratic curvature and minibatch noise covariance eΣ/B.The correspondence assumes full-rank polar inputs and retains the same schedules and decoupled weight decay.Source: §G.1 · pp. 89–90
  • The coordinate correspondence is invertible for parameter deviations in range P0, and transforms the noise covariance into the reduced coordinates.The reduction omits momentum components in ker P0 because they do not affect frozen parameter updates.Source: §G.1 · p. 90
  • The finite-horizon comparison holds under local segment assumptions, but small one-step errors alone do not guarantee a uniform-in-time approximation without stability control.The bound makes no contraction assumption and compares trajectories in the original state coordinates.Source: §G.1 · pp. 90–91
  • The local correspondence motivates diagnostic geometry, but does not prove branch-penalty predictions or transfer SignSGD scaling exponents to Muon.A quantitative guarantee for the finite Newton–Schulz variant also requires bounding its map defect.Source: §G.1 · p. 91

H Batch Size and Optimizer Compute Efficiency

  • The separate compute-efficiency experiment transfers hyperparameters tuned at 256K tokens to larger batches and compares compute needed to reach the same loss.It evaluates MuonW, KL-SOAP-W, ShampooW, AdamW, and LionW, using both the same optimizer and MuonW as reference.Source: H Batch Size and Optimizer Compute… · p. 92

H.1 Model, data, and training budgets

The appendix describes the training setup and transfer protocol for comparing optimizers across batch sizes. In the coefficient-transfer search, auxiliary-only EMA transfer is selected for KL-SOAP-W and ShampooW, while the token-horizon reference curves use prescribed transfers without independent retuning.

H.1 Model, data, and training budgets

  • The nanochat experiment uses a depth-8, width-512 decoder and FineWeb-Edu shards 0–238 for training, with shard 239 reserved for validation.Each optimizer is tuned at a 256K-token reference batch and a horizon of eight training tokens per scaling parameter.Source: §H.1 · p. 92

H.2 Anchor search

  • The anchor search uses a matrix learning-rate and weight-decay grid followed by one-hyperparameter-at-a-time searches, producing 178 completed observations.Search ranges are extended when the best observed value lies at a boundary.Source: §H.2 · p. 92, Table 57 · p. 92

H.3 Transfer across the token horizon

  • The token-horizon reference curves apply prescribed learning-rate and weight-decay transfers, without independent retuning at each token budget.The authors have not established that the combined transfer is best across token horizons or optimizers in this setting.Source: §H.3 · p. 93

H.4 Scaling moving-average coefficients across batch sizes

  • For the batch comparison, peak learning rates remain fixed while matrix weight decay increases linearly, and the same choices are applied to all five optimizers without retuning.The comparison uses 256K, 512K, 1M, and 2M-token batches at a fixed horizon of eight training tokens per scaling parameter.Source: §H.4 · p. 93
  • 0.001233 BPB for KL-SOAP-W and 0.000400 BPB for ShampooW are the mean batch-size regrets of the selected auxiliary-only EMA transfer rules.For both optimizers, the selected rule transfers auxiliary AdamW β2 and keeps other moving-average coefficients fixed.Source: §H.4 · p. 93, Table 58 · p. 94
  • At 2M tokens, auxiliary-only EMA transfer gives lower loss than also transferring matrix coefficients for both KL-SOAP-W and ShampooW.The comparison rule transfers matrix moving-average coefficients to preserve their decay in token units.Source: §H.4 · p. 93

H.5 Compute efficiency metric

This section defines compute efficiency against fitted 256K reference curves, then distinguishes optimizer-specific transfer efficiency from absolute cross-optimizer comparison.

H.5 Compute efficiency metric

  • Compute efficiency compares equivalent compute on an optimizer’s fitted 256K reference curve with the actual compute used by a larger-batch run.Equivalent compute is the compute needed on the reference curve to reach the observed loss; compute efficiency is its ratio to actual compute.Source: §H.5 · p. 94
  • A penalty is one minus compute efficiency: zero matches the 256K reference at equal loss, while a positive value means the larger batch used more compute.A penalty of 0.6 means the reference curve reaches that loss with 40 percent of the run’s compute.Source: §H.5 · p. 94
  • Figure 18 plots validation losses across token budgets at 256K and optimizer-specific fitted curves, each based on seven observations; Table 59 reports fit parameters.Each optimizer’s own 256K anchor is assigned zero penalty by convention, even if its fitted curve does not pass exactly through that observation.Source: §H.5 · p. 94, Table 59 · p. 94, Figure 18 · p. 95
  • Optimizer-specific penalties measure transfer efficiency but cannot rank absolute optimizer performance because each optimizer uses a different fitted reference curve.For a shared comparison, losses are instead mapped through the MuonW curve, preserving absolute-loss ordering at the fixed budget.Source: §H.5 · p. 95

H.6 Results

At the largest batch, ShampooW has lower absolute validation loss than KL-SOAP-W, even though KL-SOAP-W retains more efficiency against its own reference curve. These results use a batch-size and hyperparameter-transfer protocol, while the complete results section reports the observations and coefficient-transfer rankings.

H.6 Results

  • 0.1365, 0.3462, and 0.6116 are KL-SOAP-W’s own-reference penalties at 512K, 1M, and 2M, respectively—the smallest at each larger batch.At 2M, ShampooW’s own-reference penalty is 0.6320, compared with 0.7141 for MuonW, 0.7761 for LionW, and 0.8133 for AdamW.Source: §H.6 · p. 95, Table 60 · p. 96
  • The penalties combine batch-size changes with hyperparameter transfer, rather than measuring efficiency after independent retuning at each batch.The change in winner agrees qualitatively with the KL-SOAP-to-Shampoo crossover reported in Section 4.Source: §H.6 · p. 96
  • 1.1525 vs. 1.1614 validation loss at 2M: ShampooW beats KL-SOAP-W, although KL-SOAP-W has smaller shared-reference penalties at 512K and 1M.At 2M, ShampooW’s shared-reference penalty is 0.6440, compared with 0.6702 for KL-SOAP-W.Source: §H.6 · pp. 95–96, Table 60 · p. 96
  • At 2M under the shared MuonW curve, AdamW and LionW results fall outside the observed MuonW loss range and are extrapolations; the remaining common-reference values are interpolations.Optimizer-specific inversions fall within the loss ranges covered by their respective seven reference observations.Source: §H.6 · p. 96

H.7 Complete results

  • Table 61 contains the 50 observations used for reference fits and batch comparisons, while Table 62 gives coefficient-transfer rankings for KL-SOAP-W and ShampooW.Mean BSZ regret averages each rule’s gap to the best tested rule for the same optimizer and batch size over three target batches.Source: §H.7 · p. 96, Table 61 · p. 97, Table 62 · p. 98
  • Table 62 labels fixed coefficients F and EMA transfer E; ShampooW shares its matrix βM 2 coefficient with the preconditioner, so no separate preconditioner choice is listed.The table’s mean BSZ regret summarizes rule gaps across the three target batches.Source: Table 62 · p. 98

Related papers

Loading 2610.08975v1…