Newton-Schulz直交化、FLOPsを3.4倍削減
- Newton-Schulz直交化を単一演算に統合、FLOPsを3.4倍削減
- 直交化誤差を96倍減少 (3.838 → 0.040)
- 学習可能な係数で収束安定性を確保
何が起こったのか？
中国の研究チームがNewton-Schulz直交化の新しい統合フレームワークUNSOを発表した。[arXiv] 既存のNS反復を単一の多項式演算に置き換えた。核心は「意味のない項を削除し、学習可能な係数を導入」することだ。
128×512行列基準でFLOPsが2.533×10^8から8.831×10^7に低下した。[arXiv]
なぜ重要なのか？
MuonはAdamWを代替するオプティマイザーとして注目されている。GPT-2 XLを$175で学習させ、Kimi K2にも使われている。[Keller Jordan] しかしNewton-Schulz反復がボトルネックだ。
UNSOがこのボトルネックを解消した。反復を減らすのではなく、完全になくす。直交化誤差も96倍減少した。誤差の累積が学習不安定の原因だが、これを解決した。
今後どうなるか？
PyTorch 2.10にMuonが公式に含まれた。[PyTorch] NVIDIA NeMoもサポートする。[NVIDIA] UNSOが急速に吸収される可能性が高い。
よくある質問 (FAQ)
Q: UNSOを今すぐ使えるか？
A: はい。著者らがGitHubにコードを公開した。PyTorch環境で既存のMuonを代替する形で使用可能だ。ただし、プロダクション適用前に独自のベンチマークが推奨される。
Q: AdamWの代わりにMuonを使うべきか？
A: 状況による。Muonはhidden layerにのみ適用する。埋め込みや出力層にはAdamWが必要だ。2つのオプティマイザーを組み合わせるのが標準だ。
Q: 実際の学習時間はどれくらい短縮されるか？
A: 直交化段階で3.4倍速い。全体の学習時間はモデルのサイズによって異なる。直交化の比重が大きいほどUNSOの利点が大きくなる。
この記事が役に立ったら、AI Digesterを購読してください。
参考資料
- UNSO: Unified Newton Schulz Orthogonalization – arXiv (2026-02-04)
- Muon: An optimizer for hidden layers – Keller Jordan Blog (2025-01-15)
- Muon Optimizer – PyTorch Documentation (2026-01-20)