Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales
Authors: Shikai Qiu
Zixi Chen
Hoang Phan
Qi Lei
Andrew Gordon Wilson
Venue: NeurIPS 2025 poster ()
Keywords: hyperparameter transfer
second order optimizers
scaling laws
Primary area: optimization

Abstract: Several recently introduced deep learning optimizers utilizing matrix-level preconditioning have shown promising speedups relative to the current dominant optimizer AdamW, particularly in relatively small-scale experiments. However, efforts to validate and replicate their successes have reported mixed results.
To better understand the effectiveness of these optimizers at scale, in this work we investigate how to scale preconditioned optimizers via hyperparameter transfer, building on prior works such as $\mu$P. We study how the optimal learning rate and weight decay should scale with model width and depth for a wide range of optimizers, including Shampoo, SOAP, and Muon, accounting for the impact of commonly used techniques such as blocking and grafting. We find that scaling the learning rate according to $\mu$P improves transfer, but can still suffer from significant finite-width deviations that cause drifting optimal learning rates, which we show can be mitigated by blocking and explicit spectral normalization.
For compute-optimal scaling, we find scaling independent weight decay as $1/\mathrm{width}$ is nearly optimal across optimizers.
Applying these scaling rules, we show Muon, SOAP and Shampoo consistently achieve near $1.4\times$ speedup over AdamW for training Llama-architecture language models of sizes ranging from $190$M to $1.4$B, whereas the speedup vanishes rapidly with scale under incorrect scaling. Based on these results and further ablations, we argue that studying optimal hyperparameter transfer is essential for reliably comparing optimizers at scale given a realistic tuning budget.

Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales
Shikai Qiu"
Zixi Chen'
Hoang Phan
Qi Lei Andrew Gordon Wilson
New York University
Abstract
Several recently introduced deep learning optimizers utilizing matrix-level precon- ditioning have shown promising speedups relative the current dominant optimizer AdamW, particularly in relatively small-scale experiments. However, efforts to validate and replicate their successes have reported mixed results. To better under- stand the effectiveness of these optimizers at scale, in this work we investigate how scale preconditioned optimizers via hyperparameter transfer, building on prior works such as j.P. We study how the optimal learning rate and weight decay should scale with model width and depth for a wide range of optimizers, including Sham- poo, SOAP, and Muon, accounting for the impact of commonly used techniques such as blocking and grafting. We find that scaling the learning rate according to uP improves transfer, but can still suffer from significant finite-widhh deviations that cause drifting optimal learning rates, which we show can be mitigated by blocking and explicit spectral normalization. For compute-optimit scaling, find scaling independent weight decay as 1/ 1/width is nearly optimal across optimizers. Applying these scaling rules, show Muon, SOAP and Shampoo consistently achieve near 1.4x speedup over Adam for training Jlama-architecture language models of sizes ranging from 190M 1.4B, whereas the speedup vanishes rapidly with scale under incorrect scaling. Based on these results and further ablations, we argue that studying optimal hyperparameter transfer is essential for reliably comparing optimizers at scale given a realistic tuning budget.
Introduction
Developing better optimization techniques can fundamentally accelerate deep learning. While AdamW [31] continues dominate deep learning workloads, including training frontier-scale language models [16] 5T 29]. several recently introduced optimizers have demonstrated significant speedups relative AdamW, including Shampoo 19 25 3 SOAP 47 and Muon 23 These optimizers use matrix rather than elementwise preconditioners that lead to faster learning per step while incurring mild overhead. Notably, Muon's recent success in training rillioe-parammeer language models [45] points to the potential in replacing AdamW as the default optimizer at seale.
Unfortunately, subsequent works have reported mixed results in replicating these claimed speedups, especially on larger-scale experiments. For example, while Liu et al. (30) demonstrates consistent 2x speedup of Muon relative AdamW from 400M to 1B-aarrmeter language models, Wen et al. [49 show that Muon at most 1.4x more efficient than well-tunee AdamW, W, and this speedup quickly diminishes to 1.1x for 1.2B models, with SOAP following similar trend. Similarly, Semenov et al. (39 shows the speedups of matrix-preconditinn
optimizers over AdamW decrease with scale. While the different experiment setups make it difficult fully reconcile these studies, we hypothesize that the lack a robust and consistent way of choosing hyperparameters, particularly for the larger models where careful tuning is impractical, is likely the main reason for these inconsistent findings.
* Equal contribution. Correspondence to sq21298nyu. edu, zc21570nyu. edu, andreugw@cims .wy...nn.ueeemmm edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Muon-Adam Adam#SSampoo
Shampoo
SQAP
190M
340M
640M
1.48
1.4
3.2
LO55
1.2
SP
3.0
1.5
1.4
Muon
Shampoo
1.3
2.6
Adam
uP
2
LR Mūltíplier
LR'MuL
22
10"
10"
Compute (FLOPs)
4.0
SP
190M
380M
640M
1.48
Compute multiplier
1.4
2.0
Lo5s
2
1.2
μP
1.0
3.0
2.0
21
0.8
-0-
1P
1/Widh WD
27
104 10a
104
10P
10a 10
10a 104 10a
Training Steps
Compute (FLOPs)
Figure Hyperparameter transfer crucial for achieving good performance with matrix-preconditoone optimizers across scales. (Lefft) We derive yP scaling for stabilizing the optimal learning rate across widths. Fixing the base learning rate, jP yields consistent training dynamics across widths for transformers Open- WebText and that wider always better, whereas the standard parameterization (SP) leads instability and unstable optima (Section 3}. Muon-Adam uses Muon hidden layers Adam embedding and readout layers. Adam#Shampoo is Shampoo with Adam-grafting. Shampoo and SOAP use block size of 128. (Right) Combining jP with 1/ //width independent weight decay, Muon, SOAP and Shampoo achieve consistent speedups around 1.4x well-tuned AdamW in training 190M to (.4B-paramete models on Fen. By contrast, SP rapidly deteriorates the performance all three optimizers they scale (Section 4..
Indeed, key lesson from scaling up deep learning over recent years is that hyperparameters such as initialization, learning rate, weight decay, and training horizon must be carefully co-varied as Furthermore,
the compute budget increases achieve efficient scaling (2 24 52
the small loss-vs cos-vs-compute scaling exponent in language modeling, estimated to be as low as 0.05 [24 30), means that a 2% change in the loss can translate to 40% change in the estimated compute making good hyperparameters critical for estimating compute efficiencies between approaches.
To accurately quantify and fully realize the advantage of more advanced optimizers scale, is essential to understand how to optimally choose their hyperparameters when tuning becomes infeasible. In this work, we study scaling rules for optimal hyperparameters for a range of matrix- preconditioned optimizers, and demonstrate their critical impact on how performance scales with compute. We summarize our main findings as follows:
1. We study hyperparameter transfer for matrix preconditioned optimizers over both model width and depth. With a simple and general procedure, we show how derive the Maximal Update Parameterization (/aP) 52 for a wide range of optimizers, including Shampoo, SOAP, and Muon, accounting for the impact of commonly used techniques such as blocking and grafting 3 43 and extend to depth scaling for residual networks, e.g., the transformer.
2. We show yP significantly improves learning rate transfer compared to unscaled learning rate (Figure left). However, the optimal learning rate be far from converged at realistic widths with certain matrix-precondiiioneee normalization (Figure 2}. Fortunately, demonstrate such fiith deviations effectively mitigated by blocked preconditioning and spectral normalization (26 (Figure 3}
I
particularly when used with RMS- -MS-based update
3. For compute-optimal training, we find 1) μP can achieve near-optimal transfer, 2) spectral normalization reduces learning rate sensitivity, 3) scaling the independent weight decay as 1 /width as suggested in Xiao [50) is optimal across optimizers, and the compute- optimal model size is larger for matrix-preconie matrix-preeonditioned optimizers compared to Adam.
4. We show good hyperparameter transfer is crucial for realizing the benefit matrix- preconditioning scale. Using j&P and ((1/D) weight decay, Muon, SOAP and Shampoo
If L C dC/C - d.L/C.

---

consistently achieve near 1.4x speedups over Adam W in training 190M to parameeee language models, and the speedups quickly vanish with incorrect scaling (Figure[rrgght)
We make our code and wandb experiment logs available here
2 Background
Matrix- Pradiix-Preoondiioonnedd Optimizers. While most deep learning optimizers apply either no pre- conditioning (SGD) or only elementwise preconditioning (Adam, Adafactor [41], Lion [33, etc.), matrix-preconditioned optimizers apply matrix-level preconditioning to the gradient of 2D parameter matrices. Let W IR=x denote the weight matrix of given layer and G E [R= its gradient. The update rules typically take the form of AW & -Pr -P:- GP for positive definite matrices Pt, PR that depend on the past iterates. A representative example Shampoo 19
where
AW -1 (T. + f) G(R <n
(I)
where 7) is the learning rate, EI,R are regularization parameters, er. and En are positive exponents, and L. R are exponential moving averages of GG and G G respectively over the trajectory. The Shampoo preconditioner can be interpreted as a one-step factorder-factored approximation of the empirical Fisher using power iteration [35]. The original Shampoo uses ez ER = 1/4 [19] while later works found Eg = ER 1/2 to perform better [35] 43). A key weakness of the original Shampoo is that the matrix inversion scales cubically with the width of layer and can be unstable. To partitions the parameter into sub-blocks, reducing the peak memory and compute overhead of the preconditioners. Grafting B 43 1 normalizes the
make Shampoo more practical, blocking 43 3
Frobenius norm of the Shampoo update by that of another optimizer, such as Adam, improving the stability. Besides Shampoo, K-FAC [18] and Preconditioned Stochastic Gradient Descent [28] also leverage structured-matric approximations to the (empirical) Fisher as preconditioners.
Taking inspiration from Shampoo, several new optimizers have been recently introduced. Bernstein and Newhouse 77 showed that without the exponential moving average, Shampoo simplifies to
AW VV where G = UEVT is the SVD of the gradient matrix. The Muon optimizer [23) leverages this insight, using an iterative method to approximate UV directly, and can be viewed alternatively as steepest descent in the spectral norm. Muon has gained traction due to its algorithmic simplicity and robust performance. SOAP 47|, I. another stampoo-inspired algorithm, recognizes that Shampoo effectively applies Adafactor in the preconditioner's 's eigenbasis and builds on this insight by applying Adam in this eigenbasis instead. AdaMuon 44 applies Adam on top of Muon's 's orthogonalized gradient. We provide detailed update rules in Appendix A
Infinit- Limits and Hyperparameter Transfer. The Maximal- Update Parameterization (μP) [52] pioneered series of work on hyperparameter (HP) transfer: how HPs such as learning rate, initialization, and weight decay should scale as a function of width to allow stable and maximal feature learning in every layer the width (e.g., the transformer's embedding dimension) approaches infinity. Empirically, under jP,, HPs tuned on a small model can be transferred to models that are orders of magnitude larger without re-tuning [54]. Extending this idea to depth scaling, subsequent works have proposed to multiply the residual block output in transformers or other residual networks by 1/L" and scale the learning rate accordingly to ensure the magnitude of feature learning is stable across depth, where L is the depth and a E [0.5, 1] [57 9 10 14]. Depth- μP [57] follows a = 0.5 in order maximize diversity of feature across the layers, while CompleteP [14] uses a 1 to ensure nontrivial feature learning in each layer, demonstrating better results in training deep transformers.
54 53
Prior works primarily explored hyperparameter transfer in the most widely used optimizers like SGD, Adam, and Adafactor [17]. while analogous prescriptions for matrix matrix-precondd preconditioned methods remain underexplored, with the exception of Ishikawa and Karakida 21 who derived μP for vanilla Shampoo (without blocking or grafting) and K-FAC. Several recent large-scale investigations of Muon have resorted to heuristics for scaling its learning rate with width, such as directly using jPP for Adam 440 or matching the update RMS of Adam [30], both are incorrect as we show in Appendix
Finite-Width Scaling via the Spectral Norm. Yang et al. [56] showed that one can recover juP's initialization and learning rate scaling by constraining the spectral norm of the initialization and update of each layer to O( /dout/ of appealing to -intth asymptotics, Large et al. (26] showed that explicitly normalizing the updates to have precisely / /dout/ dmn spectral norm (before being multiplied by the learning rate)
corresponding to O(1) RMS-RMS operator norm. Instead
3

---

provides worst-case stability guarantees for finite-width networks and empirically enables good learning rate transfer for Adam and SGD with a modest overhead. As Muon already normalizes the update spectral norm, its learning rate scales as @ /dutt as identified in 36 2 Scion [36] generalizes the choice of the operator norm and demonstrates good hyperparameter transfer on transformer models.
Table 1: Scaling rules for learning rate as function of the weight matrix size (dmn, dout), depth L, block size (brn, bout) and number of blocks Yea if using blocking. A multiplier 1/L: assumed to apply to the output of every residual block. R are positive exponents for Shampoo, and 0-1 indicators for preconditioning either side for SOAP. For parameters outside of the residual blocks (c.g. embedding and last layer), set L = 1. n2. denotes the correctly scaled learning rate for optimizer Q1. We provide scaling of e parameters in Appendix
Shampoo (eL, Cx)
SOAP (CL, eR)
Muon
AdaMuon
Grafting #1#Q2
(Bout/dn) (ex(+ee)) L 222 -1 0. teR
im e2./2 gea/2
dout dna
1Q1
dm
d
3 Hyperparameter Transfer for Prtrix-Precondittoondd Optimizers
In this section, we present our scaling rules for the per-layer learning rate (and regularization e when applicable) as a function of width and depth ensure consistent feature learning as the model scales. Table Summarizee the results for optimizers we experiment with, though generalization is straightforward. Bias parameters have d 1. We provide a sketch of the derivations with high-level intuitions here and present the full derivations in Appendix B Q1 #Q2 denotes grafting where we take the update direction of Q2 but scaled to the Frobenius norm of Q1. For vector v E F. we write & (g(d)) or v g(d) to mean || (g)()), where ||и||lss = VE, /d. In Section 3.4
we study co-scaling model size and training steps and investigate weight decay scaling.
3.1 Width Scaling via pP
We start with μP for scaling learning rate and other optimizer hyperparameters with width. As we will show, a simple and general procedure suffices for a wide range of optimizers, each requiring only few lines of derivation. For simplicity, we consider training I--layer (I. assumed fixed for now) MLP that outputs f(E) on an input as follows:
=o({) = c, he(E) = W, (E), (E) = d(hr(E)), L= I-
((E) = hz = wI
where weights are drawn from a Gaussian (/, of) with layerwise variances of and is an element- wise nonlinear function. We assume
R, f(E) R. he(E) E IRd for é 1, 1,..., L 1, and refer to d as the width. As shown in prior works, the conclusions derived here will hold for more general architectures consisting of linear layers and element-wise non-linearities, modulo a minor modification for the 1 //dhead instead of 1/ dhead scaling for attention logits [54] 53].
As μP cares only about the lurge-width asymptotics with training steps and batch size held at @(1), analyzing the dynamics up to the first gradient step with a batch size of one turns out to be sufficient for obtaining the correct scaling (22 53 (56., while leading to great simplifications (e.g. gradient is rank 1 and momentum can be ignored). To carry out the analysis, we track the change in layer outputs generic input ç' induced by the first gradient update training point &. We denote the change in any quantity X after this step as AX. We discuss generalization to batch size greater than one in Appendices B.9 and B.10
We first make the observa- tion that the initialization scale σe of We is independent of the optimizer (e.g.. identical to es- tablished scalings for SGD and Adam), with a2 = ((1 /dm,.e for E < T. and a = O /d). This result follows from combining the stability and feature learning conditions of μP, which O(1) and Ah, @(1), respectively [52)
Initialization and Gradient Scales are Optimizer-Indeeeedeer
state that he(g') = WxIe=1(E") We leave detailed justifications for why this holds beyond SGD and Adam to Appendix [ but the gist is that nontrivial feature learning in the last layer inevitably constrains σr to be small in for f to not blow up with d. As a direct consequence. entries of the gradient of the loss w.n.t the prtivvtivas de(E") V scale as ((1/dout,) for all 1 I f < L, since 1) for L, ôL(E') dC/df(E') = @(1) since f(E') 0(1); 2) for f L -
@(1), f(E')
4

---

52-1(5') - 6 = O(o)) = (11d)); and 3) backpropagating through the hidden layers preserve the scale of the gradient given that o2 = can now state the much simplified jPP condition for general optimizers.
@ /daut,e) for 1 < t < L. We
Simplified μP Condition for General Optimizers. To ensure each layer is subsequently updated as much possible without diverging, μP requires that produces a ( update its output for every gradient step. Let e denote the update to W due to training on & E R, this condition requires
AW @(), 8= I,..., L.
The gradient of We computed on datapoint IR is Ge(E) = Vw, )e( where 5e(g) E Reui. and (5) R.m.. To simplify notation, we will omit the layer subscripts and abbreviate x<{) as z, xx(') as x', Ó(g) 6, and G(E) as G. Let Q(G) be the update returned by the optimizer, we therefore wish to choose per-layer learning rate 1) such that n)(())z' ( Finally, we will use the fact that æ x' behaves as sum of dn i..d. correlated random variables (correlation due being produced from the same weights) and thus scaling as ((din). ) standard result on the feature kernels of wide neural networks [38 27 22 52 8) We will refer to this property as the alignment between and x', following [56
Putting it together, the μP condition reduces to choosing 7) and other hyperparameters in Q such that
nQ(G)='= = 0(1, where = òr 8 = (LL/dott) x = 0(1), x TT (ddm))
(3)
Note x ===== (dda) and 8 = (1/dou)) are implied. Given these scaling relations, solving for the learning rate (and other hyperparameters) boils down to fairly straightforward algebra. We now illustrate with a few illustrative examples and delegate full derivations for the results Table to Ap-
pendix B In Appendix B.8 we also show that the spectral norm condition ny] Q 12 = @ dout/di. /dm) 556 remains valid alternative characterization of μP for matrix- nntrix-prreoondiiionee optimizers.
Example: Shampoo. The rank-1 preconditioners are L = 8 (x TT)TT R = x (5 8) T with unique
nonzero eigenvalues d Aa (x (8 8) = O di
The preconditioned update
duat)
Q(G) = (L+ G (R+ (R+
(4)
For EL to have ((1) effects, they must balance the only scale here, namely CL,R AL,R O O(din/ /oot)) Since G = 8x aligns with the corresponding eigenvectors,
1-er-en dout
Q(G) x' = (Az (Ar+ Ez) e (AR + ER) -eR
CR 5 (x
dia
(5)
hence 1) ~ (dout/ / /dm ( These results reproduce those in Ishikawa and Karakida 21 defined relative AL,R; as done in 3 CL. should be ((1) instead.
As a Tensor Program. Observe that Equation (5) expresses the Shampoo update purely in terms of vectors with i.i.d. entries, their inner products, and scalar nonlinearities, which shows that it can be expressed as a Tensor Program [33. Thus, the full training process admits a well-defined, deterministic infinite- -nfitite-witth limit by the Master Theorem [53] J. much stronger statement than features
updates being (1) and a more fundamental reason why we should expect the learning rate to transfer. We show analogous construction for batch size greater than one in Appendix BB.0
Example: Shampoo with Blocking. Blocking partitions the gradient into blocks Gij = 6,æ] R x Xbo, where 5; E R and Tú E R are the i-th and j-th chunks of ô and æ, each preconditioned independently. Within each block, the rank-1 1 preconditioners now have unique nonzero eigenvalue Aij = (x] (0 de)
O (iinbuy dous
which sets the scale of EI,R- Thus -(exten)
((G Xr (entten) &, (oTx²))
Bun Inverting this quantity and further
Tout
dividing number of blocks along the input dimension Rin
dm/bin to account for the correlated after some algebra, and bm bout 1, Shampoo degenerates to
contributions of (Q(Gry > we get 17
(duit) din
1-ee- -RR
(out) - ((z++e))
where nont dout/ Bout. Note when et. + en = 2 Adam and we recover the Adam scaling 1/dmn
5

---

Muon
Muon Adam
AdaMuon
Shampoo
#Shampoo
#Shampoo"
SOAP
Loss
SP
1.5
Loss
27
LR Multiplier
RMS(Ah) Loss
101
109
Stop 10
Step 2000
2' 20
222222
21
29
2002
202 2
28 2
222
2 2 2:02 252 22
2 2
2222202 222
Width D
Figure 2: μP leads to better but imperfect learning rate transfer for matrix mntrix-preconditionedd optimizers. (Top) The optimal learning rate is more consistent across widths under JaP for transformers trained on WeenWeeText. We show the learning rate as the multiplier T/aae//r,, /ro, where ooo is the optimal learning rate for the base model for each optimizer. (Bottom) yP achieves lower loss in zero-shot transferring the optimal learning rate found in the hase model (1) 128) larger models (up to 1 4096) and passes the "coordinate check": RMS of the one-step feature update in the last layer is invariant width in early training (step 10), except for SOAP. We explain the imperfect transfer of μP and why it fails for SOAP in Section[3.2]. stands for Adam-grafting, Muon-Adam uses Adam for the embedding and readout, and Shampoo" uses er. = CR 1/2.
Example: Shampoo with Adam Grafting. While an explicit calculation in the style above suf- fices, the spectral norm condition (56) provides an easy alternative to reason about grafting, which normalizes the Shampoo update's Frobenius norm by that of Adam's. Since the preconditioner leaves Shampoo's s update ((1)-rank, its Frobenius norm scales identically with its spectral norm, relation that holds also for Adam due to Adam update being low stable rank (56) [33. Consequently,
||Q AShampaaamp (G) ~ [[Q (G))m directly using Adam. More generally, Q1 #Q2 should scale the learning rate as if directly using Q1 long the stable rank of both optimizers have ratio ((1).
with Adam grafting the learning rate should scale as
3.2 Empirical Validation of μP and Improving Finite-Width Transfer
We train transformers on the Web'ee dataset for 100M tokens. We use a small vocabulary of size 96 so it is practical to apply the full preconditioners on the embedding and readout layers in order to verify whether learning rate is scaled correctly for those layers where only one of the dimensions grow. We compare μP, where the per-layer learning rate is parameterized by Tjr(D) (D) = Thase D Mo/Draune Drase) - with V/ase the base learning rate, Dasse the base width, and Og derived from Table for each layer and optimizer, against the Standard Parameterization (SP where n7 (D) = 17base For example, when using Muon, aę 0 for hidden layers, 1/2 for embedding, and 1/2 for readout. We set base 128, the width at which SP matches μP. See further experiment details in Appendix D
IPP Leads to Better Transfer and Consistent Early Dynamics. In Figure we find that jP leads to more stable learning rate landscape (loss vs 77base) than SP for all optimizers as D scales from 128 to 4096. Though the optimal learning rate is not always exactly stable, as expected from experiment noise and finite-width effects, juP consistently outperforms SP when zero-shot transferring the optimal learning rate found on the base model (Figure 2 3rd row). In early training (step 10), the RMS of the one-step feature update invariant width ,P as desired (with one exception for SOAP which we will soon discuss), passing the "coordinate check" (55). but can quickly diverge in SP (Figure 2 4th row). Furthermore, Figure (bottom left) shows for fixed 1Jaae in j.P, the loss curves are consistent across widths and wider is always better, revealing convergence towards the infinite-wide maximum update limit. That is, the network output for any fixed input & at any fixed time step t converges to a deterministic value [52 g 46).
We zero-initiiiaaii the readout layer for both μP and SP, making LR and scaling their only difference.

---

Sharpoo
Shamppo
Shampao'
SOAP
AdaMuon
#Shampoo
#Shampoo*
SOAP
1.
Loss
SP
15
Losss
13
AP
Norm
2'
2'
2
Multiplier
Multiplier
RMS(AN) Loos
Vpm 1.3
121
10h4
2222
13
Wiath D
Width D
Figure 3: Blocking and explicit normalization reduce finite-width deviations and improve transfer. (Left)
With fixed block size of 128, pP consistently achieves good learning rate transfer. including for Shampoo with grafting and SOAP which otherwise have unstable optimal learning rates (Figure 2 (Right) Explicit spectral normalization (Norm) improves learning rate transfer where μP alone fails (no blocking used here).
RMS Update Normalization Leads to Inconsistent Late Dynamies. Figure 2 the optimal learning rate consistently increases with width in juP for AdaMuon, Shampoo and Shampoo with grafting, accompanied by a decreasing trend in the feature updates in late training (step 1000). We attribute these phenomena to similar update normalizations performed by these optimizers, where the RMS of the optimizer update U normalized (via Adam-grafting or elementwise scaling) to
(1) after applying matrix preconditioners. Thus U has a spectral norm θ ( /iidddot///ran (U)
where srank denotes the stable rank. At step t with batch size B, the stable rank of the update is bounded by min( D, tB) Whereas μP takes D) alone infinity, reducing this quantity @(1), for realistic finite widths and large t, B (B is typically millions of tokens for language models) it is D that bottlenecks the stable rank, especially with matrix preconditioners designed to inflate the small singular values in U. As result, with RMS- based normalization, at late times jPP can undershoot the spectral norm of the update by / /U)r >> 1, which we show transitions from ((1) to growing with 1 as training progresses for matrix peerrii-preoonitionnee optimizers (Figure 7
uP Can Fail to Model Expressive Preconditioners Even at Initialization. Figure 2 shows μP fails the coordinate check for SOAP even at the start of training. In Appendix E.2 we trace this failure again to the infinite-width limit's inability to capture realistic finite-width behavior of expressive matrix preconditioners. Specifically, the grgembasis-projecttn gradient is sparse in the infinite-widte limit but dense for realistic finite widths and large batch sizes. See detailed findings in Appendix E.2 [.2
Blocking and Spectral Normalization Mitigate Sinite-Width Deviations. Fortunately, find that both above failure modes of juP at finite-width transfer can be effectively mitigated by two simple approaches: blocking and explicit spectral normalization. When using fixed block size, typically done to reduce the Shampoo optimizer overhead [4] 3) the preconditioner only operates on fixed-sized blocks, leaving less room for non-trivial finite-D scaling, which only scales the number of independently preconditioned blocks. Figure 3 (left) shows that by applying blocking we restore good learning rate transfer jPP for SOAP, Shampoo and Shampoo2 with Adam-grafting and observe milder decrease in feature update size in late training.
Spectral normalization. as proposed in 26 normalizes the update spectral norm to exactly (rather than asymptotically) /dout/dan strictly stronger constraint than jP (56) that can accommodate precisely the kind of winthh deviations we have observed. Figure 3 (right) shows spectral normalization transfers learning rates equally well or better than μP and achieves markedly more consistent feature updates in late training. Implemented with online power iteration, this normalization requires only two matrix-vector multiplies per layer per step (26). tiny overhead for matrix-preconditioned optimizers. We provide experiment and implementation details in Appendix F
'oomentum results in at most linear scaling with
7

---

Muon
Adam
AdaMuon
Shampoo
#Shampoo
Shampoo'
SOAP
Loss
(SP)
Loss
2
2 2 2
2
LR Multiplier
RMS(An) Loss
100
30
Step
10.0. 9 6
9.
A J 8
2 800
Depth
Figure 4: Depthwise learning rate transfer is effective for all tested optimizers. We apply a 1/L residual branch multiplier (a = 1) and adjust the learning rate to ensure (1)) feature learning in each layer, following Dey et al. [14,, outperforming SP (a = 0) and stahilizing the size of carly-time (step 10) feature update Ah when transferring from 3 to 192-layer transformers on Open WebTet. We provide experiment details in Appendix D.2
3.3 Depthwise Hyperparameter Transfer
We now consider depth scaling for architectures with residual blocks (e.g. the transformer). For the MLP example, each block computes h = W1 Ig = xz-1 + ((h)). ). In general, the residual block can be more complex, such involving layer normalization and multiple matrix multiplies as in the case of a transformer. Recent work has shown that by scaling down the output of each residual block by ((1/L"), i.e. Ty -t-1 -= (he), with
[1/2, 1]. and adjusting the learning rate to ensure stable feature learning leads to well-defined depth scaling limits [57 9 [0]. Dey et al. [4] shows choosing a and ensuring ((1) feature learning within each block leads to the best learning rate transfer for transformers trained with Adam [T0 14, referred to as the CompleteP. We derive the depth scaling rules for matrix-press prenodiierconn optimizers based on the same criterion.
With our setup so far, the only change to the previous derivation in the width-only scaling case is that gradients in the residual blocks now scale Ldgar rather than doet Propagating these additional L-dependent factors to the solution of 77 and leads to the final width-depth joint scaling rules in 4 verifies the effectiveness of our depth scaling rules on models with I. Thase 3 to L 192 layers. Compared to width scaling, the optimal learning rates shift less for SP (a 0), but our scaling with a = achieves consistently lower loss and more stable feature update size when transferring optimal learning rate found on the base model.
Table [
See derivations in Appendix B Figure
Summary: For width scaling, μP outperforms SP at learning rate transfer for all optimizers, though its ffinite-width asymptotics can fail to model realistic mnite-width scaling of matrix preconditioners, leading to shifting optima. Blocking and spectral normalization mitigate these finite-width deviations. For depth scaling, generalizing CompleteP, which scales residual branches 1/L and adjusts LR to ensure ((1) feature learning per layer, is effective for all optimizers.
3.4 Hyperparameter Transfer Under Compute-Optimal Scaling
We use the FineWeb dataset tokenized with the GPT-2 tokenizer, and train each model for 20 tokens per parameter. We use fixed block size of 512 and detail the experiment setup in Appendix D.3
We now investigate how to scale the hyperparameters in the compute-optimal setup 20 24].
μP Approximately Stabilizes the Optimal Learning Rate. Despite violating juPs assumption of (1) training steps, we find juP still improves the stability of the learning rate landscape and the optimal learning rate (Figure
in contrast to SP where the optima shift significantly to the left. While the optimal learning rates aren't 't exactly invariant, transferring the optimal learning rate found on the base = 256 model leads near-optimal performance with yP up to D 2048. Spectral normalization leads to slightly better performance than μP and reduces learning rate sensitivity.
8

---

Norm
Muon
Shampoo
SOAP
4.
3.5
Loss
10x1
10
10-
10t
10-
Shampoo
Loss
4.0
SOAP
215
1020
2300
Abase Do
Wiath D
Width D
Figure 5: Learning rate and weight decay transfer on FineWeb under compute-optimal training. (Left) juP approximately stabilizes the optimal learning rate, while spectral normalization reduces learning rate sensitivity and achieves slightly better performance. Best indicates taking the minimum over all learning rates and parameterizations. (Right) Optimal independent weight decay scales like 1/D. Muon uses Adam in the embedding layer. Shampoo and SOAP block size of 512. For Shampoo, we Adam-grafting and Adam the embedding and readout, which we found to perform better than applying one-sided Shampoo.
While our results align with findings in [14 regime, the larger-scale, higher-resolutinn sweeps done in Everett et al. [17] using Adam found that μP can considerably overshoot the learning rate. Therefore, while we observe clear advantage of uP over SP, we caution that more careful scaling studies may be needed to determine how the optimal learning rate scales in the compute-optimal regime if width varies by more than factor of 10, such as by fitting a power-law correction on top of j.P.
on the effectiveness of jP in the compute-optimal
Optimal Independent Weight Decay Scales as 1/D. While uP prescribes a O(1) independent weight decay [53], Xiao [50] found X = D) empirically leads to better scaling for Adam W. In Figure 5 (right), we find that À ((1/D) (bottom) is near-optimal for all three optimizers. At this scale, the models are still relatively small that parameters and tokens scale close to linearly with
D. making À = ((1/D) D) similar to keeping AdamW's EMA timescale constant 48 5]. a finer comparison between the two approaches to future work.
We leave
Summary: μP leads to good LR transfer under compute-optimal scaling, while the optimal independent weight decay scales as 1/D. Thus, constant LR-coupled weight decay happens to work well for AdamW. W. but can be highly suboptimal for other optimizers (c.g., Muon).
Compute-Efficiency Gains of Matrix-Preconditionns Optimizers
We now evaluate the compute-efficient gains of Muon, SOAP and Shampoo over AdamW in training up to 1.4B -ar-maraetter language models and quantify the impact of good hyperparameter transfer.
49]. We use the Llama architecture, matching width and depth configurations in Wen et al. 49 covering four model sizes ranging from All models are trained on randomly
Experiment Setup. We follow a similar setup to that in Wen et al.
190M to 1.4B. The model specifications are provided in Table 3 shuffled Web tokenized using the GPT-2 tokenizer. We use context length of 1024 and batch size of 128 sequences. We extensively tune hyperparameters for each optimizer the base model, which has 32 layers, embedding dimension 512, and 190M parameters. We detail the optimizer definition and hyperparameter tuning in Appendix H As our batch size (0. 12M tokens) is small compared to typical LLM training, we expect our findings to serve as a lower bound on matrix- pettii--prccoditionnee
optimizers' speedup, considering they scale better with large batch sizes [49 58 32). We use a block size of 512 for SOAP and Shampoo to achieve good learning rate transfer under μPP
Muon, SOAP and Shampoo Achieve Consistent Speedup. Figure 6 (left) shows the scaling trend for the three optimizers, following either our proposed scaling with μP and 1/D-scaled independent weight decay or constant learning rate and weight decay (SP). We quantify the speedup of each approach with its compute multiplier, which measures the compute- -umpute--fficiency gain over Adam W (with μP and scaled weight decay) for achieving the same loss (details in Appendix H.2 Following Liu et al. (00, we count the FL.OPs the forward and backward passes, not the optimizer transform, which prior works [30 47 argue can incur minimal runtime overhead with proper implementation.

---

190M
390M
640M
1.4B
190M
380M
640M
1.4B
190M
380M
640M
1.4B
3.2
Compute mulpter
1.4
Compute mulippe
multiplrr
1.4
3.0
1.2
12
2.8
SOAP
1,0
1,0
1/Wath LRRWD
2.6
Snampoo
0.8
-O-
1/Withh WD
0.8
-
pP only
1/Wath WD ony
10
10"
10
10"
Compute (FLOPs)
Compute (FLOPs)
Compute (FLOPs)
Figure 6: With μP and 1/ D-scaled weight decay, Muon, SOAP and Shampoo achieve consistent speedups over AdamW across model sizes up to 1.4B. Incorrect scaling rules (SP or ablating either μP or weight decay scaling) degrade performance significantly with scalc. Compute multiplier is measured against Adam W with μP and scaled weight decay, which we show yields near-optimal hyperparameters for large models (Appendix H..4).
Figure 6 (right) shows Muon, SOAP and Shampoo achieve a stable speedup factor over Adam W of around I.4x, using our proposed scaling rule. By contrast, SP leads rapidly diminishing speedup.
Hyperparameter and Scaling Rule Ablations. In Appendix H. we verify that our scaling rule leads to near-optimal learning rate and weight decay for the larger models for all four optimizers by showing that perturbing them by factors of 2 either direction leads to equal or worse performance. For Muon, we further ablate μP or weight decay scaling one at a time, shown in Figure (right) and Appendix H.4] and find that both contribute significantly to achieving good performance scale. Scaling Muon's learning rate and weight decay both as 1/ /witth, an effective scaling for AdamW [50], improves over SP but still falls short of juP 1/D weight decay, decreasing the speedup to below 1.1x for the 1.4B model. These results indicate that both components of our scaling rule are necessary to achieve the best performance and consistent speedups over Adam W across scales.
Optimal Tokens Per Parameter Depends on the Optimizer. Optimizers that achieve faster convergence are likely to have smaller compute-optimal tokens-per per-paammeerr (TPP) ratio, since diminishing returns from training on more data kick in faster (if loss at convergence is unchanged). In Appendix H.5] we find the optimal TPP on Wine is 9.6 for Adam and 7.4 for Muon. We also find that the optimal learning rate follows jP while independent weight decay should be kept constant as we vary model sizes subject to the fixed FLOPs budget, consistent with Bergsma et al. [4].
Summary: Muon, SOAP and Shampoo achieve consistent speedups around 1.4 over Adam W from 190M to 1.4B parameters, but only under good hyperparameter transfer. We verify that combining μP and 1 D-saled weight decay yields near-optimal hyperparameter transfer for all tested optimizers, and both components are critical for achieving the best performance.
Discussion
Robust optimizer comparisons at scale depend much the scaling rules as on the optimizers themselves. We demonstrate that even modest deviations from optimal scaling rules incur dramatic efficiency losses, sufficient obscure meaningful differences between optimizers. Under our best- effort optimal scaling rules, matrix-pres pettrixprrconditioneed optimizers like Muon, SOAP and Shampoo prove significantly more compute-efficien than AdamW W up to 1.4B models with relatively stable efficiency gains, in contrast to the findings in Wen et al. [4]. Among the tested optimizers, we find Muon achieves top performance while being the easiest to use, tune, and scale, due to its algorithmic simplicity, small number of hyperparameters, and low compute and memory overhead.
The Maximum Update Parameterization (μP) is a crucial component of our scaling rule. While jP has a reputation of being mathematically involved, we show it can be generalized to many matrix- preconditioned optimizers with brief and simple calculation. However, we identify few settings where μP fails to model expressive preconditioners at realistic widths and leads to suboptimal learning rate transfer. Developing more robust scaling rules compatible with expressive preconditioning likely requires a stronger focus on understanding finite-width dynamics, such as in Large et al. [26], or designing optimizers with wiith guarantees built in, e.g., Muon and Scion 7 36).
Lastly, our experiments show that scaling weight decay as 1/ 1/width matters as much for the perfor- mance as scaling the learning rate. Understanding why this scaling is effective across optimizers, how weight decay shapes the training dynamics of matrix-precomee preriixpreconitiooned optimizers, and whether better scaling exists, exciting and important open questions.
10

---

Acknowledgments and Disclosure of Funding
We thank Martin Marek, Andres Potapczynski, and Sanyam Kapoor, for helpful discus- sions. This work was supported by Google's TPU Research Cloud (TRC) program: http://sssesssssrrrc.ooooooo/eee SQ thanks the support of the Two Sigma PhD Fellowship.
References
[1] Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv preprint 1rXiv:2002.11800 2020.
[2] Kwangjun Ahn, Byron Xu, Natalie Abreu, Ying Fan. Gagik Magakyan, Pratyusha Sharma, Zheng Zhan, and John Langford. Dion: Distributed orthonormalized updates. arXiv preprint arXiv:2504.05295, 2025.
[3] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.0901 2020.
[4] Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Power lines: Scaling laws for weight decay and batch size in Ilm pre-tratee pre-training. arXiv preprint arXiv:2505. 2025.
[5] Shane Bergsma, Bin Claire Zhang. Nolan Dey, Shaheer Muhammad, Gurpreet Gosal, and Joel Hestness. Scaling with collapse: Efficient and predictable training of IIm families. arXiv preprint wwxiv:2509.25087, 2025.
[6] Jeremy Bernstein and Laker Newhouse. Modular duality in deep learning. arXiv preprint arXiv:2410.21265, 2024.
[7] Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv preprint 1.Xiv:2409.20325: 2024.
[8] Blake Bordelon and Cengiz Pehlevan. Self-consistent dynamical field theory of kernel evolution wide neural networks. Advances Neural Information Processing Systems, 35:32240-32256, 2022.
[9] Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309.16620, 2023.
[10] Blake Bordelon, Hamza Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head trans- former dynamics. Advances in Neural Information Processing Systems, 37:35824-35878, 2024.
[11] Dan Busbridge, Amitis Shidani, Floris Weers, Jason Ramapuram, Etai Littwin, and Russ Webb. Distillation scaling laws. arXiv preprint arXiv:2502.08666, 2025.
[12] Jie Chen and Edmond Chow. A newton-schulz variant for improving the initial convergence in matrix sign computation. Preprimi ANL/MCS-P5059000 14, Mathematics and Computer Science Division, Argonne National Laboratory, Argonne, IL, 60439, 2014.
[13] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances neural information processing systems,
4992005-4923,,. 2023.
[14] Nolan Dey, Bin Claire Zhang, Lorenzo Noci, Mufan Li, Blake Bordelon, Shane Bergsma, Cengiz Pehlevan, Boris Hanin, and Joel Hestness. Don't lazy: Completep enables compute-efficient deep transformers. arXiv preprint 250.:2500..0668. 2025.
[15] Jimmy Ba Diederik P. Kingma. Adam: Method for Stochastic Optimization. International Conference Learning Representations (ICLR), 2015.

---

[16] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.21177 2024.
[17] Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A Alemi, Roman Novak, Peter J Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee, al. Scaling exponents across parameterizations and optimizers. arXiv preprint arXiv:2407.05872, 2024.
[18] Roger Grosse and James Martens. A Kronecker-Pactored Approximate Fisher Matrix for Convolution Layers. arXiv 1602001477 2016.
[19] Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimization. International Conference on Machine Learning, pages 1842-1850. PMLR, 2018.
[20] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint Stt::2203./55550 2022.
[21] Satoki Ishikawa and Ryo Karakida. On the parameterization of second-order optimization effective towards the infinite width. arXiv preprint arXiv:2312.12228 2023.
[22] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[23] Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cecista, Laker Newhouse, and Jeremy Bernstein, Muon: An optimizer for hidden layers in neural networks, 2024. URL
httpps://kellerjordaan.. giu......... itt/ii//pooo http///sssssoonn/
[24] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
[25] Priya Kasimbeg, Frank Schneider, Runa Eschenhagen, Juhan Bae, Chandramouli Shama Sastry, Mark Saroufim, BOYUAN FENG, Less Wright, Edward Z. Yang, Zachary Nado, Sourabh Medapati, Philipp Hennig, Michael Rabbat, and George E. Dahl. Accelerating neural network training: An analysis of the algoperf competition. In The Thirteenth International Conference on Learning Representations, 2025. URL htttps://operrevilt
[26] Tim Large, Yang Liu, Minyoung Huh, Hyojin Bahng, Phillip Isola, and Jeremy Bernstein. Scalable optimization in the modular norm. Advances in Neural Information Processing Systems, 37:73501-73548, 2024.
[27] Jaehoon Lee, Yasaman Bahri, Roman Novak, Samuel S Schoenholz, Jeffrey Pennington, and Jascha Sohl-Dickstein. Deep neural networks gaussian processes. arXiv preprint 1rX1v:1711.00165, 2017.
[28] Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE transactions on neural networks and learning systems, 29(5): 2017.
[29] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint 1943, 2024.
[30] Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, al. Muon is scalable for 11m training. arXiv preprint 2502. 16982, 2025.
[31] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rXiv:1711.05101, 2017.
[32] Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Gordon Wilson, and Micah Gold- blum. Small batch size training for language models: When vanilla sgd works, and why gradient accumulation is wasteful. arXiv preprint arX 2025.
12

---

[33] Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint arXiv:1812.06/60 2018.
[34] Sean McLeish, John Kirchenbauer, David Yu Miller, Siddharth Singh, Abhinav Bhatele, Micah Goldblum, Ashwinee Panda, and Tom Goldstein. Gemstones: A model suite for multi-faceted scaling laws. arXiv preprint erXiv:2502.0666: 2025.
[35] Depen Morwani, Itai Shapira, Nikhil Vyas. Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's 's preconditioner. arXiv preprint arXiv:2406.17748, 2024.
[36] Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with norm-constrained Imos. arXiv preprint arXiv:2502.07529, 2025.
[37] Shikai Qiu, Atish Agarwala, Jeffrey Pennington, and Lechao Xiao. Scaling collapse reveals universal dynamics in compute-optimally trained neural networks. In OPT 2024: Optimization for Machine Learning.
[38] Samuel Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep informa- tion propagation. arXiv preprint 161.v22611.01222, 2016.
[39] Andrei Semenov, Matteo Pagliardini, and Martin Jaggi. Benchmarking optimizers for large language model pretraining. arXiv preprint wrXiv:2509001440, 2025.
[40] Ishaan Shah, Anthony Polloreno, Karl Stratos, Philip Monk, Adarsh Chaluvaraju, Andrew Hojel, Andrew Ma, Anil Thomas, Ashish Tanwer, Darsh J Shah, et al. Practical efficiency of muon for pretraining. arXiv preprint arXiv:2505.02222. 2025.
[41] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596-4604. PMLR, 2018.
[42] Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D Cox, and Rameswar Panda. Power scheduler: A batch size and token number agnostic learning rate scheduler. arXiv preprint ar arXiv:2408. /3359, 2024.
[43] Hao-Jun Michael Shi, Tsung- Tisnng-Hiien Lee, Shintaro Iwasaki, Jose Gallego-Posada. Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-pat aata-parallel pytorch implementation of the distributed shampoo optimizer for training neural networks at-scale, arXiv preprint arXiv:2309. 06497, 2023.
[44] Chongjie Si, Debing Zhang, and Wei Shen. Adamuon: Adaptive muon optimizer. arXiv preprint arXiv:2507.11005 2025.
[45] Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, al. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.20534, 2025.
[46] Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales. Advances in Neural Information Processing Systems, 36: 16:1036-1060. 2023.
[47] Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. Soap: Improving and stabilizing shampoo using adam. arXiv preprint arXiv:2409.1 11321, 2024.
[48] Xi Wang and Laurence Aitchison. How to set adamw's weight decay as you scale model and dataset size. arXiv preprint arXiv:2405.1366 2024.
[49] Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang. Fantastic pretraining optimizers and where to find them. arXiv preprint arXiv:2509. 2000:226009000666..... 2025.
[50] Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling, arXiv preprint 2rXiv:2409.151566 2024.
13

---

[51] An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, al. Qwen3 technical report. arXiv preprint WrXiv:2505.09388 2025.
[52] Greg Yang and Edward J. Hu. Feature Learning in Infinite-Widhh Width Neural Networks. International Conference on Machine Learning (ICML), 2021.
[53] Greg Yang and Etai Littwin. Tensor Programs IVb: Limit. International Conference Learning Representations (ICLR), 2023.
Adaptive Optimization in the Infinite-Wide -iith
[54] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer. Advances Neural Information Processing Systems (NeurIPS), 2021.
[55] Greg Yang. Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203. 22202000006,, 2022.
[56] Greg Yang, James B. Simon, and Jeremy Bernstein. A Spectral Condition for Feature Learning. Preprint arXiv:2310.17813 2023.
[57] Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Feature learning in infinite-depth neural networks. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023.
[58] Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George Dahl, Chris Shallue, and Roger B Grosse. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019.
14

---

A Update Rules For Matrix-Preconditiond Optimizers
This section presents the update rules for the optimizers considered in our work. Let W; E RR*x denote the weight matrix of a layer at time 6, and let G, E R=xx denote its gradient. We use wy E _ and 9i Im to represent the flattened versions of the weights and gradients, respectively, The hyperparameters include the learning rate (1)), momentum coefficients (8), regularization factors (e), and inverse exponents (e). Elementwise multiplication is denoted by O and vector division is applied elementwise.
Adam
[5 adjusts the magnitude of the first moment by the second moment. The update rule is:
- (wf(wrr1)
(6)
m4 31m- 1 (1-11)99
(7)
U/ 321--1 + (1 32)9e O g:
(8)
m
fy f 1- Bi
(9)
Ut
0g 1 132
(10)
m
W/ W4-1
(11)
Vēn +
Shampoo
19
preconditions the gradient with a Kronecker-factored preconditioner. The update rule
is:
M, + M, -1-
(12)
L - 32 + (1- G
(13)
R₁ + 32R1-1 (1 G-) G.
(14)
L
L, + (1 Bb)
(15)
R₂
Rg + (1 (3)))
(16)
W441 +- We - ((L + =1) + eR)
(17)
(18)
SOAP 47 runs Adam in a rotated space. The update rule is:
M, < 3,M, (1- B,)
(19)
L - 3,L,-1 + (1- G.
(20)
R₂ + + + G2 Gt
(21)
Q < Figenvecttrs(II) (Tt)
(22)
Q Figenvectors(Pi) (R.)
(23)
G² < (Qt) G.Q.
(24)
M/ < (QE) T. M.Q.!
(25)
V, - 3 + BB O.
(26)
WeS1 + W. - - M/
VV+c (Q!))
(27)
(28)
15

---

Muon 23
utilizes Newton-Schulz to compute the matrix sign [12 of the gradient. The update rule
is:
M, + 8 + (1 -1- B1)G.
(29)
M
W411 t W₂ Newton-Schulz
(30)
M +e
(31)
AdaMuun 44 simply applies Adam on top of Muon's orthogonalized gradient.
technique that adjusts the direction of primary optimizer's update to match the step size of a reference optimizer with more stable magnitude. Let Q2(G) denote the update from the original optimizer, and Q1 (G) the update from the reference (grafted) optimizer. The grafted update, denoted as Q #Q2, is given by:
Grafting 4
[Q) Q2(G)
W4+1 We " (1Q2 +
(32)
+e
Blocking 43 47 is technique that partitions the weight matrix into sub-blocks, and updates are computed independently for each block.
B
Derivations of Maximum Update Parameterization
Table Scaling rules for the per-layer learning rate 7) damping parameter(s) as function the weight matrix size (din: dout) ), depth L, block size (bin, bout) and number of blocks reck using blocking. multiplier 1/L.s is assumed to apply to the output of every residual block. e,, are positive exponents for Shampoo, and 0-1 indicators for preconditioning either side for SOAP. For parameters outside of the residual blocks (e.g. embedding and last layer), set L AdaMuon refers to the one used in denominator when applying Adam on top the orthogonalized gradient.
!QQ; denotes the correctly scaled learning rate for optimizer Q1. The for
Shampoo (EI, ER)
SOAP (ET., en)
Muon
AdaMuon
Grafting Q #Q2
(dout/ /di)) 2- 1,2(e+tRR) fllk
ber/2 beR/2
dout
1
1)
eLTeR
dm gen/2
dm
dm
7/21
din dout Tblkk
but b Ldout
1
dia
1 don
1
dout dn
L V
dout
B.1 Initialization
Before discussing specialization to any particular optimizer, we make the observation that the initialization scale a of We is independent of the optimizer, with = 0(1/ for & < L and σg = ((/d). This result follows from the combination of stability and feature learning condition of μP. Specifically, to ensure stability at initialization, jPP requires all activations in the hidden layers have (1) entries and that the function output O(1):
he(g") Werz-1(E") @(1), L- f(E') O(1)
(33)
The first condition implies that W2 have entries of scale o, O(1/ /dn.,)) for all but the last layer, where d is the input dimension of layer (in our scalar-input setup, din, = and din,t = d for f < L), while the second condition implies that feature learning. i.e., the change in the last layer feature Ary,-((E') ((1) per step, while ensuring the resulting change in the output Af(g") is (1)) (predictions are updated without diverging), the last layer weights must be ((1d)). To see that, note Af( ) = w A is sum of d i.i.d. random variables, each with a generically non-zero mean due to the correlation between elements in and elements in A The correlation is induced by baoppaaagn causing A to be strongly aligned with wr. This will become clear Appendix B.& Therefore, given Arzz1 is O wr and thus L must be (11/d) by the Law of Large Numbers (LLNN,, if maximally initialized
1 <
O(1vdd). To then ensure non-negligible
Foz 0 is also allowed and in fact common, but requires carrying out the analysis the 2nd gradient step.
16

---

Moving forward, we omit layer subscripts and abbreviate z<E) as 2, xz'') as x', 8(5) as 6, and G(E)
G. Recall from Section 3.T initialization have
(1),
TTx' ((din), =x== O(din) 8 O
(34)
deus
dont
The gradient for single datapoint is
= V=Vw((E) = 8x E Re x xdm.
(35)
The yP condition from Section 3.1 can restated as follows: for an optimizer that returns update (G),, want the per-layer learning rate 7) such that
Choose learning rate such that nQ(G)z' = (1), where = 8x T, e
doul ) x = 0(1), x TT = (din).
(36)
Solving for 1) and any additional hyperparameters (e.g. damping) is now just algebra plus bookkeeping of how quantities scale with dn and dout-
B.2 Warmup: SGD and Adam
We start by rederiving μP for SGD and Adam in this notation.
SGD. For SGD, Q(G) = G, so
Q(GG)a n/
(37)
n8((= x²).
(38)
Since TT ((da) and each coordinate of 8 is 1 (2 @ dout ).i..
).
the vector Gz' has entries of size
Gx' = =
den doot
(39)
The μP condition n/Ga' ((1) therefore requires
1 =
dout din
(40)
recovering the result in Yang and Hu [52] Yang et al. 56 The whole derivation is just algebraic step once the scalings of x, x', and are known.
SignSGD and Adam. For SignSGD, the elementwise preconditioner is
Q(G) (VGo2+c) OG,
(41)
where 3 denotes elementwise multiplication and the square/square root are also taken elementwise. Since Gi d, and
6, = O
I = O(1),
(42)
Jout
we have
|Gil
1 dons
(43)
To keep the elementwise factor
G
(44)
GGyl +
nontrivial in the width limit, we need e to be of the same order as [Gtyl, so we reparameterize
C'
é' = O(1)
(45)
dout
17

---

With this choice, each entry of Q(G) is an (1) function of (6) with no residual dependence on di or dout. As a result, each row of Q(G) has an inner product with x' of order
Q(G)a' =(dm),
(46)
by an LLN argument. Hence
n((G)a' (zjdin),
(47)
and we need
=
θ
din
(48)
Yang and Littwin [53] formalize this using the OuterNonI in instruction and show that Q(G) has an (1) infinite-width limit after factoring out the explicit dm.
The same scaling holds for Adam: accumulating the first and second moments of the gradient over
steps does not introduce new powers of din or dout. so Adam also requires 1) = θ dmn
B.3 Shampoo and Muon
For matrix- pradiiireeondiiioonned optimizers we only need to additionally track how the preconditioners scale with width. As in Section [3T]] we ignore preconditioner accumulation and gradient momentum, which do not affect width scaling (53) in the infinite-widhh limit.
For Shampoo, the single-sample left and right preconditioners are
L = 8x =ō" = (x x)) 58
(49)
R = xōT 8x T = (5 8) xx T.
(50)
and the update is
Q Shampoo (G) = (L (L+ tLl) el. 8x T (R+ + (R+E ER
(51)
Define
= x = (din), A ===
>= = 0 doout
(52)
The vector d is an eigenvector of with eigenvalue X, and x an eigenvector of R with the same
eigenvalue A. Hence
(I.+ -28 (A+ + FL) ELS, x ( Cx = (A 1 ER)
(53)
-CR T
(54)
and therefore
QShampoo (G) = (A + =L) eh (A + (n)
cn)
-Rn ox T.
(55)
As before, we parameterize damping in units of the nonzero eigenvalue:
ER = EB, EL = é
((),,
(56)
so that
(A (L) eL = -1 + +x2)) nh 0(1 (x)=))),
(57)
(A + ER) -eR = =x - +6)) -eR = (ex))).
(58)
Thus
Q Shampon (G) = 6 - 8x
(59)
Applying QShampoo ( (G) to x' gives
Q Shampoo (G)z = - 1- -eltte)) 8(z x¹)
(60)
θ A x(ee+oe)) 0 1 (dda))
(cm))) 1-R2-eR
(61)
dd dontt
(62)
18

---

since A 1 dm dout
Therefore the learning rate should scale as
17 1- dout
(63)
As Tensor Program. Going one step further, we can expand the preconditioners to write
+ (11c = (I Po) + (A+ L Po:
(64)
(R+ ERI) ERl) ER =R
ER en ( P (A+ER) =R P2.
(65)
where P. 850//0/000 /(6")) 8) and Ps XT' / are rank-l projectors. This expresses Shampoo purely in terms of vectors and scalars (deterministic numbers as D oo). Therefore, Shampoo
can be written in terms of instructions supported in the Tensor Program [53]. We show a similar construction is possible even if batch size is larger than in Appendix B.10 This is much stronger statement than features being updated at 6 rate in width, as it shows that the training dynamics of models trained with Shampoo with pP scaling admit well-defined, deterministic infinite-wied limits, by the Master Theorem (53) This is why we can expect the loss curves in Figure (left) to be highly consistent across widths.
Blocking, Blocking reduces the cost of second-order preconditioners by partitioning W E Rdoa xdn into Rout X Rin blocks of size bout bin, where
dm
dout
Rin
nout
(66)
bin
bout
Let Giy E Xbb denote the (i.j)-th block of the gradient,
Gúj = =
(67)
where 8 Roont and zy E hm the i-th and j-th chunks of 5 and z, respectively, and similarly zzJ Rho is the j-th chunk x'.
Shampoo applied independently to each block. The blockwise preconditioners are
Lij = 8,23 2Ø = (x] 23) 6,67"
(68)
Rig 2, 6, (6. 6.) zzz
(69)
Define the block eigenvalue
Arj (o) x.,) x,) (6)
(70)
As before, 6 is cigenvector of Lig and xy is eigenvector of R.g with eigenvalue Az. The
Shampoo update on block G.j is
Q(Gig) = (Taj + EL e) 8,0e (Rij cn))
(71)
Q(G.)aa = (Nis + +EL) + ER) en (2)
(72)
The inner products now have scales,
2TJ x. = O(bm))
87 6 = O
Busay dout
(73)
so we get
Aij @ - (Dnbbag
(74)
doet"
We again parameterize damping as G = Aig and cn = EnAis with c'n = =(11, which yields
Q(Gig)a's = @ ( 0, (1(xT,1)) 2.).
(75)
19

---

The i-th chunk of the feature update is then
Ah; = Σ Q(Gi))a
(76)
j=1
Ain
1) 6
+
Ais
toz
tex)
(z)
(77)
Ain
Σ θ
(esttend
(78)
where {cijli are i.i.d. (due to permutation symmetry over j) with generally nonzero mean per entry, and we used 2z n'z (bmn). Therefore, recalling 6; = ((1/dout) and the scale of Xij. we have
teezton) bin Tin
Ah = θ 17
(79)
dout
teenten) din
+ 17
hindoge Dout"
(80)
dout
since Tinn din/Din /in. Enforcing Ahe (1) leads to
oz tea
77 =
damt dia
Lin Canleag
dout
(81)
Equivalently,
0
(($em)
daml
1- teR)
(Nmnoot) (el tex)
(82)
din
since Pout
du dant bis Dous"
When bin din and bout dout, we recover the full Shampoo scaling When EL + and bin bout 1, Shampoo degenerates to Adam and
e1, eR
77 dams
get the correct Adam scaling θ du
B.4
Muon
Muon [23 applics Newton-Sede Schulz to a normalized gradient to approximate the matrix sign. In the idealized limit of zero damping, full Shampoo with e, eR 1/4 yields the same matrix-sign direction; with nonzero damping, Shampoo and Muon differ only by scalar nonlinearity applied to the singular values, so the Shampoo learning-rate scaling carries over. We now verify this directly for a rank-1 gradient.
Write G = ōx wwov T where = 5/118112- v = =x/|<z/l22- and
o =111|211001 = =0 (v
din dout
(83)
Muon forms
G
G = |G +e
wōv
(84)
σ +
Choosing = (v
din Zour
keeps ō (1)..
The Nowton-Schulle iteration used in Muon has the form
Y =
2 (3T Y Yk), Yo G
(85)
so for rank-1 G we can check inductively that Yx = usau T with the scalar recursion
1
Sp+1 = 3/6 (3- 22) σ.
(86)
2
20

---

For any fixed It = (()), this gives s; ( and therefore
Newton-Schulz(C) ul)) uo
(87)
Indeed, this holds regardless of the coefficients used in the Newton-Schulz. Thus
8 xTx
Quu (G) ( = (11)
di dour
(88)
|8112
|0|1
and the μP condition Muon (G)= yields
17 θ
V
dout
diu
(89)
Finally, the normalization requires
(V
= =
di dout
(90)
B.5 AdaMuon
AdaMuon applies Adam top of Muon's 's Newton-Schut output 0 Newton-Schulz(( G).
4
From the previous subsection, O is rank-1 and has entries of size
|0 =
(91)
Vi Cout.
AdaMuon then applies an Adam-like elementwise rescaling to O. keep the map f + t/([t] + nontrivial as width grows, we take
c
dou e
(92)
so each entry of Q AdaMunn (G) is an (1)) function of the corresponding entry of O, with no residual scaling in (din; d. ). By the same LLN argument as in the Adam derivation, this implies
AdaaMun (G)a' (di.))
(93)
and hence the uP condition AdaMuua (G)xz' = (1) gives
7 = θ
(94)
dm
(Here is the one used in the Adam-on-top denominator; Muon's normalization in the previous subsection.)
B.6 Grafting
When using Q2 with Q learning rate grafting, the update is
|[Q( [e
Q(G)
-2(G).
(95)
Assuming both Qı (G) and (2(G) have o(1) rank, which we have shown for Shampoo but generaliza- tion other optimizers this work straightforward, ||Q ~ ||Q1 (G) fly and [[Q2(())
[[) Let n2, denote the μP learning rate for Q2, by construction noa 1Q2 llz ~
y Vduut, due to the alignment between (2(G) and x', which shows [|Q2(G)
dost
dn
should match the scale of [[] so
c= +
(96)
112
after which scaling by [[Q sets Q(G)x' Qi (G)x', learning rate should now follow 7Q:
21

---

B.7 SOAP
SOAP applies an Adam-like elementwise preconditioner the eigenbasis Shampoo's 's left and right preconditioners. Let U E R xdout and V Rdu ddu, be orthogonal matrices whose first columns are the normalized eigenvectors of L and R:
8
U1
|
|0]|
(97)
with ||ō|| v and ||>|||| vx TT. x. Then
U 8 ||8| el
(98)
V x = ||-||| e
(99)
where of and ey are the first standard basis vectors in iRdout and 1Rd Transforming the gradient into this eigenbasis,
G' U GV = U fx V = ||1||
(100)
Thus G' has a single nonzero entry at (1, 1) of magnitude
g = | |
VōTōvzTz=
din
(101)
dont
SOAP applies the Adam-like rule elementwise to G":
Adam((') (VCroz +.) + G²
(102)
Since G' is zero everywhere except at (1, 1), we have
U T
Ad)m(G") 9 g c
(103)
To keep this factor
nontrivial as width grows, we choose
g+e
6= gé' = 6 (
di done
e'=((1.
(104)
so that g g+c Itc (1))
Transforming back to the original basis,
QGDAP(G)= U (''d
(105)
U U V VT
(106)
g
8 T
(107)
g +
||8|| |
1
dx
(108)
1
||]|||
Using ||6|| 0 11) and ||^||| O Vdn), obtain
Vdout
QsOAP(G) (G) = (v dous ox T
(109)
Applying this to x',
QsoAr(G)x's - \ dont
5(x Tx')
(110)
daut
θ
(ddm))
(111)
dows
di dant
(112)
22

---

Therefore,
(osoAp(G)x'= =
7j = (V
dout din
(113)
and SOAP's damping should scale as é @
dm
dout
Blocking. As the Shampoo blocking derivation, partition W (and hence G) into nout X nin blocks of size dout x bin so that
dout
dm
Nout
bout
Ténn
bin
(114)
Let 6e E Rouut and x TR denote the i-th row-chunk and j-th column-chunk respectively, the (1,5) block of the gradient is
Gi = 8,ar] E Rebou bin
(115)
From the global scales above, at the block level we have
||>ح e( v (Vbb))
(116)
||ó:112 - @
(117)
2) x = (bin))
(118)
where z/J is the j-th chunk of x'.
Per-block Shampoo preconditioners for SOAP bases. SOAP operates in the eigenbases of the per-block Shampoo preconditioners
T.g 6,50 20 + <BTam = 1111,122,00 < Touss
(119)
R, = 6, 5, + c catbi 115.112 25 + cA Jon.
(120)
Thus:
Lis has one eigenvector
8; 41
(121)
||5,112
with eigenvalue 1|,,11216,1 cB, and bout orthogonal eigenvectors with eigenvalue cp- Rig has one eigenvector
0y1
(122)
| 11
with eigenvalue | c. and bm orthogonal eigenvectors with eigenvalue e-
Let U, and Vj be the orthogonal matrices whose first columns are us1 and 1y1, respectively. When a side is not tracked (one-sided SOAP), we simply take the corresponding basis to be the identity that side.
SOAP step inside the tracked eigenbasis. In the per-block tracked basis, the gradient transforms as
G = U.
GizVz. U,
(123)
SOAP applies an elementwise Adam-like rescaling
Hy (V @ Giy +cy++) 3 G
(124)
and then rotates back
= =
(125)
23

---

The block's contribution to the hidden update on second input z' is
A ((G
(126)
We will compute the scale of Ah, for each tracking pattern and then sum over j = -l..... Ain to obtain Ah; شك= Ah The yP condition requires
|| = (VDouu))
(127)
L.c. entries of Ah, are ((1).
Cases: both-sided, right-only, left-only, neither tracked. We use eL, eR [0, as indicators for whether the left/right side is tracked. The "neither tracked" case (el scaling.
0) recovers Adam-like
Case 1: both sides tracked (e₁, = En = 1). Transforming the rank block gradient,
Giy U, (6,z))1 ||6,ll2 || ll ei VT
(128)
where ef and ey are the first standard basis vectors in the U/ and V3 coordinates. The unique nonzero entry has magnitude
Si || ||ا O
Danu dous
bin
Ybbu Bin. dont
(129)
To keep the Adam nonlinearity nontrivial, we choose
c=
bimb dout
(130)
Then the Adam-like map sends Gig to
H₁ 0())
(131)
Rotating back,
Q(Gig)a U.H,, V2J 2 U.el (V
(132)
usi (u]1=1)
(1))
(133)
6
θ
(134)
10:1l
Tayll
θ
biu
(135)
116,112
since x. x/3 = (binn) and || 0 (vb....
Thus
| (3)
e (2VVb) bin
(136)
Summing over Nin = d bm blocks along the input dimension,
|hhill = 17 min θ @( ((/b)) 1 7
din Voss
(137)
Enforcing |Ahill =((vbout) gives
"1 +
Shoumbin din
(138)
Case 2: right-only tracked (er. ER 1). Here U.= = Irun and Vj is above. Then
G G V, = 6,2 V3 6 61 e VT so the nonzero column has entries of magnitude t nontrivial,
(139)
Vbu dout
To keep the elementwise Adam operation
Yin
(140)
doet

---

Then the first column of Hij has (1) entries. Multiplying by V 2/z,
Q(G,))= Hiy V z/j
(141)
= (column with (1)) entries). (1)12²)
(142)
(11) @ @((bm).
(143)
So
A A 112 =0 (v
(144)
Summing over vin = din/bin blocks,
| 7 r/nn v biabout) = 1y 0
Au bin
V /iuloout
(145)
n dı V
Boul bin
(146)
Setting ||hhell2 = @(vbou)) yields
θ
Jbnn
(147)
dmn
Case 3: left-only tracked (el 1, LR 0). Now Ug is above and V, = Imm Then
G3 U.U.G.G, Gij U.l 8,
|6 elf 27.
(148)
so the nonzero row has entries of magnitude θ
//huus dout
So we choose
6
θ
bout dosst
(149)
Then the first row of Hij has (1) entries. Rotating back and multiplying dty,
@(Gr))= = Uc(((h))aa) U,eY (hJ2²)
(150)
(151)
8 (bm.):
(152)
16,112
where hy is an elementwise transform of æ, coming from Adam's nonlinearity and is aligned with Ty by the Master Theorem (33, so hJ 2/j = (b..). Hence
|
(
17 (bin))
(153)
Summing over Nin d /in blocks,
| 212 17 (miubra) 17 (ddn)).
(154)
Enforcing |halll @(Vbout) gives
11 = θ
y din
bous
(155)
Case 4: neither side tracked (ez 0; Adam-like). When neither side is tracked, no rotations are applied and we simply recover Adam:
+
din (
(156)
duuu
Unified expressions. Let eL, E [0.1) indicate whether the left/right sides are tracked (for two-sided SOAP take e, eR 1, for one-sided set exactly one of them to and for Adam set both to 0). Define the block concentration factor
Tan (et, bout f by ex/2
(157)
25

---

Then the per-layer learning rate and the SOAP (Adam) inside the eigenbasis scale as
1) θ
Tha (R)
Bout dmn
(158)
dmn
Tha (er.ee))
Bout "L/ =
(159)
dous
dous
These specialize to:
Adam (eg 6R = 0): 1 (L1/dm))
c - ((1/dou:)
Left-only SOAP (E1, L 0): 17 @ dlan (VVoout/ dout).
Right-only SOAP (ez 0, ER = 1): f = 6 dm) e dout))
Two-sided SOAP (e1 = en 1): θ ( din), @VDoutt dout).
B.8 Alignment at Initialization and the Spectral Norm Condition for Matrix-Preeonditioned Optimizers
Given our explicit expressions of the optimizer update purely in terms of vectors and scalars in Appendix B so far, is easy to check that Q(G)) is always (a) aligned with 8 on the left, and (b) aligned with x' on the right. Property (a) explains why the last layer wz must have entrywise scale of (1/wiitth) explicit backpropagation shows d is aligned with wz, and the update to the last layer feature Ah is (1) and aligned with 5 and thus WL. which causes an update to the output of scale f G(width 1 OL), requiring σg = @11/width) for stability. Property (b) shows that is the spectral norm of Q(G) that governs the scale of Q(G)='. Thus, the spectral condition that nl] O dout/ dm)) [56] continues to hold for these pacrix-preconddtionee optimizers.
B.9 Larger than One Batch Size and Number of Steps
So far we have analyzed the first gradient step with a batch size of one. Now we explain why increasing batch size and the number of steps does not change the asymptotic jP width scaling conclusions, as long as both B and / are @(1) with respect to width. We illustrate this with Shampoo and SOAP as concrete examples.
Shampoo. For a batch of size B, the (per-layer) gradient is
B
G=
816) (T)
(160)
B
b=1
up to 1/B factor that does affect width exponents, where (b) E Re and 5(6) E R are the activations and backpropagated signals for the b-th example in the batch. Under the same initialization and feature-kernel assumptions as in Section 3.I each term
(b)
((6)
(161)
has the same width scaling as the single-example gradient. As result, in particular, has O(1) non- zero singular values, all of size O / d /dout) as before, and the corresponding left and right singular vectors are precisely the eigenvectors corresponding to the non-zero cigenvalues of GG and R G T G. Therefore, the nonzero spectrum of L and R remains Odd,, /dot)) and the number of nonzero eigenvalues and associated eigenvectors is @ independent of width, and aligned to non-zero singular vectors of G. As result, all width-dependent scalings remain unchanged. For completely rigorous treatment, we provide an explicit construction in Appendix [0 for expressing Shampoo with batch size greater than scalings carry over.
Tensor Program and show the learning-rate and
SOAP. SOAP applies Adam- like elementwise rescaling in the eigenbasis of the Shampoo precon- ditioner. In the single-sample analysis, we saw that in this basis, the transformed gradient G' has only one nonzero entry with magnitude
(
θ
din
(162)
dour
26

---

and that choosing
dis
(163)
domt
keeps the elementwise mapping t + t/( + +) nontrivial. With batch size / = (1), the number of nonzero entries in G' is still (1) since it only has support in the non-zero eigen-directions of the preconditioners U, V, and each entry remains at the same θ din/ /
Therefore, the SOAP learning-rate and scalings derived in the single-sample case,
(v
1) = θ
dout dan
e =
dm dout
(164)
remain valid for ((1)
I >>> steps. Similar reasoning shows why analyzing beyond / = steps leads to the same asymptotic width-scaling conclusions, long (()). In particular, accumulation in the preconditioners and momentum can be viewed as effectively increasing the batch size.
B.10 Expressing Shampoo as a Tensor Program When Batch Size Exceeds One
We now show Shampoo be expressed Tensor Program even when the batch size is larger than one. Ignoring accumulation, full Shampoo with e = en = 1/4 recovers the matrix-siga sign direction when e - O, and Muon applies Newton-Schulz approximate the same direction; thus similar Tensor Program construction applies to Muon.
For batch size B, the per-layer gradient is sum of B rank-1 terms,
80, (T)
G
(165)
Collect the batch vectors into matrices
A ((()) ...., ((B)
Reo xB
X (0(()) ,..,, (B)
Edux
(166)
so that
G -A_
(167)
B
The only scalars we will ever need are inner products between these vectors. In particular, define the B x Gram matrices
K, x AA
(168)
Ignoring EMA accumulation for the moment (for which a similar analysis applies), Shampoo uses
L GG R=GTG
(169)
which become
L - AK, R= 1 XK,X B
B B²XKa
(170)
This immediately implies rank(L) < B and rank(It) s B, so when B = θ(1) the nontrivial eigenspaces of L and Rt are d(mesdimels..
To make the dependence on inner products explicit, assume for simplicity that K₂ and K₂ are invertible, and define the whitened matrices
AK 1/2
X X
(171)
then
Ä =
(172)
27

---

If either Gram matrix is singular, one can instead interpret 1/2 pscudoinverses; this only affects null directions and does not change any of the statements below about nonzero eigenvalues. We can rewrite the preconditioners
// and K2 as Moore-Penrose
L == S
Kg 12 K2 K2
(173)
B
and
R XS Sa 1
(174)
B
Under the invertibility assumption, Á and X have orthonormal columns, so the nonzero eigenvalues of I. are exactly those St (and similarly for R and Sn). Therefore the matrix functions appearing in Shampoo satisfy
(1+ (1 [-- + Ä(S + + te,1) T
(175)
= +
(176)
where
A1 = K2 + (5= 1
K3 K2
(177)
Similarly,
(R+E = =ER IRR XARXT
(178)
with
(nI) AR= 1==KK122((50)) (
C K 1/2
(179)
These formulas use only the vectors in A and X and the scalars in K2 and K3 (through constant- size operations on B x matrices). In particular, this shows that Shampoo can be implemented using Tensor Program instructions (53), since it reduces to forming inner products, applying scalar nonlinearities to the resulting scalars, and taking linear combinations of @(1) number of vectors. When B 1, K2 and Ka scalars, and the expressions reduce to the rank-1 projector form in the main text.
Plugging these expansions into the Shampoo update.
Q Shampon (G) = (L "")==((+e)) "((=)))))
(180)
and using G 1 B
AAT, we obtain
1 AMXT B
Shampoo (G)
(181)
for some B B coefficient matrix M constructed from K Ks. Cn, Cn, en via constant-size matrix operations. Thus the update is always linear combination of the B² outer products (6')) B 1b with coefficients that are scalar functions of Gram matrices, which is exactly the Tensor Program form.
Finally, we check that the width scalings of and are unchanged when B = O(1). Under jP so each rank-1 term (bT has a initialization, ||2 = @(dm) and ||5(0) 12 θ c
nonzero singular value of size e (Vdin/ dout). Since B @(1), G has @(1) nonzero singular values of this size, so nonzero eigenvalues of Sr. and SR satisfy
A(SL) 0
dia dous
A(SR) θ
dim. dowr
(182)
Therefore, choosing
ET. =
0
din deur
O
dan dous
(183)
28

---

keeps (Sy + c = and (SR + cgl) " nontrivial without trivializing to scalar multiplication, exactly as in the
case.
With these choices, applying Shampoo to a test vector x' proceeds as Section 3.1 each inner product (((6)Ta x' is (di.)) by alignment, the right preconditioner contributes factor (nut) = on the -immenioni span of x.(), the gradient contributes the usual (din dout) factor through (0)) and the left preconditioner contributes @ (d/n (out)) on the ((1)-dimensional span of ((8), Since B = ((), summing over b = ,....., B does not change width exponents, and we obtain
Shampoo (G)x² =0 ((
1--6--10 1-EL-ER ER
din dous
(184)
Hence the μP condition Shampoo (G) = =O(1) yields the same learning-rate scaling as before,
7 )= dout diu
1-ez-en
(185)
Moreover, because Kz/d,, and dout Kg have deterministic limits for fixed B, all entries of the constant-size coefficient matrices A,, AR. M converge to deterministic limits as width grows. By the Master Theorem (53). this implies that Shampoo training dynamics under μP scaling admit well-defined deterministic infinite-width limits, explaining why the loss curves across widths can be highly consistent.
C Derivation for Depth-Scaling Rules
As discussed in the main text, for residual networks we scale the output of each block by 1/L while keeping ((1) feature learning inside each block, following Bordelon et al. [10. Dey et al. 14 This implies that gradients in each residual block scale
8 e
Ldwat
(186)
instead of θ
Thus
dout
T O
Lidhous
(187)
For Shampoo, the relevant eigenvalue becomes
A= (=Tx)) =0 0
du I dout
(188)
and the update is still
Q Shampoo (G) = EL+ -L (A + ER)
CR 8z
(189)
Parameterizing damping in units of the new À
el, E ER = = 6 O
(190)
we have
Q Shampoo (G)z 0 -leette))) -x(x-tectea)
8(x 8(xTx') x')
(191)
Todout (cdin).
(192)
Relative to the width-only case, we have replaced
8
8 +
(193)
T.
Q Shumppo = 0 (1.2lectte))11
(0()))
(((dm)) (dm)
1-ez-eR
(194)
29

---

To maintain 77Q Shampoo (G)) (1) we therefore need
7 0 () -2-2(ec
dout
1-ep
(195)
din
and the damping scales as
CL = ( = = din L
(196)
For the remaining optimizers in Table depth scaling can read off from how the update behaves under an overall rescaling G + cG with 1 L.. In particular, for rank-1 gradient G 5z we have
|| || ||ح|| e
[iv
din dont
(197)
and the RMS of entries G is (LL(oL)))))
p- + before applying Newton-Schulz Choosing = @[[]G]| 5)) keeps this normalization nontrivial. With this choice, under 8 r> 5/L. both the numerator and denominator scale by 1/L, L, so the normalized matrix and hence the Newton-Schulz output are unchanged. Therefore r has no additional L dependence, while
Muon. Muon normalizes the (momentum) gradient by |G]]
== +
i
din. dous
(198)
In other words, the Muon learning-rate scaling stays = 4 (VV2 and only depends on L.
SOAP. SOAP applies an Adam-like elementwise map & - t/(pt] + c) in rotated basis. Under G G/L, all transformed entries scale as f - t/L. so this map is invariant provided scales the same way. Thus the SOAP learning-rate scaling unchanged (no L dependence), while its picks up an extra factor L/L relative to the width-only scaling. In particular, in the blocked/one-sided setting with tracking indicators e,, ER E {0, 1}, this gives
baur ",,""" Ldoou
(199)
This pairs with the width scaling 17 =
pus din
AdaMuon. AdaMuon applies Adam Muon's 's orthogonalized gradient, which already removes the L-dependence the gradient. Therefore neither the learning rate nor Adam's 's denominator acquires
any additional dependence. In particular, r = (11/din) and c = (VV d (for the Adam denominator) remain unchanged.
D Scaling Rule Validation Experiment Details
D.1
Width Scaling
In Section 3.2 we train GPT-2 architecture transformers on the OpenWebText dataset for 100M tokens using linear decay learning rate schedule with no weight decay. fix the depth (number of transformer blocks) to be 3 and attention head dimension to 64. We use batch size of 512 sequences of length 128. We use a small vocabulary of size 96 so it is practical to apply the full preconditioners the embedding and readout layers in order verify whether learning rate is scaled correctly for those layers where only one of the dimensions grow. We compare μP, where the per-layer learning rate is parameterized by (D) "Tbase()) (D/ (DDDaase)) ) with Thase the base learning rate, Dbase the base width, and ae derived from Table for each layer and optimizer, against the Standard Parameterization (SP) where m){( (e(D) TJaase For example, when using Muon, ae 0 for hidden layers, 1/2 for embedding, and - 1/2 for readout. We set D 128, the width at which SP matches p.P.
zero-initialize the readout layer for both μP and SP, making learning rate scaling their only difference. We also zero-initializz MLP and attention output projections, as done in moodddd-nanneee -anogpt
30

---

Muon
AdaMuon
Shampoo
#Shampoo
# #Shampoo (blk)
Spectral Norm Stable Rank
10a
10"4
Training Steps
2
2
Width D
Figure 7: Stable rank and spectral norm of the weight update at each step throughout training. In early training, the stable rank well modeled as a (1)) quantity independent of width, but quickly grows training progresses unless constant block size is used (blk). jPP undershoots the spectral norm for AdaMuon and Shampoo with grafting by a factor inversely related to their stable rank.
The embedding layer initialized with a standard deviation 0.1. All other layers are initialized with 1/D variance. We use 31 = 0.9 (first moment) and 83 = 0.95 (second moment and preconditioners) for all optimizers. We set = 10 at the base model, except for Shampoo variants, which use relative 10 No weight decay is used for these experiments.
We use the GPT-2 architecture without bias and no learnable layernorm parameters. To increase the maximum embedding dimension we can experiment with, we set the MLP expansion ratio to rather than 4, following 337
D.2 Depth Scaling
use the same setup in Appendix D. but keep width at 128 and only scale depth. We use a base depth of 3, where = 0 and a = 1 matches exactly. Due to zero-initialization of the output projections,
In Section 3.3
and have identical initializations.
D.3 Compute-Optimal Scaling
We use the same setup in Appendix D.I but use the FineWet dataset tokenized with the full GPT-2 tokenizer. We set 32 = 0.98 for all experiments. Shampoo uses é,, CR 1/2, i.e., Shampoo?.
E
Finite-Width Deviations
E.1
Stable Rank is Not @(1) in Width in Late Training
Fsgor the stable rank and spectral norm of the per-step parameter update throughout training for a range of widths under fixed base learning rate in jPP. We show them for the up-projection the first MLP layer, but other hidden layers behave similarly. Unlike Adam, whose elementwise preconditioner is known to be incapable of increasing the stable rank of the update beyond ( [56], matrix-prn preetixxrreconiiioneeed updates have low, ((1) stable rank only at the beginning of training, and achieve high stable ranks that scale with width training progresses. When using fixed block size, the preconditioner is less expressive and the stable rank remains nearly indeeednddennnntt throughout training.
As result, for optimizers that always normalize the update RMS O(1) after matrix maatixxpreconditiooningg (without blocking), the spectral norm of the update decreases over time. Both Shampoo with Adam-grafting and AdaMuon apply such normalization, through grafting and elementwise scaling respectively. For these optimizers, inneiindependent learning rate ensure O(1) feature learning
31

---

(translating to (1) spectral norm for hidden layers) throughout training, since early and late training require the learning rate scale differently with width. As pP sets the learning rate based on the infinite-width limit, which effectively focuses early training, the learning rate is too small for the bulk of training, explaining the increasing trend of optimal learning rates in Figure 2 Naturally, one might think it therefore makes sense to scale the learning rate instead based on the equilibrated stable rank, which appears to scale as (widthh") for some 7 close to I, but this would necessarily lead to instability in early training. By contrast, spectral normalization produces the right update scale throughout training by making time-dependent adjustments.
Finally, note that it is the combination of high stable rank with RMS or Frobenius based normalization that breaks jP's learning rate transfer. For example, Muon and Shampoo without grafting show near- perr-perect learning rate transfer under jPP (Figure 2 Adam also shows good learning rate transfer
despite applying RMS-like normalization [55] because does not produce a high stable-rank update due to an inexpressive preconditioner [56].
E.2 SOAP
Our juP derivation for SOAP relies crucially on the observation that the gradient has ((1) rank, since tB is (1), and a result the rotated gradient and its first and second moments the preconditioner's eigenbasis are only non-zero in a ( (d)-diimesioona subspace (non-zero only in a finite block in the top left corner). While true in the infinite-width limit, for realistic finite widths and large batch sizes this picture completely fails to model the structure of SOAP's update, which densely populates the full matrix, as illustrated in Figure for an MLP up-projection layer. By applying Adam in the eigenbasis, SOAP can amplify tiny entries in the gradient to be ~ =1, yielding dense update even the moments have small stable ranks, so long as they have high matrix ranks, which is expected from large batch sizes or training steps.
Given this observation, update, but none led to satisfactory transfer. For example, modeling the spectral norm as @(v di /dout) as in a random Gaussian matrix with unit entry variance tends to decrease the maximum stable learning rate as width increases, and modeling it as ((udoo) as if it were simply Adam undershoots the optimal learning rate width increases.
experimented with alternative models for the spectral norm of SOAP's
F Spectral Normalization
We estimate the spectral norm the optimizer update using online power iteration following Large et al. [26). Specifically, we maintain estimate for the top singular vector of the optimizer update A, each step At each step, we compute yt = Ar4, estimate the spectral norm of A, as ā. ||seel || and update Vt Ut4 A,
A, ys for small 0. Therefore, only two matrix-vector
multiplies are needed per step per layer. A, varies slowly over time, as expected from the use of momentum, and o, should well approximate the instantaneous top singular vector and singular value. At initialization, Up a random unit vector. Near initialization, A, may zero for some
330
U
Figure Visualization of the absolute value of the SOAP update in the preconditioner's eigenbasis, cropped to the first 32 x 32 block. In the eigenbasis. the gradient (G"), first moment (M), and second moment (V) concentrate along a few top entries along the diagonals, but the update (U) is dense.
32

---

μP
Norm
Norm+ Readout
4.0
3.0
Loss
2
2.0
2
D
101
10a
10a
10a 10a
10
102
10a
Training Steps
Figure SOAP (no blocking) training dynamics. μPP (left) does not give consistent training dynamics. Applying spectral normalization (middle) gives more consistent training dynamics but wider models learn slightly more slowly initially. Further replacing SOAP with Adam for the readout layer gives consistent dynamics.
A. would result a zero vector. To prevent this, we layers, in which case performing V2+1 [[] to +
skip this update it would produce vector of norm less than 0.5.
Given the estimated spectral norm dt, we rescale the update A, / dous A,/ setting its spectral di
norm to the desired
dous
up to errors the spectral norm estimate. dim
RMS Normalization on the Embedding Layers. For the token and positional embedding layers, we perform RMS normalization on the full embedding matrices instead, which much cheaper with large vocabulary. The rationale is as follows: since inputs the embedding layer one-hot, the embedding layer is better viewed as a fixed-size collection of vector parameters. Normalizing the spectral norm of each vector (viewed as di 1 and dout dembd matrix) reduces to RMS normalization. Based on this reasoning, Bernstein and Newhouse [6] proposed applying RMS normalization each such vector. Here, adopt the cheaper alternative of applying RMS normalization to the full embedding matrix, which should achieve similar results.
Inconsistent SOAP Readout Dynamics Even with Spectral Normalization. Figure shows that even with spectral normalization, SOAP training dynamics are not fully consistent across widths, as wider models learn more slowly initially. We found this is due to the spectral norm overestimating the level of alignment in the readout layer for SOAP and can fixed by replacing SOAP with Adam in the readout layer. We do so for the experiment in Figure 3 (right). However, we found this fix is not necessary for good learning rate transfer or performance given long training horizons (c.g.. 20 tokens per parameter), where spectral normalization alone is sufficient. If blocking is used, then training dynamics are already consistent under jPP without spectral normalization or using Adam for the readout layer (Figure (left).
G Muon μP Scaling
Incorrect Alternative Scaling, Using the experiment setup Appendix D.I. Figure I0 (left) shows that the following learning rate scaling rules do not transfer the optimal learning rate, only the / /dott/ dn scaling (first column) does:
Muon-Kimi multiplying the learning rate by 7 = 0.2 max(din domr) by Liu et al. (30 We show two versions of Muon-Kimil one where no further width-dependent scaling is applied on top of the 7 factor (Muon-Kimi ((1)), the other where one applies the correct Adam jP scaling (d,,') on top of y (Muon- Mimi ((d=.)). The latter reflects the purpose of matching Adam RMS, which is to make Adam learning rate directly transferable to Muon. However, both led to poor transfer. The latter approach fails because Liu et al. 30) assumes the Muon update is full stable rank, thus differing from jPP
2. Muon O(d=))): Shah et al. 40 directly uses Adam's 's jPP scaling rule by Shah al. which undershoots the learning rate in all layers.
33

---

Muon
Muon-Kimi eldp.)
eld
AdaMuon w/ RMS AllsAnligmmet
Loss
1.4
Loss
1.3
1.3
Multiplier
Multiplier
RMS(Ah)
RMS(AA)
10
100
2'
2'
2'2"22200
202
2
2
2
222
2
2
2222122222
2
2
22
2'
2
2"
2
21
2
2122
Width
Width
Step 10
Step 1000
Figure 10: Alternative learning rate scaling for Muon. (Left) Learning rate scaling proposed in Liu et al. [30] (Muon-Kimi) (immi and Shah et al. [40] (Muon O(dm')) ) do not transfer optimal learning rates. (Right) For AdaMuon, adding explicit RMS- RMS-alignment according original implementation 44 does not impact learning rate transfer results and the optimal learning rates continue to shift right.
Note, however, Liu et al. [30] also scales batch size, training horizon, and number of layers together with width, in which case the optimal learning rate scaling may not follow μP. Furthermore, their Adam learning rate is determined from empirical power law fits, rather than scaling as juP's ((1/dm). These differences may explain why they still observe good learning rate transfer from Adam.
RMS Alignment. The original AdaMuon optimizer performs explicit RMS alignment, where the update -1 0.2 don ||or|l, where Ô, is obtained from applying Adam on top of the
orthogonalized gradient. Since Adam already normalizes the RMS of the update to o(1) every step, doing so has little impact on how the learning rate should be scaled. particular, jPP scaling remains " 1/d,n Therefore, did not use RMS alignment in experiments in Figure Figure T0 (right) shows adding this explicit RMS alignment does not address the shifting optima observed in Figure
H Scaling Law Experiments
For the scaling experiments Section we use the FineWeb dataset with sequence length 1024 and batch size 128. We the Llama-2 architecture without bias or learnable RMSNorm parameters, following modded-nanogpi The models are trained with 20 tokens per parameter, a linear decay learning rate schedule, and gradient clipping of 1.
Params
Seq Len
Embedding Dim
FFN Dim
# Layers
# Heads
190M
1024
512
2048
32
380M
1024
768
3072
32
12
640M
1024
1024
4096
32
16
1.4B
1024
1536
6144
32
24
Table Architecture specifications for each model size.
H.1
Base model tuning
Optimizer choices. We pick the optimizers that perform the best in the OpenWebText experiments:
Muon: we find that using Adam for embedding and readout layers performs better than
Muon, as observed by previous works [30 23
Shampoo: we find that Shampoo2 (er, = en = 1/2) with Adam grafting and Adam applied
to the embedding and readout layers performs well. We found one-sided Shampoo
34

---

embedding and readout layers [ to perform significantly worse. We use blocking with block size 512.
SOAP: we use spectrally normalized SOAP since the optimal learning rates are more stable than regular SOAP. We use one-sided SOAP for embedding and readout layers. We use blocking with block size 512.
Hyperparameter tuning. We tune the main hyperparameters of each optimizer to ensure a fair com- parison. We divide hyperparameters into subsets and tune each of hyperparameters sequentially, At each stage of tuning, previously selected hyperparameters remain fixed, except learning rates, which we retunne in every stage. We show the optimal hyperparameter values found the end of each stage in Table g
with links to the full tuning sweeps.
Optimizer
LR
LR mult.
B1
82
Warmup
WD
wandb
Best Loss
4e-3
0.9
0.98
0.2
2e-4
link
3.23635
Adam
4e-3
0.9
0.98
0.2
2e-4
link
3,23635
8e-3
1.6
0.95
3e-3
1.6e-4
link
3.18302
Muon
8e-3
1.6
0.95
3e-3
3.2c-4
link
3.17795
2e-3
32
0.9
0.98
5e-2
2e-4
link
3.18254
Shampoo 2e-3
32
0.95
0.98
5e-2
2e-4
link
3.17796
2e-3
32
0.95
0.98
5e-2
2e-4
link
3.17796
3.2e-2
0.9
0.999
6.25e-3
2c-4
link
3.1751
SOAP
3.2e-2
0.9
0.999
6.25e-3
2e-4
lank
3.1751
Table 4: Tuning sweeps for each optimizer considered in scaling experiments. LR stands for learning rate. LR mult is the multiplier applied for embedding and readout layers when Adam is used, applicable only to hybrid optimizers. 31 and 32 are used differently for different optimizers, as specified in Appendix [A] Warmup is the fraction of tokens used for warmup. WD is the independent weight decay. We share the Weight&Biases view for the sweep in wandb. Bolded values are optimal values found in the sweeps. We ensure the optima are not at the boundary of the sweep grids. We sweep LR, LR mult, WD over powers of two centered around reasonable base values, except $1,2, which are swept over 0.9, 0.95, 0.98, 0.999. Full wandb sweeps available through the links.
H.2 Compute Multiplier Estimation
In Figure 6 we estimate the compute multiplier by dividing the estimated Adam compute achieve the target Toss over the actual compute spent by the optimizer. That is,
CAdam
Compute Multiplier(Copx)
(200)
Copt.
where CAdam is the estimated compute required by Adam to achieve the loss achieved by optimizer opt with Copt FLOPs. We estimate Adam's 's compute using linear interpolation between loss-compute pairs in log-log space. If the loss is smaller than Adam's 's smallest loss, use the closest two points to extrapolate (scipy www.iinoaateeee...mmmm.. interrpi with f wel.valueeeexrrppollloo
The compute is estimated by counting the FLOPs for all operations in transformer. We include the attention dot product FLOPs that scale quadratically with context length, which slightly more accurate than the 6 D formula from 24]. We follow this estimation approach from Levanter
H.3 Near-Optimal Hyperparameter Transfer
We scale the learning rate using μP and the independent weight decay as 1/width as we scale up the model to transfer hyperparameters found on the base model. Here we verify that the scaled learning rates and weight decays are nearly optimal. In Figure I1 we demonstrate how the loss does not improve from our scaled hyperparameters if we change the learning rate, weight decay, or learning rate schedule. We estimate there is 0.1% noise floor in the final loss across runs with the exact same hy perparameesss due to non-determinism, differrnnces below this floor should be ignored. Note,
35

---

though, there is slowly decreasing trend of optimal learning rates, potentially due to the increasing training horizon [7]. This suggests even larger scales our learning rate scaling may break down.
Double LR
Half
Double WD
Half WD
Fix warmup ratio
Adam
Muon
1.0%
1.0%
0.8%
0.8%
Loss increase
Loss increase
0.4%
0.4%
0.2%
0.2%
-0.2%
-0.2%
(tuned)
(scaled)
1024 (scaled)
512 (tuned)
768 (scaled)
1024 (scaled)
Model width
Model width
Shampoo
SOAP
0.8%
0.8%
Loss Increase
0.6%
0.6%
0.2%
0.2%
0.0%
0.0%
-0.2%
(tuned)
(scaled) Model width
1024 (scaled)
512 (tuned)
768 (scaled) 1024 (scaled)
Model width
Figure 11: Ablation of learning rate, weight decay and warmup schedule. We have five perturba- tions to test the optimality of our scaled hyperparameters, namely halving doubling the learning rate or weight decay and using fixed warmup ratio instead fixed number of warmup tokens. Using fixed warmup ratio or tokens doesn't t have significant impact on the final loss. Perturbations on learning rate and weight decay lead to higher losses or differences below the 0.1% noise floor.
H.4 Alternative scaling approach
We expand the ablation experiments in Figure 6 to other optimizers to understand the effects of different components of the scaling rule. In Figure 12, demonstrate that only scaling the learning rate with μP or only scaling the weight decay rapidly diminishes the efficiency gains, showing both important.
H.S Optimal Token-Per-Paramete (TPP)
Matrix-based preconditioned optimizers are more data efficient, so we expect that in the compute- optimal regime, their optimal TPPs will be smaller than Adam. To verify, here estimate the optimal TPP for Muon and Adam. We follow Approach 2 from Hoffmann al. 20 to fix compute budgets and tune model sizes to find the optimal TPPs.
The tuning runs for Adam and Muon can be found in wandb sweeps (Adam and Muon). As we vary TPP while fixing the FLOPss budget, we found that using μP and fixing the independent weight decay yields optimal hyperparameters, consistent with Bergsma et 4 (one needs convert their coupled weight decay to the equivalent independent weight decay to see this agreement). The results are shown in Figure [3 and Table [5] We observe that Adam has a 1.3 times larger optimal TPP than Muon, fairly consistent with the 1.4 times more compute efficiency Muon gets over Adam Figure 6 The agreement between these two values expected from a scaling law of the form (T/Ta)-2 where N is the number of parameters and T is the number of tokens, we assume that only To (data efficiency) can be altered by the optimizer, not the loss for fixed model size given infinite data.
L E (N/N))
36

---

Adam
Muon
Shampoo
SOAP
3.2
3.0
2.8
2.6
Compute multiplier
1.5
1.0
0.5
10
19
10 10"
Compute (FLOPs)
10
10
Figure 12: Alternative scalings for various optimizers underperform. Across various optimizers, μP with 1/ 1/D-scaled weight decay yields the best performance as model size increases. Only scaling the learning rate or the weight decay, while better than scaling neither (SP), leads to worse performance, demonstrating both components are crucial.
(nrease
2.00%
Adam Muan (2:4)
4e- FLOPs
Compute multiplier
1.6
Ze- FLOPS
1.4
1.50%
Sei 19 FLOFS
3.0
Parabola Min
1.00%
2.8
Muon
0.50%
- IP + 1/Wicth WD Optimal
0.8
0.00%
16 32
10
10a0
10:2
1020
Compute (FLOPs)
Compute (FLOPs)
Figure 13: Optimal TPP for Adam and Muon on FineWeh, Optimal TPPs for Adam and Muon are estimated to be 9.6 and 7.4. Adam's optimal TPP is 1.3 times larger than Muon's With the tuned TPP, both Muon and Adam achieve more speedup over the Adam baseline (nP, 1/D weight decay, 20 TPP) shown in solid blue.
Our result shows that the optimal TPP for Adam is 9.6, much smaller than the original estimate of 20 by Hoffmann al. [20]. We believe the explanation two-fold. First, the FineWeb dataset is known to be of higher quality than the C4 dataset used by Hoffmann et al. 20]. which suggests it has number
higher information density (bits per token) and thus for fixed number of parameters of tokens the model can fit is reduced. Second, by extensively tuning the (base) hyperparameters (n) 31, (82, A), the same model converges with fewer tokens, which also reduces the optimal TPP. For example, if training all models with a very small learning rate, the optimal TPP will be very high, scaling as (11/1) as training approaches the gradient flow.
Optimizer
Width
Loss
TPP
Adam
576
3.226039
13.338798
Adam
896
2.996272
11.494327
Adam
1280
2.849975
8.783869
Muon
704
3.157175
6.607409
Muon
1024
2.943154
7.071705
Muon
1344
2.805800
7.325234
Table 5: Optimal widths for the compute budgets of the 190M, 380M, and 640M models.
37

---

SOAP
Muon
Shampoo
3.2
3.0
2.8
μP
2.6
Norm
1.6
Compute multiplier
1.4
1.2
1.0
0.8
10
10
10
10" 20
10"
10
Compute (FLOPs)
Comparing spectrally normalized optimizers with j.P. The setting is identical to Section We compare μPP scaling and spectral normalization while keeping the same 1/width weight decay scaling. In this setting, spectral normalization does not improve over p.P.
Figure 14:
H.6 Spectral Normalization.
We study the scaling performance of the spectrally normalized variants of SOAP, Muon (where spectral normalization only affects the layers using Adam given Muon updates already have unit spectral norm), and Shampoo. In Figure [4, we find the spectrally normalized variants show no advantage μP. This is surprising given that used fixed block size (512) for all three optimizers, enabling good hyperparameter transfer already under juP (Section 3.2p. For Shampoo, we find spectral normalization in fact performed worse than p.P. Figure 15 shows that this due to the optimal learning rate and weight decay shifting toward smaller values for larger models for spectrally normalized Shampoo. We hypothesize this due to the base model being tuned too close to instability, which, combined with the finding that the optimal learning rate decreases with scale on top of pP in the compute-open optimal regime, led to observed suboptimal performance for the larger models.
Double LR
Half LR
Double WD
Half WD
Fix warmup ratio
(norm)
2.5%
3.5%
1.0%
0.0%
0.0%
0.0%
-0.2%
-0.2%
-3.5%
Model
Model width
Model widhh
Figure 15: Ablations for learning rates and weight decay for spectrally normalized optimizers. SOAP and Muon have slightly shifting optimal learning rates and weight decays, partially due to the increasing horizon. Shampoo, however, shows clear shifting toward smaller learning rates and weight decays.
I Hardware
We trained all of our models on TPU-v4 and TPU-v6e, supported by the Google TPU Research Cloud program. WehTeet experiments are trained on TPU-v4-4 and Fen experiments on TPU-v6c-8, TPU-v6e-16 16, and TPU-v4-32.
38

---

Broader Impact and Limitations
Broader Impact. Our work improves the understanding of how to efficiently scale matrix- preconditioned optimization in deep learning, which has the potential reduce the cost of training machine learning models and make machine learning research and workflows more accessible.
Limitations. We perceive two main limitations of our work. First, due to limited computational budget, our experiments are relatively small in scale compared to typical training runs industry. Verifying how well our results generalize to larger models trained with more compute is an excit- ing future direction. Second, while we have shown that optimizers like Shampoo and Muon can significantly improve the compute efficiency of training, we do not investigate why they improve the efficiency. Understanding this question holds potential for designing even more efficient future optimizers.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discussed the limiations of our work in the appendix.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
39

---

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
We provide detailed justifications of the claims in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: In the appendix, we fully described our experiment setup. We also released the code for reproducing the experiments and made experiment logs public.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
40

---

(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to deseribe the particular way they provide for reproducibility. In the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We used public datasets and made code public.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
cablic/guidess httpg/gddde//eeeubiiiieesseeee
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. www.c//wwwpgpcuucclocoeooo http://gdesssseeeee http//oeeeuaiiisiiooooooooooeeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We described experiment details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
41

---

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Due to the expensive computational cost of training language models, we run each experiment only once. Additionally, our conclusions do not rely on comparing small differences in performance of specific experiments, but their overall scaling trends. Running for multiple seeds will not impact the conclusion.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We described the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
42

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impact of this work in the Appendix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), faimesss considerations (e.g., deploymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: The paper poses no risk of misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
43

---

Justification:
We used openly available code and datasets, and follow the license require-
ments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, pappesswittteeee thcode candsatsssss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/moodl as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper did not conduct experiments with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
44

---

Answer:
[NA]
The paper did not conduct crowdsourcing or experiments with human subjects. Guidelines:
Justification:
The answer NA means that the paper does involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification:
We didn't use LLM in a way impacting the core methodology of the research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM for what should or should not be described.
45


Reviews:
[{"review_id": "KfCObGLjAS", "reviewer": "Reviewer_zik3", "summary": "This paper investigates the whether second-order optimizers can outperform standard first-order methods under fixed compute budgets, and how to appropriately scale hyperparameters as model size and batch size grow. They demonstrates that with appropriate hyperparameter scaling and algorithmic adjustments, second-order methods can yield significant speedups over first-order methods when training large neural networks.", "questions": "Could the authors elaborate on why momentum can remain fixed from a theoretical perspective? Is it because the optimal momentum doesn’t depend strongly on batch noise, or simply that any changes to momentum can be compensated by LR adjustments?\n\nAny comments on how one could extend muP to damping factor, block sizes, grafting choices, and Newton-Schultz iterations?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "[Strengths]\nThe paper generalizes prior scaling analyses using the muP framework, to show that the feature learning dynamics remain order-wise constant as width and depth grow even for second-order methods.\n\nThe authors back up their theory with extensive experiments, which verify the width and depth scaling rules by training transformers of various widths and depths.\n\nThey further extend the theory to batch sizes, and show that the advantage of second-order optimizers is their ability to scale up the batch size.\n\n[Weaknesses]\nThe authors show a 2.5x speedup in terms of training tokens needed, but this does not translate to faster wall clock time. The overhead per step is not thoroughly discussed in the paper. Also, for large models the distributed training strategy is critical for minimizing the unique communication overhead of second-order optimizers. The use of stale preconditioners could mitigate this problem, but the paper does not discuss this in detail either.\n\nThere still remain some hyperparameters that are unique to second-order optimizers such as the damping factor, block sizes, grafting choices, and Newton-Schultz iterations, for which the scaling with respect to width, depth, and batch size remain unclear.\n\nThe approach still requires significant compute for the second-order updates and more memory to store preconditioners. The authors use blocking, periodic updates, and efficient factorizations to handle this, but those introduce their own trade-offs.", "was_revised": true, "final_justification": "The authors have addressed my concerns in their rebuttal and followup comments.", "rebuttal": "Thank you for the supportive review! Inspired by your comments, we provide additional results and clarifications here and will include them in the updated paper. We hope you will consider raising your score in light of our response.\n\n**Faster wall-clock time**\n\nWe would like to clarify that we did show in Figure 3 that Muon and Shampoo are significantly more efficient compared to Adam in terms of faster wall-clock time for a given target loss, provided that batch size is scaled together with model size. Based on the results in Figure 3, we estimate the following relative wall-clock time saving relative to Adam for different final losses achieved by Adam:\n\nWall-clock time savings for Openwebtext dataset:\n\n| Target Loss | Muon | Shampoo |\n| --- | --- | --- |\n| 1.38 | 3.61 | 3.22 |\n| 1.20 | 3.19 | 2.78 |\n| 1.03 | 2.73 | 2.31 |\n| 0.94 | 2.46 | 2.05 |\n\n\nWall-clock time savings for Chess dataset:\n\n| Target Loss | Muon | Shampoo |\n| --- | --- | --- |\n| 0.83 | 1.57 | 1.87 |\n| 0.78 | 1.60 | 1.81 |\n| 0.74 | 1.62 | 1.76 |\n| 0.70 | 1.66 | 1.68 |\n| 0.67 | 1.69 | 1.62 |\n| 0.65 | 1.71 | 1.57 |\n| 0.63 | 1.73 | 1.53 |\n\n\n\nAs you pointed out, efficient implementation of second-order optimizers for distributed training, such as distributing the preconditioning operation across devices, is crucial for maximizing their performance gains in practice. As we did not implement distributed versions of second-order optimizers and performed redundant computations across devices, we expect their efficiency gains in wall-clock times can be made even larger than what we estimated, making the case for second-order optimizers even more favorable. Regarding stale preconditioners, we update the preconditioner every 10 steps for Shampoo variants. We will include these details in the updated paper.\n\n**Extending µP to cover hyperparameters specific to 2nd order optimizers**\n\nWe did consider how to scale hyperparameters specific to second-order optimizers with respect to width and depth, including the damping factor in all optimizers, and the learning rate and damping factor for grafting. We applied these scaling rules in our experiments and listed them in Table 1 with derivations in Appendix B. In Appendix B, we also derived how the learning rate should scale with block size under blocking. \n\nRegarding how to scale the number of Newton-Schultz iterations, from the µP perspective, keeping the number of iterations constant is sufficient for a well-defined infinite width limit. However, this is not the only choice, as scaling the number of iterations to infinity also leads to a well-defined limit (converging to the true matrix sign via SVD), though it leads to a much higher overhead. For blocking, both keeping the block size constant and equal to the full matrix size produce well-defined infinite-width limits (Adam-like limit vs full Shampoo limit), and it is unclear which limit is preferable a priori. We agree that exploring scaling rules for these hyperparameters is an interesting future direction, and will include this discussion in the updated paper.\n\n**Efficiency considerations**\n\nWe use techniques such as blocking and periodic updates to both speed up our experiments as well as to demonstrate that our proposed scaling rules work for practical instantiations of second-order optimizers. We view this as a strength of our paper rather than a weakness, as we ultimately care about scaling rules appropriate for practical instantiations of second-order optimizers rather than their idealized implementations, which generally require different scaling rules. \n\n**Theoretical explanation for constant momentum**\n\nWe believe better theoretically understanding how to optimally scale momentum is an exciting and important future direction. Our current hypothesis for why it’s better to keep the momentum constant rather than following the SDE scaling rules is that there is a tension between keeping the first vs second moment of momentum buffers invariant to batch size. Specifically,  as shown in [1], a large value of $\\beta$ is necessary for keeping the variance of the momentum buffers low regardless of the batch size, while a decreasing value of $\\beta$ is required for keeping its first moment (mean) invariant as we scale batch size. We will elaborate on this discussion in the revision.\n\nReferences:\n1. Malladi et al. \"On the SDEs and scaling rules for adaptive gradient algorithms.\" *Advances in Neural Information Processing Systems (Neurips)*, 2022"}, {"review_id": "wypmrBsIVq", "reviewer": "Reviewer_pbPH", "summary": "The paper systematically investigates hyperparameter scaling rules for second-order optimizers (Shampoo, SOAP, Muon) across width, depth, and batch size. The paper shows: (1) under correct scaling, second-order optimizers consistently outperform Adam in compute-optimal settings with up to 2x efficiency gain, and (2) second-order optimizers change the chinchilla scaling law, requiring up to 2x less data compared to the optimal data count. These experiments were conducted using Openwebtext and Chess datasets and models with parameters < 10^8.", "questions": "- Please see my questions and comments from the weakness section above.\n- Better precision is more important for second-order optimization (which involves eigendecomposition or inverse operation). Do you think this could be a potential bottleneck when scaling up second-order methods further?\n- The authors mention that running multiple seeds won't change the conclusions in the paper. Could the author elaborate more?", "limitations": "Yes, they are provided in Appendix H.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n- The paper is well-written, clearly motivated, and well-organized. The authors provide a comprehensive overview of the background (e.g., Shampoo and SOAP) and relevant related works (e.g., Ishikawa & Karakida) in Section 2.\n- The theoretical derivation and justification in the paper look correct and are well-presented.\n- The experiments are convincing that the proposed scaling rule works across different model scales. Figure 1 shows that the standard parameterization (SP) fails to be consistent for second-order optimizers, whereas the proposed approach succeeds. This is more evident in the MLP experiment in Appendix G.\n\nWeaknesses:\n- While the experiments are convincing, at the moment, the paper does not have sufficient details to reproduce the experimental results (e.g., how the hyperparameters were tuned for baseline methods - search space for grid search), which is important for an optimization paper.\n- There are several considerations to make the empirical justification more convincing: (1) How does Adam perform under SP in Figures 1 and 2? This would clarify whether SP specifically fails for second-order methods. (2) What Chinchilla scaling conclusions would we draw using SP and second-order methods / or without grafting? This would help better understand the significance of proper scaling. Addressing these points will make the paper stronger.\n- One might argue that the paper has limited novelty, as it combines muP with existing second-order optimization algorithms. However, I believe that a systematic study like this is useful for the community.\n- The experiments only reach ~10^9 flops, where it is unclear whether the proposed approach/analysis scales to larger settings.\n\nConfidence. I set a confidence score of 3, as I am not familiar with recent works in muP theory and second-order optimization methods.\n\nMinor:\n- AdaGrad and Adam do not approximate the exact Fisher information matrix (but an empirical version).\n- Missing a space after “data” in line 107.\n- Reference to Appendix E does not exist.\n- I believe that the title of the paper is too generic at the moment.", "was_revised": true, "final_justification": "I think this is a good paper and the authors have addressed my questions/concerns in the rebuttal. Specifically:\n- Clarified hyperparameter tuning details, which are crucial for any scaling law and optimization paper\n- Provided a comparison with SP (spectral parameterization)\n- Added error bars to the experimental results\nWhile I've reviewed other reviewers' comments, I find myself somewhat disagreeing with their assessments (especially issues related to formatting and typos), though I acknowledge the validity of certain points (such as concerns about the character-level tokenizer and simplified model architecture). I don't find them sufficiently compelling to warrant lowering my evaluation of the paper's contributions.", "rebuttal": "We really appreciate your supportive review and thoughtful suggestions!\n\n**Hyperparameter tuning details**\n\nWe described our hyperparameter tuning procedure in Appendix F. For each optimizer, we determine the base learning rate by choosing its optimal value on a small model of 3 transformer blocks and an embedding dimension of 128 (Figure 1) and scaling to larger models via our scaling rules. For momentum and batch size, our results in Section 3.4 suggest $\\beta_1=0.9$ and $\\beta_2=0.95$ is near optimal for a wide range of batch sizes. We thus used these values for all other experiments unless stated otherwise.\n\n**Whether SP specifically fails for second-order methods**\n\nSP fails not only for second-order methods, but for Adam as well, as demonstrated extensively in prior works (e.g. [1]). We provide the following additional experiment to demonstrate that SP Adam indeed has a shifting optimal learning rate while µP Adam has a stable optimal learning rate:\n\nLoss under SP\n\n| Width | LR=3e-04 | LR=1e-03 | LR=3e-03 | LR=1e-02 | LR=3e-02 | LR=1e-01 | LR=3e-01 |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| 128 |1.555|1.458|**1.455**|1.457|1.466|1.467|1.47 |\n| 256 |1.379|**1.326**|1.327|1.339|1.345|1.351|1.357 |\n| 512 |1.263|**1.228**|1.257|1.271|1.288|1.306|1.356 |\n| 1024 |**1.165**|1.172|1.184|1.212|1.238|1.293|1.353 |\n| 2048 |**1.112**|1.137|1.166|1.192|1.249|1.334|1.408 |\n\n\nLoss under µP\n\n| Width | LR=3e-04 | LR=1e-03 | LR=3e-03 | LR=1e-02 | LR=3e-02 | LR=1e-01 | LR=3e-01 |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| 128 |1.555|1.458|**1.455**|1.457|1.466|1.467|1.47 |\n| 256 |1.445|1.349|**1.33**|1.333|1.337|1.349|1.355 |\n| 512 |1.371|1.272|**1.233**|1.239|1.261|1.27|1.29 |\n| 1024 |1.326|1.21|**1.167**|1.173|1.189|1.209|1.224 |\n| 2048 |1.299|1.155|**1.118**|1.121|1.149|1.167|1.18 |\n\n\n\n**Chinchilla scaling rule with SP**\n\nWe do not believe finding Chinchilla scaling rules for SP would yield useful insights. Under SP, larger models eventually diverge after the first gradient step [1], trivializing the resulting scaling laws. In the Chinchilla paper [2], they avoided this issue by manually adjusting learning rates across model sizes, presumably tuned via additional experiments, suggesting that directly using SP would lead to worse performance.\n\n\n**Novelty**\n\nWe do believe our work has important novelty. In addition to deriving scaling rules for many second-order optimizers that prior works did not consider (Muon, Grafted Shampoo, SOAP etc.) and demonstrating their empirical effectiveness, our work contains at least two other novel and impactful contributions: 1) a simple procedure for deriving µP and depth scaling for a general second-order optimizer; and 2) the observation that second-order optimizers significantly change the Chinchila scaling laws and that language models should generally be trained with 1.5x~2x fewer tokens per parameter if using Muon or Shampoo for compute-efficient scaling. We believe these findings, while ultimately simple to state, are key to fully realizing the benefits of second-order optimizers.\n\n**Experiment scale**\n\nWe would like to clarify that our experiments are scaled up to 5e18 FLOPs, as shown in Figure 3, much larger than 1e9 FLOPs that you suggested. We acknowledge that validating our results with even larger-scale experiments, such as with billion-parameter models, is an exciting future direction.\n\n**Seed**\n\nThank you for being cautious about uncertainties. We will add error bars to the figures in the camera-ready version for clarity. In our GPT-2 setup, averaging over multiple seeds won't change the conclusions as the seed-to-seed standard deviation in the final loss is small, compared to the difference across learning rates and optimizers. To demonstrate this, here we show the mean final losses and the standard deviations in the brackets for five different random seeds.\n\nLosses for the base width model:\n\n| Optimizer | LR=3e-04 | LR=1e-03 | LR=3e-03 | LR=1e-02 | LR=3e-02 | LR=1e-01 | LR=3e-01 |\n| --- | --- | --- | --- | --- | --- | --- | --- |\n| Adam |1.56 (3.91e-03)|1.46 (2.78e-03)|**1.45** (5.97e-03)|1.46 (6.47e-03)|1.47 (2.12e-03)|1.47 (9.42e-03)|1.47 (7.46e-03) |\n| Shampoo |1.49 (2.15e-03)|1.43 (2.41e-03)|**1.41** (1.45e-03)|1.41 (3.50e-03)|1.41 (2.78e-03)|1.41 (3.29e-03)|1.42 (4.18e-03) |\n| Shampoo$^2$ |2.86 (6.34e-03)|1.73 (4.59e-03)|1.40 (9.64e-04)|**1.40** (1.54e-03)|1.40 (1.91e-03)|1.40 (3.71e-03)|1.41 (3.91e-03) |\n| Muon |3.29 (6.01e-03)|1.71 (5.74e-03)|1.43 (2.18e-03)|1.41 (1.99e-03)|**1.40** (2.04e-03)|1.41 (2.34e-03)|1.41 (1.09e-03) |\n| SOAP |2.50 (2.04e-03)|1.47 (1.23e-03)|**1.43** (8.21e-03)|1.43 (1.10e-03)|1.44 (2.42e-03)|1.46 (2.61e-03)|1.54 (1.06e-02) |\n\n\n\n**Precision**\n\nWe agree that numerical precision is particularly important for second-order optimizers, especially for matrix inversion and eigendecompositions. Exploring scaling rules for second-order optimizers that minimize numerical errors in these key operations, such as via unit-scaled µP [3], is an exciting future direction.\n\n\n\nReferences:\n1. Yang, Greg, et al. \"Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.\" *Advances in Neural Information\nProcessing Systems (NeurIPS)*, 2021.\n2. Hoffmann, Jordan, et al. \"Training compute-optimal large language models.\" *Advances in Neural Information\nProcessing Systems (NeurIPS)*, 2022.\n3. Blake, Charlie, et al. \"u-$\\mu $ P: The Unit-Scaled Maximal Update Parametrization.\" *International Conference on Machine Learning (ICML) Workshop*, 2024."}, {"review_id": "PHgyVozeXI", "reviewer": "Reviewer_xRrA", "summary": "This paper investigates hyperparameter scaling for second-order optimizers (e.g., Shampoo, SOAP, Muon) in deep learning. The authors use the maximal-update parameterization (muP) framework to derive scaling rules for learning rate and other hyperparameters as model width, depth, and batch size increase. Through theoretical analysis and empirical results on MLPs and Transformers, the work demonstrates that these scaling rules enable stable training and effective hyperparameter transfer.", "questions": "1. In Figure 2 (left), the performance improvement from using the proposed muP scaling rules over the SP appears marginal for the joint width-and-depth scaling experiments. Given the clear stability gains shown in Figure 1, why do you think this does not translate to a more substantial final loss reduction?\n\n2. The muP scaling rules are derived by merely analyzing the first gradient step, with the assumption that they hold throughout training by induction. How confident are you that these scalings remain optimal in later stages of training, especially for complicated optimizers like Shampoo where the preconditioner statistics are accumulated over time, updated at each step or not, potentially altering the dynamics away from the initial state?\n\n3. The paper demonstrates that second-order optimizers alter the Chinchilla scaling law, leading to higher sample efficiency. However, these optimizers also have higher computational overhead (e.g., scaling with D) than Adam. How do you see this trade-off taking effect for extremely large models? Is there a model scale where the overhead of preconditioning might negate the benefits of improved sample efficiency, even when scaling the batch size as proposed?", "limitations": "**Limited Scale of Experiments**: The experiments, while insightful and very promising, are conducted at a relatively small scale, espeically for such a study on scaling laws. The use of a character-level tokenizer, where the embedding dimension is larger than the vocabulary size, is a significant departure from the setup of modern large language models. While compute budget is a valid constraint, the small scale of these experiments may limit the generalizability of the conclusions to state-of-the-art language models.\n\n**Simplified Model Architectures**: The analysis and experiments are restricted to simplified MLP architectures and do not consider the effects of standard Transformer components like RMSNorms, embeddings or the interactions with softmax attentions. These components are known to have a substantial impact on optimization dynamics and gradient scaling. Omitting them from the analysis leaves a critical gap in understanding how these scaling rules would apply to practical Transformer models.\n\nI believe the paper would be strengthened if the authors could address the concerns raised in the Weaknesses and Limitations sections. I am open to increasing my score upon a satisfactory author response.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "## Strengths\n**Principled Theoretical Framework**: The paper provides a theoretical contribution by extending the muP framework to analyze how hyperparameters for second-order optimizers should be scaled with model depth and width. This is a natural extension to the existing family of muP-related work.\n\n**Improved Optimization Stability**: The empirical results presented in Figure 1 suggests that the derived scaling rules lead to more stable optimization dynamics. The loss and RMS values of feature updates appears more consistent across different model scales when using the proposed parameterization compared to the standard parameterization (SP).\n\n**Comprehensive Analysis**: The derivations are thorough and cover a wide range of popular and recent second-order optimizers, including Shampoo, SOAP, and Muon. This makes the work a good reference for practitioners and researchers working with these optimization methods.\n\n## Weaknesses\n**Clarity and Focus**: The paper is quite dense and appears to sacrifice clarity for breadth. For instance, crucial experimental details (e.g., model architecture specifics, training setup) are completely omitted from the main body and deferred to the appendix. For a paper focused on scaling laws, these details are important for interpreting the results and should be more prominent.\n\n**Inclusion of Non-Essential Content**: The paper's focus could be significantly sharpened by removing content that does not directly support its core contributions. For example, Equation (2) is presented to introduce Muon but does not appear to be used in any subsequent derivations. Similarly, Section 3.4 line 227 introduces results from a reference (Malladi et al., 2022) but does not present or contrast them with the authors' own results, which can cause confusion.\n\n**Typos**: There are a few distracting typos, for example:\n- In the Abstract (line 10): \"nerual scaling laws\" should be \"neural scaling laws\", incorrect spelling.\n- In Section 2 (line 111): \"computate-optimal\" should be \"compute-optimal\", incorrect spelling.", "was_revised": true, "final_justification": "The authors successfully addressed my main concerns about the experiment settings, providing tables of results using a full vocabulary tokenizer on the Fineweb dataset with larger models. I remain uncertain about the validity of their estimated optimal tokens per parameter, which is a surprising finding, due to lack of further results. Although the paper's contribution is modest, its findings are beneficial and should be shared with the community. For these reasons, I'm raising my score to borderline accept.", "rebuttal": "Thank you for your detailed review and thoughtful suggestions. Inspired by your comments, we provide a range of new experimental results and clarifications.\n\n**On applicability to the transformer architecture**\n\nWe would like to clarify that *all experiments* in the main text are conducted on transformers that include standard components such as RMSNorms, embeddings, and softmax attention. Only the additional experiments in Appendix G use MLPs. We also note that it is standard practice in the µP literature to perform theoretical derivations with MLPs for simplicity [1,2,3], as the underlying principles do not depend on specific architecture details and have been shown to apply to transformers both in theory and in practice [1]. We will clarify this point in the updated paper.\n\n\n**Analysis beyond the first gradient step**\n\nWhile we only explicitly analyzed the µP conditions for the first gradient step, it is a well-established result in the µP literature that the same argument straightforwardly extends to any step $t$ during training via induction, so long as $t$ is constant with respect to model size [1,2,3]. Indeed, as illustrated in the right panel of Figure 1, the difference in RMS (Root Mean Square) of the feature updates remains notably more consistent under µP scaling than standard parameterization (SP) throughout the training process up to $10^4$ steps.\n\n**Overhead of preconditioners at scale**\n\nAs shown in Section 4.3, it’s crucial to scale the batch size appropriately, ideally proportional to the model width (or $P^{1/2}$ where $P$ is the number of parameters), as the model size increases to so that the overhead of second-order optimizers like Muon and Shampoo is only ever a constant fraction of the training compute. In order for this scaling not to negate the benefits of preconditioning, we need the critical batch size to scale at least as fast as $P^{1/2}$. Fortunately, our empirical result in Figure 4 (left) suggests that the critical batch size indeed scales at a rate close to this. Under this scaling, we found both Muon and Shampoo consistently outperform Adam in runtime (Figure 4, right plot). Beyond these encouraging results, whether second-order optimizers truly can outperform Adam at much larger model sizes must be ultimately determined empirically, which we believe is an exciting and important future direction.\n\n**Larger-scale, full vocabulary experiments**\n\nWe appreciate your concerns regarding the use of character-level tokenization and experiment scale. To address this, we conducted both the µP and scaling experiments using a full vocabulary tokenizer on the Fineweb dataset with larger models.\n\nWhile we regret that we cannot include images of the results here, we provide the key outcomes in tabular format below. Notably, the conclusions from the character-level experiments continue to hold under this more realistic setup, reinforcing the robustness and generality of our findings despite the initial small-scale configuration.\n\n\n*Stable optimal learning rate (c.f. Figure 1)*\n\nBased on the setup in Appendix E, this configuration differs by using the full Fineweb vocabulary, processing a total of 100 million tokens, including a 10 million token warm-up phase. The largest model in this setting has 3 transformer blocks with width up to 2048 (280M parameters in total). While only the result of Muon is presented here, we will include all optimizers in the camera-ready version. At this scale, the performance gap between µP and SP becomes more evident. \n\n\nWidth scaling for Muon: our results show that µP maintains a consistent optimal learning rate, reflected as a vertical stripe at zero, while SP demonstrates a decreasing trend in its stable learning rate range.\n\nSP:\n\n| Width | LR=1e-03 | LR=3e-03 | LR=1e-02 | LR=3e-02 | LR=1e-01 | LR=3e-01 | LR=1e+00 | LR=3e+00 | LR=1e+01 |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| 128 |10.69|10.38|9.25|7.79|6.42|5.58|5.22|**5.14**|5.36 |\n| 256 |10.62|10.15|8.56|7.21|5.90|5.19|**4.89**|4.90|5.29 |\n| 512 |10.51|9.81|7.98|6.69|5.46|4.85|**4.66**|4.77|5.29 |\n| 1024 |10.36|9.34|7.47|6.19|5.05|4.60|**4.52**|4.74|5.44 |\n| 2048 |10.13|8.69|6.97|5.71|4.73|**4.44**|4.48|4.80|5.63 |\n\n\nµP:\n\n| Width | LR=1e-03 | LR=3e-03 | LR=1e-02 | LR=3e-02 | LR=1e-01 | LR=3e-01 | LR=1e+00 | LR=3e+00 | LR=1e+01 |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| 128 |10.69|10.38|9.25|7.79|6.42|5.58|5.22|**5.14**|5.36 |\n| 256 |10.68|10.34|9.13|7.58|6.20|5.33|4.94|**4.86**|5.10 |\n| 512 |10.67|10.31|9.05|7.42|6.04|5.13|4.73|**4.68**|4.94 |\n| 1024 |10.66|10.29|8.99|7.30|5.91|4.98|4.57|**4.54**|4.86 |\n| 2048 |10.65|10.27|8.94|7.22|5.83|4.86|4.46|**4.45**|4.78 |\n\n\nAs in Figure 1, we visualize changes in the RMS of the hidden features (∆h). Additionally, we include the RMS change of the logits (∆f). From the table, we show that ∆h and ∆f only remain stable across different widths under µP, highlighting scale-dependent inconsistencies and further underscoring the limitations of SP.\n\n\nFeatures and logits change for Muon:\n\n| Width | $\\Delta f$ (SP) | $\\Delta f$ (µP) | $\\Delta h$ (SP) | $\\Delta h$ (µP) |\n| --- | --- | --- | --- | --- |\n| 128 | 1.98e-05 |1.98e-05 |5.82e-03 |5.82e-03 |\n| 256 | 2.97e-05 |2.12e-05 |6.15e-03 |6.21e-03 |\n| 512 | 4.62e-05 |2.36e-05 |6.53e-03 |6.67e-03 |\n| 1024 | 7.07e-05 |2.57e-05 |6.98e-03 |7.19e-03 |\n| 2048 | 1.08e-04 |2.79e-05 |7.62e-03 |7.89e-03 |\n\n\nDepth scaling for Muon: similar to the conclusion of width scaling, the optimal learning rate starts to shift as the model gets deeper.\n\nSP:\n\n| Depth | LR=3e-01 | LR=1e+00 | LR=3e+00 | LR=1e+01 |\n| --- | --- | --- | --- | --- |\n| 3 |5.577|5.218|**5.143**|5.359 |\n| 6 |5.528|5.153|**5.080**|5.312 |\n| 12 |5.484|5.113|**5.051**|5.313 |\n| 24 |5.423|5.054|**5.022**|5.327 |\n| 48 |5.390|5.024|**5.014**|5.359 |\n| 96 |5.358|**5.012**|5.013|5.428 |\n| 192 |5.345|**5.015**|5.033|5.586 |\n\nDepth-Scaled:\n\n| Depth | LR=3e-01 | LR=1e+00 | LR=3e+00 | LR=1e+01 |\n| --- | --- | --- | --- | --- |\n| 3 |5.577|5.218|**5.143**|5.359 |\n| 6 |5.535|5.159|**5.059**|5.318 |\n| 12 |5.483|5.120|**5.049**|5.306 |\n| 24 |5.425|5.063|**5.004**|5.291 |\n| 48 |5.384|5.019|**4.971**|5.315 |\n| 96 |5.349|4.987|**4.972**|5.330 |\n| 192 |5.324|4.971|**4.950**|5.368 |\n\n\n*Compute-optimal scaling (c.f. Figure 3)*\n\nFurthermore, we conduct full vocabulary experiments on the Fineweb dataset in the compute-optimal setting. Due to time constraints, we only compare Adam with Muon, but will include additional optimizers in the final paper. In the following tables, we find that our observations that second-order optimizers 1) reduce Chinchilla-optimal tokens per parameter by about 2x, and 2) lead to about 1.5x - 2x compute efficiency gain relative to Adam continue to hold in this more standard experiment setup. Here, we also extended the model size up to 200M parameters. The optimal token per parameter is estimated using the same fit procedure in Section 4.1.\n\n| Optimizer | Estimated optimal token per parameter |\n| --- | --- |\n| Adam | 11.2 |\n| Muon | 5.4 |\n\n\n| Adam loss | 5.477 | 5.121 | 4.935 | 4.744 | 4.544 | 4.342 | 4.160 | 4.020 | 3.908 | 3.804 |\n|------|------|------|------|------|------|------|------|------|------|------|\n| Adam compute | 1.128e+16 | 2.635e+16 | 4.931e+16 | 8.212e+16 | 1.275e+17 | 2.712e+17 | 5.226e+17 | 9.465e+17 | 1.639e+18 | 2.739e+18 |\n\n\n| Muon loss | 5.889 | 5.382 | 5.041 | 4.785 | 4.605 | 4.329 | 4.136 | 3.980 | 3.852 | 3.747 |\n|------|------|------|------|------|------|------|------|------|------|------|\n| Muon compute | 5.534e+15 | 1.293e+16 | 2.420e+16 | 4.030e+16 | 6.257e+16 | 1.331e+17 | 2.565e+17 | 4.646e+17 | 8.048e+17 | 1.345e+18 |\n\n\n\n| Target loss | Estimated Muon compute efficiency gain |\n|-------------|--------------------------------|\n| 4.660 | 1.81× |\n| 4.261 | 2.17× |\n| 3.966 | 2.55× |\n\n\n\n**Paper organization**\nYour point regarding the focus of the presentation is well-taken. To improve readability, we have deferred the experimental details to the appendix so the main text can focus on a complete, self-contained guide for practitioners interested in scaling second-order optimizers with µP. Due to space constraints, we deferred the detailed derivation for Muon scaling rules based on its update rule in Appendix B. Similarly, we apply and extend the SDE formulation by Malladi et al. for batch size scaling in Appendix D and compare its performance against our alternative, empirically determined scaling rule in Figure 2. We will use the additional page in the camera-ready version to clarify these details and findings.\n\n\nThanks again for your review. We believe the content of this response, which required a significant effort, will strengthen the paper. We will additionally correct all formatting and typos. We would really appreciate it if you would consider raising your score in light of our response. \n\n\nReferences\n1. Yang et al. \"Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.\" *Advances in Neural Information\nProcessing Systems (NeurIPS)*, 2021.\n2. Yang and Littwin. \"Tensor programs ivb: Adaptive optimization in the infinite-width limit.\" *International Conference on Learning Representations*, 2023\n3. Everett et al. \"Scaling exponents across parameterizations and optimizers.\" *arXiv preprint arXiv:2407.05872*, 2024\n4. Liu, Jingyuan, et al. \"Muon is scalable for LLM training.\" *arXiv preprint arXiv:2502.16982*, 2025."}, {"review_id": "nmN8CU5BfH", "reviewer": "Reviewer_K9rj", "summary": "The paper builds on recent empirical success of some second-order optimizers to motivate deriving scaling laws in a manner that was previously only done for first-order optimizers ($\\mu$P). They show that under the correct scaling that second-order optimizers outperform Adam across scales and require 2 x less data per parameter for compute optimality.", "questions": "In addition to responding to my comments above I have a few questions/slight comments:\n* Could the authors provide more clarity on the difference between $\\zeta$ and $\\zeta’$ (and the corresponding $x$ and $x’$)?\n* Line 155, $L$ for the loss is the same symbol as the number of layers $L$.\n* In Figure 1 what was the fixed depth that was chosen when varying the width (and what was the fixed width when varying the depth). Does this make any difference?\n* In Figure 3, what is the meaning of the circular markers compared to the dashed lines?", "limitations": "The authors do a good job at covering the paper’s limitations.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strengths\n* The motivation and quality of the writing of the paper is a strength. It is clear, given the recent successes of Muon and Shampoo (etc.) why it would be important to look into the scaling laws of these optimizers.\n* The generality of the scaling laws derived in the paper are nice. It would be “easy” to focus on just one or two optimizers, but the authors focus on a more general “bra ket” framework that applies to all optimizers as shown in the appendix.\n* The experimental results show promise and validate the theory to a certain degree. Figure 1 for example shows especially for width that the derived $\\mu$P for the second-order optimizers behaves more consistently across widths. The same is less obvious for depth, but combined with Figure 2, it does seem the scaling applies empirically to the depth as well.\n* Another strength is how the paper tries to provide insights to the reader. There are plenty of take-aways from the paper in terms of how to scale hyperparameters and what to do with batch sizes. This is useful.\n* I tried my best to read the relevant literature and it looks like this work is original enough to be an addition. I could not see other papers deriving scaling laws for second-order optimizers. The focus in the literature looks to have been on either first-order optimizers and/or “why existing second-order optimizers work”. The latter does touch on norms, but not to the extent of this paper.\n* The introduction of the relevant literature, as well as including the equations for SOAP and Shampoo, was a helpful reminder. Furthermore, I was glad to see the simple derivation of SGD in the appendix. \n### Weaknesses\n* The notation was a bit dense and at times that made it hard to understand. It might be the case that the notation is standard within the optimization of NN community, but it did make it a challenge to read the paper. \n* The derivation for batch size scaling has less empirical data to back it up, and the authors themselves highlight the limitation of the use of the SDE for scaling hyperparameters across large ranges of batch sizes. It feels a little bit like this part of the paper was less developed and was included more for completeness. (Although it is a positive that the authors acknowledged the limitations of this part.)\n* In the section on computing the efficiency gain, and the corresponding Figure 3, it is not clear whether the FLOPS include the additional cost of the second-order optimizers. The paper refers to FLOPS as $6tP$. This looks like 6 x number of tokens x number of model parameters, and would therefore not capture the cost of the additional second-order optimizers.  In the next section at line 289, the paper writes that the compute overhead of Muon and Shampoo are negligible but it would be helpful to demonstrate this. Furthermore, it is mentioned that in scaling, the compute overhead would no longer be negligible. This is a concern when linking back to the conclusion that second-order optimizers achieve a 2x efficiency gain. The paper does say a result is that there is a need to keep control the batch size. However, this makes it more challenging to compare with first-order optimizers, which do not have this need for controlling the batch size in the same way.", "was_revised": true, "final_justification": "Thanks for the response. I will keep my score as accept.", "rebuttal": "Thank you for your supportive review and constructive feedback! We appreciate that you recognize the comprehensiveness of our scaling rules analysis and its demonstrated practical benefits. We respond to your comments and questions below.\n\n**Theory for batch size scaling**\n\nThank you for recognizing our honesty in stating the limitations of our batch size scaling analysis as a strength. We believe improving our understanding of batch size scaling, both theoretically and empirically, is an exciting next step, and hope our work highlights its importance.\n\n**Overhead of Muon and Shampoo**\n\nYou are correct that we did not report the overhead of second-order optimizers in terms of FLOPs. This was due to the complexity involved in quantifying the exact number of operations required by routines such as eigendecomposition and matrix roots, which are highly dependent on their implementations. Furthermore, these operations tend to be more sequential in nature and less compute-bound compared to matrix multiplication in the forward and backward pass, making FLOPs no longer an accurate indicator of real runtime overhead. Therefore, we instead directly report wall-clock time in Figure 4 for a more realistic estimate of efficiency gains. Figure 4 shows that 1) the preconditioner overhead indeed dramatically reduces the wall-clock speedup for Shampoo as model size increases if batch size stays constant, and 2) by scaling batch size linearly with model width, both Muon and Shampoo consistency outperform Adam in runtime. To see the concrete runtime savings, please refer to our response subtitled \"Faster wall-clock time\" to Reviewer zik3.\n\nWe therefore conclude that second-order optimizers can consistently outperform Adam provided we sufficiently co-scale batch size with model size. While we agree first-order methods such as Adam do not require large batch size in principle, practitioners often train at the critical batch size (the maximum batch size that does not significantly increase the loss for a given token budget) purely for the benefit of parallelization [1,2]. We will clarify these takeaways in the update paper.\n\n**Response to Questions**\n\n- $\\xi$ refers to the sample that we take the gradient step on to update the weights, while $\\xi’$ refers to an arbitrary point in the input space. As defined in Line 166, $x$ is the shorthand for $x(\\xi)$ and $x’$ is the shorthand for $x(\\xi’)$, where $x$ denotes the features for a sample in a particular hidden layer. The arbitrariness of $\\xi’$ means the µP conditions should be satisfied for any point in the domain. We will make this clearer in the next version.\n\n- As described in Appendix F, the base model configuration is depth 3 and width 128. We did not experiment with other values. We chose these small values to demonstrate optimal learning rate is stable across a wide range of widths and depths.\n\n- In the plot in question, the circular markers represent individual data points, while the dashed line corresponds to the linear regression line in log-log space (power-law fit). We’ll revise the figure caption to make this clearer in the final version.\n\nReferences:\n1. Chowdhery, A. et al. “PaLM: Scaling Language Modeling with Pathways.” *Journal of Machine Learning Research*, 2023\n2. Meta AI Llama Team. “The Llama 3 Herd of Models.” *arXiv preprint arXiv:2407.21783*, 2024"}, {"review_id": "M2TmSHzOp6", "reviewer": "Reviewer_6YHN", "summary": "This paper studies the scaling laws of loss, data, compute, when optimising neural networks with second-order optimisers, and it derives novel rules for scaling hyperparameters.\nPrevious work had already addressed these two topics, the novelty of this work is mainly the following: for a given compute budget, it shows that less data should be used to approach the Pareto frontier of loss vs compute when using second-order optimisers. Scaling laws for hyperparameters are derived theoretically and tested empirically for more optimisers and more hyperparameters than previously done.", "questions": "NA", "limitations": "yes", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\n- It considers a good number of optimisers: Shampoo, Shampoo2, Muon, Soap.\n- It points out a possible mistake done in a previous study that did not achieve a Pareto frontier.\n\nWeaknesses:\nA few sections of the paper are not written clearly at all, and a section is missing. For example: \n- Section 3.1 does not explain how important is the assumption of random weights in deriving the results. \n- Section 3.4 refers to and describes results of a figure and a table that are not present in the main text. It seems that those were literally moved in the Appendix before submission because they did not fit in the page limit. \n- All links between main text and Appendix are broken (there was no need to split them by the way). \n- Appendix E is empty, it says “Blank”, as if the Appendix was not even checked before submission. \n- There is no explanation about how the “target loss” is picked in Figure 3.\n- Exponents and constants of the scaling laws in Figure 3 are not reported. So there is no answer to the important question given in the main text: “For any second-order optimizer that outperforms Adam… must improve, by either decreasing the constant t0 or increasing the exponent.”\n- KFAC is studied in the Appendix but never mentioned anywhere else.\n \nOther major issues:\n- The effect of width and depth scaling on second order optimisers on Transformer is very little, for some optimisers there is no effect at all. \n- I disagree with the statement “the transformer architecture is overall highly robust suboptimal learning rates.” There may be other reasons, for example that theory is derived for MLP, not Transformer, and also that second-order optimisers are often more robust to hyperparameters than first order ones.\n- Theoretical results for batch size scaling seem to be not very useful, in fact an empirical scaling law was eventually used in experiments.\n- Very small models are used for studying scaling rules of loss, data, compute, with character-level tokenisation that is very non-standard in language modelling.\n- Scaling rules usually follow asymptotic arguments, meaning that the variable that is supposed to scale tends to infinity. However, that does not seem to be the case for the aspect ratio A in Table 1.", "was_revised": true, "final_justification": "The authors addressed most of my concerns, I increased the score, however I still do not give an accept rating because the quality of the first submission was very low, so I would have to trust that they improve the final paper by a substantial amount.", "rebuttal": "Thank you for your thoughtful feedback. Inspired by your comments, we provide additional results below, as well as several clarifications. \n\n**On the novelty of this work**\n\nWe would like to clarify that the novelty of this work goes beyond identifying that second-order optimizers have different Chinchilla scaling laws. As you pointed out, we derived hyperparameter scaling rules for a number of second-order optimizers that prior works did not consider. The only existing work studying µP scaling for second-order optimizers is [1], which only addresses width scaling and only covers K-FAC and Shampoo. By contrast, we presented a different and more general approach, avoiding the push-through identity used in [1], which does not apply when $\\xi’\\neq\\xi$ in Equation 8 of our paper. Our approach allows us to derive the scaling rules for a much wider range of optimizers, including Muon, Shampoo with blocking, Shampoo with grafting, and SOAP, and easily handles both width and depth scaling. Our results have high practical relevance as 1) these optimizers have been shown to significantly outperform K-FAC and vanilla Shampoo considered in [1], and 2) joint width and depth scaling is more efficient than width-only scaling at scale as demonstrated in [2,3]. \n\n**Testing our scaling rules in larger-scale, full vocabulary experiments**\n\nWhile Figure 1 (left) shows that the scaling rules had a relatively small impact on transformers, the lack of consistent feature learning under SP in Figure (right) and the growing gap between SP and our scaling rules in Figure 2 (left) strongly suggest that the benefit of correct scaling rules will be more pronounced at larger scales. To demonstrate that, we ran experiments with even larger models using the full vocabulary version of higher quality Fineweb dataset, showing that the distinction between SP and our scaling rules indeed become more pronounced at larger scales.\n\n*Compute-optimal scaling (c.f. Figure 3)*\n\nFurthermore, we conduct full vocabulary experiments on the Fineweb dataset in the compute-optimal setting. Due to time constraints, we only compare Adam with Muon, but will include additional optimizers in the final paper. In the following tables, we find that our observations that second-order optimizers 1) reduce Chinchilla-optimal tokens per parameter by about 2x, and 2) lead to about 1.5x - 2x compute efficiency gain relative to Adam continue to hold in this more standard experiment setup. Here, we also extended the model size up to 200M parameters. The optimal token per parameter is estimated using the same fit procedure in Section 4.1.\n\n| Optimizer | Estimated optimal token per parameter |\n| --- | --- |\n| Adam | 11.2 |\n| Muon | 5.4 |\n\n| Adam loss | 5.477 | 5.121 | 4.935 | 4.744 | 4.544 | 4.342 | 4.160 | 4.020 | 3.908 | 3.804 |\n|------|------|------|------|------|------|------|------|------|------|------|\n| Adam compute | 1.128e+16 | 2.635e+16 | 4.931e+16 | 8.212e+16 | 1.275e+17 | 2.712e+17 | 5.226e+17 | 9.465e+17 | 1.639e+18 | 2.739e+18 |\n\n| Muon loss | 5.889 | 5.382 | 5.041 | 4.785 | 4.605 | 4.329 | 4.136 | 3.980 | 3.852 | 3.747 |\n|------|------|------|------|------|------|------|------|------|------|------|\n| Muon compute | 5.534e+15 | 1.293e+16 | 2.420e+16 | 4.030e+16 | 6.257e+16 | 1.331e+17 | 2.565e+17 | 4.646e+17 | 8.048e+17 | 1.345e+18 |\n\n| Target loss | Estimated Muon compute efficiency gain |\n|-------------|--------------------------------|\n| 4.660 | 1.81× |\n| 4.261 | 2.17× |\n| 3.966 | 2.55× |\n\nPlease also refer to our response to Reviewer xRrA subtitled \"Larger-scale, full vocabulary experiments\" to see our experiments regarding stable optimal learning rate.\n\n**Transformer’s robustness to learning rate**\n\nOur statement that the transformer architecture is overall highly robust to suboptimal learning rates is directly based on empirical evidence in Figure 1 (left), where the impact of the correct scaling rules is relatively small, as you also pointed out. Furthermore, our experiments in Figure 6 Appendix G show that MLPs are much less robust to suboptimal learning rates, unlike transformers, even with second-order optimizers, directly showing that the transformer architecture indeed leads to improved robustness.\n\nDespite being derived from MLPs, our scaling rules apply to transformers as the underlying principles of µP do not depend on specific architecture details and have been shown to apply to transformers both in theory and in practice [3]. We note that it is standard practice in the µP literature to perform theoretical derivations with MLPs for simplicity [3,4,5]. We will clarify this in the next version. Finally, the consistent feature learning strength across scales in Figure 1 (right) provides a direct empirical verification that our scaling rules are correct. \n\n**Significance of Random Initialization**\n\nWe follow prior works on µP in assuming the weights are randomly initialized. This choice both reflects practice and allows one to prove existence of well-defined infinite limits using the central limit theorem and the law of large numbers [4,5,6] which underlie our derivation for the scale of the gradients, activations, and their updates in Section 3.1. We note that recent work shows that the random initialization assumption is not necessary, and can be replaced by conditions on the spectral norms of the weights at initialization [6].\n\n**Why asymptotic arguments apply**\n\nWhile the aspect ratio $A$ does not grow to infinity, asymptotic arguments in µP only require the width (or depth in the case of depth-scaling) to grow to infinity in order to appeal to the central limit theorem and the law of large numbers [4,5].\n\n**Choice of target Loss in Figure 3**\nTo highlight the efficiency multipliers across different target losses, we extend results in Figure 3 into a more detailed table. Naive refers to using the same data scaling as is optimal for Adam, while ours indicate using corrected scaling we suggested that reduce the number of tokens per parameter for Muon and Shampoo.\n\nWall-clock speed ups relative to Adam on Openwebtext:\n\n| Target Loss | Adam| Muon (naive) | Shampoo(naive) | Muon (ours) | Shampoo (ours) |\n| --- | --- | --- | --- | --- | --- |\n| 1.38 | 1.00 | 3.06 | 2.69 | 3.61 | 3.22 |\n| 1.20 | 1.00 | 2.69 | 2.35 | 3.19 | 2.78 |\n| 1.03 | 1.00 | 2.29 | 1.99 | 2.73 | 2.31 |\n| 0.94 | 1.00 | 2.05 | 1.78 | 2.46 | 2.05 |\n\nWall-clock speed ups relative to Adam on Chess:\n\n| Target Loss | Adam| Muon (naive) | Shampoo (naive) | Muon (scaling) |Shampoo (scaling) |\n| --- | --- | --- | --- | --- | --- |\n| 0.83 | 1.00 | 1.42 | 1.71 | 1.57 | 1.87 |\n| 0.78 | 1.00 | 1.45 | 1.67 | 1.60 | 1.81 |\n| 0.74 | 1.00 | 1.48 | 1.63 | 1.62 | 1.76 |\n| 0.70 | 1.00 | 1.52 | 1.58 | 1.66 | 1.68 |\n\n**Exponents and Constants of Figure 3**\n\nIn Figure 3, we fit a power law for loss\n\n$$ L(C) = \\left( \\frac{C}{C_0} \\right) ^{-b} + L_0,$$\n\nwhere $C$ is the amount of training compute. We fix $L_0$ across the three optimizers so that the values of the other parameters ($C_0$ and $b$) are meaningfully comparable. We then show the fitted coefficients:\n\nOpenwebtext:\n\n| Optimizer | $L_0$ | $C_0$ | $b$ |\n| --- | --- | --- | --- |\n| Adam | 0.578 | 0.374 | 0.105 |\n| Muon | 0.578 | 0.092 | 0.100 |\n| Shampoo | 0.578 | 0.102 | 0.099 |\n\nChess:\n\n| Optimizer | $L_0$ | $C_0$ | $b$ |\n| --- | --- | --- | --- |\n| Adam | 0.538 | 7.252e-04 | 0.149 |\n| Muon | 0.538 | 5.118e-04 | 0.151 |\n| Shampoo | 0.538 | 3.099e-04 | 0.145 |\n\nWe see that only the constant factor improved significantly (by ~2x) while the exponent changed less than 10%. This suggests that second-order optimizers primarily change the constant but not the scaling exponent. We will discuss this important finding in the updated version.\n\n**Why our batch size scaling experiments are valuable**\n\nOur results on batch size scaling are valuable as they reveal that existing theoretical models for batch size scaling based on SDEs do not allow effective hyperparameter transfer in practice. In the paper, we fully acknowledge that a better theoretical understanding of batch size scaling is needed. We believe our thorough evaluation and honest presentation should be viewed as a strength rather than a weakness, as acknowledged by Reviewer K9rj.\n\n**Improving paper organization**\n\nThank you for your suggestions in improving the organization of our paper. We separately submitted the main paper and the appendix following the NeurIPS submission guidelines, which prevented the link from working properly. Per your suggestion, we will include the figures referenced in Section 3.4 to the main text using the additional page in the camera-ready revision to improve readability. We deferred those figures to the appendix due to space constraints. We apologize for the formatting error that led to an extra section in the appendix (section E) and will correct it in the revision. No content is intended for section E. \n\nReferences:\n1. Ishikawa and Karakida. \"On the Parameterization of Second-Order Optimization Effective towards the Infinite Width.\" *International Conference on Learning Representations (ICLR)*, 2024\n2. Kaplan, Jared, et al. \"Scaling laws for neural language models.\" *arXiv preprint arXiv:2001.08361*, 2020.\n3. Dey, Nolan, et al. \"Don't be lazy: CompleteP enables compute-efficient deep transformers.\" *arXiv preprint arXiv:2505.01618*, 2025.\n4. Yang et al. \"Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer.\" *Advances in Neural Information\nProcessing Systems (NeurIPS)*, 2021.\n5. Yang and Littwin. \"Tensor programs ivb: Adaptive optimization in the infinite-width limit.\" *International Conference on Learning Representations*, 2023\n6. Yang, Greg, James B. Simon, and Jeremy Bernstein. \"A spectral condition for feature learning.\" *arXiv preprint arXiv:2310.17813*, 2023.\n7. Hoffmann, Jordan, et al. \"Training compute-optimal large language models.\" *Advances in Neural Information\nProcessing Systems (NeurIPS)*, 2022."}]

Author rebuttal:

Meta-review: This paper develops hyperparameter scaling rules for second-order optimizaters such as K-FAC, muon etc with varying width, depth, and batch size. The paper also presents a systematic study of whether second-order optimizers can achieve efficiency gains over first order optimizers under compute-optimal settings. Overall the paper finds good (though not, overwhelming) evidence that second order optimizers could indeed lead to efficiency improvements. They demonstrate token efficiency improvements, improved Chichilla token-parameter scaling, and significant wall-time improvement because of larger batch scaling.
The major weakness are 1) the lack of a distributed implementation and evaluations in a distributed setting, and 2) demonstrating convincing flop-efficiency improvements when taking the optimizers into account.

I think this paper thoroughly investigates a timely question of whether second-order optimizers can really help us with training efficiency. It provides a general derivation of scaling laws for second-order optimizers that is useful, and the empirical evaluations make a compelling case overall.

Overall, most of the reviewers were pretty supportive of the paper from the onset, but sought several clarifications on the nuances. The authors presented comprehensive answers and included results from extensive computational experiments for certain queries. The reviewers seemed to be overall happy (other than reviewer 6yhn) with the responses and some even raised their scores.
Decision: Accept (poster)