SignMuon: 通信効率的な分散Muon最適化
SignMuonは、signSGDの多数決符号集約とMuonの極座標ステップフレームワークを組み合わせた1ビット、行列認識オプティマイザです。各ワーカーはニュートン・シュルツ反復により運動量の極因子を計算し、要素ごとの符号のみを送信して多数決で集約します。スペクトルノルム平滑性と有界分散確率的勾配の下で、O(1/√T)の非凸収束率を達成します。330のCIFAR-10/ResNet-50構成において、SignMuonは最高の検証精度92.15%を達成。4-GPU多数決変種は一致する有効バッチでトレーニング時間を37%削減します。nanoGPTでは、他の符号ベースのベースラインより低いパープレキシティと優れたアニータイム性能を示し、最大16GPUまでの弱スケーリングに有利です。
最近、arXivに投稿された論文「SignMuon: Communication-Efficient Distributed Muon Optimization」では、大規模ニューラルネットワークの分散トレーニングにおける通信ボトルネックを解決するための新しい最適化手法SignMuonが提案されました。この手法は、signSGDの多数決符号集約メカニズムとMuon最適化器の極座標更新フレームワークを融合し、1ビットの通信オーバーヘッドを実現しつつ、重みテンソルの行列構造を考慮した最適化を可能にします。
分散トレーニングでは、各ワーカーがまずニュートン・シュルツ反復を用いて運動量の極因子を計算し、Muon風の更新方向を求めます。その後、要素ごとの符号のみを送信し、多数決によって集約します。オプションとして、追加の通信コストなしに直交性を強化するためのローカル極座標ステップを実行することもできます。
理論解析では、スペクトルノルム平滑性と有界分散確率的勾配の仮定の下で、ℓ₁ベースの定常性尺度に対してO(1/√T)の非凸収束率が証明されています。ノイズが単峰対称分布に従う場合、M台のワーカーによる多数決は確率的項を1/√Mに削減し、signSGDと同等の理論的利点を提供します。α-βモデルでは、分散SignMuonは各反復で1回の整数和allreduceのみを必要とし、すべての直交化はローカルで行われるため、float32と比較して32倍（int8と比較して4倍）の帯域幅削減を実現します。
実験では、330のCIFAR-10/ResNet-50構成においてSignMuonが最高の検証精度92.15%を達成。4-GPU多数決変種は、一致する有効バッチで92.02%の精度を達成し、トレーニング時間を37%削減しました。nanoGPT言語モデルでは、他の符号ベースのベースラインと比較して低いパープレキシティと優れたアニータイム性能を示し、最大16GPUまでの弱スケーリングで良好な結果が得られました。
本研究は、符号通信と行列認識更新によって帯域幅要件を大幅に削減する新たなアプローチを提供し、大規模モデルトレーニングへの応用が期待されます。