SOAP: Improving and Stabilizing Shampoo using Adam

## Abstract

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 .

## 1 Introduction

With ever-increasing costs of LLM training, optimization efficiency has become a central question in the field of deep learning. Several recent works have tackled this challenge by addressing both the memory (Zhao et al., ; Wang et al., ) and compute (Anil et al., ) footprint of optimizers. In Algoperf (Dahl et al., ), a recent optimization efficiency benchmark, Shampoo (Gupta et al., ), a second-order algorithm, outperformed all other submissions, including Adam (Kingma & Ba, ), reducing wall-clock time by 28%. Higher-order preconditioning has also been applied in large-scale training runs, such as Gemini-1.5 Flash (Gemini Team, ).

The success of Shampoo has drawn increasing attention from the deep learning community. Several works have explored ways to scale Shampoo by improving its memory and compute efficiency (Wang et al., ; Anil et al., ; Shi et al., ). Other research (Morwani et al., ) has examined the theoretical foundations of Shampoo and proposed minor adjustments (such as using power $1/2$ rather than $1/4$) that align with prior empirical findings (Anil et al., ). Moreover, Morwani et al. () also showed that Shampoo with the aforementioned modifications is close to the optimal Kronecker approximation of the Adagrad (Duchi et al., ) optimizer.

Our first contribution is demonstrating that the variant of Shampoo proposed by Morwani et al. () is equivalent(^1^11Given this connection, the results of Morwani et al. () can be interpreted as showing that the eigenbasis provided by Shampoo’s preconditioner is close to the “optimal” basis for running Adafactor.) to running Adafactor (Shazeer & Stern, ; Zhai et al., ) in the eigenbasis provided by Shampoo’s preconditioner. This interpretation of Shampoo connects it to a broader family of methods (e.g. (George et al., )) that design second-order algorithms by running a first-order method in the eigenbasis provided by a second-order 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. (); Anil et al. (); Morwani et al. () as well as implementation ideas from works of Wang et al. (); Zhao et al. ().

Explicitly, we study SOAP (ShampoO with Adam in the Preconditioner’s eigenbasis) an algorithm that runs AdamW in the eigenbasis provided by Shampoo. Our main contributions are as follows:

- •
We make a formal connection between the Shampoo and the Adafactor algorithm. This insight leads us to consider the SOAP algorithm, which runs AdamW in the preconditioned space provided by Shampoo.
- •
SOAP outperforms both Shampoo and Adam in LLM pre-training tasks with model sizes 360m and 660m, even after extensive hyperparameter tuning of Shampoo.
- •
SOAP reduces the number of hyperparameters compared to Shampoo, resulting in only one additional hyperparameter compared to AdamW: preconditioning frequency.
- •
SOAP demonstrates greater robustness to large preconditioning frequency compared to Shampoo.

We should also note that while similar algorithmic variants have been discussed in the literature (e.g. see the appendix of Anil et al. ()), we are the first to systematically evaluate it.

Organization: In Section 3, we discuss related works. In Section 4, we start by showing an equivalence between Shampoo (with exponent 1/2) and running Adafactor in the eigenspace given by Shampoo, then with this equivalence as the starting point we describe SOAP. In Section 5, we provide our experimental methodology and in Section 6, we compare the performance of AdamW, Shampoo and SOAP on language modeling tasks. In Sections 7.2 and 7.3 we discuss the the space and time complexity of SOAP and how it can be improved.

Figure: Figure 1: Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo (Shi et al., ) implementation) and SOAP. In left and middle figures, Shampoo and SOAP use a preconditioning frequency of 10. The ”shorter LR schedule” plot is where we tuned the cosine decay so as to achieve the same terminal performance as AdamW. There we observe a $\geq 40\%$ reduction in the number of iterations and a $\geq 35\%$ reduction in wall clock time compared to AdamW, and approximately a 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 6 for a discussion of experimental results and ablation of batch size and Section 5 for experimental methodology.
Refer to caption: x1.png

## 2 Notation and Background

We denote the weight matrix of a neural network layer by $W\in\mathbb{R}^{m\times n}$, and the
corresponding gradient by $G\in\mathbb{R}^{m\times n}$. At a given time step $t$, these are denoted
as $W_{t}$ and $G_{t}$, respectively. For a batch of inputs at time $t$, denoted
by $B_{t}$, the loss and its gradient evaluated at $W_{t}$ are represented as
$\phi_{B_{t}}(W_{t})$ and $\nabla_{W}\phi_{B_{t}}(W_{t})$, respectively.

