mlsanityData to 7 Oct 2026

Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner

arXiv 2506.03595 ↗First posted 4 Jun 2025cs.LG cs.AI stat.MLcited by 21

Not on the map. Why?

Abstract (arXiv)

The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks. Despite its success, Shampoo relies heavily on several heuristics such as learning rate grafting and stale preconditioning to achieve performance at-scale. These heuristics increase algorithmic complexity, necessitate further hyperparameter tuning, and lack theoretical justification. This paper investigates these heuristics from the angle of Frobenius norm approximation to full-matrix Adam and decouples the preconditioner's eigenvalues and eigenbasis updates. We show that grafting from Adam mitigates the staleness and mis-scaling of the preconditioner's eigenvalues and how correcting the eigenvalues directly eliminates the need for learning rate grafting. To manage the error induced by infrequent eigenbasis computations, we propose an adaptive criterion for determining the eigenbasis computation frequency motivated by terminating a warm-started QR algorithm. This criterion decouples the update frequency of different preconditioner matrices and enables us to investigate the impact of approximation error on convergence. These practical techniques offer a principled angle towards removing Shampoo's heuristics and developing improved Kronecker-factorization-based training algorithms.

Key facts. Cited by 21 papers in the map (data through 7 Oct 2026). Builds on 47 of its 65 references that are in the map. First posted on arXiv 4 Jun 2025.

Builds on

47 of 65 references.

References · 47

  1. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  2. 2020Open Graph Benchmark: Datasets for Machine Learning on Graphscited by 2,148
  3. 2016Optimization Methods for Large-Scale Machine Learningcited by 1,140
  4. 2015Optimizing Neural Networks with Kronecker-factored Approximate Curvaturecited by 776
  5. 2023ConvNeXt V2: Co-designing and Scaling ConvNets with Masked Autoencoderscited by 703
  6. 2018fastMRI: An Open Dataset and Benchmarks for Accelerated MRIcited by 464
  7. 2024GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projectioncited by 447
  8. 2014New insights and perspectives on the natural gradient methodcited by 435
  9. 2018Shampoo: Preconditioned Stochastic Tensor Optimizationcited by 413
  10. 2025Muon is Scalable for LLM Trainingcited by 381
  11. 2019Limitations of the Empirical Fisher Approximation for Natural Gradient Descentcited by 213
  12. 2024SOAP: Improving and Stabilizing Shampoo using Adamcited by 179
35 more
  1. 2025Training Deep Learning Models with Norm-Constrained LMOscited by 177
  2. 2024Old Optimizer, New Norm: An Anthologycited by 173
  3. 2018Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasiscited by 143
  4. 2016A Kronecker-factored approximate Fisher matrix for convolution layerscited by 138
  5. 2024Adam-mini: Use Fewer Learning Rates To Gain Morecited by 134
  6. 2018Rotate your Networks: Better Weight Consolidation and Less Catastrophic Forgettingcited by 127
  7. 2022Better plain ViT baselines for ImageNet-1kcited by 97
  8. 2023A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scalecited by 69
  9. 2020Practical Quasi-Newton Methods for Training Deep Neural Networkscited by 67
  10. 2025Gradient Alignment in Physics-informed Neural Networks: A Second-Order Optimization Perspectivecited by 55
  11. 2025ASGO: Adaptive Structured Gradient Optimizationcited by 55
  12. 2024A New Perspective on Shampoo's Preconditionercited by 54
  13. 2023Kronecker-Factored Approximate Curvature for Modern Neural Network Architecturescited by 49
  14. 2015Preconditioned Stochastic Gradient Descentcited by 48
  15. 2023Benchmarking Neural Network Training Algorithmscited by 46
  16. 2018A Progressive Batching L-BFGS Method for Machine Learningcited by 45
  17. 2020Disentangling Adaptive Gradient Methods from Learning Ratescited by 39
  18. 2024LDAdam: Adaptive Optimization from Low-Dimensional Gradient Statisticscited by 35
  19. 2025Accelerating Neural Network Training: An Analysis of the AlgoPerf Competitioncited by 30
  20. 2025Structured Preconditioners in Adaptive Optimization: A Unified Analysiscited by 30
  21. 2022Gradient Descent on Neurons and its Link to Approximate Second-Order Optimizationcited by 29
  22. 2024Can We Remove the Square-Root in Adaptive Gradient Methods? A Second-Order Perspectivecited by 25
  23. 2021Tensor Normal Training for Deep Learning Modelscited by 24
  24. 2022PipeFisher: Efficient Training of Large Language Models Using Pipelining and Fisher Information Matricescited by 24
  25. 2021KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networkscited by 20
  26. 2023ASDL: A Unified Interface for Gradient Preconditioning in PyTorchcited by 20
  27. 2023Sketchy: Memory-efficient Adaptive Regularization with Frequent Directionscited by 16
  28. 2024Curvature-Informed SGD via General Purpose Lie-Group Preconditionerscited by 16
  29. 20244-bit Shampoo for Memory-Efficient Network Trainingcited by 16
  30. 2025GaLore 2: Large-Scale LLM Pre-Training by Gradient Low-Rank Projectioncited by 13
  31. 2025On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learningcited by 12
  32. 2020A Trace-restricted Kronecker-Factored Approximation to Natural Gradientcited by 11
  33. 2024Understanding and Minimising Outlier Features in Neural Network Trainingcited by 10
  34. 2024Understanding Adam Requires Better Rotation Dependent Assumptionscited by 10
  35. 2015adaQN: An Adaptive Quasi-Newton Algorithm for Training RNNscited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 3; 2026: 18.

21 papers.

Sort:

  1. 2026Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioningcited by 15
  2. 2025Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimizationcited by 13
  3. 2025SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradientscited by 6
  4. 2025Smoothing DiLoCo with Primal Averaging for Faster Training of LLMscited by 5
  5. 2026Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectorycited by 5
  6. 2026Mitigating Staleness in Asynchronous Pipeline Parallelism via Basis Rotationcited by 2
  7. 2026Spectral Scaling Laws of Muoncited by 2
  8. 2026Leveraging Second-Order Curvature for Efficient Learned Image Compression: Theory and Empirical Evidencecited by 1
  9. 2026DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solverscited by 1
  10. 2026Demystifying Manifold Constraints in LLM Pre-trainingcited by 1
  11. 2026Pro-KLShampoo: Projected KL-Shampoo with Whitening Recovered by Orthogonalizationcited by 1
  12. 2026SOAP-Bubbles: Structured Weight Uncertainty for Neural Networkscited by 1
9 more
  1. 2026SOAP, Muon, and Beyond: Pushing LLM Pretraining Scalescited by 1
  2. 2026Second-Order Muon Done Right: A Principled Marriage of Spectral Geometry and Curvaturecited by 1
  3. 2026Convergence Rate Analysis of the AdamW-style Shampoo: Unifying One-Sided and Two-Sided Preconditioningcited by 0
  4. 2026Adaptive Batch Sizes Using Non-Euclidean Gradient Noise Scales for Stochastic Sign and Spectral Descentcited by 0
  5. 2026Reparametrizing Shampoo and SOAP for Subspace Basis Updates and BFloat16 Storagecited by 0
  6. 2026Exploiting weight-space symmetries for approximating curvaturecited by 0
  7. 2026Rethinking Bregman Divergences in Kronecker-Factored Optimizerscited by 0
  8. 2026FOAM: Frequency and Operator Error-Based Adaptive Damping Method for Reducing Staleness-Oriented Error for Shampoocited by 0
  9. 2026SS-ESOAP: Self-Scaled Adaptive Preconditioning for Physics-Informed Learningcited by 0