---
title: 'SOAP: Second-Order Optimization with Adam Eigenbasis'
url: https://www.emergentmind.com/topics/second-order-optimization-with-soap-shampoo-with-adam-in-the-preconditioner-s-eigenbasis
type: topic
---

# SOAP: Second-Order Optimization with Adam Eigenbasis

Shampoo with Adam in the Preconditioner’s Eigenbasis (SOAP) is a second-order optimization algorithm for deep neural network training that blends Kronecker-factored curvature modeling, as in Shampoo, with Adam-style adaptivity by operating the adaptive moment logic directly in the second-order preconditioner’s eigenbasis. SOAP has been developed to exploit richer curvature information than Adam while maintaining computational efficiency and robustness across large-scale neural network applications.

## 1. Algorithmic Structure and Update Rule

SOAP operates through a two-stage procedure per parameter block. For a layer parameterized by weight matrix $W \in \mathbb{R}^{m \times n}$ with loss gradient $G = \nabla_W \ell$, the “vectorized” form $w \in \mathbb{R}^{mn}$ and gradient $g = \text{vec}(G)$ are preconditioned as follows:

- **Preconditioner Construction:** SOAP tracks Kronecker-factored second-moment estimators
  $$
  L \leftarrow \beta_2 L + (1-\beta_2) GG^T,\qquad R \leftarrow \beta_2 R + (1-\beta_2) G^T G
  $$
  where $L \in \mathbb{R}^{m \times m}$ and $R \in \mathbb{R}^{n \times n}$ accumulate gradient covariances along rows and columns, respectively.

- **Eigenbasis Rotation:** The matrices are (infrequently) eigendecomposed
  $$
  L = Q_L \Lambda_L Q_L^T,\qquad R = Q_R \Lambda_R Q_R^T
  $$
  providing orthogonal eigenbases $Q_L, Q_R$ and diagonal eigenvalue matrices $\Lambda_L, \Lambda_R$.

- **Gradient Projection and Adam Moments:** The gradient is rotated into the preconditioner’s eigenbasis,
  $$
  G' = Q_L^T G Q_R
  $$
  vectorized as $g'$. Adam-style moment estimates are maintained in this basis:
  $$
  m_t = \beta_1 m_{t-1} + (1-\beta_1) g'_t,\qquad v_t = \beta_2 v_{t-1} + (1-\beta_2) (g'_t)^{\odot 2}
  $$
  with bias correction as in conventional Adam.

- **Adaptive Preconditioned Step:** The update in the eigenbasis is
  $$
  n'_t = \frac{m_t}{\sqrt{v_t} + \epsilon}
  $$
  and the adaptive step is rotated back to the parameter space:
  $$
  N = Q_L n' Q_R^T
  $$
  so the parameter update is
  $$
  W_t \leftarrow W_{t-1} - \eta N
  $$

This approach can be succinctly written in vectorized form as
$$
w \leftarrow w - \eta (Q_R \otimes Q_L) \left[ \frac{m_t}{\sqrt{v_t}+\epsilon} \right]
$$
where $Q_R \otimes Q_L$ denotes the Kronecker product eigenbasis.

SOAP introduces an additional hyperparameter, the **preconditioning frequency** $K$, determining how often eigenbases $(Q_L, Q_R)$ are recomputed. Between recomputations, Adam moments are updated in the most recent basis [2409.11321][2506.03595][2509.22938][2601.20769].

## 2. Theoretical Foundations: Whitening and Curvature

SOAP is motivated by a whitening perspective. The ideal Newton step is
$$
w \leftarrow w - \eta \Sigma^{-1/2} g
$$
for the whitening matrix $\Sigma = \mathbb{E}[gg^T]$. For computationally tractable approximations:
- Adam is diagonal: $H_{\text{Adam}} = \operatorname{diag}(\mathbb{E}[gg^T])$
- Shampoo uses the Kronecker-product structure: $\Sigma \approx R \otimes L$

