A Statistical Viewpoint on Adam, Shampoo and Muon

Authors: Tetiana Parshakova, Robert M. Gower, Paulo Orenstein Date: October 20, 2026

---

Summary (Overview)

• This paper presents a unified statistical framework for understanding and designing stochastic optimization methods, based on modeling gradients as samples from a normal distribution and applying whitening using estimated covariance.
• The authors derive a Bayesian/optimization framework for estimating gradient mean and covariance, showing that both approaches converge to the same KL projection problem with closed-form updates.
• Using diagonal covariance approximations, the framework recovers a method closely related to Adam (called Centered Adam) that performs comparably in practice while admitting a simple whitening interpretation.
• Using Kronecker-factored covariance approximations, the framework yields the new MN (Matrix Normal) method, a variant of Shampoo/KL-Shampoo, whose diagonal implementation (MND) outperforms both Distributed Shampoo and KL-Shampoo at Chinchilla scale.
• Extending the framework to model the direction of gradients yields a statistical interpretation of Muon and a novel preconditioned variant called MND-Muon, which improves upon Muon in 21 of 24 paired GPT-small runs and shows enhanced stability on GPT-medium.

---

Introduction and Theoretical Foundation

Background and Motivation

Stochastic gradients are noisy estimates of the full-batch gradient, with different components exhibiting different variances. Taking a large step in a high-variance direction can lead the optimization astray. The paper's central thesis is that whitening the gradient using its covariance improves optimization:

$$\boldsymbol{w}_{t+1} = \boldsymbol{w}_t - \eta \boldsymbol{\Sigma}_t^{-1/2} \boldsymbol{g}_t \tag{1}$$

where $\boldsymbol{g}_t \sim \mathcal{N}(\boldsymbol{m}_t, \boldsymbol{\Sigma}_t)$. This "whitened SGD" rescales steps so that high-variance directions are treated more cautiously.

The Whitened Momentum Framework

In practice, estimates $\hat{\boldsymbol{\Sigma}}_t$ and $\hat{\boldsymbol{m}}_t$ are used:

$$\boldsymbol{w}_{t+1} = \boldsymbol{w}_t - \eta \hat{\boldsymbol{\Sigma}}_t^{-1/2} \hat{\boldsymbol{m}}_t \tag{2}$$

Methods fitting this format are called whitened momentum methods.

Dual Viewpoints: Optimization and Bayesian

The paper shows that estimating the mean and covariance can be approached two equivalent ways:

Optimization viewpoint: Minimize negative log-likelihood plus a KL-divergence regularizer:

$$(\hat{\boldsymbol{m}}_t, \hat{\boldsymbol{\Sigma}}_t^{\prime}) = \underset{\boldsymbol{m}, \boldsymbol{\Sigma} \succ 0}{\mathrm{argmin}} \Big\{-\log \mathcal{N}(\boldsymbol{g}_t \mid \boldsymbol{m}, \boldsymbol{\Sigma}) + \frac{\beta}{1-\beta} \mathrm{KL}\Big(\mathcal{N}(\hat{\boldsymbol{m}}_{t-1}, \hat{\boldsymbol{\Sigma}}_{t-1}) \| \mathcal{N}(\boldsymbol{m}, \boldsymbol{\Sigma})\Big)\Big\} \tag{3}$$

Bayesian viewpoint: Use a Normal-Inverse-Wishart conjugate prior:

$$p(\boldsymbol{m}, \boldsymbol{\Sigma} \mid \hat{\boldsymbol{m}}_{t-1}, \hat{\boldsymbol{\Sigma}}_{t-1}) = \mathcal{NIW}\Big(\hat{\boldsymbol{m}}_{t-1}, \frac{\beta}{1-\beta}, \frac{\beta}{1-\beta}\hat{\boldsymbol{\Sigma}}_{t-1}, \frac{\beta}{1-\beta} - (d+2)\Big) \tag{4}$$

Proposition 2.1 establishes that both approaches are equivalent to a KL projection problem with closed-form solutions:

$$\hat{\boldsymbol{m}}_t = \beta \hat{\boldsymbol{m}}_{t-1} + (1-\beta)\boldsymbol{g}_t \tag{7}$$ $$\hat{\boldsymbol{\Sigma}}_t^{\prime} = \beta \hat{\boldsymbol{\Sigma}}_{t-1} + \beta(1-\beta)(\boldsymbol{g}_t - \hat{\boldsymbol{m}}_{t-1})(\boldsymbol{g}_t - \hat{\boldsymbol{m}}_{t-1})^\top \tag{8}$$

---

Methodology

Structured Covariance Approximations

Storing a full $d \times d$ covariance matrix is infeasible for neural networks (e.g., >300 TB for a 124M-parameter model). The paper restricts $\boldsymbol{\Sigma}$ to structured classes $\mathcal{S}$ (diagonal or Kronecker-factored), solving:

