Research map / Learning theory and algorithms / Matrix and tensor methods
Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Decomposes Shampoo's preconditioner into eigenvalues and eigenbasis, shows Adam grafting compensates for stale/mis-scaled eigenvalues, and proposes an adaptive eigenbasis update frequency based on a warm-started QR termination criterion.
Problem
Shampoo relies on heuristics (learning rate grafting from Adam and stale preconditioner updates every 100 steps) that increase complexity, need extra tuning, and lack theoretical justification.
Approach
Views Shampoo through a Frobenius-norm approximation to full-matrix Adam, shows grafting compensates for mis-scaled and stale eigenvalues, and removes the need for grafting by eigenvalue correction (EShampoo); an adaptive criterion from warm-started QR decides when to recompute the eigenbasis, decoupling update frequencies across Kronecker factors.
Key results
- EShampoo with fixed F=100 matches or outperforms Shampoo with Adam grafting in steps and wall-clock time on AlgoPerf workloads (e.g., FastMRI 2536 vs 4301 steps).
- Adaptive criterion with tau=0.1, F=50 reduces wall-clock time by about 20% on Imagewoof ViT compared to updating every 100 steps.
- On Imagewoof ViT, freezing the eigenbasis after the first iteration still outperforms AdamW (final loss about 0.12 vs 0.3).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Adam grafting with EShampoo eigenvalue correction at F=100 | Shampoo · FastMRI | steps to target: 4301 ± 109 -> 2536 ± 66 | Helped |
| Use adaptive eigenbasis frequency (τ=0.1, F=50) instead of Adam grafting (F=100) | Shampoo · FastMRI | time [min]: 13.96 ± 0.44 -> 10.81 ± 0.72 | Helped |
| Replace Adam grafting with EShampoo eigenvalue correction at F=100 | Shampoo · ImageNet ViT | time [min]: 894.27 -> 894.85 | No clear effect |
| Use adaptive eigenbasis frequency (τ=0.1, F=50) instead of fixed F=100 | EShampoo · ImageNet ViT | time [min]: 894.85 -> 935.89 | Hurt |
| Replace Adam grafting with EShampoo eigenvalue correction at F=100 | Shampoo · OGBG | steps to target: 12574 ± 708 -> 8320 ± 1203 | Helped |
| Use adaptive eigenbasis frequency (τ=0.1, F=50) instead of fixed F=100 | EShampoo · OGBG | time [min]: 33.02 ± 4.05 -> 27.55 ± 3.49 | Helped |
| Use adaptive criterion (τ=0.1, F=50) rather than updating every 100 iterations | EShampoo · Imagewoof ViT | wall-clock time: 20% | Helped |
| Use stricter eigenbasis accuracy early (τ=0.01 for first 10%, then τ=0.8) rather than τ=0.8 throughout (ablation) | EShampoo · Imagewoof ViT | final loss: ≈ 0.04 | Helped |
| Freeze the eigenbasis after the first iteration (τ=0.99) (ablation) | EShampoo · Imagewoof ViT | final loss: ≈ 0.12 vs. 0.3 | Helped |
| Use Adam for 1D parameters while retaining EShampoo for 2D parameters (ablation) | EShampoo · Imagewoof ViT | final loss | No clear effect |
| Use Adam for 2D parameters and EShampoo for 1D parameters (ablation) | EShampoo · Imagewoof ViT | final loss | Hurt |
| Apply EShampoo to 1D as well as 2D parameters (ablation) | EShampoo · Imagewoof ViT | runtime: 2.9× | Hurt |
Surprising findings
- Freezing the eigenbasis after the first iteration still significantly outperforms AdamW on Imagewoof ViT.
- Using Adam for 1D parameters does not affect Shampoo's convergence or final loss.
Limitations
- Experiments focus on vision and small-scale workloads; scaling to large language models left open.
- Naive eigenvalue correction is not optimal in Frobenius norm and the basis-aware correction is only partially addressed.
- Adaptive QR with small iteration caps can be slower in wall-clock time than torch.linalg.eigh.
Experimental setup
Datasets: Imagewoof FastMRI ImageNet OGBG C4
Benchmarks: Imagewoof AlgoPerf FastMRI AlgoPerf ImageNet ViT AlgoPerf OGBG
Metrics: final training loss steps to target wall-clock time
Baselines: AdamW Shampoo with Adam grafting SOAP Shampoo^2 with S^-1 scaling
Builds on: Shampoo SOAP AlgoPerf K-FAC AdaGrad
Closest papers by meaning
- A New Perspective on Shampoo's Preconditioner (ICLR 2025 · similarity 0.92)
- 4-bit Shampoo for Memory-Efficient Network Training (NeurIPS 2024 · similarity 0.88)
- A Stable Whitening Optimizer for Efficient Neural Network Training (NeurIPS 2025 · similarity 0.88)
- Structured Preconditioners in Adaptive Optimization: A Unified Analysis (ICML 2025 · similarity 0.87)
- On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning (ICML 2025 · similarity 0.85)
- Accelerating neural network training: An analysis of the AlgoPerf competition (ICLR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.