LLMの学習を高速化したい？ Muonを使ってみませんか？

こんにちは、PredNextの徳永です。

今日はみんな大好きオプティマイザー回です。とはいえ、オプティマイザー全般はちょっともう追いきれないので、Muon関連に絞っていくつか集めました。

Muonとは

まず、Muonについて説明しておきましょう。Muonはニューラルネットワークの最適化に使われるオプティマイザーの一種で、勾配の移動平均（いわゆるmomentumですね）をNewton-Schulz iterationという手法で近似的に直交化し、その結果を使ってパラメーターを更新します。Newton-Schulzによる直交化以外の部分はMomentum SGDと同じです。

LLMの学習においては、Adamやその改良版であるRAdam/AdamWなどのオプティマイザーが主に使われてきました。毎年新しいオプティマイザーがたくさん提案されてきましたが、そのほとんどは実用されず、定番の地位は2014年に提案されたAdamやその微改良手法がずっと占めてきたのです。

星の数ほどのオプティマイザーが論文とともに消えてゆく中、Muonの提案者が取った手段は論文ではなく、nanoGPT-speedrunというコンペティションでのMuonの利用でした。こちらで圧倒的な成果を上げたことで注目を集めたからか、今年、いくつかの大型LLMの学習においてMuonが採用される事例が出始めました。私の知る限りでは、Kimi K2、GLM-4.5、INTELLECT-3の3つのモデルがMuonで学習されています（INTELLECT-3はGLM-4.5をpretrained modelとして利用しているので、当たり前といえば当たり前ですが）。

• Kimi K2: Open Agentic Intelligence
• GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models
• INTELLECT-3: Technical Report

ここまで利用例が積み重なると、Muonが優れたオプティマイザーであることは認めざるを得ないでしょう。収束までのステップ数が2/3程度になること、オプティマイザーのメモリ消費量がAdam/AdamWと比べてやや少なく済むことなど、メリットを考えると、今後もMuonの採用は進むはずです。また、Muonを使った事例として、ResNetやViTを学習したものも増えてきています。言語データだけではなく、画像データでも使える汎用性があるようです。

余談：TrinityもMuonを使って学習されているそうです。現時点で公開されているTrinity Miniは26Bモデルなので、上記3つと比べられるほど大きくありませんが、2026年1月に420BモデルであるTrinity Largeがリリースされる予定になっています。

どのMuonを使えばいいの？

Muonの実装はいくつか存在します。KellerJordan/Muonやmodded-nanogpt内部の実装、PyTorchの実装、Optaxの実装など。普通に使う分には、PyTorchの実装もしくはOptaxの実装を使うのが一番無難です。modded-nanogptの実装はweight decayが入っていない、scaling factorの計算が$\sqrt{\max(1, A/B)}$しか選べないなど、大きなモデルを学習する場合にちょっとした落とし穴があります。PyTorchの場合はadjust_lr_fn="match_rms_adamw"にして使うのが、初手としてはおすすめです。

Muonの落とし穴

Muon is Scalable for LLM Trainingの3.5.1 Ablation Studies on the Interchangeability of Pretrain and SFT Optimizersにあるように、Muonで学習したモデルをAdamWでfinetuningしたり、その逆をしたりすると、あまりうまくいかないことが知られています。継続事前学習に用いる場合などは、事前学習で用いられたオプティマイザーと同じものを用いるのが無難ですので、気をつけましょう。

Muonに関する進展

今年出てきたMuon関連の論文のうち、興味深いなと思ったものを並べました。時系列順に並べてあります。

The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm

Muonの直交化の部分にNewton-SchulzではなくPolar Expressという手法を使うことで、収束を早めることができました。筆者が7月に論文を読んだ段階ではnanoGPT-speedrunに採用されておらず、なぜ使われていないのか疑問でしたが、さっき公式リポジトリを確認してみたところ、9/29にめでたく採用されていました。

