A Quadratic Lens on Muon's Polar Update: Benefits, Pitfalls, and the Role of Newton–Schulz

Summary (Overview)

• Core contribution: This paper provides the first closed-form analysis of PolarGD—the orthogonalized update at the core of the Muon optimizer—on matrix quadratics, isolating the effect of the update direction from step-size tuning via exact line search.
• Key finding 1: PolarGD acts as gradient descent preconditioned by the current iterate ($XX^\top$) rather than by the Hessian, so its speed is governed by the conditioning of the iterate, not the Hessian. This explains both fast regimes (row-isotropic iterates) and slow ones (tiny singular values).
• Key finding 2: Finite Newton–Schulz iterations (as used in Muon) act as a damped polar update: they behave like exact polar for large singular values but like GD for tiny ones. Near rank deficiency, fewer Newton–Schulz steps converge faster than exact orthogonalization—so Newton–Schulz is not merely an approximation of the polar factor.
• Key finding 3: PolarGD is equivariant under orthogonal rotations of the problem, while SignGD is not. This makes PolarGD robust to coordinate changes where sign descent degrades (e.g., rotated Zipf-type problems).
• Practical relevance: In a 90M-parameter transformer trained with Muon, roughly a quarter of the singular values of attention matrices fall in the regime where Newton–Schulz behaves like GD rather than exact polar.

Introduction and Theoretical Foundation

The Muon optimizer (Jordan et al., 2024b) has emerged as a strong alternative to AdamW for training large neural networks, matching or outperforming it in language-model pretraining, vision, diffusion, and tabular models. For two-dimensional parameter blocks, Muon approximately orthogonalizes the momentum matrix via a few Newton–Schulz iterations, stepping along an approximation of the polar factor:

$$\mathrm{polar}(G) = UV^{\top}$$

for $G = U\Sigma_G V^{\top}$.

The same direction underlies a broader family of matrix optimizers, including stochastic spectral descent (Carlson et al., 2015a, 2016), norm-constrained methods (Pethick et al., 2025), and Shampoo (Gupta et al., 2018) with preconditioner accumulation switched off.

The central open question: When and why does orthogonalization help? Existing analyses provide worst-case convergence bounds governed by smoothness constants in chosen norms, but they do not separate the effect of the direction from the step-size scale. This paper addresses this gap by studying the core update on matrix quadratics:

$$X_{k+1} = X_k - \gamma_k \mathrm{polar}(\nabla f(X_k))$$

called PolarGD by analogy with SignGD (which replaces the gradient by its entrywise sign).

Methodological choice: To compare directions (GD, PolarGD, SignGD) on equal footing, each method takes its best scalar step via exact line search, removing learning-rate tuning from the comparison. On matrix quadratics, the resulting one-step decreases are available in closed form.

Methodology

Problem Setting

The paper analyzes the matrix quadratic:

$$f(X) = \frac{1}{2}\mathrm{tr}\big((X - X^{\star})^{\top} H (X - X^{\star})\big)$$

with $X \in \mathbb{R}^{d \times m}$, $H = H^{\intercal} \succeq 0$, and gradient $G = H(X - X^{\star})$.

Key Definitions

The three update directions (dropping index $k$):

• GD: $D_{\mathrm{GD}} := G$
• PolarGD: $D_{\mathrm{Pol}} := Q := \mathrm{polar}(G)$
• SignGD: $D_{\mathrm{Sign}} := S := \mathrm{sign}(G)$

Both $Q$ and $S$ are LMO (linear minimization oracle) directions: $Q$ for the spectral norm and $S$ for the entrywise $\ell_\infty$ norm.

Line-Search Comparison Framework

For a direction $D$ with gain $\langle G, D \rangle > 0$, the curvature under the upper bound $L$ is $\|D\|_L^2 := \langle LD, D \rangle$, giving the step:

$$\gamma_L(D) := \frac{\langle G, D \rangle}{\|D\|_L^2}$$

Proposition 1 (One-step decrease): Under the curvature bounds (2), a direction is judged by its squared gain over its curvature:

$$f(X^{+}) \leq f(X) - \frac{\langle G, D \rangle^2}{2\|D\|_L^2}$$

The polar direction guarantees a larger decrease than GD if and only if:

$$\frac{\|G\|_L}{\|Q\|_L} > \frac{\|G\|_{\mathrm{F}}^2}{\|G\|_*}$$

This requires matrix-valued curvature bounds—a scalar bound $L = \ell I_d$ never allows polar to beat GD.

Exact One-Step Decreases on Quadratics

$$\Delta_{\mathrm{GD}} = \frac{\|G\|_{\mathrm{F}}^4}{2\mathrm{tr}(G^{\top}HG)}, \quad \Delta_{\mathrm{Pol}} = \frac{\|G\|_*^2}{2\mathrm{tr}(Q^{\top}HQ)}, \quad \Delta_{\mathrm{Sign}} = \frac{\|G\|_1^2}{2\mathrm{tr}(S^{\top}HS)}$$

Empirical Validation / Results

One-Step Convergence Regimes

Theorem 1 (One-step convergence under exact line search):

Key insight: PolarGD reaches the minimizer in one step for any positive-definite Hessian when the rows of the error are orthogonal with equal norms (e.g., scaled orthogonal initialization).

Rotation Equivariance

