技術的概要：SignMuon 通信効率に優れた分散ムオン最適化

1. 問題提起

大規模ニューラルネットワーク、特に大規模言語モデル（LLM）の分散トレーニングは、通信帯域幅とレイテンシにおいて重大なボトルネックに直面しています。同期データ並列トレーニングでは、ワーカーが大量のデータ（勾配、更新、またはオプティマイザ状態）を常に同期させる必要があります。現在のアプローチには 2 つの主な制限が存在します：

1. フル精度通信: 標準的なオプティマイザは、フル精度（例：float32）の勾配または更新を送信しますが、これは帯域幅の観点から高コストです。
2. 行列構造の座標無視: signSGD などの多くの人気のある通信効率化手法は、重みテンソルを 1 次元ベクトルとして扱います。これらはパラメータ（MLP の重み、アテンション投影など）の固有の行列構造を無視しており、行列幾何学に由来する最適化の利益を見逃す可能性があります。

既存の解決策は、しばしばこれらの問題のいずれか一方のみに対処しています：符号ベースの手法は極端な圧縮を達成しますが行列の認識が欠如しており、Muon のような行列認識型オプティマイザは収束幾何学を改善しますが、フル精度同期の高通信コストを維持しています。

2. 手法：Sign-Muon

著者は、符号ベース手法の通信効率と Muon の行列認識型幾何学を統合するハイブリッドオプティマイザであるSign-Muonを提案します。

中核メカニズム

Sign-Muon は、2 つの異なるコンポーネントを組み合わせることで動作します：

1. Muon 様式の極方向: 生勾配の代わりに、各ワーカーはモーメント行列を計算し、極分解を通じて更新方向を導出します。これは、モーメント行列を直交化するニュートン・シュルツ（NS）反復を用いて局所的に近似されます。このステップにより、更新方向が重みテンソルのスペクトルノルム幾何学を尊重することが保証されます。
2. 1 ビット符号集約: フル精度の極方向を送信する代わりに、各ワーカーは計算された方向の要素ごとの符号（$\{-1, +1\}$）のみを送信します。
3. 多数決集約: 分散環境では、ワーカーはこれらの符号行列を多数決によって集約します。これは、効率的に整数のSUM all-reduce（ワーカー間での符号の合計）を実行し、それに続いて局所的な閾値処理操作（$\text{sign}(\sum S^{(m)})$）を行うことで実装されます。

アルゴリズムフロー（分散）

1. 局所計算: 各ワーカー $m$ は、確率的勾配 $G^{(m)}_t$ を計算し、局所モーメント $M^{(m)}_{t+1}$ を更新し、ニュートン・シュルツ反復を通じて極因子 $U^{(m)}_t$ を計算します。
2. 符号抽出: ワーカーは要素ごとの符号 $S^{(m)}_t = \text{sign}(U^{(m)}_t)$ を抽出します。
3. 通信: ワーカーは単一の集合通信ステップを実行します：
   • All-Reduce 変種: 整数 SUM all-reduce が符号を集約します。
   • All-Gather 変種（1 ビット）: 符号はビットにパックされ（エントリあたり 1 ビット）、All-Gather がパックされたバッファを伝播します。
4. 局所更新: 各ワーカーは局所的に多数決 $\bar{S}_t$ を計算し、パラメータを更新します：$W_{t+1} = W_t - \eta_t \bar{S}_t$。
5. オプションの局所直交化: ワーカーは、追加の通信コストを負うことなく直交性をさらに強制するために、集約された符号行列に局所的な極ステップを適用する場合があります。

通信効率

• ペイロードサイズ: この手法は、パラメータあたりの通信を 1 ビット（実装の簡素化のために int8 を使用する場合は 8 ビット）に削減します。
• 削減率: float32 の all-reduce と比較して、Sign-Muon は（ビットパックを使用する場合）32 倍の帯域幅削減、または（int8 を使用する場合）4 倍の削減を達成します。
• 局所計算: すべてのスペクトルノルム正規化と直交化（ニュートン・シュルツ）は局所的に実行され、追加のネットワーク通信コストは発生しません。

