Hiroki Naganuma

Papers

See awesome-muon-optimizer

Blog Posts

はじめに:Muon とは何か?

Muon はニューラルネットワークの隠れ層(hidden layers)専用の最適化手法です。
特に 2 次元パラメータ(線形層や畳み込み層の重み行列)を対象とし、
SGD-Momentum の更新方向を “直交化 (orthogonalization)” する ことで安定かつ高速な学習を実現します。

Muon は、次のような実績を挙げています:

これらの結果は、既存の AdamW や Shampoo などの最先端手法を上回るものであり、
「最も速い LLM トレーニングを支えるオプティマイザ」として注目を集めています。

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

この処理は、更新行列を「最も近い半直交行列」へ写像する操作に対応します。
つまり、更新行列 $\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 反復で近似するというアイデアを発展させています。

NS(Newton–Schulz)反復の選択理由

行列の直交化にはいくつかの方法がありますが、Muon は以下の理由で
Newton–Schulz 反復を採用しています。

手法 利点 欠点
SVD 分解 理論的に明確 GPU 上で非常に遅い
Coupled Newton iteration Shampoo で使用 FP32 精度が必要・遅い
Newton–Schulz iteration bfloat16 で安定・高速 近似法ゆえチューニング必要

Muon のバージョンでは、係数 $(a, b, c)$ を
勾配ベースで最適化し、5 ステップ以内で収束するように調整されています。
これにより、実用上の FLOP オーバーヘッドは 1%未満 に抑えられます。

Shampoo との関係

Shampoo は、勾配をテンソル単位で「前処理」する 2 階最適化法です。
Muon はその簡略形とも言えます。

したがって、Momentum を無効化した Muon = “瞬間的 Shampoo” と見ることも可能です。


Orthogonal-SGDM との違い

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 オーバーヘッドは次式で近似されます:

\[\text{Overhead} \le \frac{k \, n^2}{b \, n^2} = \frac{k}{b}\]

ここで $k$ は NS 反復回数(通常 5)、$b$ はバッチサイズ。
NanoGPT でも LLaMA 3 でも 1%未満の FLOP 増加 に留まります。

研究的意義:競争的タスクによる検証文化の確立

Jordan は、Muon の成功を「NanoGPT speedrunning」という
競争的タスク (competitive task) の中で検証することを強調しています。

この形式の利点:

  1. ベースラインがすでに最適化されている
  2. 誤った過大評価は次の記録で自然修正される

Muon は 2024 年 10 月に記録を樹立して以来、
12 回連続で最速記録を保持しており、7 人の研究者がこれを使用。
したがって、「Muon が AdamW より速い」ことは実証的に再現されているといえます。

参考文献