Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon

TL;DR

Gram Newton-Schulz is a mathematically-equivalent rewrite of the Newton-Schulz polar-decomposition routine that Muon uses to orthogonalize its momentum matrix. Instead of iterating on the rectangular $n \times m$ momentum matrix $\mathbf X$ , the algorithm iterates on the small square symmetric Gram matrix $\mathbf{X X^\top} \in \mathbb{R}^{n \times n}$ , reducing FLOPs by ~55% (vs symmetric-GEMM Newton-Schulz) and ~68% (vs naive Newton-Schulz) at typical Muon aspect ratios $\alpha = m/n = 4$ . Paired with custom CuTeDSL symmetric-GEMM kernels for Hopper and Blackwell, the resulting GramMuon optimizer cuts orthogonalization wall-clock time by 40–50% in trillion-parameter MoE settings (Kimi K2 / GLM-5) while preserving training quality within 0.01 validation perplexity. A “restart” strategy stabilizes the otherwise numerically-divergent half-precision iteration.

Key claims

Standard Newton-Schulz (5 iterations) on an $n \times m$ matrix with aspect ratio $\alpha = m/n$ costs $T(3\alpha+1)n^3 = (16\alpha + 5)n^3$ FLOPs using symmetric GEMMs; the cost is dominated by ~ $3\alpha n^3$ FLOPs per iteration of rectangular matmul. Most transformer weight matrices (MLP, MoE, GQA/MLA attention projections) are rectangular ( $\alpha > 1$ ) [§“Algorithmic optimizations and Symmetric GEMMs”].
The polar decomposition can be factored as $\mathrm{polar}(\mathbf X) = (\mathbf X \mathbf X^\top)^{-1/2} \mathbf X$ ; Gram Newton-Schulz approximates $(\mathbf X \mathbf X^\top)^{-1/2}$ by an iteration that uses only $n \times n$ symmetric multiplications inside the loop, plus two rectangular matmuls (one to form $\mathbf X \mathbf X^\top$ upfront, one to apply $\mathbf Q_T \mathbf X$ at the end) [§“Gram Newton-Schulz”].
For odd polynomials $p_t(x) = x h_t(x^2)$ , the composed iteration $(p_T \circ \dots \circ p_1)(\mathbf X)$ can be computed from $\mathbf X_0$ without materializing intermediate $\mathbf X_1, \dots, \mathbf X_{T-1}$ via scalar recurrences $r_t = r_{t-1} z_t^2$ , $q_t = q_{t-1} z_t$ with $z_t = h_t(r_{t-1})$ (Theorem 1) [§“Gram Newton-Schulz”, Theorem 1].
For $T=5, \alpha=4$ : Gram Newton-Schulz saves 55% of FLOPs vs standard Newton-Schulz with symmetric GEMMs, or 68% vs a typical implementation without symmetric GEMMs [§“Gram Newton-Schulz” FLOP analysis].
Naive Gram Newton-Schulz is numerically unstable in bfloat16 due to two distinct failure modes: (1) spurious negative eigenvalues in $\mathbf R_t = \mathbf X_t \mathbf X_t^\top$ — they should be $\geq 0$ in exact arithmetic, but bfloat16 noise introduces small negatives that grow exponentially because $r_t < (15/8)^2 r_{t-1}$ ; (2) eigenvector drift accumulating across iterations [§“Stability analysis”, Figs. 4–8].
Restarting — re-initializing $\mathbf R_t \gets \mathbf X_t \mathbf X_t^\top$ and $\mathbf Q_t \gets \mathbf I$ at fixed intervals (e.g. every 2–5 iterations) — bounds the magnitude of spurious eigenvalues and controls eigenvector drift; this yields the stabilized algorithm [§“Stabilized Gram Newton-Schulz”, Figs. 9–10].
GramMuon (Muon with Gram Newton-Schulz, 5 iterations, Polar Express coefficients, restart after iteration 2) achieves “effectively identical” pretraining quality to standard Muon — within 0.01 validation perplexity — while reducing the orthogonalization wall-clock step by 40–50% on trillion-parameter MoEs [§“Stabilized Gram Newton-Schulz”, §“GramMuon”].
Custom CuTeDSL symmetric-GEMM kernels for Hopper (SM90) and Blackwell (SM100) exploit the fact that $\mathbf X \mathbf X^\top$ and $\mathbf R_t^2$ are symmetric by construction, halving the cost of computing them by only writing the lower triangle and copying to the upper [§“Symmetric GEMM kernels in CuTeDSL”].

Method

The Muon optimizer uses Newton-Schulz iteration to approximate the polar decomposition $\mathrm{polar}(\mathbf X) = \mathbf U \mathbf V^\top$ of its momentum buffer, where $\mathbf X = \mathbf U \Sigma \mathbf V^\top$ . Standard Newton-Schulz repeatedly applies $\mathbf X_{t+1} = a_t \mathbf X_t + b_t \mathbf X_t \mathbf X_t^\top \mathbf X_t + c_t (\mathbf X_t \mathbf X_t^\top)^2 \mathbf X_t$ , with three matrix multiplications per iteration on the $n \times m$ matrix — at typical Muon aspect ratios ( $\alpha = 4$ ), these rectangular matmuls dominate cost.

