Muon Optimizer: How Matrix Orthogonalization and Newton-Schulz Iterations Accelerate LLM Training

Modern large language model pre-training has relied on AdamW as its default optimizer for nearly a decade. While AdamW provides robust convergence across varied architectures, its fundamental formulation treats neural network weights as flat collections of independent scalar parameters. For the 2D weight matrices that dominate Transformer architectures—including attention projections and feed-forward linear layers—this coordinate-wise treatment ignores the underlying matrix geometry and singular value spectrum.

The Muon optimizer, introduced by Keller Jordan, Jeremy Bernstein, and collaborators, and detailed in the technical report Muon is Scalable for LLM Training, introduces matrix-level orthogonalization to gradient updates. By replacing coordinate-wise variance normalization with high-order Newton-Schulz matrix iterations, Muon constrains updates to orthogonal manifolds. This structural change yields up to 1.5x improved sample efficiency, reduces optimizer memory footprints, and accelerates convergence in large-scale model pre-training.

The Coordinate-Wise Bottleneck in AdamW

In standard AdamW, first and second gradient moments are tracked independently for every scalar coordinate $(i, j)$ in a weight tensor $W \in \mathbb{R}^{m \times n}$:

$$m_t = \beta_1 m_{t-1} + (1 - \beta_1) g_t$$ $$v_t = \beta_2 v_{t-1} + (1 - \beta_2) g_t^2$$ $$u_t = \frac{m_t}{\sqrt{v_t} + \epsilon}$$

This update rule is axis-aligned. If the coordinate system is rotated, the resulting step changes. For a linear layer computing $y = W x$, what governs representation transformation is not individual matrix entries, but the matrix's singular values and principal directional components.

Under AdamW, updates can stretch along dominant singular directions while distorting weaker spectral dimensions, leading to ill-conditioned intermediate representations. While learning rate warmups and weight decay mitigate numerical instability, the optimizer remains structurally unaware of the 2D linear operator it modifies.

Orthogonalized Momentum via Polar Decomposition

Rather than normalizing each scalar element by its historical variance, Muon treats the momentum buffer $M_t$ as a matrix operator. The objective is to find the nearest orthogonal matrix $O_t$ to $M_t$ in terms of Frobenius norm:

$$\text{Ortho}(M_t) = \arg\min_O \|O - M_t\|_F \quad \text{subject to} \quad O^T O = I \text{ or } O O^T = I$$

In linear algebra, this operation corresponds to the polar decomposition or matrix sign function. If the momentum matrix has singular value decomposition $M_t = U \Sigma V^T$, the closest orthogonal matrix is simply $U V^T$. This maps every non-zero singular value directly to 1.0, ensuring that the update step applies uniform spectral energy across all orthogonal directional modes without stretching or collapsing specific subspaces.

However, executing an exact Singular Value Decomposition at every iteration across dozens of multi-billion parameter layers would introduce prohibitive computational overhead, as SVD scales with $O(\min(m^2 n, m n^2))$ complexity and maps poorly to Tensor Core architectures.

Fast GPU Orthogonalization with Newton-Schulz Iterations

To make matrix orthogonalization viable in production training loops, Muon computes the approximate polar decomposition using a high-order Newton-Schulz matrix iteration.

The iteration operates strictly through matrix multiplications (GEMMs), allowing it to execute entirely at peak Tensor Core throughput on modern accelerators like Nvidia H100s and B200s.

Starting with a normalized momentum matrix $X_0 = \frac{M_t}{\|M_t\|_F + \epsilon}$, the algorithm applies a 5th-order polynomial recurrence for a fixed number of steps (typically $N = 5$):

$$X_{k+1} = a X_k + b (X_k X_k^T) X_k + c (X_k X_k^T)^2 X_k$$

In the reference implementation developed for the modded-nanogpt speedruns, Jordan tuned the coefficients $(a, b, c) = (3.4445, -4.7750, 2.0315)$ to rapidly inflate small singular values toward unity within just five steps, avoiding the need for deeper iterations.

Because the entire routine consists of five fused matrix multiplications per parameter tensor, the wall-clock overhead of the orthogonalization step is under 2% of the overall forward-backward iteration time.

import torch

@torch.compile
def zeropower_via_newtonschulz5(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
    """
    Computes approximate polar decomposition U @ V.T using 5th-order Newton-Schulz iteration.
    """
    assert G.ndim == 2
    a, b, c = 3.4445, -4.7750, 2.0315
    X = G.bfloat16()
    X = X / (X.norm() + eps)
    
    transposed = False
    if X.size(0) > X.size(1):
        X = X.T
        transposed = True
        
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * (A @ A)
        X = a * X + B @ X
        
    if transposed:
        X = X.T
        
    return X.to(G.dtype)

Memory Efficiency and Parameter Footprint

Beyond faster convergence, Muon introduces substantial memory savings for distributed training frameworks:

• Optimizer State Reduction: Standard AdamW requires two FP32 state tensors per parameter (first moment $m_t$ and second moment $v_t$), requiring 8 bytes per parameter of optimizer memory. Muon eliminates the second-moment variance buffer entirely, requiring only a single momentum state (4 bytes per parameter in FP32, or 2 bytes in BF16).
• Activation and Gradient Balance: For a 7B parameter model, eliminating $v_t$ on hidden layers saves approximately 24 GB of GPU memory across the cluster, allowing larger per-device batch sizes or longer context packing without offloading.

The Hybrid Architecture: Muon with Auxiliary AdamW

Muon is mathematically designed for 2D transformation matrices. It is not directly applicable to 1D vectors, nor is it optimal for sparse or non-uniform embedding lookups. Modern LLM pre-training pipelines therefore employ a hybrid optimizer configuration:

• Muon Partition: Applied to all internal 2D weight matrices in Transformer blocks, including attention projection layers ($W_q, W_k, W_v, W_o$) and MLP/SwiGLU feed-forward projections ($W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}$).
• Auxiliary AdamW Partition: Applied to 1D vectors (RMSNorm gains, biases) and non-hidden parameters (token embedding tables and the unembedding LM head), where token frequencies vary by orders of magnitude.

In practice, Muon uses a higher base learning rate (such as 0.02 to 0.05) compared to AdamW (typically 3e-4 to 1e-3), reflecting the normalized spectral scale of the orthogonalized updates.

Empirical Results and Frontier Adoption

Empirical evaluations across open research benchmarks demonstrate consistent efficiency gains over AdamW:

• Sample Efficiency: In controlled pre-training runs, models optimized with Muon achieve the same cross-entropy validation loss with 30% to 40% fewer training tokens than baseline AdamW runs.
• Large-Scale Validation: The Moonlight-16B Mixture-of-Experts model, trained by Moonshot AI using Muon, matched the downstream benchmark performance of comparable models trained with AdamW while consuming roughly half the compute budget.
• Batch Size Scaling: Research from Essential AI and academic labs indicates that Muon exhibits superior scaling stability when increasing global batch sizes into millions of tokens, avoiding the learning rate collapse often observed in coordinate-wise optimizers.

As frontier labs scale compute-optimal training runs, optimizer architectures that respect underlying matrix geometry provide a direct algorithmic avenue for reducing pre-training FLOP requirements.

Sources

• Muon: An optimizer for hidden layers in neural networks (Keller Jordan)
• Muon is Scalable for LLM Training (arXiv:2502.16982)
• Modded NanoGPT Speedrun Repository (Keller Jordan et al.)
• Newton-Schulz Polar Decomposition Documentation (Modula Systems)
• Moonlight: Scaling Muon for Mixture-of-Experts Pre-Training (Moonshot AI)