A Statistical Viewpoint on Adam, Shampoo and Muon
Abstract
We provide new statistical insights into designing stochastic optimization methods. By modeling the gradient via a normal distribution, we argue that the stochastic gradient should be whitened using its covariance. We give an online method for estimating the mean and variance of the gradients based on regularized likelihood or, equivalently, through a maximum a posteriori estimate. We then show how to use structured approximations to the covariance within this same framework, which includes diagonal and Kronecker-factored approximations. The diagonal approximation results in a method that is very closely related to Adam and performs as well in practice. The Kronecker approximation results in the new MN method, which is a variant of Shampoo and KL-Shampoo. We also recover Muon by applying this same statistical framework to the gradient direction rather than the full gradient. By combining this with a whitening step, we arrive at MN-Muon, a novel preconditioned version of Muon. Our experiments on GPT-2-type models at Chinchilla scale indicate that MND, the diagonal implementation of MN, outperforms both the Distributed Shampoo and the KL-Shampoo implementations in terms of stability, final validation loss and time taken. MND-Muon, the diagonal implementation of MN-Muon, improves on Muon in 21 of 24 paired GPT-small runs and achieves a lower best validation loss on GPT-medium, where it remains stable over the six tested learning rates. However, we found that NorMuon was comparable to MND-Muon and achieved a slightly lower validation loss on GPT-medium and at GPT-small batch 512 on FineWeb.
Related
- From Geometry to Generalization: Why Row Normalization Can Beat Adam and Muon
- First-Order Efficiency for Probabilistic Value Estimation via A Statistical Viewpoint
- Minimax Optimality of Shampoo for Nonstationary Stochastic Optimization
- ShamAN-Q: Shampoo Augmented NanoQuant for Sub-1-bit LLM Weights
- Early Memory Selection for Balanced Adam
BibTeX
@misc{parshakova2026statistical,
title = {{A Statistical Viewpoint on Adam, Shampoo and Muon}},
author = {Tetiana Parshakova and Robert M. Gower and Paulo Orenstein},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2610.24839},
eprint = {2610.24839},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2610.24839},
}
RIS
TY - UNPB
TI - A Statistical Viewpoint on Adam, Shampoo and Muon
AU - Tetiana Parshakova
AU - Robert M. Gower
AU - Paulo Orenstein
PY - 2026
DA - 2026/10/17
PB - arXiv
DO - 10.48550/arxiv.2610.24839
UR - https://arxiv.org/abs/2610.24839
AB - We provide new statistical insights into designing stochastic optimization methods. By modeling the gradient via a normal distribution, we argue that the stochastic gradient should be whitened using its covariance. We give an online method for estimating the mean and variance of the gradients based on regularized likelihood or, equivalently, through a maximum a posteriori estimate. We then show how to use structured approximations to the covariance within this same framework, which includes diagonal and Kronecker-factored approximations. The diagonal approximation results in a method that is very closely related to Adam and performs as well in practice. The Kronecker approximation results in the new MN method, which is a variant of Shampoo and KL-Shampoo. We also recover Muon by applying this same statistical framework to the gradient direction rather than the full gradient. By combining this with a whitening step, we arrive at MN-Muon, a novel preconditioned version of Muon. Our experiments on GPT-2-type models at Chinchilla scale indicate that MND, the diagonal implementation of MN, outperforms both the Distributed Shampoo and the KL-Shampoo implementations in terms of stability, final validation loss and time taken. MND-Muon, the diagonal implementation of MN-Muon, improves on Muon in 21 of 24 paired GPT-small runs and achieves a lower best validation loss on GPT-medium, where it remains stable over the six tested learning rates. However, we found that NorMuon was comparable to MND-Muon and achieved a slightly lower validation loss on GPT-medium and at GPT-small batch 512 on FineWeb.
ER -