$$\underset{\boldsymbol{\Sigma} \in \mathcal{S}}{\text{minimize}} \quad \mathrm{KL}\Big(\mathcal{N}(0, \hat{\boldsymbol{\Sigma}}_t^{\prime}) \| \mathcal{N}(0, \boldsymbol{\Sigma})\Big) \tag{10}$$

which reduces to:

$$\hat{\boldsymbol{\Sigma}}_t = \underset{\boldsymbol{\Sigma} \in \mathcal{S}}{\operatorname{argmin}} f(\boldsymbol{\Sigma}), \qquad f(\boldsymbol{\Sigma}) := \log|\boldsymbol{\Sigma}| + \mathrm{tr}(\boldsymbol{\Sigma}^{-1}\hat{\boldsymbol{\Sigma}}_t^{\prime}) \tag{11}$$

Diagonal Approximation → Centered Adam

For diagonal $\mathcal{S}$, the solution is:

$$\hat{\boldsymbol{\sigma}}_t = \beta \hat{\boldsymbol{\sigma}}_{t-1} + \beta(1-\beta)(\boldsymbol{g}_t - \hat{\boldsymbol{m}}_{t-1}) \circ (\boldsymbol{g}_t - \hat{\boldsymbol{m}}_{t-1}) \tag{12}$$

This yields Centered Adam:

$$\boldsymbol{w}_{t+1} = \boldsymbol{w}_t - \eta \,\mathrm{diag}(\hat{\boldsymbol{\sigma}}_t)^{-1/2}\hat{\boldsymbol{m}}_t \tag{14}$$

Key insight: Adam (with $\beta_1 = \beta_2 = \beta$, $\varepsilon = 0$) is equivalent to:

$$\boldsymbol{w}_{t+1} = \boldsymbol{w}_t - \eta \big(\mathrm{diag}(\hat{\boldsymbol{\sigma}}_t + \hat{\boldsymbol{m}}_t \circ \hat{\boldsymbol{m}}_t)\big)^{-1/2}\hat{\boldsymbol{m}}_t \tag{13}$$

Thus Adam whitens using the uncentered second moment, whereas Centered Adam uses the centered second moment (covariance). The paper shows Centered Adam performs comparably to Adam, suggesting that explaining Adam may not require exact accounting—a closely related whitening method suffices.

Kronecker-Factored Approximation → MN Method

Assuming gradient matrices follow a Matrix Normal distribution:

$$\boldsymbol{g}_t = \operatorname{vec}(\mathbf{G}_t) \sim \mathcal{N}(\boldsymbol{m}_t, \mathbf{C}_t \otimes \mathbf{R}_t) \tag{15}$$

The whitened update becomes:

$$\mathbf{W}_{t+1} = \mathbf{W}_t - \eta \hat{\mathbf{R}}_t^{-1/2}\hat{\mathbf{M}}_t\hat{\mathbf{C}}_t^{-1/2} \tag{17}$$

Proposition 5.1 guarantees a unique global minimizer for the Kronecker-factored covariance problem, found via alternating minimization:

$$\hat{\mathbf{R}}_t^{(k+1)} = \frac{\beta}{n}\mathrm{tr}\big((\hat{\mathbf{C}}_t^{(k)})^{-1}\hat{\mathbf{C}}_{t-1}\big)\hat{\mathbf{R}}_{t-1} + \frac{\beta(1-\beta)}{n}\mathbf{X}_t(\hat{\mathbf{C}}_t^{(k)})^{-1}\mathbf{X}_t^\top \tag{19}$$ $$\hat{\mathbf{C}}_t^{(k+1)} = \frac{\beta}{m}\mathrm{tr}\big((\hat{\mathbf{R}}_t^{(k+1)})^{-1}\hat{\mathbf{R}}_{t-1}\big)\hat{\mathbf{C}}_{t-1} + \frac{\beta(1-\beta)}{m}\mathbf{X}_t^\top(\hat{\mathbf{R}}_t^{(k+1)})^{-1}\mathbf{X}_t \tag{20}$$

where $\mathbf{X}_t = \mathbf{G}_t - \hat{\mathbf{M}}_{t-1}$.

The MND method (Algorithm 1) restricts factors to diagonal matrices, giving elementwise updates with numerical stability via clipped inverses:

$$\operatorname{inv}_\varepsilon(v) = \max\{v, \varepsilon\}^{-1}, \qquad \operatorname{isqrt}_\varepsilon(v) = \max\{v, \varepsilon\}^{-1/2} \tag{23}$$

Statistical View of Muon

The Muon update uses the polar decomposition of momentum:

$$\mathbf{W}_{t+1} = \mathbf{W}_t - \eta \,\mathrm{Polar}(\hat{\mathbf{M}}_t) \tag{24}$$

Key insight: $\mathrm{Polar}(\hat{\mathbf{M}}_t)$ is the mode of a Matrix Langevin distribution $\mathrm{ML}_{m,n}(\hat{\mathbf{M}}_t)$ on the Stiefel manifold. The full Muon update emerges from a Bayesian model with:

