cs.LGOct 6, 2026

The Best Optimizer Depends on Batch Size

Authors: Xingyu Dang, Kaiyue Wen, Sadhika Malladi

Organizations: Princeton University · Stanford University · University of California, San Diego

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.

Figures & tables

Figure 1: The best optimizer changes with batch size. FineWeb validation loss from the Modded-NanoGPT optimization benchmark, pretrained on a fixed budget of 1.71.7 B tokens with decoupled weight decay; every optimizer is tuned independently at every batch size. (a) Validation loss with respect to batch size for all five optimizers. (b) The performance of three matrix optimizers on a zoomed loss axis. SOAP is the best optimizer from 128K128\mathrm{K} to 1M1\mathrm{M} tokens, ahead by up to 0.0120.012 nats, yet at 2M2\mathrm{M} it falls to third, 0.0210.021 nats behind Shampoo. Rankings cross between other pairs of optimizers too: Muon beats Shampoo by 0.00350.0035 nats at 512K512\mathrm{K} , but Shampoo beats Muon by 0.00370.0037 nats at 2M2\mathrm{M} . The HyperBall ( Wen et al., 2026a ) counterpart is Figure 3 .
MuonWη′\eta^{\prime}μ′\mu^{\prime}λ′\lambda^{\prime}
Invariant preservation
2 *κpη\kappa^{p}\etaμ\muκ1−pλ\kappa^{1-p}\lambda
Oursκη\kappa\etaμκ\mu^{\kappa}λ\lambda
Bound minimization
51κ η\sqrt{\kappa}\,\etaμ\mu–
4κ η\sqrt{\kappa}\,\eta––
Table 1: Batch-size scaling prescriptions at B′=κBB^{\prime}=\kappa B . Entries give target values; “–” means unspecified. We extend Power Lines ( Bergsma et al., 2025 ) to Muon assuming fixed momentum. The original work favors p=0p=0 for AdamW, which we also found to be effective for MuonW in our empirical sweep. Fixed Adam moments in Power Lines and fixed β1\beta_{1} in Marek et al. (2025) are empirical choices. Bu et al. (2026) ; Schaipp et al. (2025) concern SGD on convex optimization, with empirical extensions to the optimizers shown. For a moment retention coefficient β\beta near 11 and fixed exponent aa , we use βa≈1−a(1−β)\beta^{a}\approx 1-a(1-\beta) to display moment transfers as powers; Appendix B.4 gives the details and original formulas. Empirical rows are from results in Section 3.2 .
Figure 2: No single scaling rule transfers well in both settings. We transfer a tuned configuration from a reference batch to larger batches using the rules in Table 1 : (a) MuonW–AdamW language-model pretraining on 1.71.7 B tokens, with a reference batch of 128K128\mathrm{K} tokens; (b) a ResNet with GroupNorm on CIFAR-5M, with a reference batch of 256256 images. The theoretical rules follow Definitions 3 and 2 for the matrix parameters, with AdamW counterparts chosen following prior work on SDE and bound minimization. Full retuning retunes every hyperparameter at each batch by coordinate descent. In language modeling, the empirical LLM rule is consistent with Power Lines ( Bergsma et al., 2025 ) and stays closest to full retuning, while both theoretical rules fall behind no scaling from 1M1\mathrm{M} tokens on. On CIFAR-5M, the bound-minimization rule instead stays within 0.0060.006 nats of full retuning, while the empirical LLM rule does no better than no scaling. CIFAR-5M points are medians over three seeds. Appendix E gives full results.
Figure 3: The lowest-loss optimizer changes with batch size under HyperBall. Language-model pretraining with an approximately 1.71.7 B-token budget and independent tuning of every optimizer at every batch size. Left: all five optimizers. Right: Muon, KL-SOAP (labeled SOAP), and Shampoo on a zoomed loss axis. KL-SOAP has the lowest reported loss from 128K128\mathrm{K} through 1M1\mathrm{M} , with a gap of only about 0.00030.0003 to Shampoo at 1M1\mathrm{M} ; Shampoo has the lowest loss at 2M2\mathrm{M} .
Figure 4: Which scaling rule works depends on the reference batch size and on the curvature-to-noise ratio (CNR). (a)–(c) SignSGD with momentum on a three-direction NQM with a fixed sample budget. Gray: hyperparameters retuned at each batch size. Blue and orange: hyperparameters tuned at B0B_{0} (dotted line), then scaled to each batch size using the SDE and bound-minimization rules, respectively. Excess loss is terminal loss minus the retuned loss at B=1B=1 . Shaded bands show one standard error across eight evaluation groups. (d) Learning rates of SignSGD with momentum tuned for terminal loss on a one-dimensional NQM, each relative to its own value at B=1B=1 ; we control the CNR by choosing cc and hh . Dotted and dash-dotted lines mark linear and square-root scaling. See Sections C.2.1 and C.3.3 for details.
Figure 5: Optimizer crossover on a noisy quadratic. Tuned loss on a three-direction NQM, with each optimizer retuned at every batch size; “-M” denotes momentum.
Figure 6: Holding the sharpest 0.001%0.001\% of directions at 128K128\mathrm{K} removes up to 59.5%59.5\% of the 2M2\mathrm{M} branch penalty. Top: held-out loss of the 128K128\mathrm{K} base run (black) and of the fully scaled and top- 768768 held 2M2\mathrm{M} branches started at steps 1,0001{,}000 , 3,0003{,}000 , 5,0005{,}000 , 7,0007{,}000 , and 9,0009{,}000 . Bottom: held-out loss minus the control branch’s loss at each evaluation for the fully scaled, top- 1616 held, top- 768768 held, and random- 768768 held (dotted) branches, evaluated every 3232 base steps on 11 M held-out tokens and smoothed over three evaluations.
Appendix figures & tables80 assets

