Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
Authors: Runa Eschenhagen
Aaron Defazio
Tsung-Hsien Lee
Richard E. Turner
Hao-Jun Michael Shi
Venue: NeurIPS 2025 spotlight ()
Keywords: neural network training algorithms
optimization
Shampoo
Kronecker-factored
preconditioner
grafting
approximation
adaptive gradient methods
stochastic optimization
Primary area: deep_learning
TL;DR: We investigate the role of learning rate grafting and the staleness of the preconditioner in Shampoo by decoupling the updates of the eigenvalues and eigenbasis of its preconditioner.

Abstract: The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks. Despite its success, Shampoo relies heavily on several heuristics such as learning rate grafting and stale preconditioning to achieve performance at-scale. These heuristics increase algorithmic complexity, necessitate further hyperparameter tuning, and lack theoretical justification. This paper investigates these heuristics from the angle of Frobenius norm approximation to full-matrix Adam and decouples the preconditioner's eigenvalues and eigenbasis updates. We show that grafting from Adam mitigates the staleness and mis-scaling of the preconditioner's *eigenvalues* and how correcting the eigenvalues directly eliminates the need for learning rate grafting. To manage the error induced by infrequent *eigenbasis* computations, we propose an adaptive criterion for determining the eigenbasis computation frequency motivated by terminating a warm-started QR algorithm. This criterion decouples the update frequency of different preconditioner matrices and enables us to investigate the impact of approximation error on convergence. These practical techniques offer a principled angle towards removing Shampoo's heuristics and developing improved Kronecker-factorization-based training algorithms.

Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
Runa Eschenhagen Aaron Defazio
Tsung-Hsien Leef
Richard E. Turner Hao-Jun Michael Shi
'Dpartment of Engineering, University of Cambridge Fundamental AI Research, Meta Superintelligence Labs, Meta Platforms, Inc. 'The Alan Turing Institute Infrastructure Optimizations, Meta Superintelligence Labs, Meta Platforms, Inc.
Abstract
The recent success of Shampoo in the AlgoPerf contest has sparked renewed inter- optimization algorithms for training neural
est in networks. Despite its success, Shampoo relies heavily on several heuristics such as learning rate grafting and stale preconditioning achieve performance at-scale. These heuristics increase algorithmic complexity, necessitate further hyperparame- ter tuning, and lack theoretical justification. This paper investigates these heuristics from the angle of Frobenius norm approximation to full-matrix Adam and decou- ples the preconditioner's eigenvalues and eigenbasis updates. We show that grafting from Adam mitigates the staleness and mis-scaling of the preconditioner's eigen- values and how correcting the eigenvalues directly eliminates the need for learning rate grafting. To manage the error induced by infrequent eigenbasis computations, we propose an adaptive criterion for determining the eigenbasis computation fre- quency motivated by terminating a warm-started QR algorithm. This criterion decouples the update frequency of different preconditioner matrices and enables us to investigate the impact of approximation error on convergence. These practical techniques offer a principled angle towards removing Shampoo's 's heuristics and based training algorithms.
developing improved
Introduction
Structured non-diagonal, and especially Kronecker-factored, preconditioned stochastic gradient al- gorithms have been extensively studied for neural network training (Heskes, 2000; Martens, 2010; Martens & Grosse, 2015; Li, 2018; Gupta al., 2018). Despite their promise, diagonally precondi- tioned methods like Adam have remained the de facto methods for training neural networks over the past decade (Duchi et al., 201 I; Kingma & Ba, 2015). Recently, a distributed implementation of the Shampoo algorithm (Gupta et al., 2018; Anil et al., 2020; Shi et al., 2023) won the external tuning track of the AlgoPerf neural network training algorithm competition (Dahl et al., 2023; Kasimbeg al., 2025). This result has renewed interest in non-diagonally preconditioned training algorithms, inspiring methods like Muon (Jordan et al., 2024; Bernstein, 2025) and SOAP (Vyas et al., 2025a).
Correspondence to: re3936cam. ac. uk and hjmshi @meta. com.
WWor performed while an intern and external research collaborator at FAIR, Meta Platforms. Work performed while employed by Al and Systems Co-Design, Meta Platforms.
'The external tuning track requires a submission to specify hyperparameter search space, in contrast to the hyperparaameeer-tuinn--fff self-tuning track.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Vision transformer
ConvNeXt V2
le
Train loss (log scale)
10
le
10
3e
le
10
10
le
AdamW
Shampoo with grafting
50
100
50
100
Epoch
Epoch
Figure 1: Shampoo with stale preconditioner (updating the root inverse matrices every = 100 steps) without grafting for different choices of the learning rate and on Imagewoof. All tested hy- perparameter combinations underperform AdamW and, by extension, Shampoo with Adam grafting.
However, the winning Shampoo submission to the AlgoPerf competition relied on several crucial heuristies beyond Shampoo (Anil al., 2020; Shi et al., 2023). Most notably, each layer's update is re- re-scdled by the update magnitude of reference optimizer, a technique known as learning rate grafting (Agarwal et al., 2020); see Figure for an illustration of its importance. Additionally, to reduce its computational overhead, the root-inverse of the preconditioner is only re-computed every 100 steps, resulting the use of a stale preconditioner. Despite their empirical effectiveness, both heuristics lack theoretical justification and remain poorly understood.
In this paper, we investigate the role of these two heuristics by decoupling the updates of the preconditioner's cigenvalues and eigenbasis. In Section 3, we empirically demonstrate that Shampoo requires learning rate grafting in order address the staleness and mis-scaling of the preconditioner's eigenvalues. Correcting Shampoo's 's eigenvalues at every step like in SOAP (Vyas et al., 2025a) removes the need for grafting. We further formalize this intuition by comparing bounds on the update magnitude of Shampoo and full-matrix Adam.
Given the importance of controlling the approximation error of the eigenvalues, we consider the frequency of the eigenbasis updates Section 4. Motivated by a termination criterion for the QR algorithm that bounds the relative error of the Kronecker factor approximation induced by the current cigenbasis (Golub & Van Loan, 2013), we propose an adaptive method for determining the eigenbasis update frequency. Our empirical results show that the approximation error of the Kronecker factors evolves during training and impacts convergence, depending on both the training stage and parameter's properties. Using an adaptive update frequency can improve Shampoo's training efficiency, especially when more frequent eigenbasis computations accelerate convergence.
Background
We consider neural network training as standard stochastic optimization problem, where the goal to minimize the expected loss function
min - BERd
(() [((Fo(z)) y)],
where fo R + Re the neural network prediction function with parameters E Rd (flattened and concatenated into vector), (a( y) is the joint data distribution from which inputs R and targets E Rº are sampled, and é Re x R + IR is the loss function. The neural network is commonly trained using preconditioned stochastic gradient methods that update the parameters at each iteration t by
o.+1 = - 0, - Q
(2)
where gi = Vo. (ffo (x:), y:) I2 is the stochastic (mini-batch) gradient with respect to the sample (x,, yt) ~ (2( y), a; > is the step size, > 0 is the exponent, and C, E ddxd symmetric scaling matrix. The second equality in Equation (2) expresses the
positive-defiite preconditioner update using the eigendecomposition of the preconditioner matrix C, = Qc, Ac, QZ, where the eigenbasis matrix Q: dxd is orthogonal and eigenvalue matrix Ac, gdxd is diagonal.

---

We primarily focus on C₁ as an accumulation of gradient outer products, i.e., A, = E 9.g. or A, = A- + (1 32) with = 1/2, which correspond to full-matrix AdaGrad and RMSprop/Adam, respectively (Duchi et al., 2011; Kingma & Ba, 2015), although this can be generalized to other alternatives, like the natural gradient method (Amari, 1998); see Appendix A. When C, diagonal, this scaling recovers simplified versions of popular algorithms like AdaGrad (Duchi et al., 201 I) or Adam (Kingma& & Ba, 2015), ignoring additional modifications like exponential moving averages, momentum, bias corrections, etc. However, for non con-diagonal C₁, these methods require storing and possibly inverting dense d x d matrices, which is computationally prohibitive.
To address this, two structured approximations approximation, where each block captures pairwise correlations within each layer; (2) a Kronecker- factored approximation that takes mn mn matrix block for each layer and approximates it as a Kronecker product of two smaller matrices of size m X m and n X The Kronecker product approximation particularly convenient for computing the inverse-root matrix-vector product, leading to the design of algorithms like Kronecker- Factored Approximate Curvature (K-FAC) (Heskes, 2000; Martens & Grosse, 2015; Grosse & Martens, 2016; 6; Martens et al., 2018; Eschenhagen et al., 2023), Shampoo (Gupta et al., 2018; Anil et al., 2020: Shi et al., 2023), and their variants (George et al., 2018; Gao et al., 2020: Ren & Goldfarb, 2021: Feinberg et al., 2023; Duvvuri et al., 2024; Lin et al., 2024b,a), Alternative approaches, such as Hessian- frree or inexact Newton methods, avoid explicitly storing the preconditioner matrix, e.g., by leveraging Hessian-vector products (Martens, 2010; Li, 2018; Pooladzandi & Li, 2024).
commonly applied: (1) layer-wise block-diagonal
2.1 Shampoo
The Shampoo preconditioner was originally developed as a Kronecker-factorized upper bound to full-matrix AdaGrad in its regret analysis (Gupta et al., 2018). For a weight matrix W, E R with stochastic gradient G; mm= where g, vec G.). Shampoo stores symmetric positive semi-definin matrices that approximate an idealized preconditioner with L, = E| E R mxx and R. = G. R=x", and updates the preconditioner and weight matrix by
L B + (1 -2)-3))
(3)
R, 3, (1- (1
(4)
W W, W. G.R,
(5)
given 32 E [0, 1).33 The original Shampoo algorithm uses accumulate the gradient outer products, although the exponential moving average is more commonly used in practice. Therefore, we are interested in approximating the preconditioner of full-matrix Adam, given by A, A2
(1- 1-32)999 s Elg.g!")
The search direction or update in matrix form given as U. Shamppo = LL, - G or, equiva-
lently, C Shampao = (R, @ L with = 1/2 in Equation (2). We also consider Shampoo? defined as C Shampoo' R, @ L,, which a tighter approximation full-matrix AdaGrad (Morwani et al., 2025). Both updates can be generalized to higher-order tensors.
2.2 Eigenvalue-corrette Shampoo and SOAP
By leveraging Kronecker product approximation, Shampoo's preconditioner is restricted Kronecker product structure for both its eigenvectors and eigenvalues. Specifically, if we have a pre- conditioner C, R. L, with eigendecompositions L. Qr,Ar. QT, and R. QR R, Qh, then C, has eigendecomposition C,- (QR, @ Q2.) @ AL, (QR. @ lz.) T. Preconditioning St with C Shampoo and = 1/2 in its matrix form is equivalent to the matrix transformation:
Shampoo L, G.
Qu. - QQL G.Q...QR.). Qh.
We drop the subscript in the expectation. taken with respect (æt, 313 pD(æ, y).
A pseudo-ing iecudo-nvveree inverse that ignores the null space of the matrix or perturbing the matrix by regularization term cl for 0 can be used to handle the symmetric positive semi-definite case.
3

---

This can be interpreted as applying a bronneckr-ffcctrae coordinate-wisg scaling to a gradient with changed basis, i.c., G. QT, Q.Q n. scaling the transformed gradient, and converting it back to its original basis.
Instead of restricting the eigenvalues to be a Kronecker product, Liu et al. (2018) and George et al. (2018) have proposed correcting the eigenvalues by decoupling the scaling from the basis in K-FAC. This involves computing a separate scaling matrix D, EE 02 Rmxx and using it in place of the preconditioner's 's original eigenvalues mat diag( RR. & AL.) 45
Anil et al. (2020) noted that this correction can also be applied to Shampoo. Most recently, Vyas et al. (2025a) presented promising empirical results for language models using an instance of eigenvalue- corrected Shampoo called SOAP, which updates the scaling by D. B2 + (1 - f32) G2 Since then, SOAP has also been shown to perform well for physics-informed neural networks (Wang et al., 2025) and to reduce outlier features in transformers, which potentially improves quantization (He et al., 2024). We refer to our practical instantiation of eigenvalue-correcte Shampoo as EShampoo (Appendix B, Algorithm 2), which uses the same D, SOAP. See Appendix D.1 for clarification on the distinction between eigenvalue correction, EShampoo, and SOAP.
3 Learning rate grafting compensates for Shampoo's eigenvalues
Originally motivated by decoupling an optimizer's update magnitude from its direction to account for different implicit learning rate schedules, Agarwal et al. (2020) proposed learning rate grafting, a technique that combines the layer-wise update of one optimizer with the layer-wise update magnitude update
of another. For Shampoo, this means taking Shampoo's 's layer-wise update U-Shampoo and rescaling by the Frobenius norm of the grafting method's 's update U UGrattng (typically Adam):
Grafting]
W = W. + 0t IU
F
U Shampoo
U
Shampoo
A complete description of Shampoo with Adam grafting is given in Appendix B, Algorithm 3.
This approach has been critical to Shampoo's empirical success (Anil et al., 2020; Shi et al., 2023; Kasimbeg et al., 2025). As shown in Figure 1, Shampoo without grafting (and updating the root inverse matrices every F = 100 steps) underperforms Adam with many hyperparameter settings diverging. Anil et al. (2020) suggests that grafting is used to account for differences in magnitude of the eigenspectrum of the Kronecker factors for different layers, as well as the infrequent updates of the Kronecker factors and their inverse roots (see Anil et al. (2020), Appendix G). However, these claims have not been thoroughly investigated.
To focus our empirical investigation, we study Shampoo with Adam grafting and F 100, which was the winning configuration that was used in the external tuning track of the AlgoPerf competition (Kasimbeg et al., 2025). Since grafting re-scales the layer-wise update based on its magnitude, we analyze the Frobenius norm of the updates of full-matrix and diagonal Adam, Shampoo, and EShampoo. The magnitude of the eigendecomposed update in Equation (2) with wannecke--facttree Qc = QR. @ QL. is determined by the norm of the stochastic gradient G. and the eigenvalues of C₁:
Lemma 1. U QD o (Q) E mm be the generalized eigendecomposed
Fronecker-factoree update given by orthogonal matrices QL E Rmxm QR E Raxxn and dense scaling matrix D E R with p >> 0. Then we have:
(max D - G|G S [|U|| S (min Da.)) P G| 6.
(8)
1
Lemma 1 covers multiple algorithms. We can recover the idealized Adam update by setting Qz
E Rmxm Qn = I E R and D = E [[G'2] with p 1/2. The idealized Shampoo update can be recovered through the choice of D = mat diag (AR @ Ay.) with p = 1/4, L = E[GGT] =
Q and R E G = QR n T (or p 1/2 for Shampoo²). Idealized EShampoo is recovered with the choice of D = E [ )2) and p 1/2 instead.
4 0x denotes the element-wise square. "ma diag(-) takes a mn matrix and reshapes its diagonal into n matrix.

---

Vision transformer
ConvNeXt V2
Train loss (log scale)
10
Shampoo with grafting
Shampuo
Shampoo with scaling
EShampoo
10-4
F = 100
Adam W
10 10
10
10
10
Learning rate (log scale)
Learning rate (log scale)
Figure 2: Training results with different Shampoo variants and eigendecomposition frequencies F on the Imagewoof dataset. Shampoo with eigenvalue correction achieves a better training loss compared to Shampoo with Adam grafting, and the optimal learning rate for Adam transfers to both variants.
Proposition 1. Assume that Elgg") is symmetric positive definite. The magnitude of the idealized updates for full-matrix Adam, diagonal Adam, and EShampoo are all bounded by the power of the extreme eigenvalues of full-mmtriix Adam:
Amaax (Elgg P|| S |U||
Amia (Bigg ()
for all > O. However, under the simplifying assumption that E [ = 0 and Gi.; is independent from
Gk for (i.j) + (k.,)) and has bounded second moment, Amin gig ") < E[G I Amax (E (ggg)) and Shampoo has ciinensioo-dependdee bounds:
p/2 p/2 Amaax (Elgg' J) G| S ||0|| S 2/2 p/2 Amin (E(gg) G
(10)
See Appendix C for the proofs of Lemma and Proposition 1.
This highlights key issue: Shampoo's 's update magnitude can be mis-scaled relative Adam and EShampoo, especially due to limension-dependen factors. The basis does not influence the update magnitude only the eigenvalues do. While this additional scaling can be absorbed into the learning rate when only handling single matrix, it is potentially problematic when one needs to handle multiple parameters simultaneously. addition, use of stale eigenvalues can result update magnitudes that lie outside of the bounds of full-matrix Adam. This leads us to the following hypothesis:
The role of learning rate grafting in Shampoo
Learning rate grafting compensates for the scaling and staleness of Shampoo's eigenvalues.
From the Frobenius norm approximation perspective, using Shampoo? 2 (via C Shampoo" 2 approximation to full-matrix Adam compared to Shampoo (Morwani et al., 2025), which addresses the mismatch of the eigenvalues' exponent in Equation (10). Additionally, rescaling the preconditioner by S-1 Tr(R,) -1 Tr(L,) -1 ensures exactness when full-matrix Adam A, is Kronecker product (Morwani et al., 2025). The scaling 5 has also been previously introduced for Tensor Normal Training (Ren & Goldfarb, 2021, TNT) to approximate the Fisher information matrix.
yields tighter
Based on our observations, we can make several predictions:
1. With F 1, Shampoo without grafting should perform well when layer scalings are similar, but may struggle with highly variable parameter shapes.
2. Using S-1 with F should address scaling and staleness and match Shampoo with grafting. This is exact when full-matrix Adam decomposes into Kronecker product.
3. Updating an eigenvalue correction every iteration (e.g. like in SOAP) should also address scaling and staleness, matching Shampoo with grafting.
5

---

Empirical validation. We ablate different variants of Shampoo on the Imagewoof dataset with vision transformer (ViT) and ConvNeXt V2 models. We plot the final training loss after 100 epochs against the learning rate. Following the specification of the AlgoPerf Shampoo submission, we update the preconditioner every 100 steps when grafting the learning rate from Adam, and the eigenbasis when using eigenvalue correction. Our results are presented in Figure 2. We observe that all three predictions are confirmed: Shampoo" scaled by S or eigenvalue correction is able to match or even surpass Shampoo with grafting on both problems, and Shampoo without grafting can match the final loss for the Imagewoof ViT workload, but not for the ConvNeXt V2 model.
Hyperparameter transfer. The optimal learning rate of Adam transfers to Shampoo with grafting and EShampoo. However, is not apparent whether the optimal learning rate transfers to Shampoo" scaled by S Learning rate transfer is only sensible when jointly transferring since it affects the effective step size, but Shampoo does not directly add to the root of the eigenvalues, as done in Adam and other diagonal-scaalin- based optimizers. While we have not tested this, the optimal learning rate may potentially transfer when matching the effective c.
Computational and memory costs. While updating Shampoo every iteration incurs a prohibitive computational overhead, the eigenvalue correction is only slightly more expensive than Adam grafting since it requires more matrix products. EShampoo adds the same memory overhead to Shampoo as Adam grafting, specifically requiring an additional d-dimensional buffer for the second moment accumulation. One could reduce this overhead by adopting techniques such as Adam- mini, which only requires a single scalar per parameter block, defined according to the Hessian structure (Zhang et al., 2025b), and could be applied to both grafting and the eigenvalue correction. Alternatively, Liu et al. (2025) propose to scale Muon to approximately match the RMS norm of Adam's update to enable the re-use of Adam's 's hyperparameters, which may provide a practical heuristic that removes grafting without any memory overhead. Finally, an eigenvalue correction for the individual Kronecker factors can also correct for the staleness of their eigenvalues and has, in fact, been shown to also remove the need for grafting, while reducing the necessary buffer size from mn to m + n for each m X n weight matrix (Lin et al., 2025).
Applicability to other methods. The issue of stale eigenvalues extends all sparsely- connecker-factoree preconditioners such as K-FAC variants and TNT. Although we are unaware of any work that combines K-FAC or TNT with grafting, is common practice to apply the Adam update to gradients preconditioned with K-FAC (Pauloski et al., 2021; Osawa et al., 2023a,b; Eschenhagen et al., 2023), which we hypothesize could serve a similar function as grafting.
Controlling the approximation error induced by stale eigenbases
While frequent updates to the preconditioner's cigenvalues are important, the impact of periodically computing the eigenbasis remains less clear. The optimal frequency of the eigenbasis computation should be chosen to balance the time of each iteration against the method's per-iteration convergence rate, and may depend on the stage of training, layer type, and the specific Kronecker factors (L,, R2) involved. Shampoo currently uses fixed update frequency for all matrices, which ignores these distinctions and the computational cost associated with different parameter shapes.
4.1
Can we adaptively control the approximation error induced by the stale eigenbasis?
In SOAP, the initial eigenbasis is computed via an eigendecomposition during the first iteration, then is refined by a single power iteration and QR decomposition at all subsequent iterations, c.f. Algorithm 4 in Vyas et al. (2025a), attributed to Wang et al. (2024).6 The method corresponds to a single iteration of a warm-started simultaneous iteration, an extension of the power iteration for iteratively computing eigendecompositions (Golub & Van Loan, 2013). However, this approach may still yield a poor approximation to the true eigenbasis, especially when the eigenbasis changes rapidly due to the choice of 32 or the dynamics of the stochastic gradient.
To address this, we propose controlling the error induced by the eigenbasis using a warm-started QR algorithm that iterates until a relative approximation error criterion is satisfied. This approach offers three major benefits: (1) each Kronecker factor can be treated independently. allowing adaptive frequency across different layers, factors, and stages of training; (2) evaluating the criterion with the
'This approach was referred to as randomized SVD in Wang al. (2024), Appendix B.

---

Adaptive QR algorithm eigh
Adaptive eigh for different F and
QR (P = 7 0.2)
0.2
F p=
F
1.0
100
F 1000
noo-adaptive
Train loss (log scale)
10
adoptive (7 0.01)
20
adaptive (r = 0.1)
SOAP (QR, = 10,
adaptive (r 0.2)
eigh (F 1, 0.2)
adaptive (7 0.5)
AlgoPerf setting
loss 0.028.
19.3 min
0.1
loss 0.032. time 23.2 min
0.03
0.03
0
10
20
30
40
50
120
360
Wall-clock time [minutes]
Wall-clock time [minutes] (log scale)
Figure 3: All configurations are for ESShamoo. (left) On the Imagewoof ViT problem, setting the maximum number of iterations < 10 with threshold T = 0.2 for adaptive QR algorithm leads to significant increase in wall-clock time compared to using adaptive eigh. Even with / 10, adaptive eigh is faster. The default SOAP setting achieves worse final loss and also slightly slower. (right) Using the adaptive criterion determine when to skip the cigendecomposition (eigh) improves efficiency by 20% in wall-clock time compared to updating every 100 iterations (AlgoPerf setting).
last computed eigenbasis prior to the first QR iteration enables us to skip eigenbasis updates; and (3) the error in the eigenbasis be controlled through threshold E (0, 1), quantifying inexaciness.
Consider a single Kronecker factor L2 E Rmxm without loss of generality. We want to solve inexactly for orthogonal matrix Qr, such that we obtain approximation the cigendecompositionn L; QL,AL, QT., with the diagonal eigenvalues matrix AL.- the previous approximate eigenbasis Q is good approximation the current eigenbasis Q L. then the approximate eigenvalues induced by the stale eigenbasis ÅL. = QL. LiQL.., should be nearly diagonal.
We can leverage this observation to define a unified criterion for skipping the eigenbasis computation and terminating the warm-starte QR algorithm. Diagonalizing Ås. gives an approximation of L,, specifically, L. = Q L.., diag( QQ 7 We can bound the relative error of this approximation induced by the stale eigenbasis, or equivalently the relative error the Frobenius norm of the off-diagonal entries of AL, which is cheaper to compute:
Adaptive eigenbasis computation frequency criterion
LL Lell |
1QL Li diag(A IQQ,, L.QL.-IIEE
Å diag dizg(A))))l S.. (1))
||
This condition provides guarantee of the quality of the eigendecomposition approximation. If the condition is satisfied, reuse the previous eigenbasis QI. Q otherwise, refine it through the QR algorithm until the criterion is met. A complete pseudocode provided Appendix B, Algorithm 4. Note that evaluating the approximate eigenvalues ÅL. at the first step requires two matrix multiplications, which is significantly cheaper than computing step of the QR algorithm.
diag(-) takes vector/matrix and returns diagonal matrix with the vector/matrix's diagonal its diagonal.

---

Table Results on subset of the AlgoPerf workloads. We show the mean and standard error of the steps/time to the targets across the runs that reach them. See Appendix E, Table 3 for more results.
Workload
Shampoo Variant
Hits Target
Steps
Time [min]
Adam grafting (F 100)
4/5
4301 + 109
13.96 + 0.44
FastMRI
CBShampoo (F = 100)
5/5
2536 = 66
10.44 + 0.21
CBShampoo (T 0.1, F 50)
5/5
2468 + 145
10.81 + 0.72
Adam grafting (F = 100)
1/1
79907
894.27
ImageNet
CRShampoo (F 100)
1/1
76226
894.85
ViT
CEShampoo (T 0.1, F 50)
1/1
77459
935.89
Adam grafting (F 100)
2/5
12574 + 708
39.20 + 1.88
OGBG
C CBShampoo (F 100)
3/5
8320 + 1203
33.02 + 4.05
C EShamppo (T 0.1, F = 50)
5/5
7117 = 328
27.55 3.49
4.2
Does this adaptivity translate to efficiency gains?
The practical efficiency gains of adaptively controlling the approximation error depend on multiple factors. First, the efficiency of the QR algorithm is determined by how its computational cost compares standard eigendecomposition, which in turn depends both the cost of each QR iteration and the total number of iterations needed to satisfy the threshold т. Second, evaluating the criterion adds a small constant overhead independent of the actual frequency eigendecompositions.
Interestingly. we found that setting a smaller maximum number of QR iterations (1 < 10) per step results in significant slowdown in wall-clock time, as the total number of QR iterations accumulated over all training steps is higher than when using a larger maximum (7 2 10). However, even with a slightly more expensive than computing
larger number of maximum iterations, the QR algorithm eigendecompositions with torch. linalg eigh (shortened as eigh) whenever Equation (11) does not hold. The default configuration of SOAP, which computes single step of the warm-started simultaneous iteration every 10 steps, was also slightly slower in wall-clock time compared to adaptively computing eigendecompositions and reaches worse final loss. This result conflicts with the observation in Figure 7 in Vyas et al. (2025a), which may be due to differences in the types of workloads considered; see Figure 3 (left).
Because of this result, we primarily leverage the criterion for determining the frequency of calling eigh. To further bound the total possible number of eigendecompositions performed over the course of training, we evaluate the criterion at a fixed frequency F as opposed to every step. We ablate different choices of F and in Figure (right). We compare against the baseline setting of the winning AlgoPerf submission, which re-computes the eigendecomposition of all factor matrices every 100 steps on the Imagewoof ViT problem. We observe that this setting without adaptivity requires 20% more wall-clock time compared to an adaptive setting (T = 0.1, F = 50).
To test whether the results generalize to other problems, we consider subset of the AlgoPerf workloads and compare the winning Shampoo submission with Adam grafting to: (1) EShampoo with the same fixed eigenbasis computation frequency of F 100; and (2) EShampoo with 0.1 and F 50; see Table 1. First, as predicted by Section 3, EShampoo with the fixed eigenbasis update frequency matches outperforms Shampoo with Adam grafting in steps and wall-clock time, using the same hyperparameters. Second, using the adaptive criterion with 7 = 0.1 and F' = 50 matches the fixed frequeney F 100 for the FastMRI and OGBG workloads, and does slightly worse for ImageNet ViT. We also present results with other fixed and adaptive frequencies in Appendix E, Table 3. Overall, we find that the performance on the considered problems is quite robust to the eigenbasis computation frequency (c.f. Appendix E, Table 3). We expect that when Shampoo's 's convergence benefits from more frequent eigenbasis updates, adaptively determining the frequency can result in higher efficiency.
4.3
What patterns of adaptivity emerge?
The criterion in Equation (11) also provides insight into how rapidly the eigenbasis changes for different parameter types. On Imagewoof VIT, we set 0.1 and check the criterion at every
8

---

x10
Weights
nn Linerr LayerHorm
A
eigh computations
Blases
L.
0
x10
Steps
x10
x 10
Figure 4: We show the mean with standard error across preconditioners corresponding to the labels in the legends, for EShampoo with 0.01 and F 1 Imagewoof ViT. The eigenbases for biases and layer normalization parameters are changing faster than for weight matrices and linear layers, respectively.
step (F 1); see Figure We observe more frequent eigendecomposition computations early in training that trend towards constant update frequency at the end of training. When comparing different parameter types, we find that the eigenbases of preconditioners for bias terms evolve more rapidly than those for weights. Similarly. layer normalization parameters require almost 4x as many eigenbasis updates as linear layer parameters.
We also compare the number of cigendecompositions across the two Kronecker factors L, and R2 for weight matrices, and A, for biases and layer normalization parameters." The eigenbases of A, updated most frequently, followed by L., then R. The same trend holds when removing learnable layer norm parameters and for ConvNeXt V2 trained the same dataset (c.f. Appendix E, Figure 7). A similar analysis for Llama 3 model with 324 million parameters trained on 3.2 billion tokens of C4 data reveals different trend, namely, Lt evolves faster than R2 and A,, which correspond to RMS normalization parameters (c.f. Appendix E, Figure 8).
4.4 How does the error induced by the eigenbasis affect convergence?
While we have focused on controlling the error induced by the preconditioner's stale eigenbasis, our primary concern is its effect on convergence, rather than the error itself. In the Imagewoof ViT setting, we find that the approximation quality (determined by r) during early iterations far more critical Figure 5 (left). Specifically, setting
0.8 for the first 90% and T 0.01 for the last 10% of iterations yields a final loss nearly identical to using t = 0.8 throughout all of training (zs 0.09 0.1). In contrast, setting remaining 90% of iterations leads to a significantly better final loss (= 0.04).
than during later iterations;
0.01 for the first 10% and T 0.8 for the
0.99 for all of training) still significantly outperforms AdamW (== 0.12 vs. 0.3). This highlights that frequent eigenbasis computations are most beneficial early in training, consistent with previous observations on Shampoo, PSGD, and SOAP (Ishikawa Yokota, 2024; Walters al., 2025; Nestler, 2025; Vyas al., 2025b).
Remarkably, freezing the eigenbasis after the first iteration (by setting
To further investigate the discrepancy in the rate of change in the eigenbases corresponding to 1D and 2D parameters (c.f. Figure right), we compare Imagewoof ViT run with 7 0.01 for 2D parameters and no change of basis, i.c. Adam, for ID parameters and vice versa. Surprisingly, using Adam for ID parameters does not affect Shampoo's convergence final loss. Conversely, running Adam for 2D parameters and changing the basis according to = 0.01 for ID parameters closely matches Adam's convergence and final loss; see Figure 5 (right).
"Full-matrin Adam is also used weight matrices mx with mn S wtx_ maw.ppreconditionrr.de which is hyperparameter to control largest possible dimension the preconditioner.
For 5 out of 172 Kronecker factors, the eigenbasis computed twice instead of just once.

---

right: same ax)
Training Loss
0.01
1.5
80%
0.8
0.99
1.0
10%
-0%
10%
90% r 0.8 10% 0.01
0.5
ID: = 1.0, 2D: 0.01
10%
0.01, 2D: 1.0
0.0
25
50
75
100
0
25
50
75
100
Epoch
Epoch
Figure 5: All configurations are for EShampoo. (left) The error in the eigenbases is dramatically more important for early iterations. A single cigenbasis computation at the first iteration (T = 0.99) sufficient to outperform A. (right) The difference between the convergence behavior of A and EShampoo on this problem can be exclusively attributed to the eigenbases corresponding 2D parameters.
Given that 45% of the preconditioner matrices correspond to parameters and their eigenbases change more rapidly, this suggests that many cigendecompositions are not needed at all. For example, using T = 0.01 for all parameters increases runtime by 2.9x compared to only applying EShampoo to 2D parameters and Adam to ID parameters, with improvement in final loss. In practice, SOAP already uses Adam for 1D parameters order to reduce its computational and memory overhead (Vyas et al., 2025a).
Discussion and conclusion
In this paper, we demonstrate that frequently updating the eigenvalues while periodically updating the eigenbasis of Shampoo's ' preconditioner provides principled and practical approach for elim- inating grafting. remains an open question whether the same correction applies directly to the AdaGrad summation, or if more sophisticated, basis-aware eigenvalue correction is needed (c.f. Appendix D.2). We also show how controlling the approximation error induced by the stale eigen- basis can improve efficiency. In order to determine the eigenbasis computation frequency in truly manner, we must understand how approximation error impacts convergence, as well as integrate systems-level considerations, such batched kernel efficiency, into our algorithmic design. Further exploration is needed to understand how these trade-off's and techniques scale to larger models, such as large language models. While our work focuses Shampoo, the ideas limited to the AdaGrad family, and can be adapted to other methods such as K-FAC and TNT..
From theoretical perspective, major open question is how to incorporate approximation quality into regret bounds for Shampoo (c.f. Appendix D.I). Finally, while we have implicitly treated full-matrix Adam the right algorithm approximate, alternative interpretations of Shampoo may better explain Shampoo's 's effectiveness in practice (Carlson et al., 2015a,b; Benzing, 2022; Bernstein & Newhouse, 2024; Maes et al., 2024; Pethick et al., 2025; Zhang al., 2025a; Xie et al., 2025).
Acknowledgments
We thank Anna Cai, Parameswaran Raman, and Ke Sang for their in-depth review the paper. We grateful for insightful discussions with Ganesh Ajjanagadde, Rohan Anil, Anna Cai, Rong Jin, Jihao Andreas Lin, Bruno Mlodozenice, Vinay Rao, Isaac Reid, and Xingyu (Alice) Yang. We also acknowledge managerial support from Yuchen Hao, Guna Lakshminarayanan Maxim Naumov, Sandeep Parab, Chunqiang Tang, and Lin Xiao. Lastly, we thank the anonymous reviewers for their productive feedback and suggestions of experiments and Kyunghun Nam and Yushun Zhang for pointing out minor errors in a preprint this work.
Runa Eschenhagen is supported by ARM, the Cambridge Trust, and the Qualcomm Innovation Fellowship. Richard F. Turner supported by Google, Amazon, ARM, Improbable and an EPSRC Prosperity Partnership (EP/T005386/I) and the EPSRC Probabilistic AI Hub (EP/Y028783/I),
10

---

References
Naman Agarwal, Rohan Anil, Elad Hazan, Tomer Koren, and Cyril Zhang. Disentangling adaptive gradient methods from learning rates. arXiv2002.11880 11803, 2020.
Shun-Ichi Amari. Natural gradient works efficiently learning, Neural computation. 10(2), 1998.
Kang An, Yuxing Liu, Rui Pan, Shiqian Ma, Donald Goldfarb, and Tong Zhang. ASGO: Adaptive structured gradient optimization. arXiv 2503.20762, 2025.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, arXiv 2002.09018 8, 2020.
Frederik Benzing. Gradient descent neurons and its link approximate second-order optimization. In ICML,, 2022.
Albert S Berahas, Majid Jahani, Peter Richtárik, and Martin Takáč. Quasi-Newton methods for deep learning: Forget the past, just sample. arXiv 1901 100199997, 2020.
Alberto Bernacchia, Mate Lengyel, and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. In NeurlPS, 2018.
Jeremy Bernstein. Deriving Muon, 2025. URL
deriving-muus mun..
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. arXiv 2409.20325, 2024.
Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov. Better plain VIT baselines for ImageNet- lk. arXiv 2205.01580, 2022.
Raghu Bollapragada, Jorge Nocedal, Dheevatsa Mudigere, Hao-Jun Shi, and Ping Tak Peter Tang. A progressive batching L-BFGS method for machine learning. In International Conference on Machine Learning, pp. 620-629. PMLR, 2018.
Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223-31 2018.
David Carlson, Volkan Cevher, and Lawrence Carin. Stochastic spectral descent for restricted Boltzmann machines. In ICML, 2015a.
David Edwin Carlson, Edo Collins, Ya-Ping Hsich, Lawrence Carin, and Volkan Cevher. Precondi- tioned spectral descent for deep learning. In NIPS, 2015b.
George Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, Juhan Bae, Justin Gilmer, Abel L. Peirson. Bilal Khan, Rohan Anil, Mike Rabbat, Shankar Krishnan, Daniel Snider, Ehsan Amid, Kongtao Chen, Chris J. Maddison, Rakshith Vasudev, Michal Badura, Ankush Garg, and Peter Mattson. Benchmarking neural network training algorithms. arXiv 2306.07179, 2023.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. JMLR, 12(61), 2011.
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit S Dhillon. CASPR: Combining axes preconditioners through Kronecker approximation for deep learning. ICLR, 2024.
Runa Eschenhagen, Alexander Immer, Richard E. Turner, Frank Schneider, and Philipp Hennig. Kronecker-Factored Approximate Curvature for modern neural network architectures. In NeurlPS, 2023.
Vladimir Feinberg, Xinyi Chen, Jennifer Sun, Rohan Anil, and Elad Hazan. Sketchy: Memory- efficient adaptive regularization with frequent directions. In NeurIPS, 2023.
11

---

Kai-Xin Gao, Xiaolei Liu, Zheng-Hat Hai Huang, Min Wang, Zidong Wang, Dachuan Xu, and F. Yu. A trace-restricted Kronecker- cronecker-fcctored approximation to natural gradient. In AAAI Conference on Artificial Intelligence, 2020.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a Kronecker-factored eigenbasis. In NeurIPS, 2018.
Donald Goldfarb, Yi Ren, and Achraf Bahamou. Practical quasi-Newton methods for training deep neural networks. Advances in Neural Information Processing Systems, 2-2236623996 2020.
Gene H Golub and Charles F Van Loan. Matrix computations. JHU press. 2013.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang. Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits. Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia- Diego Perino, Dicuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Eytimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Kantala-Yeary Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupulcti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Ning Zhang. Olivier Duchenne, Onur Çelebi, Patrick Alrassy. Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang. Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh. Sean Bell, Scohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do, Vish Vogeti, Vitor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang. Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song, Yuchen Zhang. Yue Li. Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan. Zhengxing Chen, Zoe Papakipos. Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi. Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic. Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu,
12

---

Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching- Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer. Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Due Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanann Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Frina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torrs,, Josh Ginsburg, Junjie Wang. Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang. Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskow, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang. Russ Howes, Ruty Rinott, Sachin Mehta. Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shuqiang Zhang. Shuqiang Zhang, Sinong Wang, Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li. Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang. Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang. Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary De Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. The Llama 3 herd of models. arXiv 2407 2024.
Roger B. Grosse and James Martens. A Kronecker-factoree approximate Fisher matrix for convolution layers. In ICML, 2016.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In ICML, 2018
Bobby He, Lorenzo Noci, Daniele Paliotta, Imanol Schlag, and Thomas Hofmann. Understanding and minimising outlier features in neural network training. In NeurIPS, 2024.
Tom Heskes. On "natural" learning and pruning in multilayered perceptrons. Neural Computation, 12(4), 2000.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv 2005.00687, 2021.
13

---

Satoki Ishikawa and Rio Yokota. When does second-ord order optimization speed up training? In The Second Tiny Papers Track at ICLR, 2024.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https: //kellerjordan betooordannggthuue. io/ postssmmontt
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 ICLR, 2025.
Nitish Shirish Keskar and Albert S Berahas. adaqn: An adaptive quasi-Newton algorithm for training RNNs. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part 16, pp. 1-16. Springer, 2016.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Florian Knoll, Jure Zbontar, Anuroop Sriram, Matthew J. Muckley, Mary Bruno, Aaron Defazio, Marc Parente, Krzysztof J. Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzaly, Adriana Romero, Michael Rabbat, Pascal Vincent, James Pinkerton, Duo Wang, Nafissa Yakubova, Erich Owens, C. Lawrence Zitnick, Michael P. Recht, Daniel K. Sodickson, and Yvonne W. Lui. fastMRI: A publicly available raw k-space and DICOM dataset of knee images for accelerated MR image reconstruction using machine learning. Radiology: Artificial Intelligence, 1900077 2020. doi: 10. 148/1y4i 20201
PMID: 32076662.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. ImageNet classification with deep convolu- tional neural networks. In NIPS, 2012.
Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical Fisher approxima- tion for natural gradient descent. In NeurlPS, 2019.
Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE Transactions on Neural Networks and Learning Systems, 29(5), 2018
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? second-order perspective. In ICML, 2024a.
Wu Lin, Felix Dangel, Runa Eschenhagen, Kirill Neklyudov, Agustinus Kristiadi, Richard E. Turner, and Alireza Makhzani. Structured inverse-free natural gradient: Memory-efficient & numerically- stable KFAC for large neural nets. In ICML, 2024b.
Wu Lin. Scott C. Lowe, Felix Dangel, Runa Eschenhagen, Zikun Xu, and Roger B. Grosse. Under- standing and improving Shampoo and SOAP via Lelleack-Leibler minimization, arXiv 2509.03378. 2025.
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang. Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, and Zhilin Yang. Muon is scalable for IIm training. arXiv 2502.16982. 2025.
Xialei Liu, Mare Masana, Luis Herranz, Joost van Weijer, Antonio M. López, and Andrew D. Bagdanov. Rotate your networks: Better weight consolidation and less catastrophic forgetting. ICPR, 2018
Lucas Maes, Tianyue H. Zhang, Alexia Lolicoeur-Martineaul Ioannis Mitliagkas, Damien Scieur, Simon Lacoste-Julien, and Charles Guille-Escuret. Understanding Adam requires better rotation dependent assumptions. arXiv 2410.19964 2024.
James Martens. Deep learning via Hessian-free optimization. In ICML, 2010.
14

---

James Martens. New insights and perspectives on the natural gradient method. JMLR, 21(146), 2014.
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In ICML, 2015.
James Martens, Jimmy Вa, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In ICLR, 2018.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on Shampoo's preconditioner. In /CLR, 2025.
Lucas Nestler. HeavyBall, 2025. URL https:://gtthiiiie
Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler. ASDL: A unified interface for gradient preconditioning in PyTorch. arXiv 2305.04684, 2023a.
Kazuki Osawa, Shigang Li, and Torsten Hoefler. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. In Dawn Song, Michael Carbin, and Tianqi Chen (eds.), MysSys, 2023b.
Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian T. Foster, and Zhao Zhang. KAISA: an adaptive second-order optimizer framework for deep neural networks. In International Conference for High Performance Compuling, Networking, Storage and Analysis (SC21), 2021.
Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, and Volkan Cevher. Training deep learning models with corm-constrained LMOs, arXiv 2502.07529, 2025.
Omead Pooladzandi and Xi-Lin Li. Curvature-informed SGD via general purpose lie-group precondi- tioners. arXiv 2402.04553, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv 1910. 10683, 2023.
Yi Ren and Donald Goldfarb. Tensor normal training for deep learning models. In NewrIPS, 2021.
Thomas Robert, Mher Safaryan, Ionut- Modoranu, and Dan Alistarh. LDAdam: Adaptive optimization from low-dimensional gradient statistics. arXiv 2410.16103, 2024.
Hao-Jun Michael Shi, Tsung-Hsien Lee, Shintaro Iwasaki, Jose Gallego-Posada, Zhijing Li, Kaushik Rangadurai, Dheevatsa Mudigere, and Michael Rabbat. A distributed data-parallel PyTorch implementation of the Distributed Shampoo optimizer for training neural networks at-scale. arXiv 2309.06497, 2023.
DiJia Su, Andrew Gu, Jane Xu, Yuandong Tian, and Jiawei Zhao. GaLore 2: Large-scale LLM pre-training by gradient low-rank projection. arXiv 2504.20437, 2025.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham Kakade. SOAP: Improving and stabilizing Shampoo using Adam. In JCLR, 2025a.
Nikhil Vyas, Rosie Zhao, Depen Morwani, Mujin Kwun, and Sham Kakade, Improving SOAP using iterative whitening and Muon. https://nikhilyye github. 1o/SOAP_Muuu ___o__uon.__ pdf, 2025b.
Evan Walters, Omead Pooladzandi, and Xi-Lin Li. kron torch, 2025. URL https://githu: com/ evanattyourerrrv.cee/rooo..o wwtps//ooccooooooo.o.oooooo
Sifan Wang, Ananyae Kumar Bhartari, Bowen Li, and Paris Perdikaris. Gradient alignment in physics-informet neural networks: A second-order optimization perspective. arXiv 2502.00604, 2025.
Sike Wang, Pan Zhou, Jia Li, and Hua Huang. 4-bit Shampoo for memory-efficient network training. In NeurIPS, 2024.
15

---

Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and Saining Xie. ConvNeXt V2: Co-designing and scaling ConvNets with masked autoencoders. arXiv 2301.00808, 2023.
Shuo Xie, Tianhao Wang, Sashank Reddi, Sanjiv Kumar, and Zhiyuan Li. Structured preconditioners in adaptive optimization: A unified analysis. arXiv 2503. 2025.
Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew J. Muckley, Aaron Defazio, Ruben Stem,, Patricia Johnson, Mary Bruno, Mare Parente, Krzysztof J. Geras, Joe Katsnelson, Hersh Chandarana, Zizhao Zhang, Michal Drozdzal, Adriana Romero, Michael Rabbat, Pascal Vincent, Nafissa Yakubova, James Pinkerton, Duo Wang, Erich Owens, C. Lawrence Zitnick, Michael P. Recht, Daniel K. Sodickson, and Yvonne W. Lui. fastMRI: An open dataset and benchmarks for accelerated MRI. arXiv 181 1811.08839, 2019.
Thomas Zhang, Behrad Moniri, Ansh Nagwekar, Faraz Rahman, Anton Xue, Hamed Hassani, and Nikolai Matni. On the concurrence of fayer-wise preconditioning methods and provable feature learning. ICML,, 2025a.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P. Kingma, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. In ICLR, 2025b.
Jawwi Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. GaLore: Memory-efficient LLM training by gradient low-rank projection. arXiv 2403.03501 2024.
16

---

Table 2: Shampoo variants and their properties.
Justified from approximation perspective
Shampoo variant
Matches grafting steps
Matches grafting in compute memory
Learning rate transfer
grafting (Adam) C Shampoo
x
x
S Shhmpooo 2
X
EEShappoo
A Connecting full-matrix AdaGrad to the Fisher
We can also consider a more general class of preconditioned stochastic gradient methods with other choices of C, and p > 0 that update the parameters at each iteration t by Equation (2). For example, one could choose C₁ as approximations of the Hessian V (0) via subsampling or secant approxima- which yields the class of stochastic Newton or quasi-Newton methods (Keskar & Be-
tion with p- rahas, 2016; Bollapragada et al., 2018; Berahas et al., 2020; Goldfarb et al., 2020). Another common choice is the Fisher information matrix F, Eg (()) VoL(fo (z)) ÿ) (oL(f(z)) ")"}. which yields the natural gradient for common loss functions, generalized Gauss-Newton method with p = (Amari, 1998; Martens, 2014). A summary of different methods provided Bottou et al. (2018).
By instead summing over per-sample gradient outer products, A. can be connected to the empirical Fisher. In general, the empirical Fisher is not expected be good approximation of the Fisher information matrix (Kunstner et al., 2019). However, by replacing the accumulation with an expec- tation over the conditional distribution given by the model fa. full-matrix AdaGrad can be made equivalent to the Fisher. This equivalence reveals a natural extension of Shampoo to approximate the Fisher, called Tensor Normal Training (Ren & Goldfarb, 2021, TNT), and is closely related to the K-FAC preconditioner (Anil et al., 2020). In fact, both approximations are exactly equivalent to the (block-diagonal) Fisher for simple cases such as deep linear networks (also with weight sharing) and mean squared error loss (Bernacchia et al., 2018; Eschenhagen et al., 2023; Morwani et al., 2025). The modifications in TNT deviate from the original Shampoo update (Gupta al., 2018) because was motivated by upper bounds to full-matrix AdaGrad in its non-smooth, convex regret analysis, rather than this Fisher approximation perspective.
B Algorithms pseudocode
In this section, we provide the pseudocode for all algorithms, including idealized (Algorithm 1) and practical (Algorithm 2) aagenvalue-correcte Shampoo, Shampoo with Adam grafting (Algorithm 3), and the adaptive warm- stamtet QR algorithm (Algorithm 4). Algorithm I and Algorithm 3 present simplified versions of the algorithm ignoring common modifications like momentum or an exponential moving average over the stochastic gradient, bias corrections, and weight decay.
Note that different instances of eigenvalue-corrected Shampoo can be employed by changing the exponential moving average in Equation (12), eigenbasis computation in Equation (13), or eigenvalue correction update in Equation (14). SOAP delays the update of the eigenbasis until after the update step, approximates the eigenvalue correction (Equation (14)) by
D. = D.D + (1 (1- (1-3,) Gr
(19)
and uses a single iteration of the warm-started simultaneous iteration. Additionally, there is a discrepancy between the SOAP algorithm presented in the paper and the official implementation: in the paper, the exponential moving average of the gradient G; is used (line 4 of Algorithm 3 in Vyas et al. (2025a)), whereas the implementation computes the exponential moving average over the rotated gradient G. This can be interpreted as running Adam Shampoo's eigenspace, which shares similarities to the GaLore algorithm (Zhao et al., 2024; Su et al., 2025).
See https: soap.py#L.. py#1167.
com http://www..cccaaco//oooooooo//t htt.://ee..1111eeaaaaaaoaeeeeaa/
17

---

Algorithm 1 Idealized eigenvalue-correeeeed corrected Shampoo pseudocode
Require: Parameter W1 E R m* learning rate dy > c > 0.
1: Initialize Lo 0 E Rm Ro 0 E Roxn and D = 0 E mm
2: for /= -,,, l,..: do
3: Compute (mini-hatch) stochastic gradient: G. (ol(Jo,(.))))))
4: Update factor matrices:
3 + (1- R. 12 (11--2))Gr G2.
(12)
5: Compute orthonormal eigenbasis of the factor matrices:
Qr. = eigvec(Lr), QR. - eigvec(
(13)
6: Transform gradient basis: G. QL, G.QR.
7: Compute or update eigenvalue correction:
D arg min A - (Qn. @ Qz.) diag(vec(D))(Qnn & Qn.)le.
(14)
DERnXn
8: Compute update: W411 = We - aaaaQL @ +d11))
9:
end for
Algorithm 2 EShampoo pscudocode (as implemented for the experiments)
Require: Parameter W1 mRx, learning rate a; 0, > 0, 0, B2 [0, 1), weight decay
X 2 0, eigenbasis computation frequency F E N, and threshold T E [0, 1] for Equation (11) and Algorithm 4.
1: Initialize Mn 0 E R Lo = =0 Raxm Ro = 0 E paxa Q Im, QR, = I,, and
Do 0 mm
2: for t = 1.... do
3: Compute (mini-batch) stochastic gradient: G. Vol(fe, (x), y).
4: Compute exponential moving average of the gradient: M, = M + (1-31) G.
5: Update factor matrices:
L, 3,L, (1- R. 32R (1-)))
(15)
6:
if mod F = 0 and not Equation (11) then
7:
Compute eigenbasis factor matrices (e.g. with torch. linalg.nenne. eigh or Algorithm 4):
Qr. eigvec L./(1 -3b),,
QR. eigvec (eyvec(R//(1 R2/(1 B5)).
(16)
8:
else
Qr, = QQ Qn. QR--1
(17)
9:
end if
10:
Transform gradient basis: G. = QL, G.Qn.
Compute or update eigenvalue correction: D, = B2 .-D1-1 (1 G3)
12:
Perform bias correction:
M, M,((1 B1), = DD/(111 -3))-
Compute parameter update:
W W. as (QL. MI C...RQRROOVD.. (VD.+an')) Qh, +aw.).
(18)
end for
18

---

Algorithm 3 Shampoo with Adam grafting pseudocode
Require: Parameter W1 E Rmx", learning rate at 0, exponential moving average constant 82 E (0, 1).
1: Initialize Lo 0 E Rmxm Ro 0 E R and D 0 E RR
2: for t= 1,.... do
3: Compute (mini-batch) stochastic gradient: G. (o)((o,())),1
4: Update factor matrices:
L. B + +1 --) R. 3= R2 + -1 G.
5:
Update Adam grafting state: D, = + (1- 82 G
6: Compute matrix root inverse of the factor matrices:
L 1/4 rootinv(L,),
R. 1/4 = rootinv R.).
| D le L 1/4 R, 1/4
7:
Compute update: W 2++
W Cq
|| L, GR, 1/1
8: end for
Algorithm 4 Warm-started QR iteration with relative enror termination criterion.
Require: Matrix L = L, previous eigenbasis Q = QL relative tolerance r E [0, 1), maximum
number of iterations I.
1: 25+0 0
2: À < QQ
3: while | diag( A) S <T|||| and í do
4:
Q, QA(A)
5:
R+RQ
6:
Q
7:
i+i+l
8: end while
9:
return Q, A
In contrast, implementation of EShampoo does not delay the update of the preconditioner, the same approximation of the eigenvalue correction as SOAP, uses torch, .lliallgg eigh for the eigendecomposition (unless indicated otherwise), and computes the exponential moving average the gradient G,. Both algorithms use bias corrections and (decoupled) weight decay. We provide the pseudocode for EShampoo as it was implemented for experiments in Algorithm 2. If F o the algorithm reduces to AdamW W until iteration == F.
While we do not present experimental results here, Algorithm 4 or Equation (11) for eigh can also be used in Shampoo without an eigenvalue correction. This implementation stores the previous cigendecomposition of the Kronecker factors instead of the previous root-inverse Kronecker fac- tors. When we skip an eigendecomposition, we maintain the previous eigenbasis and replace the previous eigenvalues with the estimated eigenvalues (A, = diag (QL... L,Q1,-1) used in Equation (1 ()... Alternatively, one can compute an exponential moving average over these estimated eigenvalues done in KL-Shampoo (Lin et al., 2025), which also removes the need for grafting like the eigenvalue correction in SOAP and EShampoo (c.f. Section 3).
C Proofs
Lemma 1. = ( QQ) E Rmxx be the generalized eigendecomposed
Kronecker-fuctored update given by orthogonal matrices Qr Rmmx Q n Raxn, and dense scaling matrix D E m Then we have:
(max Da.)) - Glle & |1 & (min Di.)) -P G
(20)
19

---

Proof. Since the Frobenius norm of a matrix is invariant to orthogonal transformations and the entries
of are bounded, i.e., m D., - Di S max., Di.j- we can show that
|U F = [Q 0 QQR = =| (Ql GQR))l
E [max Da.j) (min - GQull ijj
[(max Da.)) P, (min
ig
Proposition 1. Assume that Elgg') is symmetric positive definite. The magnitude of the updates for full- Adam, diagonal Adam, and eigenvalue-corrected Shampoo are all bounded by the power of the extreme eigenvalues of full-matrii Adam:
Xmax (gl9g") "|| S |UU| Amin (Bigg) |||G||F
(21)
for all >> O. However, under the simplifying assumption that B[G] 0 and G., is independent from G for (i.j) + (ki. 1) and has bounded second moment, Amin ([gg") < EE S maa (E (gg") and Shampoo has fimensioon-depndees bounds:
m. -P/2 p/22 mmss (Elgg P| S
| S mmmmPP p p/2 Ain (Eigg)
+ (22)
P.
Proof: Note that Equation (21) holds for full-matrix Adam since |U]| P ||||2
|F[gg
max E ggggi })> ->- Amin (Eigg -) -1gll1,, where u = vec(U) and g = vec(G).
For diagonal Adam and cigenvalue-corrected Shampoo, it is sufficient to show that Xmin (Bigg D < DiJ Amax (E(gg" ")) for all i and apply Lemma To see this, note that D.. can be represented by Rayleigh quotient, i.e..
D..j = E[G f Flgg Jes.
(Adam)
Di, = B[[Q GQ = e (Qn @ E[gg" @ Qz)ei,
(EShampoo)
where eij vec E.j mm and Egg E Rmxn with
(Eaj)a. {
if (k: (i,
otherwise.
Since |es,ll2 = R R @ Qr)ei, 1| = 1, the desired bound by the extreme eigenvalues follows from the Courant-Fischer Mrrscher-Weyll min-max theorem.
To prove Equation (22), observe that under independence between components of the gradient G. the preconditioner for full-matrix Adam E[gg"] is diagonal matrix whose diagonal entries consist of E for all i,j. Hence, E(G?,, is bounded by the minimum and maximum eigenvalues, i.e., Amax (E(gg) [)]
B[G E
Due to independence, E|GG"] and E G are also diagonal. Expanding their diagonal entries gives
BEGGG")).
ElG Gyy En (mii(E[gg) "). Amax (Eigg
V i = 1,... m,
(GGTT))a
Σ
G m [Amin (Eggg D), Amax (B[gg
V I,...,
Therefore, the Shampoo preconditioner (E[GGT) @ E G)) is also diagonal with eigenvalues lying in the interval [m 1/2 / AniimmEEgg (Eigg") m. 1/ 1/2 Amax (Elgg") [D]. The desired result follows.
Note that more general bounds can derived by relaxing the assumption EF [ 0, although the bounds more complex and are not more conceptually informative.
20

---

D On the gap between the optimal and practical eigenvalue correction
Eigenvalue-correeted Shampoo shares similarities with memorryefffciint optimizers such as GaLore, which apply Adam in a low-dimensional subspace spanned by the largest singular vectors of the gradient matrix (Zhao et al., 2024; Su et al., 2025). However, GaLore relies on the assumption that the gradient resides in a low-rank subspace that evolves gradually, allowing the same optimizer state be used even as the subspace is updated. A recent method called LDAdam proposed by Robert et al. (2024) describes projection - aware method that corrects the scaling matrix through both projection-awant update and generalized error feedback mechanism when transitioning between subspaces to address this issue. We will see that the naive eigenvalue correction used in SOAP suffers from similar limitations.
D.1 Optimal eigenvalue correction in Frobenius norm
Note that we can determine the optimal eigenvalue correction by minimizing its Frobenius norm approximation full-matrix AdaGrad or Adam:
Proposition 2. Given symmetric matrix C' E Bdddd and orthogonal matrix Q E Rdxd, the optimal eigenvalue correction D* that minimizes the Frobenius norm distance is given by:
D = diag( (2'(QQQ) CQ) E min DED
-DC
(23)
The exact expression for D' depends on the form of accumulation used for computing C: note that the damping term omitted here since it does not change the optimal solution. T denotes the current iteration.
1. (Idealized) If = 1 [gg"], then D = =E [diagg(() (Q" [2] (George et al., 2018).
2. (AdaGrad) If Ei 9:g! then D;; = ET= diag 0.
3. (Adam) If (1- (1-B)) - then Di, (1 B2) Σ 37 diag diag(Qtg.)) Q+ )2 Note that C can be generated recursively via an exponential moving average Cr Rmnxm
B 32CT-1+ (1 32)grgT with C.
Proof. (Informal.) Since Q is orthogonal, ||C -| I| QQ Dllr. with D diagonal. Therefore, the optimal solution has the form D diag( Q CQ). Each case then follows by observing that C' is the expectation or weighted sum of gg" and passing into the sum or expectation.
While D is optimal in this sense, it does not guarantee anything regarding the similarity of the (root) inverse of the approximation. Using Proposition we can establish the following corollary.
Corollary 1. The optimal eigenvalue correction yields tighter Frobenius norm approximation than Shampoo and CASPR, ie..
|CC QD'Q" S - CShampooo [F. ||C CASPPR [|]).
where CASPR is the preconditioner used by the CASPR algorithm (Duvvuri et al., 2024).
Proof. (Informal.) The first inequality trivially follows from Proposition since by definition the cigenvectors of Shampoo are equal to Q = (QR @ Qz)T. T The second inequality follows from Proposition together with Lemma 3.4 in Duvvuri et al. (2024), which states that the eigenvectors of Shampoo and CASPR are identical.
This means that using the optimal eigenvalue correction (with respect to Frobenius norm approxima- tion) yields tighter approximation to the full-matrix quantity compared to Shampoo or CASPR. A remaining open theoretical question whether tighter Frobenius norm approximation can yield a tighter regret bound compared to Shampoo and CASPR. For example, one can obtain a tighter regret bound than Shampoo by only considering one-sided Shampoo, which not generally a better approximation to full-matrix AdaGrad (Xie et al., 2025; An et al., 2025).
d denotes the of d x d diagonal matrices.
21

---

F == 100
F 10
0.5
No basis correction
0.10
0.7
0.9
0.09
92
0.95
0.0292
0.0261
0.0285
0249
0.0266
0.0187
0.08
0.99
0.0200
0.0208
0.0261
0212
0.0280
0.0229
0.0218
0.0209
0.07
Train loes
0.999
0.0206
0.0263
0.0299
0.0282
0,0222
0.06
0.5
0.0839
0.0868
0.0458
0.0063
0.1071
Basis corroction
0.7
0.0387
0.0444
0.1010
0.05
0.9
0.0283
0.0276
0.0298
0.0669
0.01
0.95
0.0268
0.029G
0.0480
0.99
0.0248
0.0247
0.0237
0.0241
0.0274
0.0381
02337
0.0222
0.0252
0.02210
0.0239
0.0269
0.999
0.02844
0.03:2
0. 0.00272
0.0274
0.0358
0272
0.0261
0.0257
0.0215
0.0234
0.0251
0.02
0.5
0.95
0.99
0.909
0.5
0.95
0 0.99
0.009
As
Bs
Figure 6: Decoupling the exponential moving average for the eigenbasis (32) and eigenvalues (3s) for different eigendecomposition frequencies (F). We can observe remarkable invariance to the choice of B2 and 33. Interestingly, the correction for the change of bases seems to hurt performance overall and especially for low 82, high 83, and low F values a pattern that we might expect without the correction.
D.2 Basis-aware eigenvalue correction
Despite its optimality, computing D" for cases 2 and 3 in Proposition 2 not feasible in practice since would have to store and transform prior gradients from all previous iterations with Qr. have access to the full-matrix quantity Cr. For simplicity, we will focus on case 3, although the results generalize without loss of generality." In contrast to the optimal eigenvalue correction, which uses fixed basis matrix Qr based on the most recent statistics Lr and Rr, SOAP uses potentially different basis matrices Q based previous statisties L, and R, available at every step 1,...
D' (1 - 32) Σ 3 diag (QT9)) 02 (1 (1)
Σ
3r T- ddaag (Q..i) 62 Dr. (24)
Intuitively, this means that the eigenvalue correction is accumulated inconsistently across different coordinate systems. This lead mismatch when preconditioning GT, which is only transformed to the current coordinate system determined by Qr. When the basis remains approximately constant, i.c., Qr x Q: for all t, Equation (24) can be tight approximation and D Dr.
The naive cigenvalue correction may be approximately correct when the basis is updated infrequently, but this be poor approximation if Q does approximate the changing eigenbasis of C, thereby increasing the approximation error. For case 3, the approximation becomes potentially milder because the terms the sum in Equation (24) are down- weighte by 337 -t through the exponential moving average. Therefore, depending 82, the contribution from the eigenvalue correction statistic in previous coordinate systems might be negligible. However, this is not the case for case 2, where all terms weighted equally.
12To address case simply drop 1 32 and -
22

---

In order to address the coor--prcctcc gap between the naive eigenvalue correction (used in SOAP and EShampoo) and the optimal correction in Frobenius norm, we first observe that
T
D' = (1 B2)
ST -t diag (Q Q')
t=1
(1 -2)
31
-t diag (Q Qr)
1=1
(25)
T
diag
Q) - 32)
Σ
8 3 9.g. 1r
diag (QfCrQr)
=: diag
(cr):
Since Qr is orthogonal, we can recursively write
D ddag (Cr)
diag (97( (1- (2)rg)))
diag
(1:3))
(26)
diag
B2 diag
RRTTT- R -1) (3) diag (Q)gr) 02
where RT,T-I = Qr 10.T-1100.0--1111 is the transition matrix between different bases.
While this is the exact expression for our desired quantity, it requires keeping track of the full which is not tractable. Since explicitly
matrix Cr-1 to compute diag( - construct QT- to be close to the best Kronecker-factored basis for Cr-1 through the choice of the Shampoo preconditioner with -T-1 and RT-1, make the additional assumption that Cr-1 is approximately diagonal, i.e., Ör-1 diag(Cr (CT-1) - diag(())=que
Substituting this back into the recursive equation above, we have that
D" = diag (cr) ~ f2 diag (RT.T-1 diag comrected) R R7 + (1 B2) diag (QIgr) 02
= f2 diag RTT-- -T-1 + (1- - B22) diag (QTgr)
corrocted)
02
diag (u (ugorrened)).
(27)
This gives recursive update rule that, in contrast to the naive solution, accounts for the changes of bases between iterations. This also recovers the exact solution when Qr = Q: for all t since R = I. This correction for the changing basis has similar motivation to and resembles parts of the LDAdam algorithm (Robert et al., 2024).
We design experiment to empirically test whether the implicit approximation in Equation (24) helps in practice, and find that interestingly, the correction appears to hurt performance in the settings where we would expect it to help, see Figure 6. A satisfying explanation of this phenomenon remains an open question.
E
Additional experimental details and results
For all experiments we used 1x NVIDIA A100 80GB GPU per run, with the exception of the ImageNet ViT experiments, for which we used 4x NVIDIA A100 80GB GPUs per run. All of the experiments were conducted on an internal compute cluster and we estimate that the total required compute was around 1440 GPU hours or 60 days. We ran additional exploratory experiments not reported in this paper which increases the total compute cost of the project. The implementation of ESShapoo and all other Shampoo variants considered here including Algorithm 4 and Equation (11) for eigh is available at https: : / github.... cwwm/aacebooreeeerreeee. www.arccccotiiizzss....
23

---

VIT (no layer norm parameters)
ConvNeXt V2
2.00
A
L
1.75
1.50
# eigh computations
1.25
0.75
0.50
0.25
0.00
0
1
7
1
2
5
Steps
X10
Steps
X10
Figure 7: Same setting as in Figure 4. (left) When removing all learnable layer norm parameters from the ViT, the number eigendecomposition for L, and RG remain approximately unchanged and compensation appears to be happening. (right) With a ConvNeXt V2 architecture instead of the ViT, the overall pattern is similar, but with a more pronounced discrepancy between L₁ and R.
E.1 Imagewoof experimental details
We use the Imagewoof dataset." 13 All models are trained with cross entropy loss for 100 epochs, using learning rate schedule consisting of linear warmup for 353 steps followed by cosine decay. We a batch size of 128, randomized cropping and horizontal flips as data augmentation, and the default settings for 31 0.9 and 82 0.999. For EShampoo in Figure Figure 4, and Figure we the learning rate = 10 and = 10 10-10
For the vision transformer (VIT) model, heads, a depth of 12 layers, an MLP dimension 1536, dimension of 384, gradient clipping with threshold 1, and weight decay of 10 4, For the ConvNeXt V2 architecture (Woo et al., 2023) weight decay 0.05 and drop paths with rate 0.1.
ViT (Beyer et al., 2022) with patch size 16, 6
For Figure 2, we fix = 10-10 and sweep the following learning rates o:
Vision transformer
AdamW: a e {10-4, 4,3 10
4,6 100110----0 1.10 10 6 10 -0-3 300130-- 10102
Shampoo with grafting: a E 410
10-3)
Shampoo: a E {3- 10 6- 3,
Shampoo with trace scaling: a (0 10
EShampoo: {10-4, 3- 10 2.0
ConvNeXt V2
AdamW: a E {10-1,3 10-4,6 10 10
10-2)
Shampoo with grafting: a E (10 10-3 1..10003110111 10-21
Shampoo: a E {10-4,3
Shampoo with trace scaling: a E (10-1,3 10-4 10-3)
EShampoo: {10-4.3 10 - 3 10 3 3.
The Imagewoof dataset available at https: /tttt/ggtttttttt ctw/fastai/ iittraaimmmmmeeett...m
We fixed based a wide sweep and for AdamW: We also this when grafting from Adam.

---

X10
1.2
A
Attention I
Embodding Tt
Le
Attention R.
Embodding R.
eigh computations
1.0
R.
MLP L.
Last layer L
MLP R.
Last layer R.
0.8
0.6
0.4
0.2
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
x 10
Steps
x10
Figure 8: Llama (324M) trained on 3.2B tokens C4 data with EShampoo (F 1, 0.01). The eigenbases for L, are also in this setting consistently updated more frequently than for R. with the exception the embedding. The eigenbases of A, are less frequently updated than of L,, in contrast to the Imagewoof experiments. There appears to be no difference in the pattern for weight matrices within the attention mechanisms and the MIPss
E.2 More patterns of adaptivity
To the generality of the trends described Section 4.3, we perform the same experiment under few additional settings.
In Figure 7, we remove all learnable layer norm parameters from the vision transformer and consider the ConvNeXt V2 model trained on the same dataset (Imagewoof). The overall pattern consistent across all architectures. When removing the learnable layer norm parameters from the vision transformer, the number of cigendecompositions for L, and Rg stays roughly constant (compare Figure 4 (right) with Figure 7 (left). For the ConvNeXt V2 architecture, the discrepancy between the required updates for L, and R, is even more pronounced: on average, the eigenbases of L, has to be updated significantly more frequently than for R.
To expand beyond the vision modality and cover class imbalance, we also train Llama model with 324 million parameters 3.2 billion tokens of the C4 dataset and conduct similar analysis (Raffel et al., 2023; Grattafiori et al., 2024). We use EShampoo with F = 1 and T = 0.01.
In Figure 8 (left), compute mean and standard errors across all (single) Kronecker factors A, for RMS normalization parameters and all Kronecker factors L, and R, for linear layers and embedding blocks at every iteration. There no bias parameters in this particular model. Consistent with the experiments the vision setting, L, is updated more frequently than RG, with L, updated at almost every iteration, whereas R, updated every other iteration. Unlike the vision setting, the Kronecker factors for the normalization layers A, are initially updated a similar frequency to L,, but diminishes after the first 25% of iterations until it is closer to, but still higher than, the update frequency for R. The standard error for A, is also larger than for L, and R.
In Figure 8 (middle), we consider two subsets of the hidden layers: the four weight matrices in the attention mechanism in each transformer block and the three weight matrices in the MLP in each transformer block. We compute the statistics across all weight matrices for each subset of the hidden weight matrices. There appears to be no significant difference in the number eigendecompositions across steps between the two subsets of the weights.
In Figure 8 (right), we consider the input embedding layer and the last output layer. Due the large vocabulary size, the gradients for these two layers are blocked such that no block has a dimension larger than 8192. Then we precondition each gradient block as usual with L, and R. Here, we compute the statistics across all of these blocks. The trend for the last layer is consistent with all other linear layers in our experiments. However, the eigendecomposition frequency of L, and R, for embeddings are almost identical.
25

---

Table Results for all considered settings on subset of the AlgoPerf workloads. We show the mean and standard error of the steps and time to the targets across the runs that actually hit the targets.
Workload Shampoo Variant
Hits Targets
Steps
Time [min]
Adam grafting (F 100)
4/5
4301 = 109
13.96 + 0.44
EShampoo (F 100)
5/5
2536 66
10.44 = 0.21
FastMRI
EShampoo (F 50)
5/5
2578 86
10.86 + 0.27
EShampoo (F 10)
5/5
2311 73
14.93 + 1.97
EShampoo (F 1)
5/5
2101 + 31
35.34 0.70
C EShampoo
(T 0.1, F = 100)
5/5
2553 + 154
16.33 2.10
EShampoo
(T = 0.1, 1' 50)
5/5
2468 145
10.81
C PShampoo (T 0.1,
10)
5/5
2420 92
10.76 0.73
C PShampon
T 0.1, F 1)
5/5
2367 = 96
10.93 = 0.42
C BShamppo (T 0.01, F F=
5/5
2208 + 41
27.43 + 0.49
Adam grafting (F 100)
1/1
79907
894.27
C BBhampoo (F = 100)
1/1
76226
894.85
ImageNet
VIT
EShampoo
(F = 10)
1/1
T3237
1160.53
EShampoo
(T = 0.1, F 100)
1/1
74010
852.66
C EShamppo
0.1, /' 50)
1/1
77459
935.89
PSShampo
(T 0.01, F 10)
1/1
75841
1188.76
Adam grafting (F' 100)
2/5
12574 + 708
39.20 + 1.88
C EShampoo (F 100)
3/5
8320 = 1203
33.02 + 4.05
OGBG
EShampoo (F = 50)
5/5
7173 + 443
26.17 + 1.31
EShampoo (F = 10)
3/5
6645 357
37.55 + 1.74
EShampoo (F = =1)18
C EShampoo 0.1,, 100)
4/5
8047 369
27.60 + 1.15
C EShampoo (T 0.1, / 50)
5/5
7117 + 328
C EShampoo (T 0.1, F = 10)
5/5
7151 + 416
29.11 1.98
C EShampoo (T 0.1, F = 1)
5/5
6758 273
34.16 1.65
C EShampoo
( 0.01, F 10)
2/5
7234 361
30.15
E.3 AlgoPerf workloads
We follow the standard AlgoPerf setup and consider wall-clock time to pre-specified validation metric targets. See Dahl et al. (2023) and Kasimbeg et al. (2025) for more details on the AlgoPerf benchmark. The FastMRI dataset can be attributed Knoll al. (2020); Zbontar al. (2019), the ImageNet dataset to Krizhevsky et al. (2012), and the OGBG dataset to Hu et al. (2021).
We choose this specific subset of the workloads because (1) transformer (ImageNet ViT), an architecture we use in the small-scale experiments, (2) the FastMRI and OGBG workloads share the same hyperparameter settings with the Imagenet VIT workload, hence excluding them as a confounding factor, and (3) the B₂ used for these workloads is the smallest among all hyperparameter settings of the winning Shampoo submission, resulting the fastest moving average and thereby potentially faster changing eigenbases.
include larger scale vision
We the winning Shampoo submission with Adam grafting, F 100, and the best hyperparameter setting for each workload." 16 For EShampoo, we use the same best-performing hyperparameter setting from Shampoo but turn off learning rate grafting from Adam, and modify F and т.
The benchmark code available https:///gthhnn oot////ooommmoo/ouuuuuessiso..o. aalgoooriii 'The submission is available https://github. cree/main/pel httee//min//reviius__o http:://s_..eeerbrrrrraeg/eppppe wttp:///ww.oo/eoessssiiiiiooeieo wwt.ooneeeeeeeeeeeeeeeeeeeeeeaee httpe//eeeeeeeatttnnnnnaaa/ twning/ saaniingssammoo.... submission.
IThis wall-clock time statistic seems negatively impacted by either an issue with the AlgoPerf code the hardware setup.
Runs failed due :tstts://githh
26

---

0.45
Shampoo grafiing (F
100)
EShampoo (F 100)
0.40
EShampoo (F 10)
EShampoo (T 0.1, F 50)
Loss 4
Training
0.35
Validation
Validation target
0.30
0.73
0.72
SSIM
0.71
0.70
0
10
20
Step
x x10
Wall-clock time [minutes]
Figure FastMRI AlgoPerf workload.
NAdamW
3.0
Shampoo with grafting (F 100)
2.5
EShampoo (F 100)
EShampoo (F 10)
Loss
2.0
EShampoo
1.5
EShampoo (r 0.1, = 100)
Training
1.0
Validation
talida
0.5
0.9
0.8
Accuracy f
0.7
20
40
60
100
120
15
20
Step
x
Wall-clock time [hours]
Figure 10: ImageNet ViT AlgoPerf workload.
27

---

0.14
Shampuo with gralting 100)
EShumpoo (f 100)
0.12
EStumpoo (f 10)
0.10
EShampoo 0.1, 50)
Loss
Training
Validation
0.06
Validation target
0.04
0.02
0.40
0.35
mAP
0.30
0.25
0.20
40
Step
x10
Wall-clock time [minutes]
Figure 11: OGBG AlgoPerf workload.
28

---

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: Every claim is supported by empirical and theoretical evidence in Section and Section 4.
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 is fine to 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 address limitations where appropriate and add specific discussion of some important limiations in Section
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 a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn 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 on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be 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 low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
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 faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 complete (and correct) proof?
29

---

Answer: [Yes]
Justification: We provide proofs for all lemmas and propositions in Appendix with an appropriate level of details and formality.
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 to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: While reproducing the exact numerical results will require the source code, we provide sufficient details to reproduce all the main claims of the paper in Section 3, Section 4, and Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, 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 to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to 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 reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
30

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: All used datasets are open access. We do not provide the source code to reproduce all the experiments, but the implementation of all Shampoo variants is available at and the AlgoPerf submission code for Shampoo is available at https: com commmmlooeeooooooooooooomm commons/ submiss ons......... ons_ algori thms wttps//tteeeemmaeeeeeeeeeeeaeaet pwwpaniiinnnevioussss.muu Jeaderbooards/llggorrr...... http:///eeeeeeeeeeooooeaaaaaaaar wttppperr__oo//uiiissisiiioo external_ external_tuning/eeee Btctrrpoo_uummsssseeee submission.
https://gii /ttggtttubb.... httm///ceeoboeeeeeerrrtttttttiii
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
wwt.iiisssisoiiiooooooocimmomiil for more details. public/ wwwccc/gguies//eeee Soaeee///ooeSuuisss..eo....
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new 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. wwtpccccpppcceeeecccc..ccccccltt gww.//////////eegeeoee /CodeSubmiaiee
/tteg/guuiieeeeeeeummmiiieeeeee wwweiissssoooooooooooooooo onPolicy) for more details.
The authors should provide instructions on 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 as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
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?
Question:
Answer: Yes]
Justification:
The main section Section 3 and Section 4 provide all necessary details to understand the presented results and claims.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
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: We do report standard emrors of multiple runs for the AlgoPerf results in Table and Table 3, with exception of the ImageNet ViT workload due to the larger computational cost. However, we do not report statistics over multiple runs for the sweeps presented Figure I, Figure 2, and Figure 3 due to the large number of runs and the associated computation cost.
31

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, 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 run 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 or 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% CL, 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 out of range (e.g. negative error 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 in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer:[Yee
Justification: We do provide details on the cluster used for the experiments and an estimate of the total amount of compute in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or 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 as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't ' make it into the paper).
paper (e.g., preliminary or failed experiments that NeurlPS Code of Ethics settps//neuuriii wwtpccccc//cpccttieeeeee hta.bliiceEhhcccuudeeeeeeee
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the
Answer: [Yes]
Justification: The paper covers fundamental algorithmic research and conforms with the all points of the NeurlPC Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
32

---

Answer: [NA]
Justification: The work here is of fundamental nature and does not directly impact society. However, it might contribute to improved training algorithms which could potentially reduce the cost or improve the performance arbitrary deep learning problems, including problems that impact society.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 tied to particular applications, let alone deployments. However, if there is 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, is not needed to point out 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 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 or unintentional) misuse of the technology.
If 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
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)?
Question:
Answer: [NA]
Justification: The work covers fundamental algorithmic research.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [Yes]
Justification: We do credit the creators the datasets, model architectures, and the AlgoPerf benchmark in Appendix E.
Guidelines:
33

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if 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 (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets 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 it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 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:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooel 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 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 details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in 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 in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or 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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
34

---

Guidelines:
The answer NA means that the paper does not 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 expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non non-standard components.
Please refer to our LLM policy (https://neuria) //eerrppp.ccc/oooffeeeeee/////// for what should or should not be described.
35


Reviews:
[{"review_id": "IbtynZH65c", "reviewer": "Reviewer_aEP3", "summary": "This work studies the Shampoo optimizer and two key heuristics (learning rate grafting and preconditioning frequency) important for its competitive performance. Based on empirical results and theoretical motivation, the authors claim that the learning rate grafting addresses stale and mis-scaled eigenvalues of the preconditioner, and this need can be removed by decoupling the eigenvalue updates from the eigenbasis. The authors also propose a method tracking approximation-error of the Kronecker factors to adaptively determine eigenbasis update frequency. Their method also allows one to track which components of the network require more frequent eigendecomposition computations, showing that biases and layer normalization parameters require more frequent updates.", "questions": "* How would the authors’ proposed adaptive criterion of the QR algorithm compare to having frequency warmup (starting with frequency = 1 and roughly gradual increase to F = 50) on all of the parameters, or a fraction of the parameters (in terms of loss vs wall-clock time)? \n* Do the authors have more intuition about the observed need for frequent updating for layer norm parameters? If the LayerNorms did not have learnable parameters, would the matrix layers then require more frequent computations (the LayerNorm parameters act as some output scale for the given layer)?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n* The paper is clearly written and easy to follow. The relevant background and algorithms were presented particularly well.\n* Connecting learning rate grafting to previously known modifications of Shampoo (using Shampoo^2 with trace scaling, low PF, eigenvalue correction a la SOAP) is a nice result.\n* Investigations about determining eigenbasis computation frequency, its effect on different network parameters, and plotting tradeoffs between efficiency and performance is an important direction.\n\nWeaknesses:\n* The empirical setting is limited to two architectures on a single dataset and it would’ve been particularly interesting to see how the eigenbasis frequency results in eg. Figure 4 change for other domains, relating to previous work similarly identifying key network parameters requiring adaptivity due to issues like class imbalance for the last layer for instance [1,2].\n* The authors’ proposed adaptive QR algorithm provides efficiency gains but some other baselines may better inform whether this adaptive control based on thresholding the approximation error is necessary - please see Questions.\n\nReferences:\n\n[1] Kunstner, Frederik, et al. \"Heavy-tailed class imbalance and why adam outperforms gradient descent on language models.\" Advances in Neural Information Processing Systems 37 (2024): 30106-30148.\n\n[2] Zhao, Rosie, et al. \"Deconstructing what makes a good optimizer for language models.\" arXiv preprint arXiv:2407.07972 (2024).\n\nMinor: \n* Authors should add a reference to Appendix C after presenting Lemma 1 and Proposition 1 in the main paper.", "was_revised": true, "final_justification": "The authors have addressed my concerns regarding other domains, comparison to a frequency baseline warmup, and ablating the influence of LayerNorm parameters. I believe the empirical results are well-presented and analyses like these on Shampoo will be of interest to the community.", "rebuttal": "Thank you for your comments and raising insightful questions. We address the weaknesses and questions below. We unfortunately cannot share figures for the new experimental results due to changes in rules this year. However, we will add all new results to the paper.\n\n\n> The empirical setting is limited to two architectures on a single dataset and it would’ve been particularly interesting to see how the eigenbasis frequency results in eg. Figure 4 change for other domains, relating to previous work similarly identifying key network parameters requiring adaptivity due to issues like class imbalance for the last layer for instance [1,2].\n\nWe consider three additional dataset-architecture pairs covering different modalities like graph classification in the experiments presented in Table 1+3 (ImageNet ViT, FastMRI, and OGBG workloads from the AlgoPerf benchmark). However, we agree that our analysis is currently **limited to a single model in Figure 4**.\n\nTo address this limitation, we conduct a **similar analysis for the ConvNeXt V2** model trained on the Imagewoof dataset. We observe a similar trend: eigenbases for 1D parameters are updated most frequently, followed by eigenbases for $L_t$, lastly followed by eigenbases for $R_t$, which are updated least frequently. This order is consistent across at least these two architectures on the same dataset.\n\nTo additionally expand beyond the vision modality and cover the mentioned class imbalance, we train a **Llama 3 model with 160M parameters on 3.2B tokens from the C4 dataset** with $\\tau=0.01$ and conduct a similar analysis. We compute mean and standard errors across all (single) Kronecker factors $A_t$ for RMSNorm parameters and all $L_t$ and $R_t$ Kronecker factors for linear layers and blocks of the embedding (there are no bias parameters) at every iteration. Consistent with both experiments in the vision setting, $L_t$ is updated more often than $R_t$. On average, $L_t$ is updated almost every iteration, whereas $R_t$ is updated closer to every other iteration. The Kronecker factors $A_t$ for the normalization layers are first updated at a similar frequency to $L_t$, but, **in contrast to the vision setting**, the update frequency decreases after the first ~25% of iterations until it is closer to, but still higher than, the update frequency for $R_t$ at the end. The standard error for $A_t$ is also larger than for $L_t$ and $R_t$.\n\n\n> Authors should add a reference to Appendix C after presenting Lemma 1 and Proposition 1 in the main paper.\n\nAgreed. We have added this missing reference, thanks for pointing this out.\n\n\n> The authors’ proposed adaptive QR algorithm provides efficiency gains but some other baselines may better inform whether this adaptive control based on thresholding the approximation error is necessary - please see Questions.\n\n> How would the authors’ proposed adaptive criterion of the QR algorithm compare to having frequency warmup (starting with frequency = 1 and roughly gradual increase to F = 50) on all of the parameters, or a fraction of the parameters (in terms of loss vs wall-clock time)?\n\nThank you for the suggestion. While you are correct that we could schedule $F$, one could also **consider a schedule for $\\tau$**. We argue that this is the right comparison. Therefore, your question raises an important issue: whether scheduling $F$ or scheduling $\\tau$ is more useful for controlling efficiency.\n\nWe argue that $\\tau$ enables more fundamental control. Specifically:\n- Setting the frequency $F$ fixes the computational cost (i.e., number of eigendecompositions) across all factor matrices and implicitly determines how the approximation quality $\\tau$ varies for each Kronecker factor.\n- In contrast, setting the threshold $\\tau$ directly controls the approximation quality and implicitly leads to a varying frequency $F$ schedule for each Kronecker factor.\n\nIn general, we believe that controlling inexactness via the approximation quality threshold $\\tau$, rather than the cost (frequency of eigendecomposition $F$), better captures its potential impact on Shampoo's convergence. For instance, in the extreme case where a Kronecker factor’s eigenbasis does not change at all, $F$ provides no mechanism to skip unnecessary computations, resulting in wasted cost without convergence benefits.\n\nTo empirically validate this intuition, we conducted an **additional experiment** on the Imagewoof ViT workload using a simple two-phase schedule for $F$ and $\\tau$:\n- We set $F=1$ for the first 10% of iterations and $F=1000$ for the remaining 90%.\n- We set $\\tau=0.01$ for the first 10% of iterations and $\\tau=0.8$ for the remaining 90%.\n\nBoth schedules yielded comparable final loss. However, the $\\tau$-based schedule resulted in significantly fewer total eigendecompositions and thus reduced overall wall-clock time.\n\nWe will include a subsection in the paper discussing these results in more detail.\n\n\n> Do the authors have more intuition about the observed need for frequent updating for layer norm parameters? If the LayerNorms did not have learnable parameters, would the matrix layers then require more frequent computations (the LayerNorm parameters act as some output scale for the given layer)?\n\nThank you for the question. We do not have a satisfactory explanation for this observation.\n\nWe also **retrained the ViT without learnable LayerNorm parameters on Imagewoof** and created the same plot from Figure 4 (right). Interestingly, removing the learnable LayerNorm parameters does not lead to any significant difference in the number of eigendecompositons for $L_t$ and $R_t$."}, {"review_id": "E7YhQoClMX", "reviewer": "Reviewer_QTer", "summary": "The paper revisits practical tricks that have become standard when training with Shampoo: spacing out the eigendecomposition of its Kronecker factors and applying learning-rate grafting to match Adam’s step size. It shows that grafting compensates for eigenvalue issues: staleness between decompositions and a scale bias tied to layer shape, while the eigenvectors largely stay relevant. Replacing those stale eigenvalues each iteration with Adam-style second-moment estimates removes the need for grafting, and a simple Frobenius-norm test decides when the costlier eigenbasis must be refreshed. Experiments confirm that this “purified” variant matches or surpasses the grafted baseline with less compute.", "questions": "Could you please clarify how Eigenvalue-Corrected Shampoo meaningfully differs from SOAP? Appendix B of Anil et al. (2020) notes that the Frobenius-optimal diagonal Kronecker approximation in a fixed Shampoo eigenbasis is precisely the Adam second moment, and SOAP implements this by keeping the Shampoo basis (refreshed every F steps) while applying an Adam/Adafactor-style diagonal update at each iteration, so it seems EShampoo follows the exact same recipe.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 2, "strengths_and_weaknesses": "The paper addresses a highly motivated problem. The paper offers an interesting explanation of why grafting is needed in standard Shampoo. By separating the Frobenius-norm effect of eigenvalues from the rotational effect of eigenvectors, they show that grafting merely rescales updates that have become too small or too large because the eigenvalues are both out-of-date. This clarifies a long-standing practical mystery.\n\nWhile your analysis of grafting’s role in Shampoo and the supporting experiments are insightful, it’s still unclear what new algorithmic insight EShampoo contributes beyond SOAP, which already applies an Adam-style diagonal second moment in the fixed Shampoo eigenbasis.", "was_revised": true, "final_justification": "The grafting analysis shows clear insight making the paper important and interesting.", "rebuttal": "Thank you for the comments and clarifying questions. We will address the questions below.\n\n\n> While your analysis of grafting’s role in Shampoo and the supporting experiments are insightful, it’s still unclear what new algorithmic insight EShampoo contributes beyond SOAP, which already applies an Adam-style diagonal second moment in the fixed Shampoo eigenbasis.\n\nOur main objective in the first part of the paper (Section 3) is to demonstrate how eigenvalue correction methods (like SOAP) can be used to remove learning rate grafting which is a crucial heuristic in Shampoo -- we have not seen these observations discussed in-depth in the literature. We do not claim to propose eigenvalue-corrected Shampoo as a new algorithm and therefore **introduce it in the paper's background section** (Section 2.2); to our knowledge, eigenvalue correction was initially proposed in Liu et al. (2018) and George et al. (2018) for K-FAC and in Anil et al. (2020) for Shampoo.\n\n\n> Could you please clarify how Eigenvalue-Corrected Shampoo meaningfully differs from SOAP? Appendix B of Anil et al. (2020) notes that the Frobenius-optimal diagonal Kronecker approximation in a fixed Shampoo eigenbasis is precisely the Adam second moment, and SOAP implements this by keeping the Shampoo basis (refreshed every F steps) while applying an Adam/Adafactor-style diagonal update at each iteration, so it seems\n\nThank you for the question. We unfortunately overloaded the term _eigenvalue-corrected Shampoo_ in the paper:\n1. We use it to refer to the family of methods that apply an eigenvalue correction to Shampoo's preconditioner. Hence, SOAP is an instance of eigenvalue-corrected Shampoo.\n2. We use its abbreviation EShampoo to refer to the method we implemented for the experiments. This is also an instance of eigenvalue-corrected Shampoo.\n\nWe will clarify this terminology in our revision.\n\nFrom an algorithmic perspective, EShampoo as it is implemented in the experiments differs from SOAP in two respects (as discussed in Appendix B):\n1. SOAP delays the update of the preconditioner until after the update of the weights.\n2. SOAP uses a warm-started single iteration of the QR algorithm. EShampoo uses eigendecompositions via `torch.linalg.eigh` unless described otherwise (Figure 3 (left)).\n\nRegarding the optimality in Frobenius norm, as far as we know, **none of the practical implementations of eigenvalue corrections for Shampoo actually use the optimal solution** to approximate full-matrix Adam. SOAP/EShampoo (EShampoo referring to the method implemented for our experiments) make an additional approximation (Equation 20):\n$$(1 - \\beta_2) \\sum_{t=1}^T \\beta_2^{T-t} \\mathrm{diag}( Q_T^\\mathrm{T} g_t )^{\\odot 2} \\approx (1 - \\beta_2) \\sum_{t=1}^T \\beta_2^{T-t} \\mathrm{diag}( Q_{t}^\\mathrm{T} g_t )^{\\odot 2}.$$\n\nNote the optimal solution (left) uses a constant $Q_T$ across iterations whereas the approximation (right) uses a different eigenbasis $Q_t$ for each iteration. This approach ignores the fact that we are accumulating gradients with respect to different eigenbases, especially when $F$ is small; see Appendix D.1 for further discussion of this approximation and an alternative eigenvalue correction."}, {"review_id": "3NPvtNqQQp", "reviewer": "Reviewer_2S8j", "summary": "This paper studies the practical tricks of the optimizer Shampoo to address the critical problem of stale updates. The paper demonstrates theoretically that the stale update of Shampoo can lead to eigenevalue scale mismatch, resulting in significantly worse performance, which can be fixed by tricks like grafting and correcting eigenvalues (e.g., SOAP). Then the authors discuss the effect of stale eigenbasis and propose an adaptive way to update the eigenbasis to solve the problem. Empirical results are also provided to verify the claims.", "questions": "1. Based on the explanation for why grafting is useful for Shampoo, I am curious why don't we just use $\\sqrt{mn}$ instead of $||U_t^{Adam}||_F$ in grafting for a matrix $W\\in\\mathbb{R}^{m\\times n}$, since they should have similar scales. Have you tried this kind of trivial scaling?\n2. I want to double check what the wall-clock time stands for in e.g., Figure 3. Is it the total training time including forward and backward or any other measure? Also, could you post the proportion of the optimizer step time in the whole training process? Based on the figure, the optimizer step computation takes the major time costs, which seems a bit strange, even though matrix operations are required.\n3. The observation that $L_t$ seems to update more than $R_t$ looks interesting. Do you have more evidence for this in other models or other scales? Also, any intuition for explaining this would be welcome.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n1. The paper is well-written and basically clear for understanding.\n2. The viewpoint of eigenvalue scale mismatch for explaining the effectiveness of grafting and SOAP is interesting, providing a novel and valid perspective for the motivation of these Shampoo variants. \n3. The study of the influence of the preconditioner staleness is detailed and solid, providing a comprehensive understanding of the problem.\n\nWeakness:\n1. The theoretical perspective in Section 3 on why grafting is useful may not be that rigorous and convincing. The bound in Lemma 1 and Proposition 1 is very loose, since the largest entry value/singular value can be much larger than the smallest ones. I don't think this bound is valid for justifying that the eigenvalue scale mismatch is the problem of Shampoo without grafting. \n2. In Section 4, the adaptive update rule of the eigenbasis is claimed to propose for reducing the effort of manually tuning the update frequency. However, it turns out that the adaptive threshold $\\tau$ also needs tuning and I don't see strong evidence in that the optimal choice can transfer over models and scales. This may make the idea not that powerful.", "was_revised": true, "final_justification": "Based on the rebuttal, my major concerns are addressed. I would raise my score to 5 since the paper explores some interesting and intuitive facts about the Shampoo optimizer, and the empirical results are solid.", "rebuttal": "Thank you for your comments and interesting questions. We address the weaknesses and questions below. We unfortunately cannot share figures for the new experimental results due to changes in rules this year. However, we will add all new results to the paper.\n\n\n> The theoretical perspective in Section 3 on why grafting is useful may not be that rigorous and convincing. The bound in Lemma 1 and Proposition 1 is very loose, since the largest entry value/singular value can be much larger than the smallest ones. I don't think this bound is valid for justifying that the eigenvalue scale mismatch is the problem of Shampoo without grafting.\n\nWe agree that the bound is not tight. However, our aim is not to prove the tightest bound possible, but to demonstrate Shampoo's intrinsic dependence on the parameter dimensions $m, n$ to **motivate our experimental study** on Shampoo's layer-wise mis-scaling. While it is possible to derive tighter, more complex bounds with weaker assumptions, we felt that this would digress from the primary goal of this section.\n\n\n>In Section 4, the adaptive update rule of the eigenbasis is claimed to propose for reducing the effort of manually tuning the update frequency. However, it turns out that the adaptive threshold also needs tuning and I don't see strong evidence in that the optimal choice can transfer over models and scales. This may make the idea not that powerful.\n\nWe do not claim that using the adaptive update rule reduces the need for manual hyperparameter tuning. However, tuning the threshold $\\tau$ instead of an update frequency $F$ (both global scalar hyperparameters) enables us to directly control the approximation quality of each Kronecker factor and independently determine each factor's implicit update frequency schedule. Controlling the approximation quality directly lays the groundwork for studying its impact on convergence.\n\nFrom our perspective, two major challenges remain that prevent us from fully removing this hyperparameter (lines 309-312 in Section 5):\n1. At this stage, we do not have a way of analyzing the impact of approximation quality or inexactness on Shampoo's convergence.\n2. The wall-clock time overhead of the eigendecompositions relative to the rest of training is highly dependent on the training setup, e.g. model, training pipeline, and hardware.\n\n\n> Based on the explanation for why grafting is useful for Shampoo, I am curious why don't we just use $\\sqrt{mn}$ instead of $||U_t^{Adam}||_F$ in grafting for a matrix $W \\in \\mathbb{R}^{m \\times n}$, since they should have similar scales. Have you tried this kind of trivial scaling?\n\nThank you for the suggestion! We have not tried this re-scaling previously. We expect that it can, depending on the setting, lead to stable training but Adam(W)'s optimal learning rate will generally not transfer.\n\nTo test your idea, we **implemented this re-scaling** for the Imagewoof ViT workload. With AdamW's optimal learning rate the loss does not diverge, but the final solution is even worse than the one found with AdamW. However, by decreasing the learning rate by an order of magnitude, we can indeed reach close to the same final loss as with $||U_t^{Adam}||_F$ in this setting.\n\n\n> I want to double check what the wall-clock time stands for in e.g., Figure 3. Is it the total training time including forward and backward or any other measure? Also, could you post the proportion of the optimizer step time in the whole training process? Based on the figure, the optimizer step computation takes the major time costs, which seems a bit strange, even though matrix operations are required.\n\nThe wall-clock time is total training time including data loading, forward and backward passes, and optimizer step. We have confirmed that the main overhead for this small model and dataset are indeed the eigendecompositions, which lead to drastic differences in the wall-clock time. Therefore, the exact proportion of the optimizer step in the overall training time strongly depends on the values of $F$ and $\\tau$ (larger for small $F$/$\\tau$).\n\n\n> The observation that $L_t$ seems to update more than $R_t$ looks interesting. Do you have more evidence for this in other models or other scales? Also, any intuition for explaining this would be welcome.\n\nPlease refer to our response to Reviewer aEP3. We have performed a similar analysis to the one in Figure 4 for ConvNeXt V2 on the same dataset (Imagewoof) and a Llama 3 model with 160M parameters trained on 3.2B tokens of the C4 dataset, and found that the **pattern for $L_t$ and $R_t$ is consistent across all settings**. \n\nWe do not have a strong explanation for this phenomenon at this time. One potentially helpful interpretation is viewing $L_t$ and $R_t$ as estimating the uncentered covariance of the gradients of the loss with respect to the layer's outputs and the uncentered covariance of the layer's inputs, respectively. (See Appendix B in [1] for more details.)\n\n[1] Anil et al. [Scalable Second Order Optimization for Deep Learning](https://arxiv.org/pdf/2002.09018) (2020)"}]

Author rebuttal:

Meta-review: This paper investigates the Shampoo optimizer, with a particular focus on learning rate grafting and the eigenbasis update frequency, two tricks that are seen as crucial for strong empirical performance. The authors find that grafting addresses the problem of staleness of eigenvalues incurred by infrequent updates. The authors further propose an adaptive update criterion based on approximation error, which allows different components of a model to refresh at different rates, and they demonstrate efficiency gains without significant loss of performance. The reviewers found the paper to be well-written and clear, offering insightful connections between Shampoo heuristics and variants (such as SOAP). The empirical examination of eigenbasis update frequency and the adaptive approach, was appreciated. Weaknesses included the initially narrow empirical scope, the looseness of some theoretical bounds and the concern that the adaptive threshold still introduces a tunable hyperparameter. Rebuttal directly addressed these issues by presenting new results on ConvNeXt V2, ViT on Imagewoof, a Llama model on C4 and clarifying that threshold-based control offers more fundamental advantages than fixed-frequency schedules.

Finally all reviewers recommended accept and I also find the paper to be a solid, insightful contribution that will be of clear interest to the optimization and deep learning community, and acceptance is strongly justified.
Decision: Accept (spotlight)