Adagrad (Duchi et al., ) is an online learning second-order algorithm that maintains a preconditioner $H\in\mathbb{R}^{mn\times mn}$. If the vectorized gradient at time $t$ is denoted by $g_{t}$ (i.e., $g_{t}=\text{vec}(G_{t})\in\mathbb{R}^{mn}$), then the update of the preconditioner and the vectorized weights $w_{t}\in\mathbb{R}^{mn}$ with learning rate $\eta$ is given by

$$ $H_{t}=H_{t-1}+g_{t}g_{t}^{\top};\quad w_{t}=w_{t-1}-\eta H_{t}^{-1/2}g_{t}$ $$

Adam (Kingma & Ba, ), a widely used first-order optimization algorithm in
deep learning is a diagonal approximation of Adagrad. It maintains an exponential moving average of the gradients
$G_{t}$ (denoted as $M_{t}$) and of element-wise squared gradients $G_{t}^{2}$
(denoted as $V_{t}$) for a given weight matrix $W$. Its update rule with
learning rate $\eta$ is given by

$$ $W_{t}\leftarrow W_{t-1}-\eta\frac{M_{t}}{\sqrt{V_{t}}},$ $$

where the division is performed element-wise.

Adafactor (Shazeer & Stern, ; Zhai et al., ), a variant of Adam, replaces $V_{t}$ with
its best rank-1 approximation $V_{t}^{\prime}$ to reduce memory usage. While the
original Adafactor paper (Shazeer & Stern, ) 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., ; Zhao et al., ), whose update
with learning rate $\eta$ is given by

$$ $W_{t}\leftarrow W_{t-1}-\eta\frac{M_{t}}{\sqrt{V_{t}^{\prime}}}.$ $$

Shampoo (Gupta et al., ) is a second-order optimization algorithm that approximates Adagrad and maintains two preconditioners, $L_{t}\in\mathbb{R}^{m\times m}$ and $R_{t}\in\mathbb{R}^{n\times n}$, for a given weight matrix $W\in\mathbb{R}^{m\times n}$. The updates for the preconditioners and the
weights with learning rate $\eta$ are as follows:

$$ $L_{t}\leftarrow L_{t-1}+G_{t}G_{t}^{T};\quad R_{t}\leftarrow R_{t-1}+G_{t}^{T} G_{t};\quad W_{t}\leftarrow W_{t-1}-\eta L_{t}^{-1/4}G_{t}R_{t}^{-1/4}.$ $$

In practice, Shampoo is implemented with several other modifications such as layerwise learning rate grafting and exponents other than $-1/4$. We will use the DistributedShampoo (Shi et al., ) implementation which has these variations available as hyperparameters.

## 3 Related Work

We begin by discussing works that are closely related, including George et al. (); Anil et al. () and Zhao et al. (). Subsequently, we cover extended related works.

KFAC (Martens & Grosse, ) is a well-known second-order optimization algorithm designed
for neural networks. E-KFAC (George et al., ) builds upon KFAC in a manner analogous to
our extension of Shampoo, introducing a diagonal preconditioner that is updated
between KFAC inversion steps. However, E-KFAC’s algorithm is not identical to running
Adam in KFAC’s eigenbasis, as the diagonal preconditioner is not Adam.

Anil et al. () introduced several algorithmic and numerical improvements to
develop a practical and scalable version of Shampoo (Gupta et al., ). 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. (), 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 et al., ) 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^{T}$ and $G^{T}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 B, we study GaLore’s performance and find that our 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 on the other side. We examine the impact of one-sided
projection for SOAP in Section 7.1.

Diagonal Preconditioning based Optimizers: Other than AdamW, there are other optimizers which involve diagonal preconditoning such as Lion (Chen et al., ), Sophia (Liu et al., ), and Adafactor (Shazeer & Stern, ). Recent works of Kaddour et al. (); Zhao et al. () showed that these optimizers perform comparably to AdamW for LLM pretraining but do not surpass 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, ; Martens & Grosse, )
optimize without explicitly computing the Hessian matrix, instead employing iterative
techniques to approximate the Newton step. Other recent works (Li, ; ; Pooladzandi & Li, ) 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, ) and Shampoo
(Gupta et al., ) are two widely recognized methods in this area.

KFAC (Martens & Grosse, ) was one of the first approaches to go beyond diagonal
preconditioners in neural networks, demonstrating that a layer-wise Kronecker-factored
preconditioner approximates the layer-wise Hessian in multi-layer perceptrons (MLPs).
Subsequent works (Martens et al., ; Osawa et al., ) extended KFAC to other
architectures. Recent research (George et al., ; Gao et al., ) has
further improved trace and diagonal estimates for KFAC. Efforts to scale up KFAC
(Ba et al., ; Puiu, ; ; Eschenhagen et al., )
have focused on making the inversion step more efficient or enhancing distributed
implementations.

