Research map / Optimization / Optimizers
A Stable Whitening Optimizer for Efficient Neural Network Training
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.
Proposes SPlus, a stabilized Shampoo-family optimizer using instant-sign normalization in a cached historical eigenbasis, shape-aware scaling for learning-rate transfer, and iterate averaging, evaluated on 160M Transformers across LLM, ViT, and DiT objectives.
Problem
Shampoo-style whitening optimizers diverge when matrix inverses are cached for long intervals, have poor learning-rate transfer across widths, and suffer parameter noise at high learning rates, limiting their practical speed versus Adam for Transformer training.
Approach
Replace the cached inverse-square-root preconditioner with a bounded sign-based update inside the historical eigenbasis (instant-sign normalization), add a symmetric 2/(m+n) per-layer scale for width transfer, and apply exponential iterate averaging to the live parameters.
Key results
- SPlus reaches Adam's validation performance in about 44-58% of gradient steps and 62-83% of wallclock time on average across objectives.
- Naive Shampoo diverges when the inverse is cached beyond 25-100 steps, while SPlus stays stable at cache intervals up to 100+ steps (e.g., Shampoo Div. at frequency 100 vs SPlus steps-to-Adam 0.77).
- In the iterate-averaging ablation on LLM-Init, SPlus achieves steps-to-Adam 0.487 versus 0.729 without averaging and 0.712 for SOAP.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use instant-sign normalization instead of historical-magnitude inverse scaling | Shampoo · LLM training with cached matrix inverses | divergence frequency: > 50% of the trials diverge; SPlus eliminates divergence across the board | Helped |
| Cache eigendecomposition results for 100 steps rather than recomputing more frequently (ablation) | SPlus · LLM-Init | time-to-Adam: 3.34 -> 0.74 | Helped |
| Cache eigendecomposition results for 100 steps rather than recomputing every 5 steps (ablation) | SPlus · LLM-Init | steps-to-Adam: 0.71 -> 0.77 | Hurt |
| Add iterate averaging (ablation) | SPlus · LLM-Init | steps-to-Adam: 0.729 -> 0.487 | Helped |
| Add iterate averaging (ablation) | SPlus · LLM-10K | steps-to-Adam: 0.662 -> 0.422 | Helped |
| Add iterate averaging (ablation) | SPlus · LLM-50K | steps-to-Adam: 0.685 -> 0.348 | Helped |
| Add iterate averaging (ablation) | Adam · LLM-Init | steps-to-Adam: 1.0 -> 0.677 | Helped |
| Add iterate averaging (ablation) | Adam · LLM-10K | steps-to-Adam: 1.0 -> 0.52 | Helped |
| Add iterate averaging (ablation) | Adam · LLM-50K | steps-to-Adam: 1.0 -> 0.81 | Helped |
| Use symmetric shape-aware scaling instead of spectral scaling (ablation) | SPlus · Transformer training across layer shapes | final performance | Helped |
| Use symmetric shape-aware scaling rather than no scaling (ablation) | SPlus · Transformer training across network widths | learning-rate transfer | Helped |
| Use SPlus rather than SOAP | 160M Transformer optimizer comparison · LLM, ViT, and DiT checkpoints | steps-to-Adam | Mixed |
Surprising findings
- Shampoo diverges in over 50% of trials across learning rates and inversion frequencies in the authors' setting.
- Spectral scaling (1/m) was found harmful in Transformer MLP blocks, with the symmetric 2/(m+n) scaling performing better.
Limitations
- Memory overhead: 3nm + 2(n^2+m^2) parameters per dense layer, roughly 60% more than Adam for square matrices, and nontrivial eigendecomposition time.
- Only Transformer architectures evaluated; behavior on non-Transformer architectures unknown.
- Single trial per optimizer/objective/checkpoint/learning rate, so no error bars reported.
Experimental setup
Datasets: OpenWebText ImageNet
Benchmarks: OpenWebText validation loss (LLM) ImageNet validation loss (ViT, DiT)
Metrics: steps-to-Adam wallclock-to-Adam validation loss
Baselines: Adam Schedule-Free Adam Sophia Shampoo SOAP PSGD Muon Naive SGD
Builds on: Shampoo SOAP Adam Schedule-Free Adam Muon
Closest papers by meaning
- Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner (NeurIPS 2025 · similarity 0.88)
- ASGO: Adaptive Structured Gradient Optimization (NeurIPS 2025 · similarity 0.86)
- Adam Reduces a Unique Form of Sharpness: Theoretical Insights Near the Minimizer Manifold (NeurIPS 2025 · similarity 0.86)
- Adam-mini: Use Fewer Learning Rates To Gain More (ICLR 2025 · similarity 0.85)
- AdaFisher: Adaptive Second Order Optimization via Fisher Information (ICLR 2025 · similarity 0.85)
- Addax: Utilizing Zeroth-Order Gradients to Improve Memory Efficiency and Performance of SGD for Fine-Tuning Language Models (ICLR 2025 · similarity 0.85)
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.