3. 理論的貢献

本論文は、スペクトルノルム滑らかさと有界分散の確率的勾配の下での収束解析を提供します。

• 収束率: $\ell_1$ ベースの定常性指標に対して、Sign-Muon は$O(1/\sqrt{T})$の非凸収束率を達成します。
• ノイズ低減: 単峰対称ノイズの仮定の下、$M$ 個のワーカーにわたる多数決メカニズムは、確率的ノイズ項を$1/\sqrt{M}$の係数で削減し、分散 signSGD の理論的利点と一致します。
• 幾何学: この解析は、座標ごとの滑らかさに依存する標準的な signSGD 解析とは区別される、スペクトルノルム降下補題を利用します。スペクトル正規化により、更新方向のスペクトルノルムが最大 1 であることが保証されます。

4. 実証結果

著者は、Sign-Muon を 2 つのワークロード、CIFAR-10（ResNet アーキテクチャによる画像分類）およびnanoGPT（言語モデリング）で評価しました。

CIFAR-10（ResNet-50）

• 精度: 330 のハイパーパラメータ構成全体で、Sign-Muon は92.15% の最高検証精度を達成しました。
• 分散効率: 4 GPU 多数決変種は、**92.02%の精度（単一 GPU パフォーマンスと一致）を達成し、マッチングされた実効バッチサイズと比較してトレーニング時間を37%**削減しました。
• スケーリング: Sign-Muon は、SignAdam や SignSGD と比較して優れた弱スケーリング挙動を示しました。16 GPU において、これは有意に低いトレーニング時間のオーバーヘッド（All-Gather 設定におけるベースラインの +150–163% に対して +77.8%）を招きました。

nanoGPT

• パフォーマンス: Sign-Muon は、他の符号ベースのベースラインと比較して、低いパープレキシティと、ウォールクロック時間全体での最小パープレキシティを実行する「anytime」パフォーマンスを達成しました。
• スケーリング: 16 GPU まで望ましい弱スケーリングが観察されました。
• ハイパーパラメータ: この手法は正規化方式（フロベニウス対スペクトルスケーリング）に対して頑健であり、最小限のチューニングを必要としました（例：ニュートン・シュルツ反復 1 回でしばしば十分でした）。

メモリとオーバーヘッド

• メモリ: Sign-Muon は、SGD と同様の永続的なメモリフットプリント（1 つのモーメントバッファ）を維持し、2 つのモーメントベクトルを維持する Adam/AdamW よりも著しく低いです。
• 計算: 局所的なニュートン・シュルツ反復はステップごとの計算オーバーヘッドを追加しますが、特に帯域幅制約環境では、通信時間の大幅な削減によって相殺されます。

5. 意義と主張

本論文は、Sign-Muon が以前に明確に分かれていた 2 つの研究分野の間のギャップを成功裏に橋渡しすると主張しています：

1. 通信効率: 符号ベース手法の極端な圧縮（1 ビット）と頑健性（多数決）を維持します。
2. 行列認識: 通信効率を犠牲にすることなく、Muon（極分解）の幾何学的利点を組み込みます。

主要な主張:

• ギャップの橋渡し: Sign-Muon は、多数決符号集約を Muon の極ステップフレームワークに直接埋め込んだ最初のオプティマイザです。
• スケーラビリティ: 帯域幅が主要なボトルネックである大規模分散トレーニングのための実用的な解決策を提供し、float32 と比較して通信量を最大 32 倍削減します。
• 理論的妥当性: この手法は、座標ごとの滑らかさよりも行列構造の更新に対してより適切な指標であるスペクトルノルム滑らかさの下で、厳密な収束保証を提供します。
• 実用性: 標準的なベンチマーク（CIFAR-10、nanoGPT）における実証結果は、Sign-Muon がフル精度および他の符号ベースのベースラインのパフォーマンスに匹敵し、しばしば凌駕するだけでなく、分散環境でトレーニング時間を大幅に削減することを示しています。

著者は、直交化の局所計算が通信コストの観点から「無料」であるため、この手法はネットワーク帯域幅がボトルネックである場合に特に効果的であると強調しています。