SOAP (ShampoO with Adam in the Preconditioner's Eigenbasis)
Algorithms AIMech:SoapShampooWithAdamInPreconditionerEigenbasis DRAFT
Aliases: SOAP, SOAP Optimizer, ShampoO with Adam in the Preconditioner's eigenbasis
A second-order stochastic optimization algorithm (ShampoO with Adam in the Preconditioner's eigenbasis) that maintains Kronecker-factored gradient covariance EMAs L <- beta_2 L + (1 - beta_2) G G^T and R <- beta_2 R + (1 - beta_2) G^T G along with their eigenvector bases Q_L, Q_R (updated every f steps via power iteration and QR decomposition), projects the gradient G' = Q_L^T G Q_R and first-moment EMA M' = Q_L^T M Q_R into Shampoo's eigenbasis, continually updates an elementwise second-moment EMA V <- beta_2 V + (1 - beta_2)(G' odot G') in the rotated space at every step, applies the Adam preconditioned step N' = M' / (sqrt(V_hat) + epsilon), and rotates the update back to the original parameter space via N = Q_L N' Q_R^T.
Identity Notes & Concept Scope
Distinct from Shampoo in that SOAP establishes an exact equivalence between idealized Shampoo (with exponent 1/2 and trace scaling) and idealized Adafactor in Shampoo's eigenbasis (Claim 1), and replaces rank-1 Adafactor with full coordinatewise Adam second-moment tracking V updated at every step inside the rotated eigenbasis (Q_L, Q_R), preventing stale-preconditioner degradation when Q_L and Q_R are recomputed infrequently every f steps. Distinct from GaLore in that SOAP maintains exponential moving averages of G G^T and G^T G on both sides of a 2D layer and stores first-moment momentum M in the original parameter space before projecting into the updated eigenbasis.
Claims & Grounded Evidence
4 Claims
Mechanism
#c-r589-shampoo-adafactor-eigenbasis-equivalence
Idealized Shampoo with matrix power 1/2 and trace normalization H_hat = (L otimes R) / Trace(L), W_t <- W_{t-1} - eta L^{-1/2} G_t R^{-1/2} / Trace(L)^{-1/2}, is algebraically equivalent to running idealized Adafactor in the eigenbasis (Q_L, Q_R) of L = E_B[G_B G_B^T] and R = E_B[G_B^T G_B] because the row and column second-moment marginals A_i = e_i^T E_B[G'_B odot G'_B] 1_m and C_j = (1_n^T E_B[G'_B odot G'_B])_j in the rotated basis equal the eigenvalues lambda_i(L) and mu_j(R) (Claim 1).
Mechanism
#c-r589-rotated-adam-and-qr-power-iteration
SOAP continually updates the elementwise second-moment matrix V <- beta_2 V + (1 - beta_2)(G' odot G') at every step in the rotated coordinate basis G' = Q_L^T G Q_R while updating the eigenbases Q_L, Q_R only every f steps via a single power-iteration matrix multiplication followed by QR decomposition (Algorithm 3 and Algorithm 4), introducing only one additional hyperparameter (preconditioning frequency f) beyond AdamW and degrading significantly more slowly than Shampoo as f increases up to 100.
Application
#c-r589-llm-pretraining-speedup-and-critical-batch
On 360M and 660M parameter language model pre-training in the large-batch regime, SOAP reduces the number of training iterations by over 40% and wall-clock time by over 35% compared to AdamW (and by approximately 20% in both metrics compared to DistributedShampoo), while improving scaling toward the ideal linear critical batch size.
Limitation
#c-r589-space-and-per-step-projection-overhead
Full two-sided SOAP requires 2m^2 + 2n^2 + 3mn memory and m^3 + n^3 + 2m^2 n + 2mn^2 per-step matrix-multiplication overhead for an m x n layer (more per-step projection work than Shampoo's m^3 + n^3 + m^2 n + mn^2), requiring identity-basis fallback on huge vocabulary embedding/output dimensions or one-sided/factorized approximations to reduce memory below AdamW.
⚠️ Curated Limitations & Negative Results
- Incurs 2m^2 + 2n^2 + 3mn state memory and m^3 + n^3 + 2m^2 n + 2mn^2 per-step matrix-multiplication overhead on 2D layers unless one-sided or factorized approximations are used.
- Cannot maintain full eigenbasis rotation matrices on layers with very large vocabulary dimensions (such as LLM embedding and output projection layers), where the rotation matrix for the large dimension must be fixed to the identity.
Artificial Intelligence Ontology (AIO) Grounding
Mapped AIO Classes
| AIO Target Class | Mapping Relation | Mapping Rationale |
|---|---|---|
AIO:ArtificialNeuralNetwork
Artificial Neural Network
|
related_match | SOAP is a Kronecker-eigenbasis preconditioned Adam optimizer for training artificial neural networks. |
Evidence Sources & Primary Literature
3 Sources| Source ID | Publication / Document | Type / Access | Reference & Checksum |
|---|---|---|---|
r589-primary-pdf
|
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
|
primary paper
full_text
|
ID: https://dblp.org/rec/conf/iclr/0001MZSBJK25
sha256:d6f33bb71eab…
|
r589-arxiv-abs
|
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
|
primary paper
abstract_only
|
ID: https://arxiv.org/abs/2409.11321
sha256:6b90ff5efde8…
|
r589-metadata-cr
|
Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade
|
primary paper
abstract_only
|
ID: https://dblp.org/rec/conf/iclr/0001MZSBJK25
sha256:8ccbc2ce42e0…
|
Ontology & Grounding Topology
Open in interactive network graph →Curation Provenance
Curator: Antigravity
Method: ai_assisted
Model: Gemini
Created: 2026-10-19T11:12:00Z
Updated: 2026-10-19T11:12:00Z
Source File:
data/curated/SoapShampooWithAdamInPreconditionerEigenbasis.yaml
· Schema:
curated_concept.yaml
Explore across X-Mech
Related computational & biological Mechs:
Embedding Map
Network Graph
ProteinTraitsMech (ESM-2 pLM)
DUFMech (Domain pLM)
TraitMech
CultureMech
All 14 X-Mech Repositories