Shampoo (Gupta et al., ), another second-order optimization algorithm, is motivated by
the online learning algorithm Adagrad (Duchi et al., ). Shampoo also employs
a layer-wise Kronecker-factored preconditioner. A recent distributed implementation of
Shampoo (Shi et al., ) won an optimization efficiency benchmark
(Dahl et al., ), highlighting the practical utility of second-order methods
in deep learning. Few recent works (Duvvuri et al., ; Morwani et al., ) have provided theoretical advancements on top of Shampoo. Other works (Anil et al., ; Peirson et al., ; Lin et al., ; Wang et al., ) have proposed various strategies to improve Shampoo’s scalability. We defer
a comparison of SOAP with these methods to future work.

## 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., ):

- 1.
We use power $1/2$ instead of power $1/4$. This was already recommended in practical implementations (Anil et al., ; Shi et al., ) and a theoretical connection between optimal Kronecker approximation of Adagrad (Duchi et al., ) preconditioner and Shampoo with power $1/2$ was established in Morwani et al. ().
- 2.
We also use the scalar correction to per layer learning rates described in Ren & Goldfarb (); Morwani et al. ().
- 3.
Instead of the running average of $L$ and $R$ across time steps, we 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 in Algorithms 1 and 2.

Figure: Algorithm 1 Single step of idealized Shampoo with power $1/2$.

Figure: Algorithm 2 Single step of idealized Adafactor in Shampoo’s eigenspace.

###### Claim 1 .

Algorithms 1 and 2 are equivalent.

###### Proof.

Consider $G_{t}$ in the basis created after rotating by $Q_{L},Q_{R}$ i.e. $G^{\prime}_{t}=Q_{L}^{T}G_{t}Q_{R}$. Let the eigenvalues of $\mathbb{E}_{B_{t}}[G_{B_{t}}G_{B_{t}}^{T}]$ and $\mathbb{E}_{B_{t}}[G_{B_{t}}^{T}G_{B_{t}}]$ be given by $\lambda_{1},...,\lambda_{m}$ and $\mu_{1},...,\mu_{n}$ respectively. Algorithm 1 scales the $i,j$ coordinate by $(\lambda_{i}\mu_{j}/(\sum_{i}\lambda_{i}))^{-1/2}$, while Algorithm 2 scales them by $(A_{i}C_{j}/(\sum_{i}A_{i}))^{-1/2}$. We now show that $A_{i}=\lambda_{i}$, an analogous argument shows $C_{j}=\mu_{j}$.

$$ $\displaystyle A_{i}$ $\displaystyle=e_{i}^{T}\mathbb{E}_{B}[G^{\prime}_{B}\odot G^{\prime}_{B}] \mathbf{1}_{m}$ $\displaystyle=\mathbb{E}_{B}[\sum_{j}(G^{\prime}_{B})_{i,j}^{2}]$ $\displaystyle=\mathbb{E}_{B}[\sum_{j}(u_{i}^{T}(G_{B})v_{j})^{2}]\quad\quad( \text{Using definition of }G^{\prime})$ $\displaystyle=\mathbb{E}_{B}[||u_{i}^{T}(G_{B})||^{2}]\quad\quad\quad\quad(v_{ j}\text{ form a basis})$ $\displaystyle=\mathbb{E}_{B}[u_{i}^{T}G_{B}G_{B}^{T}u_{i}]$ $\displaystyle=\lambda_{i}\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad( \text{By definition of }\lambda_{i}\text{ and }u_{i})$ $$

∎

While these two algorithms are 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, we do not invert or compute the eigenvector decomposition of $L$ and $R$ at every step. This means that the “adaptivity” of learning rates in Shampoo is limited(^2^22We note that practical implementations of Shampoo use grafting which allows for learning rate adaptivity at every step, but this adaptivity is restricted to a single scalar per layer.) to the updates of $L$ and $R$. In contrast, with Adafactor in Shampoo’s eigenspace, the second moment estimates (i.e., $A$ and $C$ in Algorithm 2) can be updated at every step as they are computationally inexpensive. Additionally, instead of using Adafactor, we can opt(^3^33Though using AdamW over Adafactor only gives very small improvements in performance, see Figure 5 and Section 7.2. We also note that one can use any other diagonal preconditioner based optimizer in place of Adam, such as Lion (Chen et al., ), Sophia (Liu et al., ) or Schedule-Free AdamW (Defazio et al., ).) for Adam, which offers more generality. Combining these insights leads to Algorithm 3 which can be interpreted as running Adam in Shampoo’s eigenspace.

