PaperlayerView exact revision on arXiv

Source-linked AI summary

SOAP: Improving and Stabilizing Shampoo using Adam

Nikhil Vyas, Depen Morwani, Rosie Zhao, Mujin Kwun, Itai Shapira, David Brandfonbrener, Lucas Janson, Sham Kakade

arXiv:2409.11321v2cs.LGcs.AI

TL;DR

The paper addresses optimizer efficiency by connecting Shampoo with Adafactor and introducing SOAP, which runs AdamW in Shampoo’s eigenbasis. SOAP outperforms AdamW and Shampoo in language-modeling tasks and is more robust to preconditioning-frequency changes than Shampoo.

  • Problem

    Shampoo can outperform Adam but has greater hyperparameter and computational overhead, motivating a simpler efficient optimizer.

  • Method

    SOAP runs AdamW in the eigenbasis provided by Shampoo, based on a formal connection between Shampoo and Adafactor.

  • Results

    SOAP outperforms both AdamW and Shampoo in language-modeling tasks with 360m and 660m parameter models and is more robust to preconditioning-frequency changes than Shampoo.

  • Takeaways & Limitations

    SOAP provides an optimizer that combines Shampoo’s preconditioned space with AdamW updates for language-model pre-training.

  • Takeaways & Limitations

    The study focuses on models relatively small compared with recent LLMs, and its findings at larger scales remain hypothesized rather than established.

Abstract

