# When and How Does Muon Precondition? The Role of Structured Gradient Noise

> Muon's preconditioning benefits require structured gradient noise, not isotropic noise, and it softens task-frequency dependence from linear to square-root scaling.

- **Source:** [arXiv](https://arxiv.org/abs/2610.31728)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/4268cZ
- **Whiteboard:** https://picx.dev/p/4268cZ/image

## Summary

# Summary of "When and How Does Muon Precondition? The Role of Structured Gradient Noise"

## Summary (Overview)

- **Core Question**: The paper investigates when and why the Muon optimizer provides preconditioning benefits under stochastic gradient noise, particularly in the context of large language model (LLM) training.
- **Key Finding 1**: Under isotropic Gaussian noise, Muon's first-order conditional mean update reduces to a scalar multiple of the population gradient, offering no directional preconditioning advantage beyond a learning-rate rescaling.
- **Key Finding 2**: Under structured noise from mini-batch sampling in associative memory with heterogeneous task frequencies, Muon exhibits an exact balancing threshold $B_{\text{bal}}$ for two tasks, and softens task-frequency dependence from linear to square-root scaling.
- **Key Finding 3**: Muon's orthogonal equivariance protects against rotation-induced negative transfer that affects SignSGD under sparse sampling.
- **Key Finding 4**: Momentum provides a variance-reduction/tracking trade-off, motivating heterogeneous (fast-slow) momentum designs, validated by a preliminary 0.12B-parameter LLM pretraining experiment.

---

## Introduction and Theoretical Foundation

### Background and Motivation

Muon has emerged as a prominent alternative to Adam for training matrix parameters in LLMs. Its defining operation replaces a gradient or momentum matrix by an approximate **polar factor**, flattening nonzero singular values before parameter updates. Recent frontier models (DeepSeek-V4.1-Flash, Kimi-K3) have adopted Muon variants, making theoretical understanding of its benefits increasingly practical.

The key theoretical challenge: Muon acts on **stochastic mini-batch gradients**, and nonlinear normalization does not commute with expectation:

$$\mathbb{E}[\text{msgn}(G_t) | W_t] \neq \text{msgn}(\mathbb{E}[G_t | W_t])$$

### Theoretical Foundation

**Partial Polar Map**: For $G \in \mathbb{R}^{d \times d}$ with SVD $G = U\Sigma V^\top$, define:

$$\text{msgn}(G) := U \operatorname{diag}(\mathbb{1}\{\sigma_i > 0\}) V^\top$$

Key properties:
- **Positive scale invariance**: $\text{msgn}(aG) = \text{msgn}(G)$ for $a > 0$
- **Orthogonal equivariance**: $\text{msgn}(QGR^\top) = Q\,\text{msgn}(G)\,R^\top$ for orthogonal $Q, R$
- **Nullspace preservation**: $\ker(\text{msgn}(G)) = \ker(G)$

**Muon Dynamics**: $M_t = \beta M_{t-1} + (1-\beta)G_t$, $W_{t+1} = W_t - \eta\,\text{msgn}(M_t)$, where $\beta \in [0,1)$ is momentum and $\eta > 0$ is learning rate.

---

## Methodology

### 1. Isotropic Noise Analysis

The authors prove a low-SNR expansion showing that when gradient noise is isotropic Gaussian, Muon's conditional mean response collapses to a scalar multiple of the population gradient.

### 2. Associative Memory Model

The paper studies a linear associative memory with $K \geq 2$ associations $(E_j, \tilde{E}_j)$ sampled with probabilities $p_j > 0$, where $\sum_j p_j = 1$. The population loss is:

$$\mathcal{L}(W) = -\sum_{j=1}^K p_j \log S_{jj}$$

where $S_{ij} = \exp(W_{ij}) / \sum_{k=1}^K \exp(W_{kj})$ are prediction probabilities. The stochastic gradient in task coordinates yields the exact Muon update:

$$W_{t+1} = W_t + \eta\,\text{msgn}((I - S_t)Q_t)$$

where $Q = \operatorname{diag}(q)$ with $Bq \sim \text{Multinomial}(B, p)$.

### 3. Label Noise Model

Symmetric label noise with parameter $\alpha$: target label preserved with probability $c_\alpha = 1 - \alpha + \alpha/K$, nontarget label chosen with aggregate probability $\rho_\alpha = 1 - c_\alpha$.

### 4. Momentum Analysis

Frozen-state variance analysis and local tracking error bounds characterize the averaging-tracking trade-off.

---

## Empirical Validation / Results

### Isotropic Noise: No First-Order Preconditioning

**Proposition 3.1**: As $\|\bar{G}\|_F / \sigma \to 0$,

$$\mathbb{E}[\text{msgn}(\bar{G} + \sigma\Xi)] = \frac{c_d}{\sigma}\bar{G} + O_d\left(\frac{\|\bar{G}\|_F^3}{\sigma^3}\right)$$

where $c_d = d^{-2}\mathbb{E}\|\Xi\|_* > 0$. This matches SGD with learning rate $\eta_S = \eta_M c_d/\sigma$.

### Two-Task Balancing Threshold

**Theorem 4.1**: With $p_{\min} = p_2$ and $B_{\text{bal}} = \log 2 / [-\log(1 - p_{\min})]$:

- **If $B < B_{\text{bal}}$** (unequal growth): $(\Delta_1(t), \Delta_2(t))/(\eta t) \xrightarrow{a.s.} \sqrt{2}(p_1^B, 1 - p_1^B)$
- **If $B > B_{\text{bal}}$** (controlled gap): $\sup_t \mathbb{E}[e^{a|\Delta_1(t) - \Delta_2(t)|}] < \infty$ and $(\Delta_1(t) - \Delta_2(t))/t \xrightarrow{a.s.} 0$

For rare tasks, $B_{\text{bal}} \sim (\log 2)/p_{\min}$.

### Higher-Rank Frequency Dependence

**Proposition 4.2**: For task $j$ in a block with expected count $Bp_j$:
- Task appears with probability $1 - (1-p_j)^B \sim Bp_j$ as $Bp_j \to 0$
- Well-sampled blocks recover full-batch margin increment: $\eta + O_P(\eta/\sqrt{Bp_j})$

### Label Noise: Square-Root vs. Linear Frequency Dependence

**Proposition 4.3**: 
- Deterministic labels: $\text{SNR}_j^2 = Bp_j/(1-p_j)$
- Label noise: $\text{SNR}_j^2 \sim Bp_j\|\delta_j\|_2^2 / \operatorname{tr}\Gamma_j \to 0$ as $\delta_j \to 0$

**Proposition 4.4**: Near the noisy optimum, Muon's contraction coefficient scales as $\sqrt{Bp_j}$ while SGD scales linearly as $p_j$. Empirical results confirm $p_j^{0.494}$ for Muon (close to square-root prediction).

### Muon vs. SignSGD Under Sparse Sampling

**Proposition 4.5**: For a batch with no task-2 samples:
- Muon: $\Delta_2^{\text{Muon}}(t+1) - \Delta_2^{\text{Muon}}(t) = 0$ (nullspace preservation)
- SignSGD: $\Delta_2^{\text{Sign}}(t+1) - \Delta_2^{\text{Sign}}(t) = \eta\tilde{c}(\cos\theta - \sin\theta)$ (negative transfer possible when $\theta > \pi/4$)

### Momentum: Variance Reduction vs. Tracking

**Proposition 5.1**: 
- Variance-equivalent batch size: $B_{\text{eff}}(\beta) = B(1+\beta)/(1-\beta)$
- Tracking error bound: $L_S\eta\beta/(1-\beta)$, vanishing when $\eta/(1-\beta) \to 0$
- Sufficient window: $\eta \ll 1 - \beta \ll Bp_j$

### LLM Pretraining Results

| Method | Validation Loss (step 30k) |
|--------|---------------------------|
| Standard Muon | 3.0097 |
| Muon + Fast-Slow | 2.9888 |
| **Gap** | **0.0209** |

---

## Theoretical and Practical Implications

### Theoretical Implications

1. **Isotropic noise is a null case**: Samplewise spectral normalization alone does not guarantee directional preconditioning—structured noise is necessary for Muon's benefits.

2. **Sampling structure matters**: The expected task count $Bp_j$ is the controlling scale for balancing learning across heterogeneous tasks. Muon softens frequency dependence from linear ($p_j$) to square-root ($\sqrt{Bp_j}$) near noisy optima.

3. **Orthogonal equivariance is protective**: Muon's rotation-invariance prevents negative transfer to unobserved tasks, unlike coordinate-dependent SignSGD (a proxy for Adam).

4. **Momentum is a statistical tool**: The sufficient window $\eta \ll 1-\beta \ll Bp_j$ provides a sampling-based justification for heterogeneous momentum—longer memory for rare, slowly varying directions.

### Practical Implications

- **Optimizer design guidance**: The theory motivates fast-slow momentum designs that allocate longer memory to rare/low-curvature directions.
- **Batch size selection**: The threshold $B_{\text{bal}} \sim (\log 2)/p_{\min}$ provides guidance for choosing batch sizes relative to rare-task frequencies.
- **Practical validation**: The 0.12B-parameter pretraining experiment shows a 0.0209 validation-loss improvement from fast-slow momentum over standard Muon.

---

## Conclusion

The paper establishes that Muon's preconditioning benefits **require structured gradient noise** to manifest:

1. Under isotropic noise, Muon offers no first-order directional advantage over learning-rate-rescaled SGD.
2. Under structured sampling noise, Muon achieves an exact balancing threshold for two tasks and square-root frequency dependence near noisy optima—a significant softening from SGD's linear dependence.
3. Muon's orthogonal equivariance protects against rotation-induced negative transfer that affects coordinatewise methods like SignSGD.
4. Momentum provides a variance-reduction/tracking trade-off that motivates heterogeneous (fast-slow) momentum designs, with preliminary LLM pretraining evidence of practical benefit.

**Future directions**: Extending analysis to evolving, non-orthogonal representations and validating benefits across larger training scales remain important open problems.

---

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