Shampoo effectively performs a Kronecker product whitening, utilizing Kronecker-factored approximations $L^{-\frac{1}{2}}$ and $R^{-\frac{1}{2}}$ such that
$$
w \leftarrow w - \eta (R \otimes L)^{-\frac{1}{2}} g
$$
This is optimal in Frobenius norm for Kronecker approximations [2509.22938][2506.03595].

SOAP further rotates the gradient to diagonalize the Kronecker preconditioner (the eigenbasis), then applies Adam’s diagonal scaling in that rotated space. Theoretical results establish that, under the exact Kronecker structure assumption,
SOAP’s adaptive Adam step in the eigenbasis and Shampoo’s fixed-magnitude Kronecker step become identical, as proven in Theorem 1 of [2509.22938].

## 3. Pseudocode and Computational Properties

SOAP maintains the computational structure of Shampoo but incorporates additional matrix multiplications for basis rotations and Adam moment updates in the rotated space. Typical per-layer pseudocode (excluding bias-correction and 1D handling):

```
for each step t:
    # Gradient and second-moment update
    G_t = -∇W φ(W_{t-1})
    L_t = β_2 L_{t-1} + (1-β_2) G_t G_t^T
    R_t = β_2 R_{t-1} + (1-β_2) G_t^T G_t
    # Eigendecompose occasionally (every K steps)
    if t mod K == 0:
        Q_L, Λ_L = eigh(L_t)
        Q_R, Λ_R = eigh(R_t)
    # Project gradient; update Adam moments
    g'_t = (Q_L ⊗ Q_R)^T vec(G_t)
    m_t = β_1 m_{t-1} + (1-β_1) g'_t
    v_t = β_2 v_{t-1} + (1-β_2) (g'_t ∘ g'_t)
    # SOAP step
    delta_theta_t = (Q_L ⊗ Q_R) [m_t / (sqrt(v_t)+ε)]
    # Update weights
    W_t = W_{t-1} - η mat(delta_theta_t)
```

The dominant cost per preconditioner update is $\mathcal{O}(m^3 + n^3)$ for eigendecomposition; when $K$ is large, this amortizes to $\mathcal{O}((m^3+n^3)/K)$ per step. Per-step matrix multiplications for projection/reconstruction cost $\mathcal{O}(mn(m+n))$. Overall, for typical transformer layers on current GPUs, training with $K=10$ gives a throughput drop of $\approx20\%$ compared to AdamW [2409.11321]. Extra storage for the preconditioner factors and momenta increases per-layer memory overhead only moderately.

## 4. Empirical Behavior and Practical Implications

Empirical studies show SOAP matches or outperforms Adam and Shampoo in convergence speed and/or wall-clock efficiency across large-scale language modeling and learned image compression:
- In large-batch language modeling (transformers with 360 M–660 M parameters, batch 2M tokens), SOAP reduces the number of iterations by over 40% and wall-clock time by over 35% compared to AdamW, and by ∼20% compared to Shampoo [2409.11321].
- For learned image compression, SOAP yields 65–75% fewer steps and 51–65% less wall-clock time than Adam, with 2–4% BD-Rate improvements after convergence [2601.20769].
- On vision transformers and graph workloads, SOAP (known as EShampoo in [2506.03595]) attains or exceeds Shampoo with Adam-grafting in both iteration and total time to target accuracy.

Across diverse tasks, as $K$ increases (infrequent preconditioner updates), SOAP’s performance degrades much more slowly than that of naive Shampoo. For example, at $K=100$ on language modeling, SOAP still outperforms AdamW by ∼30%, whereas Shampoo loses advantage [2409.11321].

Final validation or loss plateaus are often closely matched between SOAP, Shampoo, and Adam. When Shampoo’s raw Kronecker step is too aggressive and plateaus at higher loss, SOAP frequently recovers Adam-like final loss while maintaining Shampoo’s rapid convergence [2509.22938].