Supplementary material from the paper’s appendix.

Appendix

MethodMuonWAdam / AdamW
 η′\eta^{\prime}μ′\mu^{\prime}λ′\lambda^{\prime}η′\eta^{\prime}β1′\beta_{1}^{\prime}β2′\beta_{2}^{\prime}λ′\lambda^{\prime}
Invariant preservation
Adam SDE ( Malladi et al., 2022 )–––κ η\sqrt{\kappa}\,\eta1−κ(1−β1)1-\kappa(1-\beta_{1})1−κ(1−β2)1-\kappa(1-\beta_{2})–
Completed transfer ( Mlodozeniec et al., 2026 )–––κ η\sqrt{\kappa}\,\eta1−κ(1−β1)1-\kappa(1-\beta_{1})1−κ(1−β2)1-\kappa(1-\beta_{2})κ λ\sqrt{\kappa}\,\lambda
Second-moment memory ( Marek et al., 2025 )––––β1\beta_{1}β2κ\beta_{2}^{\kappa}–
Power Lines ( Bergsma et al., 2025 ) *κpη\kappa^{p}\etaμ\muκ1−pλ\kappa^{1-p}\lambdaκpη\kappa^{p}\etaβ1\beta_{1}β2\beta_{2}κ1−pλ\kappa^{1-p}\lambda
Appendix
Table 2: Finite-change batch-size prescriptions at B′=κBB^{\prime}=\kappa B . Each cell gives the target value in terms of the reference value for its parameter group; dashes mean unspecified. Adam SDE and Completed transfer also scale ϵ′=ϵ/κ\epsilon^{\prime}=\epsilon/\sqrt{\kappa} . Our experimental SDE candidate instead uses βi′=βiκ\beta_{i}^{\prime}=\beta_{i}^{\kappa} and fixed ϵ\epsilon . ∗* : the Power Lines product constraint leaves pp free; its favored choice is p=0p=0 . Its MuonW entries are our extension with fixed momentum, while fixed Adam moments are an empirical choice adopted from the paper. Fixed β1\beta_{1} in the second-moment row is also an empirical choice adopted from the paper. †\dagger : our fixed-token-budget conversion of the heuristic in Cesista & Wen (2025) , assuming per-step update-noise variance is proportional to η2/B\eta^{2}/B with a fixed proportionality factor; the source’s empirical square-root rule is distinct from this conversion. Our Muon diffusion theorem holds ω\omega fixed independently of μ\mu ; using it for tied MuonW is an additional approximation. Our bound rule assumes an exact oracle and fixed ω\omega , before the effective step-size cap is active. Shulgin et al. (2026) use a large-horizon bound at fixed momentum for an exact oracle without Nesterov or weight decay. Section B.4 details these conditions.
Figure 7: Tuned learning rates for SGD, Newton, SignSGD, and RMSProp (left to right) on the noisy quadratic with h=(0.2,1,5)h=(0.2,1,5) , c=(10,1,0.1)c=(10,1,0.1) , θ0=(0.3,1,3)\theta_{0}=(0.3,1,3) , and T=4096T=4096 . Blue curves show per-batch tuned learning rates; orange dashed curves show linear scaling anchored at B=2B=2 . The dotted gray SGD curve is the large- TT approximation. For RMSProp, the blue curve jointly tunes (η,β2)(\eta,\beta_{2}) ; the dashed gray curve fixes β2=.9\beta_{2}=.9 , and the dotted gray curve keeps β21/B\beta_{2}^{1/B} fixed using β2(B)=.9B/2\beta_{2}(B)=.9^{B/2} . The RMSProp linear reference is anchored to the jointly tuned rate.
Figure 8: Tuned finite-budget risk versus batch size for momentum variants on the same three-coordinate noisy quadratic as Figure 7 , with T=4096T=4096 . SGD-M, Newton-M, and SignSGD-M tune both η\eta and μ\mu at each batch size; Adam tunes η\eta , β1\beta_{1} , and β2\beta_{2} . The plot uses a dense local refinement around the broad-sweep winners; the linear methods are evaluated by the exact second-moment recursion, and the SignSGD/Adam candidates are independently checked by high-chain Monte Carlo validation.
Figure 9: CNR changes RMSProp learning-rate scaling. Colors indicate CNR. Curves show learning rates tuned for terminal loss, each relative to its own value at B=1B=1 . Dotted and dash-dotted lines mark linear and square-root scaling.
Figure 10: The dependence on the curvature-to-noise ratio (CNR) survives accumulator bias correction. The two RMSProp panels divide the squared-gradient exponential moving average by 1−β2t1-\beta_{2}^{t} at update tt , with second-moment decay 0.90.9 and 0.990.99 , respectively. The SGD panel selects learning rates using the exact expected terminal loss. All panels use the same one-dimensional NQM, initial point, sample budget, noise variances, and batch sizes as Figure 9 . Learning rates are normalized separately at B=1B=1 . SGD’s relative learning-rate curves remain close to linear in batch size in these conditions, while RMSProp’s curves separate with CNR.
Figure 11: Adam shows the same CNR dependence as RMSProp. The tuning protocol, model, budget, and batch grid match Figure 9 , with Adam’s first moment enabled and bias correction applied to both moments. Colors indicate CNR; curves show learning rates tuned for terminal loss, each relative to its own value at B=1B=1 . Dotted and dash-dotted lines mark linear and square-root scaling.
Figure 12: Tuned finite-budget terminal excess risk in the 4×44\times 4 matrix prediction models. Every point uses T=4096T=4096 examples and T/BT/B updates; each optimizer is retuned separately at every batch size. Minibatch sufficient statistics are sampled exactly from the finite-support design, while the expectation over training trajectories is estimated by independent float64 validation. Shaded bands show 95%95\% Monte Carlo intervals for the exact population excess risk.
Figure 13: Higher-noise counterpart of Figure 12 . The squared-loss models use σϵ=4\sigma_{\epsilon}=4 , and the logistic teacher uses τlog=4\tau_{\mathrm{log}}=4 . The remaining settings and protocols, including the 95%95\% Monte Carlo intervals, are unchanged.
Figure 14: All 648648 CIFAR-5M scaling rules, summarized one coordinate at a time. In each panel, a point at a target batch is the arithmetic mean of the final-window validation losses over all rules that make the indicated choice for that coordinate, weighting all combinations of the other six equally. The outlined marker and dashed line mark the shared reference at B0=256B_{0}=256 . Fixed, square root, and linear denote exponents 00 , 12\tfrac{1}{2} , and 11 ; EMA preserves the half-life in examples. Lower is better; all panels share the same vertical range.
BBηM\eta_{\rm M}ηA\eta_{\rm A}μ\muβ1\beta_{1}β2\beta_{2}λM\lambda_{\rm M}λA\lambda_{\rm A}median loss
2562560.0068160.0068160.00098040.00098040.930.930.65610.65610.9999120.9999120.001250.001250.001250.001250.23310.2331
5125120.013630.013630.0027730.0027730.930.930.43050.43050.9998230.9998230.0006250.0006250.0017680.0017680.22750.2275
102410240.027260.027260.0019610.0019610.930.930.65610.65610.9999120.9999120.001250.001250.001250.001250.22990.2299
204820480.038560.038560.0078430.0078430.930.930.65610.65610.9999120.9999120.001250.001250.0035360.0035360.23180.2318
409640960.027260.027260.00098040.00098040.95050.95050.65610.65610.9985860.9985860.001250.001250.0050.0050.23370.2337
Appendix
Table 3: CIFAR-5M configurations selected by coordinate descent. Losses are medians over seeds 4242 – 4444 . At the reference batch B0=256B_{0}=256 , the hyperparameter search ties the two weight decays, λM=λA\lambda_{\rm M}=\lambda_{\rm A} . Displayed hyperparameters are rounded.
BBηM\eta_{\rm M}ηEmb.\eta_{\rm Emb.}ηOut.\eta_{\rm Out.}ηScalar\eta_{\rm Scalar}μ\muβ1\beta_{1}β2\beta_{2}λM\lambda_{\rm M}λA\lambda_{\rm A}
128K128\mathrm{K}0.01767770.01767770.4949740.4949740.002828420.002828420.01060660.01060660.950.950.80.80.9936290.9936290.0250.0250.0010.001
Appendix
Table 4: Language-model reference configuration. Batch size BB counts tokens. The auxiliary learning rates ηEmb.\eta_{\rm Emb.} , ηOut.\eta_{\rm Out.} , and ηScalar\eta_{\rm Scalar} apply to embeddings, the output projection, and scalar parameters, respectively. Learning rates are peak values, and the cooldown fraction is 0.80.8 . Displayed values are rounded.
  Parameter scalingEMA scaling
