グラム・ニュートン＝シュルツ法：Muon オプティマイザーの安定化と高速化

1. モデル品質の維持とオプティマイザーの高速化

Muon の特性と課題

• 最先端モデルでの採用： Kimi K2 Thinking や GLM-5.2 などの先進的な言語モデル訓練において、Muon は選択すべきオプティマイザーとして注目されています。
• AdamW との違い： AdamW に対し、到達する損失値まで必要なステップ数は少ないですが、各ステップのコストは高いという特徴があります。
  • このオーバーヘッドの要因は、ニュートン＝シュルツ直交化手順によるものです（旧式オプティマイザーに存在しない $O(n^3)$ 計算）。

ベンチマーク結果 (B300) Muon の優れた最適化品質により、高コストなオプティマイザーステップを正当化しています。

スケーリングに伴うオーバーヘッドの増大

• 従来手法（SGD, AdamW）： $O(mn)$ の時間（要素ごとの更新）。
  • 勾配行列 $\mathbf{G} \in \mathbb{R}^{n \times m}$ を用いた重み更新。
• 現代的手法（Muon など）： $O(mn^2)$ の時間（行列積を使用）。
  • ニュートン＝シュルツ法などの高階プレコンディショニングによる直交化手順。
• 実際の影響： トレーニング設定によりますが、エンドツーエンドの壁時計時間の 2% ～ 17% を占めることがあります。

標準的なニュートン＝シュルツ法の限界点

• 冗長な計算： $n \times m$ の行列を 10 乗する処理が多用されており、無駄な FLOP が消費されます（$2mn^2$）。
• 非正方行列の扱い難さ： 多くの重み行列は長方形 ($m \gg n$) です。MoE アーキテクチャではさらに顕著で、矩形行列積のコストが支配的です。
• 対称行列の利用不足： 中間行列に現れる対称行列の計算利点を利用せず、約半分の仕事が冗長です。
• 非最適化されたカーネル： バッチ行列乗算には cuBLAS を使用していますが、Hopper GPU アーキテクチャに対して完全に最適化されていません。

---

2. 本稿の寄与：グラム・ニュートン＝シュルツ法の導入

核心アプローチ

• アイデア： 長方形入力行列 $\mathbf{X}$ を直接反復させるのではなく、小さな正方対称グラム行列 $\mathbf{XX^\top}$ を反復させます。
• メリット： FLOP コスト削減と、対称 GEMM カーネルの効率的利用が可能になります。

主要な貢献ポイント

1. ナイーブなグラム・ニュートン＝シュルツ法：
   • 標準的な手法を数学的に同等だが、$n \times n$ 行列空間で動作するように書き換えます。
   • 専用対称行列乗算ルーチンを採用し、各反復を高速化（標準より約 68% 削減）。
2. 安定化された手法：
   • 半精度計算における数値的不安定性（見かけ上の負の固有値）を特定。
   • **「リスタート (restarting)」**戦略を採用し、中間で行列を再構築して安定性を確保します。
3. カスタムカーネルの実装：
   • CuTeDSL を用いた対称行列乗算専用カーネルを実装。
   • Hopper および Blackwell アーキテクチャ向けに最適化されています。
4. グラム・Muon (GramMuon) の発表：
   • Muon のニュートン＝シュルツルーチンを置き換える新しいオプティマイザー。
   • 直交化ステップの実行時間を 40〜50% 削減しました。

「無料のランチ」効果： 安定しており、検証パプラクスは標準版と比較して ±0.01 の範囲内で維持されています。

オープンソース実装

• ニュートン＝シュルツルーチンのドロップイン置換（数学的同等・数値的安定・最大 2 倍高速）。
• Hopper/Blackwell 向け対称行列乗算カーネルの公開。

---

3. Muon のアルゴリズム解説 recap

モデル原理

Muon は、スペクトルノルムに関する最急降下法として記述されます。更新規準は以下の通りです。

