Dion3: Full-Stack Orthogonal Updates
This presentation explores Dion3, a comprehensive systems redesign that addresses the scaling bottlenecks of Muon-family optimizers. Through coordinated algorithm, kernel, and communication innovations—including Gram Newton-Schulz reformulation, symmetry-aware GPU kernels, fractional updates, and megabatched communication—Dion3 reduces optimizer-step time by up to 6× while preserving or improving validation loss. The talk examines how mathematical equivalence, numerical stability, and hardware-aware implementation combine to make orthogonal optimization practical at scale.Script
Training large language models with orthogonal updates delivers better optimization quality, but standard Muon pays a cubic cost that grows with matrix rectangularity and distributed sharding. Dion3 cuts that optimizer-step time by up to 6 times through a coordinated redesign of algorithm, kernel, and communication.
The central mathematical insight is Gram Newton Schulz, which reformulates orthogonalization around the smaller symmetric Gram matrix. Instead of iterating on the full rectangular matrix, it computes the Gram product once, refines the symmetric inverse square root, and reconstructs the output in a single final multiplication. For typical transformer aspect ratios, this cuts floating point operations by 55 percent.
But mathematical equivalence does not guarantee numerical stability. Naive half precision Gram Newton Schulz introduces spurious negative eigenvalues and eigenvector drift, causing training to spike and diverge. The solution is restarting: after 2 iterations, the algorithm reconstructs an intermediate approximation, recomputes its Gram matrix, and begins a fresh sequence. This stabilization is not an implementation detail, it is part of the algorithmic specification.
Dion3 goes further by orthogonalizing only a fraction of the momentum matrix, selecting rows with the largest norms and applying error feedback to preserve residual information. Contrary to initial expectations, this fractional update improves validation loss. At 14 billion parameters, selecting one quarter of the rows delivers a 0.027 cross entropy improvement and a 0.7 percentage point gain in downstream accuracy.
The full stack compounds these gains. Symmetry aware CuteDSL kernels exploit the structure exposed by Gram Newton Schulz for another 1.5 to 2 times speedup. Fractional updates with selection fraction one quarter add roughly 3.7 times. Megabatching groups all matrices with the same shape into a single collective, cutting communication rounds and latency overhead. Together, these contributions reduce optimizer step time from 26 times AdamW down to approximately 4 times AdamW for a 7 billion parameter model.
Dion3 shows that orthogonal optimization can scale when algorithm, kernel, and communication are designed together. The unexpected quality gains from fractional updates suggest new optimization dynamics worth exploring across architectures and training regimes. Visit EmergentMind.com to dive deeper into this work and create your own research videos.
