Dion3: Full-Stack Orthogonal Updates
Noah Amsel, Jack Zhang, Kwangjun Ahn, Ali Naeimi, Austin Feng, Berlin Chen, Tri Dao, John Langford
cs.LG, cs.AI
2026-08-12
Dion3 cuts Muon's orthogonalization cost four ways. The key trick is orthogonalizing only a subset of the momentum matrix's rows each step, for 6.5x faster steps and lower loss.
Muon is a family of optimizers that orthogonalize the momentum matrix before applying the update, projecting the step onto an orthogonal matrix, which amounts to steepest descent in the spectral norm. It trains well, and over the past year has been picked up for pretraining models in the tens of billions of parameters. The cost sits in a single step: every iteration runs a Newton-Schulz procedure to approximate the orthogonalization, which is cubic-time, and once weights are sharded across GPUs, all-to-all communication piles on top. The paper frames the bill concretely: on a 7B model across four GH200s, Muon's optimizer step (excluding forward/backward) takes 26× as long as AdamW's.
Others have already tried to make this step cheaper. Dion builds a low-rank approximation via power iteration and orthogonalizes only that; Trion goes further and picks columns straight out of a discrete-cosine-transform matrix. Both lean on error feedback, recording the approximation error each step and compensating for it later, to hold quality. Dion3 pushes the same line one step further: if the goal is just to shrink the input, the simplest "low-rank approximation" is to pick some rows of the momentum matrix and ignore the rest.
Dion3 is not a single trick; it compresses this step on four layers at once.
The first layer is the algorithm itself. Standard Newton-Schulz iterates on the large input matrix X. Dion3 instead iterates on the small symmetric Gram matrix XX⊤ and multiplies back to X at the end. The output is mathematically identical to standard Newton-Schulz, but almost all the arithmetic lands on the much smaller n×n symmetric matrix, cutting large rectangular multiplies from ten down to two. At the typical setting (T=5, α=4) this saves 55% of the FLOPs versus standard Newton-Schulz with symmetric GEMM, or 68% versus the common implementation without it. The price is that the Gram matrix can develop spurious negative eigenvalues and go unstable, so the authors insert a restart at iteration 3 and switch to float16.
The second layer is the GPU kernel. Symmetric products like A⊤A and A²+B only need the lower triangle, mirrored to the upper. The authors write dedicated symmetric-GEMM kernels in CuteDSL, about 2× faster than cuBLAS on large matrices (Hopper, Blackwell). This stacks well with Gram Newton-Schulz, which uses more symmetric products.
The third layer is the most counterintuitive change: each step orthogonalizes only a fraction of the momentum matrix's rows. It picks the k=⌈fn⌉ rows of largest ℓ1 norm (f recommended at 1/4 or 1/8), orthogonalizes only that submatrix, and updates only those rows' weights. Error feedback then decays just the selected rows by a factor μ, leaving the unselected rows intact for the next round. At f=1 this collapses back to plain Muon. Because fewer rows update per step, the effective step size shrinks, so the learning rate must scale as η ∝ 1/√f — a transfer rule the authors derive explicitly.
The fourth layer is communication. A transformer has only a handful of distinct weight shapes, so megabatching packs every same-shape matrix into a single all-to-all, dropping the number of communication rounds from O(N/worldsize) to O(1), independent of model depth.
The speedups stack. Symmetric kernels plus Gram Newton-Schulz give about 1.5× over standard Muon; adding row selection brings the total to 3.6× at f=1/2 and 6.5× at f=1/4. Returning to the 7B example, standard Muon sits at 26× AdamW and the full stack brings it down to 4× AdamW. Still slower than AdamW, but back from prohibitively expensive to workable.
Megabatching pays off most when the step is communication-bound: on a 1B model with 8 shards on one node, the optimizer step drops from 80.7ms to 52.1ms (−35%); at 32 shards it saves only 4%, because each rank holds so few matrices that batch sizes barely change.
The quality result surprised the authors. They had set out only to show that row selection does not hurt the loss; Dion3 (f<1) actually beats NorMuon. On a 1B model over 100B tokens of ClimbMix, the best setting is f=1/8, with validation loss running below NorMuon throughout training. Scaled to 3B–14B (only 10B tokens, for cost), Dion3 (f=1/4) has lower validation loss than NorMuon at every size:
| Size | NorMuon loss | Dion3 loss | Δ | Downstream Δ (12 benchmarks) |
| 3B | 2.269 | 2.257 | −0.012 | +1.0 |
| 4B | 2.243 | 2.232 | −0.011 | −0.3 |
| 7B | 2.220 | 2.206 | −0.014 | +0.1 |
| 14B | 2.189 | 2.162 | −0.027 | +0.7 |
Loss wins at all four scales; downstream accuracy wins three and loses one, with the biggest gain at 14B (−0.027 loss, +0.7 accuracy).
Muon draws attention because it may beat AdamW on per-token training efficiency, but its per-step cost has been the main thing blocking it from scaling. Dion3 cuts that obstacle down substantially: a 6.5× faster step on large models with quality that does not regress and slightly improves, shipped as the dion package (github.com/microsoft/dion) and a drop-in replacement for Muon. For teams already using or considering Muon for large-model pretraining, this is a directly usable cost saving.
It does not pull Muon onto AdamW's level: at 7B it is still 4× AdamW. Whether Muon pays off overall depends on how much its faster convergence offsets the pricier step. Dion3 tilts that trade toward Muon without ending the argument.
The authors are fairly candid. Gram Newton-Schulz and the symmetric kernels run in half precision, introducing numerical differences from the original (experiments show no quality hit, but the differences are real). Dion3 at f=1 is not bit-identical to Muon either: row permutation, momentum-damping timing, float32 NorMuon normalization, and a custom Triton kernel all differ, though the convergence curves nearly overlap. The restart itself costs an extra 3(α−1)n³ FLOPs.
A few things are under-tested. The quality gain is measured on one dataset (ClimbMix) and one architecture family (dense transformers); mixture-of-experts models are benchmarked only for speed, not quality, and the authors flag that further work is needed to see how widely the improvement generalizes. The −0.027 loss jump at 14B rests on a single 10B-token run at limited scale. And the 4B row loses 0.3 on downstream accuracy, a reminder that loss and downstream metrics do not always agree, so loss curves alone may overstate the payoff.