The Muon Optimizer: Mathematical Foundations, Newton-Schulz Polar Iterations, and Spectral Steepest Descent in LLM Pre-Training

Transformer pre-training has long relied on AdamW as the default first-order optimizer. While AdamW provides coordinate-wise scale invariance through running estimates of first and second gradient moments, it treats every weight parameter as an isolated scalar. For modern neural networks composed primarily of two-dimensional linear projections, attention transformations, and feed-forward weight matrices, this coordinate-wise factorization neglects the structural transformation properties of matrix-vector mappings.

The Muon optimizer (Momentum Orthogonalized by Newton-Schulz), developed by Keller Jordan, Jeremy Bernstein, Laker Newhouse, and collaborators, addresses this limitation by formulating parameter updates as steepest descent under the matrix operator norm. Rather than computing expensive singular value decompositions or maintaining memory-heavy Kronecker covariance statistics, Muon computes orthogonalized momentum updates directly through iterative quintic Newton-Schulz polynomial iterations.

The Geometric Limitation of Coordinate-Wise Optimizers

Standard stochastic gradient descent minimizes loss along Euclidean contours, updating parameters along the negative gradient vector $\Delta W = -\eta \nabla \mathcal{L}(W)$. In deep linear and multi-head attention layers, weight matrices $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ act as linear operators on activation spaces. Euclidean gradient updates implicitly assume an isotropic Frobenius norm geometry $\|W\|_F = \sqrt{\sum_{i,j} W_{i,j}^2}$.

When training deep networks, the gradient matrix $G = \nabla \mathcal{L}(W)$ exhibits a highly ill-conditioned singular value spectrum. A small number of dominant singular vectors absorb most of the gradient energy, while dozens or hundreds of trailing directions receive negligible updates.

Adaptive optimizers like AdamW attempt to counteract gradient variance by normalizing each entry by the square root of its uncentered second moment:

$$W_{t} = W_{t-1} - \eta \frac{m_t}{\sqrt{v_t} + \epsilon}$$

Because the division is performed element-wise along coordinate axes, AdamW is sensitive to coordinate rotations and cannot equalize the rate of learning across orthogonal singular modes. Second-order methods such as K-FAC and Shampoo address this by computing Kronecker-factored covariance approximations $(G G^T)^{-1/4}$ and $(G^T G)^{-1/4}$. However, these algorithms require auxiliary memory buffers for large covariance matrices and periodic inversion of matrix roots, imposing non-trivial compute and communication overhead in distributed training.

Spectral Steepest Descent and the Polar Decomposition

Muon reconsiders the update step from first principles by solving for the optimal update direction under the spectral (operator) norm.

Let $M \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$ denote the classical Polyak momentum accumulator:

$$M_t = \mu M_{t-1} + \nabla \mathcal{L}(W_{t-1})$$

The goal is to find a parameter perturbation $\Delta W$ that minimizes the first-order Taylor approximation of the loss $\langle M, \Delta W \rangle = \text{Tr}(M^T \Delta W)$ subject to a constraint on the maximum spectral amplification of the update:

$$\Delta W^* = \arg\min_{\|\Delta W\|_2 \leq 1} \langle M, \Delta W \rangle$$

Here, $\|\Delta W\|_2 = \sigma_{\max}(\Delta W)$ denotes the spectral norm (the largest singular value).

Using the Singular Value Decomposition (SVD) of the momentum matrix, $M = U \Sigma V^T$, where $U \in \mathbb{R}^{d_{\text{out}} \times k}$ and $V \in \mathbb{R}^{d_{\text{in}} \times k}$ have orthonormal columns and $\Sigma = \text{diag}(\sigma_1, \sigma_2, \dots, \sigma_k)$ with $\sigma_i \geq 0$:

$$\langle M, \Delta W \rangle = \text{Tr}((U \Sigma V^T)^T \Delta W) = \text{Tr}(\Sigma U^T \Delta W V)$$

To minimize this trace inner product under the spectral constraint $\|\Delta W\|_2 \leq 1$, the optimal update aligns perfectly with the singular subspaces $U$ and $V$, setting every singular value to the maximum allowable boundary of 1:

$$\Delta W^* = -U V^T$$

The matrix $O = U V^T$ is the polar factor (or matrix sign factor) of $M$. Replacing $\Sigma$ with the identity matrix equalizes the update magnitude across all singular directions. Dominant directions with massive singular values and weakly represented directions with small singular values are updated at the exact same spectral rate, preventing representation collapse and eliminating gradient scale disparity.

Iterative Orthogonalization via Newton-Schulz Iterations

Computing the exact SVD of every 2D weight matrix at every optimization step is computationally impractical on modern GPU accelerators. Standard SVD algorithms require iterative QR or bidiagonal divide-and-conquer routines that trigger frequent host-device synchronizations and cannot leverage tensor contraction cores.

