Dion3: An Efficient Muon Optimizer Revision

• Dion3 is an optimized revision of the Muon optimizer that reduces matrix-orthogonalization costs through Gram Newton–Schulz, symmetric GPU kernels, fractional updates, and megabatched communication.
• Dion3 supports Muon and NorMuon variants, uses error feedback to preserve residual momentum during partial row or column updates, and is available as a drop-in replacement through the `dion` package.
• Experiments report comparable or improved loss and accuracy versus Muon and NorMuon, with optimizer-step speedups of up to approximately 6× and communication reductions that scale roughly with the selected update fraction.

Dion3 is a revision of the Muon optimizer that combines algorithmic, kernel-level, update-rule, and distributed-communication changes to reduce the cost of matrix orthogonalization. It replaces or augments conventional Newton–Schulz orthogonalization with Gram Newton–Schulz, exploits symmetry through CuteDSL kernels, orthogonalizes only a selected fraction of momentum rows or columns, and coalesces distributed collectives through megabatching. Dion3 supports Muon and NorMuon variants and is distributed through the dion package as a drop-in replacement for Muon. The reported experiments show comparable or improved loss relative to Muon and NorMuon, with optimizer-step reductions of up to approximately $6\times$ (Amsel et al., 12 Aug 2026).

1. Motivation and relationship to Muon

Muon maintains a matrix-valued momentum buffer. For a parameter matrix $W\in\mathbb R^{n\times m}$ and gradient $G_t$, the momentum recurrence is

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$

and the parameter update is

$W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$

possibly with weight decay and per-matrix learning-rate scaling. If $M=U\Sigma V^\top$ is an SVD, its polar factor is $\operatorname{polar}(M)=UV^\top$. The polar operation therefore preserves the singular-vector structure while replacing the singular values by $1$, distinguishing Muon from SGD with momentum.

Because exact polar decomposition is expensive, Muon approximates it using a five-step Newton–Schulz iteration. For a normalized matrix $X_0$, the standard iteration is

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$

Assuming $W\in\mathbb R^{n\times m}$0 and defining the aspect ratio $W\in\mathbb R^{n\times m}$1, one iteration requires, ignoring lower-order operations,

• $W\in\mathbb R^{n\times m}$2 FLOPs for $W\in\mathbb R^{n\times m}$3;
• $W\in\mathbb R^{n\times m}$4 FLOPs for $W\in\mathbb R^{n\times m}$5;
• $W\in\mathbb R^{n\times m}$6 FLOPs for multiplication back by $W\in\mathbb R^{n\times m}$7.

Thus $W\in\mathbb R^{n\times m}$8 iterations cost

$W\in\mathbb R^{n\times m}$9

For $G_t$0, this becomes $G_t$1 across 15 matrix multiplications. Muon consequently has cubic complexity for square matrices and a particularly strong dependence on $G_t$2 for rectangular transformer matrices. MLP, MoE, and attention-projection matrices can have $G_t$3, with increasingly asymmetric matrices occurring in fine-grained MoEs.

Sharding introduces a second cost. Under FSDP or related schemes, the shards of a momentum matrix must generally be assembled before orthogonalization and scattered afterward. A conventional implementation processes batches of approximately $G_t$4 matrices, producing $G_t$5 communication rounds per step for $G_t$6 matrices. The resulting overhead includes full-matrix communication volume, collective-launch latency, synchronization, and poor interconnect utilization for small messages.

Dion3 addresses these costs at four levels:

1. Gram Newton–Schulz: operates primarily on a smaller symmetric Gram matrix.
2. CuteDSL symmetric kernels: avoid recomputing the redundant triangular half of symmetric products.
3. Fractional row or column updates: orthogonalize and update only a selected fraction of the momentum matrix.
4. Megabatched communication: packs same-shaped matrices into a small number of distributed collectives.

NorMuon extends Muon with row-wise second-moment normalization. For $G_t$7, it maintains

$G_t$8

forms

$G_t$9

and rescales the result to preserve the Frobenius norm. Dion3 supports both Muon and NorMuon variants.

2. Gram Newton–Schulz orthogonalization

For $M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$0 with $M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$1,

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$2

Gram Newton–Schulz forms the symmetric Gram matrix

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$3

and approximates its inverse square root. The final output is

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$4

where $M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$5 approximates $M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$6.