Figure: Algorithm 3 Single step of SOAP for a $m\times n$ layer. Per layer, we maintain four matrices: $L\in\mathbb{R}^{m\times m},R\in\mathbb{R}^{n\times n}$ and $V,M\in\mathbb{R}^{m\times n}$. For simplicity we ignore the initialization and other boundary effects such as bias correction. Hyperparameters: Learning rate $\eta$, $\text{betas}=(\beta_{1},\beta_{2})$, epsilon ${\epsilon}$, and preconditioning frequency $f$. An implementation of SOAP is available at [https://github.com/nikhilvyas/SOAP](https://github.com/nikhilvyas/SOAP).

Figure: Algorithm 4 Eigenvectors function, implemented using power iteration and QR decomposition. Inputs: PSD matrix $P$ and estimate of eigenvectors $Q$. If the estimate was exact we would have $P=QDQ^{T}$ where $D$ is the diagonal matrix with eigenvalues.

We now describe some additional implementation details:

- 1.
Algorithm 3 describes the behavior of the algorithm for 2D layers. Following Zhao et al. (), for 1D layers we run standard AdamW. This reduces the overhead as compared to standard implementations of Shampoo which solve an eigenvector problem for 1D layers too.
- 2.
Following Wang et al. (), we compute eigenvectors of $L$ (and $R$) using one step of power method (Algorithm 4). This requires doing one matrix multiplication followed by QR decomposition. QR decomposition is faster (Documentation, ) 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 transformers, maintaining the eigenvectors would be space and time prohibitive. For such dimensions we fix the rotation matrix ($Q_{L}$ or $Q_{R}$) to be identity. Note that if we fix both $Q_{L}$ and $Q_{R}$ 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.

###### Claim 1 .

###### Proof.

## 5 Experimental Methodology

Hyperparameter tuning: We begin with hyperparameter values suggested by prior research for both AdamW and Distributed Shampoo (e.g., $\beta_{2}=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 we vary the selected hyperparameter alongside the learning rate. The purpose of these sweeps is to demonstrate that our default hyperparameter settings are near-optimal, disregarding potential interactions between two non-learning-rate hyperparameters. A detailed discussion of the hyperparameter sweeps is provided in Appendix A.

Figure: 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 to Section 5.
Refer to caption: x2.png

Throughput Measurement: We evaluate the throughput of each optimizer by measuring the number 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— as the overhead of Shampoo/SOAP is compared against multiple gradient accumulation steps— it is important to note that the overhead of Shampoo/SOAP can be amortized across layers by distributing the updates across multiple GPUs. This technique is employed in the distributed implementation of Shampoo (Shi et al., ). A comprehensive comparison of distributed implementations of these algorithms is left to future work.

Efficiency Benefits: 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,.625,.75$ and $.875$ fraction of the training data and fit a scaling law of the form $a+bN^{-\beta}$ through the final losses obtained, where $N$ represents the number of training points and $a,b,\beta$ 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 2. Here, the horizontal lines represent the final losses of AdamW and Shampoo.

## 6 Language Modeling Experiments

In this section we focus on empirically comparing AdamW, DistributedShampoo, and SOAP on language modeling tasks.

Figure: Figure 3: Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo (Shi et al., ) implementation) and SOAP. Shampoo and SOAP use preconditioning frequency of 10. We observe a $\geq 40\%$ reduction in the number of iterations and a $\geq 35\%$ reduction in wall clock time compared to AdamW, and approximately a $20\%$ reduction in both metrics compared to Shampoo. See Figure 1 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 improvement and experimental methodology.
Refer to caption: x3.png

### 6.1 Measuring Efficiency Benefits

In Figure 1 (left and middle) and Figure 3 we show train loss curves of 360m and 660m models with 2m token batch size for AdamW, Shampoo, and SOAP, where SOAP outperforms the other two. To directly calculate the efficiency benefit of SOAP, we also run SOAP with cosine decay for a shorter lr schedule, as shown in Figures 1 and 3. This allows us to approximate the following efficiency benefits (when setting batch size to 2m and preconditioning frequency to 10): $\geq 40\%$ reduction in number of iterations and $\geq 35\%$ reduction in wall clock time as compared to AdamW; $\approx 20\%$ reduction in iterations and wall clock time as compared to Shampoo. Precise efficiency benefit calculations are presented in  Figure 2(left and middle).

### 6.2 Effect of Frequency of Finding Eigenvectors/Inverse

In Figure 1 (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 are quite close in performance.
- •
At higher frequencies, the performance of both SOAP and Shampoo degrades but SOAP’s performance degrades significantly slower than Shampoo’s.

### 6.3 Effect of Batch Size

Figure: Figure 4: Comparing performance of tuned runs for AdamW, Shampoo (using DistributedShampoo (Shi et al., ) implementation) and SOAP for token batch size of 256k. Shampoo and SOAP use preconditioning frequency of 80. We observe a $\geq 25\%$ reduction in the number of iterations compared to AdamW, and approximately a 10% reduction compared to Shampoo. See Figure 2 (right) for wall-clock time improvement and Section 5 for detailed calculation of efficiency improvement.
Refer to caption: x4.png

In this section, we examine the impact of batch size on the performance of the Shampoo and SOAP optimizers. Specifically, we reduce the batch size by a factor of 8, from 2m to 256k. To maintain the same FLOPS overhead for the eigenvector decomposition steps as in the 2m setting, we increase the preconditioning frequency by a factor of 8, from 10 to 80. In Figure 4, we present the optimal runs for each optimizer. Our results show that SOAP consistently outperforms both Shampoo and AdamW, demonstrating a reduction of 25% or more in the number of iterations compared to AdamW, and approximately a 10% reduction compared to Shampoo. In Figure 2 (right), we show that SOAP also improves in wall-clock time by $\geq 15\%$ over AdamW and approximately $10\%$ over Shampoo. Note that we present these results as a preliminary analysis for small batch size runs. It is quite likely that our increase in preconditioning frequency by a factor of 8 is not optimal and a better trade-off is achievable. Furthermore, the overhead of SOAP can likely be ameliorated by doing $L$ and $R$ updates in lower precision (instead of fp32).

We also note that the decrease in efficiency improvements at smaller batch sizes for second-order methods is consistent with prior works (Zhang et al., ; Ishikawa & Yokota, ).

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

### 7.1 One Sided Eigenbasis

As described in Section 3, Zhao et al. () have an algorithm similar to ours. One of the differences 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 $Q_{R}=I_{n}$ in Algorithm 3 and if $m>n$ we set $Q_{L}=I_{m}$. Doing this leads to a reduction in space usage as well as reduction of optimizer time overhead, which is discussed in Sections 7.2.1 and 7.3.1.

In Figure 5, it is evident that the one-sided projection results in slightly reduced performance compared to the original SOAP optimizer. However, it 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.

Figure: Figure 5: Performance of variants of SOAP which improve space usage or time overhead. 1. SOAP (factorized): Uses Adafactor instead of Adam in Shampoo’s eigenbasis and 2. SOAP (one-sided): Uses $Q=I$ (i.e. no rotation) on the large side of weight matrix and 3. SOAP (factorized, one-sided): Combines both of these changes. We observe that while using Adafactor instead of Adam causes a negligible increase in loss, using the one-sided variant causes a larger increase. However, the one-sided variant also has much larger reduction in time and space overhead. For computational benefits of these variants see Sections 7.2 and 7.3.
Refer to caption: x5.png

### 7.2 Space usage of SOAP

For a $m\times n$ matrix where $m>n$ we require

$$ $2m^{2}\text{ (for $L,Q_{L}$)}+2n^{2}\text{ (for $R,Q_{R}$)}+3mn\text{ (for gradient, $M,V$)}$ $$

space usage(^4^44One $mn$ is for storing the gradients, this can be avoided (as long as there is no gradient accumulation) by applying gradients along with backprop (Lv et al., ) but this is not implemented by default in standard deep learning frameworks such as PyTorch. Hence we will include this term in all of our calculations.) (beyond weights and activations), specifically for $L,Q_{L},R,Q_{R},\text{momentum }(M)$, AdamW’s second order estimate ($V$), and the gradient. This is the same space usage as DistributedShampoo while AdamW uses $3mn$.

#### 7.2.1 Improving space usage of SOAP

The most direct way to reduce memory is using low precision to store the $L,R,Q_{L},Q_{R},V$ matrices, which is done by Dettmers et al. (); Wang et al. (). Orthogonal to the low precision approaches, there are two algorithmic approaches to improving the space usage of SOAP:

- •
Using Adafactor instead of Adam as the diagonal preconditioner after rotating by $Q_{L}$ and $Q_{R}$. This reduces the space usage by $mn$.
- •
Using one sided version of SOAP (Section 7.1). This reduces space usage from $2m^{2}+2n^{2}+3mn$ to $2\min(m,n)^{2}+3mn$.
- •
Combining these approaches yields space usage of $2\min(m,n)^{2}+2mn$.

For standard transformer architectures the last variant which combines the two approaches would yield less space usage overall compared to AdamW (which uses $3mn$).

We try these approaches in Figure 5. We observe that using Adafactor instead of AdamW yields very small reductions in performance while using one-sided preconditioner results in larger reductions. 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., ) and for language model training (Zhao et al., ). Further space reduction beyond Adafactor has been studied in the Adalomo (Lv et al., ), GaLore (Zhao et al., ), and AdaMeM (Vyas et al., ) papers.

### 7.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 eigenbasis). Let us first analyze the first one. For SOAP per step for a layer of size $m\times n$ we have an overhead of

$$ $m^{3}\text{ (updating $L$)}+n^{3}\text{ (updating $R$)}+(2m^{2}n+2mn^{2})\text { (projecting and projecting back on both sides)}.$ $$

We note that this is more than the overhead of Shampoo which is $m^{3}+n^{3}+m^{2}n+n^{2}m$. This can be observed in Figure 2 (bottom, right) but not in the other figures since there the second type of overhead is the dominant term.

The second type of overhead is due to changing the preconditioner for Shampoo (or for SOAP, preconditioner’s eigenbasis i.e. $Q_{L}$ and $Q_{R}$). The DistributedShampoo (Shi et al., ) implementation of Shampoo uses a direct call to torch.linalg.eigh for this. Following Wang et al. () we use Algorithm 4 which uses power iteration based approach which calls torch.linalg.qr. We note that torch.linalg.qr is faster than torch.linalg.eigh (Documentation, ). In Figure 6 (right) we see that using power iteration based approach (torch.linalg.qr) performs as well as fresh eigenvector decomposition (torch.linalg.eigh).

Figure: Figure 6: (Left) Depicting the overhead of SOAP over AdamW as a function of preconditioning frequency (Right) Comparing the performance of SOAP with torch.linalg.eigh for computing the eigenvectors with Algorithm 4, which uses torch.linalg.qr. Note that torch.linalg.qr is computationally more efficient than torch.linalg.eigh (as mentioned in Documentation ()); however, both seem to have comparable performance throughout the preconditioning frequency spectrum.
Refer to caption: x6.png

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/(\text{preconditioning frequency})$, approaching zero. However, empirical results (Figure 6 left) show that the overhead approaches an asymptote greater than zero. This is attributable to the additional matrix multiplications required to update $L$, update $R$, project the gradient, and reproject the gradient (for each layer) in the optimizer. Currently, these operations are performed in float32; reducing the precision of these operations, as proposed in Wang et al. (), could lower this asymptote.

#### 7.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,Q_{L},Q_{R},V$ matrices (Dettmers et al., ; Wang et al., ), 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.linalg.qr does not support precision lower than float32. 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 (Section 7.2) as the diagonal preconditioner after rotating by $Q_{L}$ and $Q_{R}$. In this version of SOAP the overhead can be improved by from $m^{3}+n^{3}+2m^{2}n+2n^{2}m$ to $m^{3}+n^{3}+m^{2}n+n^{2}m+\max(m,n)^{2}\min(m,n)+\min(m,n)^{3}$ by merging the project and project back steps for the smaller dimension.
- •
Using one sided version of SOAP (Section 7.1). This reduces overhead from $m^{3}+n^{3}+2m^{2}n+2n^{2}m$ to $\min(m,n)^{3}+2\min(m,n)^{2}\max(m,n)$.
- •
Combining these approaches yields an overhead of $\min(m,n)^{2}\max(m,n)+2\min(m,n)^{3}$

Using one-sided version also reduces the second type of overhead from a calls to torch.linalg.qr on a $m\times m$ and a $\ n\times n$ matrix to only a single call to $\min(m,n)\times\min(m,n)$ matrix.

## 8 Discussion and Future Work

We study an optimizer called SOAP: ShampoO with Adam in the Preconditioner’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. For future work, we would like to explore further improvements to the design of SOAP, in particular, related to using lower precision for the preconditioners as well as a better distributed implementation. We would also like to explore the performance of SOAP on other domains such as vision.

## Appendix A Experimental Setup

Many aspects of our setup such as models are the same as in Zhao et al. (). We will restate those details verbatim for completeness.

We train language models on C4 tokenized with the T5 tokenizer (Raffel et al., ) and report results in terms of validation loss.

##### Models.

We start from the OLMo codebase (Groeneveld et al., ) and train decoder-only transformer models of two sizes: 210m, 360m, and 660m, where the parameter count refers to non-embedding parameters. The models have widths of 1024, 1024, and 1408 and depths of 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, ). We use RoPE positional encodings (Su et al., ). Attention heads are always dimension 64. We use PyTorch default LayerNorm. We use QK layer norm (Dehghani et al., ). Following Wortsman et al. () we do not learn biases for the linear layers or LayerNorms. We train in mixed precision with bfloat16.

