arXiv (ML)AI
SignMuon:通信効率的な分散Muon最適化手法
SignMuon: Communication-Efficient Distributed Muon Optimization
この記事についてAIに質問する →
日本語要約青い用語にマウスを合わせると解説が表示されます
大規模なニューラルネットワークの分散学習は、高精度の勾配通信とテンソルの行列構造を無視する座標ごとの最適化手法によってボトルネックに直面しています。このような課題に対して、研究者らは「Sign-Muon」という新しい最適化手法を提案しました。これは1ビット精度でありながら行列の構造を意識した最適化器で、signSGDの多数決符号集約とMuonの極座標ステップフレームワークを組み合わせています。
Sign-Muonの動作メカニズムは次の通りです。各ワーカーはニュートン・シュルツ反復を用いてモーメンタムの極因子を取得することで、Muon形式の方向を形成します。その後、符号のみを送信し、多数決によって集約します。さらにオプションの局所極座標ステップにより、追加の通信コストなしで直交性を強制できます。理論的には、スペクトラルノルム平滑性と有界分散確率勾配の条件下で、このアルゴリズムは非凸最適化に対して$\mathcal{O}(1/\sqrt{T})$の収束速度を達成します。複数ワーカー間の多数決投票により、確率的項の大きさを$1/\sqrt{M}$に削減でき、signSGDと同等の性能を実現します。
実験結果では、CIFAR-10とResNet-50を用いた330通りの構成において、Sign-Muonは最高の検証精度92.15%を達成しました。4GPUを用いた多数決投票バリアントでは92.02%の精度を達成しながら、同じ有効バッチサイズで37%の学習時間削減を実現しています。さらに通信効率の面でも顕著な改善があり、float32比較で32倍、int8比較で4倍の帯域幅削減を達成します。言語モデルのnanoGPTでの評価でも、他の符号ベースの手法を上回るperplexityを達成し、16GPUまでのスケーリングでも良好な性能を保持しています。