The reformulation follows from expressing each odd Newton–Schulz polynomial as

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$7

where

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$8

The scalar recurrence is

$M_t=\mu M_{t-1}+G_t,\qquad M_0=0,$9

$W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$0

$W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$1

The composition of the Newton–Schulz polynomials is then $W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$2, with $W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$3 approaching $W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$4. The matrix recurrence is

$W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$5

$W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$6

$W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$7

In exact arithmetic, Gram Newton–Schulz is mathematically equivalent to applying the original Newton–Schulz polynomial sequence directly to $W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$8.

Computational scaling

The Gram formulation requires two large rectangular products:

1. formation of $W_{t+1}=W_t-\eta\,\operatorname{polar}(M_t),$9;
2. final multiplication $M=U\Sigma V^\top$0.

The iterative operations act on $M=U\Sigma V^\top$1 matrices. With symmetric matrix multiplication, the paper gives the naive Gram cost as

$M=U\Sigma V^\top$2

compared with

$M=U\Sigma V^\top$3

for standard Newton–Schulz using symmetric products. At $M=U\Sigma V^\top$4, the costs are equal, so Dion3 falls back to standard Newton–Schulz for square matrices. For $M=U\Sigma V^\top$5, Gram Newton–Schulz is cheaper, with asymptotic dependence changing from

$M=U\Sigma V^\top$6

to

$M=U\Sigma V^\top$7

For $M=U\Sigma V^\top$8 and $M=U\Sigma V^\top$9, representative of a $\operatorname{polar}(M)=UV^\top$0-expanded MLP, the reported FLOP reductions are 55% relative to standard Newton–Schulz with symmetric GEMMs and 68% relative to an implementation without symmetric GEMMs.

Numerical stability

Naive Gram Newton–Schulz can be unstable in half precision because finite-precision formation of $\operatorname{polar}(M)=UV^\top$1 can produce small negative eigenvalues even though an exact Gram matrix is positive semidefinite. The inverse-square-root iteration behaves poorly on such eigenvalues: under

$\operatorname{polar}(M)=UV^\top$2

a negative $\operatorname{polar}(M)=UV^\top$3 can grow in magnitude instead of converging. Eigenvector drift caused by roundoff can amplify the problem.

The stabilized implementation uses a restart with the recommended five iterations:

1. run two Gram iterations;
2. explicitly reconstruct the intermediate matrix;
3. form a fresh Gram matrix;
4. run the remaining three iterations.

If $\operatorname{polar}(M)=UV^\top$4 is the output after two iterations, the restart forms

$\operatorname{polar}(M)=UV^\top$5

and reinitializes the inverse-square-root iteration. This adds approximately

$\operatorname{polar}(M)=UV^\top$6

FLOPs but limits error accumulation. The reported stability analysis, assuming spurious eigenvalues as low as approximately $\operatorname{polar}(M)=UV^\top$7, identifies restarting after iteration 2 as the best condition-number-control strategy for the five-step Polar Express coefficients.

The production implementation additionally uses float16 rather than bfloat16 by default, normalizes the input by $\operatorname{polar}(M)=UV^\top$8 with $\operatorname{polar}(M)=UV^\top$9, applies a safety factor such as scaling the input argument by $1$0, and avoids explicitly materializing some $1$1 additions in numerically sensitive paths. The stabilized method preserves training quality, whereas naive Gram Newton–Schulz can produce loss spikes, infinities, and divergence.

3. Symmetric CuteDSL kernels

The Gram matrices and intermediate matrices satisfy

$1$2

A conventional GEMM nevertheless computes both triangular halves. Dion3’s CuteDSL kernels compute only the lower triangle and copy it to the upper triangle. For a symmetric output $1$3, the kernel computes $1$4 only for $1$5, then writes

$1$6

The kernels use a triangular tile scheduler, an epilogue that writes each computed lower-triangular tile and its transpose, special handling for diagonal tiles, and fusion of expressions such as

$1$7

and matrix-quadratic operations. The implementation targets NVIDIA Hopper and Blackwell GPUs and uses $1$8 tile operation with architecture-specific cluster and thread-block configurations.

The kernels are implemented in roughly 160 lines by wrapping standard CuteDSL GEMM abstractions. For sufficiently large matrices, they achieve approximately $1$9 the speed of cuBLAS GEMM for symmetric products, with or without an epilogue addition. Their largest benefit arises in Gram Newton–Schulz because that algorithm performs many small symmetric operations.

