Muon はニューラルネットワークの隠れ層(hidden layers)専用の最適化手法です。
特に 2 次元パラメータ(線形層や畳み込み層の重み行列)を対象とし、
SGD-Momentum の更新方向を “直交化 (orthogonalization)” する ことで安定かつ高速な学習を実現します。
Muon は、次のような実績を挙げています:
これらの結果は、既存の AdamW や Shampoo などの最先端手法を上回るものであり、
「最も速い LLM トレーニングを支えるオプティマイザ」として注目を集めています。
Muon は「MomentUm Orthogonalized by Newton-Schulz」の略です。
SGD-Momentum の更新 $\mathbf{G}$ に対して、
Newton–Schulz 反復を適用して更新行列を直交化します。
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
この処理は、更新行列を「最も近い半直交行列」へ写像する操作に対応します。
つまり、更新行列 $\mathbf{G}$ の特異値分解 $\mathbf{G} = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top$ において、
$\mathbf{\Sigma}$ を単位行列に置き換える操作と同等です。
SGD-Momentum や Adam が生成する更新行列は、
実際には 非常に高い条件数(condition number) を持っています。
つまり、更新方向がほぼ低ランクで、特定の少数の方向に支配されてしまっているのです。
この状態では、「珍しい方向(rare directions)」の情報が抑圧され、
学習の多様性が失われます。
Muon の直交化ステップは、この隠れた方向のスケールを引き上げることで、
よりバランスの取れた学習を実現します。
🔍 要するに:Muon は「SGD の偏った更新を均等に整える」役割を持つ。
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 はこれらの系譜の中で、理論的に正当化され、実装効率も高いという点で特異的です。
特に、Bernstein & Newhouse (2024) の分析に基づき、
Shampoo の行列逆平方根計算を NS 反復で近似するというアイデアを発展させています。
行列の直交化にはいくつかの方法がありますが、Muon は以下の理由で
Newton–Schulz 反復を採用しています。
| 手法 | 利点 | 欠点 |
|---|---|---|
| SVD 分解 | 理論的に明確 | GPU 上で非常に遅い |
| Coupled Newton iteration | Shampoo で使用 | FP32 精度が必要・遅い |
| Newton–Schulz iteration | bfloat16 で安定・高速 | 近似法ゆえチューニング必要 |
Muon のバージョンでは、係数 $(a, b, c)$ を
勾配ベースで最適化し、5 ステップ以内で収束するように調整されています。
これにより、実用上の FLOP オーバーヘッドは 1%未満 に抑えられます。
Shampoo は、勾配をテンソル単位で「前処理」する 2 階最適化法です。
Muon はその簡略形とも言えます。
したがって、Momentum を無効化した Muon = “瞬間的 Shampoo” と見ることも可能です。
Orthogonal-SGDM(Tuddenham et al., 2022)は、
SVD を用いて勾配を直交化 → Momentum を適用する手法でした。
しかし、実装が遅く、調整済み SGD-Momentum に劣る結果が報告されています。
Muon は Momentum の適用順序を逆転 し、
さらに SVD → Newton–Schulz に置き換えることで高速化を実現しています。
各層パラメータ $\mathbf{W} \in \mathbb{R}^{n \times n}$ に対して、
Muon の FLOP オーバーヘッドは次式で近似されます:
ここで $k$ は NS 反復回数(通常 5)、$b$ はバッチサイズ。
NanoGPT でも LLaMA 3 でも 1%未満の FLOP 増加 に留まります。
Jordan は、Muon の成功を「NanoGPT speedrunning」という
競争的タスク (competitive task) の中で検証することを強調しています。
この形式の利点:
Muon は 2024 年 10 月に記録を樹立して以来、
12 回連続で最速記録を保持しており、7 人の研究者がこれを使用。
したがって、「Muon が AdamW より速い」ことは実証的に再現されているといえます。