Muon Optimizer: Newton–Schulz直交化の定義・理論・実装上の注意
2026年10月の再サーベイ: 研究史・重要文献・最近の進展を整理した新版を追加しました(2026年10月3日基準)。以下は従来のメモです。定義・適用条件の訂正は新版を参照してください。
更新方針(2026-09-30): Muon は発展中の手法です。以下では、定義と線形代数上の事実、 提案者が報告した実験結果、著者自身が「推測」としている説明を区別します。特定のモデル・ 計算環境での速度記録を、あらゆる学習タスクでの優位性とは解釈しません。
Papers
主要な一次資料は、提案者による Muonの定義・実験・未解決問題 と、 行列に割り当てるノルムから最急降下を捉える Old Optimizer, New Norm です。
Blog Posts
- 「Modular Manifolds」(Jeremy Bernstein, Thinking Machines, 2025)
- 「Muon: An optimizer for hidden layers in neural networks」(Keller Jordan et al.)
はじめに:Muon とは何か?
Muon はニューラルネットワークの隠れ層(hidden layers)専用の最適化手法です。
特に 2 次元パラメータ(線形層や畳み込み層の重み行列)を対象とし、
SGD-Momentum の更新方向を “直交化 (orthogonalization)” することで、提案者の実験では
学習の安定性とtime-to-targetの改善を狙います。
提案者の記事の実験結果では、次の値が報告されています:
- CIFAR-10 における 94% 精度達成までの速度記録を 3.3 → 2.6 A100 秒 に短縮
- NanoGPT speedrun における validation loss 3.28 達成速度を 1.35 倍高速化
- 774M〜1.5B パラメータへのスケールに成功
- GPT-2 XL レベル性能を持つモデルを 13.3 → 10 時間 (8xH100) で到達
これらは提案者がまとめた特定の speedrun と小〜中規模事前学習の結果です。比較対象、 ハイパーパラメータ、停止条件、ハードウェアを固定しない横断比較には使えません。 提案者の記事 も、スケール、分散実装、 fine-tuning や強化学習への適用を未解決問題として挙げています。
Muon の定義
Muon は「MomentUm Orthogonalized by Newton-Schulz」の略です。
SGD-Momentum の更新 $\mathbf{G}$ に対して、
Newton–Schulz 反復を適用して更新行列を直交化します。
数式的定義(概念図)
\[\mathbf{G}_{\text{new}} = \text{NewtonSchulz}_5(\mathbf{G})\]Newton–Schulz 反復は、次のように定義される行列近似法です:
def newtonschulz5(G, steps=5, eps=1e-7):
a, b, c = (3.4445, -4.7750, 2.0315)
X = G.bfloat16()
X /= (X.norm() + eps)
if G.size(0) > G.size(1):
X = X.T
for _ in range(steps):
A = X @ X.T
B = b * A + c * A @ A
X = a * X + B @ X
if G.size(0) > G.size(1):
X = X.T
return X
標準的な収束型Newton–Schulz反復は、適切な初期スペクトル条件の下で、非零特異値を1へ 近づけて極分解因子 $\mathbf{U}\mathbf{V}^\top$ を計算できます。入力がfull rankなら、この因子は Frobenius距離で最も近い半直交行列です。
一方、上のコードの係数は 有限5ステップで良い近似を得る目的で調整されています。 $a+b+c=0.701\ne1$ なので特異値1は不動点ではなく、この係数の反復を無限に続ければ polar factorへ収束する、という意味ではありません。厳密なSVD置換でもありません。 また階数欠損行列では零特異値は反復で零のままで、得られる $UV^\top$ は部分等長写像です。 完全な半直交行列への拡張、およびFrobenius最近点は一般に一意ではありません。 導出と係数の設計は提案者による説明を参照してください。
なぜ更新を直交化するのか?
SGD-Momentum や Adam が生成する更新行列は、提案者が調べた Transformer 等では 高い条件数を示したと報告されています。これは一般定理ではなく観察です。
提案者の仮説では、この状態で「珍しい方向(rare directions)」の情報が相対的に抑圧され、 有用な更新方向を十分に使えない可能性があります。これは性能低下を必然とする定理ではありません。
Muon の直交化が小さい特異値方向を相対的に持ち上げるという説明は線形代数的には自然ですが、 それが性能向上の因果機構であるという部分は提案者も speculation と明記しています。
🔍 要するに:Muon は、momentum更新の特異値を有限回の行列反復で揃える方向へ変換する。
技術的背景と研究的系譜
Muon は突発的に生まれたアイデアではなく、
過去 10 年にわたる 直交化ベースの最適化研究の流れ の上に立っています。
| 年 | 手法 | 内容・位置づけ |
|---|---|---|
| 2015 | Stochastic Spectral Descent (Carlson et al.) | 勾配を SVD で直交化して RBM を最適化 |
| 2018 | Shampoo (Gupta et al.) | 勾配をテンソルごとに前処理する 2 階法 |
| 2020 | Orthogonal-SGDM (Tuddenham et al.) | SVD 直交化+ Momentum 適用(ただし遅い) |
| 2024 | Old Optimizer, New Norm (Bernstein & Newhouse) | Shampoo における NS 反復を提案 |
| 2024 | Muon (Jordan) | Momentum 更新を NS 反復で直交化する実用最適化法 |
Muon はこれらの系譜の中で、ノルムに基づく一次法としての動機とGPU上の実装効率を組み合わせています。
特に、Bernstein & Newhouse (2024) の分析に基づき、
Shampoo の行列逆平方根計算を NS 反復で近似するというアイデアを発展させています。
NS(Newton–Schulz)反復の選択理由
行列の直交化にはいくつかの方法がありますが、Muon は以下の理由で
Newton–Schulz 反復を採用しています。
| 手法 | 利点 | 欠点 |
|---|---|---|
| SVD 分解 | 理論的に明確 | GPU 上で非常に遅い |
| Coupled Newton iteration | Shampoo で使用 | FP32 精度が必要・遅い |
| Newton–Schulz iteration | bfloat16 で安定・高速 | 近似法ゆえチューニング必要 |
Muon のバージョンでは、係数 $(a, b, c)$ を
勾配ベースで最適化し、5 ステップ以内で収束するように調整されています。
提案者の対象モデルと実装では、小さいwall-clock overheadが報告されています。ただし
比率は行列形状、バッチサイズ、通信方式、カーネル実装に依存します。
Shampoo との関係
Shampoo は、勾配をテンソル単位で「前処理」する 2 階最適化法です。
Muon はその簡略形とも言えます。
- Shampoo の更新式から累積項を除去すると、直交化された勾配が得られる
- Momentum を加えて直交化を後段に持ってくると、Muon の更新式に一致する
したがって、Momentum を無効化した Muon = “瞬間的 Shampoo” と見ることも可能です。
Orthogonal-SGDM との違い
Orthogonal-SGDM(Tuddenham et al., 2022)は、
SVD を用いて勾配を直交化 → Momentum を適用する手法でした。
しかし、実装が遅く、調整済み SGD-Momentum に劣る結果が報告されています。
Muon は Momentum の適用順序を逆転 し、
さらに SVD → Newton–Schulz に置き換えることで、提案者の実装では直交化処理を高速化しています。
実装上のポイント
- Muon は 2D パラメータ(線形層や Conv 層)専用。
- Conv 層では最後の 3 軸を flatten して使用。
- スカラー・ベクトルパラメータは AdamW で更新するのが推奨。
- Transformer では:
- Embedding 層と出力層(Classifier)は AdamW で更新
- Q/K/V パラメータは個別に Muon を適用する方が安定。
実験的知見
- 提案者が試した範囲では Nesterov Momentum が通常のmomentumより少し良かった
- 提案者が試したTransformerでは Q/K/V 分離 が良かった
- NanoGPT speedrun において 35% 高速化
- 1.5B規模までの事前学習例が報告されている
いずれも提案者の実験報告の範囲であり、 データセットやチューニング予算を変えた場合まで保証する結果ではありません。
パフォーマンスとオーバーヘッド
行列 $\mathbf{W}\in\mathbb{R}^{n\times m}$($m\le n$)に対し、提案者の実装では Newton–Schulz 1回あたりの行列積コストを概ね
\[2(2nm^2+m^3)\]FLOPs と見積もっています。総学習コストに占める割合は、反復回数だけでなく層の形状、 sequence length、batch size、並列化に依存するため、単純に $k/b$ とはなりません。 詳細はMuonのruntime analysisを参照してください。
研究的意義:競争的タスクによる検証文化の確立
Jordan は、Muon の成功を「NanoGPT speedrunning」という
競争的タスク (competitive task) の中で検証することを強調しています。
この形式の利点:
- ベースラインがすでに最適化されている
- 誤った過大評価は次の記録で自然修正される
speedrun は強いベースラインを共有できる有用な実験環境ですが、記録回数は時間とともに変わり、 そのまま一般的なoptimizer順位を意味しません。再現時は同一コード、同一停止条件、同一計算資源で AdamW等を再調整して比較する必要があります。
参考文献
- Gupta et al. (2018) Shampoo: Preconditioned stochastic tensor optimization
- Bernstein & Newhouse (2024) Old Optimizer, New Norm
- Anil et al. (2020, 2024a,b) Scalable 2nd-order optimization; Shampoo with no accumulation
- Tuddenham et al. (2022) Orthogonalising Gradients to Speed Up Neural Network Optimisation
- Carlson et al. (2015–2016) Stochastic Spectral Descent
- Jordan (2024) Muon: An Optimizer for Hidden Layers in Neural Networks