TT Lab
ブログ

한국어 English 日本語 中文

はじめる
ブログ

ブログ

Muon 解説: 勾配行列を直交化するオプティマイザ

ステータス: ドラフト(2026-10-31)。 論文の数値はarXivの原典(本文は同一版のHTML)を直接開き、表や図の番号とともに引用しました。特異値の図(第3節)は筆者がノートPCのCPUで直接測定した値です。学習の実測は 第2部にあります。 作成ルール: 実測データが基本であり、測定できなかったものは「未測定」、確認できなかったものは「未確認」と表示します。

目次

1. オプティマイザは「どれだけ、どちらへ」を決定する

学習の1ステップは、勾配 $G$ を見て重み $W$ をわずかに動かす作業です。最も単純なSGDは W ← W − η·G そのため勾配が大きい方は大きく、小さい方は小さく動きます。Adamはここに座標ごとに「これまでこの座標の勾配がどれくらい大きかったか」を記憶しておき、その大きさで割ります。そのため座標ごとの歩みが均等になります。LLMの事前学習が数年間にわたりAdamW(Adam + 分離型 weight decay)を基本として使ってきた理由は、この安定性です。

しかし、Adamは重みを 数値のまとまりとして見ます。Transformerの重みの大部分は行列であり、行列には座標ではなく 方向という構造があります。Muonはここから出発します。

2. 行列として見るとアップデートがいくつかの方向に偏っている

どのような行列でも特異値分解(SVD)で G = U·Σ·Vᵀ のように表すことができます。U·Vの列は互いに直交する方向であり、Σの対角成分(特異値)はその方向の強さです。言語モデルの勾配・モメンタム行列は特異値が少数の成分に偏っている場合が多いことが知られています(Keller JordanのMuonに関する記事では「ほぼ低ランク(low rank)」と表現されています)。すると、SGDであれAdamであれ実際のアップデートは少数の強い方向だけに動くことになり、残りの方向はほとんど学習されません。

Muonの処方はシンプルです。 特異値をすべて1に変えた U·Vᵀ をアップデートとして用いる 方向はそのままにして強さを均等にするのです。 U·Vᵀ は元の行列に最も近い「直交行列」であるため、この演算は直交化と呼ばれます。

3. SVDなしでの直交化:Newton-Schulz法を5回

ステップごとに大きな行列に対してSVDを行うのは低速です。Muonは行列積のみによって同じ処理を近似します。Xをフロベニウスノルムで割って特異値を1以下にした後、次の反復を5回行います。

def newton_schulz5(G, steps=5, eps=1e-7):
    a, b, c = 3.4445, -4.7750, 2.0315      # Keller Jordan 이 고른 계수
    X = G.bfloat16()
    if G.size(0) > G.size(1):               # 세로로 긴 행렬은 눕혀서 계산
        X = X.T
    X = X / (X.norm() + eps)                # 특잇값을 1 이하로
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X                   # 특잇값 s 가 a·s + b·s³ + c·s⁵ 로 바뀐다
    if G.size(0) > G.size(1):
        X = X.T
    return X

1回の反復は特異値ごとに多項式 a·s + b·s³ + c·s⁵ を適用します(方向U·Vはそのまま)。係数が小さいsを素早く大きくするように選ばれているため、0近傍の小さな特異値も5回繰り返せば1近傍まで上がります。その代わり、正確に1には収束しません。係数が「素早く大きくすること」のために精度を少し妥協した値だからです。

実際に測定してみました。512×2048の行列を作り、特異値が1、1/2、1/3、…と減少するようにし(勾配のように少数の方向に偏った形)、上記の関数をbf16で実行しながら反復ごとに特異値を測定しました(コード:ソースデータの節)。

Newton–Schulz 반복이 특잇값을 1 근처로 모은다 512×2048 합성 행렬(특잇값 1/i), bf16, 실측 0 0.25 0.5 0.75 1 1.25 1 4 16 64 256 512 특잇값 순위(로그 눈금) 특잇값 NS 5회 NS 3회 NS 1회 입력(정규화)
図:Newton-Schulz反復に応じた特異値
反復回数 最大の特異値 最小の特異値 中央値 上位10方向のエネルギーストック
0(正規化のみ) 0.776 0.0024 0.0048 93.7%
1 1.075 0.0083 0.017 84.8%
3 1.202 0.098 0.196 19.7%
5 1.205 0.679 1.112 2.5%
正確な U·Vᵀ 1.000 1.000 1.000 1.95%