BatchLossMatrix LRMatrix WDAuxiliary LRMuon μ\muAdam β1\beta_{1}Adam β2\beta_{2}
256K3.26553Square rootFixedSquare rootFixedFixedEMA
512K3.28004Square rootSquare rootSquare rootFixedFixedEMA
1M3.31183FixedLinearFixedFixedFixedEMA
2M3.37977FixedLinearFixedFixedFixedEMA
Appendix
Table 5: Lowest validation loss observed independently at each target batch size. The choice at one batch size is not required to transfer to another batch size.
Figure 15: All 216216 language-model scaling rules, summarized one coordinate at a time. Each point is the arithmetic mean of the target-batch validation losses over all rules using the indicated choice on that coordinate, weighting all combinations of the other five choices equally. Fixed, square root, and linear denote exponents 0,12,10,\tfrac{1}{2},1 ; E/EMA preserves decay over the complete token budget. Lower is better; each panel uses its own focused vertical range. Every rule shares the reference run at B0=128KB_{0}=128\mathrm{K} , drawn as the outlined marker and the dashed horizontal line. The full grid is reported in Section E.3 .
  Matrix hyperparameter configurationAuxiliary hyperparameter configuration
BatchStepsLRWDμ\muEmb. LROut. LRScalar LRβ1\beta_{1}β2\beta_{2}
128K130000.01767770.0250.950.4949740.002828420.01060660.80.993629
256K65000.01767770.050.950.4949740.002828420.01060660.80.987299
512K32500.01767770.10.950.4949740.002828420.01060660.80.97476
1M16250.01767770.20.950.4949740.002828420.01060660.80.950156
2M8130.01767770.40.950.4949740.002828420.01060660.80.902854
Appendix
Table 6: Resolved numerical hyperparameters for the selected common rule. Values are rounded for readability. Auxiliary weight decay is 0.001 at every batch size, and the learning rate cooldown occupies the final 80 percent of training.
Only change from selected rule256K512K1M2MMean BSZ regret
Selected common rule3.269693.281923.311833.379770.00151
Matrix LR to square root3.271403.295733.350263.463580.03595
Matrix LR to linear3.283253.339743.449003.664920.12493
Matrix WD to fixed3.270303.289063.327873.402660.01318
Matrix WD to square root3.268623.283343.317003.386620.00460
Auxiliary LR to square root3.267633.281193.316453.410010.00953
Appendix
Table 7: One coordinate changes around the selected common rule. Every row changes exactly one entry and keeps the other five entries fixed. Mean BSZ regret is the equally weighted mean over the four target batch-size choices 256K,512K,1M,2M256\mathrm{K},512\mathrm{K},1\mathrm{M},2\mathrm{M} of each row’s loss gap to the best tested rule at the same batch size.
Figure 16: Mean and maximum BSZ regret for all 216216 transfer rules. At each batch size, BSZ regret is the rule’s validation loss minus the lowest validation loss among the tested rules at that same batch size. Each point represents one rule; the mean (horizontal axis, equal weights) and maximum (vertical axis) are taken over the four target batch-size choices 256K,512K,1M,2M256\mathrm{K},512\mathrm{K},1\mathrm{M},2\mathrm{M} tokens. The star has the lowest mean BSZ regret, and the cross has the lowest maximum BSZ regret. Lower is better; both axes use logarithmic scales.
RankParameter scalingEMA scalingValidation lossMean BSZ regret
 ηmat\eta_{\rm mat}λmat\lambda_{\rm mat}ηaux\eta_{\rm aux}λaux\lambda_{\rm aux}μ\muβ1\beta_{1}β2\beta_{2}64128512102420484096Δ‾\overline{\Delta}