## 5. Practical Guidelines and Hyperparameter Selection

SOAP requires no learning rate grafting or unusual schedule tuning. Hyperparameters:
- Learning rate $\eta$ and Adam’s $\beta_1, \beta_2$ are chosen as in AdamW; e.g., $\beta_1=0.9$ or $0.95$, $\beta_2=0.99$ or $0.95$ for large batches, and $\epsilon=10^{-8}$.
- Preconditioning frequency $K$: $K=1$ is equivalent to Shampoo; $K\sim10$ is typical for large models; $K$ can be increased to reduce overhead at some accuracy loss.
- Trace normalization and damping as in Shampoo to stabilize scale.
- 1D parameters (biases, LayerNorm) are preconditioned with vanilla Adam.

Integration is straightforward: a single optimizer swap suffices in existing codebases [2409.11321][2601.20769]. Eigenbasis updates can be accelerated using power-iteration or warm-started QR routines, and $Q_L, Q_R$ can be stored in reduced precision (float16 or bfloat16) [2409.11321][2506.03595].

SOAP also supports adaptive scheduling of eigenbasis updates via a tolerance on the off-diagonal energy in the current basis, further reducing unnecessary preconditioner computations without increasing error [2506.03595].

## 6. Theoretical Equivalence and Gradient Alignment

SOAP is theoretically equivalent to running Adam (or Adafactor) in the rotated basis defined by the Shampoo preconditioner for layerwise gradients [2509.22938][2409.11321]. In the idealized Kronecker world, it is fully equivalent to Shampoo by performing elementwise adaptation in the approximately whitened basis. The per-step update magnitude in SOAP is provably bounded between the maximal and minimal eigenvalue scalings from full-matrix Adam, eliminating the $m^{-1/4}n^{-1/4}$ shrinkage observed in Shampoo [2506.03595].

From the gradient–whitening viewpoint, SOAP achieves two-stage approximate whitening: blockwise rotation (Shampoo), then diagonal whitening (Adam) in the rotated space. This construction allows SOAP to resolve intra-step and inter-step gradient conflicts, aligning rate and distortion gradients better than diagonal optimizers [2601.20769]. The result is more regular update directions, stabilizing deep neural network training and benefiting downstream applications such as quantization in compression models (outlier suppression by >50%) [2601.20769].

## 7. Extensions, Limitations, and Empirical Robustness

SOAP’s main cost is preconditioner eigendecomposition; for large $m, n$, this can dominate per-step compute. However, with moderate $K$, the additional wall-clock time over AdamW is limited and continues to decrease as batch sizes increase. Adaptive criteria for basis updates can reduce this further [2506.03595].

Empirical evidence confirms robustness of SOAP to a wide range of $K$, whereas Shampoo degrades sharply with infrequent updates. Adaptive eigenbasis updates allow for accurate step computation while reducing frequency of expensive decompositions. SOAP also eliminates the need for learning rate grafting and other heuristics introduced to compensate for Shampoo’s scaling errors, directly correcting scaling and staleness in the eigenvalues themselves.

SOAP is compatible with distributed training, and the preconditioner computations can be amortized over devices. The optimizer is applicable to a range of architectures, including transformers and convolutional nets, and tasks such as language modeling, vision, and compression, with no task-specific modifications required.

---

**Key References:**
- "SOAP: Improving and Stabilizing Shampoo using Adam" [2409.11321]
- "Understanding SOAP from the Perspective of Gradient Whitening" [2509.22938]
- "Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner" [2506.03595]
- "Leveraging Second-Order Curvature for Efficient Learned Image Compression: Theory and Empirical Evidence" [2601.20769]

Source: https://www.emergentmind.com/topics/second-order-optimization-with-soap-shampoo-with-adam-in-the-preconditioner-s-eigenbasis