Dion3: Full-Stack Orthogonal Updates

# Dion3: Full-stack orthogonal updates
Source: [https://arxiv.org/html/2608.11612](https://arxiv.org/html/2608.11612)
Jack ZhangPrinceton UniversityKwangjun Ahn11footnotemark:1NVIDIAAli NaeimiIndependent ResearcherAustin Feng11footnotemark:1Yale UniversityBerlin ChenPrinceton UniversityTri DaoPrinceton UniversityJohn LangfordMicrosoft Research

###### Abstract

The Muon optimizer incurs a significant overhead cost due to its cubic\-time Newton\-Schulz orthogonalization step\. When weights are sharded, communication overhead compounds this computational cost, eroding the benefits of Muon in many settings\. We present Dion3, a revision of Muon that targets this overhead at every level of the stack\. Our Gram Newton\-Schulz algorithm reduces the FLOP cost of orthogonalization, our CuteDSL kernels accelerate it by exploiting symmetry, and our megabatching strategy reduces communication overhead\. Moreover, we propose a simple change to the update rule that cuts costs even further: selecting only a fraction of the momentum matrix’s rows to orthogonalize at each step\. This update rule improves on Dion \(another “compressed” version of Muon\), in both speed and performance\[[2](https://arxiv.org/html/2608.11612#bib.bib5)\]\. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to6×6\\times\. Dion3 is available via thedionpackage111[https://github\.com/microsoft/dion](https://github.com/microsoft/dion)as a drop\-in replacement for Muon\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x1.png)Figure 1:Optimizer step time of Muon \(excluding forward/backward pass\) relative to AdamW for a 7B\-parameter language model trained on four GH200s\. Each colored bar adds one of our contributions on top of the previous one\. Together they reduce Muon’s cost from26×26\\timesAdamW to just4×4\\times\.###### Contents

1. [1Introduction](https://arxiv.org/html/2608.11612#S1)
2. [2The Challenge of Scaling Muon](https://arxiv.org/html/2608.11612#S2)1. [2\.1Muon and NorMuon Recap](https://arxiv.org/html/2608.11612#S2.SS1) 2. [2\.2Standard Newton\-Schulz](https://arxiv.org/html/2608.11612#S2.SS2) 3. [2\.3Scalability of Muon](https://arxiv.org/html/2608.11612#S2.SS3)
3. [3Comparison with Related Work](https://arxiv.org/html/2608.11612#S3)1. [3\.1Improving Newton\-Schulz](https://arxiv.org/html/2608.11612#S3.SS1) 2. [3\.2Orthogonalizing a Smaller Matrix](https://arxiv.org/html/2608.11612#S3.SS2)
4. [4Gram Newton\-Schulz](https://arxiv.org/html/2608.11612#S4)1. [4\.1Runtime of Naive Gram Newton\-Schulz](https://arxiv.org/html/2608.11612#S4.SS1) 2. [4\.2Stabilizing Gram Newton\-Schulz](https://arxiv.org/html/2608.11612#S4.SS2)
5. [5Symmetric GEMM Kernels in CuteDSL](https://arxiv.org/html/2608.11612#S5)
6. [6The Dion3 Update Rule](https://arxiv.org/html/2608.11612#S6)
7. [7Megabatching and Communication](https://arxiv.org/html/2608.11612#S7)
8. [8Experiments](https://arxiv.org/html/2608.11612#S8)1. [8\.1Model Quality Is Preserved](https://arxiv.org/html/2608.11612#S8.SS1) 2. [8\.2Dion3 Accelerates the Optimizer](https://arxiv.org/html/2608.11612#S8.SS2)
9. [9Conclusion](https://arxiv.org/html/2608.11612#S9)
10. [References](https://arxiv.org/html/2608.11612#bib)
11. [AAlternative Experiments on Gram Newton\-Schulz](https://arxiv.org/html/2608.11612#A1)1. [A\.1Setup](https://arxiv.org/html/2608.11612#A1.SS1) 2. [A\.2Kernelized Gram Newton\-Schulz preserves quality](https://arxiv.org/html/2608.11612#A1.SS2) 3. [A\.3Kernelized Gram Newton\-Schulz speeds up the optimizer step](https://arxiv.org/html/2608.11612#A1.SS3)
12. [BStability of Gram Newton\-Schulz](https://arxiv.org/html/2608.11612#A2)1. [B\.1Instability of Naive Gram Newton\-Schulz](https://arxiv.org/html/2608.11612#A2.SS1) 2. [B\.2Stabilizing Gram Newton\-Schulz by Restarting](https://arxiv.org/html/2608.11612#A2.SS2)
13. [CKernel Implementation Details](https://arxiv.org/html/2608.11612#A3)1. [C\.1Symmetric GEMM Kernel Details](https://arxiv.org/html/2608.11612#A3.SS1) 2. [C\.2Implementation Strategy in Code](https://arxiv.org/html/2608.11612#A3.SS2) 3. [C\.3Kernel Optimizations for Standard Newton\-Schulz](https://arxiv.org/html/2608.11612#A3.SS3)
14. [DAdditional Experiments](https://arxiv.org/html/2608.11612#A4)1. [D\.1Architecture and Optimization Setup](https://arxiv.org/html/2608.11612#A4.SS1) 2. [D\.2Finer\-Grained Timing Metrics](https://arxiv.org/html/2608.11612#A4.SS2) 3. [D\.3Benchmarking all\-to\-all communications](https://arxiv.org/html/2608.11612#A4.SS3) 4. [D\.4Ablations](https://arxiv.org/html/2608.11612#A4.SS4)
15. [ECase Studies of End\-to\-End Training Time](https://arxiv.org/html/2608.11612#A5)

## 1Introduction

Muon is becoming the method of choice for training frontier LLMs like Kimi K2 and GLM\-5\[[21](https://arxiv.org/html/2608.11612#bib.bib8),[13](https://arxiv.org/html/2608.11612#bib.bib33)\]\. Compared to AdamW, Muon needs fewer optimizer steps to reach a given loss, but each step is more expensive\. This overhead is due to Muon’s Newton\-Schulz orthogonalization procedure, a cubic\-time matrix operation not present in older optimizers\. As model size increases, the overhead of computing each Muon step grows rapidly\. Moreover, Muon is more difficult than traditional optimizers to parallelize, introducing additional communication overhead in distributed settings\. Thus, scaling Muon presents a range of algorithmic and systems\-level challenges that diminish its effectiveness in the most demanding settings and hinder its adoption to new problems\.

This paper gives a comprehensive answer to these challenges\. We provide a full\-stack solution that performs well across a wide range of model sizes, architectures, cluster sizes, and parallelism strategies\. We package our contributions into a single optimizer, Dion3, built on four improvements that each reduce the cost of Muon’s orthogonalization step and compound when used together:

1. 1\.Gram Newton\-Schulz, a mathematically equivalent reformulation of Newton\-Schulz that iterates on the small symmetric Gram matrix and cuts the FLOP cost of each orthogonalization dramatically \([Section˜4](https://arxiv.org/html/2608.11612#S4)\)\.
2. 2\.CustomGPU kernels for symmetric matrix multiplicationwritten in CuteDSL, which speed up Gram Newton\-Schulz and take maximal advantage of its symmetric structure \([Section˜5](https://arxiv.org/html/2608.11612#S5)\)\.
3. 3\.A newoptimizer update rulethat subsamples the rows or columns of the momentum matrix before orthogonalizing it to make each step faster \([Section˜6](https://arxiv.org/html/2608.11612#S6)\)\. Our update rule is simpler and faster than Dion\[[2](https://arxiv.org/html/2608.11612#bib.bib5)\]while matching or improving the optimization quality of Muon and NorMuon\.
4. 4\.Megabatched communication, which eliminates a significant source of overhead by reducing the number of rounds of communication per optimizer step to a small constant\.

Our improvements combine to accelerate the optimizer step at every level of the stack\. The GPU kernels at the lowest level speed up Gram Newton\-Schulz, which in turn reduces the FLOP cost of executing our optimizer’s update rule\. At the top level, the compression achieved by our update rule joins with the megabatch strategy to limit the communication cost\. Each of our contributions is of independent value, but combining them gives Dion3 maximum flexibility to handle different parallelism strategies, architectures, and model sizes efficiently\. Moreover, Dion3 benefits from synergies between them that further reduce computational cost\. Gram Newton\-Schulz uses more symmetric multiplications than standard Newton\-Schulz does, compounding the benefits of our CuteDSL kernels\. Likewise, Gram Newton\-Schulz is especially fast for matrices with highly asymmetric shapes—like those produced by our update rule’s subsampling strategy\. Surprisingly, Dion3 appears to benefit the update*quality*as well as the computational footprint in some settings, even with a strong optimized baseline\.

To realize these gains in practice, we provide open\-source implementations in the form of two interoperable packages:gram\-newton\-schulz222[https://github\.com/Dao\-AILab/gram\-newton\-schulz](https://github.com/Dao-AILab/gram-newton-schulz)is a drop\-in replacement for Muon’s Newton\-Schulz routine that also incorporates our CuteDSL kernels, while thedionpackage implements our optimizer update rule \(along with baselines like Muon and NorMuon\), and handles parallelism in the distributed setting \(megabatching, FSDP, DDP, etc\.\)\. In all, Dion3 allows practitioners working in a wide range of settings to realize the benefits of Muon and related optimizers with almost no overhead\.

## 2The Challenge of Scaling Muon

### 2\.1Muon and NorMuon Recap

The Muon optimizer\[[19](https://arxiv.org/html/2608.11612#bib.bib1)\]is best described as steepest\-direction descent with respect to the spectral norm\[[5](https://arxiv.org/html/2608.11612#bib.bib2)\]\. At a given training step, let𝑾∈ℝn×m\{\\bm\{W\}\}\\in\\mathbb\{R\}^\{n\\times m\}be a weight matrix and let𝑮\{\\bm\{G\}\}be the gradient of the loss with respect to𝑾\{\\bm\{W\}\}\. The Muon update rule is

𝑴←μ​𝑴\+𝑮𝑾←𝑾−η​polar⁡\(𝑴\)\\displaystyle\\begin\{split\}\{\\bm\{M\}\}&\\leftarrow\\mu\{\\bm\{M\}\}\+\{\\bm\{G\}\}\\\\ \{\\bm\{W\}\}&\\leftarrow\{\\bm\{W\}\}\-\\eta\\operatorname\{polar\}\(\{\\bm\{M\}\}\)\\end\{split\}\(1\)whereμ\\muis the momentum coefficient,η\\etais the learning rate, and𝑴\{\\bm\{M\}\}is the momentum matrix \(with𝑴0:=𝟎\{\\bm\{M\}\}\_\{0\}:=\\bm\{0\}\)\. In many ways, Muon resembles basic stochastic gradient descent \(SGD\) with momentum\. Its key innovation is thepolar\\operatorname\{polar\}operation, which is defined as follows:

###### Definition 1\(Polar Decomposition\)\.

If𝐗=𝐔​𝚺​𝐕⊤\{\\bm\{X\}\}=\{\\bm\{U\}\}\\mathbf\{\\Sigma\}\{\\bm\{V\}\}^\{\\top\}is the singular value decomposition \(SVD\) of a matrix, thenpolar⁡\(𝐗\)=𝐔​𝐕⊤\\operatorname\{polar\}\(\{\\bm\{X\}\}\)=\{\\bm\{U\}\}\{\\bm\{V\}\}^\{\\top\}\.

This operation—the “orthogonalization step”—balances the spectrum of the update, ensuring that it has full numerical rank\.

NorMuon\[[23](https://arxiv.org/html/2608.11612#bib.bib12)\]is a popular variant of Muon that aims to combine its strengths with those of Adam\. While Muon’s update has a balanced spectrum, it does not have balanced row norms; each step could greatly affect some neurons while almost neglecting others\. Following Adam, NorMuon uses the second moment of the \(orthogonalized\) gradients to adaptively adjust the stepsize for each neuron, helping all neurons to learn at every step\. A final rescaling ensures that the overall magnitude of the update \(as measured in the Frobenius norm\) matches that of Muon:

𝑴←μ​𝑴\+𝑮,𝑶←polar⁡\(𝑴\)𝒗i←β2​𝒗i\+\(1−β2\)⋅1m​∑j𝑶i​j2,𝑶^i​j←𝑶i​j𝒗i\+ϵ𝑾←𝑾−η​‖𝑶‖𝖥‖𝑶^‖𝖥​𝑶^\\displaystyle\\begin\{split\}\{\\bm\{M\}\}&\\leftarrow\\mu\{\\bm\{M\}\}\+\{\\bm\{G\}\},\\qquad\{\\bm\{O\}\}\\leftarrow\\operatorname\{polar\}\(\{\\bm\{M\}\}\)\\\\ \{\\bm\{v\}\}\_\{i\}&\\leftarrow\\beta\_\{2\}\{\\bm\{v\}\}\_\{i\}\+\(1\-\\beta\_\{2\}\)\\cdot\\frac\{1\}\{m\}\\sum\\nolimits\_\{j\}\{\\bm\{O\}\}\_\{ij\}^\{2\},\\qquad\\widehat\{\{\\bm\{O\}\}\}\_\{ij\}\\leftarrow\\frac\{\{\\bm\{O\}\}\_\{ij\}\}\{\\sqrt\{\{\\bm\{v\}\}\_\{i\}\}\+\\epsilon\}\\\\ \{\\bm\{W\}\}&\\leftarrow\{\\bm\{W\}\}\-\\eta\\frac\{\\\|\{\\bm\{O\}\}\\\|\_\{\\mathsf\{F\}\}\}\{\\\|\\widehat\{\{\\bm\{O\}\}\}\\\|\_\{\\mathsf\{F\}\}\}\\widehat\{\{\\bm\{O\}\}\}\\end\{split\}\(2\)The cost of NorMuon’s extra normalization steps is negligible compared to that of computingpolar⁡\(𝑴\)\\operatorname\{polar\}\(\{\\bm\{M\}\}\), so NorMuon’s benefits come almost for free\. Our methods, which speed up the computation ofpolar⁡\(𝑴\)\\operatorname\{polar\}\(\{\\bm\{M\}\}\), apply to both Muon and NorMuon alike\.

### 2\.2Standard Newton\-Schulz

The main challenge to scaling Muon is the need to computepolar⁡\(𝑴\)\\operatorname\{polar\}\(\{\\bm\{M\}\}\)for each weight matrix at every step\. Sincepolar⁡\(⋅\)\\operatorname\{polar\}\(\\cdot\)is expensive to compute exactly, Muon uses the Newton\-Schulz method to approximate it\. Newton\-Schulz is an iterative method based on matrix polynomials\. Beginning with𝑿0\{\\bm\{X\}\}\_\{0\}, each iteration improves the approximation𝑿t≈polar⁡\(𝑿0\)\{\\bm\{X\}\}\_\{t\}\\approx\\operatorname\{polar\}\(\{\\bm\{X\}\}\_\{0\}\)according to the update rule

𝑿t\+1=at​𝑿t\+bt​𝑿t​𝑿t⊤​𝑿t\+ct​\(𝑿t​𝑿t⊤\)2​𝑿t\.\{\\bm\{X\}\}\_\{t\+1\}=a\_\{t\}\{\\bm\{X\}\}\_\{t\}\+b\_\{t\}\{\\bm\{X\}\}\_\{t\}\{\\bm\{X\}\}\_\{t\}^\{\\top\}\{\\bm\{X\}\}\_\{t\}\+c\_\{t\}\\left\(\{\\bm\{X\}\}\_\{t\}\{\\bm\{X\}\}\_\{t\}^\{\\top\}\\right\)^\{2\}\{\\bm\{X\}\}\_\{t\}\.We can interpret Newton\-Schulz by understanding how it affects the singular value decomposition\. Let𝑿0=𝑼​𝚺​𝑽⊤\{\\bm\{X\}\}\_\{0\}=\{\\bm\{U\}\}\\mathbf\{\\Sigma\}\{\\bm\{V\}\}^\{\\top\}be the SVD, where𝑼⊤​𝑼=𝑽⊤​𝑽=𝑰\{\\bm\{U\}\}^\{\\top\}\{\\bm\{U\}\}=\{\\bm\{V\}\}^\{\\top\}\{\\bm\{V\}\}=\{\\bm\{I\}\}and𝚺\\mathbf\{\\Sigma\}is diagonal with positive entries called the singular values\. A direct computation shows

𝑿1=𝑼​\(a1​𝚺\+b1​𝚺3\+c1​𝚺5\)​𝑽⊤=𝑼​p1​\(𝚺\)​𝑽⊤\{\\bm\{X\}\}\_\{1\}=\{\\bm\{U\}\}\\left\(a\_\{1\}\\mathbf\{\\Sigma\}\+b\_\{1\}\\mathbf\{\\Sigma\}^\{3\}\+c\_\{1\}\\mathbf\{\\Sigma\}^\{5\}\\right\)\{\\bm\{V\}\}^\{\\top\}=\{\\bm\{U\}\}p\_\{1\}\(\\mathbf\{\\Sigma\}\)\{\\bm\{V\}\}^\{\\top\}wherep1​\(x\):=a1​x\+b1​x3\+c1​x5p\_\{1\}\(x\):=a\_\{1\}x\+b\_\{1\}x^\{3\}\+c\_\{1\}x^\{5\}\. Since𝑼\{\\bm\{U\}\}and𝑽\{\\bm\{V\}\}have orthonormal columns andp1​\(𝚺\)p\_\{1\}\(\\mathbf\{\\Sigma\}\)is diagonal, the right\-hand side of this equation must be the SVD of𝑿1\{\\bm\{X\}\}\_\{1\}\. By extension,𝑿T\{\\bm\{X\}\}\_\{T\}also has the same singular vectors𝑼\{\\bm\{U\}\}and𝑽\{\\bm\{V\}\}as𝑿0\{\\bm\{X\}\}\_\{0\}, and its singular values have been transformed according to the composition of polynomials\(pT∘⋯∘p1\)​\(𝚺\)\(p\_\{T\}\\circ\\cdots\\circ p\_\{1\}\)\(\\mathbf\{\\Sigma\}\)\. If we normalize the input matrix𝑿0=𝑿/‖𝑿‖𝖥\{\\bm\{X\}\}\_\{0\}=\{\\bm\{X\}\}/\\\|\{\\bm\{X\}\}\\\|\_\{\\mathsf\{F\}\}, then all singular values of𝑿0\{\\bm\{X\}\}\_\{0\}lie in\[0,1\]\[0,1\]\.\[[19](https://arxiv.org/html/2608.11612#bib.bib1)\]identified a sequence of degree\-5 odd polynomials for which\(pT∘⋯∘p1\)​\(x\)≈1\(p\_\{T\}\\circ\\cdots\\circ p\_\{1\}\)\(x\)\\approx 1on\[0,1\]\[0,1\]\. Therefore,

𝑿T=𝑼\(pT∘⋯∘p1\)\(𝚺\)𝑽⊤≈𝑼𝑽⊤=:polar\(𝑿0\)\{\\bm\{X\}\}\_\{T\}=\{\\bm\{U\}\}\(p\_\{T\}\\circ\\cdots\\circ p\_\{1\}\)\(\\mathbf\{\\Sigma\}\)\{\\bm\{V\}\}^\{\\top\}\\approx\{\\bm\{U\}\}\{\\bm\{V\}\}^\{\\top\}=:\\operatorname\{polar\}\(\{\\bm\{X\}\}\_\{0\}\)
Algorithm 1Standard Newton\-Schulz1:

𝑿∈ℝn×m\{\\bm\{X\}\}\\in\\mathbb\{R\}^\{n\\times m\}with

n≤mn\\leq m, coefficients

\{\(at,bt,ct\)\}t=15\\\{\(a\_\{t\},b\_\{t\},c\_\{t\}\)\\\}\_\{t=1\}^\{5\}
2:

𝑿←𝑿/\(‖𝑿‖𝖥\+ϵ\)\{\\bm\{X\}\}\\leftarrow\{\\bm\{X\}\}\\,/\\,\(\\\|\{\\bm\{X\}\}\\\|\_\{\\mathsf\{F\}\}\+\\epsilon\)⊳\\trianglerightNormalize sing\. vals\. to\[0,1\]\[0,1\];ϵ=10−7\\epsilon=10^\{\-7\}

3:

𝑿←bfloat16​\(𝑿\)\{\\bm\{X\}\}\\leftarrow\\texttt\{bfloat16\}\(\{\\bm\{X\}\}\)⊳\\trianglerightCast to half precision for speed

4:for

t=1,…,5t=1,\\ldots,5do⊳\\trianglerightApplypt​\(𝑿\)p\_\{t\}\(\{\\bm\{X\}\}\)

5:

𝑨←𝑿​𝑿⊤\{\\bm\{A\}\}\\leftarrow\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}
6:

𝑩←bt​𝑨\+ct​𝑨2\{\\bm\{B\}\}\\leftarrow b\_\{t\}\{\\bm\{A\}\}\+c\_\{t\}\{\\bm\{A\}\}^\{2\}
7:

𝑿←at​𝑿\+𝑩​𝑿\{\\bm\{X\}\}\\leftarrow a\_\{t\}\{\\bm\{X\}\}\+\{\\bm\{B\}\}\{\\bm\{X\}\}
8:endfor

9:return

𝑿\{\\bm\{X\}\}

[Algorithm˜1](https://arxiv.org/html/2608.11612#alg1)gives the standard implementation of Newton\-Schulz\. We now analyze its runtime in FLOPs to help us understand its performance bottlenecks\. We count only the cubic\-time matrix multiplication operations, ignoring the lower\-order scalar multiplications and matrix additions\. For clarity, we letTTdenote the number of iterations, remembering that Muon typically setsT=5T=5\. We also assume without loss of generality thatn≤mn\\leq mand define the aspect ratioα=m/n≥1\\alpha=m/n\\geq 1\. Intuitively,α\\alphameasures how asymmetric the shape of the matrix is, withα=1\\alpha=1being square andα≫1\\alpha\\gg 1being very asymmetric\.

Each iteration has three steps\. Each step contains a single matrix multiplication \(𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\},𝑨2\{\\bm\{A\}\}^\{2\},𝑩​𝑿\{\\bm\{B\}\}\{\\bm\{X\}\}\) costing, respectively,2​m​n22mn^\{2\},2​n32n^\{3\}, and2​m​n22mn^\{2\}FLOPs for a total cost ofT​\(4​m​n2\+2​n3\)=2​T​\(2​α\+1\)​n3T\(4mn^\{2\}\+2n^\{3\}\)=2T\(2\\alpha\+1\)n^\{3\}FLOPs\. WhenT=5T=5, the cost is\(20​α\+10\)​n3\(20\\alpha\+10\)n^\{3\}spread across 15 matrix\-matrix multiplications \(GEMMs\)\.

### 2\.3Scalability of Muon

While Muon comfortably outperforms traditional optimizers like Adam on a per\-step basis, two factors immediately make it more difficult to scale:

- •Super\-linear complexity\.SGD and Adam perform only inexpensive element\-wise operations like addition and scalar multiplication; therefore, the computational complexity of these methods scales linearly with the number of parameters\. Orthogonalization is substantially more expensive; for ann×nn\\times nweight matrix, it requiresO​\(n3\)O\(n^\{3\}\)time\. Although sparse MoE architectures keep most weight matrices smaller, they also reduce overall model FLOPs, increasing the relative cost of the optimizer step\[[11](https://arxiv.org/html/2608.11612#bib.bib3)\]compared to the forward and backward passes\. Moreover, MoEs still include large weight matrices in their dense layers\[[9](https://arxiv.org/html/2608.11612#bib.bib4)\]\.
- •Distributed training\.There are additional challenges when weights are sharded across GPUs, as is standard in distributed training\. Element\-wise optimizers like Adam can update each shard separately, but Muon must gather the shards together to computepolar\\operatorname\{polar\}on the full matrix\. Early implementations of Muon duplicated the orthogonalization step across devices, but this increases its already considerable cost\.\[[11](https://arxiv.org/html/2608.11612#bib.bib3)\]propose using all\-to\-all communication along the sharding dimension, with each device processing a different weight in parallel\.\[[2](https://arxiv.org/html/2608.11612#bib.bib5)\]implemented this strategy in PyTorch FSDP2, and\[[24](https://arxiv.org/html/2608.11612#bib.bib6)\]examined its compute\-communication overlap characteristics\. While this approach makes the overhead of distributed Newton\-Schulz manageable in some settings, it does not fully resolve the scalability limitations\.

In[Section˜6](https://arxiv.org/html/2608.11612#S6), we introduce a variant of Muon that addresses these two challenges by subselecting \(that is, compressing\) the momentum matrix before orthogonalizing it\.[Section˜7](https://arxiv.org/html/2608.11612#S7)describes a megabatching strategy and a flexible, pip\-installable package that gracefully handle the distributed setting\. Our other contributions are inspired by the analysis in the previous section, which reveals two shortcomings of the standard Newton\-Schulz algorithm:

- •Symmetric Matrix Multiplication\.The matrices𝑨=𝑿​𝑿⊤\{\\bm\{A\}\}=\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}and𝑩=bt​𝑨\+ct​𝑨2\{\\bm\{B\}\}=b\_\{t\}\{\\bm\{A\}\}\+c\_\{t\}\{\\bm\{A\}\}^\{2\}computed at each iteration of Newton\-Schulz are symmetric by definition, but standard Newton\-Schulz does not exploit this structure\. We can compute the lower triangular part of these matrices in the usual way and simply copy the results to the upper triangular part\. This technique halves the cost of computing𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}and𝑨2\{\\bm\{A\}\}^\{2\}, giving an overall total ofT​\(3​α\+1\)​n3T\(3\\alpha\+1\)n^\{3\}FLOPs\.[Section˜5](https://arxiv.org/html/2608.11612#S5)describes custom CuteDSL kernels that implement this technique\.
- •Dependence on Aspect Ratio\.Newton\-Schulz’s runtime is dominated by the large rectangular matrix multiplications needed to compute𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}and𝑩​𝑿\{\\bm\{B\}\}\{\\bm\{X\}\}, which together cost3​α​n33\\alpha n^\{3\}FLOPs per iteration even when using symmetric matrix multiplications\. A typical implementation withT=5T=5requires 10 of these expensive rectangular multiplications\. This strong dependence onα\\alphais unfortunate\. Most of the weight matrices in transformer architectures are rectangular, including the MLP weights, MoE weights, and attention projection weights when using GQA or MLA\. Furthermore, we observe that the latest MoE architectures are trending towards finer\-grained, sparser experts, meaning that the aspect ratios of their hidden dimensions to intermediate dimensions are increasing as well\[[21](https://arxiv.org/html/2608.11612#bib.bib8),[16](https://arxiv.org/html/2608.11612#bib.bib13),[34](https://arxiv.org/html/2608.11612#bib.bib14),[31](https://arxiv.org/html/2608.11612#bib.bib15)\]\. Thus, at large scales, pretraining time would benefit greatly from an algorithm that uses fewer rectangular multiplications and more small symmetric ones\. We develop such an algorithm in[Section˜4](https://arxiv.org/html/2608.11612#S4)\.

#### 2\.3\.1Kimi’s Success

Muon has been scaled successfully by Moonshot AI\[[26](https://arxiv.org/html/2608.11612#bib.bib7),[21](https://arxiv.org/html/2608.11612#bib.bib8)\]\. Their success was enabled by an alignment of several factors:

1. 1\.Earlier releases of PyTorch and Megatron\-LM used DP\-sharding strategies for optimizer states that were, by chance, favorable for Muon\[[27](https://arxiv.org/html/2608.11612#bib.bib9)\]\. Model and optimizer states were stored in large contiguous flat buffers, and data\-parallel shards were produced by splitting this buffer\. As a result, only tensors crossing a DP boundary required an additional gather\. These advantageous strategies, however, have been deprecated in more recent releases\.
2. 2\.They adopt a fine\-grained MoE architecture with only a single dense layer \(even fewer than in DeepSeek\-V3\[[9](https://arxiv.org/html/2608.11612#bib.bib4)\]\), so most matrices remain small even at the one\-trillion\-parameter scale, keeping the Newton\-Schulz overhead manageable\.
3. 3\.Their main distributed training strategy combines pipeline parallelism with expert parallelism, which naturally distributes Muon’s computation across devices with minimal communication\.

In sum, Muon has been successfully scaled, but its success relies on a subtle alignment of architectural, parallelism, and framework factors\. For Muon to serve as a general\-purpose replacement for Adam, it would benefit from cheaper, more flexible scaling properties that relax these constraints\. This is the goal of the present work\.

## 3Comparison with Related Work

As Muon has gained popularity, successive work has sought to improve it in several ways\. Most of these proposals \(e\.g\. NorMuon\) modify Muon’s update rule so as to reach a given loss in fewer training steps; however, they use the same Newton\-Schulz routine described above and generally suffer from the same scaling challenges\.

### 3\.1Improving Newton\-Schulz

A few papers*have*attempted to improve the Newton\-Schulz step itself—e\.g\., by optimizing the sequence of polynomials\(at,bt,ct\)\(a\_\{t\},b\_\{t\},c\_\{t\}\)or the normalization step\[[3](https://arxiv.org/html/2608.11612#bib.bib10),[15](https://arxiv.org/html/2608.11612#bib.bib11),[6](https://arxiv.org/html/2608.11612#bib.bib17)\]—but they retain the form of[Algorithm˜1](https://arxiv.org/html/2608.11612#alg1)\. In contrast, our Gram Newton\-Schulz algorithm departs from this form\. Since its output is still mathematically identical to the standard version, it remains compatible with nearly all varieties of Muon, including prior improvements to Newton\-Schulz\.

Gram Newton\-Schulz is closely akin to a method proposed in Appendix J of\[[3](https://arxiv.org/html/2608.11612#bib.bib10)\]\. Both aim to reduce the FLOP cost of Newton\-Schulz, both form the Gram matrix to reduce the number ofm×nm\\times nmatrix multiplications, and both are mathematically identical to standard Newton\-Schulz\. However, our work supersedes\[[3](https://arxiv.org/html/2608.11612#bib.bib10)\]in several ways\. First, the precise formulas of Gram Newton\-Schulz are different and, we believe, more stable\. Second, we use symmetric matrix multiplication kernels; the opportunity to use these kernels more is an essential advantage of Gram Newton\-Schulz not studied previously\. Third, we undertake a thorough stability analysis and provide practical recommendations that allow Gram Newton\-Schulz to be used in practice with minimal ad\-hoc hyperparameter tuning\. The iterative part of the method, which approximates the inverse square root of the Gram matrix𝑸T≈𝑹0−1/2:=\(𝑿​𝑿⊤\)−1/2\{\\bm\{Q\}\}\_\{T\}\\approx\{\\bm\{R\}\}\_\{0\}^\{\-1/2\}:=\(\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\)^\{\-1/2\}, generalizes and speeds up the method of\[[22](https://arxiv.org/html/2608.11612#bib.bib24)\]\(see also\[[17](https://arxiv.org/html/2608.11612#bib.bib25), Eq\. 7\.18 \]\), but Lakić does not consider the symmetric kernels, the stability issues, or the application topolar⁡\(𝑿\)\\operatorname\{polar\}\(\{\\bm\{X\}\}\)\.

The idea of exploiting symmetry to reduce the arithmetic cost of computing𝑨⊤​𝑨\{\\bm\{A\}\}^\{\\top\}\{\\bm\{A\}\}has appeared before, both in standard linear algebra packages like BLAS and in the context of Muon’s Newton\-Schulz routine\[[30](https://arxiv.org/html/2608.11612#bib.bib26),[33](https://arxiv.org/html/2608.11612#bib.bib23),[25](https://arxiv.org/html/2608.11612#bib.bib27)\]\. However, our results show that the true potential of this trick is only realized when combined with Gram Newton\-Schulz, which requires more general symmetric operations like𝑨​𝑩\+β​𝑪\{\\bm\{A\}\}\{\\bm\{B\}\}\+\\beta\{\\bm\{C\}\}\. Furthermore, our symmetric matrix multiplication kernels are written in CuteDSL, exploiting the advanced features of NVIDIA’s Hopper and Blackwell GPUs\.

### 3\.2Orthogonalizing a Smaller Matrix

Our update rule is part of a second line of work that reduces the cost of the Newton\-Schulz step by shrinking the size of its input, an approach pioneered by Dion\[[2](https://arxiv.org/html/2608.11612#bib.bib5)\]\. Dion constructs a low\-rank approximation of the momentum matrix and orthogonalizes only this approximation\. If𝑴≈𝑴​𝑽^⋅𝑽^⊤\{\\bm\{M\}\}\\approx\{\\bm\{M\}\}\\widehat\{\{\\bm\{V\}\}\}\\cdot\\widehat\{\{\\bm\{V\}\}\}^\{\\top\}is a rank\-kkapproximation with𝑽^⊤​𝑽^=𝑰∈ℝk×k\\widehat\{\{\\bm\{V\}\}\}^\{\\top\}\\widehat\{\{\\bm\{V\}\}\}=\{\\bm\{I\}\}\\in\\mathbb\{R\}^\{k\\times k\}, then

polar⁡\(𝑴\)≈polar⁡\(𝑴​𝑽^​𝑽^⊤\)=polar⁡\(𝑴​𝑽^\)​𝑽^⊤\.\\operatorname\{polar\}\(\{\\bm\{M\}\}\)\\approx\\operatorname\{polar\}\(\{\\bm\{M\}\}\\widehat\{\{\\bm\{V\}\}\}\\widehat\{\{\\bm\{V\}\}\}^\{\\top\}\)=\\operatorname\{polar\}\(\{\\bm\{M\}\}\\widehat\{\{\\bm\{V\}\}\}\)\\widehat\{\{\\bm\{V\}\}\}^\{\\top\}\.\(3\)Because the dimensions of𝑴​𝑽^\{\\bm\{M\}\}\\widehat\{\{\\bm\{V\}\}\}are much smaller than those of𝑴\{\\bm\{M\}\}, this approximation greatly reduces the runtime of Newton\-Schulz\. In Dion,𝑽^\\widehat\{\{\\bm\{V\}\}\}is found by a warm\-started power\-iteration procedure\. If𝑽^\\widehat\{\{\\bm\{V\}\}\}spans the top\-kkright singular vectors of𝑴\{\\bm\{M\}\}, then this approximation gives the optimal rank\-kkapproximation ofpolar⁡\(𝑴\)\\operatorname\{polar\}\(\{\\bm\{M\}\}\)\.\[[2](https://arxiv.org/html/2608.11612#bib.bib5)\]found that approximating𝑽^\\widehat\{\{\\bm\{V\}\}\}from only a single warm\-started power\-iteration suffices for good downstream performance\. A key element of Dion’s success is*error feedback*, a technique that helps offset the error introduced by low\-rank approximation in future iterations\. We adopt error feedback and describe it fully in[Section˜6](https://arxiv.org/html/2608.11612#S6)\.

Trion\[[29](https://arxiv.org/html/2608.11612#bib.bib19)\]adopts the main ideas ofDion, but changes how the low\-rank approximation is computed\. Instead of constructing𝑽^\\widehat\{\{\\bm\{V\}\}\}via power iteration, it selectskkcolumns from the discrete cosine transform matrix\. Surprisingly, this simpler approximation performs even better thanDion, suggesting that finding a*good*low\-rank approximation is not essential; the error\-feedback mechanism compensates even for large differences between𝑴\{\\bm\{M\}\}and𝑴​𝑽^​𝑽^⊤\{\\bm\{M\}\}\\widehat\{\{\\bm\{V\}\}\}\\widehat\{\{\\bm\{V\}\}\}^\{\\top\}\.

We push this simplification further by using the simplest possible “low\-rank approximation” of𝑴\{\\bm\{M\}\}: we selectkkof its rows or columns and do not operate on the others\. This achieves the same savings in the runtime of Newton\-Schulz as Dion and Trion, but makes all the other steps of the algorithm \(construction of𝑽\{\\bm\{V\}\}, error feedback, etc\.\) much easier to implement, especially in the distributed setting\.

Besides low\-rank approximation, block orthogonalization is an alternative way to shrink the size of the input to Newton\-Schulz\. This approach partitions the momentum matrix into blocks—each typically corresponding to a shard in weight\-sharded settings—and orthogonalizes each block separately\. Several papers have proposed versions of this approach\[[7](https://arxiv.org/html/2608.11612#bib.bib20),[20](https://arxiv.org/html/2608.11612#bib.bib16),[32](https://arxiv.org/html/2608.11612#bib.bib21)\]\. In particular,\[[20](https://arxiv.org/html/2608.11612#bib.bib16)\]found success by periodically alternating block\-wise steps with full steps, a method they callMuonBP\.

Both low\-rank and block orthogonalization methods use Newton\-Schulz as a black box\. Therefore, they are trivial to combine with improvements to Newton\-Schulz\.

## 4Gram Newton\-Schulz

Our first contribution is an orthogonalization algorithm that uses fewer rectangular matrix multiplications than standard Newton\-Schulz\. Instead of iterating directly on the large input matrix𝑿\{\\bm\{X\}\}, we iterate on the small symmetric Gram matrix𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\. The output of our algorithm is mathematically identical to that of standard Newton\-Schulz, but it is significantly cheaper to compute\.

At a high level, our strategy is based on the following formula\. If𝑿∈ℝn×m\{\\bm\{X\}\}\\in\\mathbb\{R\}^\{n\\times m\}withn≤mn\\leq m, thenpolar​\(𝑿\)=\(𝑿​𝑿⊤\)−1/2​𝑿\\mathrm\{polar\}\(\{\\bm\{X\}\}\)=\(\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\)^\{\-1/2\}\{\\bm\{X\}\}\. Rather than use an iterative method to approximate𝑿T≈polar​\(𝑿\)\{\\bm\{X\}\}\_\{T\}\\approx\\mathrm\{polar\}\(\{\\bm\{X\}\}\)directly, we instead:

1. 1\.Compute then×nn\\times nGram matrix𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}
2. 2\.Use an iterative method to approximate𝑸T≈\(𝑿​𝑿⊤\)−1/2\{\\bm\{Q\}\}\_\{T\}\\approx\(\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\)^\{\-1/2\}
3. 3\.Output𝑸T​𝑿\{\\bm\{Q\}\}\_\{T\}\{\\bm\{X\}\}

Step 2—which comprises almost all of the algorithm’s wall clock runtime and FLOP cost—works entirely with smalln×nn\\times nsymmetric matrices\. This version uses just two rectangular matrix multiplications:𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}in the beginning and𝑸T​𝑿\{\\bm\{Q\}\}\_\{T\}\{\\bm\{X\}\}at the end\. It also synergizes well with our symmetric GEMM kernels \([Section˜5](https://arxiv.org/html/2608.11612#S5)\)\. Because we use more symmetric multiplications than before, these kernels provide an even greater speedup\. Since our algorithm works on then×nn\\times nGram matrix of𝑿\{\\bm\{X\}\}, we call it “Gram Newton\-Schulz”\.

How can we ensure that the output matches that of standard Newton\-Schulz? Recall that Newton\-Schulz outputs\(pT∘⋯∘p1\)​\(𝑿\)\(p\_\{T\}\\circ\\cdots\\circ p\_\{1\}\)\(\{\\bm\{X\}\}\), where eachptp\_\{t\}is an odd polynomialp​\(x\)=a​x\+b​x3\+c​x5p\(x\)=ax\+bx^\{3\}\+cx^\{5\}\. Any odd polynomial can be rewritten in the formp​\(x\)=x​h​\(x2\)p\(x\)=xh\(x^\{2\}\), wherehhis a lower\-degree polynomial with the same coefficients, likeh​\(x\)=a\+b​x\+c​x2h\(x\)=a\+bx\+cx^\{2\}\. Intuitively, ifp​\(x\)≈1p\(x\)\\approx 1, thenh​\(y\)=p​\(y1/2\)​y−1/2≈y−1/2h\(y\)=p\(y^\{1/2\}\)y^\{\-1/2\}\\approx y^\{\-1/2\}, so the Newton\-Schulz polynomials implicitly provide a way to approximate inverse square roots\. If we use this approximation in step 2 of Gram Newton\-Schulz, then its output will match that of standard Newton\-Schulz\.

Formally, Gram Newton\-Schulz is based on the following theorem\. In effect, it shows how to compute𝑿T\{\\bm\{X\}\}\_\{T\}from𝑿0\{\\bm\{X\}\}\_\{0\}without ever constructing the intermediate values𝑿1,…,𝑿T−1\{\\bm\{X\}\}\_\{1\},\\ldots,\{\\bm\{X\}\}\_\{T\-1\}:

###### Theorem 2\.

Ifpt​\(x\)=x​ht​\(x2\)p\_\{t\}\(x\)=xh\_\{t\}\(x^\{2\}\)for allt∈\{1,…,T\}t\\in\\\{1,\\ldots,T\\\}, then\(pT∘⋯∘p1\)​\(x\)=qT​x\(p\_\{T\}\\circ\\cdots\\circ p\_\{1\}\)\(x\)=q\_\{T\}x, whereqTq\_\{T\}is defined by the iterationr0=x2,q0=1,r\_\{0\}=x^\{2\},\\,\\,q\_\{0\}=1,and

zt=ht​\(rt−1\),rt=rt−1​zt2,qt=qt−1​zt\\displaystyle z\_\{t\}=h\_\{t\}\(r\_\{t\-1\}\),\\qquad\\qquad r\_\{t\}=r\_\{t\-1\}z\_\{t\}^\{2\},\\qquad\\qquad q\_\{t\}=q\_\{t\-1\}z\_\{t\}for allt∈\{1,…,T\}t\\in\\\{1,\\ldots,T\\\}\.

###### Proof\.

Definex0=xx\_\{0\}=xandxt=pt​\(xt−1\)x\_\{t\}=p\_\{t\}\(x\_\{t\-1\}\)fort∈\{1,…,T\}t\\in\\\{1,\\ldots,T\\\}\. We will show by induction thatrt=xt2r\_\{t\}=x\_\{t\}^\{2\}andqt=xt/x0q\_\{t\}=x\_\{t\}/x\_\{0\}for alltt\. The base caset=0t=0holds by the definitionr0=x2,q0=1r\_\{0\}=x^\{2\},q\_\{0\}=1\. Now assume the hypothesis holds fort−1t\-1\. By assumption,

xt=pt​\(xt−1\)=xt−1​ht​\(xt−12\)x\_\{t\}=p\_\{t\}\(x\_\{t\-1\}\)=x\_\{t\-1\}h\_\{t\}\(x\_\{t\-1\}^\{2\}\)By the inductive hypothesis,ht​\(xt−12\)=ht​\(rt−1\)=zth\_\{t\}\(x\_\{t\-1\}^\{2\}\)=h\_\{t\}\(r\_\{t\-1\}\)=z\_\{t\}, soxt=xt−1​ztx\_\{t\}=x\_\{t\-1\}z\_\{t\}\. Squaring both sides,

xt2=xt−12​zt2=rt−1​zt2=rtx\_\{t\}^\{2\}=x\_\{t\-1\}^\{2\}z\_\{t\}^\{2\}=r\_\{t\-1\}z\_\{t\}^\{2\}=r\_\{t\}If we instead divide both sides byx0x\_\{0\}and apply the other part of the inductive hypothesis, we get

xtx0=xt−1x0​zt=qt−1​zt=qt\\frac\{x\_\{t\}\}\{x\_\{0\}\}=\\frac\{x\_\{t\-1\}\}\{x\_\{0\}\}z\_\{t\}=q\_\{t\-1\}z\_\{t\}=q\_\{t\}Thus, both parts of the hypothesis hold fortt\. Finally, conclude\(pT∘⋯∘p1\)​\(x\)=xT=qT​x0\(p\_\{T\}\\circ\\cdots\\circ p\_\{1\}\)\(x\)=x\_\{T\}=q\_\{T\}x\_\{0\}\. ∎

Note that, as an immediate corollary of the proof,qt=xt/x0→1/x0=\(x02\)−1/2q\_\{t\}=x\_\{t\}/x\_\{0\}\\to 1/x\_\{0\}=\\left\(x\_\{0\}^\{2\}\\right\)^\{\-1/2\}\. In effect, this shows that𝑸T→\(𝑿​𝑿⊤\)−1/2\{\\bm\{Q\}\}\_\{T\}\\to\(\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\)^\{\-1/2\}\.

To obtain our initial version of Gram Newton\-Schulz, we simply lift the iteration from[Theorem˜2](https://arxiv.org/html/2608.11612#Thmtheorem2)to matrices\. As in standard Newton\-Schulz, each matrix operation preserves singular vectors\. Therefore, each singular value of𝑹t\{\\bm\{R\}\}\_\{t\},𝑸t\{\\bm\{Q\}\}\_\{t\}, and𝒁t\{\\bm\{Z\}\}\_\{t\}evolves independently of the others according to the scalar iteration described above\. Note that while this algorithm is mathematically equivalent to standard Newton\-Schulz, it is not yet practical due to numerical instability\. The only difference between our proposed method \([Algorithm˜3](https://arxiv.org/html/2608.11612#alg3)\) and this naive version is the presence of what we call a “restart” at the beginning of iteration 3 of the loop\. We will motivate this change below \([Section˜4\.2](https://arxiv.org/html/2608.11612#S4.SS2)\)\.

Algorithm 2Naive Gram Newton\-Schulz1:

𝑿∈ℝn×m\{\\bm\{X\}\}\\in\\mathbb\{R\}^\{n\\times m\}with

n≤mn\\leq m, coefficients

\{\(at,bt,ct\)\}t=15\\\{\(a\_\{t\},b\_\{t\},c\_\{t\}\)\\\}\_\{t=1\}^\{5\}
2:

𝑿←𝑿/\(∥𝑿∥𝖥\+ϵ\)\{\\bm\{X\}\}\\leftarrow\{\\bm\{X\}\}\\,/\\,\(\\lVert\{\\bm\{X\}\}\\rVert\_\{\\mathsf\{F\}\}\+\\epsilon\)⊳\\trianglerightNormalize sing vals to\[0,1\]\[0,1\],ϵ=10−7\\epsilon=10^\{\-7\}

3:

𝑹0=𝑿​𝑿⊤\{\\bm\{R\}\}\_\{0\}=\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}
4:

𝑸0=𝑰\{\\bm\{Q\}\}\_\{0\}=\{\\bm\{I\}\}
5:for

t=1,…,5t=1,\\ldots,5do

6:

𝒁t←at​𝑰\+bt​𝑹t−1\+ct​𝑹t−12\{\\bm\{Z\}\}\_\{t\}\\leftarrow a\_\{t\}\{\\bm\{I\}\}\+b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}⊳\\trianglerightApplyht​\(𝑹t−1\)h\_\{t\}\(\{\\bm\{R\}\}\_\{t\-1\}\)

7:

𝑸t←𝑸t−1​𝒁t\{\\bm\{Q\}\}\_\{t\}\\leftarrow\{\\bm\{Q\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}
8:

𝑹t←𝒁t​𝑹t−1​𝒁t\{\\bm\{R\}\}\_\{t\}\\leftarrow\{\\bm\{Z\}\}\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}
9:endfor

10:return

𝑸5​𝑿\{\\bm\{Q\}\}\_\{5\}\{\\bm\{X\}\}

### 4\.1Runtime of Naive Gram Newton\-Schulz

We now calculate the FLOP count of this new algorithm to show how its runtime improves on standard Newton\-Schulz\. There are four matrix multiplications per iteration; if we use our symmetric GEMM kernel, these costn3n^\{3\}FLOPs each\. The initialization \(𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\) and output \(𝑸5​𝑿\{\\bm\{Q\}\}\_\{5\}\{\\bm\{X\}\}\) steps costm​n2mn^\{2\}and2​m​n22mn^\{2\}, respectively, since𝑸5​𝑿\{\\bm\{Q\}\}\_\{5\}\{\\bm\{X\}\}is not symmetric\. Computing𝑸1=𝑸0​𝒁1\{\\bm\{Q\}\}\_\{1\}=\{\\bm\{Q\}\}\_\{0\}\{\\bm\{Z\}\}\_\{1\}is free since𝑸0=𝑰\{\\bm\{Q\}\}\_\{0\}=\{\\bm\{I\}\}, and we can skip computing𝑹5=𝒁5​𝑹4​𝒁5\{\\bm\{R\}\}\_\{5\}=\{\\bm\{Z\}\}\_\{5\}\{\\bm\{R\}\}\_\{4\}\{\\bm\{Z\}\}\_\{5\}; together this saves3​n33n^\{3\}FLOPs\. Thus, the total cost isT⋅4​n3\+3​m​n2−3​n3=\(4​T\+3​α−3\)​n3T\\cdot 4n^\{3\}\+3mn^\{2\}\-3n^\{3\}=\(4T\+3\\alpha\-3\)n^\{3\}FLOPs\.

Compare this to standard Newton\-Schulz’sT​\(3​α\+1\)​n3T\(3\\alpha\+1\)n^\{3\}FLOPs when using symmetric GEMMs\. Whenα=1\\alpha=1, they are equal\. Whenα\>1\\alpha\>1, Gram Newton\-Schulz is cheaper, often significantly so\. For a typical Muon application \(T=5,α=4T=5,\\alpha=4333Transformers’ MLP blocks typically have an intermediate dimension4×4\\timesthe model’s hidden dimension\.\), it saves55%of the FLOPs used by standard Newton\-Schulz with symmetric GEMMs, or68%compared to a typical implementation without symmetric GEMMs\. For largerTTandα\\alpha, the savings are even greater, as the leading order term isO​\(\(T\+α\)​n3\)O\(\(T\+\\alpha\)n^\{3\}\)instead ofO​\(T​α​n3\)O\(T\\alpha n^\{3\}\)\.

In practice, whenα=1\\alpha=1, we fall back to standard Newton\-Schulz with our symmetric GEMMs \([Section˜C\.3](https://arxiv.org/html/2608.11612#A3.SS3)\), since it launches fewer GEMMs and has a faster wall clock time\.

### 4\.2Stabilizing Gram Newton\-Schulz

As written,[Algorithm˜2](https://arxiv.org/html/2608.11612#alg2)destabilizes Muon and can even diverge\. The leading cause of this instability is the introduction of spurious negative eigenvalues in the Gram matrix𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}due to large rounding errors in half\-precision arithmetic\. These negative eigenvalues theoretically should not exist in a Gram matrix like𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}\. The main loop of Gram Newton\-Schulz approximates inverse square root of positive numbers, but it diverges for negative inputs\. We provide a thorough theoretical analysis and numerical experiments showing the impact of spurious negative eigenvalues in the Gram matrix in[Section˜B\.1](https://arxiv.org/html/2608.11612#A2.SS1)\.

In practice, we can fully mitigate this instability using a restarting strategy\. Instead of running all five iterations of the loop in a single pass and outputting𝑸5​𝑿\{\\bm\{Q\}\}\_\{5\}\{\\bm\{X\}\}, we run only the first two iterations and compute𝑿2=𝑸2​𝑿\{\\bm\{X\}\}\_\{2\}=\{\\bm\{Q\}\}\_\{2\}\{\\bm\{X\}\}\. We then “restart” the algorithm treating𝑿2\{\\bm\{X\}\}\_\{2\}as the new input; we construct the Gram matrix𝑿2​𝑿2⊤\{\\bm\{X\}\}\_\{2\}\{\\bm\{X\}\}\_\{2\}^\{\\top\}, initialize𝑸2=𝑰\{\\bm\{Q\}\}\_\{2\}=\{\\bm\{I\}\}, and proceed with the remaining three iterations of the loop\. This resets any spurious negative eigenvalues to near zero and restores commutativity of𝑿\{\\bm\{X\}\},𝑸\{\\bm\{Q\}\}, and𝑹\{\\bm\{R\}\}at the cost of3​\(α−1\)​n33\(\\alpha\-1\)n^\{3\}FLOPs\. We find this strategy is sufficient to preserve training quality\.

To select the best iteration after which to restart, we sweep for the restart location that provides the best bound on the condition number of𝑸t\{\\bm\{Q\}\}\_\{t\}, assuming that forming the Gram matrix introduces spurious eigenvalues as negative as−4⋅10−4\-4\\cdot 10^\{\-4\}\. We walk through an example of how to automate this sweep in[Section˜B\.2](https://arxiv.org/html/2608.11612#A2.SS2)\.

[Algorithm˜3](https://arxiv.org/html/2608.11612#alg3)presents our stabilized, training\-ready version of Gram Newton\-Schulz, which uses the restart strategy as well as a slight reformulation of the intermediate polynomials and a switch tofloat16instead ofbfloat16\. The latter two changes are motivated in[Sections˜B\.2\.3](https://arxiv.org/html/2608.11612#A2.SS2.SSS3)and[B\.2\.2](https://arxiv.org/html/2608.11612#A2.SS2.SSS2.Px2)\. Our algorithm is implemented in a pip\-installable package called[gram\-newton\-schulz](https://github.com/Dao-AILab/gram-newton-schulz), which uses the Polar Express coefficients\[[3](https://arxiv.org/html/2608.11612#bib.bib10)\]\.

Algorithm 3Stabilized Gram Newton\-Schulz1:

𝑿∈ℝn×m\{\\bm\{X\}\}\\in\\mathbb\{R\}^\{n\\times m\}with

n≤mn\\leq m, coefficients

\{\(at,bt,ct\)\}t=15\\\{\(a\_\{t\},b\_\{t\},c\_\{t\}\)\\\}\_\{t=1\}^\{5\}
2:

𝑿←𝑿/\(∥𝑿∥𝖥\+ϵ\)\{\\bm\{X\}\}\\leftarrow\{\\bm\{X\}\}\\,/\\,\(\\lVert\{\\bm\{X\}\}\\rVert\_\{\\mathsf\{F\}\}\+\\epsilon\)⊳\\trianglerightNormalize sing vals to\[0,1\]\[0,1\],ϵ=10−7\\epsilon=10^\{\-7\}

3:

𝑿←float16​\(𝑿\)\{\\bm\{X\}\}\\leftarrow\\texttt\{float16\}\(\{\\bm\{X\}\}\)⊳\\trianglerightCast to half precision for speed

4:

𝑹0←𝑿​𝑿⊤\{\\bm\{R\}\}\_\{0\}\\leftarrow\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}
5:

𝑸0←𝑰\{\\bm\{Q\}\}\_\{0\}\\leftarrow\{\\bm\{I\}\}
6:for

t=1,…,5t=1,\\ldots,5do

7:if

t=3t=3then⊳\\trianglerightRestart to stabilize

8:

𝑿←𝑸2​𝑿\{\\bm\{X\}\}\\leftarrow\{\\bm\{Q\}\}\_\{2\}\{\\bm\{X\}\}
9:

𝑹2←𝑿​𝑿⊤\{\\bm\{R\}\}\_\{2\}\\leftarrow\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}
10:

𝑸2←𝑰\{\\bm\{Q\}\}\_\{2\}\\leftarrow\{\\bm\{I\}\}
11:endif

12:

𝒁t←bt​𝑹t−1\+ct​𝑹t−12\{\\bm\{Z\}\}\_\{t\}\\leftarrow b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}
13:

𝑸t←𝑸t−1​𝒁t\+at​𝑸t−1\{\\bm\{Q\}\}\_\{t\}\\leftarrow\{\\bm\{Q\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}\+a\_\{t\}\{\\bm\{Q\}\}\_\{t\-1\}⊳\\triangleright𝑸t=𝑸t−1​ht​\(𝑹t−1\)\{\\bm\{Q\}\}\_\{t\}=\{\\bm\{Q\}\}\_\{t\-1\}h\_\{t\}\(\{\\bm\{R\}\}\_\{t\-1\}\)

14:

\(𝐑𝐙\)t←𝑹t−1​𝒁t\+at​𝑹t−1\(\\mathbf\{RZ\}\)\_\{t\}\\leftarrow\{\\bm\{R\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}\+a\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}
15:

𝑹t←𝒁t​\(𝐑𝐙\)t\+at​\(𝐑𝐙\)t\{\\bm\{R\}\}\_\{t\}\\leftarrow\{\\bm\{Z\}\}\_\{t\}\(\\mathbf\{RZ\}\)\_\{t\}\+a\_\{t\}\(\\mathbf\{RZ\}\)\_\{t\}⊳\\triangleright𝑹t=𝑹t−1​ht​\(𝑹t−1\)2\{\\bm\{R\}\}\_\{t\}=\{\\bm\{R\}\}\_\{t\-1\}h\_\{t\}\(\{\\bm\{R\}\}\_\{t\-1\}\)^\{2\}

16:endfor

17:

𝑿←𝑸5​𝑿\{\\bm\{X\}\}\\leftarrow\{\\bm\{Q\}\}\_\{5\}\{\\bm\{X\}\}
18:return

𝑿\{\\bm\{X\}\}

## 5Symmetric GEMM Kernels in CuteDSL

Our second contribution is a set of custom GPU kernels for the operations𝑨​𝑩\{\\bm\{A\}\}\{\\bm\{B\}\}andα​𝑨​𝑩\+β​𝑪\\alpha\{\\bm\{A\}\}\{\\bm\{B\}\}\+\\beta\{\\bm\{C\}\}that assume𝑨​𝑩\{\\bm\{A\}\}\{\\bm\{B\}\}and𝑪\{\\bm\{C\}\}are symmetric\. These*symmetric GEMM*kernels compute the lower triangle of the output matrix in the usual way and copy the results to the upper triangle \([Figure˜2](https://arxiv.org/html/2608.11612#S5.F2), left\), saving about half the floating point operations used by standard matrix multiplication routines\. Our kernels allow us to take advantage of the symmetric structure of Gram Newton\-Schulz\. As noted above, they also accelerate standard Newton\-Schulz, but their impact is far greater when combined with Gram Newton\-Schulz\. We target the Hopper and Blackwell GPU architectures\.[Figure˜2](https://arxiv.org/html/2608.11612#S5.F2)\(right\) shows that our kernels achieve superb performance on both architectures, significantly outperforming the standard GEMM routine from cuBLAS across a range of matrix sizes\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x2.png)

![Refer to caption](https://arxiv.org/html/2608.11612v1/x3.png)

Figure 2:Left:Symmetric GEMM computes256×256256\\times 256tiles from the lower triangle and main diagonal, then transposes and copies each lower tile to the corresponding upper tile\.Right:Our CuteDSL symmetric GEMM kernels benchmarked against cuBLAS GEMM kernels on Hopper and Blackwell GPUs\. Input matrices𝑨,𝑩,𝑪\{\\bm\{A\}\},\{\\bm\{B\}\},\{\\bm\{C\}\}have dimensionsn×nn\\times n\. For large enoughnn, our kernels achieve a∼2×\\sim 2\\timesspeedup over cuBLAS, both with and without an epilogue addition of𝑪\{\\bm\{C\}\}\.Most GEMM kernels have the following form:

1. 1\.Scheduler: The output matrix is divided into tiles\. A schedule is created that assigns each tile to a group of workers\.
2. 2\.Computing each output tile of𝑨​𝑩\{\\bm\{A\}\}\{\\bm\{B\}\}orα​𝑨​𝑩\+β​𝑪\\alpha\{\\bm\{A\}\}\{\\bm\{B\}\}\+\\beta\{\\bm\{C\}\}requires three steps: 1. \(a\)Prologue: The rows of𝑨\{\\bm\{A\}\}and columns of𝑩\{\\bm\{B\}\}needed for the current tile are loaded in from general memory \(high\-bandwidth memory\) to shared memory \(SRAM\)\. 2. \(b\)Matrix\-Multiply Accumulate \(MMA\): The rows and columns are multiplied and written to the register file \(Hopper\) or tensor memory \(Blackwell\)\. 3. \(c\)Epilogue: Additional tensors needed for fused operations \(e\.g\.𝑪\{\\bm\{C\}\},α\\alpha,β\\beta\) are loaded, fused operations are executed, and the final output is written from the register file to shared memory and then to general memory\.

Our symmetric GEMM kernels differ from the standard one only in their schedulers and epilogues\.

##### Triangular Scheduler

In the standard GEMM, the entire output matrix is partitioned into work tiles that are load balanced and evenly partitioned amongst clusters of thread blocks, where thread blocks in the same cluster can access the same shared memory and are therefore scheduled to run together\. Each cluster then computes its assigned work tiles in succession\. Our tile scheduler in the symmetric GEMM is almost identical, except that it only partitions work tiles from the lower triangle of the matrix \(including the main diagonal\); work tiles in the upper triangle are unassigned\. This*triangular scheduler*ensures that clusters are load balanced and no unnecessary work is performed\.

##### Epilogue: Writing to the Transposed Tile

In the GEMM epilogue, when the computed values of the lower triangle \(excluding the main diagonal\) are written to their assigned tile in general memory \(HBM\), they are also written to their transposed tile location in the upper triangle\. The left panel of[Figure˜2](https://arxiv.org/html/2608.11612#S5.F2)illustrates this process\.

We describe the details of our symmetric kernels and of our kernelized implementation of standard Newton\-Schulz in[Appendix˜C](https://arxiv.org/html/2608.11612#A3)\.

## 6The Dion3 Update Rule

Our third contribution is a variant of Muon that avoids orthogonalizing the full momentum matrix, making the optimizer even faster\. At each step, we select only a fraction of the rows from the momentum matrix and perform the Muon update \(including the orthogonalization\) as if the other rows did not exist\. In words, the main steps of the Dion3 update rule are as follows:

1. 1\.Selecta fraction of the rows \(or columns\) of the momentum matrix\.
2. 2\.Orthogonalizethe selected submatrix using Gram Newton\-Schulz\.
3. 3\.Updatethe selected rows of the weight matrix using the output of the previous step\. Do not update the other rows\.
4. 4\.Error FeedbackDecay*only*the selected rows of the momentum matrix by a multiplicative factor\.

Unlike our other contributions, this procedure changes the optimization trajectory, but experiments in[Section˜8\.1](https://arxiv.org/html/2608.11612#S8.SS1)show that this change is benign\. When*all*rows are selected \(i\.e\., the fraction equals 1\), our update rule reduces to Muon in exact arithmetic\. Note that an earlier version of this update rule appeared in an unpublished report under the name Dion2\. As it was not formally published, we drop the old name and present it here as part of Dion3\. We now describe each step of the optimizer in detail\.[Algorithm˜4](https://arxiv.org/html/2608.11612#alg4)gives the full pseudocode\.

Algorithm 4Dion3 update rule \(single weight matrix\)1:Weight

𝑾∈ℝn×m\{\\bm\{W\}\}\\in\\mathbb\{R\}^\{n\\times m\}, gradient

𝑮\{\\bm\{G\}\}, momentum buffer

𝑴\{\\bm\{M\}\}, row\-wise 2ndmoment buffer

𝒗∈ℝn\{\\bm\{v\}\}\\in\\mathbb\{R\}^\{n\}, fraction

f∈\(0,1\]f\\in\(0,1\], momentum

μ\\mu, decay

β2\\beta\_\{2\}, learning rate

η\\eta\.

2:

𝑴←𝑴\+𝑮\{\\bm\{M\}\}\\leftarrow\{\\bm\{M\}\}\+\{\\bm\{G\}\}
3:

𝒮←\\mathcal\{S\}\\leftarrowindices of the

k=⌈f​n⌉k=\\lceil fn\\rceilrows of

𝑴\{\\bm\{M\}\}with largest

ℓ1\\ell\_\{1\}norm⊳\\trianglerightSelection

4:

𝑶←polar⁡\(𝑴​\[𝒮,:\]\)\{\\bm\{O\}\}\\leftarrow\\operatorname\{polar\}\(\{\\bm\{M\}\}\[\\mathcal\{S\},:\]\)⊳\\trianglerightOrthogonalizek×mk\\times msubmatrix

5:

𝒗i←β2​𝒗i\+\(1−β2\)​1m​∑j𝑶i​j2\{\\bm\{v\}\}\_\{i\}\\leftarrow\\beta\_\{2\}\{\\bm\{v\}\}\_\{i\}\+\(1\-\\beta\_\{2\}\)\\tfrac\{1\}\{m\}\\sum\_\{j\}\{\\bm\{O\}\}\_\{ij\}^\{2\}for

i∈𝒮i\\in\\mathcal\{S\}⊳\\trianglerightOptional: NorMuon steps…

6:

𝑶^i​j←𝑶i​j/\(𝒗i\+ϵ\)\\widehat\{\{\\bm\{O\}\}\}\_\{ij\}\\leftarrow\{\\bm\{O\}\}\_\{ij\}/\(\\sqrt\{\{\\bm\{v\}\}\_\{i\}\}\+\\epsilon\)
7:

𝑶^←𝑶^​∥𝑶∥𝖥/∥𝑶^∥𝖥\\widehat\{\{\\bm\{O\}\}\}\\leftarrow\\widehat\{\{\\bm\{O\}\}\}\\,\\lVert\{\\bm\{O\}\}\\rVert\_\{\\mathsf\{F\}\}/\\lVert\\widehat\{\{\\bm\{O\}\}\}\\rVert\_\{\\mathsf\{F\}\}
8:

𝑾←\(1−η⋅wd\)​𝑾\{\\bm\{W\}\}\\leftarrow\(1\-\\eta\\cdot\\texttt\{wd\}\)\{\\bm\{W\}\}⊳\\trianglerightOptional: weight decay

9:

𝑾i​j←𝑾i​j−η​𝑶^i​j\{\\bm\{W\}\}\_\{ij\}\\leftarrow\{\\bm\{W\}\}\_\{ij\}\-\\eta\\,\\widehat\{\{\\bm\{O\}\}\}\_\{ij\}for

i∈𝒮i\\in\\mathcal\{S\}⊳\\trianglerightUpdate selected rows

10:

𝑴i​j←μ​𝑴i​j\{\\bm\{M\}\}\_\{ij\}\\leftarrow\\mu\\,\{\\bm\{M\}\}\_\{ij\}for

i∈𝒮i\\in\\mathcal\{S\}⊳\\trianglerightError Feedback

##### Selection

Our update rule introduces a new hyperparameterf∈\(0,1\]f\\in\(0,1\]that controls the fraction of rows \(or columns\) to select\. We think offfas a compression factor;f=1f=1corresponds to Muon, and decreasingffspeeds up the algorithm\. We recommendf=1/4f=\\nicefrac\{\{1\}\}\{\{4\}\}orf=1/8f=\\nicefrac\{\{1\}\}\{\{8\}\}\. We could subselect either the rows or the columns\. When the weight matrix is sharded across one of these dimensions, we subselect along that same dimension\. Otherwise, we pick whichever dimension is smaller\. For brevity, we refer only to row selection throughout this section\.

We use a simple selection strategy: we pick thek=⌈f​n⌉k=\\lceil fn\\rceilrows of largestℓ1\\ell\_\{1\}norm\. Initial experiments showed that optimization quality when using random selection was not much worse, so other selection strategies may perform well\. In the distributed setting, selecting the top rows across all shards would require an extra synchronization round and a ragged all\-to\-all, so we simply select the top\-fffraction of rows from each shard\. True global selection is available as an option in our package\.

##### Orthogonalization

The benefit of our selection strategy is realized in this step\. For an input of sizen×mn\\times m, each matrix multiplication in Newton\-Schulz costs eitherm​n2mn^\{2\}orn3n^\{3\}FLOPs with our symmetric kernels\. Selection reduces the dimensions tof​n×mfn\\times m, slashing the cost of these multiplications by a factor of at least1/f21/f^\{2\}\. This benefit is compounded by using Gram Newton\-Schulz, which is faster than standard Newton\-Schulz precisely when the aspect ratioα=m/n\\alpha=m/nis large\. Happily, selection increasesα\\alphaby a factor of1/f1/f\.

Our optimizer also benefits the*communication*cost of Muon in the distributed setting\. Rather than assemble the entire momentum matrix on a single device to perform the orthogonalization, we only assemble the selected rows, reducing the communication volume by a factor of1/f1/f\. We describe our communication strategy in detail in[Section˜7](https://arxiv.org/html/2608.11612#S7)\.

##### Weight Update

We experiment with two flavors of the update rule: Muon and NorMuon \(corresponding to \([1](https://arxiv.org/html/2608.11612#S2.E1)\) and \([2](https://arxiv.org/html/2608.11612#S2.E2)\), respectively\)\. For the Muon flavor, we use each row of the orthogonalized submatrix𝑶=polar⁡\(𝑴​\[𝒮,:\]\)\{\\bm\{O\}\}=\\operatorname\{polar\}\(\{\\bm\{M\}\}\[\\mathcal\{S\},:\]\)to update the matching row of the weight matrix:𝑾​\[𝒮,:\]←𝑾​\[𝒮,:\]−η​𝑶\{\\bm\{W\}\}\[\\mathcal\{S\},:\]\\leftarrow\{\\bm\{W\}\}\[\\mathcal\{S\},:\]\-\\eta\{\\bm\{O\}\}\. The NorMuon flavor is similar, but with extra steps that update the row\-wise second moments and use them to rescale𝑶\{\\bm\{O\}\}; see[Algorithm˜4](https://arxiv.org/html/2608.11612#alg4)\. In both flavors, if weight decay is enabled, it is applied to all rows of𝑾\{\\bm\{W\}\}before updating the selected rows\.

Our initial implementation of Dion3 suffered from a nefarious numerical issue that caused it to underperform Muon and NorMuon in final loss\. Normally, the compiler fuses the weight update into a single operation𝑾←\(1−η⋅wd\)​𝑾−η​𝑶\{\\bm\{W\}\}\\leftarrow\(1\-\\eta\\cdot\\texttt\{wd\}\)\{\\bm\{W\}\}\-\\eta\{\\bm\{O\}\}\. Tensors are cast frombfloat16tofloat32, the arithmetic is performed, and the result is cast back tobfloat16\. When we introduced the row selection operator, it broke this automatic fusion, creating multiple rounds of upcasting and downcasting\. To restore the original numerical behavior, we implemented a custom Triton kernel for[˜8](https://arxiv.org/html/2608.11612#alg4.l8)and[9](https://arxiv.org/html/2608.11612#alg4.l9)of[Algorithm˜4](https://arxiv.org/html/2608.11612#alg4)\. Furthermore, we found that NorMuon’s normalization step should be performed infloat32to avoid similar issues\.

##### Error Feedback

We adapt the*error feedback*mechanism from Dion\[[2](https://arxiv.org/html/2608.11612#bib.bib5)\]to our simpler selection\-based strategy for approximating𝑴\{\\bm\{M\}\}\. The usual rule for updating the momentum matrix is𝑴←𝑴\+𝑮;𝑴←μ​𝑴\{\\bm\{M\}\}\\leftarrow\{\\bm\{M\}\}\+\{\\bm\{G\}\};\\,\{\\bm\{M\}\}\\leftarrow\\mu\{\\bm\{M\}\}, which ensures that𝑴\{\\bm\{M\}\}is an exponential moving average of the gradients across iterations\. The main idea of error feedback is to decay only the component of𝑴\{\\bm\{M\}\}that was captured by our approximation\. Decompose𝑴\{\\bm\{M\}\}into two parts, the selected part𝑴^\\widehat\{\{\\bm\{M\}\}\}—which matches𝑴\{\\bm\{M\}\}at the selected rows and contains zeros in all other rows—and the unselected part𝑴−𝑴^\{\\bm\{M\}\}\-\\widehat\{\{\\bm\{M\}\}\}\. \(In Dion,𝑴^\\widehat\{\{\\bm\{M\}\}\}would be a low\-rank approximation of𝑴\{\\bm\{M\}\}\.\) Instead of the usual rule𝑴←μ​𝑴\{\\bm\{M\}\}\\leftarrow\\mu\{\\bm\{M\}\}, error feedback does𝑴←μ​𝑴^\+\(𝑴−𝑴^\)\{\\bm\{M\}\}\\leftarrow\\mu\\widehat\{\{\\bm\{M\}\}\}\+\(\{\\bm\{M\}\}\-\\widehat\{\{\\bm\{M\}\}\}\)\. That is, it decays the selected rows by a factor ofμ\\mubut leaves the other rows alone\.

In effect, error feedback adds an extra\(1−μ\)​\(𝑴−𝑴^\)\(1\-\\mu\)\(\{\\bm\{M\}\}\-\\widehat\{\{\\bm\{M\}\}\}\)to the momentum\. By boosting the residual component, this extra term nudges future iterations to select rows that were ignored by the current iteration\. Even if a given row of𝑮\{\\bm\{G\}\}is consistently small, error feedback allows it to build up in𝑴\{\\bm\{M\}\}over many iterations, so it eventually gets selected and participates in the update\. When the approximation is exact \(𝑴=𝑴^\{\\bm\{M\}\}=\\widehat\{\{\\bm\{M\}\}\}\), we recover standard momentum\.

##### Finite Precision

Atf=1f=1every row is selected each step, and Dion3 reduces to Muon/NorMuon up to four implementation details: the selection step permutes the rows, momentum damping is applied after forming the update444That is, Newton\-Schulz sees𝑴\+𝑮\{\\bm\{M\}\}\+\{\\bm\{G\}\}rather thanμ​𝑴\+𝑮\\mu\{\\bm\{M\}\}\+\{\\bm\{G\}\}\. Since𝑴\{\\bm\{M\}\}is initialized to zero, these matrices differ by a constant factor ofμ\\mu, which is removed by orthogonalization anyway\., the NorMuon normalization steps run infloat32, and the update step uses our custom Triton kernel\. We confirmed that, despite these differences, the convergence curve of our optimizer matches that of Muon / NorMuon almost exactly \(see[Section˜D\.4](https://arxiv.org/html/2608.11612#A4.SS4)\)\.

##### Kernel launch minimization

Dion3’s row selection process necessarily requires more kernel launches than standard Muon does\. At small scale, the cost of these extra kernel launches can dominate\. Since the kernel dependency structure is the same from round to round, we use a standard CUDA graph capture and replay to avoid this overhead and realize the full potential of our optimizer on smaller\-scale models\.

## 7Megabatching and Communication

Our final contribution concerns the distributed setting\. Under fully\-sharded data parallelism \(FSDP\), each momentum matrix is partitioned across GPUs, so it must first be assembled onto a single device before Newton\-Schulz can run\. This requires an all\-to\-all communication along the sharded dimension and a second all\-to\-all to scatter back the result\. As discussed in[Section˜2\.3](https://arxiv.org/html/2608.11612#S2.SS3), the cost of these operations is a major obstacle to scaling Muon\[[11](https://arxiv.org/html/2608.11612#bib.bib3),[2](https://arxiv.org/html/2608.11612#bib.bib5)\]\. To help users easily adopt Dion3 in a wide range of scenarios, we implement it indion, a pip\-installable package that gracefully manages this communication for FSDP2, DDP, and mixed sharding strategies with minimal overhead\. The package supports our kernelized implementation of Gram Newton\-Schulz and both flavors of the Dion3 update rule as well as Muon, NorMuon, and standard Newton\-Schulz, all of which share the same distributed backend\. We believedionprovides a full\-stack, production\-ready solution for distributed orthogonal optimization\.

A key feature of our implementation ismegabatching\. While the update rule of[Section˜6](https://arxiv.org/html/2608.11612#S6)reduces the communication*volume*\(the number of bits that must be assembled and scattered\), megabatching reduces the number of*rounds*of communication\. A naive implementation of Muon performs the update in batches ofworld\_sizematrices at a time; each matrix in the batch is assembled and orthogonalized in parallel on a different GPU without any GPU sitting idle\. For a model withNNorthogonalized matrices, this requiresO​\(N/world\_size\)O\(N/\\texttt\{world\\\_size\}\)rounds of communication \(all\-to\-alls\) per optimizer step\. Unfortunately, launching and synchronizing each round incurs overhead\. Since this overhead is independent of the communication volume, our update rule does not mitigate it\. Furthermore, each peer\-to\-peer message is small—just one shard of one parameter\. These messages may fail to saturate the interconnect, so bandwidth utilization is poor\. We demonstrate these effects in[Section˜D\.3](https://arxiv.org/html/2608.11612#A4.SS3)\.

Our solution is megabatching; we group all matrices of the same shape into a*single*batch\. For a given shape, the local momentum shards are packed into one all\-to\-all, assembled, orthogonalized as a batch, and scattered back together\. Transformers contain only a handful of distinct weight shapes, so megabatching reduces the number of communication rounds toO​\(1\)O\(1\), independent of model depth\. We implement this megabatching strategy in thedionpackage, where it powers Muon, NorMuon, and Dion3 alike\. Note that megabatching only changes*when*data moves; the assembly, momentum\-state layout, checkpoint format, and orthogonalization math are all the same\.

##### Benchmarking

We compare megabatching against the previous communication strategy by measuring the wall\-clock time of the optimizer step\. We hold the model, data, and optimizer \(Muon\) fixed, and time both strategies back\-to\-back on the same GPUs\. We test two model sizes \(1B and 14B parameters\) and two FSDP configurations \(8 or 32 GPUs\)\. Results are shown in[Table˜1](https://arxiv.org/html/2608.11612#S7.T1)\. For the 14B model, the computational cost of Newton\-Schulz dominates, so reducing communication overhead has little impact\. For the 1B model with 32 shards, each rank holds only one or two matrices of each shape anyway \(see[Table˜4](https://arxiv.org/html/2608.11612#A4.T4)\), so the batch size is about the same for both strategies\. However, for the 1B model with 8 shards, the optimizer is communication\-bound and each rank holds many matrices\. Here, megabatching has a major impact, reducing step time by35%35\\%\.

Table 1:Effect of megabatched communication on the per\-GPU optimizer step time \(Muon, median over steps\)\. “Batching” orthogonalizes matrices in groups whose size equals the shard count; “megabatch” uses a single batch per shape group\. When the weights are small \(so Newton\-Schulz is cheap\) and each rank holds several matrices, megabatching substantially reduces optimizer time\. Each node has four GH200s\.
##### Compressed Data Parallelism

There is a prospect for further savings at larger scales, where data parallelism spans multiple pods in a data\-center network \(DCN\) or hybrid sharding is used\[[35](https://arxiv.org/html/2608.11612#bib.bib29)\]\. A common deployment combines model parallelism within an inter\-chip interconnect \(ICI\) with pure data parallelism across pods\[[4](https://arxiv.org/html/2608.11612#bib.bib28)\]\. Because DCN bandwidth is typically far smaller than ICI bandwidth, shrinking the volume of data\-parallel communication is especially valuable in this regime\. Here, Dion3 allows for compressed data\-parallel synchronizations analogous to those ofDion\[[2](https://arxiv.org/html/2608.11612#bib.bib5), §3\.3\]\. Ordinarily, gradients are averaged across data\-parallel replicas at every step so that the momentum buffers, and hence the weights, stay consistent\. With Dion3, however, we need not synchronize the entire momentum matrix at every step, only the selected submatrix𝑴​\[𝒮,:\]\{\\bm\{M\}\}\[\\mathcal\{S\},:\]\. When the rows𝒮\\mathcal\{S\}can be selected without first synchronizing𝑴\{\\bm\{M\}\}\(e\.g\., by picking𝒮\\mathcal\{S\}at random\), this strategy performs an identical update at a fraction of the communication cost\.

## 8Experiments

### 8\.1Model Quality Is Preserved

Of our four contributions, only our update rule explicitly changes the optimizer trajectory\. However, Gram Newton\-Schulz and our symmetric GEMM kernels introduce numerical differences due to finite precision arithmetic\. Experiments in[Section˜A\.2](https://arxiv.org/html/2608.11612#A1.SS2)show that these differences do not affect training\. We also confirmed that, as expected, Dion3 with selection fractionf=1f=1matches Muon/NorMuon almost exactly \(see[Figure˜22](https://arxiv.org/html/2608.11612#A4.F22)\)\.

In contrast, Dion3 withf<1f<1is a genuinely new optimizer\. In this section, we demonstrate that Dion3 achieves a slightly*better*loss than the baseline\. We did not set out to improve training quality, so this result is a small but pleasant surprise\. At a minimum, it shows that our subselection strategy does no harm to the loss\.

We train 1B\-parameter dense transformers on 100B tokens of the ClimbMix dataset\[[10](https://arxiv.org/html/2608.11612#bib.bib30)\], using fully\-sharded data parallelism \(FSDP\) with MXFP8 weights\. \(Below, we also train at larger scales\.\) As always, Newton\-Schulz runs in half precision arithmetic\. The models use grouped\-query and sliding\-window attention, but are*not*mixtures of experts\. We report the cross\-entropy loss \(CE\) of next\-token prediction on a held\-out validation split of ClimbMix\. For further details about our setup, see[Section˜D\.1](https://arxiv.org/html/2608.11612#A4.SS1)\. For a fair comparison, we begin by tuning the baseline\. In this section, we compare NorMuon to the NorMuon flavor of Dion3, as initial experiments gave it a slight edge over Muon\. We find that the optimal learning rate and momentum coefficient for NorMuon areη=0\.01\\eta=0\.01andμ=0\.9\\mu=0\.9, respectively\.

##### Learning\-rate transfer rule

Since Dion3 updates only a fraction of the rows in each step, it reduces the effective step size\. This in turn changes the optimal learning rate\. To understand this effect, we sweep both the fractionf∈\{1/2,1/4,1/8,1/16\}f\\in\\\{\\nicefrac\{\{1\}\}\{\{2\}\},\\nicefrac\{\{1\}\}\{\{4\}\},\\nicefrac\{\{1\}\}\{\{8\}\},\\nicefrac\{\{1\}\}\{\{16\}\}\\\}and the learning rateη\\etaof Dion3, along with our earlier sweep for NorMuon \(f=1f=1\)\. The results in[Figure˜3](https://arxiv.org/html/2608.11612#S8.F3)show a clear trend: the optimal learning rate for a givenffscales asη∝1/f\\eta\\propto\\sqrt\{1/f\}\. We can explain this transfer rule as follows\. For Muon, the size of the update is‖η⋅polar⁡\(𝑴\)‖𝖥=η​n\\\|\\eta\\cdot\\operatorname\{polar\}\(\{\\bm\{M\}\}\)\\\|\_\{\\mathsf\{F\}\}=\\eta\\sqrt\{n\}, wherennis the smaller dimension\. For Dion3, it is‖η′⋅polar⁡\(𝑴​\[𝒮,:\]\)‖𝖥=η′​f​n\\\|\\eta^\{\\prime\}\\cdot\\operatorname\{polar\}\(\{\\bm\{M\}\}\[\\mathcal\{S\},:\]\)\\\|\_\{\\mathsf\{F\}\}=\\eta^\{\\prime\}\\sqrt\{fn\}\. For these to match, we must setη′=η/f\\eta^\{\\prime\}=\\eta/\\sqrt\{f\}\. We do not bother retuning the momentum coefficient; we simply copy the optimal setting from NorMuon \(μ=0\.9\\mu=0\.9\) to Dion3\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x4.png)

![Refer to caption](https://arxiv.org/html/2608.11612v1/x5.png)

Figure 3:Left:Final validation loss for 1B\-parameter models trained on 100B tokens of ClimbMix with NorMuon or Dion3, as a function of row\-selection fractionffand learning rateη\\eta\. NorMuon isf=1f=1\. Bolded cell in each row shows bestη\\eta\. Optimalη\\etas track the lineη​f=0\.01\\eta\\sqrt\{f\}=0\.01, shown in red on a log\-log scale\. Not shown:f=1,η=0\.03f=1,\\eta=0\.03\(loss =2\.2372\.237\) andf=1/32,η=0\.04f=\\nicefrac\{\{1\}\}\{\{32\}\},\\eta=0\.04\(loss=2\.191=2\.191\)\.Right:The same data reduced to two dimensions: validation loss as a function ofη​f\\eta\\sqrt\{f\}\. Final loss is minimized whenη​f≈0\.01\\eta\\sqrt\{f\}\\approx 0\.01\.
##### Dion3 improves the loss

[Figure˜3](https://arxiv.org/html/2608.11612#S8.F3)contains another remarkable finding: Dion3 withf<1f<1actually outperforms NorMuon when tuned correctly\. Indeed, the lowest loss is achieved atf=1/8f=\\nicefrac\{\{1\}\}\{\{8\}\}\.[Figure˜4](https://arxiv.org/html/2608.11612#S8.F4)shows the validation loss curve for each optimizer \(excludingf=1/4f=\\nicefrac\{\{1\}\}\{\{4\}\}\) at its best learning rate\. It shows a clear gap throughout training between NorMuon and the Dion3 variants, which confirms this finding\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x6.png)Figure 4:Validation loss curves for 1B\-parameter models trained on 100B tokens of ClimbMix with NorMuon or Dion3 atf=1/2,1/8f=\\nicefrac\{\{1\}\}\{\{2\}\},\\nicefrac\{\{1\}\}\{\{8\}\}, or1/16\\nicefrac\{\{1\}\}\{\{16\}\}\. Each optimizer uses its best learning rate from[Figure˜3](https://arxiv.org/html/2608.11612#S8.F3)\. Parentheses show final loss\. All Dion3 variants track below the fully\-tuned NorMuon baseline throughout training, finishing about0\.010\.01points lower\.To strengthen this finding, we scale up the model size to between 3B and 14B parameters\. For reasons of cost, we now train on 10B tokens instead of 100B\. We compare Dion3 withf=1/4f=\\nicefrac\{\{1\}\}\{\{4\}\}against NorMuon, each using their best learning rate as identified above\. For each scale and optimizer,[Table˜2](https://arxiv.org/html/2608.11612#S8.T2)reports both validation loss on ClimbMix and downstream accuracy on a suite of 12 standard benchmarks\. In validation loss, Dion3 outperforms NorMuon at every scale, achieving the largest improvement \(−0\.027\-0\.027\) at the largest scale \(14B\)\.[Figure˜5](https://arxiv.org/html/2608.11612#S8.F5)shows that, as for the 1B parameter model, validation loss is lower throughout training when using Dion3\. In downstream accuracy, Dion3 wins at three of four scales, with an improvement of 0\.7 percentage points at 14B\.

Table 2:Dion3 \(f=1/4f\\\!=\\\!\\nicefrac\{\{1\}\}\{\{4\}\},η=0\.02\\eta\\\!=\\\!0\.02\) versus NorMuon \(η=0\.01\\eta\\\!=\\\!0\.01\) across model sizes trained on 10B tokens of ClimbMix: final validation loss \(cross\-entropy; lower is better\) and downstream accuracy \(macro\-average over 12 standard benchmarks: ARC easy/challenge, BoolQ, COPA, HellaSwag, LAMBADA, MMLU, OpenBookQA, PIQA, RTE, TruthfulQA, WinoGrande; higher is better\)\. The winner of each comparison is bolded\.![Refer to caption](https://arxiv.org/html/2608.11612v1/x7.png)Figure 5:Validation loss over training at 14B for Dion3 \(f=1/4f\\\!=\\\!\\nicefrac\{\{1\}\}\{\{4\}\},η=0\.02\\eta\\\!=\\\!0\.02\) versus NorMuon \(η=0\.01\\eta\\\!=\\\!0\.01\), trained on 10B tokens of ClimbMix\. Dion3 tracks below NorMuon throughout\.

### 8\.2Dion3 Accelerates the Optimizer

We now measure how much Dion3 reduces the cost of Muon’s optimizer step\. Starting with standard Muon, we successively add our improvements—symmetric kernels, Gram Newton\-Schulz, and fractional updates withf=0\.5f=0\.5orf=0\.25f=0\.25—stacking each on top of the previous ones\. \(All versions use megabatching\.\) We also compare against plain AdamW\. We run each version of the optimizer on models of various sizes either on 1 GPU or using FSDP on 4 GPUs\. We use CUDA events to measure the GPU time of the optimizer step\. For fidelity, we run full training steps with forward and backward passes \(on synthetic data\) but these are excluded from the recorded timings\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x8.png)Figure 6:Optimizer step time \(excluding forward/backward pass\) relative to standard Muon across model scales when training on 1 GH200 \(left\) and FSDP over 4 GH200s \(right\)\. Each colored line adds one of our contributions on top of the previous one; AdamW is shown for reference\. Lines show median over 25 steps and bands show interquartile range\. Adding each of our contributions consistently reduces runtime; together they achieve a6×6\\timesspeedup for larger models\.Results are shown in[Figure˜6](https://arxiv.org/html/2608.11612#S8.F6)\. A subset of these results appears in[Figure˜1](https://arxiv.org/html/2608.11612#S0.F1), and the corresponding timings for the NorMuon family are given in[Section˜D\.4](https://arxiv.org/html/2608.11612#A4.SS4)\([Figure˜23](https://arxiv.org/html/2608.11612#A4.F23)\)\. Each part of Dion3 yields a consistent speedup across scales and parallelism configurations\. Our symmetric kernels and Gram Newton\-Schulz give a combined speedup of1\.5×1\.5\\timesor more over standard Muon\. The fractional update rule cuts the runtime dramatically\. As expected, its impact grows with model size, as the cubic computational cost of Newton\-Schulz increasingly dominates the other operations\. For sufficiently large models, settingf=1/2f=\\nicefrac\{\{1\}\}\{\{2\}\}andf=1/4f=\\nicefrac\{\{1\}\}\{\{4\}\}gives further2×2\\timesand3\.7×3\.7\\timesreductions respectively, for overall speedups of3\.6×3\.6\\timesand6\.5×6\.5\\times\. Thus, while Dion3 is still slower than AdamW, the gap is*significantly*smaller\.

In[Section˜D\.2](https://arxiv.org/html/2608.11612#A4.SS2), we report CPU time and communication volume in addition to GPU time, showing that CPU overhead is small and that fractional updates reduce communication by a factor of1/f1/f\. The communication savings of our fractional update rule can be crucial in some settings, though the setting of[Figure˜6](https://arxiv.org/html/2608.11612#S8.F6)is compute\-bound: neither communication nor host time is exposed\.

In[Appendix˜A](https://arxiv.org/html/2608.11612#A1), we conduct an additional suite of benchmarks on a wider range of architectures, including open source models like Gemma and mixtures\-of\-experts\. On architectures like these, whose weight matrices have higher aspect ratios \(α=8\\alpha=8instead of44\), Gram Newton\-Schulz and the symmetric kernels alone achieve a speedup of2×2\\times\(see[Figures˜9](https://arxiv.org/html/2608.11612#A1.F9)and[10](https://arxiv.org/html/2608.11612#A1.F10)\)\.

## 9Conclusion

For LLMs trained with Muon, estimates suggest that optimizer step time accounts for between 1% \(e\.g\.,\[[1](https://arxiv.org/html/2608.11612#bib.bib37)\]\) and 17% of total training time \(see[Appendix˜E](https://arxiv.org/html/2608.11612#A5)\)\. While Muon has been scaled successfully, doing so has demanded particular choices of architecture and parallelism strategy, along with*significant*engineering effort\. As interest in Muon continues to spread, this paradigm is too brittle; a flexible, general\-purpose remedy for Muon’s overhead is needed\.

Dion3 allows practitioners to easily realize the benefits of Muon without paying the high cost of its orthogonalization step—even for large\-scale models in highly distributed settings\. Our Gram Newton\-Schulz algorithm and CuteDSL kernels speed up Muon by1\.5×1\.5\\timesfor dense models and2×2\\timesfor MoEs, a rare case of free lunch performance\. Megabatching has an equally large effect in certain distributed settings\. Dion3’s fractional update rule \(f=1/4f=\\nicefrac\{\{1\}\}\{\{4\}\}\) provides an*additional*3\.7×3\.7\\timesspeedup\. Though this update rule changes the optimizer trajectory, we find that it actually improves training quality in our setting\. This improvement is unexpected, as we designed Dion3 to cheaply approximate Muon rather than improve upon it\. Further work is needed to determine how widely this improvement generalizes, but we find support in a recent result similar to our own:\[[18](https://arxiv.org/html/2608.11612#bib.bib36)\]showed that randomly masking blocks of the update improves the trajectory of SGD with momentum\. Overall, Dion3 as implemented in ourdionandgram\-newton\-schulzpackages provides the tools to make orthogonal optimizers practical and accessible in a wide range of settings\.

##### Acknowledgements

The authors thank Zichong Li for contributing an implementation of NorMuon to thedionrepo\. NA is supported by NSF award 2234660\.

## References

- \[1\]J\. Abadji, M\. Abdin, C\. Adams, E\. Alcaide, M\. Altun, M\. Artoni, J\. Bao, U\. Barar, V\. Bekiaris, A\. Bessonov, B\. Bütikofer, J\. Chang, Y\. Chen, D\. Chernenkov, Y\. Chi, F\. Christianos, F\. Christopoulou, R\. Ciocoiu, T\. Cohen, Y\. Coppel, D\. Emelianenko, B\. Fergerson, B\. Fitzgerald, M\. Gallé, A\. Golonzovskyi, G\. Grigorev, Y\. Hao, C\. Hensel, J\. Huenermann, Y\. Ji, S\. Joshi, E\. Kant, K\. Khandpur, S\. Kim, V\. Kirichenko, U\. Kocasarac, I\. Kochik, I\. Komarov, C\. Kong, A\. Koul, F\. Lacroix, S\. Laktionov, W\. Long, Q\. Malartic, V\. Markovtsev, A\. Marques, R\. McHardy, C\. Mocholí, D\. Monakhov, A\. Morris, M\. Muller, C\. Mürtz, R\. Nabel, T\. Nguyen, R\. Novosel, S\. Ozog, A\. Patankar, A\. Petrov, A\. Piché, A\. Pignet, T\. Poncu, P\. Potter, A\. Rakowski, P\. Ritschard, J\. Roberts, J\. Rowell, P\. Sarna, P\. Savalle, U\. Sazanovich, N\. Shapovalov, A\. Shevchenko, M\. Shilkov, A\. Sokol, M\. Soliman, J\. Stephenson, V\. Storchan, D\. Tantaru, A\. Tyurin, A\. Wälchli, P\. Wang, J\. Yang, R\. Zayashnikov, A\. Z\. Martin, N\. Zinov, C\. Bercier, J\. Caldeira, M\. Garcia, T\. George, K\. Gharzai, G\. Hitchcock, C\. Klingenberg, I\. Pinto, V\. Randery, N\. Smith, A\. Sugako, and J\. Warner\(2026\)Laguna m\.1/xs\.2 technical report\.External Links:2605\.27605,[Link](https://arxiv.org/abs/2605.27605)Cited by:[§9](https://arxiv.org/html/2608.11612#S9.p1.1)\.
- \[2\]K\. Ahn, B\. Xu, N\. Abreu, Y\. Fan, G\. Magakyan, P\. Sharma, Z\. Zhan, and J\. Langford\(2025\)Dion: distributed orthonormalized updates\.External Links:2504\.05295,[Link](https://arxiv.org/abs/2504.05295)Cited by:[item 3](https://arxiv.org/html/2608.11612#S1.I1.i3.p1.1),[2nd item](https://arxiv.org/html/2608.11612#S2.I1.i2.p1.1),[§3\.2](https://arxiv.org/html/2608.11612#S3.SS2.p1.12),[§3\.2](https://arxiv.org/html/2608.11612#S3.SS2.p1.3),[§6](https://arxiv.org/html/2608.11612#S6.SS0.SSS0.Px4.p1.13),[§7](https://arxiv.org/html/2608.11612#S7.SS0.SSS0.Px2.p1.4),[§7](https://arxiv.org/html/2608.11612#S7.p1.1)\.
- \[3\]N\. Amsel, D\. Persson, C\. Musco, and R\. M\. Gower\(2026\)The polar express: optimal matrix sign methods and their application to the muon algorithm\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=yRtgZ1K8hO)Cited by:[§A\.2](https://arxiv.org/html/2608.11612#A1.SS2.p1.1),[§B\.2\.1](https://arxiv.org/html/2608.11612#A2.SS2.SSS1.p2.14),[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p1.1),[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p2.3),[§4\.2](https://arxiv.org/html/2608.11612#S4.SS2.p4.1)\.
- \[4\]J\. Austin, S\. Douglas, R\. Frostig, A\. Levskaya, C\. Chen, S\. Vikram, F\. Lebron, P\. Choy, V\. Ramasesh, A\. Webson, and R\. Pope\(2025\)How to scale your model\.Google DeepMind\.External Links:[Link](https://jax-ml.github.io/scaling-book/)Cited by:[§7](https://arxiv.org/html/2608.11612#S7.SS0.SSS0.Px2.p1.4)\.
- \[5\]J\. Bernstein\(2025\)Deriving Muon\.External Links:[Link](https://jeremybernste.in/writing/deriving-muon)Cited by:[§2\.1](https://arxiv.org/html/2608.11612#S2.SS1.p1.3)\.
- \[6\]T\. Boissin, T\. Massena, F\. Mamalet, and M\. Serrurier\(2025\)Turbo\-muon: accelerating orthogonality\-based optimization with pre\-conditioning\.External Links:2512\.04632,[Link](https://arxiv.org/abs/2512.04632)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p1.1)\.
- \[7\]V\. Boreiko, Z\. Bu, and S\. Zha\(2025\)Towards understanding of orthogonalization in muon\.InTiny Titans: The next wave of On\-Device Learning for Foundation Models \(TTODLer\-FM\),External Links:[Link](https://openreview.net/forum?id=4vzhqq5hpX)Cited by:[§3\.2](https://arxiv.org/html/2608.11612#S3.SS2.p4.1)\.
- \[8\]F\. L\. Cesista, J\. You, and K\. Jordan\(2025\)Squeezing 1\-2% efficiency gains out of muon by optimizing the newton\-schulz coefficients\.External Links:[Link](http://leloykun.github.io/ponder/muon-opt-coeffs/)Cited by:[§A\.2](https://arxiv.org/html/2608.11612#A1.SS2.p1.1),[§A\.2](https://arxiv.org/html/2608.11612#A1.SS2.p2.1)\.
- \[9\]DeepSeek\-AI, A\. Liu, B\. Feng, B\. Xue, B\. Wang, B\. Wu, C\. Lu, C\. Zhao, C\. Deng, C\. Zhang, C\. Ruan, D\. Dai, D\. Guo, D\. Yang, D\. Chen, D\. Ji, E\. Li, F\. Lin, F\. Dai, F\. Luo, G\. Hao, G\. Chen, G\. Li, H\. Zhang, H\. Bao, H\. Xu, H\. Wang, H\. Zhang, H\. Ding, H\. Xin, H\. Gao, H\. Li, H\. Qu, J\. L\. Cai, J\. Liang, J\. Guo, J\. Ni, J\. Li, J\. Wang, J\. Chen, J\. Chen, J\. Yuan, J\. Qiu, J\. Li, J\. Song, K\. Dong, K\. Hu, K\. Gao, K\. Guan, K\. Huang, K\. Yu, L\. Wang, L\. Zhang, L\. Xu, L\. Xia, L\. Zhao, L\. Wang, L\. Zhang, M\. Li, M\. Wang, M\. Zhang, M\. Zhang, M\. Tang, M\. Li, N\. Tian, P\. Huang, P\. Wang, P\. Zhang, Q\. Wang, Q\. Zhu, Q\. Chen, Q\. Du, R\. J\. Chen, R\. L\. Jin, R\. Ge, R\. Zhang, R\. Pan, R\. Wang, R\. Xu, R\. Zhang, R\. Chen, S\. S\. Li, S\. Lu, S\. Zhou, S\. Chen, S\. Wu, S\. Ye, S\. Ye, S\. Ma, S\. Wang, S\. Zhou, S\. Yu, S\. Zhou, S\. Pan, T\. Wang, T\. Yun, T\. Pei, T\. Sun, W\. L\. Xiao, W\. Zeng, W\. Zhao, W\. An, W\. Liu, W\. Liang, W\. Gao, W\. Yu, W\. Zhang, X\. Q\. Li, X\. Jin, X\. Wang, X\. Bi, X\. Liu, X\. Wang, X\. Shen, X\. Chen, X\. Zhang, X\. Chen, X\. Nie, X\. Sun, X\. Wang, X\. Cheng, X\. Liu, X\. Xie, X\. Liu, X\. Yu, X\. Song, X\. Shan, X\. Zhou, X\. Yang, X\. Li, X\. Su, X\. Lin, Y\. K\. Li, Y\. Q\. Wang, Y\. X\. Wei, Y\. X\. Zhu, Y\. Zhang, Y\. Xu, Y\. Xu, Y\. Huang, Y\. Li, Y\. Zhao, Y\. Sun, Y\. Li, Y\. Wang, Y\. Yu, Y\. Zheng, Y\. Zhang, Y\. Shi, Y\. Xiong, Y\. He, Y\. Tang, Y\. Piao, Y\. Wang, Y\. Tan, Y\. Ma, Y\. Liu, Y\. Guo, Y\. Wu, Y\. Ou, Y\. Zhu, Y\. Wang, Y\. Gong, Y\. Zou, Y\. He, Y\. Zha, Y\. Xiong, Y\. Ma, Y\. Yan, Y\. Luo, Y\. You, Y\. Liu, Y\. Zhou, Z\. F\. Wu, Z\. Z\. Ren, Z\. Ren, Z\. Sha, Z\. Fu, Z\. Xu, Z\. Huang, Z\. Zhang, Z\. Xie, Z\. Zhang, Z\. Hao, Z\. Gou, Z\. Ma, Z\. Yan, Z\. Shao, Z\. Xu, Z\. Wu, Z\. Zhang, Z\. Li, Z\. Gu, Z\. Zhu, Z\. Liu, Z\. Li, Z\. Xie, Z\. Song, Z\. Gao, and Z\. Pan\(2025\)DeepSeek\-v3 technical report\.External Links:2412\.19437,[Link](https://arxiv.org/abs/2412.19437)Cited by:[item 1](https://arxiv.org/html/2608.11612#A5.I1.i1.p1.2),[Appendix E](https://arxiv.org/html/2608.11612#A5.SSx2.p1.8),[1st item](https://arxiv.org/html/2608.11612#S2.I1.i1.p1.2),[item 2](https://arxiv.org/html/2608.11612#S2.I3.i2.p1.1)\.
- \[10\]S\. Diao, Y\. Yang, Y\. Fu, X\. Dong, D\. Su, M\. Kliegl, Z\. Chen, P\. Belcak, Y\. Suhara, H\. Yin, M\. Patwary, Y\. C\. Lin, J\. Kautz, and P\. Molchanov\(2025\)Nemotron\-CLIMB: clustering\-based iterative data mixture bootstrapping for language model pre\-training\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems Datasets and Benchmarks Track,External Links:[Link](https://openreview.net/forum?id=aBlqKPkc4a)Cited by:[§8\.1](https://arxiv.org/html/2608.11612#S8.SS1.p3.2)\.
- \[11\]Essential AI\(2025\)Layer Sharding for Large\-Scale Training with Muon\.External Links:[Link](https://www.essential.ai/research/infra)Cited by:[1st item](https://arxiv.org/html/2608.11612#S2.I1.i1.p1.2),[2nd item](https://arxiv.org/html/2608.11612#S2.I1.i2.p1.1),[§7](https://arxiv.org/html/2608.11612#S7.p1.1)\.
- \[12\]Gemma Team, A\. Kamath, J\. Ferret, S\. Pathak, N\. Vieillard, R\. Merhej, S\. Perrin, T\. Matejovicova, A\. Ramé, M\. Rivière, L\. Rouillard, T\. Mesnard, G\. Cideron, J\. Grill, S\. Ramos, E\. Yvinec, M\. Casbon, E\. Pot, I\. Penchev, G\. Liu, F\. Visin, K\. Kenealy, L\. Beyer, X\. Zhai, A\. Tsitsulin, R\. Busa\-Fekete, A\. Feng, N\. Sachdeva, B\. Coleman, Y\. Gao, B\. Mustafa, I\. Barr, E\. Parisotto, D\. Tian, M\. Eyal, C\. Cherry, J\. Peter, D\. Sinopalnikov, S\. Bhupatiraju, R\. Agarwal, M\. Kazemi, D\. Malkin, R\. Kumar, D\. Vilar, I\. Brusilovsky, J\. Luo, A\. Steiner, A\. Friesen, A\. Sharma, A\. Sharma, A\. M\. Gilady, A\. Goedeckemeyer, A\. Saade, A\. Feng, A\. Kolesnikov, A\. Bendebury, A\. Abdagic, A\. Vadi, A\. György, A\. S\. Pinto, A\. Das, A\. Bapna, A\. Miech, A\. Yang, A\. Paterson, A\. Shenoy, A\. Chakrabarti, B\. Piot, B\. Wu, B\. Shahriari, B\. Petrini, C\. Chen, C\. L\. Lan, C\. A\. Choquette\-Choo, C\. Carey, C\. Brick, D\. Deutsch, D\. Eisenbud, D\. Cattle, D\. Cheng, D\. Paparas, D\. S\. Sreepathihalli, D\. Reid, D\. Tran, D\. Zelle, E\. Noland, E\. Huizenga, E\. Kharitonov, F\. Liu, G\. Amirkhanyan, G\. Cameron, H\. Hashemi, H\. Klimczak\-Plucińska, H\. Singh, H\. Mehta, H\. T\. Lehri, H\. Hazimeh, I\. Ballantyne, I\. Szpektor, I\. Nardini, J\. Pouget\-Abadie, J\. Chan, J\. Stanton, J\. Wieting, J\. Lai, J\. Orbay, J\. Fernandez, J\. Newlan, J\. Ji, J\. Singh, K\. Black, K\. Yu, K\. Hui, K\. Vodrahalli, K\. Greff, L\. Qiu, M\. Valentine, M\. Coelho, M\. Ritter, M\. Hoffman, M\. Watson, M\. Chaturvedi, M\. Moynihan, M\. Ma, N\. Babar, N\. Noy, N\. Byrd, N\. Roy, N\. Momchev, N\. Chauhan, N\. Sachdeva, O\. Bunyan, P\. Botarda, P\. Caron, P\. K\. Rubenstein, P\. Culliton, P\. Schmid, P\. G\. Sessa, P\. Xu, P\. Stanczyk, P\. Tafti, R\. Shivanna, R\. Wu, R\. Pan, R\. Rokni, R\. Willoughby, R\. Vallu, R\. Mullins, S\. Jerome, S\. Smoot, S\. Girgin, S\. Iqbal, S\. Reddy, S\. Sheth, S\. Põder, S\. Bhatnagar, S\. R\. Panyam, S\. Eiger, S\. Zhang, T\. Liu, T\. Yacovone, T\. Liechty, U\. Kalra, U\. Evci, V\. Misra, V\. Roseberry, V\. Feinberg, V\. Kolesnikov, W\. Han, W\. Kwon, X\. Chen, Y\. Chow, Y\. Zhu, Z\. Wei, Z\. Egyed, V\. Cotruta, M\. Giang, P\. Kirk, A\. Rao, K\. Black, N\. Babar, J\. Lo, E\. Moreira, L\. G\. Martins, O\. Sanseviero, L\. Gonzalez, Z\. Gleicher, T\. Warkentin, V\. Mirrokni, E\. Senter, E\. Collins, J\. Barral, Z\. Ghahramani, R\. Hadsell, Y\. Matias, D\. Sculley, S\. Petrov, N\. Fiedel, N\. Shazeer, O\. Vinyals, J\. Dean, D\. Hassabis, K\. Kavukcuoglu, C\. Farabet, E\. Buchatskaya, J\. Alayrac, R\. Anil, Dmitry, Lepikhin, S\. Borgeaud, O\. Bachem, A\. Joulin, A\. Andreev, C\. Hardin, R\. Dadashi, and L\. Hussenot\(2025\)Gemma 3 technical report\.External Links:2503\.19786,[Link](https://arxiv.org/abs/2503.19786)Cited by:[§A\.1](https://arxiv.org/html/2608.11612#A1.SS1.p1.1)\.
- \[13\]GLM\-5 Team, A\. Zeng, X\. Lv, Z\. Hou, Z\. Du, Q\. Zheng, B\. Chen, D\. Yin, C\. Ge, C\. Huang, C\. Xie, C\. Zhu, C\. Yin, C\. Wang, G\. Pan, H\. Zeng, H\. Zhang, H\. Wang, H\. Chen, J\. Zhang, J\. Jiao, J\. Guo, J\. Wang, J\. Du, J\. Wu, K\. Wang, L\. Li, L\. Fan, L\. Zhong, M\. Liu, M\. Zhao, P\. Du, Q\. Dong, R\. Lu, Shuang\-Li, S\. Cao, S\. Liu, T\. Jiang, X\. Chen, X\. Zhang, X\. Huang, X\. Dong, Y\. Xu, Y\. Wei, Y\. An, Y\. Niu, Y\. Zhu, Y\. Wen, Y\. Cen, Y\. Bai, Z\. Qiao, Z\. Wang, Z\. Wang, Z\. Zhu, Z\. Liu, Z\. Li, B\. Wang, B\. Wen, C\. Huang, C\. Cai, C\. Yu, C\. Li, C\. Hu, C\. Zhang, D\. Zhang, D\. Lin, D\. Yang, D\. Wang, D\. Ai, E\. Zhu, F\. Yi, F\. Chen, G\. Wen, H\. Sun, H\. Zhao, H\. Hu, H\. Zhang, H\. Liu, H\. Zhang, H\. Peng, H\. Tai, H\. Zhang, H\. Liu, H\. Wang, H\. Yan, H\. Ge, H\. Liu, H\. Chu, J\. Zhao, J\. Wang, J\. Zhao, J\. Ren, J\. Wang, J\. Zhang, J\. Gui, J\. Zhao, J\. Li, J\. An, J\. Li, J\. Yuan, J\. Du, J\. Liu, J\. Zhi, J\. Duan, K\. Zhou, K\. Wei, K\. Wang, K\. Luo, L\. Zhang, L\. Sha, L\. Xu, L\. Wu, L\. Ding, L\. Chen, M\. Li, N\. Lin, P\. Ta, Q\. Zou, R\. Song, R\. Yang, S\. Tu, S\. Yang, S\. Wu, S\. Zhang, S\. Li, S\. Li, S\. Fan, W\. Qin, W\. Tian, W\. Zhang, W\. Yu, W\. Liang, X\. Kuang, X\. Cheng, X\. Li, X\. Yan, X\. Hu, X\. Ling, X\. Fan, X\. Xia, X\. Zhang, X\. Zhang, X\. Pan, X\. Zou, X\. Zhang, Y\. Liu, Y\. Wu, Y\. Li, Y\. Wang, Y\. Zhu, Y\. Tan, Y\. Zhou, Y\. Pan, Y\. Zhang, Y\. Su, Y\. Geng, Y\. Yan, Y\. Tan, Y\. Bi, Y\. Shen, Y\. Yang, Y\. Li, Y\. Liu, Y\. Wang, Y\. Li, Y\. Wu, Y\. Zhang, Y\. Duan, Y\. Zhang, Z\. Liu, Z\. Jiang, Z\. Yan, Z\. Zhang, Z\. Wei, Z\. Chen, Z\. Feng, Z\. Yao, Z\. Chai, Z\. Wang, Z\. Zhang, B\. Xu, M\. Huang, H\. Wang, J\. Li, Y\. Dong, and J\. Tang\(2026\)GLM\-5: from vibe coding to agentic engineering\.External Links:2602\.15763,[Link](https://arxiv.org/abs/2602.15763)Cited by:[§A\.1\.1](https://arxiv.org/html/2608.11612#A1.SS1.SSS1.p2.11),[§1](https://arxiv.org/html/2608.11612#S1.p1.1)\.
- \[14\]A\. Grattafiori, A\. Dubey, A\. Jauhri, A\. Pandey, A\. Kadian, A\. Al\-Dahle, A\. Letman, A\. Mathur, A\. Schelten, A\. Vaughan, A\. Yang, A\. Fan, A\. Goyal, A\. Hartshorn, A\. Yang, A\. Mitra, A\. Sravankumar, A\. Korenev, A\. Hinsvark, A\. Rao, A\. Zhang, A\. Rodriguez, A\. Gregerson, A\. Spataru, B\. Roziere, B\. Biron, B\. Tang, B\. Chern, C\. Caucheteux, C\. Nayak, C\. Bi, C\. Marra, C\. McConnell, C\. Keller, C\. Touret, C\. Wu, C\. Wong, C\. C\. Ferrer, C\. Nikolaidis, D\. Allonsius, D\. Song, D\. Pintz, D\. Livshits, D\. Wyatt, D\. Esiobu, D\. Choudhary, D\. Mahajan, D\. Garcia\-Olano, D\. Perino, D\. Hupkes, E\. Lakomkin, E\. AlBadawy, E\. Lobanova, E\. Dinan, E\. M\. Smith, F\. Radenovic, F\. Guzmán, F\. Zhang, G\. Synnaeve, G\. Lee, G\. L\. Anderson, G\. Thattai, G\. Nail, G\. Mialon, G\. Pang, G\. Cucurell, H\. Nguyen, H\. Korevaar, H\. Xu, H\. Touvron, I\. Zarov, I\. A\. Ibarra, I\. Kloumann, I\. Misra, I\. Evtimov, J\. Zhang, J\. Copet, J\. Lee, J\. Geffert, J\. Vranes, J\. Park, J\. Mahadeokar, J\. Shah, J\. van der Linde, J\. Billock, J\. Hong, J\. Lee, J\. Fu, J\. Chi, J\. Huang, J\. Liu, J\. Wang, J\. Yu, J\. Bitton, J\. Spisak, J\. Park, J\. Rocca, J\. Johnstun, J\. Saxe, J\. Jia, K\. V\. Alwala, K\. Prasad, K\. Upasani, K\. Plawiak, K\. Li, K\. Heafield, K\. Stone, K\. El\-Arini, K\. Iyer, K\. Malik, K\. Chiu, K\. Bhalla, K\. Lakhotia, L\. Rantala\-Yeary, L\. van der Maaten, L\. Chen, L\. Tan, L\. Jenkins, L\. Martin, L\. Madaan, L\. Malo, L\. Blecher, L\. Landzaat, L\. de Oliveira, M\. Muzzi, M\. Pasupuleti, M\. Singh, M\. Paluri, M\. Kardas, M\. Tsimpoukelli, M\. Oldham, M\. Rita, M\. Pavlova, M\. Kambadur, M\. Lewis, M\. Si, M\. K\. Singh, M\. Hassan, N\. Goyal, N\. Torabi, N\. Bashlykov, N\. Bogoychev, N\. Chatterji, N\. Zhang, O\. Duchenne, O\. Çelebi, P\. Alrassy, P\. Zhang, P\. Li, P\. Vasic, P\. Weng, P\. Bhargava, P\. Dubal, P\. Krishnan, P\. S\. Koura, P\. Xu, Q\. He, Q\. Dong, R\. Srinivasan, R\. Ganapathy, R\. Calderer, R\. S\. Cabral, R\. Stojnic, R\. Raileanu, R\. Maheswari, R\. Girdhar, R\. Patel, R\. Sauvestre, R\. Polidoro, R\. Sumbaly, R\. Taylor, R\. Silva, R\. Hou, R\. Wang, S\. Hosseini, S\. Chennabasappa, S\. Singh, S\. Bell, S\. S\. Kim, S\. Edunov, S\. Nie, S\. Narang, S\. Raparthy, S\. Shen, S\. Wan, S\. Bhosale, S\. Zhang, S\. Vandenhende, S\. Batra, S\. Whitman, S\. Sootla, S\. Collot, S\. Gururangan, S\. Borodinsky, T\. Herman, T\. Fowler, T\. Sheasha, T\. Georgiou, T\. Scialom, T\. Speckbacher, T\. Mihaylov, T\. Xiao, U\. Karn, V\. Goswami, V\. Gupta, V\. Ramanathan, V\. Kerkez, V\. Gonguet, V\. Do, V\. Vogeti, V\. Albiero, V\. Petrovic, W\. Chu, W\. Xiong, W\. Fu, W\. Meers, X\. Martinet, X\. Wang, X\. Wang, X\. E\. Tan, X\. Xia, X\. Xie, X\. Jia, X\. Wang, Y\. Goldschlag, Y\. Gaur, Y\. Babaei, Y\. Wen, Y\. Song, Y\. Zhang, Y\. Li, Y\. Mao, Z\. D\. Coudert, Z\. Yan, Z\. Chen, Z\. Papakipos, A\. Singh, A\. Srivastava, A\. Jain, A\. Kelsey, A\. Shajnfeld, A\. Gangidi, A\. Victoria, A\. Goldstand, A\. Menon, A\. Sharma, A\. Boesenberg, A\. Baevski, A\. Feinstein, A\. Kallet, A\. Sangani, A\. Teo, A\. Yunus, A\. Lupu, A\. Alvarado, A\. Caples, A\. Gu, A\. Ho, A\. Poulton, A\. Ryan, A\. Ramchandani, A\. Dong, A\. Franco, A\. Goyal, A\. Saraf, A\. Chowdhury, A\. Gabriel, A\. Bharambe, A\. Eisenman, A\. Yazdan, B\. James, B\. Maurer, B\. Leonhardi, B\. Huang, B\. Loyd, B\. D\. Paola, B\. Paranjape, B\. Liu, B\. Wu, B\. Ni, B\. Hancock, B\. Wasti, B\. Spence, B\. Stojkovic, B\. Gamido, B\. Montalvo, C\. Parker, C\. Burton, C\. Mejia, C\. Liu, C\. Wang, C\. Kim, C\. Zhou, C\. Hu, C\. Chu, C\. Cai, C\. Tindal, C\. Feichtenhofer, C\. Gao, D\. Civin, D\. Beaty, D\. Kreymer, D\. Li, D\. Adkins, D\. Xu, D\. Testuggine, D\. David, D\. Parikh, D\. Liskovich, D\. Foss, D\. Wang, D\. Le, D\. Holland, E\. Dowling, E\. Jamil, E\. Montgomery, E\. Presani, E\. Hahn, E\. Wood, E\. Le, E\. Brinkman, E\. Arcaute, E\. Dunbar, E\. Smothers, F\. Sun, F\. Kreuk, F\. Tian, F\. Kokkinos, F\. Ozgenel, F\. Caggioni, F\. Kanayet, F\. Seide, G\. M\. Florez, G\. Schwarz, G\. Badeer, G\. Swee, G\. Halpern, G\. Herman, G\. Sizov, Guangyi, Zhang, G\. Lakshminarayanan, H\. Inan, H\. Shojanazeri, H\. Zou, H\. Wang, H\. Zha, H\. Habeeb, H\. Rudolph, H\. Suk, H\. Aspegren, H\. Goldman, H\. Zhan, I\. Damlaj, I\. Molybog, I\. Tufanov, I\. Leontiadis, I\. Veliche, I\. Gat, J\. Weissman, J\. Geboski, J\. Kohli, J\. Lam, J\. Asher, J\. Gaya, J\. Marcus, J\. Tang, J\. Chan, J\. Zhen, J\. Reizenstein, J\. Teboul, J\. Zhong, J\. Jin, J\. Yang, J\. Cummings, J\. Carvill, J\. Shepard, J\. McPhie, J\. Torres, J\. Ginsburg, J\. Wang, K\. Wu, K\. H\. U, K\. Saxena, K\. Khandelwal, K\. Zand, K\. Matosich, K\. Veeraraghavan, K\. Michelena, K\. Li, K\. Jagadeesh, K\. Huang, K\. Chawla, K\. Huang, L\. Chen, L\. Garg, L\. A, L\. Silva, L\. Bell, L\. Zhang, L\. Guo, L\. Yu, L\. Moshkovich, L\. Wehrstedt, M\. Khabsa, M\. Avalani, M\. Bhatt, M\. Mankus, M\. Hasson, M\. Lennie, M\. Reso, M\. Groshev, M\. Naumov, M\. Lathi, M\. Keneally, M\. Liu, M\. L\. Seltzer, M\. Valko, M\. Restrepo, M\. Patel, M\. Vyatskov, M\. Samvelyan, M\. Clark, M\. Macey, M\. Wang, M\. J\. Hermoso, M\. Metanat, M\. Rastegari, M\. Bansal, N\. Santhanam, N\. Parks, N\. White, N\. Bawa, N\. Singhal, N\. Egebo, N\. Usunier, N\. Mehta, N\. P\. Laptev, N\. Dong, N\. Cheng, O\. Chernoguz, O\. Hart, O\. Salpekar, O\. Kalinli, P\. Kent, P\. Parekh, P\. Saab, P\. Balaji, P\. Rittner, P\. Bontrager, P\. Roux, P\. Dollar, P\. Zvyagina, P\. Ratanchandani, P\. Yuvraj, Q\. Liang, R\. Alao, R\. Rodriguez, R\. Ayub, R\. Murthy, R\. Nayani, R\. Mitra, R\. Parthasarathy, R\. Li, R\. Hogan, R\. Battey, R\. Wang, R\. Howes, R\. Rinott, S\. Mehta, S\. Siby, S\. J\. Bondu, S\. Datta, S\. Chugh, S\. Hunt, S\. Dhillon, S\. Sidorov, S\. Pan, S\. Mahajan, S\. Verma, S\. Yamamoto, S\. Ramaswamy, S\. Lindsay, S\. Lindsay, S\. Feng, S\. Lin, S\. C\. Zha, S\. Patil, S\. Shankar, S\. Zhang, S\. Zhang, S\. Wang, S\. Agarwal, S\. Sajuyigbe, S\. Chintala, S\. Max, S\. Chen, S\. Kehoe, S\. Satterfield, S\. Govindaprasad, S\. Gupta, S\. Deng, S\. Cho, S\. Virk, S\. Subramanian, S\. Choudhury, S\. Goldman, T\. Remez, T\. Glaser, T\. Best, T\. Koehler, T\. Robinson, T\. Li, T\. Zhang, T\. Matthews, T\. Chou, T\. Shaked, V\. Vontimitta, V\. Ajayi, V\. Montanez, V\. Mohan, V\. S\. Kumar, V\. Mangla, V\. Ionescu, V\. Poenaru, V\. T\. Mihailescu, V\. Ivanov, W\. Li, W\. Wang, W\. Jiang, W\. Bouaziz, W\. Constable, X\. Tang, X\. Wu, X\. Wang, X\. Wu, X\. Gao, Y\. Kleinman, Y\. Chen, Y\. Hu, Y\. Jia, Y\. Qi, Y\. Li, Y\. Zhang, Y\. Zhang, Y\. Adi, Y\. Nam, Yu, Wang, Y\. Zhao, Y\. Hao, Y\. Qian, Y\. Li, Y\. He, Z\. Rait, Z\. DeVito, Z\. Rosnbrick, Z\. Wen, Z\. Yang, Z\. Zhao, and Z\. Ma\(2024\)The llama 3 herd of models\.External Links:2407\.21783,[Link](https://arxiv.org/abs/2407.21783)Cited by:[§A\.1](https://arxiv.org/html/2608.11612#A1.SS1.p1.1),[Appendix E](https://arxiv.org/html/2608.11612#A5.SSx2.p1.8)\.
- \[15\]E\. Grishina, M\. Smirnov, and M\. Rakhuba\(2026\)Accelerating newton\-schulz iteration for orthogonalization via chebyshev\-type polynomials\.External Links:2506\.10935,[Link](https://arxiv.org/abs/2506.10935)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p1.1)\.
- \[16\]W\. Guo, M\. Mishra, X\. Cheng, I\. Stoica, and T\. Dao\(2026\)SonicMoE: accelerating moe with io and tile\-aware optimizations\.External Links:2512\.14080,[Link](https://arxiv.org/abs/2512.14080)Cited by:[2nd item](https://arxiv.org/html/2608.11612#S2.I2.i2.p1.5)\.
- \[17\]N\. J\. Higham\(2008\)Functions of matrices: Theory and Computation\.edition,Society for Industrial and Applied Mathematics,\.External Links:[Document](https://dx.doi.org/10.1137/1.9780898717778),[Link](https://epubs.siam.org/doi/abs/10.1137/1.9780898717778)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p2.3)\.
- \[18\]T\. Joo, W\. Xia, C\. Kim, M\. Zhang, and E\. Ie\(2026\)On surprising effectiveness of masking updates in adaptive optimizers\.External Links:2602\.15322,[Link](https://arxiv.org/abs/2602.15322)Cited by:[§9](https://arxiv.org/html/2608.11612#S9.p2.4)\.
- \[19\]K\. Jordan, Y\. Jin, V\. Boza, J\. You, F\. Cesista, L\. Newhouse, and J\. Bernstein\(2024\)Muon: an optimizer for hidden layers in neural networks\.External Links:[Link](https://kellerjordan.github.io/posts/muon/)Cited by:[§2\.1](https://arxiv.org/html/2608.11612#S2.SS1.p1.3),[§2\.2](https://arxiv.org/html/2608.11612#S2.SS2.p1.22)\.
- \[20\]A\. Khaled, K\. Ozkara, T\. Yu, M\. Hong, and Y\. Park\(2026\)MuonBP: faster muon via block\-periodic orthogonalization\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=mHouLSUQP5)Cited by:[§3\.2](https://arxiv.org/html/2608.11612#S3.SS2.p4.1)\.
- \[21\]Kimi Team, Y\. Bai, Y\. Bao, Y\. Charles, C\. Chen, G\. Chen, H\. Chen, H\. Chen, J\. Chen, N\. Chen, R\. Chen, Y\. Chen, Y\. Chen, Y\. Chen, Z\. Chen, J\. Cui, H\. Ding, M\. Dong, A\. Du, C\. Du, D\. Du, Y\. Du, Y\. Fan, Y\. Feng, K\. Fu, B\. Gao, C\. Gao, H\. Gao, P\. Gao, T\. Gao, Y\. Ge, S\. Geng, Q\. Gu, X\. Gu, L\. Guan, H\. Guo, J\. Guo, X\. Hao, T\. He, W\. He, W\. He, Y\. He, C\. Hong, H\. Hu, Y\. Hu, Z\. Hu, W\. Huang, Z\. Huang, Z\. Huang, T\. Jiang, Z\. Jiang, X\. Jin, Y\. Kang, G\. Lai, C\. Li, F\. Li, H\. Li, M\. Li, W\. Li, Y\. Li, Y\. Li, Y\. Li, Z\. Li, Z\. Li, H\. Lin, X\. Lin, Z\. Lin, C\. Liu, C\. Liu, H\. Liu, J\. Liu, J\. Liu, L\. Liu, S\. Liu, T\. Y\. Liu, T\. Liu, W\. Liu, Y\. Liu, Y\. Liu, Y\. Liu, Y\. Liu, Z\. Liu, E\. Lu, H\. Lu, L\. Lu, Y\. Luo, S\. Ma, X\. Ma, Y\. Ma, S\. Mao, J\. Mei, X\. Men, Y\. Miao, S\. Pan, Y\. Peng, R\. Qin, Z\. Qin, B\. Qu, Z\. Shang, L\. Shi, S\. Shi, F\. Song, J\. Su, Z\. Su, L\. Sui, X\. Sun, F\. Sung, Y\. Tai, H\. Tang, J\. Tao, Q\. Teng, C\. Tian, C\. Wang, D\. Wang, F\. Wang, H\. Wang, H\. Wang, J\. Wang, J\. Wang, J\. Wang, S\. Wang, S\. Wang, S\. Wang, X\. Wang, Y\. Wang, Y\. Wang, Y\. Wang, Y\. Wang, Y\. Wang, Z\. Wang, Z\. Wang, Z\. Wang, Z\. Wang, C\. Wei, Q\. Wei, H\. Wu, W\. Wu, X\. Wu, Y\. Wu, C\. Xiao, J\. Xie, X\. Xie, W\. Xiong, B\. Xu, J\. Xu, L\. H\. Xu, L\. Xu, S\. Xu, W\. Xu, X\. Xu, Y\. Xu, Z\. Xu, J\. Xu, J\. Xu, J\. Yan, Y\. Yan, H\. Yang, X\. Yang, Y\. Yang, Y\. Yang, Z\. Yang, Z\. Yang, Z\. Yang, H\. Yao, X\. Yao, W\. Ye, Z\. Ye, B\. Yin, L\. Yu, E\. Yuan, H\. Yuan, M\. Yuan, S\. Yuan, H\. Zhan, D\. Zhang, H\. Zhang, W\. Zhang, X\. Zhang, Y\. Zhang, Y\. Zhang, Y\. Zhang, Y\. Zhang, Y\. Zhang, Y\. Zhang, Y\. Zhang, Y\. Zhang, Z\. Zhang, H\. Zhao, Y\. Zhao, Z\. Zhao, H\. Zheng, S\. Zheng, L\. Zhong, J\. Zhou, X\. Zhou, Z\. Zhou, J\. Zhu, Z\. Zhu, W\. Zhuang, and X\. Zu\(2025\)Kimi k2: open agentic intelligence\.External Links:2507\.20534,[Link](https://arxiv.org/abs/2507.20534)Cited by:[Appendix E](https://arxiv.org/html/2608.11612#A5.SSx1.p2.8),[§1](https://arxiv.org/html/2608.11612#S1.p1.1),[2nd item](https://arxiv.org/html/2608.11612#S2.I2.i2.p1.5),[§2\.3\.1](https://arxiv.org/html/2608.11612#S2.SS3.SSS1.p1.1)\.
- \[22\]S\. Lakić\(1998\)On the computation of the matrix k\-th root\.Zeitschrift für Angewandte Mathematik und Mechanik78\(3\),pp\. 167–172\.External Links:[Link](https://onlinelibrary.wiley.com/doi/abs/10.1002/%28SICI%291521-4001%28199803%2978%3A3%3C167%3A%3AAID-ZAMM167%3E3.0.CO%3B2-R)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p2.3)\.
- \[23\]Z\. Li, L\. Liu, C\. Liang, W\. Chen, and T\. Zhao\(2026\)NorMuon: making muon more efficient and scalable\.InForty\-third International Conference on Machine Learning,External Links:[Link](https://openreview.net/forum?id=m1IRWFAMsa)Cited by:[§2\.1](https://arxiv.org/html/2608.11612#S2.SS1.p3.3)\.
- \[24\]J\. Lim, S\. Lee, D\. Kim, T\. Kim, E\. Park, J\. Lee, J\. Lee, J\. Lee, W\. T\. Cheung, D\. Choi, J\. Her, J\. Huh, H\. Jung, C\. Kang, B\. Kim, M\. Kim, T\. Kim, Y\. Kim, H\. Kweon, H\. Lee, K\. Lee, D\. Oh, Y\. Park, B\. Ryu, and D\. Weon\(2025\)Motif 2 12\.7B technical report\.External Links:2511\.07464,[Link](https://arxiv.org/abs/2511.07464)Cited by:[2nd item](https://arxiv.org/html/2608.11612#S2.I1.i2.p1.1)\.
- \[25\]T\. Lin\(2025\)Flash\-muon: an efficient implementation of muon optimizer\.External Links:[Link](https://github.com/nil0x9/flash-muon)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p3.2)\.
- \[26\]J\. Liu, J\. Su, X\. Yao, Z\. Jiang, G\. Lai, Y\. Du, Y\. Qin, W\. Xu, E\. Lu, J\. Yan, Y\. Chen, H\. Zheng, Y\. Liu, S\. Liu, B\. Yin, W\. He, H\. Zhu, Y\. Wang, J\. Wang, M\. Dong, Z\. Zhang, Y\. Kang, H\. Zhang, X\. Xu, Y\. Zhang, Y\. Wu, X\. Zhou, and Z\. Yang\(2025\)Muon is scalable for llm training\.External Links:2502\.16982,[Link](https://arxiv.org/abs/2502.16982)Cited by:[§A\.1](https://arxiv.org/html/2608.11612#A1.SS1.p2.7),[§2\.3\.1](https://arxiv.org/html/2608.11612#S2.SS3.SSS1.p1.1)\.
- \[27\]J\. Liu\(2025\)A proof of concept for Distributed Muon\.GitHub\.Note:Pull Request \#1428, NVIDIA/Megatron\-LMAccessed: 2026\-07\-10External Links:[Link](https://github.com/NVIDIA/Megatron-LM/pull/1428)Cited by:[item 1](https://arxiv.org/html/2608.11612#S2.I3.i1.p1.1)\.
- \[28\]A\. Lozhkov, L\. Ben Allal, L\. von Werra, and T\. Wolf\(2024\)FineWeb\-edu: the finest collection of educational content\.Hugging Face\.External Links:[Link](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu),[Document](https://dx.doi.org/10.57967/hf/2497)Cited by:[§A\.1](https://arxiv.org/html/2608.11612#A1.SS1.p1.1)\.
- \[29\]I\. Modoranu, M\. Safaryan, E\. Schultheis, M\. Ryabinin, A\. Chumachenko, and D\. Alistarh\(2026\)Trion: FFT\-based dynamic subspace selection for low\-rank adaptive optimization of LLMs\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=TkHjRwbMNl)Cited by:[§3\.2](https://arxiv.org/html/2608.11612#S3.SS2.p2.4)\.
- \[30\]L\. Newhouse, D\. Goldberg, and R\. Ruiz\(2024\)Faster symmetric matrix multiplication with ThunderKittens\.External Links:[Link](https://www.lakernewhouse.com/assets/writing/faster-symmul-with-thunderkittens.pdf)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p3.2)\.
- \[31\]OpenAI, S\. Agarwal, L\. Ahmad, J\. Ai, S\. Altman, A\. Applebaum, E\. Arbus, R\. K\. Arora, Y\. Bai, B\. Baker, H\. Bao, B\. Barak, A\. Bennett, T\. Bertao, N\. Brett, E\. Brevdo, G\. Brockman, S\. Bubeck, C\. Chang, K\. Chen, M\. Chen, E\. Cheung, A\. Clark, D\. Cook, M\. Dukhan, C\. Dvorak, K\. Fives, V\. Fomenko, T\. Garipov, K\. Georgiev, M\. Glaese, T\. Gogineni, A\. Goucher, L\. Gross, K\. G\. Guzman, J\. Hallman, J\. Hehir, J\. Heidecke, A\. Helyar, H\. Hu, R\. Huet, J\. Huh, S\. Jain, Z\. Johnson, C\. Koch, I\. Kofman, D\. Kundel, J\. Kwon, V\. Kyrylov, E\. Y\. Le, G\. Leclerc, J\. P\. Lennon, S\. Lessans, M\. Lezcano\-Casado, Y\. Li, Z\. Li, J\. Lin, J\. Liss, L\. \(\. Liu, J\. Liu, K\. Lu, C\. Lu, Z\. Martinovic, L\. McCallum, J\. McGrath, S\. McKinney, A\. McLaughlin, S\. Mei, S\. Mostovoy, T\. Mu, G\. Myles, A\. Neitz, A\. Nichol, J\. Pachocki, A\. Paino, D\. Palmie, A\. Pantuliano, G\. Parascandolo, J\. Park, L\. Pathak, C\. Paz, L\. Peran, D\. Pimenov, M\. Pokrass, E\. Proehl, H\. Qiu, G\. Raila, F\. Raso, H\. Ren, K\. Richardson, D\. Robinson, B\. Rotsted, H\. Salman, S\. Sanjeev, M\. Schwarzer, D\. Sculley, H\. Sikchi, K\. Simon, K\. Singhal, Y\. Song, D\. Stuckey, Z\. Sun, P\. Tillet, S\. Toizer, F\. Tsimpourlas, N\. Vyas, E\. Wallace, X\. Wang, M\. Wang, O\. Watkins, K\. Weil, A\. Wendling, K\. Whinnery, C\. Whitney, H\. Wong, L\. Yang, Y\. Yang, M\. Yasunaga, K\. Ying, W\. Zaremba, W\. Zhan, C\. Zhang, B\. Zhang, E\. Zhang, and S\. Zhao\(2025\)Gpt\-oss\-120b & gpt\-oss\-20b model card\.External Links:2508\.10925,[Link](https://arxiv.org/abs/2508.10925)Cited by:[2nd item](https://arxiv.org/html/2608.11612#S2.I2.i2.p1.5)\.
- \[32\]Z\. Tang, T\. Xu, Y\. Saad, and Y\. Xi\(2026\)Hierarchical muon: tiled newton\-schulz updates for efficient muon optimization\.External Links:2606\.27216,[Link](https://arxiv.org/abs/2606.27216)Cited by:[§3\.2](https://arxiv.org/html/2608.11612#S3.SS2.p4.1)\.
- \[33\]B\. Xu\(2025\)Transpose one of the MLP matrices \+ add Triton kernel for symmetric matmul\.GitHub\.Note:Pull Request \#109, KellerJordan/modded\-nanogptAccessed: 2026\-07\-10External Links:[Link](https://github.com/KellerJordan/modded-nanogpt/pull/109)Cited by:[§3\.1](https://arxiv.org/html/2608.11612#S3.SS1.p3.2)\.
- \[34\]A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv, C\. Zheng, D\. Liu, F\. Zhou, F\. Huang, F\. Hu, H\. Ge, H\. Wei, H\. Lin, J\. Tang, J\. Yang, J\. Tu, J\. Zhang, J\. Yang, J\. Yang, J\. Zhou, J\. Zhou, J\. Lin, K\. Dang, K\. Bao, K\. Yang, L\. Yu, L\. Deng, M\. Li, M\. Xue, M\. Li, P\. Zhang, P\. Wang, Q\. Zhu, R\. Men, R\. Gao, S\. Liu, S\. Luo, T\. Li, T\. Tang, W\. Yin, X\. Ren, X\. Wang, X\. Zhang, X\. Ren, Y\. Fan, Y\. Su, Y\. Zhang, Y\. Zhang, Y\. Wan, Y\. Liu, Z\. Wang, Z\. Cui, Z\. Zhang, Z\. Zhou, and Z\. Qiu\(2025\)Qwen3 technical report\.External Links:2505\.09388,[Link](https://arxiv.org/abs/2505.09388)Cited by:[§A\.1](https://arxiv.org/html/2608.11612#A1.SS1.p1.1),[2nd item](https://arxiv.org/html/2608.11612#S2.I2.i2.p1.5)\.
- \[35\]Y\. Zhao, A\. Gu, R\. Varma, L\. Luo, C\. Huang, M\. Xu, L\. Wright, H\. Shojanazeri, M\. Ott, S\. Shleifer,et al\.\(2023\)PyTorch FSDP: experiences on scaling fully sharded data parallel\.Proceedings of the VLDB Endowment16\(12\),pp\. 3848–3860\.Cited by:[§7](https://arxiv.org/html/2608.11612#S7.SS0.SSS0.Px2.p1.4)\.

## Appendix AAlternative Experiments on Gram Newton\-Schulz

This section presents additional experimental results about Gram Newton\-Schulz and the symmetric GEMM kernels\. Complementing[Section˜8](https://arxiv.org/html/2608.11612#S8), these experiments study a wider range of architectures but at smaller scale, using only one GPU at a time\.

### A\.1Setup

We train four architectures: Llama\-430M, Qwen\-600M, Gemma\-1B, and a custom MoE architecture with 1B parameters, of which∼20%\{\\sim\}20\\%are active\[[14](https://arxiv.org/html/2608.11612#bib.bib31),[34](https://arxiv.org/html/2608.11612#bib.bib14),[12](https://arxiv.org/html/2608.11612#bib.bib32)\]\. We train on FineWeb\-Edu\[[28](https://arxiv.org/html/2608.11612#bib.bib35)\]\. The number of training tokens for each dense model is given by the Chinchilla scaling law; for MoE\-1B it is twice the Chinchilla scaling law with respect to the active parameters\. We use a cosine learning rate scheduler with the base learning rates given in[Table˜3](https://arxiv.org/html/2608.11612#A1.T3)\.

We use Muon on most matrix parameters—the attention layer’s projection matrices \(𝑾Q,𝑾K,𝑾V\{\\bm\{W\}\}\_\{Q\},\{\\bm\{W\}\}\_\{K\},\{\\bm\{W\}\}\_\{V\}\), the projection following attention \(𝑾O\{\\bm\{W\}\}\_\{O\}\), the SwiGLU MLP weights \(𝑾MLP,UP\{\\bm\{W\}\}\_\{\\mathrm\{MLP,UP\}\},𝑾MLP,GATE\{\\bm\{W\}\}\_\{\\mathrm\{MLP,GATE\}\},𝑾MLP,DOWN\{\\bm\{W\}\}\_\{\\mathrm\{MLP,DOWN\}\}\), and the token router of the MoE \(𝑾router\{\\bm\{W\}\}\_\{\\mathrm\{router\}\}\)—but exclude the embedding and unembedding layers\. Although we are not in the distributed setting, we use megabatching \([Section˜7](https://arxiv.org/html/2608.11612#S7)\), as batching Newton\-Schulz’s GEMM operations makes them faster\. As is standard for Muon, we adjust the learning rate by scaling the update for each weight matrix based on its dimensions\. For these experiments, we found that Moonshot AI’s strategy of scaling by0\.2​max⁡\(fan​\_​out,fan​\_​in\)0\.2\\sqrt\{\\max\(\\mathrm\{fan\\\_out\},\\mathrm\{fan\\\_in\}\)\}yields the best loss curves\[[26](https://arxiv.org/html/2608.11612#bib.bib7), §2\.2\]\.

Table 3:Architecture, base learning rate, and training\-token budget of each model\.#### A\.1\.1Splitting the Weights

We draw special attention to the fact that we split𝑾MLP,UP\{\\bm\{W\}\}\_\{\\mathrm\{MLP,UP\}\}from𝑾MLP,GATE\{\\bm\{W\}\}\_\{\\mathrm\{MLP,GATE\}\}and orthogonalize them separately\. Ordinary implementations of SwiGLU MLPs concatenate these into a single weight matrix; however, their contributions to the activation are fundamentally different, so gradients are different too\. We find that orthogonalizing them separately improves the final loss; for example, in Llama\-430M, we observe an improvement of≈0\.2\\approx 0\.2in perplexity\. For MoE architectures—whose intermediate size is typically smaller than the hidden size—separating them also reduces the FLOP cost of orthogonalization by a factor of two for standard Newton\-Schulz and even more for Gram Newton\-Schulz\. Likewise, while earlier implementations of Muon orthogonalized the combined matrix\[𝑾Q​\|𝑾K\|​𝑾V\]\\begin\{bmatrix\}\{\\bm\{W\}\}\_\{Q\}\\,\|\\,\{\\bm\{W\}\}\_\{K\}\\,\|\\,\{\\bm\{W\}\}\_\{V\}\\end\{bmatrix\}, we orthogonalize each piece separately\.

We are also aware that in some settings, including pretraining of GLM\-5, Muon benefits from splitting the Multi\-Latent Attention weights \(𝑾U​Q\{\\bm\{W\}\}^\{UQ\},𝑾U​K\{\\bm\{W\}\}^\{UK\}, and𝑾U​V\{\\bm\{W\}\}^\{UV\}\) by attention head before orthogonalizing\[[13](https://arxiv.org/html/2608.11612#bib.bib33)\]\. This choice is principled, since the actual function computed by attention treats each head separately; it does not see the concatenated weights\. Inspired by this, we experimented with splitting𝑾Q\{\\bm\{W\}\}\_\{Q\},𝑾K\{\\bm\{W\}\}\_\{K\},𝑾V\{\\bm\{W\}\}\_\{V\}, and𝑾O\{\\bm\{W\}\}\_\{O\}by attention head to formHHmatrices each of sizedH×d\\tfrac\{d\}\{H\}\\times d, whereddis the embedding dimension andHHis the number of heads\. However, this design led to higher losses throughout training, so we did not adopt it\.

Still, we believe that there are other settings like GLM\-5 where this strategy works well\. Such cases would benefit*immensely*from Gram Newton\-Schulz, since the aspect ratio of these weight matrices would be the number of headsHH\. For a standard attention weight like𝑾Q\{\\bm\{W\}\}\_\{Q\}withH=16H=16andT=5T=5, Gram Newton\-Schulz on the little matrices would use𝟖𝟎×\\mathbf\{80\\times\}fewer FLOPs than orthogonalizing the big matrix\!

### A\.2Kernelized Gram Newton\-Schulz preserves quality

We first verify that switching from standard Newton\-Schulz to our kernelized implementation of Gram Newton\-Schulz does not affect training quality\. We try each implementation with both the Polar Express coefficients\[[3](https://arxiv.org/html/2608.11612#bib.bib10)\]and those of\[[8](https://arxiv.org/html/2608.11612#bib.bib34)\]\. As[Figure˜7](https://arxiv.org/html/2608.11612#A1.F7)shows, the loss curves are identical in all cases, and the final validation perplexity is preserved to within0\.010\.01\.[Figure˜7](https://arxiv.org/html/2608.11612#A1.F7)uses a Hopper GPU, but we got the same results on Blackwell\.

We see loss preserved as follows, when both using the Polar Express coefficients and the coefficients derived by\[[8](https://arxiv.org/html/2608.11612#bib.bib34)\]:

![Refer to caption](https://arxiv.org/html/2608.11612v1/x9.png)Figure 7:When training with Muon on a Hopper GPU, switching from standard Newton\-Schulz to Gram Newton\-Schulz preserves the validation perplexity throughout training \(up to a 0\.01 difference in GNS’s favor\)\. Setup is described in[Section˜A\.1](https://arxiv.org/html/2608.11612#A1.SS1)\.
### A\.3Kernelized Gram Newton\-Schulz speeds up the optimizer step

##### Newton\-Schulz Performance

We observe that our method speeds up the runtime of the Newton\-Schulz step by1\.51\.5–2×2\\times\.[Figure˜8](https://arxiv.org/html/2608.11612#A1.F8)reports these speed\-ups for each model, benchmarked on both H100 and B300 GPUs\. As expected, savings are greatest for weights that are highly rectangular, like Gemma’s MLP weights \(α=8\\alpha=8\) and MoE\-1B’s expert weights \(α=4\\alpha=4\)\. Note that these experiments use standard Newton\-Schulz as the fallback whenm=nm=n\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/blog_png/BJy11AZoWe.png)

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/blog_png/Bk1YJCZj-x.png)

Figure 8:Newton\-Schulz time per model weight for \(1\) standard Newton\-Schulz implemented in pure PyTorch, \(2\) standard Newton\-Schulz with our symmetric GEMM kernels, and \(3\) Gram Newton\-Schulz with our kernels\. We test four architectures on a single Hopper \(top\) or Blackwell \(bottom\) GPU\. Square weights \(𝑾Q\{\\bm\{W\}\}\_\{Q\}or𝑾K\{\\bm\{W\}\}\_\{K\}\) benefit from the kernels\. Rectangular weights—especially those with a high aspect ratio \(e\.g\. Up/Gate in Gemma\-1B\)—additionally benefit from Gram Newton\-Schulz\.
##### End\-to\-End Optimizer Performance

[Figure˜9](https://arxiv.org/html/2608.11612#A1.F9)shows the end\-to\-end wall clock time of the Muon optimizer step for each version of Newton\-Schulz, along with that of AdamW\. For Muon, these timings include updating the momentum matrix, learning rate scaling, applying the weight update, and the AdamW updates for weights not assigned to Muon \(such as the embedding layer and the vector\-valued weights\)\. Our method yields a1\.31\.3–2×2\\timesspeedup, with both Gram Newton\-Schulz and the kernels contributing significantly\. As before, the impact of Gram Newton\-Schulz is greatest for architectures with highly rectangular weights\. Due to its MLP’s higher8×8\\timesaspect ratio, Gemma\-1B sees the largest speedup\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/blog_png/optim_step.png)Figure 9:Optimizer step time for AdamW and Muon with three different Newton\-Schulz routines on a single H100\. Gram Newton\-Schulz with our symmetric kernels gives a1\.31\.3–2×2\\timesspeedup over standard Newton\-Schulz\. Timings include matrix splitting and recombination for QKV and MLP, learning rate scaling, weight updates, and the scalar optimizer \(AdamW\) step for non\-2D weights\.
##### Estimating Gram Newton\-Schulz time in Kimi K2

Kimi K2 is a trillion parameter sparse, fine\-grained MoE model with384384experts per layer, a hidden size of71687168, and a small expert intermediate dimension of20482048\. Since models are trending towards finer\-grained MoE architectures and Kimi K2 was trained with Muon, this is a perfect setting to benchmark Gram Newton\-Schulz\. Due to Kimi’s sophisticated pipeline parallel strategy, the optimizer steps of many of its weights are completely hidden behind the backward pass of the next pipeline stage\. We estimate that orthogonalization steps of following weights are exposed:216216expert up/gate/down weights of shape2048×71682048\\times 7168and11dense up/gate/down weight of shape7168×184327168\\times 18432\. Therefore, we estimate the exposed Newton\-Schulz time by orthogonalizing weights of these shapes on a single GPU\.[Figure˜10](https://arxiv.org/html/2608.11612#A1.F10)shows the results: our method yields a2×2\\timesspeedup\. As before, both our kernels and the Gram Newton\-Schulz algorithm contribute significantly\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/blog_png/B1w5VCZjbe.png)

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/blog_png/HkQAVRZoWe.png)

Figure 10:Estimated exposed Newton\-Schulz time for one step of Kimi K2 with pipeline parallelism\. We measure the runtime of the exposed operations on a single Hopper \(top\) or Blackwell \(bottom\) GPU\. Gram Newton\-Schulz with kernels is2×2\\timesfaster than standard Newton\-Schulz\.

## Appendix BStability of Gram Newton\-Schulz

### B\.1Instability of Naive Gram Newton\-Schulz

We trained Llama\-430M with Muon using Naive Gram Newton\-Schulz \([Algorithm˜2](https://arxiv.org/html/2608.11612#alg2)\)\. The loss curve is shown in[Figure˜11](https://arxiv.org/html/2608.11612#A2.F11)\. Clearly, training is very unstable\. Not only do we observe loss spikes, but eventually, the outputs of Gram Newton\-Schulz become full of Infs\. The problem is due to floating point arithmetic\. While Gram Newton\-Schulz is mathematically equivalent to standard Newton\-Schulz in exact arithmetic, it behaves differently in finite precision, especially in half precision\. As half precision is essential for good performance, this makes Naive Gram Newton\-Schulz unworkable in practice\.

We now analyze the source of the instability in order to motivate our solution\. A Jupyter notebook reproducing the experiments in this section is available[here](https://github.com/NoahAmsel/PolarExpress/blob/appF-stability/gram_newton_schulz_stability.ipynb)\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/gns_instability/llama_430_no_reset.png)Figure 11:Naive Gram Newton\-Schulz used to train Llama\-430M in half precision\. Numerical instability wrecks training\.#### B\.1\.1Tracking Eigenvalues of Intermediate Matrices

To understand how the matrices in Naive Gram Newton\-Schulz evolve and why they diverge, we track their eigenvalues and singular values\. Recall that the entries of any matrix are upper bounded by its largest singular value, so if we can bound the singular values, we can prevent blowups\.

As a baseline, we start by running Naive Gram Newton\-Schulz in fullfloat64precision for88steps to simulate its behavior in exact arithmetic\. We use a synthetic128×512128\\times 512matrix with an exponentially decaying spectrum\. To make our plots more readable, the experiments in this section use the coefficients\(at,bt,ct\)=\(158,−108,38\)\(a\_\{t\},b\_\{t\},c\_\{t\}\)=\(\\tfrac\{15\}\{8\},\-\\tfrac\{10\}\{8\},\\tfrac\{3\}\{8\}\)at every iteration, but our conclusions will generalize to the coefficients used in practice\. If𝑿0=𝑼​𝚺​𝑽⊤\{\\bm\{X\}\}\_\{0\}=\{\\bm\{U\}\}\{\\bm\{\\Sigma\}\}\{\\bm\{V\}\}^\{\\top\}is the SVD of the input matrix, then the intermediate matrices of[Algorithm˜2](https://arxiv.org/html/2608.11612#alg2)\(𝑹t\{\\bm\{R\}\}\_\{t\},𝑸t\{\\bm\{Q\}\}\_\{t\},𝒁t\{\\bm\{Z\}\}\_\{t\}\) are square symmetric with eigenvectors𝑼\{\\bm\{U\}\}\. We therefore plot the diagonal entries of𝑼⊤​𝑹t​𝑼\{\\bm\{U\}\}^\{\\top\}\{\\bm\{R\}\}\_\{t\}\{\\bm\{U\}\}and𝑼⊤​𝑸t​𝑼\{\\bm\{U\}\}^\{\\top\}\{\\bm\{Q\}\}\_\{t\}\{\\bm\{U\}\}against the corresponding singular values in𝚺\{\\bm\{\\Sigma\}\}to track how each evolves according to the polynomial update rules—or diverges from them\. Even though Gram Newton\-Schulz does not need to compute𝑿1,…,𝑿T−1\{\\bm\{X\}\}\_\{1\},\\ldots,\{\\bm\{X\}\}\_\{T\-1\}, we do so here for demonstration using the formula𝑿t=𝑸t​𝑿0\{\\bm\{X\}\}\_\{t\}=\{\\bm\{Q\}\}\_\{t\}\{\\bm\{X\}\}\_\{0\}and plot𝑼⊤​𝑿t​𝑽\{\\bm\{U\}\}^\{\\top\}\{\\bm\{X\}\}\_\{t\}\{\\bm\{V\}\}against𝚺\{\\bm\{\\Sigma\}\}\. The top panel of[Figure˜12](https://arxiv.org/html/2608.11612#A2.F12)shows that the eigenvalues evolve as predicted by[Theorem˜2](https://arxiv.org/html/2608.11612#Thmtheorem2)\. Initially,x0∈\[0,1\]x\_\{0\}\\in\[0,1\],r0=x02∈\[0,1\]r\_\{0\}=x\_\{0\}^\{2\}\\in\[0,1\]andq0=1q\_\{0\}=1\. As the algorithm progresses,rt→1r\_\{t\}\\to 1,qt=xt/x0→1/x0q\_\{t\}=x\_\{t\}/x\_\{0\}\\to 1/x\_\{0\}, andxt→1x\_\{t\}\\to 1\. Forx0≈1x\_\{0\}\\approx 1, convergence is faster; forx0≪1x\_\{0\}\\ll 1, it is slower\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x10.png)

![Refer to caption](https://arxiv.org/html/2608.11612v1/x11.png)

Figure 12:Evolution of eigenvalues of𝑹t\{\\bm\{R\}\}\_\{t\},𝑸t\{\\bm\{Q\}\}\_\{t\}, and𝑿t\{\\bm\{X\}\}\_\{t\}in Naive Gram Newton\-Schulz with\(at,bt,ct\)=\(15/8,10/8,3/8\)\(a\_\{t\},b\_\{t\},c\_\{t\}\)=\(\\nicefrac\{\{15\}\}\{\{8\}\},\\nicefrac\{\{10\}\}\{\{8\}\},\\nicefrac\{\{3\}\}\{\{8\}\}\)as a function of the corresponding singular value of𝑿0\{\\bm\{X\}\}\_\{0\}\.Top:float64\.Bottom:bfloat16\.Now we rerun the experiment inbfloat16arithmetic \([Figure˜12](https://arxiv.org/html/2608.11612#A2.F12), bottom\)\. The first few iterations look correct, but by step 7, we see noisy and unexpected behavior \(rt<0r\_\{t\}<0,xt\>1x\_\{t\}\>1\)\. By step 8, the spectrum diverges completely, and by step 10 the algorithm returns Infs\. We now describe two key causes of divergence: spurious negative eigenvalues in the Gram matrix𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}, and eigenvector drift\.

#### B\.1\.2Spurious Negative Eigenvalues

The leading cause of divergence is the presence of negative eigenvalues in the Gram matrix due to half\-precision arithmetic\. These negative eigenvalues blow up after too many iterations of Gram Newton\-Schulz\.

By construction,rt=xt2≥0r\_\{t\}=x\_\{t\}^\{2\}\\geq 0, so𝑹t\{\\bm\{R\}\}\_\{t\}should be positive semidefinite \([Theorem˜2](https://arxiv.org/html/2608.11612#Thmtheorem2)\), but[Figure˜12](https://arxiv.org/html/2608.11612#A2.F12)shows that is has negative eigenvalues\. In fact, even𝑹0\{\\bm\{R\}\}\_\{0\}has tiny negative eigenvalues introduced in the first matrix multiplication𝑿0​𝑿0⊤\{\\bm\{X\}\}\_\{0\}\{\\bm\{X\}\}\_\{0\}^\{\\top\}\. Because𝑿0\{\\bm\{X\}\}\_\{0\}has many singular values that are nearly zero \(as is the case in Muon\),𝑹0=𝑿0​𝑿0⊤\{\\bm\{R\}\}\_\{0\}=\{\\bm\{X\}\}\_\{0\}\{\\bm\{X\}\}\_\{0\}^\{\\top\}has many eigenvalues that are*numerically equal*to zero\. Inbfloat16, even a slightly negative number can be numerically equal to zero\. Later computations can introduce additional negative eigenvalues into𝑹t\{\\bm\{R\}\}\_\{t\}as well\. These eigenvalues represent nothing about the original problem; they are just an artifact of floating point arithmetic\. Therefore, we call them “spurious eigenvalues”\.

These spurious negative eigenvalues start small, but[Figure˜12](https://arxiv.org/html/2608.11612#A2.F12)shows that their magnitude grows quickly\. Recall the update rule:

rt=rt−1​zt2=rt−1​ht​\(rt−1\)2r\_\{t\}=r\_\{t\-1\}z\_\{t\}^\{2\}=r\_\{t\-1\}h\_\{t\}\(r\_\{t\-1\}\)^\{2\}\(4\)For our choice of coefficients,ht​\(x\)=158−108​x\+38​x2h\_\{t\}\(x\)=\\tfrac\{15\}\{8\}\-\\tfrac\{10\}\{8\}x\+\\tfrac\{3\}\{8\}x^\{2\}\.[Figure˜13](https://arxiv.org/html/2608.11612#A2.F13)plots this update rule\. As it shows,rt<\(158\)2​rt−1r\_\{t\}<\\left\(\\tfrac\{15\}\{8\}\\right\)^\{2\}r\_\{t\-1\}\. Thus, if anyrt<0r\_\{t\}<0, the spurious negative eigenvalues grow exponentially, diverging to−∞\-\\infty\! This sets off a chain reaction that causes𝑸t\{\\bm\{Q\}\}\_\{t\}and𝑿t\{\\bm\{X\}\}\_\{t\}to diverge as well\. This problem cannot be fixed by choosing different polynomials; while the main loop of Gram Newton\-Schulz approximates the inverse square root ofr0\>0r\_\{0\}\>0, it diverges for negative inputs\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x12.png)Figure 13:When𝑹t\{\\bm\{R\}\}\_\{t\}evolves according to \([4](https://arxiv.org/html/2608.11612#A2.E4)\), negative eigenvalues diverge to−∞\-\\infty\.To prove that the tiny spurious negative eigenvalues of𝑹0\{\\bm\{R\}\}\_\{0\}suffice to cause divergence, we rerun the experiment infloat64precision, but cast𝑹0\{\\bm\{R\}\}\_\{0\}fromfloat64tobfloat16and then back tofloat64to induce a small floating point error\. As[Figure˜14](https://arxiv.org/html/2608.11612#A2.F14)shows, this is enough to cause a blowup\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x13.png)Figure 14:Evolution of eigenvalues of𝑹t\{\\bm\{R\}\}\_\{t\},𝑸t\{\\bm\{Q\}\}\_\{t\}, and𝑿t\{\\bm\{X\}\}\_\{t\}when all operations usefloat64except𝑹0=𝑿0​𝑿0⊤\{\\bm\{R\}\}\_\{0\}=\{\\bm\{X\}\}\_\{0\}\{\\bm\{X\}\}\_\{0\}^\{\\top\}, which usesbfloat16\.
#### B\.1\.3Eigenvector Drift

Spurious negative eigenvalues are not the only source of numerical instability\. If we take the input matrix𝑿0\{\\bm\{X\}\}\_\{0\}to have no small singular values \(i\.e\., all≥0\.017\\geq 0\.017\), then we do not observe any negative eigenvalues in𝑹t\{\\bm\{R\}\}\_\{t\}, but𝑿t\{\\bm\{X\}\}\_\{t\}still fails to converge\. The culprit seems to be “eigenvector drift”\.

In exact arithmetic, the eigenvectors of all intermediate matrices match𝑼\{\\bm\{U\}\}, the left singular vectors of𝑿0\{\\bm\{X\}\}\_\{0\}, but in finite precision they do not\. We demonstrate this effect by measuring how far𝑼⊤​𝑹t​𝑼\{\\bm\{U\}\}^\{\\top\}\{\\bm\{R\}\}\_\{t\}\{\\bm\{U\}\},𝑼⊤​𝑸t​𝑼\{\\bm\{U\}\}^\{\\top\}\{\\bm\{Q\}\}\_\{t\}\{\\bm\{U\}\}, and𝑼⊤​𝑿t​𝑽\{\\bm\{U\}\}^\{\\top\}\{\\bm\{X\}\}\_\{t\}\{\\bm\{V\}\}are from being diagonal matrices\. That is, we take the Frobenius norm of the off\-diagonal entries as a fraction of the total Frobenius norm\.[Figure˜15](https://arxiv.org/html/2608.11612#A2.F15)shows that after several iterations, the eigenvectors / singular vectors of𝑸t\{\\bm\{Q\}\}\_\{t\}and𝑿t\{\\bm\{X\}\}\_\{t\}have drifted significantly from those of𝑿0\{\\bm\{X\}\}\_\{0\}\. At the same time, the eigen*values*of𝑸t\{\\bm\{Q\}\}\_\{t\}\(and by extension, those of𝑿t\{\\bm\{X\}\}\_\{t\}\) diverge from where they would be in exact arithmetic\. The growing eigenvalues of𝑸t\{\\bm\{Q\}\}\_\{t\}seem to spill into one another\. The strength of this effect is less consistent than that of negative eigenvalues, but it is still harmful\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x14.png)Figure 15:As the eigenvectors drift, the spectral norms of𝑹t\{\\bm\{R\}\}\_\{t\},𝑸t\{\\bm\{Q\}\}\_\{t\}, and𝑿t\{\\bm\{X\}\}\_\{t\}diverge\.

### B\.2Stabilizing Gram Newton\-Schulz by Restarting

As we have seen, if we run Gram Newton\-Schulz for more than a few iterations, the spurious negative eigenvalues of𝑹t\{\\bm\{R\}\}\_\{t\}diverge to negative infinity and𝑸t\{\\bm\{Q\}\}\_\{t\}blows up\. Our solution is simple: run Gram Newton\-Schulz for only a few iterations\. For instance, rather than using Gram Newton\-Schulz to compute𝑿T\{\\bm\{X\}\}\_\{T\}directly, we can use it to compute𝑿5\{\\bm\{X\}\}\_\{5\}in a stable manner\. While𝑿5\{\\bm\{X\}\}\_\{5\}is not a good approximation topolar⁡\(𝑿0\)\\operatorname\{polar\}\(\{\\bm\{X\}\}\_\{0\}\), it is closer than where we started\. Now we apply Gram Newton\-Schulz a second time on the input𝑿5\{\\bm\{X\}\}\_\{5\}to compute𝑿10\{\\bm\{X\}\}\_\{10\}stably\. This process can be repeated to reach any desired𝑿T\{\\bm\{X\}\}\_\{T\}\. This restarting technique sacrifices some of the performance gains of Gram Newton\-Schulz, but it still offers a significant speedup over standard Newton\-Schulz\.

[Figure˜16](https://arxiv.org/html/2608.11612#A2.F16)repeats the experiment from above with a restart every five iterations\. While𝑹t\{\\bm\{R\}\}\_\{t\}develops some negative eigenvalues, unlike before, the growth of these eigenvalues is controlled\. Each time we restart, we re\-initialize𝑹t=𝑿t​𝑿t⊤\{\\bm\{R\}\}\_\{t\}=\{\\bm\{X\}\}\_\{t\}\{\\bm\{X\}\}\_\{t\}^\{\\top\}, eliminating any large negative eigenvalues\. By restarting regularly, we prevent any from growing too large\. As expected,𝑸t\{\\bm\{Q\}\}\_\{t\}resets to the identity at iterations5,10,15,205,10,15,20, and2525\. Therefore, the eigenvalues of𝑸t\{\\bm\{Q\}\}\_\{t\}never grow beyond≈12\\approx 12, despite the negative eigenvalues in𝑹t\{\\bm\{R\}\}\_\{t\}\. Since the eigenvalues of𝑸t\{\\bm\{Q\}\}\_\{t\}remain controlled, those of𝑿t=𝑸t​𝑿t−5\{\\bm\{X\}\}\_\{t\}=\{\\bm\{Q\}\}\_\{t\}\{\\bm\{X\}\}\_\{t\-5\}stay at or below11\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x15.png)Figure 16:Restarting prevents the divergence of𝑹t\{\\bm\{R\}\}\_\{t\}in half\-precision\.Restarting also helps control eigenvector drift\. We repeat the experiment from[Figure˜15](https://arxiv.org/html/2608.11612#A2.F15)on the same matrix \(with all singular values\>0\.017\>0\.017\) with a restart after step55\. Diagonalization error always remains≤0\.05\\leq 0\.05, and the maximum eigenvalues now closely track their predicted values\. Note that we always measure eigenvector drift relative to the original input𝑿0\{\\bm\{X\}\}\_\{0\}, not the restarted𝑿5\{\\bm\{X\}\}\_\{5\}\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x16.png)Figure 17:Restarting curbs eigenvector drift\.#### B\.2\.1When to Restart: Polar Express Coefficients for Muon

We now derive the optimal restarting schedule for a given total number of iterationsTT\. To prevent𝑿t=𝑸t​𝑿0\{\\bm\{X\}\}\_\{t\}=\{\\bm\{Q\}\}\_\{t\}\{\\bm\{X\}\}\_\{0\}from blowing up, we must control the condition number of𝑸t\{\\bm\{Q\}\}\_\{t\}, even when𝑹0\{\\bm\{R\}\}\_\{0\}has spurious negative eigenvalues\. \(Becauseqt→1/x0≥1q\_\{t\}\\to 1/x\_\{0\}\\geq 1, this also controls the maximum eigenvalue of𝑸t\{\\bm\{Q\}\}\_\{t\}\.\) The growth of𝑸t\{\\bm\{Q\}\}\_\{t\}in turn depends on the size of the spurious negative eigenvalues and the specific sequence of polynomials defined by\{\(at,bt,ct\)\}t=1T\\\{\(a\_\{t\},b\_\{t\},c\_\{t\}\)\\\}\_\{t=1\}^\{T\}\. Fixing a desired number of restarts, we sweep over all possible restart schedules and pick the one that minimizes the maximum condition number of𝑸t\{\\bm\{Q\}\}\_\{t\}acrosstts\.

For the application to Muon, we analyze five iterations of the Polar Express coefficients\[[3](https://arxiv.org/html/2608.11612#bib.bib10)\]\. In the experiments of the previous section, we observe that the most negative spurious eigenvalue of𝑹0\{\\bm\{R\}\}\_\{0\}is about−4⋅10−4\-4\\cdot 10^\{\-4\}\. Therefore, we simulate Gram Newton\-Schulz in full precision with Polar Express coefficients and track how the eigenvalues of𝑹t\{\\bm\{R\}\}\_\{t\}and𝑸t\{\\bm\{Q\}\}\_\{t\}evolve when𝑹0\{\\bm\{R\}\}\_\{0\}has eigenvalues in the range\[−4⋅10−4,1\]\[\-4\\cdot 10^\{\-4\},1\]\. The left panels of[Figure˜18](https://arxiv.org/html/2608.11612#A2.F18)show that without restarting, they blow up\. We then repeat the simulation with one restart, sweeping all possible choices of when to restart\. Every time we restart and form𝑹=𝑿​𝑿⊤\{\\bm\{R\}\}=\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}, we subtract4⋅10−4​𝑰4\\cdot 10^\{\-4\}\{\\bm\{I\}\}to simulate potentially dangerous corruption of the eigenvalues due to floating point error\. As the right panel shows, restarting after the*second*iteration provides the best bounds, ensuring that the eigenvalues of𝑹t\{\\bm\{R\}\}\_\{t\}stay well above−0\.4\-0\.4and those of𝑸t\{\\bm\{Q\}\}\_\{t\}stay below≈100\\approx 100for all iterations\.[Figure˜19](https://arxiv.org/html/2608.11612#A2.F19)shows that for Gram Newton\-Schulz with Polar Express coefficients and one restart after the second iteration, all eigenvalues of𝑿t\{\\bm\{X\}\}\_\{t\}converge stably to11\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x17.png)

![Refer to caption](https://arxiv.org/html/2608.11612v1/x18.png)

Figure 18:Left:Min/max eigenvalue of𝑹t\{\\bm\{R\}\}\_\{t\}and𝑸t\{\\bm\{Q\}\}\_\{t\}without restarts\.𝑹0\{\\bm\{R\}\}\_\{0\}starts with a negative eigenvalue at−4⋅10−4\-4\\cdot 10^\{\-4\}\.Right:Minimum eigenvalue of𝑹t\{\\bm\{R\}\}\_\{t\}and condition number of𝑸t\{\\bm\{Q\}\}\_\{t\}if restart is placed after iteration11,22,33, or44\. Restarting after iteration22best controls the size of𝑸t\{\\bm\{Q\}\}\_\{t\}\.![Refer to caption](https://arxiv.org/html/2608.11612v1/x19.png)Figure 19:Gram Newton\-Schulz with Polar Express coefficients and a restart after22iterations converges stably\.
#### B\.2\.2Further Precautions

While restarting greatly improves stability, it is not absolutely foolproof\. A second or third restart may be required if running for more than five iterations or using a numerically sensitive set of coefficients\. Moreover, the usual numerical precautions for standard Newton\-Schulz still apply\.

##### Safety Factors

Most choices of Newton\-Schulz polynomials are designed to converge only when the input lies in\[0,1\]\[0,1\]; any singular values larger than11may diverge rapidly, as[Figure˜20](https://arxiv.org/html/2608.11612#A2.F20)shows\. Even when𝑿0\{\\bm\{X\}\}\_\{0\}is properly normalized, singular values greater than11can arise due to numerical error\. This problem affects standard Newton\-Schulz too, so the Polar Express polynomials are typically adjusted according to the formulap~t​\(x\)=pt​\(x/1\.02\)\\tilde\{p\}\_\{t\}\(x\)=p\_\{t\}\(x/1\.02\)\. This ensures convergence even for singular values as large as1\.021\.02\. When using Gram Newton\-Schulz, roundoff errors like this can worsen due to computations like𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}, which do not have such numerical buffers; however, we have never seen this happen when using our recommended setup \(float16arithmetic with restarting after22iterations\)\. It is wise to be conservative in the choice of safety factor, for instance, by replacing1\.021\.02with1\.051\.05\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/figures/gns_instability/X_final_unbounded.png)Figure 20:Theoretical behavior of both standard and Gram Newton\-Schulz \(Polar Express coefficients\) when𝑿0\{\\bm\{X\}\}\_\{0\}has singular values slightly above one\.
##### Float16 vs BFloat16 in Newton\-Schulz

In addition, we argue for usingfloat16instead ofbfloat16to implement Newton\-Schulz\. Both use1616bits and thus run with the same performance; however, the distribution of the1616bits across the mantissa and exponent bits differs\. Compared tobfloat16,float16represents values from a narrower range, but it has greater precision within that range\. For Newton\-Schulz, however, the range offloat16\(roughly6\.1⋅10−56\.1\\cdot 10^\{\-5\}to6\.5⋅1046\.5\\cdot 10^\{4\}\) suffices because the magnitudes of its intermediate matrices are not too large\. On certain test matrices, we see more accuratepolar⁡\(𝑿\)\\operatorname\{polar\}\(\{\\bm\{X\}\}\)approximations withfloat16, but in practice, we have not found a case where the training quality is meaningfully different betweenfloat16andbfloat16\. Still, we default tofloat16\.

##### High Accuracy Setting

In the application to Muon, we do not need to computepolar⁡\(𝑿\)\\operatorname\{polar\}\(\{\\bm\{X\}\}\)to high accuracy, and[Section˜A\.2](https://arxiv.org/html/2608.11612#A1.SS2)shows that Muon with Gram Newton\-Schulz yields effectively identical results to Muon with standard Newton\-Schulz in terms of training quality\. However, when high accuracy is desired, the usual warnings about forming the Gram matrix apply\. Since forming𝑿​𝑿⊤\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}immediately squares the condition number, Gram Newton\-Schulz may not be appropriate in these cases\.

#### B\.2\.3Computing Matrix Quadratics

A key step in Gram Newton\-Schulz is computing the matrix quadratic𝒁t←at​𝑰\+bt​𝑹t−1\+ct​𝑹t−12\{\\bm\{Z\}\}\_\{t\}\\leftarrow a\_\{t\}\{\\bm\{I\}\}\+b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}\. Standard Newton\-Schulz implicitly computes matrix quadratics too, but PyTorch implementations typically avoid adding the identity explicitly\. Rather, they compute𝑿​\(at​𝑰\+bt​𝑨\+ct​𝑨2\)\{\\bm\{X\}\}\(a\_\{t\}\{\\bm\{I\}\}\+b\_\{t\}\{\\bm\{A\}\}\+c\_\{t\}\{\\bm\{A\}\}^\{2\}\)in two steps, each with a single GEMM:

1. 1\.𝑩←bt​𝑨\+ct​𝑨2\{\\bm\{B\}\}\\leftarrow b\_\{t\}\{\\bm\{A\}\}\+c\_\{t\}\{\\bm\{A\}\}^\{2\}
2. 2\.𝑿←at​𝑿\+𝑩​𝑿\{\\bm\{X\}\}\\leftarrow a\_\{t\}\{\\bm\{X\}\}\+\{\\bm\{B\}\}\{\\bm\{X\}\}

Our symmetric GEMM kernel can compute matrix quadratics in one launch, fusing the addition ofat​𝑰a\_\{t\}\{\\bm\{I\}\}by addingata\_\{t\}to all diagonal entries of the output when they are at the register level\. This optimization completely obviates any I/O operations needed for theat​𝑰a\_\{t\}\{\\bm\{I\}\}addition, typically outspeedinggemm\_symmetric\(A, B, C, c\_t, b\_t\) \+ a\_t \* I, which would require loading𝑰\{\\bm\{I\}\}from general memory to shared memory to registers\. Once𝒁t\{\\bm\{Z\}\}\_\{t\}is assembled, Gram Newton\-Schulz can perform its three subsequent multiplications without the need for further fused additions\.

However, our tests show that addingat​𝑰a\_\{t\}\{\\bm\{I\}\}explicitly can be less stable than distributing it into those later three multiplications\. If we stress\-test our method by ignoring some of our own advice—restarting after three iterations instead of two, using a Polar Express safety factor of1\.021\.02instead of1\.051\.05, and computing the quadratic withat​𝑰a\_\{t\}\{\\bm\{I\}\}explicitly—we observe instability\. Interestingly, this instability disappears if we use non\-symmetric GEMMs \(either from PyTorch or Quack\) instead of our symmetric kernels; however, if we force symmetry after calling standard PyTorch GEMMs, we see instability again\. We conclude that fusing\+at​𝑰\+a\_\{t\}\{\\bm\{I\}\}into a symmetric GEMM is numerically unfavorable; this is not just a kernel bug\.

We believe this effect can be explained as follows\. While the fused kernel computesat​𝑰\+bt​𝑹t−1\+ct​𝑹t−12a\_\{t\}\{\\bm\{I\}\}\+b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}infloat32arithmetic under the hood, the result𝒁t\{\\bm\{Z\}\}\_\{t\}is rounded back down tofloat16at the end of the GEMM\. Future computations like𝑸t​𝒁t\{\\bm\{Q\}\}\_\{t\}\{\\bm\{Z\}\}\_\{t\}suffer from this loss of precision inata\_\{t\}\. In contrast, if theat​𝑰a\_\{t\}\{\\bm\{I\}\}term is handled implicitly, all arithmetic involvingata\_\{t\}takes place infloat32\. Therefore, it is more stable to computeat​𝑸t\+𝑸t​\(bt​𝑹t−1\+ct​𝑹t−12\)a\_\{t\}\{\\bm\{Q\}\}\_\{t\}\+\{\\bm\{Q\}\}\_\{t\}\\left\(b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}\\right\)than𝑸t​\(at​𝑰\+bt​𝑹t−1\+ct​𝑹t−12\)\{\\bm\{Q\}\}\_\{t\}\\left\(a\_\{t\}\{\\bm\{I\}\}\+b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}\\right\)\.

We reiterate that in all our experiments, this instability can be avoided entirely by restarting correctly or using a safety factor of1\.051\.05\. Out of an abundance of caution, we rearrange Naive Gram Newton\-Schulz \([Algorithm˜2](https://arxiv.org/html/2608.11612#alg2)\) to avoid addingat​𝑰a\_\{t\}\{\\bm\{I\}\}explicitly\. We change

1. 1\.𝒁t←at​𝑰\+bt​𝑹t−1\+ct​𝑹t−12\{\\bm\{Z\}\}\_\{t\}\\leftarrow a\_\{t\}\{\\bm\{I\}\}\+b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}// Applyht​\(𝑹t−1\)h\_\{t\}\(\{\\bm\{R\}\}\_\{t\-1\}\)
2. 2\.𝑸t←𝑸t−1​𝒁t\{\\bm\{Q\}\}\_\{t\}\\leftarrow\{\\bm\{Q\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}
3. 3\.𝑹t←𝒁t​𝑹t−1​𝒁t\{\\bm\{R\}\}\_\{t\}\\leftarrow\{\\bm\{Z\}\}\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}

to

1. 1\.𝒁t←bt​𝑹t−1\+ct​𝑹t−12\{\\bm\{Z\}\}\_\{t\}\\leftarrow b\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}\+c\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}^\{2\}
2. 2\.𝑸t←𝑸t−1​𝒁t\+at​𝑸t−1\{\\bm\{Q\}\}\_\{t\}\\leftarrow\{\\bm\{Q\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}\+a\_\{t\}\{\\bm\{Q\}\}\_\{t\-1\}
3. 3\.\(𝐑𝐙\)t←𝑹t−1​𝒁t\+at​𝑹t−1\(\\mathbf\{RZ\}\)\_\{t\}\\leftarrow\{\\bm\{R\}\}\_\{t\-1\}\{\\bm\{Z\}\}\_\{t\}\+a\_\{t\}\{\\bm\{R\}\}\_\{t\-1\}
4. 4\.𝑹t←𝒁t​\(𝐑𝐙\)t\+at​\(𝐑𝐙\)t\{\\bm\{R\}\}\_\{t\}\\leftarrow\{\\bm\{Z\}\}\_\{t\}\(\\mathbf\{RZ\}\)\_\{t\}\+a\_\{t\}\(\\mathbf\{RZ\}\)\_\{t\}

This change fixes all collected training examples in which symmetric GEMMs were less stable than non\-symmetric GEMMs\.

## Appendix CKernel Implementation Details

### C\.1Symmetric GEMM Kernel Details

We implement all of our symmetric GEMM kernels with square cluster work tiles\. Hopper uses cluster size\(2,1\)\(2,1\)and thread block tile size\(128,256\)\(128,256\), and Blackwell uses cluster size\(2,1\)\(2,1\)and 2\-CTA collaboration, in which the 2 thread blocks in the cluster collaborate on the same big\(256,256\)\(256,256\)tile\. Notably, highly optimized custom GEMM kernels on Hopper typically use Ping Pong Scheduling, in which the MMA of tileiiand the epilogue of tilei−1i\-1are overlapped in two consumer warp groups\.555[https://pytorch\.org/blog/cutlass\-ping\-pong\-gemm\-kernel/](https://pytorch.org/blog/cutlass-ping-pong-gemm-kernel/)However, Ping Pong Scheduling uses more registers at once, and\(128,256\)\(128,256\)is too large of a tile size for Ping Pong Scheduling, leading to register spillage\. This is much slower than standard single producer warp, single consumer warp scheduling\. Thus, our Hopper symmetric kernels do not use Ping Pong Scheduling\. Blackwell GEMM kernels have no explicit conception of Ping Pong Scheduling, since by default in cuBLAS and most kernel libraries, two accumulators are kept in the new tensor memory hierarchy, and MMA is computed on one accumulator while the epilogue is computed on the other\.

As a small implementation detail, note that the main diagonal of256×256256\\times 256cluster work tiles is part of the work assigned by the triangular scheduler\. Since their transposed locations are identical to their current locations, we only write those values to general memory once—writing twice can cause inaccurate values or NaNs\.

### C\.2Implementation Strategy in Code

There are only two differences between the symmetric GEMM kernel and the standard GEMM kernel: the triangular scheduler and the transposed tile write in the epilogue\. We design our symmetric kernel such that it is abstracted around the standard GEMM kernel to enable lightweight but maximally performant GEMM epilogue fusions\. Using these abstractions, we are able to implement the symmetric GEMM kernel for both Hopper and Blackwell in just 160 lines, while achieving state\-of\-the\-art performance\.

We override the standard tile scheduler with our triangular scheduler and wrap the symmetric GEMM class around the GEMM with activation class\. GEMM with activation itself is a wrapper around the Blackwell and Hopper default GEMMs\. It supports writing two output tensors—the standard GEMM output \(the preactivation\) and the standard GEMM output with an activation function such as SwiGLU or ReLU applied \(the postactivation\)\. We define the activation function to be the identity and the postactivation tensor to be the inplace transpose of the preactivation tensor\. Then, when the GEMM with activation class writes to the postactivation, it is really writing to the upper triangle with a transposed layout—this is exactly the intent of the symmetric GEMM kernel\. We override the epilogue of GEMM with activation just to ensure we do not write twice to the diagonal tiles, for the correctness reasons mentioned previously\.

### C\.3Kernel Optimizations for Standard Newton\-Schulz

Using just optimized CuteDSL kernels, we can accelerate standard Newton\-Schulz with two changes\.

##### Symmetric Matrix Multiplication

As discussed above, the matrices𝑨=𝑿​𝑿⊤\{\\bm\{A\}\}=\{\\bm\{X\}\}\{\\bm\{X\}\}^\{\\top\}and𝑩=bt​𝑨\+ct​𝑨2\{\\bm\{B\}\}=b\_\{t\}\{\\bm\{A\}\}\+c\_\{t\}\{\\bm\{A\}\}^\{2\}computed at each iteration of Newton\-Schulz are symmetric by definition\. Therefore, we use our symmetric GEMM kernels for these operations, reducing their FLOP cost by half\.

##### Fused GEMM \+ Add

The typical way to implement the non\-symmetric multiplication𝑿←at​𝑿\+𝑩​𝑿\{\\bm\{X\}\}\\leftarrow a\_\{t\}\{\\bm\{X\}\}\+\{\\bm\{B\}\}\{\\bm\{X\}\}is to usetorch\.baddbmm, which calls cuBLAS under the hood\. However, we offer a much faster implementation of this “Fused GEMM \+ Add” operation for Hopper\. Unlike cuBLAS, our “Fused GEMM \+ Add” supports Ping Pong Scheduling for Hopper, which better hides the epilogue addition ofat​𝑿a\_\{t\}\{\\bm\{X\}\}\.

## Appendix DAdditional Experiments

### D\.1Architecture and Optimization Setup

We now describe details of the experiments in[Section˜8\.1](https://arxiv.org/html/2608.11612#S8.SS1)\. \(The benchmarks in[Section˜8\.2](https://arxiv.org/html/2608.11612#S8.SS2)are similar, but with minor architectural differences, synthetic data, and GH200s\.\)

##### Architecture\.

The models are decoder\-only dense transformers \(not mixtures of experts\)\. They use grouped\-query attention, sliding\-window attention \(window20482048\), rotary position embeddings, and RMSNorm, at sequence length81928192over a vocabulary of100,352100\{,\}352tokens\. Weights are stored and updated in MXFP8\.[Table˜4](https://arxiv.org/html/2608.11612#A4.T4)lists the per\-scale dimensions\. Each run uses a single node of eight B200s\.

Table 4:Model dimensions by scale\. All share sequence length81928192, vocabulary100,352100\{,\}352, sliding window20482048, grouped\-query attention, and rotary position embeddings\.
##### Optimization\.

Matrix parameters are updated by the orthogonalizing optimizer under study—Muon, NorMuon, or their filtered variants\. For matrix parameters, we multiply the base learning rate bydout/din\\sqrt\{d\_\{\\text\{out\}\}/d\_\{\\text\{in\}\}\}\. Embeddings, the language\-model head, norm scales, and biases are updated by AdamW \(β1=0\.9\\beta\_\{1\}=0\.9,β2=0\.95\\beta\_\{2\}=0\.95\) at the base learning rate, without thedout/din\\sqrt\{d\_\{\\text\{out\}\}/d\_\{\\text\{in\}\}\}scaling applied to the matrix parameters \(the language\-model head, in particular, also uses the base rate\)\. We apply weight decay0\.010\.01to the matrix parameters only; the AdamW \(non\-matrix\) groups use no weight decay\. The sequence length is81928192\. The learning\-rate schedule is warmup–stable–cooldown: a200200\-step linear warmup from2×10−32\\times 10^\{\-3\}to the peak learning rate, a constant phase, and a cooldown over the final25%25\\%of training back down to2×10−32\\times 10^\{\-3\}\. All reported held\-out cross\-entropies are taken at the end of cooldown\.

### D\.2Finer\-Grained Timing Metrics

The main text reports GPU \(device\) time, the metric of interest for the orthogonalization compute\. Here we report three per\-step metrics side by side:

- •GPU \(device\) time, read from CUDA events between two stream markers aroundoptimizer\.step\(\);
- •CPU \(host\) time, fromtime\.perf\_counter\(\)stopped immediately afterstep\(\)returns and before any device sync — cost of kernel launch, Python, and any host\-side syncs the optimizer performs;
- •Communication volume, the number of bytes moved by the megabatch collectives \(all\_to\_all,all\_gather,reduce\_scatter\) at each step, measured by tallying the receive\-side payload per\-GPU\.

As always, Newton\-Schulz runs in half precision arithmetic\. Results for a 14B model sharded over eight B200s and a 7B model on four GH200s appear in[Table˜5](https://arxiv.org/html/2608.11612#A4.T5)and[Table˜6](https://arxiv.org/html/2608.11612#A4.T6), respectively\. Unlike in[Section˜8\.2](https://arxiv.org/html/2608.11612#S8.SS2), we show all four combinations of kernel and Newton\-Schulz variant\. We see that Gram Newton\-Schulz*without*symmetric kernels achieves up to a1\.2×1\.2\\timesspeedup\. Otherwise, GPU timings are consistent with the results presented in[Section˜8\.2](https://arxiv.org/html/2608.11612#S8.SS2)\. The NorMuon family has higher costs overall due to the overhead of its extra normalization steps\. Our contributions still yield similar speedups: up to1\.5×1\.5\\timesfor symmetric kernels and Gram Newton\-Schulz,3\.3×3\.3\\timesforf=1/2f=\\nicefrac\{\{1\}\}\{\{2\}\}, and5\.6×5\.6\\timesforf=1/4f=\\nicefrac\{\{1\}\}\{\{4\}\}\.[Table˜5](https://arxiv.org/html/2608.11612#A4.T5)shows that the CPU cost is negligible—just∼0\.1\{\\sim\}0\.1ms across configurations\. This near\-constant host cost is a direct consequence of CUDA\-graph capture\. Row selection, error feedback, and Gram Newton\-Schulz issue many small operations and kernel launches\. Without capture, the host would dispatch each one individually, and this per\-launch overhead would dominate, especially for the filtered configurations, whose device work is small\. Capture and replay instead collapse the entire step into a single graph launch, so the host issues one call regardless of how many kernels the step contains or how much they compute\. The reported host time is therefore just the fixed cost of that one launch \(∼0\.1\{\\sim\}0\.1ms\), essentially independent of the configuration\. Finally, communication volume is directly proportional to the fractionff, as expected\.

Table 5:Fine\-grained per\-step metrics for the distributed 14B / 8\-GPU benchmark: GPU \(device\) time from CUDA events, CPU \(host\) time fromperf\_counterstopped before the device sync, and per\-GPU communication volume of one step\. We distinguish the GPU time between the Muon and NorMuon families of optimizers; CPU time and communication volume are reported once because neither depends on the family: host time is the fixed CUDA\-graph launch cost, and communication volume is set by the parameter shapes and the fractionff\.Table 6:Replication of[Table˜5](https://arxiv.org/html/2608.11612#A4.T5)for a 7B\-parameter model sharded over a node of four GH200s with NVLink\.
### D\.3Benchmarking all\-to\-all communications

In[Section˜7](https://arxiv.org/html/2608.11612#S7), we claim that reducing the number of communication rounds via megabatching is beneficial in two ways: \(1\) each round adds a fixed overhead not dependent on the size of the payload, and \(2\) small messages fail to saturate the interconnect\. Here, we verify both claims directly with a standalone microbenchmark of NCCLall\_to\_allover NVLink, independent of any optimizer\. On a single node of four H100s \(SXM, NVLink\-4 mesh, NCCL 2\.29\), we sweep the size of the payload per rank from11KiB to11GiB\. We time each all\-to\-all with CUDA events \(10 warmup, 50 timed iterations\), taking the per\-iteration time to be the maximum over ranks, and measure the bandwidth\. We confirmed that the transport is a pure NVLink peer\-to\-peer communication, with no PCIe or network fallback and that the measured bandwidths agree with thenccl\-testsalltoall\_perftool to within1%1\\%\.

[Figure˜21](https://arxiv.org/html/2608.11612#A4.F21)shows the results\. Both effects discussed in[Section˜7](https://arxiv.org/html/2608.11612#S7)are clearly visible\. The latency floor is∼25​μ\{\\sim\}25\\,\\mus \(left panel\), of which∼15\{\\sim\}15–17​μ17\\,\\mus is host\-side dispatch \(as measured separately withperf\_counterand no device synchronization\)\. Without megabatching, we pay this cost over and over\. Bandwidth climbs by more than an order of magnitude as the per\-link payload grows \(right panel\), so coalescing many small all\-to\-alls into one megabatch increases the effective bandwidth of each transfer\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x20.png)Figure 21:NCCLall\_to\_allover NVLink on4×4\\timesH100\.Left:median all\-to\-all time versus payload per rank \(sender\) over 50 trials\. For small payloads, runtime is dominated by a fixed latency\.Right:bus bandwidth versus payload per link \(sender→\\rightarrowreceiver\)\. Bandwidth is near zero for small messages and does not reach8080–90%90\\%of its peak until the per\-link payload is∼16\{\\sim\}16–3232MiB\. Atworld\_size=4\\texttt\{world\\\_size\}=4, a256256KiB per\-link payload achieves only∼10%\{\\sim\}10\\%of the peak bandwidth\.
### D\.4Ablations

##### Tuning the baseline

The experiments in[Section˜8\.1](https://arxiv.org/html/2608.11612#S8.SS1)compare Dion3 to a NorMuon baseline\.[Table˜7](https://arxiv.org/html/2608.11612#A4.T7), below, demonstrates that this baseline is properly tuned for a fair comparison\.

Table 7:Tuning hyperparameters \(learning rateη\\eta, momentumμ\\mu\) of NorMuon at 1B\. Final validation loss on 100B tokens of ClimbMix\. Winner:η=0\.01\\eta=0\.01,μ=0\.9\\mu=0\.9\(bold\)\.

##### Dion3 with no compression

As a control, we run Dion3 withf=1f=1, which should reduce to plain NorMuon\.[Figure˜22](https://arxiv.org/html/2608.11612#A4.F22)confirms this at 1B; the two curves coincide throughout training, and their final validation losses differ by0\.00050\.0005\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x21.png)Figure 22:Validation loss at 1B for NorMuon and for Dion3 atf≈1f\\approx 1, on 100B tokens of ClimbMix\. As expected, Dion3 tracks NorMuon throughout, ending at2\.21412\.2141versus2\.21462\.2146\.
##### Timings with NorMuon family

[Section˜8\.2](https://arxiv.org/html/2608.11612#S8.SS2)compares the runtime of Muon to the Muon version of Dion3\. Here, in[Figure˜23](https://arxiv.org/html/2608.11612#A4.F23), we repeat this benchmark with the NorMuon family\. Results are quite alike, but due to the extra overhead of NorMuon’s normalization steps, it benefits slightly less from our contributions, which target the orthogonalization step\.

![Refer to caption](https://arxiv.org/html/2608.11612v1/x22.png)Figure 23:Optimizer step time \(excluding forward/backward pass\) relative to standard Muon across model scales when training on 1 GH200 \(left\) and FSDP over 4 GH200s \(right\)\. Each colored line adds one of our contributions on top of the previous one; AdamW is shown for reference\. Lines show median over 25 steps and bands show interquartile range\. Compare to[Figure˜6](https://arxiv.org/html/2608.11612#S8.F6)\.

## Appendix ECase Studies of End\-to\-End Training Time

The share of end\-to\-end training time taken up by Newton\-Schulz can vary widely depending on the training setup\. To explain this variability, we analyze two idealized scenarios\. In one, standard Newton\-Schulz takes 2% of training time; in the other it takes 17%\.

### Case Study 1: Standard Newton\-Schulz takes 2% of Kimi K2 training time

The following analysis gives a very optimistic estimate of the optimizer’s wall clock time\. We assume an efficient training infrastructure with highly optimized pipeline parallelism\. Moreover, we assume that the optimizer step of each pipeline stage is completely hidden behind the backward pass of the next pipeline stage\.

Kimi K2 Thinking is a1\.11\.1trillion parameter model with3232billion active parameters\. It has11dense layer followed by6060MoE layers\[[21](https://arxiv.org/html/2608.11612#bib.bib8)\]\. It is pretrained with256256\-GPU model parallel groups,1616\-way pipeline parallelism,1616\-way expert parallelism within each pipeline stage, and a huge batch size of6767million tokens\.

We use a single H100 to approximate the share of each training step’s runtime occupied by Newton\-Schulz in this setting under the following assumptions:

1. 1\.The training cluster consists of256256nodes of eight H100s each \(20482048GPUs in total\), connected with NDR 400 Gb/s InfiniBand inter\-node \(8 NICs per node, 1:1 NIC\-to\-GPU ratio\) and NVLink 4\.0 intra\-node\. This is the size of the cluster used to train DeepSeek\-V3, with upgraded hardware\[[9](https://arxiv.org/html/2608.11612#bib.bib4)\]\.
2. 2\.Training inbfloat16attains40%40\\%model flop utilization \(MFU\), which is typical for MoEs at this scale on H100s\.
3. 3\.The only non\-overlapped optimizer time is that of the last pipeline stage to complete its backward pass \(i\.e\., pipeline stage11of1616\)\. The optimizer steps of pipeline stages22to1616are fully hidden behind the backwards of stages11to1515\.
4. 4\.Pipeline stage11contains the dense layer and33MoE layers\.

Under these assumptions, the optimal way to partition the Newton\-Schulz work of pipeline stage11is as follows\. Each of the1616GPUs in pipeline stage11’s expert parallel group gets

384​experts/layer×3​MoE layers16​GPUs=72​experts/GPU=216​expert up\-gate\-down/GPU\.\\frac\{384\\text\{ experts/layer\}\\times 3\\text\{ MoE layers\}\}\{16\\text\{ GPUs\}\}=72\\text\{ experts/GPU\}=216\\text\{ expert up\-gate\-down/GPU\}\.Each of the1616GPUs has its own unique expert weights, so no communication is needed\. Pipeline stage11also contains the dense MLP’s three7168×184327168\\times 18432weights \(up/gate/down\) and three shared experts\. Orthogonalizing the dense MLP weights is the dominant cost, so they are sent to three different GPUs; the shared experts are split amongst the remaining1313GPUs\. Thus, the orthogonalization time for pipeline stage11is the time it takes for a single GPU to orthogonalize216216expert up/gate/down weights and11dense up/gate/down weight\. As benchmarked in[Figure˜10](https://arxiv.org/html/2608.11612#A1.F10), standard Newton\-Schulz implemented in PyTorch takes315315ms to do this\. Per our assumption, stage11’s orthogonalization step is the only one that is not overlapped\.

Having estimated the Newton\-Schulz time, we now estimate the end\-to\-end wall clock time of an entire Kimi K2 global training step\. We use the standard estimate of6​N​B6NBFLOPs for the forward and backward pass, whereNNis the number of active parameters andBBis the global batch size\. Given

- •Active parameters:N=32⋅109N=32\\cdot 10^\{9\}
- •H100 peak:P=989⋅1012P=989\\cdot 10^\{12\}FLOP/s
- •Model flop utilization:MFU=40%\\text\{MFU\}=40\\%
- •Cluster size:G=2048G=2048GPUs
- •Global batch size:B=67⋅106B=67\\cdot 10^\{6\}tokens

we have

sec/batch=6​N×BP×MFU×G=15\.9\\text\{sec/batch\}=\\frac\{6N\\times B\}\{P\\times\\text\{MFU\}\\times G\}=15\.9
Thus, Newton\-Schulz takes approximately315​ms15900​ms\+315​ms=1\.9%\\frac\{315\\text\{ ms\}\}\{15900\\text\{ ms\}\+315\\text\{ ms\}\}=1\.9\\%of total pretraining wall clock time in this setting\.

### Case Study 2: Standard Newton\-Schulz takes 17% of Llama3\-70B SFT time

Llama3\-70B is an 80\-layer dense model with hidden size81928192, intermediate size2867228672, and grouped query attention with𝑾K,𝑾V\{\\bm\{W\}\}\_\{K\},\{\\bm\{W\}\}\_\{V\}of size1024×81921024\\times 8192and𝑾Q,𝑾O\{\\bm\{W\}\}\_\{Q\},\{\\bm\{W\}\}\_\{O\}of size8192×81928192\\times 8192\[[14](https://arxiv.org/html/2608.11612#bib.bib31)\]\. Supervised finetuning \(SFT\) typically uses small batch sizes, ranging from3232to256256sequences\[[9](https://arxiv.org/html/2608.11612#bib.bib4)\]\.

We analyze the following SFT case:

1. 1\.Training uses3232H100s across44nodes \(8 GPUs per node\)\.
2. 2\.Training inbfloat16hits40%40\\%MFU\.
3. 3\.Weights are sharded evenly across GPUs using FSDP, and the exposed Newton\-Schulz time is that of80​layers/32​GPUs≈3​layers80\\text\{ layers\}/32\\text\{ GPUs\}\\approx 3\\text\{ layers\}\. Each layer has33MLP weights \(up, gate, down\), and the attention weights𝑾Q\{\\bm\{W\}\}\_\{Q\},𝑾K\{\\bm\{W\}\}\_\{K\},𝑾V\{\\bm\{W\}\}\_\{V\}, and𝑾O\{\\bm\{W\}\}\_\{O\}\.

According to our benchmarking, standard Newton\-Schulz of

- •nine8192×286728192\\times 28672weights takes739739ms,
- •six8192×81928192\\times 8192weights takes156156ms,
- •six1024×81921024\\times 8192weights takes2\.322\.32ms,

totaling897897ms\. Given

- •Parameters:N=70⋅109N=70\\cdot 10^\{9\}
- •H100 peak:P=989⋅1012P=989\\cdot 10^\{12\}FLOP/s
- •Model flop utilization:MFU=40%\\text\{MFU\}=40\\%
- •Cluster size:G=32G=32GPUs
- •Global batch size:B=64​sequences×2048​tokens/sequence=131,072B=64\\text\{ sequences\}\\times 2048\\text\{ tokens/sequence\}=131\{,\}072tokens

we have

sec/batch=6​N×BP×MFU×G=4\.35​s\\text\{sec/batch\}=\\frac\{6N\\times B\}\{P\\times\\text\{MFU\}\\times G\}=4\.35\\text\{ s\}
Thus, Newton\-Schulz takes approximately897​ms4350​ms\+897​ms=17%\\frac\{897\\text\{ ms\}\}\{4350\\text\{ ms\}\+897\\text\{ ms\}\}=17\\%of total SFT wall clock time in this setting\.