SOAP: Improving and Stabilizing Shampoo using Adam for Language Modeling
Authors: Nikhil Vyas
Depen Morwani
Rosie Zhao
Itai Shapira
David Brandfonbrener
Lucas Janson
Sham M. Kakade
Venue: ICLR 2025 Poster (2025)
Keywords: Shampoo
Adam
Second Order Optimizer
Primary area: optimization
TL;DR: We design and empirically study a new second order optimizer called SOAP, which runs Adam in the eigenbasis provided by Shampoo.

Abstract: There is growing evidence of the effectiveness of Shampoo, a higher-order preconditioning method, over Adam in deep learning optimization tasks. However, Shampoo's drawbacks include additional hyperparameters and computational overhead when compared to Adam, which only updates running averages of first- and second-moment quantities. This work establishes a formal connection between Shampoo (implemented with the 1/2 power) and Adafactor --- a memory-efficient approximation of Adam --- showing that Shampoo is equivalent to running Adafactor in the eigenbasis of Shampoo's preconditioner. This insight leads to the design of a simpler and computationally efficient algorithm: **S**hampo**O** with **A**dam in the **P**reconditioner's eigenbasis (SOAP).
With regards to improving Shampoo's computational efficiency, the most straightforward approach would be to simply compute Shampoo's eigendecomposition less frequently. Unfortunately, as our empirical results show, this leads to performance degradation that worsens with this frequency. SOAP mitigates this degradation by continually updating the running average of the second moment, just as Adam does, but in the current (slowly changing) coordinate basis. Furthermore, since SOAP is equivalent to running Adam in a rotated space, it introduces only one additional hyperparameter (the preconditioning frequency) compared to Adam. We empirically evaluate SOAP on language model pre-training with 360m and 660m sized models. In the large batch regime, SOAP reduces the number of iterations by over 40\% and wall clock time by over 35\% compared to AdamW, with approximately 20\% improvements in both metrics compared to Shampoo. An implementation of SOAP is available at https://github.com/nikhilvyas/SOAP.