入力では512方向中上位10個がエネルギーの93.7%を占めていました。5回繰り返した後は2.5%になり、すべての方向がほぼ同じ強さ(0.68–1.21)になりました。正確な U·Vᵀ(1.95%)にかなり近いです。コストは行列積15回であり、bf16で動作するためGPU上で安価であるというのがMuonの実用的な核心です。実際の学習で測定した負荷については第2部で扱います(スループット5.4%減、バッチを4倍に拡大すると1.5%減)。

4. 全体アルゴリズムと「どこに」使うのか

Keller Jordanの公開実装(ブログがリンクした muon.py)に従うと、1ステップは次のようになります。

  1. モメンタム: M ← β·M + (1−β)·G, β=0.95。更新にはネステロフ形式の (1−β)·G + β·M を使用します。
  2. 直交化: O = NS5(위의 값).
  3. スケール補正: O ← O · sqrt(max(1, 행/열))。この行はブログ本文にはなく、コードにのみ存在します。
  4. 更新: W ← W − η·O(デフォルト η=0.02)。

重要なのは 適用範囲です。Muonは隠れ層の2次元行列(アテンションのQ·K·V·O、MLPの2つの行列)にのみ使用します。埋め込み、出力層(lm_head)、正規化層の利得、バイアスはAdamWのままにします。埋め込みは一度に数行のみ勾配を受け取るテーブル(ルックアップ)であるため「行列の方向」という図式が当てはまらず、1次元パラメータには直交化すべき行列が存在しないためです。Q·K·Vを1つの行列にまとめた実装の場合は、3つに分けて個別に直交化することを推奨します。

スケール補正の規約が論文ごとに異なる

直交行列は、要素1つの大きさ(RMS)が「行列の形状」によって決まります。そのため、「更新をどれだけ大きくするか」を決定する行が論文ごとに異なり、同じ「学習率0.02」であっても意味が変わってきます。

規約 補正 使用箇所 学習率の感覚
Keller Jordan sqrt(max(1, 行/列)) オリジナル実装、modded-nanogpt 0.02前後
Moonlight 0.2·sqrt(max(行, 列)) Moonlight, Essential AI AdamWと同じ値(1e-3前後)
次元比 sqrt(行/列) Scion, MUON+, Spectral Allocation 0.005–0.06

Moonlight(arXiv 2502.16982)の 0.2·sqrt(max) は、更新のRMSをAdamWの典型的な値である0.2に合わせ、AdamW用に調整された学習率とweight decayをそのまま流用するための仕組みです。512×512の行列であれば、Keller規約の0.02はMoonlight規約の約4.4e-3に相当します(計算)。論文の学習率を移植する際に最も間違いやすいポイントです。

5. 2025–2026年の動向

Moonlight: スケールを拡大するためのweight decayとスケール補正

Moonlightは、Muonを大規模モデルに適用するには2つの要素が必要であることを示しました。分離型weight decay(W ← W − η(O + λW))を追加しないと、重みと出力の規模が増え続け、行列の形状ごとに更新サイズが異なる問題を 0.2·sqrt(max) で調整する必要があるという点です。計算最適条件においてAdamWの約52%のFLOPsで同等の性能に到達したと報告されています(図1a、表3)。

大バッチにおけるデータ効率の維持(2505.02222)

Essential AIの「Practical Efficiency of Muon for Pretraining」では、バッチサイズを拡大した場合を取り上げています。バッチを大きくするとステップ数が減り高速化しますが、ある地点(臨界バッチサイズ)を超えると、同じ損失に到達するために必要なトークンが増加します。この論文では、同じ損失に到達するまでにAdamWとMuonが消費したトークンの比率をバッチサイズ別にプロットし、Muonはトークンを10–15%少なく消費し、その優位性がバッチが大きくなっても維持、あるいは拡大すると述べています(図2・3)。原文の表現は "retaining data efficiency at large batch sizes, far beyond the so-called critical batch size" です。「臨界バッチサイズがより大きい」という点推定に頼らず、トークンの比率が減衰しないという点に根拠を置いています。第2部でこの主張を小規模に再検証します。

