How Bregman Divergences Shape Shampoo
cs.LG
Submitted: 2026-10-06
Updated: 2026-10-06
Code: https://github.com/KellerJordan/modded-nanogpt
License: http://creativecommons.org/licenses/by/4.0/
The gist: Understanding the principles behind Shampoo has recently guided the development of more effective neural network optimizers.
Abstract
Understanding the principles behind Shampoo has recently guided the development of more effective neural network optimizers. These methods learn a preconditioner by optimizing the Frobenius or Kullback-Leibler (KL) divergence against the gradient second moment. In this work, we investigate how the choice of divergence shapes preconditioning, which remains unclear and blocks further improvements. To do so, we develop a unified Bregman divergence framework that connects all popular divergences, allowing us to study them jointly. Through empirical spectral analysis of gradient second moments, we examine how divergence choice shapes Kronecker approximation and interacts with finite-sample error in preconditioning. We find that some divergences can better compensate for finite-sample underestimation of the empirical second moment, helping explain the differing behavior of their corresponding Shampoo variants. We further validate this explanation through GPT-2 pretraining experiments. By connecting divergence choice to practical training behavior, we believe our framework provides principled guidance for understanding the foundations of, and further improving, Shampoo.
Sources
- Disentangling Adaptive Gradient Methods from Learning Rates
- Scalable Second Order Optimization for Deep Learning
- Cross-validation based Nonlinear Shrinkage
- Benchmarking Neural Network Training Algorithms
- Clarifying Shampoo: Adapting Spectral Descent to Stochasticity and the Parameter Trajectory
- Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory
- A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks