技術要約: DeMuon - グラフ上の行列最適化のための分散型Muon

問題定義

本論文は、通信グラフ上における完全分散型の行列最適化の問題に取り組んでいる。目的は、以下のマルチエージェント最適化問題を解くことである：
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \frac{1}{N} \sum_{i=1}^N f_i(X)$$
ここで、$N$ 個のエージェントは有向グラフ $\mathcal{G}=(\mathcal{V}, \mathcal{E})$ を介して接続されており、各ノード $i$ はローカルな目的関数 $f_i$ を保持している。行列変数をベクトル化することが多い従来の分散型最適化手法とは異なり、本研究は行列の構造を維持することに焦点を当て、特に中央集権的な Muon アルゴリズム（Jordan et al., 2024）の手法を活用している。目標は、中央サーバーなしで、局所的な目的関数の不均一性を緩和しつつ、グローバルな収束とコンセンサス（合意）を達成することである。

手法

本論文では、2つのアルゴリズム、DeMuon（DeMuonの分散型バリアント）および DeMuon-A（加速型分散DeMuon）を提案している。

1. DeMuon: 分散型Muon

DeMuonは、スペクトルノルムに基づく更新を勾配トラッキングのフレームワークに統合する。このアルゴリズムは、3つのシーケンス、すなわちローカル勾配推定値 $\{M^{(k)}\}$、グローバル勾配推定値 $\{V^{(k)}\}$、およびローカルイテレート $\{X^{(k)}\}$ を保持する。

• ローカルモーメンタム: 各ノードは、確率的勾配の指数移動平均を用いて、ローカル推定値 $M^{(k)}_i$ を更新する。
• 勾配トラッキング: ノードは、近傍からの情報を集約することで $V^{(k)}_i$ を更新し、グローバルな勾配を追跡する。これにより、局所的な不均一性にもかかわらず、勾配の方向に関するコンセンサスを確保する。
• 行列直交化: ローカルイテレート $X^{(k+1)}_i$ は、近傍の更新を集約することで更新される。具体的には、Muonの手法の中核であるスペクトルノルム ($\|\cdot\|$) を用いて、追跡された勾配に対して行列符号関数 ($\text{msgn}(\cdot)$) を適用する。これは、ノルム制約付き線形部分問題を解くものである。

DeMuonの理論的保証:
有界なノイズおよびスペクトルノルムと核ノルム（nuclear norm）に関する勾配リプシッツ連続性の仮定の下で、著者らは以下を確立している：

• コンセンサス誤差: 誤差 $\|X^{(k)} - \mathbf{1} \otimes \bar{X}^{(k)}\|$ は $O(\frac{\sqrt{N}\lambda \eta}{1-\lambda})$ で抑えられる。ここで $\lambda$ はグラフの混合率を特徴付ける。
• 定常性誤差: 平均定常性誤差は、適切に調整されたステップサイズ $\eta$ とモーメンタム $\theta$ のもとで、$O(K^{-1/4})$ のレートで収束する。これは、中央集権的なモーメンタム付き正規化SGDの最適収束レートと一致する。

2. DeMuon-A: 加速型分散DeMuon

DeMuon-Aは、マルチエクストラポレーション技術（He et al., 2026b）を分散型行列設定へと拡張したものである。

• メカニズム: 目的関数の高次平滑性を利用するために、1イテレーションあたり $q = p-1$ 回のローカル・エクストラポレーション・ステップを実行し、点 $Z^{(k,s)}_i$ を生成する。
• 仮定: これは、関数 $f_i$ の $p$ 次導関数に関する高次のリプシッチ条件（スペクトルノルムおよび核ノルムによって誘導される）を必要とする。
• 収束: これらの条件下において、DeMuon-Aは改善された定常性誤差レート $O(K^{-p/(3p+1)})$ およびコンセンサス誤差レート $O(K^{-(2p+1)/(3p+1)})$ を達成する。

主な貢献

1. アルゴリズムの提案: 著者らは、Muonのスペクトルノルムに基づく直交化と勾配トラッキングを取り入れた、初の完全分散型行列最適化手法である DeMuon を導入した。
2. 収束解析: 著者らは、定常性とコンセンサスの誤差の両面において、DeMuonおよびDeMuon-Aに対する厳密なグローバル収束保証を提供している。特筆すべきは、次元フリーな収束（行列の次元 $m, n$ に依存しない）を確立するために、行列マルチンゲールモーメント不等式を用いている点である。
3. 加速: 本論文は、高次の平滑性を活用するためにマルチエクストラポレーションを統合した DeMuon-A を提案し、ベースとなるDeMuonアルゴリズムよりも優れた収束レートを実現している。
4. 実証的検証: 分散ディープラーニング（Transformerの学習）および分散行列回帰に関する数値実験により、提案手法が既存の分散型ベースラインに対して実用的な優位性を持つことを示している。

実験結果

著者らは、DeMuonを3つのベースライン（DSGD：バニラな分散型SGD、DSGD-C：クリップ付きSGD、DSGD-N：勾配トラッキングを用いた正規化SGD）と比較評価した。

• 設定: 実験は、Multi30kデータセットを用いて学習された7.4MパラメータのGPTモデルを用い、3つのグラフトポロジー（Complete、Directed Exponential、Ring）において実施された。
• パフォーマンス: DeMuonは、すべてのトポロジーにおいて、訓練損失、検証損失、およびパープレキシティの観点から、一貫してすべてのベースラインを上回った。
• コンセンサス: DeMuonのコンセンサス誤差はベースラインと同等またはそれ以下であり、エージェント間の効果的な合意を示している。
• 詳細: Complete Graphトポロジーにおいて、DeMuonは最終的な検証損失 0.240（パープレキシティ 1.27）を達成し、これはDSGD-N（0.302）やDSGD（7.031）よりも大幅に低い値である。

意義と主張

本論文は、DeMuonが、中央集権的なディープラーニングにおけるMuonの経験的な成功と、分散型最適化の要件との間の溝を埋めることに成功したと主張している。

• 次元フリーな収束: 行列の次元 $(m, n)$ に明示的に依存する場合がある最近の研究（例：Zhang et al., 2026a）とは異なり、著者らは、行列マルチンゲールモーメント不等式を使用することで、行列の次元に依存しない収束保証が可能になったと主張している。
• 実用的な利点: 実験結果は、Muonの特定の構造（スペクトルノルム制約と行列直交化）が、分散環境においても明確な実用的利点を提供することを示唆している。これは、行列構造が固有である大規模言語モデルの学習において特に顕著である。
• 理論的厳密性: 本研究は、適切な勾配トラッキングと適切なノルムベースの更新が行われる限り、行列最適化におけるモーメンタムと正規化の恩恵が、中央のコーディネーターなしの完全分散環境においても維持できることを確立している。

著者らは、分散型ディープラーニングや、分散PCA、行列補完といった古典的な問題への応用を動機として、高度な行列最適化技術を分散型の世界へと持ち込むためのステップとして本研究を位置づけている。