##### Algorithms.

We use the standard Pytorch implementation of AdamW (Paszke et al., ), the DistributedShampoo Shi et al. () implementation of Shampoo. We implement ourselves SOAP and GaLore starting from an older version of Pytorch implementation of AdamW and the official GaLore implementation Zhao et al. ().

##### Default hyperparameters.

We use $\beta_{1}=0.95$, as we found it to outperform $\beta_{1}=0.9$ in our sweeps for the 360m model. Following Wortsman et al. () we use decoupled weight decay with coefficient $110-4$ and z-loss with coefficient $110-4$. We use the default value of ${\epsilon}=1e-8$ in AdamW (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.1x$ the maximum learning rate.

##### Default hyperparameters for DistributedShampoo

Shi et al. () state that they find the optimal exponent to be either $-1/2$ or $-1.82/4\approx-1/2.2$. Our preliminary findings were similar to this. Hence we set the default values of exponent to be $-1/2.5$ for both 1D and 2D parameters. We set ${\epsilon}_{\text{shampoo}}=$110-12$$ and $\beta_{\text{shampoo}}=0.95$ based on our initial set of experiments on the 210m model.

##### Default hyperparameters for GaLore

GaLore introduces an additional hyperparameter called scale ($\alpha$) since due to low rank updates the overall update magnitude decreases. Since we are running a full rank version of GaLore we set $\alpha=1$.

##### 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 $\approx$ 2m. We default to training models for the approximately “chinchilla optimal” number of tokens that is $\approx$20 times the number of parameters. Explicitly, this means for our default batch size of 2m, the 210m models are trained for 1600 steps or $\approx$ 3.3b tokens. The 360m models are trained for 3200 steps, the 660m models are trained for 6400 steps.

### A.1 Sweeping over hyperparameters

AdamW, 2m batch size: Starting from the default hyperparameters above we do the following sweeps:

- 1.
We sweep over learning rate in $\{.1,.0316,.01,\ldots,$3.1610-4$\}$.
- 2.
(360m) We sweep over the cross product of best 3 learning rates and $\beta_{1}\in\{0.9,0.95,0.99\}$.
- 3.
(360m) We sweep over the cross product of best 3 learning rates and $\beta_{2}\in\{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 we only sweep over learning rate for the 660m model with 2m batch size.

DistributedShampoo, 2m batch size: Starting from the default hyperparameters above we do the following sweeps:

- 1.
We sweep over learning rate in $\{.1,.0316,.01,\ldots,$3.1610-4$\}$.
- 2.
(360m) We sweep over over the cross product of best 3 learning rates from above and ${\epsilon}_{\text{shampoo}}\in\{$110-11$,$110-12$,$110-13$\}$.
- 3.
(360m) We sweep over over the cross product of best 3 learning rates from above and $\beta_{\text{shampoo}}\in\{.9,.95,.975\}$.
- 4.
Let $e_{1},e_{2}$ denote the exponents used in DistributedShampoo for 1D and 2D parameters respectively. We also sweep over the cross product of best 3 learning rates from above and $(e_{1},e_{2})$ in $\{(2,2),(2.5,2.5),(3,3),(2,4)\}$.

These sweeps did not yield any significant improvement in performance ($<.004$) for the 360m model. Hence we 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,.0316,.01,\ldots,$3.1610-4$\}$.

AdamW, 256k batch size:
For the 360m model with 256 batch size we start from the default hyperparameters and do the following sweeps:

- 1.
We sweep over learning rate in $\{.1,.0316,.01,\ldots,$3.1610-4$\}$.
- 2.
We sweep over the cross product of best 3 learning rates and $\beta_{2}\in\{0.95,0.99\}$.

In the second sweep we observe small improvements in performance by using $\beta_{2}=.99$, so our final numbers use $\beta_{2}=.99$. This (small) improvement in performance by using a larger $\beta_{2}$ at smaller batch sizes was also observed by Porian et al. (); Zhao et al. ().

DistributedShampoo, 256k batch size:
For the 360m model with 256 batch size we start from the default hyperparameters and do the following sweeps:

- 1.
We sweep over learning rate in $\{.1,.0316,.01,\ldots,$3.1610-4$\}$.
- 2.
We sweep over the cross product of best 3 learning rates and $(\beta_{2},\beta_{\text{shampoo}})\in\{(.95,.95),(.99,.99)\}$.

In the second sweep we observe small improvements in performance by using $\beta_{2}=\beta_{\text{shampoo}}=.99$, so our final numbers use $\beta_{2}=\beta_{\text{shampoo}}=.99$.

SOAP, 256k batch size:
For the 360m model with 256 batch size we start from the default hyperparameters and do the following sweeps:

- 1.
We sweep over learning rate in $\{.1,.0316,.01,\ldots,$3.1610-4$\}$.
- 2.
We sweep over the cross product of best 3 learning rates and $\beta_{2}\in\{.95,.99\}$.

In the second sweep we observe small improvements in performance by using $\beta_{2}=.99$, so our final numbers use $\beta_{2}=.99$.

Preconditioning frequency sweeps: For the preconditioning frequency experiments of SOAP and Shampoo ( Figure 1 (right)), for each frequency we do a 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.

360m and 660m shorter runs: For each of the shorter runs of 360m and 660m models for the SOAP optimizer (Figure 2), we did learning rate sweep over the best 3 learning rates found for the standard length run. Other hyperparameters are set to their optimal values obtained using the standard length run.

Warmup: The warmup duration for the 360m and 660m models were 600 and 1200 steps respectively. 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 durations were 600, 750, 900 and 1050 steps for 0.5, 0.625, 0.75 and 0.875 runs respectively.

## Appendix B 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:

- 1.
We swept the cross product over learning rate ($3.16e-4,1e-3,3.16e-3,1e-2$), preconditioning frequency ($10,50,200$), both sided and one sided versions. Frequency 200 had the best results matching the observation of Zhao et al. ().
- 2.
We did a cross product sweep over learning rate ($3.16e-4,1e-3,3.16e-3,1e-2$), both sided and one sided versions with $\beta_{2}=.99$ instead of $.95$ and preconditioning frequency 200.
- 3.
We did a cross product sweep over learning rate ($3.16e-4,1e-3,3.16e-3,1e-2$), both sided and one sided versions, preconditioning frequency ($50,200$) with $\beta_{1}=.9$ instead of $.95$.

The best performing run among all of these achieved a final loss of 3.12 while the best Shampoo run achieved a final loss of 3.10.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: (1) What's the performance of SOAP in other tasks such as vision and graph tasks?

(2) why can SOAP perform better than Shampoo? Is there any rationale behind this?

(3) What's the performance of SOAP in fine-tuning settings?
soundness: 2
strengths: (1)  This paper is well-organized.

(2) This work has significant implications for large language model training.
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.
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.

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

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

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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - L.350: n -> In (typo)
- Did you observe any influence of SOAP on training instabilities (e.g. grad norm slikes / rise or loss spikes) in language model pretraining?
- 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.
soundness: 3
strengths: This paper draw a novel connection between Shampoo and Adafactor, which has not been done before (to the best of my knowledge). 
The paper is very well written and clear in the description of its method. The authors also provide code.
In 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.
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. 
Building 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. 
The new optimizer SOAP only introduces one additional hyperparameter (the preconditioning frequency) compared to Adam.
In language modeling experiments up to model sizes of 660m parameters the authors demonstrate that SOAP outperforms Adam and Shampoo on pre-training tasks.
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. 
I believe that showing good performance compared to AdamW on the 1B to 7B parameter scale, would strongly support the adoption of SOAP. 
Also, it would be interesting how the new optimizer SOAP performs in combination with newly emerging Linear Recurrent Models (like Mamba).

---

confidence: 5
contribution: 2
initial_rating: 6
presentation: 3
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: 
```python
low_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)
```
https://github.com/jiaweizzhao/GaLore/blob/2cc66f88cce189e505affbb91042a8e77f5bf4e9/galore_torch/galore_projector.py#L39 

* 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?

* Why does higher frequency lead to worse results? Intuitively, it should provide better and more up-to-date approximation to the eigenspace?

[1] Liang, Kaizhao, et al. "Memory-efficient llm training with online subspace descent." arXiv preprint arXiv:2408.12857 (2024).
soundness: 3
strengths: * Interesting connection and insights between Adafactor in Shampoo’s eigenspace and SOAP
* The name is funny
* Experiment results seem pretty promising
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.
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?
* 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.
* Please also do some downstream evaluations on checkpoints output by different optimizers, just for the sake of completeness.

---

confidence: 2
contribution: 4
initial_rating: 8
presentation: 4
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)?

