SOAP: IMPROVING AND STABILIZING SHAMPOO USING ADAM

Introduction

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., 2024; Anil et al., 2020; Shi et al., 2023). Other research (Morwani et al., 2024) 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., 2020).

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 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 to Shampoo, resulting in only one additional hyperparameter compared to AdamW: preconditioning frequency.
• SOAP demonstrates greater robustness to large preconditioning frequency compared to Shampoo on language model pre-training tasks.

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.

Adam (Kingma & Ba, 2015), 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

where the division is performed element-wise.

Adafactor (Shazeer & Stern, 2018; Zhai et al., 2022), a variant of Adam, replaces $V_t$ with its best rank-1 approximation $V'_t$ 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 $\eta$ is given by

Shampoo (Gupta et al., 2018b) 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:

Algorithms

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

• 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 a theoretical connection between optimal Kronecker approximation of Adagrad (Duchi et al., 2011b) preconditioner and Shampoo with power $1/2$ was established in Morwani et al. (2024).
• We also use the scalar correction to per layer learning rates described in Ren & Goldfarb (2021); Morwani et al. (2024).
• Instead of the running average of $L$ and $R$ across time steps, we use dataset averages.

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

• Sample batch $B_t$.
• $G_t \in \mathbb{R}^{m \times n} \leftarrow -\nabla_W \phi_{B_t} (W_t)$
• $L \leftarrow \mathbb{E}_B [G_B G_B^T]$ (Where the expectation is over a random batch B.)
• $R \leftarrow \mathbb{E}_B [G_B^T G_B]$
• $\hat{H} \leftarrow (L \otimes R) / \operatorname{tr}(L)$
• $W_{t+1} \leftarrow W_t - \eta \hat{H}^{-1/2} G_t = W_t - \eta L^{-1/2} G_t R^{-1/2} / \operatorname{tr}(L)^{-1/2}$

Algorithm 2 Single step of idealized Adafactor in Shampoo's eigenspace.

