クレッシーがAofei Temple、QbitAIからレポート | WeChat公式アカウントQbitAI

AdamW よりも 48% 低い計算能力要件で、OpenAI エンジニアによって提案されたミューオン トレーニング最適化アルゴリズムが Dark Side of the Moon チームによってさらに進化しました。

研究チームはミューオン法のスケーリング則を発見し、改良を加えて、ミューオンがより大きなモデルにも同様に適用できることを証明した。

パラメータ数が最大 1.5B であるさまざまな Llama アーキテクチャ モデルでは、改良された Muon の計算能力要件は AdamW のわずか 52% です。

チームはまた、改良された最適化アルゴリズムとともにオープンソース化された DeepSeek アーキテクチャに基づく 160 億の MoE モデルをトレーニングしました。

Muon の技術ブログが最初に公開されたときは、主に小規模なモデルとデータセットを対象としていたため、次の 3 つの疑問が未解決のまま残っていました。

• Muonはより大規模なトレーニングに使用できますか?
• Muon は大規模な GPU クラスターで使用できますか?
• Muonは微調整や強化学習にも適用できるのでしょうか？Moonlightチームは実験を通して、これらすべてに「はい」という答えを示しました。

このニュースが報じられると、ミューオンの著者たちは大いに興奮しました。筆頭著者のケラー・ジョーダン氏は、これがミューオンの大規模実験における最初の成功報告であると述べ、チームを祝福しました。

もう一人の貢献者であり、当時ミュオンの大規模実験を担当していたハイパーボリックラボの共同設立者兼CTOであるユチェン・ジン氏も、ミュオン・ダークネスチームの成果はミュオンの勝利であると述べた。

MuonにAdamWの機能を導入

Moonlight チームの取り組みを紹介する前に、まずは Muon がどのような技術であるかを理解しましょう。

これはニューラル ネットワークの隠し層用の 2D パラメータ オプティマイザーであり、主な作者は OpenAI ディープラーニング チームの Keller Jordan です。

この研究結果は昨年12月8日に発表され、ケラー氏は昨年12月にOpenAIに加わった。

Muon の中心的なアイデアは、勾配更新行列を直交化することでパラメータ更新が局所最小値に陥るのを防ぎ、モデルがより多様な特徴表現を学習できるようにすることです。

Muon は 94% の精度で、A100 上の CIFAR-10 のトレーニング時間を 3.3 秒から 2.6 秒に短縮しました。

しかし、当時、ミュオンのチームは小規模なモデルとデータセットでのみ実現可能性を証明しており、より大規模なモデルに適用できるかどうかは不明でした。

現在、Moon チームが行った改良のおかげで、Muon はより大規模なモデルやデータセットにも同様に応用できることが証明されています。

モデル自体に関しては、チームは AdamW のいくつかの機能を取り入れ、具体的には 2 つの側面で Muon に移植しました。

まず、重みの減衰メカニズムが導入され、減衰係数を持つ項が重みの更新式に追加されました。

著者らがこれを行ったのは、Muon を大規模なトレーニングに直接適用すると、モデルの重みとレイヤー出力の大きさが増加し続け、最終的に bf16 の高精度表現範囲を超え、モデルのパフォーマンスが低下することがわかったためです。

8 億パラメータのモデルを 1000 億トークン (最適な計算予算の約 5 倍) にトレーニングする間、チームは AdamW、重み減衰のない Muon、重み減衰のある Muon を比較しました。

結果は、重み減衰を備えた Muon がオーバーフィッティング段階で最良の結果を達成し、重み減衰の必要性を検証したことを示しています。

2 番目の改善点は、異なる形状行列のパラメータの更新の大きさを一定に保ち、AdamW の更新の大きさと一致するように、 Muon のパラメータ更新スケールを調整することです。

Muon の 1 つの特徴は、形状 [A,B] の行列パラメータの場合、その理論的な更新の大きさが sqrt(1/max(A,B)) になることです。

その結果、異なる形状の行列ではパラメータの更新量に大きな違いが生じます。例えば、MLPのような幅の広い行列では更新量が小さすぎるのに対し、各ヘッドを独立した行列として扱うと更新量が大きすぎます。