ところで、前に記事を書いた頃、nanoGPT-speedrunは3分をやっと切ったばかりでしたが、今はもう2.248分まで記録が縮まっています。この分だと来年には2分を切りそうで、進化のスピードが速すぎて、感情の方が追いつきません。

MuLoCo: Muon is a practical inner optimizer for DiLoCo

DiLoCoとMuonの共演です。こんなのもうお茶碗で3膳くらいはいけますね。なんとさらにみんな大好きな量子化までついてきます。これで今日はお茶碗4膳いける。

DiLoCoは分散学習のための手法で、外側オプティマイザーと内側オプティマイザーの2つを使って学習を行います。手法の概要については以前に個人ブログに書いたのでそちらを参照してください。内側オプティマイザーにはAdamWが用いられることが多いのですが、これをMuonに変えると、大方の予想通り、収束が早まります。それだけではなく、外側オプティマイザーに共有する勾配情報を2bitに量子化しても、内側オプティマイザーがMuonの場合は、損失の値がほとんど変化せず、通信量を大幅に削減できることがわかりました。ちなみに、AdamWだとかなり性能が劣化してしまいます。なぜMuonだと勾配の量子化が可能なのかさっぱりわかりませんが、うまくいくということで、飲み込んでしまうしかありません。

大規模で高速なネットワークというのは非常にお金がかかるので、MuLoCoのような手法は重要性を増していくことでしょう。Prime Intellectのように、すでにDiLoCoを実用しているところもあります。

Accelerating Newton-Schulz Iteration for Orthogonalization via Chebyshev-type Polynomials

Muonで使われているNewton-Schulzは高速化のために途中で打ち切られているため、近似的な結果しか得られません。近似精度を上げるためにRemezのアルゴリズムを使って最適な係数を求めてみたのがこの研究です。この「最適」は、最悪ケースでの近似誤差が最も小さくなる、という意味です。iteration回数が4の場合は提案手法の方が明らかに収束が早まるのですが、iteration回数が5になると、違いがほとんどなくなってしまいます。ただ、直交化の精度を上げると学習率を上げられる、という話もあるので、ハイパーパラメータ探索をしたら、違う結果が得られるかもしれません。

ちなみに、Remezのアルゴリズムは、東京科学大の岡崎先生のブログ記事で知りました。昔は、いろいろな人がブログ記事を書いていて、そういうところから仕入れた知識というのはこうして今でも役に立っています。ブログサービス自体のサ終などの理由で読めない記事が年々増えていくことをさみしく感じます。

NorMuon: Making Muon more efficient and scalable

Muonの収束を早めるために2次のモーメントを導入します。ただし、自明な実装方法ではメモリ消費量が大きく増えてしまうため、2次モーメントは列単位で値を保持します。実験では、Muonに対し、1.1Bモデルで11%程度、5.4Bモデルで5%程度、収束が早くなったそうです。現在nanoGPT-speedrunに採用されているのはこの手法とPolar Expressの組み合わせです。

Cautious Weight Decay

本論文はMuonそのものについての研究ではありませんが、Muonを使った実験が載っています。Cautious Weight Decayと似たアイデアを用いた論文として、Cautious Optimizers: Improving Training with One Line of Codeがあります。Cautious Optimizersでは、gradientと更新方向の符号が一致する場合にのみ、weightを更新します。こちらの解説は以前に個人ブログに書きました。また、Muonに適用してもnanoGPT-speedrunに対しては効果がなかったことも書きました。

Cautious Weight Decayは、重みの符号と更新方向が一致する場合にのみWeight Decayを適用します。論文中にMuonに対する適用結果が記載されており、ViT-S/16, ResNet-50, ViT-B/16を使った画像分類の実験や、OLMo 1Bを使ったテキスト系タスクの実験において、収束が一貫して早まり、また、最終的な精度も少しよくなることが報告されています。nanoGPT-speedrunにもすでに採用されています。

Beyond the Ideal: Analyzing the Inexact Muon Update