Muon bypasses SVD by computing the polar factor through polynomial iterations. The Newton-Schulz iteration is an iterative matrix method designed to compute the matrix sign function and polar decomposition using purely matrix-matrix multiplications.

Given an initial matrix $X_0$ normalized such that its spectral norm is bounded strictly below $\sqrt{3}$:

$$X_0 = \frac{M}{\|M\|_F + \epsilon}$$

The classical third-order Newton-Schulz iteration updates $X$ via:

$$X_{k+1} = \frac{1}{2} X_k (3 I - X_k^T X_k)$$

To achieve faster numerical convergence within fewer steps, the Muon implementation uses a tuned fifth-order (quintic) polynomial iteration:

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

For rectangular matrices where $d_{\text{out}} > d_{\text{in}}$, the transpose order is inverted to operate on the smaller inner dimension $X_k^T X_k$, minimizing intermediate FLOPs.

Using polynomial optimization over the singular value interval $[0, 1]$, the coefficients are set to:

$$a = 3.4445, \quad b = -4.7750, \quad c = 2.0315$$

These specific coefficients maximize the rate at which small singular values are inflated toward 1 in early iterations while maintaining strict numerical stability. Running exactly $N = 5$ iterations produces an approximately orthogonal matrix $O \approx U V^T$.

Because each iteration consists solely of GEMM operations (general matrix multiplies), the entire orthogonalization executes directly in bfloat16 precision on GPU Tensor Cores with virtually zero kernel launch overhead.

Operator Norm Scaling and Width Transfer

To maintain consistent gradient dynamics when scaling model width (hidden dimension $d$), Muon scales updates using the root-mean-square (RMS) operator norm:

$$\|W\|_{\text{RMS} \to \text{RMS}} = \max_{x \neq 0} \frac{\|W x\|_{\text{RMS}}}{\|x\|_{\text{RMS}}} = \sqrt{\frac{d_{\text{in}}}{d_{\text{out}}}} \|W\|_2$$

When applying the orthogonalized update $O_t$, Muon scales the learning rate by the aspect ratio of the layer:

$$W_t = W_{t-1} - \eta_t \cdot \sqrt{\max(1, d_{\text{out}} / d_{\text{in}})} \cdot O_t$$

This scaling formulation ensures that the expected RMS change in activations remains constant regardless of the model's width $d_{\text{model}}$ or intermediate expansion factor $d_{\text{ffn}}$. Consequently, optimal learning rate hyperparameters tuned on small prototype models transfer directly to multi-billion parameter configurations without retuning.

Memory Footprint and Hybridization

Muon is designed specifically for two-dimensional weight tensors in linear, projection, and attention layers. It is not suitable for one-dimensional parameters (layer normalization gains, biases) or high-cardinality embedding tables where rows are updated sparsely.

Production implementations employ a hybrid optimization topology:

1. 2D Hidden Weights (Muon): Multi-head attention projections ($W_q, W_k, W_v, W_o$) and feed-forward layers ($W_{\text{gate}}, W_{\text{up}}, W_{\text{down}}$) are updated via Muon.
2. 1D and Embedding Parameters (AdamW): Token embeddings, positional embeddings, normalization gains, and final classification heads are updated via standard AdamW.

This division yields substantial memory savings over pure AdamW baselines. AdamW maintains two full-precision float32 state tensors per parameter (first moment $m_t$ and second moment $v_t$), requiring 8 bytes of optimizer state per parameter. Muon maintains only a single momentum buffer $M_t$ (which can be stored in bfloat16 or float32), reducing optimizer state memory on 2D weights by 50 percent.

Pre-Training Efficiency and Convergence Dynamics

Empirical evaluations across language modeling benchmarks demonstrate that Muon achieves the same pre-training validation perplexity with roughly 1.3x to 1.5x fewer training tokens compared to well-tuned AdamW baselines.

In the community NanoGPT speedrun benchmarks (training a 124M parameter transformer on FineWeb-Edu to a validation loss of 3.28), replacing AdamW with Muon reduced total GPU wall-clock training time from several minutes to under 90 seconds on an 8x H100 node. The computational overhead of the five Newton-Schulz matrix multiplications accounts for less than 2 percent of total step execution time, which is offset by the faster loss decay per optimization step.

By treating neural network weights as geometric linear operators rather than independent coordinate collections, Muon establishes a practical middle ground: capturing the spectral benefits of second-order optimization while preserving the computational simplicity and speed of first-order methods.

Sources

• Muon: An optimizer for hidden layers in neural networks (Keller Jordan)
• Deriving Muon (Jeremy Bernstein)
• Understanding Muon (Laker Newhouse)
• Muon is Scalable for LLM Training (arXiv:2502.16982)
• The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm (arXiv:2505.16932)
• Old Optimizer, New Norm: An Anthology (arXiv:2409.20325)
• Modded NanoGPT (Keller Jordan et al.)
• Functions of Matrices: Theory and Computation (Nicholas J. Higham)