• Likelihood: $\mathbf{G}_t \sim \mathrm{MN}_{m,n}(\mathbf{D}_t, \mathbf{I}_m, \mathbf{I}_n)$ (isotropic noise)
• Prior: $\mathbf{D}_t \sim \mathrm{ML}_{m,n}\big(\frac{\beta}{1-\beta}\hat{\mathbf{M}}_{t-1}\big)$

The posterior mode gives exactly Polar$(\hat{\mathbf{M}}_t)$, recovering Muon.

MN-Muon generalizes this by using an anisotropic likelihood with covariance factors:

$$p(\mathbf{G}_t \mid \mathbf{D}_t, \mathbf{R}_t, \mathbf{C}_t) = \mathrm{MN}_{m,n}(\mathbf{R}_t^{1/2}\mathbf{D}_t\mathbf{C}_t^{1/2}, \mathbf{R}_t, \mathbf{C}_t) \tag{34}$$

The joint optimization problem (43) couples direction and covariance, solved by alternating block updates (Proposition 7.1). The diagonal implementation (MND-Muon, Algorithm 2) uses closed-form factor updates via the function:

$$\Psi_q(u, v) := \Big(\frac{\sqrt{v^2 + 4qu} - v}{2q}\Big)^2 \tag{45}$$

---

Empirical Validation / Results

Centered Adam vs. Adam (Figure 3)

Final validation losses for Adam and Centered Adam on FineWeb are nearly identical across all learning rates at both GPT-small and GPT-medium scales, confirming that Centered Adam's whitening interpretation suffices to explain Adam's practical effectiveness.

MND vs. Shampoo Variants (GPT-small, Figure 4)

Training on FineWeb-10B, 2.5B tokens, batch 512, three seeds:

MND vs. Shampoo Variants (GPT-medium, Figure 5)

Training on FineWeb, 7.09B tokens, batch 768, three seeds: Centered MND achieves the lowest mean final validation loss, followed closely by uncentered MND. Both established Shampoo implementations attain higher losses than all three MND variants.

MND-Muon vs. Baselines (GPT-small, Figure 6)

FineWeb, batch 512, 2.5B tokens:

MND-Muon improves on Muon in 21 of 24 paired GPT-small runs.

MND-Muon vs. Baselines (GPT-medium, Figure 7)

FineWeb, 7.09B tokens, batch 768, seed 42, six learning rates:

MND-Muon maintains validation losses between 3.014–3.016 across learning rates 0.02–0.06 (range 0.0024 nats/token), while Muon becomes unstable at 0.045 (loss 6.295) and NorMuon at 0.09 (loss 6.071).

---

Theoretical and Practical Implications

Theoretical Contributions

1. Unified framework: Adam, Shampoo, KL-Shampoo, and Muon can all be viewed as instances of whitened momentum methods with different covariance structures.
2. Bayesian foundation: The mean/variance estimation follows from a principled Bayesian recursion (Normal-Inverse-Wishart prior) or equivalently a regularized likelihood optimization.
3. Novel methods: The framework naturally generates new algorithms (Centered Adam, MND, MND-Muon) with clear statistical interpretations.
4. Muon's statistical meaning: Muon's polar decomposition is shown to be the MAP estimate under a Matrix Langevin model, providing theoretical justification for its empirical success.

Practical Implications

• MND offers a drop-in replacement for Shampoo with better stability, lower final validation loss, and reduced time.
• MND-Muon improves robustness of Muon to learning rate choices, particularly valuable at smaller batch sizes where gains over Muon are largest.
• The diagonal implementations avoid the numerical instability of full matrix inverses in low precision (BF16), making them practical on modern hardware.
• The framework suggests that "explaining Adam" may not require exact accounting—a closely related whitening method achieves comparable performance with a simpler interpretation.

---

Conclusion

Main Takeaways

The paper postulates that stochastic optimization methods can be unified through the lens of whitened momentum: estimating gradient mean and covariance, then whitening the mean before stepping. Different structured covariance approximations recover and generalize existing methods:

• Diagonal covariance → Centered Adam (comparable to Adam)
• Kronecker-factored covariance → MN/MND (outperforms Shampoo variants)
• Direction modeling → MND-Muon (improves on Muon)

Limitations and Future Directions

1. Theoretical justification: The whitened momentum framework's effectiveness is mostly empirical; deeper theoretical analysis on toy models or function classes would strengthen the argument.
2. Scope of experiments: Only GPT-2 models at Chinchilla scale were tested; more models, datasets, and scales are needed.
3. Gaussian assumption: Gradients in language models may not be Gaussian (Zhang et al., 2020; Kunstner et al., 2023); exploring other distributions (e.g., heavy-tailed) could yield further insights.
4. Full MN method: The full-factor MN underperformed due to numerical instability of matrix inversion in BF16; a stable, GPU-friendly matrix inversion could make it competitive.
5. NorMuon comparison: NorMuon was comparable to or slightly better than MND-Muon in some settings, suggesting further investigation into their relationship is warranted.