Dion3: Full-Stack Orthogonal Updates

2026-08-11

AI Summary

This work addresses the computational inefficiencies of the Muon optimizer, which stem from the cubic-time complexity of its orthogonalization step and communication overhead induced by weight sharding. To overcome these limitations, the authors propose a full-stack optimization strategy comprising the Gram-Newton-Schulz orthogonalization algorithm, high-performance kernels implemented via symmetry-aware CuteDSL, a megabatching communication scheme, and a refined update mechanism that orthogonalizes only selected rows of the momentum matrix. The resulting method matches or exceeds Muon’s convergence performance while reducing per-step optimization time by up to 6×. Released as an open-source, plug-and-play alternative, this approach significantly outperforms its predecessor, Dion.

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. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6x. Dion3 is available via the dion package (https://github.com/microsoft/dion) as a drop-in replacement for Muon.

Problem

• orthogonalization
• optimization overhead
• communication cost
• Newton-Schulz iteration
• weight sharding

Innovation

• orthogonal optimization
• Newton-Schulz iteration
• parameter sharding
• megabatching
• momentum compression