$$ \begin{align*} \mathbf{M}k &= \mu \mathbf{M}{k-1} + \mathbf{G}k \ \mathbf{W}{k+1} &= \mathbf{W}_k - \eta \operatorname{polar}(\mathbf{M}_k) \end{align*} $$

ここで $\operatorname{polar}(\mathbf{X})$ は極分解であり、その計算は Newton-Schulz 法で近似されます。

アルゴリズム 1: 標準的なニュートン＝シュルツ法

入力: $\mathbf{X} \in \mathbb{R}^{n \times m}$ ($n \leq m$)

// 特異値を [0, 1] に正規化。epsilon = 1e-7
X <- X / (||X||_F + epsilon) 
X <- bfloat16(X)                // 速度のため半精度へキャスト

If m < n:
    X <- X^T                    // X X^T をより安くするためのトリック

For t = 1, ..., 5:              // p_t(X) を適用
    A <- X X^T                  // n x n の対称行列
    B <- b_t * A + c_t * A^2    // 対結合
    X <- a_t * X + B X         // 長方形行列を更新

If m < n:
    X <- X^T                    // トリックを元に戻す

Return X

数学的解釈 (特異値の多項式変換)

• $\mathbf{X}_0$ の SVD を $\mathbf{U} \mathbf{\Sigma} \mathbf{V}^\top$ とすると、反復は対角行列 $\mathbf{\Sigma}$ に対して多項式 $p_t(x)$ を適用することになります。
• 最終的に特異ベクトル $\mathbf{U}, \mathbf{V}$ は維持され、特異値のみが変換されます。これにより $\mathbf{X} \to \operatorname{polar}(\mathbf{X})$ に収束します。

---

4. 実行時間分析と最適化の余地

標準法の実行時間 ($T=5$)

• 1 反復あたりの FLOP: $2(2\alpha + 1)n^3$ （$\alpha = m/n$）
  • $\mathbf{X} \mathbf{X}^\top$: $2mn^2$
  • $\mathbf{A}^2$: $2n^3$
  • $\mathbf{B} \mathbf{X}$: $2mn^2$

限界点と対策

1. 対称行列乗算の活用：
   • 標準実装では対称性が利用されず、対称 GEMM カーネルを使用することでコストを削減可能。
2. 長方形性の依存（$\alpha$）：
   • MoE など長大な重みを持つモデルほど効果が大きく、現代のアーキテクチャで特に重要です。

最適化後の効果 (Hopper)

• 対称 GEMM カーネル使用により総コスト削減。
• Ping Pong スケジューリングによりエピローグ加算を隠蔽し、顕著な高速化を実現。

---

5. グラム・ニュートン＝シュルツ法の提案手法

戦略：高価な長方形乗算の削減

• 目標： $n \times m$ の反復ではなく、小さな正方対称グラム行列 $\mathbf{XX^\top}$ を反復させる。
• 手順:
  1. グラム行列 $\mathbf{X} \mathbf{X}^\top$ を計算 ($n \times n$)。
  2. この小さない平方行列で反復法を実行して $(\mathbf{XX^\top})^{-1/2}$ を近似。
  3. 最後に $\mathbf{Q}_T \mathbf{X}$ を計算して直交化を適用。

アルゴリズム 2: ナイーブなグラムのニュートン＝シュルツ法

入力: $\mathbf{X} \in \mathbb{R}^{n \times m}$ ($n \leq m$)

// 特異値を [0, 1] に正規化
X <- X / (||X||_F + epsilon) 
R_0 <- X X^T                 // 正方対称行列のみで操作開始
Q_0 <- I                     // 単位行列

For t = 1, ..., 5:
    // h_t(R_{t-1}) を適用 (全小さない平方行列空間内)
    Z_t <- a_t * I + b_t * R_{t-1} + c_t * R_{t-1}^2 
    
    Q_t <- Q_{t-1} * Z_t      // 小さな対称乗算
    R_t <- Z_t * R_{t-1} * Z_t // 小さな対称乗算

