# A Stable Whitening Optimizer for Efficient Neural Network Training

arXiv: https://arxiv.org/abs/2506.07254 · HTML: https://mlsanity.com/paper/2506.07254/

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Builds on 42 of its 60 references that are in the map.
- First posted on arXiv 8 Jun 2025.

## Builds on (most cited first, 30 of 42)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Adafactor: Adaptive Learning Rates with Sublinear Memory Cost](https://mlsanity.com/paper/1804.04235/) (2018)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [signSGD: Compressed Optimisation for Non-Convex Problems](https://mlsanity.com/paper/1802.04434/) (2018)
- [Latte: Latent Diffusion Transformer for Video Generation](https://mlsanity.com/paper/2401.03048/) (2024)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Shampoo: Preconditioned Stochastic Tensor Optimization](https://mlsanity.com/paper/1802.09568/) (2018)
- [Analyzing and Improving the Training Dynamics of Diffusion Models](https://mlsanity.com/paper/2312.02696/) (2023)
- [Don't Decay the Learning Rate, Increase the Batch Size](https://mlsanity.com/paper/1711.00489/) (2017)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Old Optimizer, New Norm: An Anthology](https://mlsanity.com/paper/2409.20325/) (2024)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)

## Built on by (most cited first, 18 of 18)

- [What Really Matters in Matrix-Whitening Optimizers?](https://mlsanity.com/paper/2510.25000/) (2025)
- [Preconditioned Norms: A Unified Framework for Steepest Descent, Quasi-Newton and Adaptive Methods](https://mlsanity.com/paper/2510.10777/) (2025)
- [Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement](https://mlsanity.com/paper/2602.22681/) (2026)
- [MuonAll: Muon Variant for Efficient Finetuning of Large Language Models](https://mlsanity.com/paper/2511.06086/) (2025)
- [Decoupling Variance and Scale-Invariant Updates in Adaptive Gradient Descent for Unified Vector and Matrix Optimization](https://mlsanity.com/paper/2602.06880/) (2026)
- [Beyond Muon: MUD (MomentUm Decorrelation) for Faster Transformer Training](https://mlsanity.com/paper/2603.17970/) (2026)
- [Evolution of Optimization Methods: Algorithms, Scenarios, and Evaluations](https://mlsanity.com/paper/2604.12968/) (2026)
- [Beyond Adam: SOAP and Muon for Faster, Label-Efficient Training of Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2607.02499/) (2026)
- [Leveraging Second-Order Curvature for Efficient Learned Image Compression: Theory and Empirical Evidence](https://mlsanity.com/paper/2601.20769/) (2026)
- [FOAM: Frequency and Operator Error-Based Adaptive Damping Method for Reducing Staleness-Oriented Error for Shampoo](https://mlsanity.com/paper/2606.02365/) (2026)
- [Convergence Rate Analysis of the AdamW-style Shampoo: Unifying One-Sided and Two-Sided Preconditioning](https://mlsanity.com/paper/2601.07326/) (2026)
- [Discovering Data Manifold Geometry through Geometric Properties](https://mlsanity.com/paper/2602.02611/) (2026)
- [Error whitening: Why Gauss-Newton outperforms Newton](https://mlsanity.com/paper/2605.11316/) (2026)
- [Muse: Representation Geometry of Muon Beyond Normalized Momentum](https://mlsanity.com/paper/2607.14536/) (2026)
- [MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning](https://mlsanity.com/paper/2608.05088/) (2026)
- [Blog: Survey of Optimizers](https://mlsanity.com/paper/2608.28557/) (2026)
- [Minimax Optimality of Shampoo for Nonstationary Stochastic Optimization](https://mlsanity.com/paper/2610.16923/) (2026)
- [How Inner Optimizers Shape Replica Scaling](https://mlsanity.com/paper/2610.19064/) (2026)

## Abstract (arXiv)

In this work, we take an experimentally grounded look at neural network optimization. Building on the Shampoo family of algorithms, we identify and alleviate three key issues, resulting in the proposed SPlus method. First, we find that naive Shampoo is prone to divergence when matrix-inverses are cached for long periods. We introduce an alternate bounded update combining a historical eigenbasis with instantaneous normalization, resulting in across-the-board stability and significantly lower computational requirements. Second, we adapt a shape-aware scaling to enable learning rate transfer across network width. Third, we find that high learning rates result in large parameter noise, and propose a simple iterate-averaging scheme which unblocks faster learning. To properly confirm these findings, we introduce a pointed Transformer training benchmark, considering three objectives (language modelling, image classification, and diffusion modelling) across different stages of training. On average, SPlus is able to reach the validation performance of Adam within 44-58% of the gradient steps and 62-83% of the wallclock time.