[Pretraining 2026] Newton–Muon: Decoding the "Black Box" of Matrix Orthogonalization

TL;DR

The Muon optimizer has taken the LLM world by storm, but why its "gradient orthogonalization" works so well has remained a mystery. Newton–Muon provides the answer: Muon is actually an incomplete Newton method. By adding a simple right-preconditioner based on the second moment of input activations ($Z{Z}^{o}p$), Newton–Muon achieves 6% faster convergence and 4% lower wall-clock time than the original Muon in GPT-2 speedruns.

Problem: The Hidden Anisotropy in Muon

Standard Muon updates weights using $W\leftarrow W-\eta \cdot \mathrm{msgn}(G)$. While this successfully captures the "direction" of the gradient by discarding singular values, it makes a massive, implicit assumption: that the input data to each layer is isotropic (perfectly uniform in all directions).

In reality, data activations in LLMs are famously anisotropic (highly skewed). This skewness distorts the gradient, making the optimization "see" a warped landscape. Standard Muon is essentially "blind" to this data geometry.

Methodology: The Triplet Quadratic Surrogate

The authors introduce a principled framework called the Triplet Quadratic Surrogate. Instead of just looking at the gradient $G$, they model the loss change $J(Q)$ around a weight matrix $W$ using three components:

1. G: The Gradient matrix (First-order info).
2. H: Output-space curvature (The Hessian's left side).
3. Z: Input activations (The Hessian's right side).

The Newton-Muon Update Rule

Through a "one-step descent analysis," the authors prove that if you assume the difference between your current weights and the optimal weights is isotropic, the optimal update is: $W\leftarrow W-\eta \cdot \mathrm{msgn}(G(Z{Z}^{o}p{)}^{-1})$

This formula is beautiful in its simplicity: it uses $(Z{Z}^{o}p{)}^{-1}$ to "whiten" the input data before applying the matrix sign.

Figure: The core Newton-Muon update rule derived from the triplet surrogate.

Experiments: Breaking the Speed Records

The researchers tested Newton–Muon on the Modded-NanoGPT speedrun benchmark, which is the gold standard for measuring how fast an optimizer can reach a specific validation loss.

Key Findings:

• Convergence Efficiency: Newton–Muon reaches target loss in 6% fewer steps.
• Wall-Clock Gains: Even with the overhead of computing matrix inverses, it reduces total training time by 4%.
• Robustness to "Spikes": In synthetic tests where data has "spiked" directions (high condition numbers), Newton–Muon's performance remains constant, while Muon and AdamW degrade significantly.

Figure: Newton–Muon outperforms Muon and AdamW on both GPT-2 and CIFAR-10 tasks.

Why It Works: A Physical Intuition

Think of optimization as trying to find the bottom of a bowl. If the bowl is perfectly round, any direction works. If the bowl is a narrow "crevice" (anisotropic), you'll bounce off the walls.

• AdamW tries to fix this element-wise (scalar).
• Muon tries to fix this by making the gradient matrix "square" (orthogonal).
• Newton–Muon realizes that the "crevice" is often caused by the input data. By right-preconditioning with the inverse of the data's second moment, it effectively "rounds out" the crevice before taking a step.

Implementation & Practicality

You might worry that inverting a matrix $(Z{Z}^{o}p{)}^{-1}$ is too slow. The authors solve this with three clever tricks:

1. Periodic Refreshes: They only re-compute the inverse every 16 or 32 steps.
2. Newton–Schulz Iteration: They use high-performance polynomial approximations for the matrix sign.
3. Symmetric Compute: They use custom Triton kernels to exploit matrix symmetry, cutting the math workload in half.

Conclusion & Future Look

Newton–Muon proves that we haven't reached the limit of first-order optimizers. By bringing a "local second-order view" to matrix optimization, we can train models faster without the massive memory overhead of full Natural Gradient Descent.

Limitations: The method currently relies on an "isotropic proxy" for weights, which might not hold true in very late-stage training. Future research into estimating weight displacement (${\Sigma }_W$) could unlock even more speed.

Primary Source: The Newton–Muon Optimizer, Du & Su, University of Pennsylvania (2026).