# A Statistical Viewpoint on Adam, Shampoo and Muon

> Whitening gradients with estimated covariance unifies Adam, Shampoo, and Muon, yielding MND that outperforms Shampoo and MND-Muon that improves Muon's stability.

- **Source:** [arXiv](https://arxiv.org/abs/2610.24839)
- **Published:** 2026-10-24
- **Permalink:** https://picx.dev/p/sgyj0f
- **Whiteboard:** https://picx.dev/p/sgyj0f/image

## Summary

# 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:

| Method | Mean Final Validation Loss |
|---|---|
| **MND (centered)** | **3.373** |
| MND (uncentered) | 3.376 |
| MND ($\beta_2$) | 3.409 |
| KL-Shampoo | 3.444 |
| Distributed Shampoo | 3.494 |

### 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:

| Method | Mean Final Validation Loss |
|---|---|
| **NorMuon** | **3.327** |
| **MND-Muon** | **3.329** |
| Muon | 3.336 |
| AdamW | 3.381 |
| KL-Shampoo | 3.449 |
| Distributed Shampoo | 3.501 |

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:

| Method | Best Final Validation Loss | Learning Rate |
|---|---|---|
| **NorMuon** | **3.007** | 0.045 |
| **MND-Muon** | **3.014** | 0.045 |
| Muon | 3.017 | 0.04 |

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.

---

_Markdown view of https://picx.dev/p/sgyj0f, served by PicX — AI-generated visual whiteboard summaries of research papers._