Dion3: Full-Stack Orthogonal Updates

TL;DR

Dion3 is a revision of the Muon optimizer that targets the Newton-Schulz orthogonalization overhead at every level of the stack. It combines the Gram Newton-Schulz algorithm (FLOP reduction from operating on the small symmetric Gram matrix rather than the rectangular momentum), CuteDSL symmetric-GEMM kernels (exploiting A @ A.T structure), and a megabatching strategy for communication overhead, plus a novel update rule that selects only a fraction of the momentum matrix’s rows to orthogonalize per step. The result matches or improves on Muon’s training loss while reducing optimizer step time by up to 6×. Ships as a drop-in Muon replacement in the dion package.

Key claims

Muon’s Newton-Schulz orthogonalization is a significant per-step overhead due to cubic-time cost, and communication overhead compounds this cost when weights are sharded — eroding Muon’s benefits in distributed settings [Abstract].
The Gram Newton-Schulz algorithm reduces the FLOP cost of orthogonalization by moving the dominant computation to the small symmetric Gram matrix X @ X.T [Abstract].
Custom CuteDSL kernels accelerate orthogonalization by exploiting symmetry of the intermediate matrices [Abstract].
A megabatching strategy reduces the communication overhead associated with distributed orthogonalization [Abstract].
A simple change to the update rule — selecting only a fraction of the momentum matrix’s rows to orthogonalize at each step — cuts costs further and improves on Dion (the prior “compressed” Muon variant) in both speed and performance [Abstract].
Dion3 matches or improves on the loss achieved by Muon while reducing optimizer step time by up to 6× [Abstract].
Ships in the dion package as a drop-in replacement for Muon [Abstract].

Method

Dion3 is the “full stack” consolidation of three previously separate lines of Muon-overhead work by an overlapping author group. The Newton-Schulz step in Muon iterates X_{t+1} = a·X + b·(X X^T)X + c·(X X^T)^2 X on the rectangular d_out × d_in momentum matrix; at typical aspect ratios these rectangular matmuls dominate cost, and under FSDP the sharded momentum must be reconstructed before the iteration runs.

The paper packages four levers into one optimizer. (1) Algorithmic: Gram Newton-Schulz (from the Amsel/Zhang/Chen/Dao blog) rewrites the iteration to run on the small d_out × d_out symmetric Gram matrix X X^T, with only an upfront X X^T formation and a final Q_T X application touching the rectangular dimension. Restart stabilizes the otherwise divergent bfloat16 iteration. (2) Kernel: CuteDSL symmetric-GEMM kernels for Hopper/Blackwell write only the lower triangle of X X^T and its powers, halving compute on those steps. (3) Communication: megabatching amortizes the communication overhead of collecting the sharded momentum. (4) Update rule: selecting only a fraction of momentum rows to orthogonalize per step is the descendant of Dion’s power-iteration compression and Dion2’s submatrix-selection scheme (from the microsoft/dion codebase), but with a simpler formulation the paper claims strictly dominates prior “compressed” variants.

Results

Optimizer step time reduced by up to 6× vs Muon [Abstract].
Training loss matches or improves on Muon [Abstract].
Improves on Dion in both speed and training performance [Abstract].

The abstract does not report scale-specific numbers or the training workloads used; the referenced dion package is the primary code artifact.

Why it’s interesting

Dion3 is the convergence point of three separately-filed lines of Muon-overhead work: Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon contributed the Gram-Newton-Schulz algorithm + CuteDSL symmetric-GEMM kernels, microsoft/dion — distributed Muon, Dion2, Dion, and NorMuon orthonormal optimizers contributed the Dion/Dion2 compression update-rule lineage and the distributed / Triton-kernel implementation, and the resulting package is shared under one Muon-drop-in API — 6× step-time speedup is the strongest end-to-end claim yet for the “make distributed Muon actually affordable” line the Training stability at scale page has been tracking. That the same authorship cluster (Amsel + Zhang + Ahn + Dao + Langford) chose to package all four levers as one optimizer, rather than shipping them piecewise, is the interesting move — it implicitly bets that the levers compose multiplicatively rather than saturating each other.

The novel row-selection update rule is the paper’s proper delta over the prior work: Dion selected a low-rank factor via power iteration + error feedback, Dion2 selected an α-fraction submatrix, and Dion3 selects rows only — a strict simplification that the abstract claims dominates both. Whether this stacks with Delphi: Scaling Laws That Extrapolate 300× Past the Fit‘s Hyperball-style Frobenius-sphere retraction (AdamH) or Controlled LLM Training on Spectral Sphere‘s SSO spectral-sphere constraint is the natural follow-on: Dion3 attacks the wall-clock cost, those attack the quality per unit compute, and they occupy orthogonal enough axes to plausibly stack.