技術的サマリー：Muon 最適化のための Muown - 行ノルム制御

問題定義

Muon は、トランスフォーマーの線形層の行列構造を活用することで AdamW を凌駕し、大規模言語モデル（LLM）の前学習における競争力のあるオプティマイザとして浮上してきた。しかし、最近の観察では、Muon のスケーリング時の挙動は重み減衰に敏感であり、bfloat16 でのオーバーフローやアテンション対数オッズの発散といった数値的不安定性に陥りやすいことが示されている。

特定された重要な課題は、Muon を非結合型（decoupled）の重み減衰なしで使用する場合、学習中に重み行列のスペクトルノルム（$\|W_t\|_{S_\infty}$）が一貫して上方にドリフトする現象である。このドリフトは、安定した特徴学習に必要なスペクトル条件（$\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$）に違反する。非結合型の重み減衰は重み行列を均一に縮小させることでこれを緩和できるが、これは成長を駆動する特定の方向に関わらず全スペクトルをスケーリングする非標的的な介入であり、学習初期段階での収束を遅らせる可能性がある。本論文は、このスペクトルノルムドリフトの具体的な経験的駆動要因を特定し、より標的的な制御メカニズムを提案することを目的としている。

手法：Muown

著者は、集約的な制約ではなく再パラメータ化戦略を通じてスペクトルノルムドリフトに対処する、Muon のドロップイン代替となるMuownを導入する。

1. スペクトルノルムの分解

核心的な洞察は、スペクトルノルムを 2 つの因子に分解することである（命題 1）：
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
ここで：

• $g_t = \|W_t\|_{row}$ は行の大きさ（行のユークリッドノルム）のベクトルである。
• $D_t$ は行正規化された行列（$W_t = \text{Diag}(g_t)D_t$）である。
• $C_t = D_t D_t^\top$ は単位ノルム行のグラム行列である。
• $P_t$ は正規化された行の大きさの対角行列である。
• $\lambda_{\max}(P_t C_t P_t)$ は行コヒーレンス因子を表す。

5 億パラメータモデルにおける経験的分析により、行の大きさ因子（$\|g_t\|_\infty$）がスペクトルノルムドリフトの主要な駆動要因であり、$\|W_t\|_{S_\infty}$ と連動して移動することが明らかになった。対照的に、行コヒーレンス因子は学習を通じて良好に振る舞い、有界である。

2. 再パラメータ化と最適化

この診断に基づき、Muown は重み行列 $W$ を以下のように再パラメータ化する：
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
ここで、$g \in \mathbb{R}^m$ は明示的な行の大きさベクトルを表し、$R \in \mathbb{R}^{m \times n}$ は方向成分を表す。

• 方向成分（$R$）： Muon は変更を加えずに $R$ に適用され、スペクトルノルム幾何学（最急スペクトル降下）の下で処理される。
• 大きさ成分（$g$）： 行の大きさ $g$ は明示的なオプティマイザ変数として昇格される。これらは $\ell_\infty$ 幾何学の下で更新され、これは対角行列の作用素ノルムに対する自然な双対ノルムである。
• 更新則： $g$ の更新には、$\ell_\infty$ 幾何学下での平滑化された SignSGD として機能する Adam オプティマイザが用いられ、確率的ノイズを処理する。一方、$R$ の更新は、ニュートン・シュルツ直交化を用いた標準的な Muon 手順に従う。

このアプローチはオプティマイザの状態内で暗黙的に実装され、モデル構造のフォワードパスは変更されない。

主要な貢献

1. スペクトルドリフトの診断： 本論文はスペクトルノルムを分解し、Muon 下でのスペクトルノルム成長の経験的駆動要因として最大行大きさ（$\|g_t\|_\infty$）を特定し、安定した行コヒーレンス因子と区別する。
2. Muown アルゴリズム： 行の大きさを $\ell_\infty$ 幾何学の下で更新される明示的な学習可能変数として扱い、Muon の方向更新を維持する新しいオプティマイザ。これは、重み減衰の無差別な縮小やスペクトル再パラメータ化の直交性事前分布を回避する。
3. 収束保証： 著者は、Muown が決定論的および確率的の両方の領域で最適な非凸収束率を達成することを証明する。分析は、更新幾何学と整合する混合双対ノルム（大きさには $\ell_1$、方向には Schatten-1 ノルム）の下で行われる。理論的 bound は、Muon と比較して実効的な確率的ノイズ係数の減少を示唆する。
4. 経験的検証： 1 億 2400 万から 27 億パラメータまでのモデルサイズにおける FineWeb-Edu での広範な実験により、一貫した改善が実証された。

実験結果

• パープレキシティの改善： Muown は、テストされたすべてのモデルスケール（1 億 2400 万から 27 億）および学習率において、Muon、SOAP、AdamW、Lion よりも一貫して低いパープレキシティを達成する。
• ハイパーパラメータへの頑健性：
  • 学習率： Muown は異なるモデル幅にわたって準最適学習率のプレートを広げ、スケールごとの学習率チューニングへの感度を低減する。
  • 重み減衰： Muown は重み減衰のチューニングの必要性を軽減する。重み減衰なし（$\lambda=0$）でも良好に機能するのに対し、Muon は発散や最適でない性能を防ぐために慎重なチューニング（最適 $\lambda=0.1$）を必要とする。
• 安定性： Muown は、Muon で観察される体系的なスペクトルノルムドリフトを防ぐ。高学習率領域（例：$4 \times 10^{-3}$）において、Muon はしばしば発散するが、Muown は安定し、準最適の状態を維持する。
• 勾配特性： 分析により、Muown は学習を通じて Muon よりも高い正規化されたミニバッチ勾配の実効ランクを維持し、より低い確率的ノイズ係数を示すことが示された。
• 計算オーバーヘッド： この手法はステップ時間のオーバーヘッドを無視できる程度にしか導入しない。オプティマイザがランク間でシャード化される場合（分散環境）、Muown と Muon のステップ時間は区別がつかなくなる。メモリオーバーヘッドは最小限（行数に比例）であり、レイヤーあたり数メガバイトのみを追加する。

意義と主張

本論文は、Muon で観察されるスペクトルノルムドリフトがオプティマイザの不可避な性質ではなく、制御されていない行の大きさの成長の結果であると主張する。この駆動要因を分離し、明示的なオプティマイザ変数として昇格させることで、Muown は以下の成果を達成する：

• 収束の改善： 積極的な重み減衰を必要とせず、パープレキシティと安定性の向上。
• 理論的妥当性： 更新の混合性質を明示的に考慮する幾何学を考慮したフレームワークにおける最適な収束率。
• 実用的効率： 分散環境において Muon と同等のランタイムを維持しつつ、ハイパーパラメータ選択に対する優れた頑健性を提供するドロップイン代替手段。

著者は、Muown を、特徴学習の理論的安定性要件と大規模 LLM 学習の実践的ニーズの間のギャップを埋める解決策として位置づけ、非結合型重み減衰に対するより標的的な代替案を提供する。