The Best Optimizer Depends on Batch Size

Xingyu Dang; Kaiyue Wen; Sadhika Malladi

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

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.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

128\mathrm{K}

1\mathrm{M}

tokens, ahead by up to

0.012

nats, yet at

2\mathrm{M}

it falls to third,

0.021

nats behind Shampoo. Rankings cross between other pairs of optimizers too: Muon beats Shampoo by

0.0035

nats at

512\mathrm{K}

, but Shampoo beats Muon by

0.0037

nats at

2\mathrm{M}

. The HyperBall ( Wen et al., 2026a ) counterpart is Figure 3 .

Invariant preservation
MuonW	$\eta^{\prime}$	$\mu^{\prime}$	$\lambda^{\prime}$
2 *	$\kappa^{p}\eta$	$\mu$	$\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^{\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=0

for AdamW, which we also found to be effective for MuonW in our empirical sweep. Fixed Adam moments in Power Lines and fixed

\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

1

and fixed exponent

a

, we use

\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.7

B tokens, with a reference batch of

128\mathrm{K}

tokens; (b) a ResNet with GroupNorm on CIFAR-5M, with a reference batch of

256

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

1\mathrm{M}

tokens on. On CIFAR-5M, the bound-minimization rule instead stays within

0.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.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

128\mathrm{K}

through

1\mathrm{M}

, with a gap of only about

0.0003

to Shampoo at

1\mathrm{M}

; Shampoo has the lowest loss at

2\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

B_{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=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=1

; we control the CNR by choosing

c

and

h

. 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\%

of directions at

128\mathrm{K}

removes up to

59.5\%

of the

2\mathrm{M}

branch penalty. Top: held-out loss of the

128\mathrm{K}

base run (black) and of the fully scaled and top-

768

held

2\mathrm{M}

branches started at steps

1{,}000

3{,}000

5{,}000

7{,}000

, and

9{,}000

. Bottom: held-out loss minus the control branch’s loss at each evaluation for the fully scaled, top-

16

held, top-

768

held, and random-

768

held (dotted) branches, evaluated every

32

base steps on

1

M held-out tokens and smoothed over three evaluations.

Appendix figures & tables80 assets

Supplementary material from the paper’s appendix.

Appendix

Method	MuonW			Adam / AdamW
	$\eta^{\prime}$	$\mu^{\prime}$	$\lambda^{\prime}$	$\eta^{\prime}$	$\beta_{1}^{\prime}$	$\beta_{2}^{\prime}$	$\lambda^{\prime}$
Invariant preservation
Adam SDE ( Malladi et al., 2022 )	–	–	–	$\sqrt{\kappa}\,\eta$	$1-\kappa(1-\beta_{1})$	$1-\kappa(1-\beta_{2})$	–
Completed transfer ( Mlodozeniec et al., 2026 )	–	–	–	$\sqrt{\kappa}\,\eta$	$1-\kappa(1-\beta_{1})$	$1-\kappa(1-\beta_{2})$	$\sqrt{\kappa}\,\lambda$
Second-moment memory ( Marek et al., 2025 )	–	–	–	–	$\beta_{1}$	$\beta_{2}^{\kappa}$	–
Power Lines ( Bergsma et al., 2025 ) *	$\kappa^{p}\eta$	$\mu$	$\kappa^{1-p}\lambda$	$\kappa^{p}\eta$	$\beta_{1}$	$\beta_{2}$	$\kappa^{1-p}\lambda$

Appendix

Table 2: Finite-change batch-size prescriptions at

B^{\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

\beta_{i}^{\prime}=\beta_{i}^{\kappa}

and fixed

\epsilon

*

: the Power Lines product constraint leaves

p

free; its favored choice is

p=0

. Its MuonW entries are our extension with fixed momentum, while fixed Adam moments are an empirical choice adopted from the paper. Fixed

\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

\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)

c=(10,1,0.1)

\theta_{0}=(0.3,1,3)

, and

T=4096

. Blue curves show per-batch tuned learning rates; orange dashed curves show linear scaling anchored at

B=2

. The dotted gray SGD curve is the large-

T

approximation. For RMSProp, the blue curve jointly tunes

(\eta,\beta_{2})

; the dashed gray curve fixes

\beta_{2}=.9

, and the dotted gray curve keeps

\beta_{2}^{1/B}

fixed using

\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=4096

. SGD-M, Newton-M, and SignSGD-M tune both

\eta

and

\mu

at each batch size; Adam tunes

\eta

\beta_{1}

, and

\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=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-\beta_{2}^{t}

at update

t

, with second-moment decay

0.9

and

0.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=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=1

. Dotted and dash-dotted lines mark linear and square-root scaling.

Figure 12: Tuned finite-budget terminal excess risk in the

4\times 4

matrix prediction models. Every point uses

T=4096

examples and

T/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\%

Monte Carlo intervals for the exact population excess risk.

Figure 13: Higher-noise counterpart of Figure 12 . The squared-loss models use

\sigma_{\epsilon}=4

, and the logistic teacher uses

\tau_{\mathrm{log}}=4

. The remaining settings and protocols, including the

95\%

Monte Carlo intervals, are unchanged.

Figure 14: All

648

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

B_{0}=256

. Fixed, square root, and linear denote exponents

0

\tfrac{1}{2}

, and

1

; EMA preserves the half-life in examples. Lower is better; all panels share the same vertical range.

$B$	$\eta_{\rm M}$	$\eta_{\rm A}$	$\mu$	$\beta_{1}$	$\beta_{2}$	$\lambda_{\rm M}$	$\lambda_{\rm A}$	median loss
$256$	$0.006816$	$0.0009804$	$0.93$	$0.6561$	$0.999912$	$0.00125$	$0.00125$	$0.2331$
$512$	$0.01363$	$0.002773$	$0.93$	$0.4305$	$0.999823$	$0.000625$	$0.001768$	$0.2275$
$1024$	$0.02726$	$0.001961$	$0.93$	$0.6561$	$0.999912$	$0.00125$	$0.00125$	$0.2299$
$2048$	$0.03856$	$0.007843$	$0.93$	$0.6561$	$0.999912$	$0.00125$	$0.003536$	$0.2318$
$4096$	$0.02726$	$0.0009804$	$0.9505$	$0.6561$	$0.998586$	$0.00125$	$0.005$	$0.2337$

Appendix

Table 3: CIFAR-5M configurations selected by coordinate descent. Losses are medians over seeds

42

–

44

. At the reference batch

B_{0}=256

, the hyperparameter search ties the two weight decays,

\lambda_{\rm M}=\lambda_{\rm A}

. Displayed hyperparameters are rounded.

$B$	$\eta_{\rm M}$	$\eta_{\rm Emb.}$	$\eta_{\rm Out.}$	$\eta_{\rm Scalar}$	$\mu$	$\beta_{1}$	$\beta_{2}$	$\lambda_{\rm M}$	$\lambda_{\rm A}$
$128\mathrm{K}$	$0.0176777$	$0.494974$	$0.00282842$	$0.0106066$	$0.95$	$0.8$	$0.993629$	$0.025$	$0.001$

Appendix

Table 4: Language-model reference configuration. Batch size

B

counts tokens. The auxiliary learning rates

\eta_{\rm Emb.}

\eta_{\rm Out.}

, and

\eta_{\rm Scalar}

apply to embeddings, the output projection, and scalar parameters, respectively. Learning rates are peak values, and the cooldown fraction is

0.8

. Displayed values are rounded.

		Parameter scaling			EMA scaling
Batch	Loss	Matrix LR	Matrix WD	Auxiliary LR	Muon $\mu$	Adam $\beta_{1}$	Adam $\beta_{2}$
256K	3.26553	Square root	Fixed	Square root	Fixed	Fixed	EMA
512K	3.28004	Square root	Square root	Square root	Fixed	Fixed	EMA
1M	3.31183	Fixed	Linear	Fixed	Fixed	Fixed	EMA
2M	3.37977	Fixed	Linear	Fixed	Fixed	Fixed	EMA

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

216

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,\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

B_{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 configuration			Auxiliary hyperparameter configuration
Batch	Steps	LR	WD	$\mu$	Emb. LR	Out. LR	Scalar LR	$\beta_{1}$	$\beta_{2}$
128K	13000	0.0176777	0.025	0.95	0.494974	0.00282842	0.0106066	0.8	0.993629
256K	6500	0.0176777	0.05	0.95	0.494974	0.00282842	0.0106066	0.8	0.987299
512K	3250	0.0176777	0.1	0.95	0.494974	0.00282842	0.0106066	0.8	0.97476
1M	1625	0.0176777	0.2	0.95	0.494974	0.00282842	0.0106066	0.8	0.950156
2M	813	0.0176777	0.4	0.95	0.494974	0.00282842	0.0106066	0.8	0.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 rule	256K	512K	1M	2M	Mean BSZ regret
Selected common rule	3.26969	3.28192	3.31183	3.37977	0.00151
Matrix LR to square root	3.27140	3.29573	3.35026	3.46358	0.03595
Matrix LR to linear	3.28325	3.33974	3.44900	3.66492	0.12493
Matrix WD to fixed	3.27030	3.28906	3.32787	3.40266	0.01318
Matrix WD to square root	3.26862	3.28334	3.31700	3.38662	0.00460
Auxiliary LR to square root	3.26763	3.28119	3.31645	3.41001	0.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

256\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

216

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

256\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.

Rank	Parameter scaling				EMA scaling			Validation loss						Mean BSZ regret
	$\eta_{\rm mat}$	$\lambda_{\rm mat}$	$\eta_{\rm aux}$	$\lambda_{\rm aux}$	$\mu$	$\beta_{1}$	$\beta_{2}$	64	128	512	1024	2048	4096	$\overline{\Delta}$
1	S	F	S	F	F	F	F	0.23370	0.23136	0.23373	0.22778	0.23290	0.23683	0.00492
2	S	F	F	S	F	F	E	0.22970	0.23049	0.23440	0.23301	0.23772	0.23346	0.00533
3	L	F	S	S	F	E	E	0.23435	0.23346	0.22739	0.22892	0.23043	0.24545	0.00553
4	S	F	S	L	F	F	F	0.23374	0.23289	0.23135	0.23350	0.23206	0.23692	0.00561
5	S	F	S	F	F	F	E	0.23201	0.23276	0.23128	0.23387	0.23190	0.23986	0.00581

Appendix

Table 8: Complete CIFAR-5M validation losses for all

648

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

42

; 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/256}

. Bold marks the selected common rule. Ordering and BSZ regret use unrounded losses; displayed values have five decimal places.

Rank	Parameter scaling				EMA scaling			Terminal validation loss				Mean BSZ regret
	$\eta_{\rm mat}$	$\lambda_{\rm mat}$	$\eta_{\rm aux}$	$\lambda_{\rm aux}$	$\mu$	$\beta_{1}$	$\beta_{2}$	256K	512K	1M	2M	$\overline{\Delta}$
1	F	L	F	F	F	F	E	3.26969	3.28192	3.31183	3.37977	0.00151
–	F	L	F	S	F	F	E	–	–	3.30815	3.37707	–
–	F	L	F	L	F	F	E	–	–	3.30896	3.37751	–
2	F	L	F	F	F	F	F	3.26952	3.28214	3.31272	3.38241	0.00240
–	F	L	F	S	F	F	F	–	–	3.30956	3.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.

Panel of Figure 26 — Algorithm 1 WD and HyperBall matrix wrappers

Panel of Figure 27 — Algorithm 2 Adam and Lion raw directions

Panel of Figure 28 — Algorithm 3 Muon raw direction with the Newton–Schulz map

Panel of Figure 29 — Algorithm 4 KL-SOAP raw direction

Panel of Figure 30 — Algorithm 5 Distributed Shampoo raw direction

Method	Tuned coordinates	Count
MuonW	$\eta_{\rm M},\mu,\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}$	7
KL-SOAP-W	$\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	$\eta_{\rm M},\lambda_{\rm M},\eta_{\rm A},\beta_{1},\beta_{2},f_{\rm cd}$	6
AdamW	$\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	$\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.

$B$	Method	Steps	$\eta_{\rm M}$	$\mu$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\lambda_{\rm M}$	$L_{\rm val}$
				Momentum / preconditioner
128K	MuonW	13000	0.0177	0.95	–	–	–	0.025	3.2656
128K	KL-SOAP-W	13000	0.0045	–	0.965	0.974	0.948	0.1	3.2536
128K	ShampooW	13000	0.005	–	0.9	0.9	–	0.05	3.2650
128K	AdamW	13000	0.00075	–	0.94	0.949	–	0.3	3.3151
128K	LionW	13000	$5\!\times\!10^{-5}$	–	0.9	0.99	–	8.6	3.3176

Appendix

Table 11: Final WD hidden-matrix hyperparameters and terminal validation loss.

$B$	Method	$\eta_{\rm Emb.}$	$\eta_{\rm Out.}$	$\eta_{\rm Scalar}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$\lambda_{\rm A}$
		Auxiliary learning rates			Adam coefficients
128K	MuonW	0.495	0.00283	0.0106	0.8	0.994	0.8	0.001
128K	KL-SOAP-W	0.15	0.00156	0.005	0.434	0.995	0.4	0
128K	ShampooW	1.2	0.0125	0.04	0.52	0.995	0.6	0
128K	AdamW	0.849	0.00884	0.0283	0.717	0.995	0.4	0
128K	LionW	0.287	0.00299	0.00958	0.859	0.995	0.6	0

Appendix

Table 12: Final WD auxiliary-Adam hyperparameters.

$B$	Method	Steps	$\eta_{\rm M}$	$\mu$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$L_{\rm val}$
				Momentum / preconditioner
128K	MuonH	13000	0.009	0.95	–	–	–	3.2629
128K	KL-SOAP-H	13000	0.009	–	0.929	0.974	0.948	3.2585
128K	ShampooH	13000	0.009	–	0.922	0.853	–	3.2660
128K	AdamH	13000	0.009	–	0.94	0.949	–	3.3013
128K	LionH	13000	0.00636	–	0.929	0.987	–	3.2948

Appendix

Table 13: Final HyperBall hidden-matrix hyperparameters and terminal validation loss.

$B$	Method	$\eta_{\rm Emb.}$	$\eta_{\rm Out.}$	$\eta_{\rm Scalar}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$
		Auxiliary learning rates			Adam coefficients
128K	MuonH	0.3	0.00313	0.01	0.8	0.994	0.4
128K	KL-SOAP-H	0.3	0.00313	0.01	0.434	0.994	0.4
128K	ShampooH	1.2	0.0125	0.04	0.52	0.995	0.4
128K	AdamH	0.424	0.00442	0.0141	0.717	0.995	0.4
128K	LionH	0.3	0.00313	0.01	0.717	0.994	0.4

Appendix

Table 14: Final HyperBall auxiliary-Adam hyperparameters.

$\eta_{\rm M}$	$\mu$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.009	0.95	0.3	0.8	0.99363	0.4	3.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=128K

$\eta_{\rm M}$	$\mu$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.018	0.95	0.6	0.85858	0.975	0.4	3.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=512K

$\eta_{\rm M}$	$\mu$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025456	0.95	0.6	0.9	0.95125	0.4	3.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=1M

$\eta_{\rm M}$	$\mu$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.036	0.95	0.6	0.9	0.93442	0.4	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.009	0.92929	0.974	0.94801	0.3	0.43431	0.99363	0.4	3.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.018	0.92929	0.85858	0.9	0.6	0.8	0.975	0.8	3.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025456	0.92929	0.86565	0.86565	0.6	0.8	0.95125	0.6	3.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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.036	0.92929	0.6561	0.87841	0.6	0.8	0.90725	0.6	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.009	0.92222	0.85294	1.2	0.52	0.9955	0.4	3.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.018	0.92929	0.9	0.84853	0.8	0.98232	0.4	3.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025456	0.92929	0.848	0.84853	0.8	0.97075	0.4	3.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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.036	0.92929	0.72488	0.6	0.8	0.94435	0.4	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.009	0.94	0.94903	0.42426	0.71716	0.9955	0.4	3.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.018	0.92929	0.9	0.6	0.8	0.975	0.4	3.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.018	0.96464	0.9025	0.3	0.8	0.95863	0.25	3.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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025456	0.96464	0.90725	0.42426	0.8	0.9213	0.25	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.006364	0.92929	0.98726	0.3	0.71716	0.99363	0.4	3.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.018	0.92929	0.95	0.42426	0.8	0.98333	0.4	3.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.036	0.9	0.9025	0.6	0.9	0.97562	0.4	3.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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.036	0.9	0.90725	0.6	0.9	0.95363	0.2	3.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=2M

$\eta_{\rm M}$	$\mu$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.017678	0.95	0.025	0.49497	0.8	0.99363	0.8	3.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=128K

$\eta_{\rm M}$	$\mu$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025	0.95	0.05	0.7	0.8	0.96464	0.7	3.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=512K

$\eta_{\rm M}$	$\mu$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025	0.95	0.1	0.98995	0.8	0.93106	0.8	3.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=1M

$\eta_{\rm M}$	$\mu$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.025	0.95	0.1	0.7	0.85858	0.90725	0.4	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0045	0.96464	0.974	0.94801	0.1	0.15	0.43431	0.9955	0.4	3.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.012728	0.9	0.85858	0.85858	0.15	0.84853	0.8	0.975	0.4	3.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.006364	0.92929	0.86565	0.86565	0.6	0.6	0.8	0.95125	0.6	3.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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.036	0.92929	0.82805	0.87841	0.4	0.6	0.8	0.93442	0.4	3.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=2M

$\eta_{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.005	0.05	1.2	0.52	0.9955	0.6	3.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=128K

$\eta_{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0070711	0.15	0.84853	0.8	0.98232	0.6	3.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=512K

$\eta_{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0070711	0.3	0.84853	0.8	0.97075	0.6	3.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=1M

$\eta_{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.01	0.5	0.6	0.8	0.94435	0.6	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.00075	0.94	0.94903	0.3	0.84853	0.71716	0.9955	0.4	3.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0015	0.9	0.85858	0.6	0.84853	0.85858	0.975	0.4	3.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0015	0.85858	0.93106	0.8	0.6	0.85858	0.9025	0.4	3.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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0021213	0.9	0.86884	1.8	3.3941	0.92929	0.96721	0.6	3.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=2M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
$5\!\times\!10^{-5}$	0.9	0.98996	8.6	0.28728	0.85858	0.9955	0.6	3.32190
$2.5\!\times\!10^{-5}$	–	–	–	–	–	–	–	3.36648
$3.5355\!\times\!10^{-5}$	–	–	–	–	–	–	–	3.33947
$7.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=128K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0001	0.95	0.96	0.05	0.21213	0.71716	0.95	0.2	3.39255
$5\!\times\!10^{-5}$	–	–	–	–	–	–	–	3.44386
$7.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=512K

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.00014142	0.95	0.94371	0.15	0.3	0.6	0.9025	0.2	3.44463
$7.0711\!\times\!10^{-5}$	–	–	–	–	–	–	–	3.50421
$10\!\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=1M

$\eta_{\rm M}$	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\lambda_{\rm M}$	$\eta_{\rm Emb.}$	$\beta_{1}$	$\beta_{2}$	$f_{\rm cd}$	$L_{\rm val}$
0.0002	0.9	0.88855	0.3	0.3	0.8	0.95363	0.2	3.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=2M

Panel of Figure 76 — Algorithm 6 Directional branching

anchor	fully scaled	top- $16$	top- $64$	top- $128$	top- $256$	top- $768$	random- $768$
$1{,}000$	$+68.7$	$+58.9$	$+55.9$	$+53.7$	$+42.2$	$+41.6$	$+70.4$
$2{,}000$	$+50.9$	$+52.1$	$+45.8$	$+39.4$	$+38.6$	$+30.5$	$+50.8$
$3{,}000$	$+36.1$	$+34.3$	$+27.6$	$+28.6$	$+21.9$	$+18.3$	$+36.0$
$4{,}000$	$+40.3$	$+29.3$	$+27.3$	$+25.5$	$+23.2$	$+19.0$	$+39.7$
$5{,}000$	$+28.8$	$+21.2$	$+18.4$	$+17.0$	$+14.8$	$+11.7$	$+28.6$
$6{,}000$	$+25.1$	$+21.9$	$+16.7$	$+14.1$	$+12.5$	$+10.2$	$+24.6$

Appendix

Table 55: Branch penalty at

2\mathrm{M}

for every anchor and held subspace , in units of

10^{-3}

nats: the validation loss at the end of a

1{,}024

-step branch minus that of the

128\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

2\mathrm{M}

branch minus that of the

128\mathrm{K}

control branch from the same anchor, evaluated every

32

base steps on

1

M held-out tokens and smoothed over three evaluations, for the fully scaled, top-

r

held (

r=16,64,128,256,768

), and random-

768

held (dotted) branches. Each panel has its own vertical range. The endpoints in Table 55 are measured on

10.5

M validation tokens instead. Top-

128

was not run at anchors

7{,}000

9{,}000

, and

11{,}000

	Matrix hyperparameter configuration
Optimizer	$\eta_{\rm M}$	$\lambda_{\rm M}$	$\mu$ / $\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	Aux. LR mult.
MuonW	0.02	0.0473004	0.95	–	–	1.41421
KL-SOAP-W	0.0127279	0.05	0.95	0.9	0.9	1.41421
ShampooW	0.005	0.2	0.9	0.9	–	2
AdamW	0.00178381	0.6	0.9	0.95	–	1
LionW	0.000336359	12	0.9	0.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

\beta_{1}^{\rm M}

for the other optimizers. The coefficient

\beta_{\rm sh}

applies only to KL-SOAP-W; ShampooW uses

\beta_{2}^{\rm M}

for its preconditioner. The auxiliary learning rate multiplier acts on both

\eta_{\rm Emb.}

and

\eta_{\rm Out.}

Optimizer	Completed points	Initial BPB	Final BPB	Improvement
MuonW	47	1.0730	1.0711	0.0019
KL-SOAP-W	23	1.0715	1.0712	0.0003
ShampooW	40	1.0777	1.0750	0.0027
AdamW	25	1.1094	1.1080	0.0014
LionW	43	1.1510	1.1369	0.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.

Optimizer	Transfer rule	512K BPB	1M BPB	2M BPB
KL-SOAP-W	Auxiliary-only EMA	1.0812	1.1035	1.1614
	Invariant-preserving	1.0867	1.1594	1.3735
ShampooW	Auxiliary-only EMA	1.0847	1.1039	1.1525
	Invariant-preserving	1.0844	1.1044	1.6558

Appendix

Table 58: Coefficient transfer across batch sizes. Auxiliary-only EMA transfers auxiliary Adam

\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

\beta_{2}^{\rm M}

and KL-SOAP-W’s preconditioner coefficient

\beta_{\rm sh}

, preserving their decay in token units. Both rules keep

\beta_{1}^{\rm M}

and

\beta_{1}

fixed. EMA transfer means

q(B)=q_{0}^{\kappa}

, where

\kappa=B/B_{0}

, with

B_{0}=256\mathrm{K}

tokens.

Optimizer	Irreducible BPB	Amplitude	Exponent	RMSE
MuonW	0.983714	0.025145	0.672383	0.002624
KL-SOAP-W	0.992348	0.018555	0.805057	0.001106
ShampooW	1.011685	0.011108	0.907607	0.006164
AdamW	0.987890	0.030050	0.778845	0.005971
LionW	0.918750	0.095158	0.492865	0.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
$B$	Optimizer	$L_{\rm val}$ (BPB)	Own reference	Muon reference
256K	MuonW	1.0712	0.0000	0.0000
	KL-SOAP-W	1.0722	0.0000	0.0699
	ShampooW	1.0753	0.0000	0.1163
	AdamW	1.1076	0.0000	0.4361
	LionW	1.1365	0.0000	0.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	$\tau/N$	$B$ (tokens)	$S(B)$	$L_{\rm val}$ (BPB)
MuonW	1	262,144	224	1.3264
	2	262,144	448	1.1932
	4	262,144	896	1.1197
	8	262,144	1,792	1.0712
	12	262,144	2,688	1.0492
	16	262,144	3,584	1.0366

Appendix

Table 61: Complete observations across token budgets and batch sizes used in the compute efficiency analysis.

Rank	$\beta_{1}^{\rm M}$	$\beta_{2}^{\rm M}$	$\beta_{\rm sh}$	$\beta_{1}$	$\beta_{2}$	512K	1M	2M	Mean BSZ regret
KL-SOAP-W
1	F	F	F	F	E	1.0812	1.1035	1.1614	0.00123
2	F	F	F	F	F	1.0787	1.1033	1.1650	0.00153
3	F	E	F	F	F	1.0821	1.1042	1.1659	0.00327
4	F	E	F	F	E	1.0829	1.1045	1.1660	0.00367
5	F	F	F	E	E	1.0793	1.1069	1.2031	0.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

\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

May 9, 2026cs.LG

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

Training large language models requires optimization algorithms that are not only statistically effective, but also computationally and memory efficient at extreme scale. Although Adam remains the dominant optimizer for large-scale language-model pretraining and fine-tuning, recent work has revisited nearly every component of the optimization stack: adaptive moment estimation, decoupled weight decay, memory footprint, curvature approximation, sign-based updates, large-batch stability, low-rank gradient structure, and matrix-wise orthogonalized updates. This survey reviews optimizer design for large language models through a systems-and-optimization lens. We organize the literature into classical first-order optimizers, adaptive optimizers, memory-efficient variants, second-order and curvature-aware methods, sign-based and discovered optimizers, low-rank and projection-based methods, and matrix-based optimizers such as Muon. We also discuss benchmarking methodology, including hyperparameter fairness, scale dependence, wall-clock efficiency, token efficiency, memory overhead, and downstream evaluation. We argue that optimizer research for LLMs is entering a new phase: moving from single-algorithm speedup claims toward rigorous, scale-aware comparisons that jointly evaluate convergence, stability, memory, and implementation complexity.

Aditya Ranganath

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

Deep Learning Optimization Memory-Efficient Optimization Matrix Optimization LLM Training

Ask PDF

Jun 15, 2026cs.LG

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

Matrix based optimizers such as Muon can substantially speed up language model pretraining, but their gains over AdamW are observed to shrink as model size and data scale grow when using standard constant decoupled weight decay. We propose Hyperball, a simple optimizer wrapper that addresses this issue. Given a base optimizer such as Adam or Muon, Hyperball sets the Frobenius norms of weight matrices and their corresponding optimizer updates to fixed constants. On Qwen3 style models up to 1.2B parameters, Muon Hyperball achieves 20--30% token equivalent speedup over weight decay baselines. Hyperball also improves learning rate transfer across widths and depths compared to decoupled weight decay. This method is motivated by prior theory showing that training with weight decay leads to an equilibrium weight norm that only depends on the training hyperparameters. Through this mechanism, the weight decay then decides the angular learning rate, i.e. how fast the direction of the weight matrix changes.

Kaiyue Wen, Xingyu Dang, Kaifeng Lyu +2

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

Language Model Pretraining Weight Decay Neural Network Optimization Hyperparameter Transfer Constrained Optimization

Ask PDF

Jul 30, 2026cs.LG

Towards joint scaling laws with optimal batch size schedules

Modern deep learning typically keeps the batch size static throughout training, thus overlooking the joint effect of learning rate and batch size on the training dynamics. In this paper, we study the deep learning dynamics through the lens of convex optimization and derive a joint characterization of loss in terms of both schedules, applicable to general optimizers and model architectures. This characterization yields a closed-form optimal batch size schedule for any prescribed learning rate schedule, and further leads to joint scaling laws that consistently outperform static batch size baselines, highlighting the significance of dynamic batch size schedule in large language model training.

Jiaxiang Li, Zhiqi Bu, Shiyun Xu

Meta · Independent researcher

Deep Learning Optimization Language Model Scaling Laws Neural Network Training Dynamics LLM Training

Ask PDF