MUON+: 直交化の後の追加正規化(2602.21545)

直交化を行えば行・列の大きさも均等になりそうですが、実際に使用する5回の反復では、かえって行・列のノルムの不均衡が大きくなる可能性があるという指摘です(セクション2.1、図2·3)。5回の反復は正確な U·Vᵀ ではないという点(上の表で特異値が0.68–1.21に分散していたこと)と根を同じくする問題です。著者らは、この不均衡が安定して使用できるステップの最大サイズを縮小させると分析しました。MUON+は直交化結果の列ノルムで割り、続いて行ノルムで割るステップを1つ追加します。状態の追加はなく、コード数行の変更で済みます。60Mから7BまでのGPTおよびLLaMAのいずれにおいてもMuonより検証PPLが低く(LLaMA-60M 25.75 → 25.25、表2)、同じ損失に到達するまでのステップ数を22.4–37.1%削減できたと報告されています(表5)。我々の実験モデルとほぼ同等規模であるLLaMA-60Mでは、改善効果が約0.5 PPLと小さかった点もあわせて記憶に留めておく価値があります。

Spectral Allocation: 方向ごとに異なるステップ(2608.25990)

Muonはすべての方向に同じステップを与えます。この論文では、学習中のモメンタム行列の特異方向ごとに「未知のバッチにおいて損失を最も大きく減少させるステップ」を直接測定しました。最も上の方向は安定性の境界(Edge of Stability)に位置しており、残りの方向は数倍大きなステップにも耐えられるという結果が得られました。そこで、上の方向にはMuonのサイズを維持し、残りの方向をγ倍に拡大した「SAMuon」を提案し、同じ損失に対してトークンを13–24%削減できたとしています(表7)。この観点から、「Muon > Adam > SGD」という順序も説明されます。SGDは強い方向にステップを集中させ、Adamはそれを座標単位でのみ緩和し、Muonは方向単位で均等化します。この論文の表5において、AdamWとMuonの損失の差は、バッチ1024シーケンスで0.048、4096シーケンスで0.110となり、バッチが大きくなるほど広がりました。

大規模な比較: SOAP, Muon, and Beyond(2607.20548)

NVIDIAのこの論文では、8Bのdense、30B・72BのMoEモデルを1~3Tトークンで学習させ、AdamW、Muon、SOAPを比較しました。MuonとSOAPはAdamWを一貫して上回り、グローバルバッチを100Mトークンまで拡大しても耐えましたが、AdamWは悪化しました(要旨、5.3節)。SOAPの大バッチでの不安定さを修正したKL-SOAPがMuonを「一貫してわずかに」上回りました(図7・8)。ただし、ここでの公平な比較とは、オプティマイザごとに学習率を個別に探索したのではなく、更新のRMSを合わせてスケーリングしたものであると著者らは述べています。8Bハイブリッドモデルの数学的平均では、むしろAdamWが高かったという点(表5、84.04対80.79)も原文に記載されています。

小規模モデルのレシピもMuon系へ

2026年に発表された小規模モデルの事前学習レシピの2本も、オプティマイザとしてMuon系を使用しています。IMU-1(2602.02522、430M)はNorMuonにCautious Weight DecayとμPを、Puro-2B(2608.27370)はMuonHを使用しました。どちらも埋め込み・出力層・1次元パラメータにはAdamWを使用しています。この流れについては 第4回で改めて取り上げます。

6. まとめ、そして確認したいこと

論文では数億~千数百万(※原文ママ、数千億)パラメータにおいてこの結論が導き出されました。それでは、8GBのノートPC用GPU1枚で動かせる3,400万パラメータのモデル、1億トークンでも同じ現象が起きるでしょうか? 学習率をオプティマイザごとに個別に探索して公平に合わせた場合、どれほどの差が残るでしょうか? 大バッチで差が開くでしょうか? 第2回で実際に測定します。

原資料および参考文献

ログインするといいねできます

コメント

まだコメントはありません。

ログインするとコメントできます