The reported performance effects include:

• $X_0$0–$X_0$1 speedup of the Newton–Schulz routine with kernelized Gram Newton–Schulz;
• $X_0$2–$X_0$3 end-to-end Muon optimizer-step speedup on a single H100 across tested architectures;
• approximately $X_0$4 speedup for exposed orthogonalization work estimated for Kimi K2;
• approximately $X_0$5 speedup on higher-aspect-ratio models such as Gemma-1B and the tested MoE model.

In quality experiments on Hopper and Blackwell, replacing standard Newton–Schulz with stabilized, kernelized Gram Newton–Schulz preserved validation perplexity to within approximately $X_0$6.

4. Fractional row and column updates

Dion3’s fractional update rule avoids orthogonalizing the full momentum matrix. Let

$X_0$7

be the updated momentum matrix. For a fraction $X_0$8, define

$X_0$9

Dion3 selects a set

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$0

which by default consists of the $$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$1 rows with largest $$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$2 norm:

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$3

It forms

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$4

computes

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$5

and updates only the selected rows:

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$6

Unselected rows receive no orthogonalized update during that step, although weight decay remains applied to all rows. Columns can be selected instead of rows. The implementation selects along the sharded dimension when possible; otherwise it chooses the smaller dimension.

When $$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$7, Dion3 reduces to Muon or NorMuon in exact arithmetic, apart from implementation-level effects such as operation ordering and finite-precision behavior.

Error feedback

Dion3 decays only the selected component of momentum. Define the selected matrix embedded in the full shape by

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$8

The momentum update is

$$X_{t+1}=a_tX_t+b_tX_tX_t^\top X_t +c_t(X_tX_t^\top)^2X_t.$$9

or equivalently,

$W\in\mathbb R^{n\times m}$00

Selected rows are damped by $W\in\mathbb R^{n\times m}$01, whereas unselected rows are retained without damping. Residual momentum therefore accumulates and can cause previously ignored rows to be selected in later steps. This mechanism is analogous to Dion’s error feedback, but Dion uses a low-rank factorization while Dion3 uses direct row or column selection.

Compute and communication reduction

For a selected matrix of shape $W\in\mathbb R^{n\times m}$02, the relevant multiplication costs become approximately

$W\in\mathbb R^{n\times m}$03

and

$W\in\mathbb R^{n\times m}$04

The rectangular portions can therefore decrease by roughly $W\in\mathbb R^{n\times m}$05, while communication volume falls approximately linearly:

$W\in\mathbb R^{n\times m}$06

Selection also changes the aspect ratio to

$W\in\mathbb R^{n\times m}$07

making Gram Newton–Schulz relatively more attractive. The paper recommends $W\in\mathbb R^{n\times m}$08 or $W\in\mathbb R^{n\times m}$09. Random selection performed only somewhat worse than top-$W\in\mathbb R^{n\times m}$10-norm selection in initial experiments, suggesting that error feedback may be a major contributor to the approximation’s behavior.

Because a full Muon update has Frobenius norm

$W\in\mathbb R^{n\times m}$11

whereas a selected update has norm

$W\in\mathbb R^{n\times m}$12

matching update magnitudes gives

$W\in\mathbb R^{n\times m}$13

The experiments found the corresponding relation

$W\in\mathbb R^{n\times m}$14

for the 1B NorMuon/Dion3 experiments.

5. Megabatched distributed communication

Conventional distributed Muon processing handles approximately $W\in\mathbb R^{n\times m}$15 matrices per communication round on $W\in\mathbb R^{n\times m}$16 GPUs, requiring $W\in\mathbb R^{n\times m}$17 all-to-all rounds for $W\in\mathbb R^{n\times m}$18 matrices. Every round incurs fixed latency and synchronization overhead, while small per-link payloads fail to saturate the interconnect.

A four-H100 NCCL microbenchmark connected by NVLink reports:

• an all-to-all latency floor of approximately $W\in\mathbb R^{n\times m}$19;
• approximately $W\in\mathbb R^{n\times m}$20–$W\in\mathbb R^{n\times m}$21 attributable to host-side dispatch;
• approximately 10% of peak bandwidth for a 256 KiB per-link payload;
• 80–90% of peak bandwidth when per-link payloads reach approximately 16–32 MiB.