Theorem 2: For orthogonal $P \in \mathbb{R}^{d \times d}$, $R \in \mathbb{R}^{m \times m}$, and $\tilde{H} := P^{\top}HP$: GD and PolarGD are equivariant under orthogonal changes of basis (errors satisfy $E_k = P\tilde{E}_k R^{\top}$, losses coincide for all $k$). SignGD does not have this property.

Zipf experiment: On a Zipf-type regression model of language data (diagonal $H$ with $\lambda_k(H) \propto k^{-\zeta}$), PolarGD and SignGD both perform well on diagonal problems, but when $H$ is randomly rotated, SignGD degrades significantly while PolarGD is unaffected (Figure 2).

Implicit Preconditioning

The polar direction can be rewritten as:

$$Q = \mathrm{polar}(HX) = (H\Sigma H)^{\dagger/2} HX =: A(X)X, \quad A(X) := (H\Sigma H)^{\dagger/2} H$$

Proposition 2: In the full-row-rank regime ($\mathrm{rank}(X) = d$):

$$A(X)\Sigma A(X)^{\top} = I_d, \quad A(X)^{\top}A(X) = \Sigma^{-1}$$

The singular values of $A(X)$ are $\{1/\sigma_i(X)\}_{i=1}^d$, so:

$$\kappa(A(X)) = \kappa(X) = \sqrt{\kappa(\Sigma)}$$

Theorem 3 (Exact progress coefficient): With exact line search:

$$\frac{f(X^{+})}{f(X)} = 1 - p(X), \quad p(X) := \frac{\|HX\|_*^2}{\mathrm{tr}(H)\mathrm{tr}(HXX^{\top})}$$

with bounds:

$$\max\left\{\frac{\lambda_{\min}(H)}{\mathrm{tr}(H)}, \frac{4\kappa(X)}{(1+\kappa(X))^2}\right\} \leq p(X) \leq 1$$

The second lower bound behaves like $4/\kappa(X)$ for ill-conditioned iterates—the Kantorovich factor of GD with $\kappa(H)$ replaced by $\kappa(X)$. The first lower bound can be as small as $\approx 1/(\kappa d)$, and is tight: a near-rank-deficient family exists where GD is nearly one-step exact while PolarGD contracts at a slow fixed rate (Corollary 5, Figure 4).

Newton–Schulz as Damped Polar

Finite Newton–Schulz steps map each normalized singular value $s$ to $\phi_q(s)$. Small singular values are multiplied by a fixed factor $\ell_q$ (growing geometrically with $q$; $\ell_5 \approx 485$ for Muon's coefficients), while those above $\approx 1/\ell_q$ are mapped close to 1.

Key finding: Finite Newton–Schulz acts like GD on singular values below the damping scale $1/\ell_q$ and like exact polar above it. Near rank deficiency, fewer steps converge faster than exact polar (Figure 4)—the approximation error is not merely a degradation factor.

Transformer Measurements

In a 90M-parameter transformer trained with Muon on 1.8B tokens:

Percentage of singular values below the damping scale $1/\ell_5 \approx 2 \cdot 10^{-3}$. Attention matrices are much harder to orthogonalize than MLP matrices; the effect is largest in early layers. About 16% of the update's squared Frobenius norm would be placed on these directions by exact polar, vs. ~5% by five Newton–Schulz steps.

Theoretical and Practical Implications

1. Direction vs. step size: The paper cleanly separates the effect of update direction from step-size choice via line search, showing that one-step preferences need not persist along trajectories (Figure 3 shows alternating local winners between GD and PolarGD along the GD path).

2. Preconditioning perspective: PolarGD is a self-preconditioned method whose effective conditioning is that of the iterate ($\kappa(X)$), not the Hessian. This explains both its robustness to ill-conditioned Hessians near row-isotropic iterates and its failure near rank deficiency.

3. Newton–Schulz as a feature, not a bug: The finite Newton–Schulz iterations in Muon provide damping of tiny singular values that can accelerate convergence near rank deficiency. This suggests treating the number of Newton–Schulz steps (or damping level) as a design choice of the optimizer, not merely an accuracy setting.

4. Rotation equivariance matters: PolarGD's equivariance under orthogonal transformations makes it robust to coordinate-system choices, a property sign descent lacks—relevant for understanding why Muon performs well on language models where the natural coordinate system is not aligned with the Hessian eigenvectors.

5. Caveats: Exact line search alone does not make PolarGD uniformly better than GD; its advantage depends on how the iterate covariance evolves along the trajectory. Greedy one-step choices (line search) can be worse than tuned fixed steps for PolarGD.

Conclusion

The paper isolates the geometry of the orthogonalized update at the core of Muon by analyzing PolarGD on matrix quadratics. PolarGD benefits from rotation equivariance and implicit preconditioning by the iterate, but is not uniformly better than GD or SignGD: it is fast near row-isotropic iterates and slow near rank deficiency, where exact polar normalization overweights tiny singular values. Finite Newton–Schulz softens the update exactly in that regime and can converge faster than exact polar, making it a genuine design element rather than a mere approximation. The trajectory results rely on quadratic structure, but Proposition 1 only requires curvature bounds, and equivariance plus Newton–Schulz softening hold for any objective—including the momentum matrix Muon orthogonalizes. Transformer measurements confirm that the regime where exact polar and Newton–Schulz differ occurs in practice, suggesting future work on adaptive Newton–Schulz step counts or damping levels as optimizer hyperparameters.