Dion3: フルスタック直交更新

要旨

Muonオプティマイザは、三次時間のNewton-Schulz直交化ステップにより、大きなオーバーヘッドコストを伴う。重みがシャーディングされると、通信オーバーヘッドがこの計算コストをさらに悪化させ、多くの設定においてMuonの利点を損なう。我々は、スタックのあらゆるレベルでこのオーバーヘッドに対処するMuonの改訂版であるDion3を提案する。我々のGram Newton-Schulzアルゴリズムは直交化のFLOPコストを削減し、CuteDSLカーネルは対称性を活用してこれを高速化し、メガバッチ戦略は通信オーバーヘッドを削減する。さらに、各ステップでモーメンタム行列の一部の行のみを直交化対象として選択するという、コストをさらに削減する単純な更新規則の変更を提案する。この更新規則は、Muonの別の「圧縮」版であるDionを、速度と性能の両方で改善する。全体として、Dion3はMuonが達成する損失に匹敵するかそれを改善しつつ、オプティマイザのステップ時間を最大6倍削減する。Dion3はdionパッケージ（https://github.com/microsoft/dion）から利用可能であり、Muonのドロップイン代替として提供される。

English

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.