Published as a conference paper at ICLR 2025
SOAP: I IMPROVINE AND STABILIZING SHAMPOO US- ING ADAM FOR LANGUAGE MODELING
Depen Morwani: Harvard Universsiy
Nikhil Vyas* Harvard University
Rosie Zhaot
Itai Shapira i
Harvard University E
Harvard University
David Brandfonbrener
Lucas Janson
Kempner Institute Harvard University
Harvard University
Sham Kakade
Kempner Institute at Harvard University
ABSTRACT
There is growing evidence of the effectiveness of Shampoo, a higher-order pre- condition conditioning method, over Adam deep learning optimization tasks. How- ever, Shampoo's 's drawbacks include additional hyperparameters and computa- tional overhead when compared to Adam, which only updates running averages of first- and second-momets quantities. This work establishes formal connec- tion between Shampoo (implemented with the 1/2 power) and Adafactor - a memory-efficient approximation of Adam showing that Shampoo is equivalent memmone
to running Adafactor in the eigenbasis of Shampoo's 's preconditioner. This insight leads the design of a simpler and computationally efficient algorithm: ShampoO with Adam in the Preconditioner's s eigenbasis (SOAP). With regards to improving Shampoo's computational efficiency, the most straightforward approach would be to simply compute Shampoo's 's eigendecomposition less frequently. Unfortunately, as our empirical results show, this leads to performance degradation that worsens with this frequency. SOAP mitigates this degradation by continually updating the running average of the second moment, just as Adam does, but in the current (slowly changing) coordinate basis. Furthermore, since SOAP is equivalent to running Adam in a rotated space, it introduces only one additional hyperparam- eter (the preconditioning frequency) compared to Adam. We evaluate SOAP on language model pre-training, with experiments on 360m and 660m sized models. In the large batch regime, SOAP reduces the number of iterations by over 40% and wall clock time by over 35% compared to Adam W, with approximately 20% im- provements in both metrics compared to Shampoo. An implementation of SOAP is
s available at https:: /gtpg/gtttttt////////////te com/ miwhiikhhiiio Shanhhilv www.syaaas.. S
PETTRDDUCTON IOROODUCTION
With w over-increasing costs of LLM training, optimization efficiency has become central question
in the field of deep learning. Several recent works have tackled this challenge by addressing both the memory (Zhao et al. 2024a; Wang al. 2024) and compute (Anil al.) 2020) footprint of
optimizers. In Algoperf (Dahl et al. 2023). recent optimization efficiency benchmark, Sham- optimmen poo (Gupta et al., al.]|20| 20) second-order algorithm, outperformed all other submissions, including
Adam (Kingma & Ba 2015), reducing wall-clock time by 28% MML.Commons 2024). Higher-
(icecondiin order Dresonditioning has also been applied in large-scale training runs, such as 2emini-1.3 Flash
(Gemini Team 2
The success of Shampoo has drawn increasing attention from the deep learning community. Several works have explored ways to scale Shampoo by improving
prrooviigg its memory and compute efficiency
"Equal contribution. Correspondence to mikhileg earvard.edu.. edd..
1 Equal contribution.

---

Published as a conference paper at ICLR 2025
Wang et al. 2024 Anil al. 2020 Shi al. 2023). Other research (Morwani al. 2024 has examined the theoretical foundations of Shampoo and proposed minor adjustments (such as using power 1/2 rather than 1 1/4) that align with prior empirical findings (Anil et al. 2020). Moreover, Morwani et al. (024) also showed that Shampoo with the aforementioned modifications close to the optimal Kronecker approximation of the Adagrad (Duchi et al. 201 optimizer.
Our first contribution is demonstrating that the variant of Shampoo proposed by Morwani et al. (2024) is equivalen running Adafactor (Shazeer & Stern 2018 Zhai al. 2022) the eigenba- sis provided by Shampoo's preconditioner. This interpretation of Shampoo connects it to a broader family of methods (e.g. (George et al. 2018)) that design second-orde algorithms by running a first- order method in the eigenbasis provided by second-orde method. Building on this insight, we can explore a broader design space for combining first and second order methods. Many of our design choices are a synthesis of conceptual ideas from prior works of George et al. (2018); 8): AAnil et al. (2020); Morwani et al. (2024) as well as implementation ideas from works of Wang al. (2024);
Zhao et al. (2024a).
Explicitly, we study SOAP (ShampoO with Adam the Preconditioner's eigenbasis) an algorithm that runs AdamW in the eigenbasis provided by Shampoo. Our main contributions are follows:
We make a formal connection between the Shampoo and the Adafactor algorithm. This insight leads to consider the SOAP algorithm, which runs AdamW in the preconditioned space provided by Shampoo.
SOAP outperforms both Shampoo and Adam in language model pre-training tasks with model sizes 360m and 660m, even after extensive hyperparameter tuning of Shampoo.
SOAP reduces the number of hyperparameters compared Shampoo, resulting only one additional hyperparameter compared to Adam W: preconditioning frequency.
SOAP demonstrates greater robustness to large preconditioning frequency compared to Shampoo on language model pre-training tasks.
We should also note that while similar algorithmic variants have been discussed in the literature (e.g. the appendix Anil et al. (2020)), we the first to systematically evaluate it.
Organization: In Section 3] we discuss related works. In Section 4 we start by showing an equiv- alence between Shampoo (with exponent 1/2) and rumning Adafactor in the eigenspace given by Shampoo, then with this equivalence as the starting point we describe SOAP. In Section 5 we pro- vide experimental methodology and in Section [6] we compare the performance of Adam W, Shampoo and SOAP on language modeling tasks. In Appendices B.2 and B.3 we discuss the the space and time complexity of SOAP and how it can be improved. In Appendix C we show that efficiency benefits of SOAP over Adam are maintained for longer duration runs where #tokens 100 model size.
2 NOTATION AND BACKGROUND
We denote the weight matrix of a neural network layer by W E R and the corresponding gradient by G Rmxn, At given time step I, these are denoted as W, and G., respectively. For batch of inputs at time I, denoted by B,, the loss and its gradient evaluated at W, are represented as бB. (Wr) and V W ΦB, (W,), respectively.
Adagrad (Duchi et al. 2011b) is an online learning second-order algorithm that maintains precon- ditioner H guanxmn If the vectorized gradient at time t is denoted by gt (i.c., St = vec(G.) e Rmm). then the update of the preconditioner and the vectorized weights wt Rmn with learning rate 77 is given by
H, H,-1 9:9 = w:-1 ,,,,, gt
Adam (Kingma & Ba 2015), a widely used first-order optimization algorithm in deep learning is diagonal approximation of Adagrad. It maintains an exponential moving average of the gradients
'Given this connection, the results of Morwani al. 20224) can he interpreted as showing that the eigenbasis provided by Shampoo's preconditioner is close the "optimal" hasis for running Adafactor.
2

---

Published as a conference paper at ICLR 2025
660m, 2m batch size Preconditioning #requency=10
660m, 2m batch size Preconditioning Frequency=10
360m, 2m batch size
3.2
3.2
3.1
3.1
2.92
2.90
Voan Locs
2.9
2.9
2.80
2.86
2.7
2.7
2.82
2.6
2.6
1600
3200
4800
6400
0.25
0.5
0.75
1.0
10
100
Training Steps
Wall Time (scaled by AdamIN
Presonditioning Frequency
Adami
Shampoo
SOAP
SOAP (shorter UR schedule)
Figure 1: Comparing performance of tuned runs for AdamW, Shampoo (using DistributedSham- poo (Shi et al.] 2023) implementation) and SOAP. In left and middle figures, Shampoo and SOAP preconditioning frequency of 10. The "shorter LR schedule" plot is where tuned the cosine decay so as to achieve the same terminal performance as Adam W. There we observe a 2 40% re- duction in the number of iterations and 35% reduction in wall clock time compared to Adam W, and approximately 20% reduction in both metrics compared to Shampoo. In the right figure we ablate preconditioning frequency and observe a slower degradation of performance of SOAP as compared to Shampoo. See Section O for a discussion of experimental results and ablation of batch size and Section for experimental methodology.
G, (denoted as M,) and of element-wise squared gradients Gū (denoted as Vt) for a given weight matrix W. Its update rule with learning rate is given by
M VV.
W, + W 17
where the division is performed element-wise.
Adafactor (Shazeer & Stern 2018 Zhai et al. 2022 a variant of Adam, replaces V with its best rank-1 approximation V. to reduce memory usage. While the original Adafactor paper (Shazeer & Stern] 2018) proposed additional modifications, such as changes to the learning rate schedule, we focus on the version of Adafactor proposed in recent works (Zhai et al. 2022 Zhao et al. 2024c), whose update with learning rate ry is given by
M VV!
W, < W
Shampoo (Gupta et al. 2018b) is second-orded optimization algorithm that approximates Adagrad and maintains two preconditioners, L, Rmxm and RG E gaxx, for a given weight matrix W Rmxn, The updates for the preconditioners and the weights with learning rate are follows:
L L, G,G.; R₂ Rd + G. G.; W < W nT. 1/4 G. 1/4
In practice Shampoo is implemented with several other modifications such as layerwise learning rate grafting and exponents other than -144. We will use the DistribuedShampoo (Shi et al. 2023) implementation which has these variations available as hyperparameters.
3 RELATED WORK
We begin by discussing works that are closely related, including George et al. (2018): Anil et al. 2020 and Zhao et al. (2024a). Subsequently,
extended related works.
KFAC Martens & Grosse 2013) is a well-known second-order optimization algorithm designed for neural networks. E-KFAC (George al. 2018) builds upon KFAC in manner analogous to
3

---

Published as a conference paper at ICLR 2025
extension of Shampoo, introducing diagonal preconditioner that is updated between KFAC inver- sion steps. However, E-KFAC's algorithm is not identical to running Adam in KFAC's eigenbasis, as the diagonal preconditioner is not Adam.
Amil et al. (2020) introduced several algorithmic and numerical improvements to develop a practical and scalable version of Shampoo (Gupta et al. 2018b). Notably, they empirically found that using an exponent of 1/2 outperforms the original exponent of 1/4 in Shampoo. Of particular interest to our work is Appendix B of Anil et al. (2020), where, inspired by E-KFAC. they describe an algorithm that is essentially equivalent to SOAP for 2D layers. However, no experiments were provided, and the authors claimed that unpublished experiments showed no empirical improvement over Shampoo. This discrepancy between our findings may be due to some of the implementation details of SOAP.
GaLore (Zhao al. 2024a) was recently proposed as a method to reduce Adam's memory footprint by maintaining momentum in a low-rank subspace derived from the singular value decomposition (SVD) of the gradients. Their algorithm's full-rank version bears similarity to ours, with some notable distinctions. Firstly, their projection subspace is determined by the SVD of the current gradient, while we maintain an exponential moving average of GG and GH G. Secondly, we retain momentum in the original space and project it onto the preconditioned space, whereas they maintain it in the preconditioned space and do not rotate it each time the preconditioned space is updated. In Appendix [ we study GaLore's 's performance and find that modifications are necessary for improving upon Shampoo. Moreover, their method only projects one side of a layer using the eigenbasis while using the identity basis the other side. We examine the impact of one-sided projection for SOAP in Appendix B.1
Diagonal Preconditioning based Optimizers: Other than AdamW, there are other optimizers which involve diagonal preconditoning such as Lion (Chen et al. (023). Sophia (Liu et al. 2024), and Adafactor (Shazeer & Stern 2018). Recent works Kaddour et al. (2023); Zhao al. (202cc)
showed that these optimizers perform comparably to Adam for LLM pretraining but do not sur- pass it. This suggests the need to explore non-diagonal preconditioners. We discuss prior works on non-diagonal preconditioners below.
Second-Order Optimization:
Research on second-order optimization in deep learning is generally divided into two categories: Hessian-free methods and methods that estimate the Hessian.
Hessian-Free Methods: Hessian-free approaches (Martens) 2010 Martens & Grosse 2015) op-
timize without explicitly computing the Hessian matrix, instead employing iterative techniques to
approximate the Newton step. Other recent works [Li 2018 2024 Pooladzandi & Li] (024] have
focused on designing iterative preconditioners to improve the convergence specifically for stochastic optimization algorithms.
Hessian Estimation Methods: These methods maintain an efficient approximation of the Hessian for neural networks. KFAC (Martens & Grosse 2015) and Shampoo (Gupta et al. 2018b) are two widely recognized methods in this area.
KFAC Martens & Grosse 2015) was one of the first approaches to go beyond diagonal precondi- tioners neural networks, demonstrating that layer-wise cronecker-factorel preconditioner ap- proximates the layer-wise wise Hessian in multi-layer perceptrons (MLPs). Subsequent works (Martens) iet al. 2018 Osawa et al.] (019)) extended KFAC to other architectures. Recent research (George
et al. 2018 Gao et al. [021] has further improved trace and diagonal estimates for KFAC. Efforts
scale up KFFAC (Ba et al. 2017: Puiu 2022 2023 Eschenhagen al. 2023) have focused on
making the inversion step more efficient or enhancing distributed implementations.
Shampoo (Gupta et al. 2018b), another second-order optimization algorithm, is motivated by the online learning algorithm Adagrad (Duchi et al.) 2011a). Shampoo also employs a layer-wise Cronecker-factoree preconditioner. A recent distributed implementation of Shampoo (Shi et al., 2023) won an optimization efficiency benchmark (Dahl et al., 2023), highlighting the practical util- ity of second-order methods in deep learning. Few recent works (Duvvuri et al., 2024| Morwani et al. 2024) have provided theoretical advancements on top of Shampoo. Other works (Anil et al.]
2020 Peirson et al.) 2022 Lin et al. 20224] Waag et al., 2024) have proposed various strategies to improve Shampoo's scalability. We defer comparison of SOAP with these methods to future work.
4

---