1SFSFFFF0.233700.231360.233730.227780.232900.236830.00492
2SFFSFFE0.229700.230490.234400.233010.237720.233460.00533
3LFSSFEE0.234350.233460.227390.228920.230430.245450.00553
4SFSLFFF0.233740.232890.231350.233500.232060.236920.00561
5SFSFFFE0.232010.232760.231280.233870.231900.239860.00581
Appendix
Table 8: Complete CIFAR-5M validation losses for all 648648 scaling rules at six target batches, ordered by mean BSZ regret Δ‾\overline{\Delta} . This is the equally weighted mean over the six target batch-size choices of the loss gap to the best tested rule at the same batch size. Each loss is cross-entropy in nats averaged over the final ten evaluations of one run with seed 4242 ; lower is better. Batch sizes count examples. Rule codes are F (fixed), S (square-root), L (linear), and E (retention preservation in processed examples), with q(B)=q(256)B/256q(B)=q(256)^{B/256} . Bold marks the selected common rule. Ordering and BSZ regret use unrounded losses; displayed values have five decimal places.
RankParameter scalingEMA scalingTerminal validation lossMean BSZ regret
 ηmat\eta_{\rm mat}λmat\lambda_{\rm mat}ηaux\eta_{\rm aux}λaux\lambda_{\rm aux}μ\muβ1\beta_{1}β2\beta_{2}256K512K1M2MΔ‾\overline{\Delta}
1FLFFFFE3.269693.281923.311833.379770.00151
–FLFSFFE––3.308153.37707–
–FLFLFFE––3.308963.37751–
2FLFFFFF3.269523.282143.312723.382410.00240
–FLFSFFF––3.309563.38014–
Appendix
Table 9: Terminal validation losses for the 216 original language-model rules and their auxiliary weight-decay extensions. Each group contains fixed (F), square-root (S), and linear (L) auxiliary decay; E denotes retention preservation over the complete run. Groups retain the original ranking by mean BSZ regret, the equally weighted mean over the four target batch-size choices of the loss gap to the original grid minimum at the same batch size. Dashes mark unmeasured losses and unavailable four-batch ranks and BSZ regrets. Bold marks the common rule selected from the original grid.
Algorithm 1 WD and HyperBall matrix wrappers
Algorithm 2 Adam and Lion raw directions
Algorithm 3 Muon raw direction with the Newton–Schulz map
Algorithm 4 KL-SOAP raw direction
Algorithm 5 Distributed Shampoo raw direction
MethodTuned coordinatesCount
MuonWηM,μ,λM,ηA,β1,β2,fcd\eta_{\rm M},\mu,\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}7
KL-SOAP-WηM,β1M,β2M,βsh,λM,ηA,β1,β2,fcd\eta_{\rm M},\beta_{1}^{\rm M},\beta_{2}^{\rm M},\beta_{\rm sh},\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}9
ShampooWηM,λM,ηA,β1,β2,fcd\eta_{\rm M},\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}6
AdamWηM,β1M,β2M,λM,ηA,β1,β2,fcd\eta_{\rm M},\beta_{1}^{\rm M},\beta_{2}^{\rm M},\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}8
LionWηM,β1M,β2M,λM,ηA,β1,β2,fcd\eta_{\rm M},\beta_{1}^{\rm M},\beta_{2}^{\rm M},\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}8
Appendix
Table 10: Coordinates tuned in the WD campaign.
    Momentum / preconditioner  
BBMethodStepsηM\eta_{\rm M}μ\muβ1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}λM\lambda_{\rm M}LvalL_{\rm val}
128KMuonW130000.01770.95–––0.0253.2656
128KKL-SOAP-W130000.0045–0.9650.9740.9480.13.2536
128KShampooW130000.005–0.90.9–0.053.2650
128KAdamW130000.00075–0.940.949–0.33.3151
128KLionW130005 ⁣× ⁣10−55\!\times\!10^{-5}–0.90.99–8.63.3176
Appendix
Table 11: Final WD hidden-matrix hyperparameters and terminal validation loss.
  Auxiliary learning ratesAdam coefficients  
