|
Dark Side of the MoonとDeepSeekが再び衝突した。 前回は論文で、両社がほぼ相次いで改良されたアテンション機構をリリースしました。「DeepSeek NSAとの衝突：Kimi Yang Zhilin氏署名の新アテンションアーキテクチャMoBAがリリース、コードも公開」と「たった今！DeepSeekのLiang Wenfeng氏が新アテンションアーキテクチャNSAに自ら署名」をご覧ください。 今回はオープンソースです。 先週の金曜日、DeepSeek は 5 つのコード リポジトリを今週オープンソース化すると公式発表しましたが、そのリポジトリは深夜に Dark Side of the Moon によって奪われてしまいました。 昨日、Dark Side of the Moon は、AdamW オプティマイザーの 2 倍の計算効率を誇る Muon オプティマイザーの改良版をリリースしました。 研究チームによると、オリジナルのMuonオプティマイザーは小規模な言語モデルの訓練では優れた性能を発揮するが、大規模モデルへの拡張性はまだ実証されていない。そこで研究チームは、Muonを拡張するために重要な2つの手法を特定した。
これらの技術により、Muonはハイパーパラメータの調整を必要とせずに大規模な学習に直接使用できます。スケーリング則の実験では、Muonは最適な学習を計算するAdamWの2倍の計算効率を示すことが示されています。 これらの改良を基に、Dark Side of the MoonはMoonを用いて5.7兆トークンで学習した3B/16BパラメータのMixture-of-Expert（MoE）モデル「Moonlight」を発表しました。このモデルは新たなパレートフロンティアを確立しました。つまり、同じ学習予算で、全てのパフォーマンス指標において同時にこのモデルを上回るモデルは存在しないということです。 以前のモデルと比較して、Moonlight はより少ないトレーニング FLOP でより優れたパフォーマンスも達成しました。 下図に示すように、本研究ではスケーリング則の調査を行い、MuonとロバストなAdamWベースラインを比較することで、Muonの優れた性能を実証しました。Muonは、AdamWの約52%のトレーニングFLOPで、同等のトレーニング性能を達成しました。 Dark Side of the Moon は、メモリが最適化され、通信効率の高い Muon 実装コードをオープンソース化しただけでなく、将来の研究をサポートするための事前トレーニング、命令の調整、中間チェックポイントもリリースしました。 論文のタイトルは「MUON は LLM トレーニングにスケーラブルです」です。
研究紹介ミューオンを拡張する Muon オプティマイザーは、2024 年に OpenAI の研究者 Keller Jordan らによって提案されました。彼らの研究によると、小規模なトレーニングでは Muon が AdamW を大幅に上回るパフォーマンスを示しました。 しかし、Dark Side of the Moonは、これをスケールアップしてより大きなモデルを学習させ、より多くのトークンを使用すると、パフォーマンスの向上が徐々に減少することを発見しました。重みとレイヤー出力のRMS値は増加し続け、最終的にbf16の高精度範囲を超え、モデルのパフォーマンスに悪影響を与える可能性があることが分かりました。 この問題に対処するために、Dark Side of the Moon は Muon に標準的な AdamW (Loshchilov 他、2019) の重み減衰メカニズムを導入しました。 このメカニズムを調査するため、研究者らはMuonにおいて重み減衰の有無で実験を行いました。8億個のパラメータと1000億トークン（最適なトレーニングトークン数の約5倍）でモデルを学習しました。図2は、AdamW、重み減衰なしの元のMuon、重み減衰ありのMuonを用いて学習したモデルの検証損失曲線を示しています。 結果によると、元の Muon は初期段階ではより速く収束しますが、一部のモデルの重みは時間の経過とともに大きくなりすぎて、モデルの長期的なパフォーマンスに影響を及ぼす可能性があります。 この問題は重み減衰を追加することで解決されました。結果は、重み減衰を適用したMuonが元のMuonとAdamWよりも性能が高く、検証損失が低いことを示しています。式3は、λが重み減衰比である式です。 一貫性のあるRMS更新。研究者たちは、AdamとAdamWの重要な特性として、更新されたRMSが1付近に維持されることを発見しました。しかし、月の暗黒面では、Muonの更新されたRMSは、以下の補題1に示すように、パラメータ行列の形状に応じて変化することを発見しました。 異なる形状の行列にわたって一貫した RMS 更新を維持するために、本研究では、ミューオン行列の更新を拡張して、補題 1 で説明した効果を打ち消します。 実用アプリケーションでは、研究者はAdamWとMuonを組み合わせて非行列パラメータを扱うことがよくあります。本論文では、最適化ハイパーパラメータ（学習率η、重み減衰λ）を行列パラメータと非行列パラメータ間で共有できるようにすることを目的とします。 そこで彼らは、Muonの更新されたRMSをAdamWのRMSに近い範囲に調整することを提案しました。彼らは以下の調整を用いて、Muonの更新されたRMSをこの範囲にスケーリングしました。 分散ミューオン Dark Side of the Moonチームは、ZeRO-1をベースにした分散ソリューション「Distributed Muon」も提案しました。Distributed Muonは、データ並列性（DP）に基づいて最適化状態を分割する点でZeRO-1を踏襲し、通常のZeRO-1 AdamW最適化器と比較して2つの追加操作を導入しています。アルゴリズム1は、Distributed Muonの実装を示しています。 実験RMS一貫性 すべてのマトリックスパラメータ更新のRMS値がAdamWのRMSと一致することを確認するために、研究チームはパラメータ更新のRMSを制御する2つの方法を試し、AdamWのみを使用したベースラインのRMSと比較しました。 大規模モデル学習では様々な予期せぬ状況が発生する可能性があるため、研究チームは学習の初期段階におけるMuonの影響をテストしました。行列次元の差が大きくなるにつれて、RMS更新の不一致の問題がより顕著になります。研究チームはモデルアーキテクチャを微調整し、Swiglu MLPを標準的な2層MLPに置き換え、行列パラメータの形状を[H, 2.6H]から[H, 4H]に変更しました。 チームはモデルの損失を評価し、主要なパラメータ、特に形状 [H, H] の注目クエリ重みと形状 [H, 4H] の MLP 重みの RMS を監視しました。 実験結果 (表 1 を参照) は、Update Norm と Adjusted LR の両方がベースライン手法よりも優れていることを示しています。また、Adjusted LR の計算コストが低いため、後続の実験では Adjusted LR が選択されました。 ミューオンのスケーリング則 AdamW との公平な比較を確実にするために、チームは Llama アーキテクチャに基づくさまざまなモデルに Muon を拡張しました。 Muon の場合、RMS が AdamW と一致するため、チームは AdamW のハイパーパラメータを直接再利用しました。 実験結果によると、近似したスケーリング則曲線は、最適な計算設定下では、MuonがAdamWに匹敵する性能を達成するために、学習FLOPの約52%しか必要としないことを示しています。これは、大規模言語モデルの学習におけるMuonの効率性をさらに実証しています。 Muonを使用した事前トレーニング モデル アーキテクチャにおける Muon のパフォーマンスを評価するために、チームは DeepSeek-V3-Small アーキテクチャを使用して Moonlight モデルを最初から事前トレーニングしました。 Moonlightモデルは合計5.7兆トークンで学習されました。しかし、1.2兆トークンの学習段階では、DeepSeek-V3-Small（1.33兆トークンで学習された2.4B/16BパラメータのMoEモデル）およびMoonlight-A（Moonlightと同じ設定だが、AdamWオプティマイザーを使用）と比較しました。表4に示すように、Moonlightは言語、数学、コーディングなどのタスクにおいてMoonlight-Aを大幅に上回り、Moonlightのスケーリング優位性を実証しました。 完全な学習後、MoonlightはLLAMA3-3B、Qwen2.5-3B、Deepseek-v2-Liteといった同規模のオープンソースモデルと比較されました。その結果、Moonlightは同じトークン数で学習したモデルよりも優れた性能を示し、より大きなパラメータスケールを持つモデルとも競争力があることが示されました。 さらに、研究チームは、Muon が、特に MoE モデルにおいて、モデルの重みの更新をより「多様化」できることも発見しました。 ファインチューニングフェーズでは、事前学習フェーズとファインチューニングフェーズの両方でMuonを使用すると、AdamWを組み合わせて使用するよりもモデルのパフォーマンスが向上します。ただし、ファインチューニングと事前学習で使用される最適化手法に一貫性がない場合、この利点は目立たなくなります。 詳細は原論文をご参照ください。 |