Published as a conference paper at ICLR 2025
4 ALGORITHM
4.1 THEORY
We begin by describing an equivalence between Shampoo and running Adafactor in the eigenbasis of the Shampoo preconditioner. For simplicity we omit momentum but the equivalence also holds with momentum. For this equivalence we use Shampoo with the following modifications from the original Shampoo optimizer (Gupta et al. 2018b):
1. We use power 1/2 instead of power 1/4. This was already recommended in practical implementations (Anil et al. 2020: Shi et al. 2023) and theoretical connection between optimal Kronecker approximation of Adagrad (Duchi et al. ( preconditioner and Shampoo with power 1/2 was established in Morwani et al. (2024).
2. We also the scalar correction to per layer learning rates described in Ren & Goldfarb (2021); Morwani et al. (2024).
3. Instead of the running average of L and R across time steps, use dataset averages.
With these changes in place (first occurrence of these changes is highlighted in red in the algorithm below) we formally define the two algorithms whose equivalence we show Algorithms andd
Algorithm 1 Single step of idealized Shampoo with power 1/2.
1: Sample batch Br.
2: G, E R
V whb, (W,)
3: L + EB EB[GnGT [Where the expectation is over a random batch B.}
4: +- Eg
5: A + L R/ Tracee L)
6: W. t Wi-1 Hu 1/2G, = W - -nni-nnnn----RR/// /Tacce/TTTTeeTT (L) -)-
Algorithm 2 Single step of idealized Adafactor in Shampoo's eigenspace.
Sample batch By.
2: G, m + V (W,)
3: L Gn
4: R Eg
5: Q1 + (igenvectors(L)
6: Qn + eigenvectors(f)
7: G, -QG,
8: [Idealized version of code for Adafactor taking G² to be the gradient)
9: G²n + QT,GG QR
A Eg(G"n 1p where G = Q.RRRRR
C E, G 3 Ga
12: V/ ACT 1.A [Elementwise division}
G.¹ G; (Elementwise division and square root)
VVite 14: G OH (Projecting back to original space)
15: W, + Wi- G
Claim 1. Algorithms Mand 2 are equivalent.
Proof. Consider Ge in the basis created after rotating by Qr, QR ie. Gj = Q.QQR. Let the eigenvalues of En, GG GT, and Ex, [GT, Gn.] be given by Aı,... Am and pi, respectively. Algorithm scales the i coordinate by (Aag/(Z; -1/2, while Algorithm 2 scales them by (A,C,/(2 EE A,)) We now show that A, >i. analogous argument shows C ptj-

---

Published as a conference paper at ICLR 2025
Algorithm 3 Single step of SOAP for X layer. Per layer, we maintain four matrices: E Rmxm, R I and V, M E IRm"x", For simplicity we ignore the initialization and other boundary effects such as bias correction. Hyperparameters: Learning rate 7, betas (B1, B2), ep- silon e, and preconditioning frequency f. An implementation of SOAP is available at status/ BEEEu//SOAPFSS0
1: Sample batch Br.
2: G Rmxx t- VWoB.
3: G" + GE
4: + 3,M + -1-
5: M' <- M-MMQn
6: (Now we "run" Adam on G*}
7: V + 32V + (1 (G' @ G") {Elementwise multiplication]
8: N' M'
(Elementwise division and square root)
Vitte
9: (Now that we have preconditioned by Adam in the rotated space, we go back to the original space.]
10: N + Q
W + WWWNN n.V
12: (End of gradient step, we now update and R and possibly also Ql and Qn.
L+ 3 (1 32 GGT
14: R
GT
15: ff%f==0 then
16: QL + (igenveetttosss(..01
17: Qn B...iiivettors
18: end if
A E. G' Gm
EB Σ (l
EBl Σ (uT(Ga) (GB)u,)
(Using definition of G')
EElllu (GB) (GB))
(vy form basis)
Eglu[G GBGBu.
A
(By definition of X and u;)
While these two algorithms equivalent in their idealized forms, practical considerations reveal some differences. Firstly, the algorithms differ when using running averages instead of dataset averages. Secondly, and more significantly in practice, decomposition of L and at every step. This means that the "adaptivity" of learning rates in Shampoo is limited to the updates of L and R. In contrast, with Adafactor in Shampoo's 's eigenspace, the second moment estimates (i.e., A and C' A Algorithmm2 2 can be updated at every step as they are computationally inexpensive. Additionally, instead of using Adafactor, can opf" for for Adam, which offers more generality. Combining these insights leads to Algorithm 3 which can be interpreted as running Adam in Shampoo's eigenspace.
do not invert compute the eigenvector
We note that practical implementations of Shampoo use grafting which allows for learning rate adaptivity every step. but this adaptivity is restricted to single scalar per layer.
'Though using AdamW over Adafactor only gives very small improvements in performance, see Figure and Appendix B.2 We also note that one can use other diagonal preconditioner based optimizer in place of Adam, such Lion (Chen al. 2023) Sophia (Liu et al. 2024 or Schedule- Free AdamW (Defazio et al.
2024

---

Published as a conference paper at ICLR 2025
Algorithm 4 Eigenvectors function, implemented using power iteration and QR decomposi- tion. Inputs: PSD matrix P and estimate of eigenvectors Q. If the estimate was exact we would have P = QDQ" where D is the diagonal matrix with eigenvalues.
S PQ 2: Q - QR(S)
We now describe some additional implementation details:
1. Algorithm 3 descrihes the behavior of the algorithm for 2D layers. Following Zhao et al. (2024a), for ID layers we run standard AdamW. This reduces the overhead as compared to standard implementations of Shampoo which solve an eigenvector problem for ID layers too.
2. Following Wang et al. (024), compute eigenvectors of L (and R) using one step of
power method (Algorthm 4 This requires doing one matrix multiplication followed by QR decomposition. QR decomposition is faster (Documentation (024) than standard eigenvector decomposition in PyTorch. For the first iteration, eigenvectors are initialized by doing a standard eigenvector decomposition.
3. For layers with huge dimensions such as the first and last layer in language modeling trans- formers, maintaining the eigenvectors would be space and time prohibitive. For such di- mensions we fix the rotation matrix (QL or QR) to be identity. Note that if we fix both QL and QR to be identity for a 2D layer, we would recover Adam.
4. Algorithm 3 omits bias correction and weight decay for simplicity, but these are used in the actual implementation, identical to their use in AdamW.
The main focus of the next sections will be to explore the empirical performance of this algorithm and its variations. In Appendices B.2 and B.3 we discuss the the space and time complexity of SOAP and how it can be improved.
5 EXPERIMENTAL METHODOLOGG
Hyperparameter tuning: We begin with hyperparameter values suggested by prior research for both AdamW and Distributed Shampoo (e.g., 32 0.95). Initially, we conduct a learning rate sweep to determine the optimal learning rate for each optimizer. Once the optimal learning rate is identified, we perform two-dimensional sweeps for each of the remaining hyperparameters, where vary the selected hyperparameter alongside the learning rate. The purpose of these sweeps to demonstrate that our default hyperparameter settings are near-optimal, disregarding potential interactions between two non-learning-- rate hyperparameters. A detailed discussion of the hyperpa- rameter sweeps is provided in Appendix A
We evaluate the throughput of each optimizer by measuring the num- ber of tokens processed per second. At present, we perform these measurements on a single H100 GPU and utilize gradient accumulation to accommodate large batch sizes. While this approach may seem to disadvantage AdamW-al as the overhead of Shampoo/SOAP is compared against mul- tiple gradient accumulation steps- is important to note that the overhead of Shampoo/SOAP can be amortized across layers by distributing the updates across multiple GPUs. This technique employed in the distributed implementation of Shampoo (Shi al. 2023). comprehensive comparison of distributed implementations of these algorithms is left to future work.
Throughput Measurement:
Simply running SOAP for the same duration as Shampoo and AdamW cannot be directly used to calculate the efficiency benefit (in terms of training steps or wall-clock time) of using SOAP since we use a cosine schedule. Therefore, we run SOAP on .5,55 and .875 fraction of the training data and fit scaling law of the form + 6N through the final losses obtained, where N' represents the number of training points and a, b, 8 are the parameters of the fit. We show these points and the corresponding scaling laws obtained in Figure 2 This scaling law is then used to calculate the efficiency benefit in terms of training steps and wallclock time as shown in Figure Here, the horizontal lines represent the final losses of AdamW and Shampoo.
Efficiency Benefits:

---

Published as a conference paper at ICLR 2025
360m, 2m batch size
660m, 2m batch size
360m, 256k batch size
2.76
2:90
2.94
2:88
2.92
2.86
2.90
2.84
2.00
2.70
2.06
2.82
2.68
2.84
2.80
2.82
2.78
0.5
0.75
1.0
0.565
0.8
1.0
0.72
0.875
Training Steps (scaled)
Training Stops (scaled)
Training Steps (scaled)
2.76
2.90
2.94
2.74
2.88
2.92
2.86
2.90
2.72
Test
2.84
2.88
2.86
2.82
2.84
2.80
2.66
2.78
0.52
0.78
1.01.06
1.12
1.0 1.11
Total Wall Time (scalled by AdamiW
Total Wall Time (scaled by AdamWW
Total Wall Time (scaled by AdamiW)
AdamiW
SOAP
Shampoo
Figure 2: Precise efficiency benefits of SOAP over AdamW and Shampoo for 360m (at 256k and 2m batch size) and 660m (at 2m batch size) model. For the precise methodology, refer Section
LANGUAGE MODELING EXPERIMENTS
In this section focus on empirically comparing AdamW, DistributedShampoo and SOAP on language modeling tasks.
6.1 MEASURING EFFICIENCY BENEFITS
In Figure T (left and middle) and Figure we show train loss curves for AdamW, Shampoo, and SOAP on 360m and 660m models with 2m token batch size and "chinchila--ppimaa" i.. 20x model size number of tokens. In these plots we observe that SOAP outperforms the other two optimizers. To directly calculate the efficiency benefit of SOAP, we also run SOAP with cosine decay for shorter Ir schedule, as shown in Figures [ and 3 This allows us to approximate the following efficiency benefits (when batch size is set to 2m and preconditioning frequency to 10); > 40% reduction in the number of iterations and > 35% reduction in wall clock time compared to AdamW: W: zs 20% reduction in iterations and wall clock time as compared to Shampoo. Precise efficiency benefit calculations
presented in Figure left and middle). In Appendix C we show that efficiency benefits of SOAP over AdamW W are maintained for longer duration runs where #tokens 100 x model size.
8

---

Published as a conference paper at ICLR 2025
360m, 2m batch size, Preconditioning Frequency-18 10
3.4
3,4
3.2
3.2
Thain Locs
3.2
3.0
3.0
2.8
2.9
800
1600
2400
3200
0.25
0.5
0.75
1.0
Training Steps
Wall Time Iscaled by AdemW))
AdamW
Shampoo
SOAP
SOAP (shorter LR schedule)
Figure 3: Comparing performance tuned runs for AdamW, Shampoo (using DistributedSham- poo (Shi et al. (2023) implementation) and SOAP. Shampoo and SOAP use preconditioning fre- quency of 10. We observe 2 40% reduction the number of iterations and 2 35% reduction in wall clock time compared to Adam W, and approximately a 20% reduction in both metrics compared to Shampoo. See Figure [ for 660m results, Sections 6.2 and [6.3 for ablations of preconditioning frequency and batch size respectively, and Section 5 for detailed calculation of efficiency improve- ment and experimental methodology.
6.2 EFFECT OF FREQUENCY OF FINDING FICENVECTOORS/INVRSS
In Figure I (right), we compare SOAP and Shampoo with respect to preconditioning frequency. We observe the following:
For all frequencies we tried from 1 to 100, both optimizers outperform AdamW.
At frequency 1, SOAP and Shampoo quite close in performance.
At bigher frequencies, the performance of both SOAP and Shampoo degrades but SOAP's performance degrades significantly slower than Shampoo's.
6.3 SOAP IMPROVES THE CRITICAL BATCH SIZE
When scaling up batch sizes, the ideal outcome is that doubling the batch size results halving the number of training steps needed to achieve the same performance. The batch size at which this ideal scaling starts to break down is referred to by McCandlish et al. (2018) as the critical batch size. As models and datasets grow larger, it becomes increasingly important develop optimizers that support larger critical batch sizes, thereby reducing the serial runtime of a training run. this baseline setup
subsection, compare the critical batch sizes of AdamW and SOAP. Relative of 2 million batch size, when we decrease the batch size by a factor of k, we increase the precon- ditioning frequency by the same factor. This ensures that the FLOPS and wall clock multiplicative overhead for the eigenvector decomposition steps remains consistent with the 2 million batch size setting.
We start by training 360 million parameter model with batch size of 256k for "Chinchilla- optimal" number of tokens (20 times the model size) using AdamW, achieving loss of 2.842. This value is set as the target loss for our comparisons. In Figure g (left), we show the number of steps AdamW and SOAP require to reach this target loss as we vary the batch size. SOAP consistently requires fewer steps across all batch sizes, with the multiplicative benefits becoming more pronounced at larger batch sizes. Additionally, we compare these results to the ideal scenario (dashed line) of linear scaling, where doubling the batch size halves the number of steps. SOAP more closely follows the linear scaling trend compared to AdamW, indicating that has higher critical batch size in this setup.

---

Published as a conference paper at ICLR 2025
360m
360m, 256k batch size
Critical Batch Size for AdamW SOAP
Preconditioning Frequency=80
3.4
25000
12500
Estimated Taining Steps
311
3125
256
512
1024
2048
5000
10000
15000
20000
Batch
Training Stops
AdamiV
SOAP (shorter schedule)
Linear
Scaling
Figure 4: (left) Comparing the critical batch size of Adam vs SOAP. We can see that SOAP im- proves the critical batch size, by being much closer to the ideal linear scaling with batch size as compared to AdamW. (right) Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo (Shi et al., 2023) implementation) and SOAP for token batch size of 256k. Shampoo and SOAP use preconditioning frequency of 80. We observe Z 25% reduction in the number iterations compared to A, and approximately a 10% reduction compared to Sham- poo. See Figure (right) efficiency improvement.
(right) for wall-clock time improvement and Section 5 for detailed calculation of
Figure (right), we present the optimal runs for each optimizer (including Shampoo) at smallest batch size consider: 256k. SOAP outperforms both Shampoo and AdamW, reducing the number of iterations by 25% compared to AdamW, and by approximately 10% compared to demonstrate that SOAP also achieves 2 15% over Adam and around 10% over Shampoo. We note
Shampoo. Furthermore, in Figure 2 (right, bottom),
wall-clock time improvement that these results are a preliminary analysis for smaller batch size runs. Our approach of keeping the product of batch size and preconditioning frequency constant may not be optimal, and better trade-off could likely be found. Furthermore, SOAP's overhead could potentially be reduced by performing L and R updates in lower precision (instead of fp32). Finally, the diminished efficiency gains of second-order methods at smaller batch sizes consistent with prior findings (Zhang al., 2019 Ishikawa & Yokota 2024).
DISCUSSION AND LIMITATIONS
We study an optimizer called SOAP: ShampoO with Adam in the Preconditioner's 's eigenbasis. We show that SOAP outperforms both AdamW and Shampoo in language modeling tasks and show that it is more robust to changes in preconditioning frequency than Shampoo. While we have explored many factors such as batch size (Section 663) and training duration (Appendix C) we acknowledge that our study focuses on relatively small scale compared to recent L.MMs Touvron et al. (2023) which are two orders of magnitude bigger. We hypothesize that our findings on the performance of SOAP would generalize to larger scales due to its theoretical foundation. SOAP's robustness is also supported by the fact that SOAP is equivalent to running Adam in a rotated space, and Adam has proven be effective across scale and tasks. However, this hypothesis remains to be validated.
For future work, we aim to improve the design of SOAP further, particularly by exploring the use of lower precision for preconditioners and optimizing its distributed implementation. Additionally, we are interested in testing SOAP's performance in other domains, such as vision, to evaluate its performance across different types of tasks.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
SK, DM, and RZ acknowledges support from the Office of Naval Research under award N0001422 1-2377 and the National Science Foundation Grant under award #IIS 2229881. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to estab- lish the Kempner Institute for the Study of Natural and Artificial Intelligence. NV, DM and RZ are supported by Simons Investigator Fellowship, NSF grant DMS-213411 DARPA grant W91 INF2011110 ladd DOE grant DE-SC0022199. LJ acknowledges funding from the National Science Foundation RMS-2134157.

---

Published as a conference paper at ICLR 2025
REFERENCES
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint 2002:2002.090018 2020.
Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using kronecker-factored approximations. International Conference on Learning Representations, 2017. URL/attps://ppeeeeeee/
net
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang. Hieu Pham, Xuanyi Dong, Thang Luong. Cho-Jui Hsieh, Yifeng Lu, and Quoe Le. Sym- bolic discovery of optimization algorithms. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances Neu- ral Information Processing Systems 36: Annual Conference on Neural Information Pro- cessing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL http:: 1439bb9955e33 3
www.ppsccppaappee fwwes/ pwtps//ppppeee22/.3/.. htaap/rr/////hsht/// otg.a//eee.....aat-ttotto------- html
George E. 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, 2023.
Aaron Defazio, Xingyu Yang. Harsh Mehta, Konstantin Mishchenko, Ahmed Khaled, and Ashok Cutkosky. The road less scheduled. CoRR, abs/2405. 15682, 2024. doi: 10.48550/ARXIV2440 15682. URL hhtt:s:s::
wwwg/oog//.o......o 48550/arXiii. 2405. 15682
Mostafa Dehghami, Josip Djolonga, Basil Mustafa. Piotr Padlewski. Jonathan Heek, Justin Gilmer, Andreas Peter Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to billion parameters. In International Conference on Machine Learning, pp. 7480-7585 7512. PMLR, 2023.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In The Tenth International Conference Learning Representationss ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. wtt./ttpsss///eeeeeeee..eer Forum?id-shpppp//
Documentation.
torch borch.inalg.co documentation.
linalg.eigh. html 2024.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61):2111-2159, 201 la. URL
http:// wwtpppppppp..w///aaap.. http:///uuhhhaaaaaaaa html
John C. Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Mach. Learn. Res., 22:222212219 101 doi: 10.5555/195300 2021068. URL httpt:::////........tt //dl acm, wwt. doi/ / 10. 5555 000/5/5/1155348.0000 2021068
Sai Surya Duvvuri, Fnu Devvrit, Rohan Anil, Cho-Jui Hsieh, and Inderjit Dhillon. Combining axes preconditioners through kronecker approximation for deep learning. The Tweifth International Conference on Learning Representations, 2024, httpttts::/ppenneviiiii net forum?
1id8899h2DDV13
Runa Eschenhagen, Alexander Immer, Richard E Turner, Frank Schneider, and Philipp Hen- nig. Kronecker-factored approximate curvature for modern neural network architectures. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https:
Ex33JEKS5333
Kaixin Gao, Xiaolei Liu, Zhenghai Huang, Min Wang. Zidong Wang, Dachuan Xu, and Fan Yu. A trace-restricted kronecker-factored approximation natural gradient. Proceedings of the AAAI Conference on Artificial Intelligence, (599)7511-7522 May 2021. doi: 10. wtw.0///aai.......
16921. URL htR.hhtts::///t///o http::///wggaaaaaao...oo.a.a.aia
12

---

Published as a conference paper at ICLR 2025
Google
Gemini Team. Gemini 1.5: Unlocking multimodal understanding across millions of to- www.looogpp..... com/
kens of context. https:// gemini v1 report. pdf 2024. [Online; accessed M02My))2)22..
mwdia/ediaaeeeee gemini/
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approx- imate natural gradient descent kronecker factored eigenbasis. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa Desa-Bianchi, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Infor- mation Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp-
9573-9583, 2018, URL http:: htttpss:/pprooeee///ttoooo ngs neuri ipeuripp... ps. wwt.ccccpppppe httppapprr2218/hhaae www.bbabbbbbttttteeeeaeea html
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, et al. Olmo: Accelerating the science of language models. arXiv preprint arXiv:2402.00 000888000 2024.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo:
Preconditioned stochastic tensor opti- mization. In International Conference Machine Learning, pp. 1842-1850. PMLR, 201 10.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning. ICML 2018, Stockholmsmüsssun, Stockholm, Sweden, July 10- 15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1837-1845. PMLR, 2018b. URL
m.rr wttpspps///s//pssaaaaaa htmm
Dan Hendrycks and Kevin Gimpel. Gaussian error
linear units (gelus). arXiv preprint
1rXi::1606.08455, 2016.
Jordan Hoffmaan, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models. CoRR, abs/2203. 15556, 2022. doi: 10.48550/ http://.11.1455///// www.:riwwww............o... 15551
1RX1V.2203.15556. URL
Satoki Ishikawa and Rio Yokota. When does second-order optimizzainn speed up training? In The Second Tiny Papers Track at ICLR 2024, 2024. URL
net. nwt//frumm forum?
id=NLr FEs5ZNb
Jean Kaddour, Oscar Key, Piotr Nawrot, Pasquale Minervini, and Matt J. Kusner. No train no gain: Revisiting efficient training algorithms for transformer-based language models. In Al- ice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NewrIPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023. URL htt.::tp/// /papess//ppppppppppppepee nips www.:pppp..c........ppppppppppp. wwwes: httppp/ee.ileee/pp////////eeee 51 211 ottp::/eeooeeooaaaaaaaaaaaaar-ee Cowraactt-oooffeeneeee html
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings. 2015. htth:
6980
Xi-Lin Li. Preconditioned stochastic gradient descent. IEEE Transactions Neural Networks and Learning Systems, 2955):1454-1466 2018. doi: 10.110/TNN.SS22172222993
Xi-Lin Li. Stochastic hessian fittings with lie groups, 2024. URL htt.:http////xiiieee abs/
2402.11858
Wu Lin, Felix Dangel, Runa Eschenhagen, Juhan Bae, Richard E. Turner, and Alireza Makhzani. Can we remove the square-root in adaptive gradient methods? A second-order perspective. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scar- lett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning. volume 235 Proceedings of Machine Learning Research, pp. 29949 29949-22973 29973. PMLR, 21-27 Jul 2024. J ptttpss:/prroceedie igs. ngs ml w.wpre.s.....////
html
13

---

Published as a conference paper at ICLR 2025
Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Tweifth International Conference Learning Representations, 2024.
[d==3xDee@Nool
Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimization with adaptive learning rate. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meet- ing, August 11-16, 2024, pp. 12486-12502 Association for Computational Linguisties, 2024a. nd ngs-anlcclnn.... 742
URL https://act http://aclaaahooocoooooc
Kai Lv, Yuqing Yang. Tengxiao Liu, Qipeng Guo, and Xipeng Qiu. Full parameter fine-tuning for large language models with limited resources. In Lun- Wei Ku, Andre Martins, and Vivek Sriku- mar (eds.). Proceedings of the 62nd Annual Meeting of the Association for Computational Lin- guistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 16, 2024. pp. 8187- 8198. Association for Computational Linguisties, 2024b. URL https: / http:///laaatthooo......o 445
James Martens. Deep learning via hessian-free optimization. In Proceedings of the 27th Interna- fional Conference on International Conference on Machine Learning, ICML'10, pp. 735-742, Madison, WI, USA, 2010. Omnipress. ISBN 978 7781605589777.
James Martens and Roger Grosse. Optimizing neural networks with fratocder-factored approxi- mate curvature. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 14408-2417 Lille, France, 07-09 Jul 2015. PMLR. URL https: http::/:///prceeeenng/ .llr.
press/v37 wwt///3///aa/////. htt//m/aarenemm.em/mm html
James Martens, Jimmy Ba, and Matt Johnson. Stonecker-factoree curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018. URL
hittps://openrevile wwtp://wwwweeetteeeeeeeeee
Forum? feerum?iddee wtweee/ooorumm1ddyytokoom d=HYMTKQZAAb
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. CoRR, abs/1812.06162 2018. URL http:::/aaa.ii.ogg//bbb/////////
061 62
benchmark MLCommons. Mlc benchaaa competition. httpss///i.ccomooso.rr////22///2
mlc-
it 2024. Accessed: 2024-10-01.
Depen Morwani, Itai Shapira, Nikhil Vyas, Eran Malach, Sham Kakade, and Lucas Janson. A new perspective on shampoo's 's preconditioner. arXiv preprint arXiv:2406. 17748, 2024.
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order optimization using kronecker-factored approximate curva- ture for deep convolutional neural networks. In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 12351-12359, 2019. doi: 10.1 10.1109/CVPP.2219..0222
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Abel Peirson, Ehsan Amid, Yatong Chen, Vladimir Feinberg, Manfred K Warmuth, and Rohan Anil. Fishy: Layerwise fisher approximation for higher-order neural network optimization. In Has it Trained Yer? NeurIPS 2022 Workshop, 2022. URL. httts::/
net net//formm?
id=cScb-RrBOc
Omead Pooladzandi and Xi-Lin Li. Curvature-informedd SGD via general purpose lie-group precon- ditioners, 2024. URL
jxRnVpF
Tomer Porian, Mitchell Wortsman, Jenia Jitsev, Ludwig Schmidt, and Yair Carmon. Resolving discrepancies in compute-optimal scaling of language models. CoRR, abs/2406. 191446 2024. doi: 10.48555/ARXXI..4466.. 19146. (ttt:tt:::// doi. /ttpo//oo....ooooo/oo////o.oo 40w.11...05500000........ 2406.
1 6
14

---

Published as a conference paper at ICLR 2025
Constantin Octavian Puiu. Randomized k-facs: Speeding up k-fac with randomized numerical linear algebra. In Hujun Yin, David Camacho, and Peter Tino (eds.), Intelligent Data Engineering and Automated Learning- IDEAL 2022, pp. 411-422, Cham, 2022. Springer International Publishing. ISBN 978-3-031-21753-1
Constantin Octavian Puiu. Brand new f-facs: Speeding up k-fac with online decomposition updates, 2023. BtLpptts:///xxet
org. abs 2ttpg//bbb/22100843333
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. Journal of machine learning research, 21 (140): 2020.
Yi Ren and Donald Goldfarb. Tensor normal training for deep learning models. In M. Ran- zato, A. Beygelzimer. Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 26040-26052. Curran Associates, Inc., 2021. URL. https: ngs neurips cc/ http://ppppeepapeeaap/eeeeeppppe wwwes//appe////22/ 2ataa332266600000011111000000000
pdff
Nikhil Sardana, Jacob Portes, Sasha Doubov, and Jonathan Frankle. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.nee, 2024. pet/foruue ii 1tmmt TDUu
CtRLhttps:/ /tttpps:/oopenrevee
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost.
In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4603-4611. PMLR, 2018. URL
http::/proceeddnng .mw...ssm...................m... mlr. wwwss/ //0/s//B0///// 2t000/hhaeee18.... html
Hao-Jun Michael Shi, Tsung- Lee, Shintaro Iwasaki, Jose Callego-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. CoRR, abs/2309.066 0664977 2023. doi: 0044550/ARXX..33000006644 //doi. 48550/arXix
URL https:
Jianlin Su. Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurélien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models. CoRR, abs/2302. 13971, 2023. doi: 10.4850/ARXIV.230...333 13971. URL https://doi www.ooo........o.oooooo 48550.4855/8850 ariiv. 2302. 13971
Nikhil Vyas, Depen Morwani, and Sham M. Kakade. Adamem: Memory efficient momentum for adafactor. In 2nd Workshop on Advancing Neural Network Training: Computational Ef- ficiency, Scalability, and Resource Optimization (WANT@ICML 2024), 2024. URL https:
Sike Wang. Jia Li, Pan Zhou, and Hua Huang. 4-bit shampoo for memory-efficient network training. CoRR, abs/2405. 2024. doi: 00448500/RXXI..200..88 10. 00.44855//aaa0
URL httpss://// org/
2405. 18144
Mitchell Wortsman, Peter J Liu, Lechao Xiao, Katie E Everett, Alexander A Alemi, Ben Adlam, John D Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Kornblith. Small-scale proxies for large-scale transformer training instabilities. In The Tweifth International Confer- ence on Learning Representations, 2024. URL https:://///........ http:s:://ppeeeviee...ee. / fortt//frmmmiim (d800pmxXbZ)
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Or- leans, LA, USA, June 18-24, 2022, pp. 1204 1213. IEEE, 2022. doi: 10.110///V 01179, URL htR.htts:: /ooi http://..01/////25555550111 CVWR/26RR22228 2ET028...022..0 2022. 01179
15

---

Published as a conference paper at ICLR 2025
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E. Dahl, Christopher J. Shallue, and Roger B. Grosse. Which algorithmic choices matter which batch sizes? insights from noisy quadratic model. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Pro- cessing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 8194-8205, 2019, URL https: httpsp////.pceeed.....e.e.rrrr. @0eacd98391 167432633222 7 ae 2b
wc/paperr/c11//aaee html
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuan- dong Tian. Galore: Memory-efficient LLM training by gradient low-rank projection. CoRR, abs/2403.03500 2024a. doi: 10. 00448500ARXXI.2403.0055 URL https::/doi.cr// org/10..
48550/arxiv. 2403. 03507
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang. Anima Anandkumar, and Yuandong Tian. (code) galore: Memory-efficient LLM training by gradient low-rank projection. https:
//itthub.com/jiaweeiiiiii GaLore 2024b.
Rosie Zhao, Depen Morwani, David Brandfonbrener, Nikhil Vyas, and Sham M. Kakade. Decon- structing what makes a good optimizer for language models. CORR, doi: TT44550/ARXXIV22407.777777 URL phttps://do 485500 2407.
2024c
07972
16

---

2. (360m) We sweep over the cross product of best 3 learning rates and 31 E (0.9, 0.95, 0.99). Published as a conference paper at ICLR 2025
A EXPERIMENTAL SETUP
Many aspects of our setup such as models are the same as in Zhao et al. (022c). We will restate
those details verbatim for completeness.
We train language models on C4 tokenized with the T5 tokenizer (Rffell et al. 2020) and report results in terms of validation loss.
Models. We start from the OLMo codebase (Groeneveld et al., 2024) and train decoder-only trans- former models of three sizes: 210m, 360m, and 660m, where the parameter count refers to non- embedding parameters. The models have widths of 1024, 1024, and I 4008 and depths 12, 24, 24. We used the 210m model to explore various ablations, most of our reported results are on 360m and 660m. The MLP hidden dimension is 4x of the width. The activation function is GeLU (Hendrycks) & Gimpel 2016). We use RoPE positional encodings (Su et al.) 2024p. Attention heads are always dimension 64. We use PyTorch default LayerNorm. We use QK layer norm (Dehghani et al. (023), Following Wortsman et al. (2024) do not learn biases for the linear layers LayerNorms. We train in mixed precision with b6
Algorithms. We use the standard Pytorch implementation of Adam (Paszke et al. 2019), the DistributedShmppor (hi et al. (2023) implementation of Shampoo. We implement ourselves SOAP and GaLore starting from an older version Pytorch implementation Adam W and the official GaLore implementation Zhao al. (024b))
Default hyperparameters. We use 31 0.95, sweeps for the 360m model. Following Wortsman et al. (2024) we use decoupled weight decay with coefficient le-4 and z-loss with coefficient 1e-4. We the default value of = le- in Adam W (actual or when used for grafting). SOAP and GaLore. We use warmup followed by cosine decay as our scheduler. We start the warmup and end the cosine decay at 0.1z the maximum learning rate.
found it outperform A1 0.9 in
Default hyperparameters for DistributedShampoo Shi et al. (2023 state that they find the op- timal exponent to be either 1122 or - -1822/4 s 1.2. Our preliminary findings were similar this. Hence we set the default values of exponent to be 1/2.5 for both ID and 2D parameters. We set Eshmpoo = le- 12 and Sshampoo = 0.95 based on our imitial set of experiments on the 210m
model.
Default hyperparameters for GaLore GaLore introduces an additional hyperparameter called scale (a) since due to low rank updates the overall update magnitude decreases. Since we are running a full rank version of GaLore we set a a=
Token counts. For all of our runs we use a sequence length of 1024. For all models (except in Section 6.3), we use a token batch size of 2048k s 2m. We default to training models for the approximately "chinchilla optimal" (Hoffmann et al. (022) number of tokens that is s=20 times the number of parameters. Explicitly, this means for our default batch size of 2m, the 21 Omm models are trained for 1600 steps or &e 3.3b tokens. The 360m models are trained for 3200 steps, the 660m models are trained for 6400 steps.
A.1 SWEEPING OVER HYPERPARAMETERS
Adam 2m batch size: Starting from the default hyperparameters above do the following sweeps:
1. We sweep over learning rate in {.1, .0316, .0.,....,3 3.1..,,,.66e4,,
3. (360m) We sweep over the cross product of best 3 learning rates and 32 E (0.9, 0.95, 0.99).
The last two of the sweeps did not yield any benefit for the 360m model with 2m batch size hence only sweep over learning rate for the 660m model with 2m batch size.
17

---

Published as a conference paper at ICLR 2025
DistributedShamppn, 2m batch size: Starting from the default hyperparameters above do the following sweeps:
1. We sweep over learning rate in (.1,. 01,, 3. 3.16e-14 43.
2. (360m) We sweep over over the cross product of best learning rates from above and Eshampoo (le-11, 1e-12, le-13).
3. (360m) We sweep over over the cross product of best 3 learning rates from above and Sshumpoo E .975).
4. Let e1, 62 denote the exponents used in DistributedShampoo for ID and 2D parameters respectively. We also sweep over the cross product best learning rates from above and (21, in ((2 2), (2.5, 2.5), (3, 3), 4))
These sweeps did not yield any significant improvement in performance (< .004) for the 360m model. Hence only sweep over the learning rate for the 660m model.
SOAP, 2m batch size: Starting from the default hyperparameters above we sweep over learning rate in {.1, 00060 3.16..,...6e44)
AdamW, W, 256k batch size: For the 360m model with 256 batch size we start from the default hyperparameters and do the following sweeps:
We sweep over learning rate {.1, .0316, .01, 3. 4}.
2. We sweep over the cross product of best learning rates and 82 E 0.995.0.99.
In the second sweep we observe small improvements in performance by using 3g = .99, so our final
numbers B2 .9.. This (small) improvement in performance by using larger B2 smaller batch sizes was also observed by Porian et al. [2024); Zhao et al. (2024c).
DistributedShamppo,, 256k batch size: For the 360m model with 256 batch size start from the default hyperparameters and do the following sweeps:
1. We sweep over learning rate in 0316,, 01,....3 16.-16).. 4).
2. We sweep over the cross product of best learning rates and (32, Pshampoo) e {[.95, .95), (.99, .99)].
In the second sweep we observe small improvements in performance by using B3 = Sshampoo .999 so our final numbers use B2 3shampoo .999
SOAP, 256k batch size: For the 360m model with 256 batch size we start from the default hyper- parameters and do the following sweeps:
3.16e-4]. 1. We sweep over learning rate in {.1, .0316, 3.6644,
2. We sweep over the cross product of best learning rates and 32 E (.95, .99..
In the second sweep we observe small improvements in performance by using 32 .99, our final numbers use 32- .99.
Preconditioning frequency sweeps: For the preconditioning frequency experiments of SOAP and Shampoo Figure [] (right)), for each frequency we do learning rate sweep over the best 3 learning rates found at preconditioning frequency 10. Other hyperparameters are set to their optimal values obtained using the precondition frequency 10 sweeps.
For each of the shorter runs of 360m and 660m models for the SOAP optimizer (Fi we did learning rate sweep over the best learning rates found for the standard length run. Other hyperparameters are set to their optimal values obtained using the standard length run.
360m and 660m shorter runs:
Warmup: The warmup duration for the 360m and 660m models were 600 and 1200 steps respec- tively. For the shorter runs (Figure 2), for 360m model, the warmup durations were 400, 400, 500 and 525 steps for 0.5, 0.625, 0.75 and 0.875 runs respectively. For the 660m model, the warmup du- rations were 600, 750, 900 and 1050 steps for 0.5, 0.625, 0.75 and 0.875 runs respectively. For 360m model with 256k batch size (Section 6.3) we warmup for 4000 steps (total steps is 25000).
18

---

Published as a conference paper at ICLR 2025
B FURTHER EFFICIENCY IMPROVEMENTS
In this section, we discuss space and time complexity of SOAP and provide an overview of potential avenues for further space and compute efficiency improvements in SOAP.
B.1 ONE SIDED EIGENBASIS
As described in Section Zhao et al. 2024a) have an algorithm similar to ours. One of the differ- ences is that they only project the smaller side of the layer using the eigenbasis while using identity as the rotation matrix for the larger side i.e. if m < n. we set Qn = I, Algorithm and m > n set Q2 m- Doing this leads to reduction in space usage well as reduction of optimizer time overhead, which is discussed in Appendices B.2.1 and B.3.1
In Figure 5 is evident that the one-sided projection results in slightly reduced performance com- pared to the original SOAP optimizer. However, still performs on par with, or marginally better than, Shampoo, while maintaining greater computational efficiency. Further investigation into the potential for these variants to surpass the computational efficiency of original SOAP optimizer is left for future work.
360m, 2m batch size Preconditioning Frequency
660m, 2m batch size Preconditioning Frequency 10
3.3
3.1
3.2
3.0
Train Loss
3.1
2.9
2.8
2.6
1000
2000
3000
2000
4000
6000
Train Steps
Trann Steps
AdamW
SOAP
SOAP (ane-sided)
Shampoo
SOAP (factorized)
SOAP (factorized, one-sided)
Figure 5: Performance of variants of SOAP which improve space usage time overhead. 1. SOAP (factorized): Uses Adafactor instead of Adam in Shampoo's 's eigenbasis and 2. SOAP (one-sided): Uses Q = (i.c. rotation) the large side of weight matrix and 3. SOAP (factorized, one- sided): Combines both these changes. We observe that while using Adafactor instead of Adam negligible increase in loss, using the one-sided variant
larger increase. However, the one-sided variant also has much larger reduction in time and space overhead. For computational benefits of these variants see Appendices B.2 and B.3
B.2 SPACE USAGE OF SOAP
For m matrix where m > we require
2m2 (for L, Q1) 2n² (for R, QR) + 3mn (for gradient, M,V) space usage (beyond weights and activations), specifically for L, QL. R, QR. momentum (M), AdamW's second order estimate (V), and the gradient. This the same space usage as Distributed- Shampoo while AdamW uses 3mm..
'One mn is for storing the gradients. this can be avoided (as long as there is no gradient accumulation) by applying gradients along with backprop (L.v et al. (024b) but this is not implemented by default in standard deep learning frameworks such PyTorch. Hence will include this term in all of our calculations.
19

---

Published as a conference paper at ICLR 2025
B.2.1 IMPROVING SPACE USAGE OF SOAP
The most direct way reduce memory is using low precision store the L, R, Q Qn: V matri- ces, which is done by Dettmers et al. (2022): Wang et al. (2024). Orthogonal to the low precision approaches, there are two algorithmic approaches improving the space usage of SOAP:
Using Adafactor instead of Adam as the diagonal preconditioner after rotating by Qr. and QR. This reduces the space usage by mn.
Using one sided version of SOAP (Appendix B.1 This reduces space usage from 2m + 2n² + 3mn 2 min n) + 3mm.
Combining these approaches yields space usage of 2 min(m, n2² + 2mn.
For standard transformer architectures the last variant which combines the approaches would yield less space usage overall compared to A (which uses 3mm).
We try these approaches in Figure 5 We observe that using Adafactor instead of Adam yields very small reductions in performance while using one-sided preconditioner results in larger reduc- tions. Nonetheless even after combining these two approaches the resulting optimizer outperforms AdamW while having a smaller space requirement than AdamW. Regarding space usage we also note that Adafactor (with momentum added back) itself utilizes only 2mn space usage and has been shown to perform comparable to AdamW for VIT training (Zhai et al. (202)) and for language model training (Zhao et al. 2024c). Further space reduction beyond Adaiactor has been studied in the Adalomo (L. et al. 2024a), GaLore (Zhao et al.) 2024a), and AdaMeM (Vyas et al., 2024)
papers.
B.3 TIME OVERHEAD OF SOAP
There are two types of overhead of Shampoo and SOAP over AdamW: the overhead per step and the overhead when changing the preconditioner (or for SOAP, the preconditioner's 's eigenbasis). Let us first analyze the first one. For SOAP per step for a layer of size m x n we have an overhead
m² (updating +)+n² (updating +2m 2012mmn+222 ++2mm² (projecting and projecting back on both sides).
We note that this is more than the overhead of Shampoo which m + n m2n + n2m. This can be observed in Figure 2 (bottom, right) but not in the other figures since there the second type of overhead the dominant term.
The second type of overhead is due to changing the preconditioner for Shampoo (or for SOAP, pre- conditioner's eigenbasis i.e. Qr and QR). The DistributedShampoo (Shi et al. 2023) implementa- tion of Shampoo uses a direct call to torch. www.iinaall 1g. eigh gh for this. Following Wang et al. (0044) we use Algorithm 4 which uses power iteration based approach which calls torch. linalgi.llg.nnoo qr. We note that torch... 1 linalg..... qr is faster than torch. 1 linalg. e eigh (Documentation 2024) In Figure 6 (right) we see that using power iteration based approach (torch. linannnll g. qr) per- forms as well as fresh eigenvector decomposition (torch.linal lw..ll...gggg)h.
Effect of frequency on overhead: In Figure 6 (left), we observe that the overhead decreases as the preconditioning frequency increases, i.e., the frequency of invoking Algorithm 4 If the only additional computation occurred in Algorithm 4 we would expect the overhead to scale as 1.0/ (preccondiiionnn frequency), approaching zero. However, empirical results (Figure (left show that the overhead approaches an asymptote greater than zero. This is attributable to the additional matrix multiplications required update L, update R. project the gradient, and reproject the gradi- ent (for each layer) in the optimizer. Currently, these operations are performed in float32; reducing the precision of these operations, as proposed in Wang al. (2024). could lower this asymptote.
B.3.1 IMPROVING TIME OVERHEAD OF SOAP
The per step overhead of SOAP can be reduced by using low precision to store the L, R, Q2, QR. V matrices (Dettmers et al. 2022 Wang et al.) 2024h, which in turn will speed up computation done using these matrices. This approach cannot be used for reducing the overhead for the preconditioner update in popular deep learning frameworks such as Pytorch since torch. 1 linal allg qr does not
20

---

Published as a conference paper at ICLR 2025
Frequancy vs
2.94
Sovemard in Co timg Adammi
2.92
2.00
2.54
2.82
19
32
100
Precoddiisnin frequency
Prccond doning Frequancy
Figure 6: (Left) Depicting overhead of SOAP over AdamW as a function precondition- ing frequency (Right) Comparing the performance of SOAP with torch. www.allanngiii.......... for computing the eigenvectors with Algorithm 4 which uses torch. www.aniiinngg........ qr. Note that torch. linalgg.... qr is computationally more efficient than torch. linalgg...... eigh (as men- tioned in Documentation (2024)); however, both seem to have comparable performance throughout the preconditioning frequency spectrum.
support precision lower than float 32. Orthogonal to the low precision approach we can improve the per step time overhead of SOAP by the following algorithmic approaches:
Using Adafactor instead of Adam (Appendix B.2 tating by Qr and QR- In this version of SOAP the overhead can be improved by from m² +n 22222+ 2 m² m # max((,, n12 m mmin(m, n)* by merging the project and project back steps for the smaller dimension.
the diagonal preconditioner after
n)
Using one sided version of SOAP (Appendix B.I). This reduces overhead from m n² + 2m'n + 2n²m to min (, n) + 2 min (m, n max(m, nt).
Combining these approaches yields an overhead of min (m, n)² max(m, n)) + min(m, n),
Using one-sided version also reduces the second type of overhead from a calls to torch. www.llnallll gr. qr m and a nxn matrix only single call to min(m, n) x min(m, n) matrix.
C LONGER DURATION RUN
Chinchilla scaling laws (offmaan et al. 2022) suggest that is compute optimal to use tokens which approximately 20x the models size, which what have been using for standard runs. But many recent L.LMs such as the LLaMA (Touvron et al. 2023) series of models are trained on much larger token counts. This can be take account the computational ence (Sardana et al., 2024) or to create models which are usable or finetunable by downstream users. In Figure we train language model with AdamW on 100x model size token count. We then train the same model with SOAP for 50x, 75x, and 100x token counts to approximate the efficiency benefits. We find efficiency benefits (> 40%) similar to those observed in Figure 2 for Adam W runs with 20x token counts.
during infer-
D GALORE
We tried GaLore for 210m model, and while it outperformed AdamW it performed worse than Shampoo. Hence we do not try GaLore for higher model sizes.
Hyperparameter sweeps: We did the following sweeps:
21

---

Published as a conference paper at ICLR 2025
360m, batch size, long with 33.5b tokens 100x model size Preconditioning Frequency=10
2.71
2.71
2.70
2.70
Train Loss
2.68
2.67
2.67
2.66
2.66
2.65
2.65
0.5
0.6
0.7
0.8
1.0
0.5
0.6
0.7
1.0
Total Training Steps (Scaled)
Total Wall Time (Scaled by AdamW)
AdamW
SOAP
Figure 7: Total training steps (Left) and total wall clock time versus final test loss for long runs (#tokens = "chinchilla" tokens 100x model size).
1. We swept the cross product over learning rate (3.16e 1e 3, le 2). precon- ditioning frequency (10, 50, 200), both sided and one sided versions. Frequency 200 had the best results matching the observation of Zhao et al. (2024a).
2. We did cross product sweep over learning rate (3.16e 4, 1e 3, 3. 1e 2), both sided and one sided versions with 32 .99 instead of .95 and preconditioning frequency 200.
3. We did cross product sweep learning rate (3.16e 4, le 3. 33 le 2), both sided and one sided versions, preconditioning frequency (50, 200) with B1 = .9 instead of .95.
The best performing among all of these achieved final loss of 3.12 while the best Shampoo run achieved final loss 3.
22


Reviews:
[{"review_id": "apoLy84jrg", "reviewer": "Reviewer_kaMq", "summary": "This paper introduces SOAP (ShampoO with Adam in the Preconditioner’s eigenbasis), an optimizer designed to improve the efficiency and stability of Shampoo by integrating Adam updates within Shampoo's eigenbasis. SOAP simplifies Shampoo by reducing its hyperparameters and computational overhead while maintaining its benefits.", "questions": "(1) What's the performance of SOAP in other tasks such as vision and graph tasks?\n\n(2) why can SOAP perform better than Shampoo? Is there any rationale behind this?\n\n(3) What's the performance of SOAP in fine-tuning settings?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "(1)  This paper is well-organized.\n\n(2) This work has significant implications for large language model training.", "weaknesses": "(1) The evaluation in this paper is limited, as it only tests the proposed SOAP on 360M and 660M language models. It is unclear whether SOAP can scale effectively to larger models, such as 1B or 7B, or how it performs in fine-tuning settings.\n\n(2) Given that SOAP is intended as a general optimizer, it would be valuable to assess its performance across diverse tasks, including vision and graph tasks, to better understand its generalization.\n\n(3) While the paper establishes a theoretical connection between Shampoo and Adafactor that motivates SOAP, it remains unclear why SOAP outperforms both Shampoo and Adam.\n\n(4) Since SOAP is inspired by the connection between Adafactor and Shampoo, it is essential to include Adafactor as a baseline to fully evaluate SOAP's effectiveness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ItpBmMf5F3", "reviewer": "Reviewer_C64v", "summary": "The paper establishes a connection between Shampoo and Adafactor and shows that Shampoo with the ½ power is equivalent to running Adafactor in the eigenbasis of Shampoo’s preconditioner. \nBuilding on this insight the paper proposes to continually update the running average of the second moment, just as Adam, but in the slowly changing coordinate basis. \nThe new optimizer SOAP only introduces one additional hyperparameter (the preconditioning frequency) compared to Adam.\nIn language modeling experiments up to model sizes of 660m parameters the authors demonstrate that SOAP outperforms Adam and Shampoo on pre-training tasks.", "questions": "- L.350: n -> In (typo)\n- Did you observe any influence of SOAP on training instabilities (e.g. grad norm slikes / rise or loss spikes) in language model pretraining?\n- It would be helpful for others adopting SOAP if the authors provide learning curves for other common metrics monitored during pre-training (e.g. parameter norms, gradienten norms, validation loss, etc.)? This would help compering other training setups to your suggested one.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper draw a novel connection between Shampoo and Adafactor, which has not been done before (to the best of my knowledge). \nThe paper is very well written and clear in the description of its method. The authors also provide code.\nIn general the conducted experiments are very thorough (especially the experiments on batch size scaling and efficiency) and demonstrate the benefits of SOAP compared to Adam or Shampoo. They authors also provide detailed analysis of the memory consumption of SOAP in the Appendix.", "weaknesses": "The main weakness of this paper (which is also acknowledged by the authors) are the rather small scale experiments in language model pre-training compared to current state of the art or production models. \nI believe that showing good performance compared to AdamW on the 1B to 7B parameter scale, would strongly support the adoption of SOAP. \nAlso, it would be interesting how the new optimizer SOAP performs in combination with newly emerging Linear Recurrent Models (like Mamba).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DdLjhc7kla", "reviewer": "Reviewer_vXc9", "summary": "In this paper, the authors propose SOAP, ShampoO with Adam in the Preconditioner’s eigenbasis. They starts with an insight: Shampoo is equivalent to running Adafactor in the eigenbasis of Shampoo’s preconditioner. SOAP continually updates the running average of the second moment as Adam does.In the large batch regime, SOAP reduces the number of iterations by over 40% and wall clock time by over 35% compared to AdamW, with approximately 20% improvements in both metrics compared to Shampoo.", "questions": "* In theory, how is this different from running GaLore under \"full\" mode + adam/adafactor? Please refer to the following code snippet  in official galore repo for the context: \n```python\nlow_rank_grad = torch.matmul(self.ortho_matrix[0].t().to(full_rank_grad.device.type), full_rank_grad) @ self.ortho_matrix[1].t().to(full_rank_grad.device.type)\n```\nhttps://github.com/jiaweizzhao/GaLore/blob/2cc66f88cce189e505affbb91042a8e77f5bf4e9/galore_torch/galore_projector.py#L39 \n\n* What happens if you use power iterations to approximate SVD in galore, does that still work? it feels like that should be included as a baseline too. [1] seems to suggest it works?\n\n* Why does higher frequency lead to worse results? Intuitively, it should provide better and more up-to-date approximation to the eigenspace?\n\n[1] Liang, Kaizhao, et al. \"Memory-efficient llm training with online subspace descent.\" arXiv preprint arXiv:2408.12857 (2024).", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Interesting connection and insights between Adafactor in Shampoo’s eigenspace and SOAP\n* The name is funny\n* Experiment results seem pretty promising", "weaknesses": "* I am a bit confused by the main contribution of this paper overall, algorithmetic insights on connection between adafactor and shampoo are interesting, hence I am recommending acceptance. However, beyond that, the rest seems to be a special case of galore [see my comments in the question below.]. Maybe a more efficient implementation?\n* As far as I know, in pytorch, SVD was also implemented via power iteration but at cuda kernel level. The difference proposed here seems to be how many iterations it runs.\n* Please also do some downstream evaluations on checkpoints output by different optimizers, just for the sake of completeness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2nAVGPd3l7", "reviewer": "Reviewer_SFbv", "summary": "I would like to preface my review: I am not an expert in this area. I have published a single paper in this area a very long time ago, but have not touched optimization since. I have asked several of my colleagues, who are among the most qualified researchers to review this paper, and they responded with very positive feedback. Still, I ask the AC to down-weight my review and defer to the other reviewers who may be more familiar with this space.\n\nThis work proposes a new optimization algorithm, called SOAP, that extends several of the (largely untested) concepts presented in Appendix B of the second Shampoo paper. In particular, they make changes to the computation of the Shampoo preconditioners---using dataset averages rather than running averages and computing the eigenvectors of L and R every $f$ steps---that make this algorithm now work.\n\nAdditionally, the authors prove a correspondence between the Shampoo parameter update and running Adafactor in the eigenbasis of Shampoo's preconditioner.  They then extend this to running Adam in Shampoo's preconditioner eigenbasis, and provide experimental results that demonstrate the effectiveness of their optimization algorithm.", "questions": "The authors write Soap is equivalent to running Adam in a rotated space; however, is transformation to the eigenbasis $G' \\leftarrow Q_L^TGQ_R$ necessarily unitary? What is the specific operation in SO(d)?\n\nSeveral nits: \n\nUsing $\\phi$ rather than $\\mathcal{L}$ or $L$ to denote the loss seems unusual (although perhaps this is to distinguish from the Shampoo preconditioner $L$).", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper is exceptionally well written and easy to follow. Prior and related work are excellently summarized, and the contribution of work in the context of the overarching field is clear. The experimental results evaluate on a very relevant setting (LLM pre-training), the baselines are appropriate (direct comparison with Shampoo and AdamW), and the results demonstrate that SOAP is quite effective. I also appreciate the discussion of the hyperparameter sweeps in the appendix -- not conducing this comparison has caused some unfortunate situations in the past within the optimization community.\n\nAs an aside, the naming is excellent (shampoo, rinse, lather, and now soap...).", "weaknesses": "To be frank, I do not see any substantial weaknesses. One could argue that it doesn't go far enough in certain directions, but these areas are acknowledged by the paper and \"left to future work\" (which I believe is appropriate given this paper's scope). This is, in my view, a very good paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)