# KrADagrad: Kronecker Approximation-Domination Gradient Preconditioned Stochastic Optimization

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 17 of its 47 references that are in the map.
- First posted on arXiv 30 May 2023.

## Builds on (most cited first, 17 of 17)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Gradient Episodic Memory for Continual Learning](https://mlsanity.com/paper/1706.08840/) (2017)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [An Empirical Investigation of Catastrophic Forgetting in Gradient-Based Neural Networks](https://mlsanity.com/paper/1312.6211/) (2013)
- [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)
- [ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning](https://mlsanity.com/paper/2006.00719/) (2020)
- [Avalanche: an End-to-End Library for Continual Learning](https://mlsanity.com/paper/2104.00405/) (2021)
- [Practical Quasi-Newton Methods for Training Deep Neural Networks](https://mlsanity.com/paper/2006.08877/) (2020)
- [Reproducibility Report: La-MAML: Look-ahead Meta Learning for Continual Learning](https://mlsanity.com/paper/2102.05824/) (2021)
- [On the Factory Floor: ML Engineering for Industrial-Scale Ads Recommendation Models](https://mlsanity.com/paper/2209.05310/) (2022)
- [Why Approximate Matrix Square Root Outperforms Accurate SVD in Global Covariance Pooling?](https://mlsanity.com/paper/2105.02498/) (2021)
- [Tensor Normal Training for Deep Learning Models](https://mlsanity.com/paper/2106.02925/) (2021)
- [Continual Learning Through Synaptic Intelligence](https://mlsanity.com/paper/1703.04200/) (2017)

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

- [SWAN: SGD with Normalization and Whitening Enables Stateless LLM Training](https://mlsanity.com/paper/2412.13148/) (2024)
- [ShamAN-Q: Shampoo Augmented NanoQuant for Sub-1-bit LLM Weights](https://mlsanity.com/paper/2609.38521/) (2026)
- [Minimax Optimality of Shampoo for Nonstationary Stochastic Optimization](https://mlsanity.com/paper/2610.16923/) (2026)

## Abstract (arXiv)

Second order stochastic optimizers allow parameter update step size and direction to adapt to loss curvature, but have traditionally required too much memory and compute for deep learning. Recently, Shampoo [Gupta et al., 2018] introduced a Kronecker factored preconditioner to reduce these requirements: it is used for large deep models [Anil et al., 2020] and in production [Anil et al., 2022]. However, it takes inverse matrix roots of ill-conditioned matrices. This requires 64-bit precision, imposing strong hardware constraints. In this paper, we propose a novel factorization, Kronecker Approximation-Domination (KrAD). Using KrAD, we update a matrix that directly approximates the inverse empirical Fisher matrix (like full matrix AdaGrad), avoiding inversion and hence 64-bit precision. We then propose KrADagrad^star, with similar computational costs to Shampoo and the same regret. Synthetic ill-conditioned experiments show improved performance over Shampoo for 32-bit precision, while for several real datasets we have comparable or better generalization.