既存のMuonの理論解析は直交化部分にSVDを仮定していましたが、実際のMuonはNewton-Schulzを使っているので、そこで理論と現実にギャップが生まれる可能性があります。この論文では、線形最小化オラクルという解析手法を用いて、より現実に近い解析を行っています。その結果、オラクルが不正確であればあるほど（Newton-Schulzでの直交化が理想と遠ければ遠いほど）最適なステップ幅、つまり、最適な学習率が小さい、ということを示しました。言い換えると、しっかりと直交化すればするほど、学習率を上げることができるという主張です。学習率は大きいほうが収束が早まり、最終的な損失も小さくなることが知られているので、大きな学習率を使えるのはうれしいことです。nanoGPT-speedrunのコードを使って実験してみたところ、Newton-Schulzの繰り返し回数によって最適な学習率が変化し、解析結果と実験結果が同じ傾向を示したそうです。

MuonAll: Muon Variant for Efficient Finetuning of Large Language Models

Muonは行列に対して使うものなので、1DパラメーターはAdamWで最適化するのが慣例ですが、1Dパラメーターを敢えて対角行列化してMuonで最適化してみたらどうなるだろう、という研究です。いくつかのネットワークに対してSFTしていますが、AdamWが勝ったりMuonが勝ったりMuonAllが勝ったりで、正直、傾向が見えません。Adam/AdamWで学習されたネットワークに対して継ぎ足しでMuonで学習してもうまくいかない、というのは上述の通りなので、実験設定があまりよくないなと思います。しかし、Muonのみで最適化ができる、というのはおもしろいアイデアです。pretrainingに使った場合にどうなるかも見てみたいところです。

Turbo-Muon: Accelerating Orthogonality-Based Optimization with Pre-Conditioning

Muonは直交化の際に行列積を10〜15回程度実行するので、1ステップあたりにかかる時間がAdamWよりも伸びてしまうことが知られています。Turbo-Muonでは、Newton-Schulzのループ実行回数を減らすことでMuonを高速化します。単純に行列積の回数を減らすだけだと直交化が不十分になってしまうので、前処理に変更を加えることで直交化の誤差を減らします。

直交化する対象となる行列を$X_0$とすると、Muonでは、スケーリングファクターとして$\frac{1}{\|X_0\|}$を使います。$\|X_0\|$は$X_0$のフロベニウスノルム、つまり、行列の各要素の二乗和の平方根です。

Turbo-Muonでは、$A_0 = X_{0}^{T} X_0$ として、$\frac{1}{\sqrt{\|A_0\|}}$ をスケーリングファクターとして用います。

違いはほぼこれだけ（実装上は、不要な再計算を避けるためにちょっとした工夫があります）ですが、Turbo-Muonでは直交化の誤差が大幅に減少し、その分、ループの回数を減らせるので、結果としてMuonよりも高速になります。1.3BモデルをnanoGPT-speedrunの設定で学習したところ、Muonに対して8〜10%程度の高速化が見られたそうです。1.3Bは通常のnanoGPT-speedrunの設定よりは大きなモデルです。行列が大きくないと効果が薄いので、大きめのモデルを使ったのでしょう。

おわりに

Muonに関連する研究をいくつか取り上げてみました。個人的には、直交化の精度を上げることで学習率を上げられるというのはおもしろく感じました。おもしろいだけでなく、今後、役に立つ可能性を感じています。

オプティマイザーには、まだまだ改善の余地が残っているはずです。The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newtonによると、Gauss-Newton法を使って最適化した場合、ステップ数で計算すると、Muonの1/16程度の回数で同程度の損失関数の値が得られるそうです。Gauss-Newton法そのものの計算が重たいので、1ステップあたりの実行にかかる時間が大幅に増え、実際に速くはならないわけですが、よりよい近似的手法を発見できればより速いオプティマイザーが作れる余地がありそうだ、という期待が持てます。2026年もオプティマイザーの進展に期待しましょう。

お仕事募集中です

2025年12月現在、スケジュールにまだ若干の空きがあります。機械学習が関わる案件ならなんでもお待ちしております。ご興味のある方はお問い合わせフォームからご連絡ください。