from arXiv · show

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: $\textbf{S}$hampo$\textbf{O}$ with $\textbf{A}$dam in the $\textbf{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

SOAP is motivated by the efficiency benefits of higher-order preconditioning and connects Shampoo to Adafactor through Shampoo’s eigenbasis. Building on that connection, SOAP runs AdamW in the preconditioned space and is evaluated against AdamW and Shampoo.

  • Motivation: Optimization efficiency is increasingly important as LLM training costs rise, while Shampoo has outperformed Adam in efficiency benchmarks.Shampoo reduced wall-clock time by 28% in Algoperf.
  • Connection: Shampoo with exponent 1/2 is formally equivalent to running Adafactor in the eigenbasis supplied by Shampoo’s preconditioner.This interpretation connects Shampoo to methods that run first-order optimizers in second-order eigenbases.
  • Method: SOAP runs AdamW in the eigenbasis provided by Shampoo’s preconditioner.SOAP stands for ShampoO with Adam in the Preconditioner’s eigenbasis.
  • Results: SOAP outperforms both Shampoo and AdamW on language-model pre-training tasks with 360m and 660m models.The comparison follows extensive hyperparameter tuning of Shampoo.
  • Practical design: SOAP has only one additional hyperparameter relative to AdamW: preconditioning frequency.This reduces the hyperparameter burden relative to Shampoo.
  • Novelty: The authors systematically evaluate an algorithmic variant that had previously been discussed without reported experiments.They identify this systematic evaluation as a contribution.

2 NOTATION AND BACKGROUND

The background introduces matrix and gradient notation, then situates Adam, Adafactor, and Shampoo as progressively structured approximations or preconditioned optimization methods. It also identifies the implementation variants used for Shampoo.

  • Notation: A neural-network layer uses weight matrix W ∈ R^{m×n} and gradient G ∈ R^{m×n}, with batch loss and gradient evaluated at W_t.The batch at time t is denoted B_t.
  • Adagrad: Adagrad maintains a full preconditioner H ∈ R^{mn×mn} for vectorized gradients and weights.The vectorized gradient is g_t = vec(G_t).
  • Adam: Adam is a diagonal approximation of Adagrad that maintains exponential moving averages of gradients and element-wise squared gradients.These averages are denoted M_t and V_t.
  • Adafactor: Adafactor reduces memory usage by replacing Adam’s second-moment estimate with its best rank-1 approximation.The cited version focuses on this memory-efficient modification.
  • Shampoo: Shampoo approximates Adagrad with separate left and right preconditioners L_t ∈ R^{m×m} and R_t ∈ R^{n×n}.The implementation used in this work exposes layerwise grafting and exponent choices as hyperparameters.

3 RELATED WORK

Related work spans second-order optimizers, scalable Shampoo variants, diagonal alternatives, and methods that combine changing preconditioners with diagonal updates. SOAP is positioned among these approaches while retaining a distinct Adam-in-eigenbasis formulation.

  • Related eigenbasis methods: E-KFAC updates a diagonal preconditioner between KFAC inversion steps, but its diagonal preconditioner is not Adam.This distinguishes E-KFAC from running Adam in KFAC’s eigenbasis.
  • Shampoo variants: Anil et al. found that Shampoo’s exponent 1/2 outperformed its original exponent 1/4, while an essentially SOAP-equivalent variant lacked reported experiments.The prior work reported unpublished experiments claiming no improvement over Shampoo.
  • Memory-efficient optimizers: GaLore reduces Adam’s memory footprint using a low-rank gradient-SVD subspace, whereas SOAP maintains exponential moving averages of G G^T and G^T G.The methods differ in how their projection subspaces are determined.
  • Diagonal preconditioning: Diagonal optimizers such as Lion, Sophia, and Adafactor perform comparably to AdamW for LLM pre-training but do not surpass it in cited studies.This motivates exploring non-diagonal preconditioners.
  • Second-order optimization: Second-order optimization research includes Hessian-free methods and methods that estimate the Hessian efficiently.KFAC and Shampoo are recognized examples of Hessian-estimation methods.
  • KFAC: KFAC introduced layerwise Kronecker-factored preconditioning beyond diagonal preconditioners, while later work extended and scaled it.These extensions target different architectures and improved curvature estimates.
  • Shampoo: Shampoo is a layerwise Kronecker-factored second-order method whose distributed implementation won an optimization-efficiency benchmark.The benchmark result highlights its practical utility.

4 ALGORITHM

SOAP formalizes a connection between Shampoo with power 1/2 and Adafactor, then uses that insight to run AdamW in Shampoo’s eigenbasis. Its practical design updates Adam moments every step while refreshing eigenvectors periodically, adding only preconditioning frequency beyond AdamW.

  • Shampoo–Adafactor equivalence: Shampoo with power 1/2 is equivalent to running Adafactor in the eigenbasis supplied by Shampoo’s preconditioner.The equivalence is established for idealized algorithms and also holds with momentum.
  • Practical algorithm: SOAP introduces one additional hyperparameter compared to AdamW: preconditioning frequency.The algorithm also uses learning rate, Adam betas, and epsilon, while practical implementation details include power iteration and special handling for some layers.
  • SOAP: SOAP runs AdamW in the eigenbasis provided by Shampoo’s preconditioner.The gradient and momentum are rotated into the current eigenbasis, Adam moments are updated there, and the update is rotated back.
  • Practical algorithm: SOAP updates second-moment estimates every step while updating Shampoo eigenvectors only at the chosen preconditioning frequency.This preserves inexpensive per-step adaptivity in a slowly changing coordinate basis rather than recomputing eigendecompositions every step.
  • Practical algorithm: Practical Shampoo adaptivity is limited when eigendecompositions are infrequent, although grafting provides a single scalar learning-rate adjustment per layer.SOAP instead uses Adam’s more general diagonal adaptation in the rotated space.
  • Implementation details: For very large dimensions, SOAP fixes the corresponding rotation matrix to identity; fixing both rotations recovers Adam.For one-dimensional layers, the implementation uses standard AdamW to reduce overhead.

5 EXPERIMENTAL METHODOLOGY

The experiments compare AdamW, DistributedShampoo, and SOAP on language-model optimization using tuned hyperparameters, throughput measurements, and scaling-law estimates of efficiency.

  • Hyperparameter tuning: Hyperparameter tuning begins with prior-recommended values, then sweeps learning rate and pairs it with each remaining hyperparameter.The procedure is intended to demonstrate the effect of extensive tuning rather than relying on a single configuration.
  • Throughput measurement: Throughput is measured as tokens processed per second on a single H100 GPU with gradient accumulation for large batches.The methodology accounts for optimizer overhead relative to accumulated gradient steps and possible amortization across layers.
  • Efficiency calculation: Efficiency benefits are estimated by fitting a + bN^-β to SOAP runs using 0.5, 0.625, 0.75, and 0.875 fractions of the training data.The fitted scaling law is used because cosine schedules prevent direct efficiency comparisons from equal-duration runs.

6 LANGUAGE MODELING EXPERIMENTS

Language-model experiments evaluate SOAP against AdamW and Shampoo across model sizes, batch sizes, and eigenvector-update frequencies. SOAP generally achieves better efficiency, retains an advantage in longer runs, and degrades more slowly as preconditioning frequency increases.

  • Efficiency benefits: ≥40% fewer iterations and ≥35% less wall-clock time are observed versus AdamW, while SOAP improves both metrics by approximately 20% versus Shampoo.These results use 360m and 660m models with 2m-token batches and preconditioning frequency 10.
  • Preconditioning frequency: For frequencies from 1 to 100, both SOAP and Shampoo outperform AdamW, while SOAP’s performance degrades significantly more slowly at higher frequencies.At frequency 1, the two methods are close; increasing frequency harms both methods but affects Shampoo more strongly.
  • Critical batch size: SOAP follows ideal linear scaling more closely than AdamW, indicating a higher critical batch size in this setup.The critical batch size marks where doubling batch size no longer halves the required training steps.
  • Critical batch size: SOAP consistently requires fewer steps than AdamW across batch sizes, with larger multiplicative benefits at larger batch sizes.The comparison targets the AdamW loss of 2.842 for a 360m model trained on a Chinchilla-optimal token count.
  • Small-batch comparison: At 256k-token batches, SOAP reduces iterations by 25% versus AdamW and approximately 10% versus Shampoo.The same comparison reports at least 15% wall-clock improvement over AdamW and around 10% over Shampoo, but is described as preliminary.
  • Extended training: SOAP maintains its advantage over AdamW in extended runs using 100 × model-size tokens.The paper also notes that one-sided projection variants can remain comparable to or marginally better than Shampoo with greater computational efficiency.

7 FURTHER EFFICIENCY IMPROVEMENTS

SOAP’s efficiency can be improved through factorized or one-sided variants that trade some performance for lower space and time overhead. Its overhead also depends on preconditioning frequency and persistent matrix operations.

  • Space and time variants: One-sided SOAP slightly reduces performance but remains on par with, or marginally better than, Shampoo while improving computational efficiency.Further work is needed to determine whether these variants can surpass original SOAP’s efficiency.
  • Space usage: Replacing Adam with Adafactor reduces SOAP’s space usage by mn, while combining it with one-sided preconditioning reduces usage to 2 min(m, n)^2 + 2mn.The factorized change causes only a negligible performance loss, whereas one-sided preconditioning causes a larger increase in loss.
  • Space usage: For standard transformer architectures, the combined factorized and one-sided variant uses less space than AdamW.Its space usage is 2 min(m, n)^2 + 2mn, compared with AdamW’s 3mn.
  • Time overhead: SOAP’s per-step overhead exceeds Shampoo’s because it includes projecting and reprojecting gradients on both sides of each layer.For an m × n layer, SOAP adds 2m^2n + 2mn^2 to the m^3 + n^3 terms.
  • Time overhead: Power-iteration eigenvector updates using torch.linalg.qr perform comparably to fresh torch.linalg.eigh decomposition across the preconditioning-frequency spectrum.torch.linalg.qr is described as computationally more efficient than torch.linalg.eigh.
  • Time overhead: Increasing preconditioning frequency lowers SOAP’s overhead, but the overhead approaches a nonzero asymptote from recurring matrix updates and gradient projections.Lower-precision matrix operations could reduce this asymptote.

8 DISCUSSION AND FUTURE WORK

The paper presents SOAP as an optimizer that places AdamW in Shampoo’s preconditioner eigenbasis and reports advantages over AdamW and Shampoo in language modeling. Future work targets lower precision, distributed implementation, and evaluation beyond language modeling.

  • Discussion: SOAP runs AdamW in the eigenbasis provided by Shampoo’s preconditioner.The method is named ShampoO with Adam in the Preconditioner’s eigenbasis.
  • Discussion: SOAP outperforms both AdamW and Shampoo in language modeling tasks and is more robust to preconditioning-frequency changes than Shampoo.The paper identifies lower-precision preconditioners, better distributed implementation, and vision experiments as future directions.

9 DISCUSSION AND LIMITATIONS

The evaluation focuses on relatively small models compared with recent large language models, so whether SOAP generalizes to larger scales remains unvalidated. The paper also identifies broader implementation and domain evaluations as future work.

  • Limitations: The study is relatively small compared with recent LLMs that are two orders of magnitude larger, leaving larger-scale generalization unvalidated.The authors hypothesize that SOAP’s findings may generalize, but explicitly state that this remains to be validated.
  • Future work: Future work will test lower-precision preconditioners, improve distributed implementation, and evaluate SOAP on domains such as vision.These efforts aim to assess performance across different task types.

A EXPERIMENTAL SETUP

The experiments train decoder-only transformer language models on C4 using three model sizes and mixed-precision training. Comparisons use standard AdamW, DistributedShampoo, SOAP, and GaLore under specified optimization and token-budget settings.

  • Data and evaluation: Experiments use C4 tokenized with the T5 tokenizer and report validation loss.The sequence length is 1024.
  • Models: The study trains decoder-only transformers with 210m, 360m, and 660m non-embedding parameters.Most reported results use the 360m and 660m models, while the 210m model supports ablations.
  • Algorithms: The optimizer comparison uses PyTorch AdamW, DistributedShampoo, and self-implemented SOAP and GaLore.The implementations are based on PyTorch AdamW and the official GaLore implementation where applicable.
  • Hyperparameters: Default training uses β1 = 0.95, decoupled weight decay and z-loss coefficients of 1e-4, ϵ = 1e-8, and warmup followed by cosine decay.These settings apply to the reported optimizer setup unless otherwise specified.
  • Hyperparameters: DistributedShampoo uses a default exponent of −1/2.5, with ϵshampoo = 1e-12 and βshampoo = 0.95.These values were based on initial experiments and prior findings favoring an exponent near −1/2.
  • Training budget: Runs generally use a 2m-token batch and approximately 20 times the model size in training tokens.The default schedules correspond to 1600, 3200, and 6400 steps for the 210m, 360m, and 660m models.

A.1 SWEEPING OVER HYPERPARAMETERS

The experiments tune learning rates and optimizer-specific hyperparameters across model sizes, batch sizes, and preconditioning frequencies. Several later sweeps provide little or no benefit beyond learning-rate tuning, while β2 = .99 improves performance slightly in selected settings.

  • Learning-rate and optimizer sweeps: Learning-rate sweeps are performed for AdamW, DistributedShampoo, and SOAP across the tested batch-size and model-size settings.The candidate learning rates range from .1 to 3.16e−4.
  • Learning-rate and optimizer sweeps: For 360m models with 2m batch size, additional DistributedShampoo sweeps over epsilon, βshampoo, and exponents produce no significant improvement below .004.Consequently, only learning rate is swept for the 660m model.
  • Preconditioning-frequency sweeps: Preconditioning-frequency experiments retune learning rates at each frequency while holding other hyperparameters at values obtained at frequency 10.The experiments also vary β1 and β2 for the 360m model.
  • Shorter runs and warmup: Shorter SOAP runs for 360m and 660m models reuse the best standard-run learning rates and other optimal hyperparameters, with separately specified warmup durations.Warmup durations vary by model size and shortened-run fraction.

B GALORE

The appendix evaluates GaLore against AdamW and Shampoo, then reports a hyperparameter search whose best run remains slightly worse than the best Shampoo run. GaLore is therefore not evaluated at larger model sizes.

  • GaLore comparison: GaLore outperforms AdamW on the 210m model but performs worse than Shampoo, so it is not tested at higher model sizes.This is the stated basis for excluding GaLore from larger-model experiments.
  • Hyperparameter search: The hyperparameter search varies learning rate, preconditioning frequency, sidedness, β2, and β1.The tested preconditioning frequencies include 10, 50, and 200.
  • Hyperparameter search: 3.12 final loss is achieved by the best run in the search, compared with 3.10 for the best Shampoo run.The reported comparison favors Shampoo on final loss.
Loading 2409.11321v2…