Research map / Learning theory and algorithms / Matrix and tensor methods
A New Perspective on Shampoo's Preconditioner
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.
Shows that the square of Shampoo's Kronecker preconditioner equals one step of power iteration for the optimal Kronecker product approximation of the Adagrad or Gauss-Newton matrix, and empirically studies batch-gradient and empirical-Fisher approximations.
Problem
Shampoo's Kronecker-product preconditioner is widely used but its theoretical foundations, and how close its approximation is to the optimal Kronecker product approximation of the Adagrad or Gauss-Newton matrix, are not well understood.
Approach
The square of Shampoo's approximation equals a single power-iteration step (starting from identity) for the optimal Frobenius-norm Kronecker product approximation of H; the authors show empirically this is close to optimal, and that batch-gradient and empirical-Fisher tricks have limited effect on Hessian approximation quality at batch size 1.
Key results
- One step of power iteration (Shampoo squared) closely tracks the optimal Kronecker approximation in cosine similarity across MNIST logistic regression, CIFAR-5M ResNet18, and ImageNet ConvNeXt-T, while the original Shampoo upper bound is significantly worse.
- For binomial logistic regression the squared Shampoo GN estimate exactly matches H_GN since the Kronecker-structured H is rank-1.
- Using empirical Fisher (real labels) at batch size 1 gives approximation quality close to the optimal Kronecker approximation, but quality degrades at large batch sizes (e.g., 256).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the square of Shampoo’s approximation instead of the Shampoo approximation | Shampoo · MNIST, CIFAR-5M, ImageNet; Gauss–Newton and Adagrad preconditioners | cosine similarity to the true matrix | Helped |
| Use the square of Shampoo’s approximation instead of the K-FAC approximation | Hessian approximation · Datasets and architectures in Appendix experiments | cosine similarity to the true Gauss–Newton matrix | Helped |
| Average sampled-label gradients across the batch (ablation) | Shampoo Hessian approximation · Gauss–Newton matrix | expected approximation | No clear effect |
| Replace sampled labels with real labels at batch size 1 (ablation) | Shampoo Hessian approximation · Gauss–Newton matrix | cosine similarity to the optimal Kronecker approximation | Helped |
| Increase batch size while using real labels (ablation) | Shampoo Hessian approximation · Gauss–Newton matrix | approximation quality | Hurt |
| Use the squared Shampoo estimate under exact Kronecker structure | Shampoo Hessian approximation · Binomial logistic regression | cosine similarity to the Gauss–Newton matrix | Helped |
Surprising findings
- Using real labels (empirical Fisher) at batch size 1 tracks the optimal Kronecker approximation well throughout training, even though it differs from the Fisher with sampled labels.
- ViT experiments show a relatively larger gap between Shampoo squared and the optimal Kronecker approximation than other architectures.
Limitations
- The gap between Shampoo squared and the optimal Kronecker approximation depends on problem structure and is larger for ViT.
- Conditions under which H is close to a Kronecker product remain an open question.
- Results are empirical on a limited set of datasets and architectures, with theory mainly for the direction of the approximation.
Experimental setup
Datasets: MNIST CIFAR-5M ImageNet
Benchmarks: MNIST CIFAR-5M ImageNet
Metrics: cosine similarity between approximation and true GN Hessian
Baselines: Shampoo K-FAC Optimal Kronecker approximation
Builds on: Shampoo Adagrad K-FAC
Closest papers by meaning
- Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner (NeurIPS 2025 · similarity 0.92)
- Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-Sided and Two-Sided Preconditioning (ICML 2026 · similarity 0.91)
- Shampoo: Preconditioned Stochastic Tensor Optimization (ICML 2018 · similarity 0.91)
- Structured Preconditioners in Adaptive Optimization: A Unified Analysis (ICML 2025 · similarity 0.88)
- 4-bit Shampoo for Memory-Efficient Network Training (NeurIPS 2024 · similarity 0.86)
- The Expected Loss of Preconditioned Langevin Dynamics Reveals the Hessian Rank (AAAI 2024 · similarity 0.83)
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.