BBMethodηEmb.\eta_{\rm Emb.}ηOut.\eta_{\rm Out.}ηScalar\eta_{\rm Scalar}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}λA\lambda_{\rm A}
128KMuonW0.4950.002830.01060.80.9940.80.001
128KKL-SOAP-W0.150.001560.0050.4340.9950.40
128KShampooW1.20.01250.040.520.9950.60
128KAdamW0.8490.008840.02830.7170.9950.40
128KLionW0.2870.002990.009580.8590.9950.60
Appendix
Table 12: Final WD auxiliary-Adam hyperparameters.
    Momentum / preconditioner 
BBMethodStepsηM\eta_{\rm M}μ\muβ1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}LvalL_{\rm val}
128KMuonH130000.0090.95–––3.2629
128KKL-SOAP-H130000.009–0.9290.9740.9483.2585
128KShampooH130000.009–0.9220.853–3.2660
128KAdamH130000.009–0.940.949–3.3013
128KLionH130000.00636–0.9290.987–3.2948
Appendix
Table 13: Final HyperBall hidden-matrix hyperparameters and terminal validation loss.
  Auxiliary learning ratesAdam coefficients 
BBMethodηEmb.\eta_{\rm Emb.}ηOut.\eta_{\rm Out.}ηScalar\eta_{\rm Scalar}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}
128KMuonH0.30.003130.010.80.9940.4
128KKL-SOAP-H0.30.003130.010.4340.9940.4
128KShampooH1.20.01250.040.520.9950.4
128KAdamH0.4240.004420.01410.7170.9950.4
128KLionH0.30.003130.010.7170.9940.4
Appendix
Table 14: Final HyperBall auxiliary-Adam hyperparameters.
ηM\eta_{\rm M}μ\muηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0090.950.30.80.993630.43.26346
0.0045–––––3.29906
0.006364–––––3.27465
0.012728–––––3.27016
0.018–––––3.29629
–0.975––––3.27140
Appendix
Table 15: Hyperparameter ablation for MuonH at batch size B=128KB=128K .
ηM\eta_{\rm M}μ\muηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0180.950.60.858580.9750.43.27768
0.009–––––3.31502
0.012728–––––3.28879
0.025456–––––3.28412
0.036–––––3.31009
–0.975––––3.28664
Appendix
Table 16: Hyperparameter ablation for MuonH at batch size B=512KB=512K .
ηM\eta_{\rm M}μ\muηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0254560.950.60.90.951250.43.30679
0.012728–––––3.34188
0.018–––––3.31511
0.036–––––3.31483
0.050912–––––3.33885
–0.975––––3.31714
Appendix
Table 17: Hyperparameter ablation for MuonH at batch size B=1MB=1M .
ηM\eta_{\rm M}μ\muηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0360.950.60.90.934420.43.37098
0.018–––––3.40605
0.025456–––––3.38018
0.050912–––––3.38112
0.072–––––3.40408
–0.96464––––3.37684
Appendix
Table 18: Hyperparameter ablation for MuonH at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0090.929290.9740.948010.30.434310.993630.43.25728
0.0045–––––––3.29786
0.006364–––––––3.27114
0.012728–––––––3.26005
0.018–––––––3.27850
–0.96464––––––3.25890
Appendix
Table 19: Hyperparameter ablation for KL-SOAP-H at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0180.929290.858580.90.60.80.9750.83.26964
0.009–––––––3.30243
0.012728–––––––3.28109
0.025456–––––––3.27408
0.036–––––––3.29737
–0.975––––––3.29182
Appendix
Table 20: Hyperparameter ablation for KL-SOAP-H at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0254560.929290.865650.865650.60.80.951250.63.29536
0.012728–––––––3.32611
0.018–––––––3.30198
0.036–––––––3.30162
0.050912–––––––3.32237
–0.95––––––3.30005
Appendix
Table 21: Hyperparameter ablation for KL-SOAP-H at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0360.929290.65610.878410.60.80.907250.63.35062
0.018–––––––3.38592
0.025456–––––––3.36252
0.050912–––––––3.35583
0.072–––––––3.38261
–0.95––––––3.35416
Appendix
Table 22: Hyperparameter ablation for KL-SOAP-H at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0090.922220.852941.20.520.99550.43.26617
0.0045––––––3.30292
0.006364––––––3.27724
0.012728––––––3.26965
0.018––––––3.29102
–0.9725–––––3.27934
Appendix
Table 23: Hyperparameter ablation for ShampooH at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0180.929290.90.848530.80.982320.43.27578
0.009––––––3.31081
0.012728––––––3.28723
0.025456––––––3.28195
0.036––––––3.30230
–0.96464–––––3.28247
Appendix
Table 24: Hyperparameter ablation for ShampooH at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0254560.929290.8480.848530.80.970750.43.29641
0.012728––––––3.33305
0.018––––––3.30733
0.036––––––3.30217
0.050912––––––3.32518
–0.96464–––––3.30594
Appendix
Table 25: Hyperparameter ablation for ShampooH at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0360.929290.724880.60.80.944350.43.33928
0.018––––––3.37605
0.025456––––––3.34886
0.050912––––––3.34763
0.072––––––3.37374
–0.95–––––3.34432
Appendix
Table 26: Hyperparameter ablation for ShampooH at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0090.940.949030.424260.717160.99550.43.30144
0.0045––––––3.32743
0.006364––––––3.30696
0.012728––––––3.31867
0.018––––––3.37667
–0.96464–––––3.30978
Appendix
Table 27: Hyperparameter ablation for AdamH at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0180.929290.90.60.80.9750.43.33227
0.009––––––3.36880
0.012728––––––3.34160
0.025456––––––3.34476
0.036––––––3.38899
–0.96464–––––3.35050
Appendix
Table 28: Hyperparameter ablation for AdamH at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0180.964640.90250.30.80.958630.253.37114
0.006364––––––3.45697
0.009––––––3.41310
0.012728––––––3.38258
0.025456––––––3.37975
0.036––––––3.41154
Appendix
Table 29: Hyperparameter ablation for AdamH at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0254560.964640.907250.424260.80.92130.253.48723
0.012728––––––3.52999
0.018––––––3.49427
0.036––––––3.50275
0.050912––––––3.54708
–0.975–––––3.50189
Appendix
Table 30: Hyperparameter ablation for AdamH at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0063640.929290.987260.30.717160.993630.43.29369
0.003182––––––3.33450
0.0045––––––3.30646
0.009––––––3.31381
0.012728––––––3.35250
–0.96464–––––3.30461
Appendix
Table 31: Hyperparameter ablation for LionH at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0180.929290.950.424260.80.983330.43.33575
0.009––––––3.37584
0.012728––––––3.34782
0.025456––––––3.34516
0.036––––––3.38208
–0.96464–––––3.37497
Appendix
Table 32: Hyperparameter ablation for LionH at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0360.90.90250.60.90.975620.43.40475
0.018––––––3.43203
0.025456––––––3.40845
0.050912––––––3.43559
0.072––––––3.48983
–0.95–––––3.66610
Appendix
Table 33: Hyperparameter ablation for LionH at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0360.90.907250.60.90.953630.23.55651
0.018––––––3.60350
0.025456––––––3.56984
0.050912––––––3.57392
0.072––––––3.62350
–0.92929–––––3.64818
Appendix
Table 34: Hyperparameter ablation for LionH at batch size B=2MB=2M .
ηM\eta_{\rm M}μ\muλM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0176780.950.0250.494970.80.993630.83.26547
0.00625––––––3.27885
0.0088388––––––3.27196
0.0125––––––3.26909
0.025––––––3.27012
0.035355––––––3.28074
Appendix
Table 35: Hyperparameter ablation for MuonW at batch size B=128KB=128K .
ηM\eta_{\rm M}μ\muλM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0250.950.050.70.80.964640.73.27788
0.0125––––––3.28710
0.017678––––––3.28196
0.035355––––––3.28059
0.05––––––3.28835
–0.96464–––––3.27811
Appendix
Table 36: Hyperparameter ablation for MuonW at batch size B=512KB=512K .
ηM\eta_{\rm M}μ\muλM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0250.950.10.989950.80.931060.83.31108
0.0125––––––3.32240
0.017678––––––3.31407
0.035355––––––3.31274
0.05––––––3.32043
0.070711––––––3.33872
Appendix
Table 37: Hyperparameter ablation for MuonW at batch size B=1MB=1M .
ηM\eta_{\rm M}μ\muλM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0250.950.10.70.858580.907250.43.36617
0.0125––––––3.37566
0.017678––––––3.36828
0.035355––––––3.37356
0.05––––––3.38618
–0.975–––––3.37805
Appendix
Table 38: Hyperparameter ablation for MuonW at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00450.964640.9740.948010.10.150.434310.99550.43.25388
0.00225––––––––3.26253
0.003182––––––––3.25685
0.006364––––––––3.25702
0.009––––––––3.26629
–0.975–––––––3.25474
Appendix
Table 39: Hyperparameter ablation for KL-SOAP-W at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0127280.90.858580.858580.150.848530.80.9750.43.27594
0.006364––––––––3.28188
0.009––––––––3.27597
0.018––––––––3.28161
0.025456––––––––3.28798
–0.96464–––––––3.96353
Appendix
Table 40: Hyperparameter ablation for KL-SOAP-W at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0063640.929290.865650.865650.60.60.80.951250.63.29629
0.003182––––––––3.30629
0.0045––––––––3.30860
0.009––––––––3.30143
0.012728––––––––3.30874
–0.96464–––––––3.32143
Appendix
Table 41: Hyperparameter ablation for KL-SOAP-W at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0360.929290.828050.878410.40.60.80.934420.43.38357
0.018––––––––3.39916
0.025456––––––––3.41082
0.050912––––––––3.41118
0.072––––––––3.43941
–0.95–––––––3.41115
Appendix
Table 42: Hyperparameter ablation for KL-SOAP-W at batch size B=2MB=2M .
ηM\eta_{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.0050.051.20.520.99550.63.26453
0.0025–––––3.26990
0.0035355–––––3.26949
0.0070711–––––3.26794
0.01–––––3.27269
–0.025––––3.26934
Appendix
Table 43: Hyperparameter ablation for ShampooW at batch size B=128KB=128K .
ηM\eta_{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00707110.150.848530.80.982320.63.28179
0.0035355–––––3.28292
0.005–––––3.28384
0.01–––––3.28310
0.014142–––––3.28916
–0.075––––3.28211
Appendix
Table 44: Hyperparameter ablation for ShampooW at batch size B=512KB=512K .
ηM\eta_{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00707110.30.848530.80.970750.63.30873
0.0035355–––––3.31717
0.005–––––3.31121
0.01–––––3.31060
0.014142–––––3.31547
–0.2––––3.31070
Appendix
Table 45: Hyperparameter ablation for ShampooW at batch size B=1MB=1M .
ηM\eta_{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.010.50.60.80.944350.63.35964
0.005–––––3.36559
0.0070711–––––3.36361
0.014142–––––3.37268
0.02–––––3.38900
–0.4––––3.36664
Appendix
Table 46: Hyperparameter ablation for ShampooW at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.000750.940.949030.30.848530.717160.99550.43.31446
0.000375–––––––3.33480
0.00053033–––––––3.32473
0.0010607–––––––3.31801
0.0015–––––––3.32305
–0.97––––––3.32536
Appendix
Table 47: Hyperparameter ablation for AdamW at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00150.90.858580.60.848530.858580.9750.43.33876
0.00075–––––––3.36489
0.0010607–––––––3.34867
0.0021213–––––––3.35025
0.003–––––––3.39202
–0.92929––––––3.41100
Appendix
Table 48: Hyperparameter ablation for AdamW at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00150.858580.931060.80.60.858580.90250.43.40136
0.00075–––––––3.44254
0.0010607–––––––3.41807
0.0021213–––––––3.41238
0.003–––––––3.50500
–0.9––––––3.41518
Appendix
Table 49: Hyperparameter ablation for AdamW at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00212130.90.868841.83.39410.929290.967210.63.47554
0.0010607–––––––3.52875
0.0015–––––––3.49515
0.003–––––––3.49297
0.0042426–––––––3.52116
–0.95––––––3.69887
Appendix
Table 50: Hyperparameter ablation for AdamW at batch size B=2MB=2M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
5 ⁣× ⁣10−55\!\times\!10^{-5}0.90.989968.60.287280.858580.99550.63.32190
2.5 ⁣× ⁣10−52.5\!\times\!10^{-5}–––––––3.36648
3.5355 ⁣× ⁣10−53.5355\!\times\!10^{-5}–––––––3.33947
7.0711 ⁣× ⁣10−57.0711\!\times\!10^{-5}–––––––3.33920
0.0001–––––––3.37269
–0.92929––––––3.32372
Appendix
Table 51: Hyperparameter ablation for LionW at batch size B=128KB=128K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00010.950.960.050.212130.717160.950.23.39255
5 ⁣× ⁣10−55\!\times\!10^{-5}–––––––3.44386
7.0711 ⁣× ⁣10−57.0711\!\times\!10^{-5}–––––––3.41694
0.00014142–––––––3.39765
0.0002–––––––3.44060
–0.975––––––3.41967
Appendix
Table 52: Hyperparameter ablation for LionW at batch size B=512KB=512K .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.000141420.950.943710.150.30.60.90250.23.44463
7.0711 ⁣× ⁣10−57.0711\!\times\!10^{-5}–––––––3.50421
10 ⁣× ⁣10−510\!\times\!10^{-5}–––––––3.47023
0.0002–––––––3.48006
0.00028284–––––––3.55216
–0.96464––––––3.48061
Appendix
Table 53: Hyperparameter ablation for LionW at batch size B=1MB=1M .
ηM\eta_{\rm M}β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}λM\lambda_{\rm M}ηEmb.\eta_{\rm Emb.}β1\beta_{1}β2\beta_{2}fcdf_{\rm cd}LvalL_{\rm val}
0.00020.90.888550.30.30.80.953630.23.55681
0.0001–––––––3.64900
0.00014142–––––––3.59357
0.00028284–––––––3.55522
0.0004–––––––3.63287
–0.92929––––––3.64348
Appendix
Table 54: Hyperparameter ablation for LionW at batch size B=2MB=2M .
Algorithm 6 Directional branching
anchorfully scaledtop- 1616top- 6464top- 128128top- 256256top- 768768random- 768768
1,0001{,}000+68.7+68.7+58.9+58.9+55.9+55.9+53.7+53.7+42.2+42.2+41.6+41.6+70.4+70.4
2,0002{,}000+50.9+50.9+52.1+52.1+45.8+45.8+39.4+39.4+38.6+38.6+30.5+30.5+50.8+50.8
3,0003{,}000+36.1+36.1+34.3+34.3+27.6+27.6+28.6+28.6+21.9+21.9+18.3+18.3+36.0+36.0
4,0004{,}000+40.3+40.3+29.3+29.3+27.3+27.3+25.5+25.5+23.2+23.2+19.0+19.0+39.7+39.7
5,0005{,}000+28.8+28.8+21.2+21.2+18.4+18.4+17.0+17.0+14.8+14.8+11.7+11.7+28.6+28.6
6,0006{,}000+25.1+25.1+21.9+21.9+16.7+16.7+14.1+14.1+12.5+12.5+10.2+10.2+24.6+24.6
Appendix
Table 55: Branch penalty at 2M2\mathrm{M} for every anchor and held subspace , in units of 10−310^{-3} nats: the validation loss at the end of a 1,0241{,}024 -step branch minus that of the 128K128\mathrm{K} control branch from the same anchor. A dash marks a branch that was not run.
Figure 17: Branch-penalty trajectories at every anchor. Held-out loss of each 2M2\mathrm{M} branch minus that of the 128K128\mathrm{K} control branch from the same anchor, evaluated every 3232 base steps on 11 M held-out tokens and smoothed over three evaluations, for the fully scaled, top- rr held ( r=16,64,128,256,768r=16,64,128,256,768 ), and random- 768768 held (dotted) branches. Each panel has its own vertical range. The endpoints in Table 55 are measured on 10.510.5 M validation tokens instead. Top- 128128 was not run at anchors 7,0007{,}000 , 9,0009{,}000 , and 11,00011{,}000 .
 Matrix hyperparameter configuration 
OptimizerηM\eta_{\rm M}λM\lambda_{\rm M}μ\mu / β1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}Aux. LR mult.
MuonW0.020.04730040.95––1.41421
KL-SOAP-W0.01272790.050.950.90.91.41421
ShampooW0.0050.20.90.9–2
AdamW0.001783810.60.90.95–1
LionW0.000336359120.90.99–1.41421
Appendix
Table 56: Final anchor hyperparameter configurations at a global batch of 256K tokens and eight tokens per scaling parameter. The first-moment column uses μ\mu for MuonW and β1M\beta_{1}^{\rm M} for the other optimizers. The coefficient βsh\beta_{\rm sh} applies only to KL-SOAP-W; ShampooW uses β2M\beta_{2}^{\rm M} for its preconditioner. The auxiliary learning rate multiplier acts on both ηEmb.\eta_{\rm Emb.} and ηOut.\eta_{\rm Out.} .
OptimizerCompleted pointsInitial BPBFinal BPBImprovement
MuonW471.07301.07110.0019
KL-SOAP-W231.07151.07120.0003
ShampooW401.07771.07500.0027
AdamW251.10941.10800.0014
LionW431.15101.13690.0141
Appendix
Table 57: Anchor search observations. Initial BPB is the selected result after the first matrix learning rate and weight decay grid. Final BPB is the selected result after the subsequent coordinate searches and boundary extensions.
OptimizerTransfer rule512K BPB1M BPB2M BPB
KL-SOAP-WAuxiliary-only EMA1.08121.10351.1614
 Invariant-preserving1.08671.15941.3735
ShampooWAuxiliary-only EMA1.08471.10391.1525
 Invariant-preserving1.08441.10441.6558
Appendix
Table 58: Coefficient transfer across batch sizes. Auxiliary-only EMA transfers auxiliary Adam β2\beta_{2} and keeps all other coefficients fixed; it minimizes mean BSZ regret across the three target batches for each optimizer. The invariant-preserving rule also transfers matrix β2M\beta_{2}^{\rm M} and KL-SOAP-W’s preconditioner coefficient βsh\beta_{\rm sh} , preserving their decay in token units. Both rules keep β1M\beta_{1}^{\rm M} and β1\beta_{1} fixed. EMA transfer means q(B)=q0κq(B)=q_{0}^{\kappa} , where κ=B/B0\kappa=B/B_{0} , with B0=256KB_{0}=256\mathrm{K} tokens.
OptimizerIrreducible BPBAmplitudeExponentRMSE
MuonW0.9837140.0251450.6723830.002624
KL-SOAP-W0.9923480.0185550.8050570.001106
ShampooW1.0116850.0111080.9076070.006164
AdamW0.9878900.0300500.7788450.005971
LionW0.9187500.0951580.4928650.013589
Appendix
Table 59: Parameters of the fitted compute loss curves.
Figure 18: Validation losses across token budgets when the batch isze is 256K and the fitted loss–compute curve for each optimizer. Each curve is fit only to that optimizer’s seven observations.
Figure 19: Compute efficiency penalties at a fixed horizon of eight tokens per scaling parameter. The left panel fits and uses a separate reference curve for every optimizer. The right panel maps every loss through the MuonW curve and therefore retains absolute loss differences. Lower is better.
   Compute-efficiency penalty
BBOptimizerLvalL_{\rm val} (BPB)Own referenceMuon reference
256KMuonW1.07120.00000.0000
 KL-SOAP-W1.07220.00000.0699
 ShampooW1.07530.00000.1163
 AdamW1.10760.00000.4361
 LionW1.13650.00000.5872
Appendix
Table 60: Compute efficiency penalties at the eight token per parameter horizon. Own curve penalty uses the compute loss curve fitted separately for each optimizer. Muon curve penalty maps every observed loss through the single MuonW curve.
Optimizerτ/N\tau/NBB (tokens)S(B)S(B)LvalL_{\rm val} (BPB)
MuonW1262,1442241.3264
 2262,1444481.1932
 4262,1448961.1197
 8262,1441,7921.0712
 12262,1442,6881.0492
 16262,1443,5841.0366
Appendix
Table 61: Complete observations across token budgets and batch sizes used in the compute efficiency analysis.
Rankβ1M\beta_{1}^{\rm M}β2M\beta_{2}^{\rm M}βsh\beta_{\rm sh}β1\beta_{1}β2\beta_{2}512K1M2MMean BSZ regret
KL-SOAP-W
1FFFFE1.08121.10351.16140.00123
2FFFFF1.07871.10331.16500.00153
3FEFFF1.08211.10421.16590.00327
4FEFFE1.08291.10451.16600.00367
5FFFEE1.07931.10691.20310.01563
Appendix
Table 62: Complete beta-transfer rule rankings. F keeps a coefficient fixed, and E applies EMA transfer. For ShampooW, the preconditioner shares the matrix β2M\beta_{2}^{\rm M} coefficient, so no separate preconditioner choice is shown. Mean BSZ regret averages the gap to the best tested rule for the same optimizer and batch size over the three target batches.

Explore similar work

CardsList

  1. Navigating LLM Valley: From AdamW to Memory-Efficient and Matrix-Based Optimizers

    May 9, 2026Aditya RanganathDeep Learning OptimizationMemory-Efficient Optimization

    Center for Applied Scientific Computing Lawrence Livermore National Laboratory Livermore, CA, USA

  2. Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization

    Jun 15, 2026Kaiyue Wen, Xingyu Dang, Kaifeng Lyu +2Language Model PretrainingWeight Decay

    †Stanford University. · ‡Princeton University. · §Tsinghua University.

  3. Towards joint scaling laws with optimal batch size schedules

    Jul 30, 2026Jiaxiang Li, Zhiqi Bu, Shiyun XuDeep Learning OptimizationLanguage Model Scaling Laws

    Meta · Independent researcher