Return Q_5 * X               // 最終的に長方形行列への適用

性能比較

• ナイーブ法: $(4T + 3\alpha - 3)n^3$ の FLOP。
• 標準手法との差:
  • 対称 GEMM なし比較：68% の節約。
  • 対称 GEMM あり比較（最適化後）：55% の節約。

---

6. ナイーブなグラムのニュートン＝シュルツ法の不安定性と安定化

問題点：数値的不安定性

有限精度（bfloat16）下では、標準的な手法とは異なり以下の問題が発生します。

図 2: Llama-430M のトレーニング結果 ナイーブなグラムのニュートン＝シュルツ法では損失スパイクが発生し、最終的に出力が Inf に発散します。

発散の原因

1. 見かけ上の負の固有値：
   • 浮動小数点誤差によりグラム行列に負の固有値が生じる。
   • 更新規準（例：$r_t = r_{t-1} z_t^2$）により、負の絶対値が指数関数的に増幅される。
2. 固有ベクトルのドリフト：
   • 有限精度で中間行列の固有ベクトルが元の左特異ベクトルからずれる。
   • これによりスペクトルノルムが発散する。

図 5 & 8: BFloat16 での固有値の進化と、ドリフトによる発散の仕組み。

解決策：リスタート (Restarting) 戦略

アルゴリズムの途中で（例：反復 2 回後）グラム行列を再構築し、負の固有値やドリフトをリセットします。

アルゴリズム 3: 安定化されたグラムのニュートン＝シュルツ法

// ... プログレッシング ...
R_0 <- X X^T
Q_0 <- I

For t = 1, ..., 5:
    If t = 3:                      // リスタートポイント（適応的に決定可能）
        X <- Q_2 * X               // 現在の近似を適用済みとして保存
        R_2 <- X X^T               // 新しいグラム行列で再初期化
        Q_2 <- I                   // 単位行列にリセット
    
    Z_t <- b_t * R_{t-1} + c_t * R_{t-1}^2
    Q_t <- Q_{t-1} * Z_t + a_t * Q_{t-1}   // 安定性のために恒等加算を明示
    (RZ)_t <- R_{t-1} * Z_t + a_t * R_{t-1}
    R_t <- Z_t * (RZ)_t + a_t * (RZ)_t

// Post-processing
X <- Q_4 * X

実装上の注意点

• 精度設定:

  float16

  を

  bfloat16

  より推奨する。グラム行列計算の数値誤差を制御できるため。
• サファティファクター (Safety Factor): 係数を調整して特異値の対応範囲を広げる（例：入力 $/ 1.05$）。
• フューズドクォadratiks: フューズドカーネルで対称 GEMM を実装する場合、$\gamma \mathbf{I}$ の明示的な加算は不要で、暗黙的な処理の方が安定することがある。

---

7. 実装詳細とカーネル最適化戦略

CuTeDSL による対称 GEMM カーネル

CuTeDSL (Quack ライブラリ) を用いて Hopper/Blackwell 向けに最適化した高速カーネルを実装。

• 三角形スケジューラー： 下三角のみをスレッドブロック間で分散し、上三角は転置/コピーで埋めることで負荷バランスを取ります。
• エピローグ最適化： 値の両方の三角への書き込みと、対角要素の冗長性回避（NaN 防止）を制御します。

高速化サマリー

比較対象	ハードウェア	性能向上率	特徴
標準ニュートン＝シュルツ (対称カーネル有)	Hopper	~25%	フューズド演算、Ping Pong スケジューリングによる隠蔽
グラム・ニュートン＝シュルツ	Hopper / Blackwell	最大 50% (長方形重みの場合)	アルゴリズム的構造変化による FLOP削減
Kimi K2 ベンチマーク	Trillion params MoE	2 倍速	パイプライン並列化設定下での実証