Gram Newton-Schulz exploits the identity $\mathrm{polar}(\mathbf X) = (\mathbf X \mathbf X^\top)^{-1/2} \mathbf X$ . The authors prove (Theorem 1) that for any sequence of odd polynomials $p_t(x) = x h_t(x^2)$ — which is the form Newton-Schulz update polynomials take — the composed iteration can be computed via scalar recurrences acting on $x^2$ rather than $x$ . Lifted to matrices, this becomes: maintain $\mathbf R_t = \mathbf X_t \mathbf X_t^\top$ (an $n \times n$ symmetric matrix) and $\mathbf Q_t$ (an $n \times n$ matrix tracking $\prod_s \mathbf Z_s$ with $\mathbf Z_s = h_s(\mathbf R_{s-1})$ ); update $\mathbf R_t = \mathbf Z_t \mathbf R_{t-1} \mathbf Z_t$ and $\mathbf Q_t = \mathbf Q_{t-1} \mathbf Z_t$ ; output $\mathbf Q_T \mathbf X_0$ . Only the upfront $\mathbf X \mathbf X^\top$ and final $\mathbf Q_T \mathbf X$ touch the rectangular dimension $m$ .

Naive lifting is numerically unstable in bfloat16: $\mathbf R_0$ has spurious negative eigenvalues from rounding, which the iteration amplifies because the update polynomial satisfies $|r_t| \lesssim (15/8)^2 |r_{t-1}|$ for negative inputs. The fix is to restart: after a few iterations, replace $\mathbf R_t \gets \mathbf X_t \mathbf X_t^\top$ and $\mathbf Q_t \gets \mathbf I$ , which truncates accumulated negative-eigenvalue magnitude. The Stabilized Gram Newton-Schulz algorithm in the post does this once (after iteration 2) over the 5-iteration Polar Express schedule.

Custom CuTeDSL kernels for symmetric matrix multiplication (Hopper SM90, Blackwell SM100) write only the lower triangle of $\mathbf X \mathbf X^\top$ , $\mathbf R_t \mathbf Z_t$ , and $\mathbf R_t^2$ before mirroring — halving compute for these operations and giving Gram Newton-Schulz a larger speedup than would come from FLOP reduction alone.

Results

FLOP reduction: For $T=5, \alpha=4$ (typical Muon weight matrices), Gram Newton-Schulz with symmetric GEMMs uses $(17 + 3 \cdot 4)n^3 = 29n^3$ FLOPs vs $(16 \cdot 4 + 5)n^3 = 69n^3$ for standard Newton-Schulz with symmetric GEMMs — a 58% reduction. Crossover with standard NS happens at $\alpha = 1$ (square matrices), where the methods are equivalent and the system falls back to standard NS + symmetric GEMMs to launch fewer kernels [§“Gram Newton-Schulz” FLOP table].
Wall-clock on real models: 40–50% reduction in the Muon orthogonalization step’s runtime on trillion-parameter MoEs in the style of Kimi K2 [§“GramMuon”]. Figure 1 (referenced) shows AdamW vs Muon optimizer-step wall-clock on Llama at increasing sizes on B300.
Training-quality preservation: GramMuon and Muon match within 0.01 validation perplexity over Llama-style pretraining runs [§“Stabilized Gram Newton-Schulz”]. The naive (unrestarted) variant exhibits catastrophic loss spikes and Inf outputs (Fig. 2).
Stability: With restart-after-2 over 5 Polar Express iterations, $\mathbf Q_t$ ‘s largest singular value remains bounded ≤ ~12 and $\mathbf X_t$ ‘s singular values stay below 1, vs unbounded divergence in the naive variant (Figs. 9–10).

Why it’s interesting

Gram Newton-Schulz is the textbook example of the IO-Aware Kernel Design template applied to optimizer-state computation rather than to attention. Like FlashAttention, Flash-KMeans, and FlashSampling, it (i) starts from an algebraic restructuring that is mathematically identical to the original, (ii) targets a specific GPU bottleneck (here: the cost of large rectangular matmuls relative to small symmetric ones on Hopper/Blackwell), and (iii) ships a hardware-aware kernel co-designed with the algorithm. The mechanism — exploit the symmetric structure of intermediate products that standard implementations were computing with general-purpose GEMMs — is the same lever FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling pulls for backward attention via 2-CTA MMA. The IO-aware-kernel concept page now has an example from a fourth domain (after attention, k-means, and LM-head sampling); the open question on the concept page about “where else does the template apply” gets a direct datapoint.

It’s also a concrete counterpoint to the optimizer-side stability narrative on Training stability at scale. SSO (Controlled LLM Training on Spectral Sphere) wants to change what Muon’s Newton-Schulz computes; Gram Newton-Schulz wants to compute exactly the same thing, just faster. The two papers occupy orthogonal axes — quality (SSO) vs. wall-clock (GramMuon) — and could in principle compose. The “free lunch” framing here is rare and credible because the output is bit-for-bit identical up to floating-point error, validated by perplexity-matching at scale.

Finally: the failure-mode analysis (spurious negative eigenvalues amplified to $\infty$ , eigenvector drift in half precision) is a model of how to write up a numerical-stability investigation — diagnose with synthetic matrices in float64 vs bfloat16, plot the spectrum evolution, identify the polynomial growth factor exactly, propose a minimal fix (restart). Generalizable scaffolding for anyone debugging a half-precision matrix iteration.