さらに、この大きさは AdamW と一致しないため、ハイパーパラメータの設定が困難になります。

さまざまな行列パラメータの更新の大きさがAdamWと一致し、一貫性があることを保証するために、著者らはいくつかの改善スキームを試し、最終的に各パラメータの形状に基づいて学習率を直接調整することを選択しました。

0.2 という値は実験的に決定された定数であり、Muon の更新スケールを AdamW に合わせるために使用されます。

Muon をより大規模な訓練に利用するためには、Muon 自体の改良に加えて、分散訓練環境への拡張が必要となります。

Muon では、直交化更新を計算するために完全な勾配行列が必要ですが、既存の分散トレーニング フレームワーク (ZeRO-1、Megatron-LM など) では、オプティマイザーの状態を要素ごとに異なるデバイスに独立して分割できることを前提としているため、Muon を直接サポートすることはできません。

この問題に対処するために、著者らは分散 Muon の並列化戦略を提案しました。

ZeRO-1 の上に 2 つの追加操作が導入されています。

• まず、各データ並列グループ内で勾配集約通信が実行され、散在する勾配スライスが完全なマトリックスに結合されます。
• 次に、集約された勾配行列に基づいて直交化更新が並列に計算され、ローカル パラメータに対応する部分のみが保持されます。

この実装により、元の Muon アルゴリズムの数学的特性を可能な限り維持しながら、メモリ フットプリントと通信オーバーヘッドが最小限に抑えられます。

ミューオン拡張の実現可能性の証明

上記の Muon の改良に基づいて、著者らは次のような結果を達成しました。Llama アーキテクチャの一連の高密度モデルに対して、Muon と AdamW のモデル スケーリング比較実験を実施しました。

結果によると、最適な計算予算では、Muon のサンプル効率は AdamW の 1.92 倍であり、同等のパフォーマンスを達成するには AdamW のトレーニング FLOPS の 52% しか必要ありません。

この発見は、大規模なトレーニングにおける Muon の効率性の利点を裏付けています。

この基盤を基に、著者らは、DeepSeek-V3-Small アーキテクチャに基づく改良された Muon アーキテクチャを使用してMoonlightモデルをトレーニングしました。

Moonlight は、合計 152.9 億のパラメータと 22.4 億のアクティベーション パラメータ、および 5.7T のトレーニング トークン数を持つ MoE モデルです。

同じサイズとデータ量のモデルと比較すると、Moonlight は、英語の理解と推論 (MMLU、TriviaQA、BBH)、コード生成 (HumanEval、MBPP)、数学的推論 (GSM8K、MATH、CMATH)、中国語の理解 (C-Eval、CMMLU) など、さまざまなタスクで大幅に優れたパフォーマンスを実現します。

大規模なデータセットでトレーニングされた高密度モデルと比較しても、Moonlight は強力な競争力を発揮します。

いくつかのよく知られている言語モデルとの比較により、Moonlight はパフォーマンスとトレーニングの予算の面でパレート最適解を前進させることがわかります。

（注：パレート面は経済学と経営学における概念であり、複数の目的間の最適なバランスを実現する多目的意思決定問題に対するすべての可能な最適解の集合を表します。パレート面上の各点は、1つの目的の改善が必然的に別の目的の犠牲を伴うことを意味し、複数の目的間で達成される最適なトレードオフを表します。）

Muon の行列パラメータ更新の基本的なメカニズムをさらに分析するために、著者らは、さまざまなトレーニング段階で Muon と AdamW によってトレーニングされたモデルのパラメータ行列の特異値スペクトルを比較しました。

結果は、Muon最適化行列が、すべての層とパラメータにおいてAdamWよりも一貫して高い特異エントロピーを持つことを示しました。これは、直交化を通じてより多様な表現を学習するというMuonの直感を経験的に検証するものです。

最後に、著者らはMoonモデルを基に、微調整段階におけるMuonの効果についても調査しました。その結果、事前学習段階と微調整段階の両方でMuonを使用することで最良の結果が得られることが示されました。

技術レポート: https://github.com/MoonshotAI... コード: https://github.com/MoonshotAI... Moonlight モデル: https://huggingface.co/moonsh...