実験結果の検証

• モデル品質: 検証パプラクスは標準 Muon と±0.01 の範囲内で維持（「無料のランチ」）。
• スケーリング則: 高次長方形な重み行列を持つモデル（MoE の MLP など）において最も性能が向上します。

---

8. 結論と利用法

メインメッセージ

• グラム・ニュートン＝シュルツ法は、トレーニング品質を維持しつつ、一般的なアーキテクチャにおいてオプティマイザーステップを最大 2 倍まで高速化することを示しました。
• Muon の標準実装へのドロップイン置換として提供可能で、Open Source 実装が公開されています。

調整が必要なハイパーパラメータ

• リスタートする反復数： 一般的に 2 回目でリスタートするのが推奨（Polar Express 係数の場合）。
  • オートチューニングスクリプトにより、多項式列に最適化されたリスタートポイントも提案可能です。

引用情報

@misc{GramNewtonSchulz,
  title     = {Gram Newton-Schulz},
  author    = {Jack Zhang and Noah Amsel and Berlin Chen and Tri Dao},
  year      = {2026},
  url       = {https://dao-ailab.github.io/blog/2026/gram-newton-schulz/}
}

参考文献

1. Keller et al., "Muon: An optimizer for hidden layers in neural networks," 2024.
2. Bernstein, "Deriving Muon," 2025.
3. Less Wright and Hoque, "CUTLASS Ping-Pong GEMM Kernel," PyTorch Blog, 2024.
4. Amsel et al., "The Polar Express," ICLR, 2026.
5. Grishina et al., "Accelerating Newton-Schulz Iteration...", arXiv:2506.10935, 2025.
6. Ahn et al., "Dion: Distributed Orthonormalized Updates," arXiv:2504.05295, 2025.
7. Liu et al., "Muon is Scalable for LLM Training," arXiv:2502.16982, 2025.
8. Kimi Team, "Kimi K2," arXiv:2507.20534, 2026.
9. Grattafiori et al., "The Llama 3 Herd of Models," arXiv:2407.21783, 2024.
10. GLM-5 Team et al., "GLM-5," arXiv:2602.15763, 2026.
11. Yang et al., "Qwen3 Technical Report," arXiv:2505.09388, 2025.
12. OpenAI et al., "gpt-oss...", arXiv:2508.10925, 2025.
13. DeepSeek-AI et al., "DeepSeek-V3 Technical Report," arXiv:2412.19437, 2025.
14. Zhong et al., "DPO Meets PPO," arXiv:2404.18922, 2025.
15. Pethick et al., "Training Deep Learning Models with Norm-Constrained LMOs," arXiv:2502.07529, 2025.
16. Vyas et al., "SOAP...", arXiv:2409.11321, 2025.
17. Frans et al., "A Stable Whitening Optimizer...", arXiv:2506.07254, 2025.
18. Gupta et al., "Shampoo," ICML, 2018.
19. Gemma Team et al., "Gemma 3 Technical Report," arXiv:2503.19786, 2025.
20. Newhouse et al., "Faster Symmetric Matrix Multiplication with ThunderKittens."
21. Lin, "Flash-Muon," GitHub, 2025.
22. Yang et al., "PRISM...", arXiv:2601.22137, 2026.
23. Merrill, "Critical Batch Size Revisited," AllenAI Blog, 2025.

補遺：トレーニング時間分析ケーススタディ

• ケース 1 (Kimi K2): 楽観的・パイプライン並列化設定下で、標準ニュートン＝シュルツ法は事前トレーニングの約 1.7% - 2.2% を占める。
• ケース 2 (Llama3-70B SFT): 小バッチサイズ・4-GPU クラスターでは、約 17% を占める。

結論: ニュートン＝シュルツステップの最適化は、特に低精度体制や小バッチ設定においてボトルネック解消に大きく寄与する。