Several nits: 

Using $\phi$ rather than $\mathcal{L}$ or $L$ to denote the loss seems unusual (although perhaps this is to distinguish from the Shampoo preconditioner $L$).
soundness: 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.

As an aside, the naming is excellent (shampoo, rinse, lather, and now soap...).
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.

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

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

Meta-review: This paper proposes SOAP, a novel optimization method that builds on the insight that Shampoo is equivalent to running Adafactor in the eigenbasis of Shampoo’s preconditioner. SOAP updates the running average of the second moment in a manner similar to Adam, achieving significant efficiency improvements in large-batch training. 

The writing is clear, the method is well-described, and the authors provide thorough experiments, particularly on batch size scaling and efficiency, which highlight the benefits of SOAP. 

However, the paper’s experimental setup has notable limitations. The scale of the experiments is relatively small, the settings are not practical for real-world applications, and there is a lack of downstream task evaluations or fine-tuning results. More practical results are particularly concerning, as lower loss does not always correlate with better performance, especially in fine-tuning scenarios. Moreover, the experiments are verified on the limited types of tasks, limiting the generalizability of the results.

Despite these weaknesses, the paper introduces a highly novel idea that bridges the gap between first-order and second-order optimizers effectively. Given its theoretical contribution and potential practical impact, I recommend weak acceptance.

Decision: Accept (Poster)