• Sample batch $B_t$.
• $G_t \in \mathbb{R}^{m \times n} \leftarrow -\nabla_W \phi_{B_t} (W_t)$
• $L \leftarrow \mathbb{E}_B [G_B G_B^T]$
• $R \leftarrow \mathbb{E}_B [G_B^T G_B]$
• $Q_L \leftarrow \text{Eigenvectors}(L)$
• $Q_R \leftarrow \text{Eigenvectors}(R)$
• $G'_t \leftarrow Q_L^T G_t Q_R$
• (Idealized version of Adafactor taking $G'_t$ to be the gradient)
• $G'_B \leftarrow Q_L^T G_B Q_R$
• $A = \mathbb{E}_B [G'_B \odot G'_B] \mathbf{1}_n$ where $G'_B = Q_L^T G_B Q_R$
• $C = \mathbf{1}_m^T \mathbb{E}_B [G'_B \odot G'_B]$
• $\hat{V}_t = (A C^T)/(\mathbf{1}_m^T A)$ (Elementwise division)
• $G''_t \leftarrow (G'_t)/(\sqrt{\hat{V}_t} + \epsilon)$ (Elementwise division and square root)
• $G'''_t \leftarrow Q_L G''_t Q_R^T$
• $W_{t+1} \leftarrow W_t - \eta G'''_t$

Proposition: Algorithms 1 and 2 are equivalent.

Proof: Consider $G_t$ in the basis created after rotating by $Q_L, Q_R$ i.e. $G'_t = Q_L^T G_t Q_R$. Let the eigenvalues of $\mathbb{E}_B[G_B G_B^T]$ and $\mathbb{E}_B[G_B^T G_B]$ be given by $\lambda_1, \dots, \lambda_m$ and $\mu_1, \dots, \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$.

$$\begin{aligned} A_i &= e_i^T \mathbb{E}_B [G'_B \odot G'_B] \mathbf{1}_n \\ &= \mathbb{E}_B [\sum_j (G'_B)_{i, j}^2] \\ &= \mathbb{E}_B [\sum_j (u_i^T G_B v_j)^2] \quad \text{(Using definition of } G'\text{)} \\ &= \mathbb{E}_B [\|u_i^T G_B\|^2] \quad \text{(}v_j \text{ form a basis)} \\ &= \mathbb{E}_B [u_i^T G_B G_B^T u_i] \\ &= \lambda_i \quad \text{(By definition of } \lambda_i \text{ and } u_i\text{).} \end{aligned}$$

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 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 for Adam, which offers more generality. Combining these insights leads to Algorithm 3 which can be interpreted as running Adam in Shampoo's eigenspace.

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$, betas $= (\beta_1, \beta_2)$, epsilon $\epsilon$, and preconditioning frequency $f$.

• Sample batch $B_t$.
• $G \in \mathbb{R}^{m \times n} \leftarrow -\nabla_W \phi_{B_t}(W_t)$
• $G' \leftarrow Q_L^T G Q_R$
• $M \leftarrow \beta_1 M + (1 - \beta_1) G$
• $M' \leftarrow Q_L^T M Q_R$
• {Now we "run" Adam on $G'$}
• $V \leftarrow \beta_2 V + (1 - \beta_2) (G' \odot G')$ (Elementwise multiplication)
• $N' \leftarrow M' / (\sqrt{V} + \epsilon)$ (Elementwise division and square root)
• {Now that we have preconditioned by Adam in the rotated space, we go back to the original space.}
• $N \leftarrow Q_L N' Q_R^T$
• $W \leftarrow W - \eta N$
• {End of gradient step, we now update $L$ and $R$ and possibly also $Q_L$ and $Q_R$.}
• $L \leftarrow \beta_2 L + (1 - \beta_2) G G^T$
• $R \leftarrow \beta_2 R + (1 - \beta_2) G^T G$
• if $t \% f == 0$ then
  • $Q_L \leftarrow \text{Eigenvectors}(L, Q_L)$
  • $Q_R \leftarrow \text{Eigenvectors}(R, Q_R)$
• end if

We now describe some additional implementation details:

• Algorithm 3 describes the behavior of the algorithm for 2D layers. Following Zhao et al. (2024a), 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.
• Following Wang et al. (2024), 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 than standard eigenvector decomposition in PyTorch. For the first iteration, eigenvectors are initialized by doing a standard eigenvector decomposition.
• 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.
• Algorithm 3 omits bias correction and weight decay for simplicity, but these are used in the actual implementation, identical to their use in AdamW.

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.

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., 2023).

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 + b N^{-\beta}$ through the final losses obtained, where $N$ represents the number of training points and $a, b, \beta$ are the parameters of the fit. This scaling law is then used to calculate the efficiency benefit in terms of training steps and wallclock time. Here, the horizontal lines represent the final losses of AdamW and Shampoo.

Language Modeling Experiments

Further Efficiency Improvements

One Sided Eigenbasis

Doing this leads to a reduction in space usage as well as reduction of optimizer time overhead.

Space Usage of SOAP

For a $m \times n$ matrix, we require

space usage.

The most direct way to reduce memory is using low precision to store the $L, R, Q_L, Q_R, V$ matrices. 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 $m n$.
• Using one sided version of SOAP. This reduces space usage from $2m^2 + 2n^2 + 3m n$ to $2 \min(m, n)^2 + 3m n$.
• Combining these approaches yields space usage of $2 \min(m, n)^2 + 2m n$.

We try these approaches and 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.

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

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. 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 as the diagonal preconditioner after rotating by $Q_L$ and $Q_R$.
• Using one sided version of SOAP. This reduces overhead from $m^3 + n^3 + 2m^2n + 2n^2m$ to $\min(m, n)^3 + 2 \min(m, n)^2 \max(m, n)$.
• Combining these approaches yields overhead of $2\min(m, n)^3 + \min(m, n)^2 \max(m, n)$.