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

---

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.