REVIEW 15 cited by
A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
Signed reviews
read the original abstract
Shampoo is an online and stochastic optimization algorithm belonging to the AdaGrad family of methods for training neural networks. It constructs a block-diagonal preconditioner where each block consists of a coarse Kronecker product approximation to full-matrix AdaGrad for each parameter of the neural network. In this work, we provide a complete description of the algorithm as well as the performance optimizations that our implementation leverages to train deep networks at-scale in PyTorch. Our implementation enables fast multi-GPU distributed data-parallel training by distributing the memory and computation associated with blocks of each parameter via PyTorch's DTensor data structure and performing an AllGather primitive on the computed search directions at each iteration. This major performance enhancement enables us to achieve at most a 10% performance reduction in per-step wall-clock time compared against standard diagonal-scaling-based adaptive gradient methods. We validate our implementation by performing an ablation study on training ImageNet ResNet50, demonstrating Shampoo's superiority over standard training recipes with minimal hyperparameter tuning.
Forward citations
Cited by 15 Pith papers
-
On The Concurrence of Layer-wise Preconditioning Methods and Provable Feature Learning
For two feature-learning models with anisotropic inputs, KFAC-style layer-wise preconditioning provably recovers features better than SGD and matches ridge regression in the single-index case.
-
Muon Meets Mamba: Spectral Optimization for State Space Models
On Mamba-2 130M, Muon on the output projection alone gives lower validation loss and better token efficiency than Muon on the input projection or on both, across corpora and budgets.
-
OneShot: Index-in-Ranking with Neural Scoring for Large-Scale Retrieval
OneShot trains hierarchical item codebooks jointly with the ranking loss, enabling nonlinear neural scoring in billion-scale retrieval and reporting +20% recall, 10x fewer dense-ranked items, and live Instagram gains.
-
PoLoRA: A Preconditioned Orthogonalized LoRA Optimizer
A preconditioned, product-aware spectral optimizer for LoRA reaches Adam's final loss in 1.2–1.7x fewer steps on instruction-tuning, with ~3% per-step overhead.
-
Muse: Representation Geometry of Muon Beyond Normalized Momentum
The matrix shape given to Muon's orthonormalization step is a genuine optimizer axis: square-ish reshapes match native Muon, skinnier reshapes interpolate toward normalized SGD with momentum.
-
MatrixFSDP: communication-free matrix optimizers under ZeRO-3 parameter sharding
MatrixFSDP places whole 2D weight matrices on single ZeRO-3 owner ranks so matrix optimizers run locally without optimizer-step collectives, preserving ZeRO-3 memory while achieving up to 54.6x optimizer-step speedup ...
-
Understanding and Improving Shampoo and SOAP via Kullback-Leibler Minimization
A KL-divergence view of Shampoo and SOAP yields KL-Shampoo and KL-SOAP, with KL-Shampoo outperforming both baselines in LLM pretraining without Adam grafting.
-
DHO$_2$: Accelerating Distributed Hybrid Order Optimization via Model Parallelism and ADMM
DHO2 distributes FOSI's Hessian eigenvector computation with model parallelism and an ADMM-like update, claiming lower per-GPU memory and faster time-to-solution in distributed DNN training.
-
Spectral-factorized Positive-definite Curvature Learning for NN Training
The paper derives a Riemannian update rule for the spectral factors of a positive-definite preconditioner, making arbitrary matrix roots fast and numerically stable for low-precision NN training.
-
SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training
SWAN, a stateless optimizer combining gradient normalization and whitening, matches or beats Adam on LLaMA pretraining through 1.3B parameters with roughly half the memory and reported 2x token efficiency.
-
Memory-Efficient 4-bit Preconditioned Stochastic Optimization
A 4-bit Shampoo optimizer that quantizes Cholesky factors and adds error feedback matches 32-bit Shampoo's accuracy at a fraction of the memory.
-
MARS: Unleashing the Power of Variance Reduction for Training Large Models
MARS combines STORM-style variance reduction with preconditioned updates and reports improved token efficiency over AdamW on GPT-2 pretraining, though the analyzed and evaluated variants differ.
-
Second-Order Muon Done Right: A Principled Marriage of Spectral Geometry and Curvature
GO-MUON combines K-FAC-style second-moment factors with Muon's polar map in a matched spectral oracle, and a four-step refresh schedule trades factor freshness for compute without claiming denoising.
-
Learning by solving differential equations
Runge-Kutta optimizers adapted with momentum, preconditioning, or adaptive learning rates can close the large-batch generalization gap and match Adam on small MLP workloads.
-
Materials Learning Algorithms (MALA): Scalable Machine Learning for Electronic Structure Calculations in Large-Scale Atomistic Simulations
MALA predicts electron densities and energies from local atomic environments using trained neural networks, reaching system sizes beyond standard DFT.
Discussion (0). Continue with ORCID to comment.