Dion3 groups all matrices with the same shape into a single batch. For each shape group, it packs local momentum shards, performs one all-to-all to assemble the matrices, orthogonalizes the batch, and performs one all-to-all or equivalent operation to scatter the results. Since transformers contain only a small number of distinct matrix shapes, the number of communication rounds becomes approximately $W\in\mathbb R^{n\times m}$22 with respect to model depth rather than $W\in\mathbb R^{n\times m}$23.

Megabatching does not change the mathematical update, state layout, checkpoint format, or orthogonalization procedure. It changes the timing and organization of tensor movement. On GH200 systems, the reported Muon optimizer-step measurements were:

Model and configuration	Ordinary batching	Megabatching	Change
1B, 1 node, 8 FSDP shards	80.7 ms	52.1 ms	$W\in\mathbb R^{n\times m}$24
1B, 4 nodes, 32 FSDP shards	61.9 ms	59.3 ms	$W\in\mathbb R^{n\times m}$25
14B, 1 node, 8 FSDP shards	144.0 ms	140.8 ms	$W\in\mathbb R^{n\times m}$26
14B, 4 nodes, 32 FSDP shards	94.7 ms	89.1 ms	$W\in\mathbb R^{n\times m}$27

The largest benefit occurs for the 1B, 8-shard case, where communication dominates and each rank holds many matrices. For the 14B model, Newton–Schulz computation dominates, limiting the effect of communication coalescing.

Dion3 also permits compressed data-parallel synchronization. When row selection can be performed locally, replicas need synchronize only the selected submatrix rather than the full momentum matrix, reducing data-parallel communication by approximately $W\in\mathbb R^{n\times m}$28. Local per-shard selection avoids global top-$W\in\mathbb R^{n\times m}$29 synchronization but is not equivalent to global top-$W\in\mathbb R^{n\times m}$30 selection.

6. Empirical results

Optimization quality

The principal quality experiments trained dense decoder-only transformers on ClimbMix. The 1B experiment used 100B tokens, FSDP, MXFP8 weights, sequence length $W\in\mathbb R^{n\times m}$31, and tuned NorMuon as the baseline with

$W\in\mathbb R^{n\times m}$32

Dion3 used the learning-rate transfer rule

$W\in\mathbb R^{n\times m}$33

For $W\in\mathbb R^{n\times m}$34, the best loss occurred near the predicted learning-rate-transfer line. Dion3 variants consistently tracked below NorMuon and finished approximately $W\in\mathbb R^{n\times m}$35 cross-entropy points lower, with the best result at $W\in\mathbb R^{n\times m}$36. The paper presents this as an experimental result rather than a universal theoretical guarantee.

The 3B–14B experiments trained for 10B tokens:

Model	NorMuon loss	Dion3 loss, $W\in\mathbb R^{n\times m}$37	Loss difference	NorMuon accuracy	Dion3 accuracy	Accuracy difference
3B	2.269	2.257	$W\in\mathbb R^{n\times m}$38	53.9%	54.9%	$W\in\mathbb R^{n\times m}$39
4B	2.243	2.232	$W\in\mathbb R^{n\times m}$40	55.2%	54.9%	$W\in\mathbb R^{n\times m}$41
7B	2.220	2.206	$W\in\mathbb R^{n\times m}$42	56.0%	56.1%	$W\in\mathbb R^{n\times m}$43
14B	2.189	2.162	$W\in\mathbb R^{n\times m}$44	57.4%	58.1%	$W\in\mathbb R^{n\times m}$45

Dion3 had lower validation loss at every reported scale and higher downstream accuracy at three of four scales. At 14B, the accuracy improvement was $W\in\mathbb R^{n\times m}$46 percentage points. In the $W\in\mathbb R^{n\times m}$47 ablation, Dion3 and NorMuon produced essentially identical curves, ending at 2.2141 and 2.2146, respectively.

Optimizer-step speed

The speed benchmarks successively added symmetric kernels, Gram Newton–Schulz, fractional updates, and megabatching. Timings excluded forward and backward passes and used CUDA-event measurements over 25 steps. Relative to standard Muon on larger models, the reported effects were:

• symmetric kernels plus Gram Newton–Schulz: at least $W\in\mathbb R^{n\times m}$48 faster;
• fractional updates with $W\in\mathbb R^{n\times m}$49: an additional approximately $W\in\mathbb R^{n\times m}$50;
• fractional updates with $W\in\mathbb R^{n\times m}$51: an additional approximately $W\in\mathbb R^{n\times m}$52;
• complete speedups of $W\in\mathbb R^{n\times m}$53 for $W\in\mathbb R^{n\times m}$54 and $W\in\mathbb R^{n\times m}$55 for $W\in\mathbb R^{n\times m}$56.

The abstract summarizes the practical result as up to $W\in\mathbb R^{n\times m}$57, reflecting variation across measurements.

For a 7B model on four GH200 GPUs, the measured Muon GPU times were:

Configuration	Optimizer-step time
Standard Newton–Schulz, PyTorch	360.6 ms
Standard plus symmetric kernels	274.4 ms
Gram Newton–Schulz, PyTorch kernels	296.6 ms
Gram Newton–Schulz plus CuteDSL kernels	217.7 ms
Fractional update, $W\in\mathbb R^{n\times m}$58	101.5 ms
Fractional update, $W\in\mathbb R^{n\times m}$59	59.6 ms

For a 14B model on eight B200 GPUs, Muon time decreased from 153.8 ms with standard PyTorch operations to 34.8 ms with Gram Newton–Schulz, CuteDSL kernels, and $W\in\mathbb R^{n\times m}$60. Communication decreased from 5479 MB to 1370 MB, consistent with approximately proportional scaling in $W\in\mathbb R^{n\times m}$61.

The full-stack 7B GH200 example reports Muon at approximately $W\in\mathbb R^{n\times m}$62 AdamW optimizer-step time before Dion3’s optimizations and approximately $W\in\mathbb R^{n\times m}$63 AdamW afterward.

7. Implementation, limitations, and scope

Dion3 is distributed through the open-source dion package as a drop-in Muon replacement. The package supports Muon, NorMuon, Dion3 fractional updates, standard and Gram Newton–Schulz, FSDP2, DDP, mixed sharding, megabatched collectives, and compressed data-parallel synchronization. A separate gram-newton-schulz package provides a drop-in replacement for Muon’s Newton–Schulz routine.

Important implementation details include CUDA graph capture and replay, custom Triton kernels for fused weight decay and selected-row updates, float32 NorMuon normalization, half-precision Newton–Schulz, top-$W\in\mathbb R^{n\times m}$64-norm row selection by default, local per-shard selection in distributed settings, megabatching by matrix shape, fallback to standard Newton–Schulz for square matrices, and explicit splitting of naturally distinct matrices such as SwiGLU up and gate weights before orthogonalization.

Splitting SwiGLU up and gate matrices separately improved Llama-430M perplexity by approximately $W\in\mathbb R^{n\times m}$65. Attention projections were not split by head because this worsened loss in the reported tests, although the paper notes that head-wise splitting may be beneficial in architectures whose computation treats heads independently.

The principal limitations are:

• Empirical rather than universal fractional-update quality: Dion3 improved loss in the reported ClimbMix experiments, but the generality of this behavior remains unresolved.
• Condition-number effects: forming $W\in\mathbb R^{n\times m}$66 squares the condition number, making Gram Newton–Schulz less suitable for high-accuracy polar decomposition.
• Half-precision stability: naive Gram Newton–Schulz can diverge, and the recommended restart schedule is tied to the iteration count and coefficient sequence.
• Local-versus-global selection: independent per-shard selection avoids synchronization but differs from global top-$W\in\mathbb R^{n\times m}$67 selection.
• Launch overhead: small models can become launch-bound because fractional selection introduces additional operations and kernel launches; CUDA graphs are important in this regime.
• Hardware and shape dependence: Gram Newton–Schulz benefits most from large aspect ratios, CuteDSL kernels target Hopper and Blackwell, and megabatching is most effective when many small matrices are communicated.

Dion3’s contribution is therefore a coordinated systems-and-algorithmic design rather than a single orthogonalization formula. Gram Newton–Schulz reduces the dependence on rectangular matrix dimensions, CuteDSL kernels reduce redundant symmetric arithmetic, fractional updates reduce computation and communication through a controlled approximation with error feedback, and megabatching reduces collective latency. The reported evidence supports matching or improving Muon and NorMuon loss in the tested settings while substantially reducing optimizer-step time, but the fractional update’s behavior outside those settings remains an empirical question.