月之暗面和 DeepSeek 這次又「撞車」了。

上次是論文，兩家?guī)缀跚昂竽_放出改進版的注意力機制，可參考《撞車 DeepSeek NSA，Kimi 楊植麟署名的新注意力架構(gòu) MoBA 發(fā)布，代碼也公開》、《剛剛！DeepSeek 梁文鋒親自掛名，公開新注意力架構(gòu) NSA》。

這次是開源。

上周五，DeepSeek 剛剛官宣這周要連續(xù)開源 5 個代碼庫，卻被月之暗面深夜截胡了。

昨天，月之暗面搶先一步開源了改進版 Muon 優(yōu)化器，比 AdamW 優(yōu)化器計算效率提升了 2 倍。

團隊人員表示，原始 Muon 優(yōu)化器在訓(xùn)練小型語言模型方面表現(xiàn)出色，但其在擴展到更大模型方面的可行性尚未得到證明。因此，團隊人員確定了兩種對擴展 Muon 至關(guān)重要的技術(shù)：

• 添加權(quán)重衰減：對擴展到更大模型至關(guān)重要。
• 一致的 RMS 更新：在模型更新上執(zhí)行一致的均方根。

這些技術(shù)使得 Muon 能夠在大規(guī)模訓(xùn)練中直接使用，而無需調(diào)整超參數(shù)。Scaling law 實驗表明，與計算最優(yōu)訓(xùn)練的 AdamW 相比，Muon 的計算效率提升了 2 倍。

基于這些改進，月之暗面推出了 Moonlight，這是一個 3B/16B 參數(shù)的 Mixture-of-Expert（MoE）模型，使用 Muon 進行了 5.7 萬億 tokens 的訓(xùn)練。該模型刷新了當前的「帕累托前沿」，換句話說，在相同的訓(xùn)練預(yù)算下，沒有其他模型能在所有性能指標上同時超越它。

與之前的模型相比，Moonlight 也以更少的訓(xùn)練 FLOPs 獲得了更好的性能。

如下圖所示，該研究進行了 Scaling law 研究，將 Muon 與強大的 AdamW 基線進行了比較，結(jié)果展示了 Muon 的卓越性能。Muon 實現(xiàn)了與 AdamW 訓(xùn)練相當?shù)男阅?，同時僅需要大約 52% 的訓(xùn)練 FLOP。

月之暗面不但開源了內(nèi)存優(yōu)化且通信高效的 Muon 實現(xiàn)代碼，并且還發(fā)布了預(yù)訓(xùn)練、指令調(diào)優(yōu)以及中間檢查點，以支持未來的研究。

論文《 MUON IS SCALABLE FOR LLM TRAINING 》。

• 論文地址：https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
• 代碼地址：https://github.com/MoonshotAI/Moonlight
• 模型地址：https://huggingface.co/moonshotai/Moonlight-16B-A3B

研究介紹

擴展 Muon 

Muon 優(yōu)化器為 OpenAI 研究者 Keller Jordan 等人在 2024 年提出的，他們的研究表明在小規(guī)模訓(xùn)練中 Muon 的表現(xiàn)顯著優(yōu)于 AdamW。

但月之暗面發(fā)現(xiàn)，當將其擴展到訓(xùn)練更大模型并使用更多 token 時，模型性能提升逐漸減弱。他們觀察到，權(quán)重和層輸出的 RMS 值持續(xù)增長，最終超出了 bf16 的高精度范圍，這可能會損害模型的性能。

為了解決這個問題，月之暗面在 Muon 中引入了標準的 AdamW（Loshchilov 等人，2019）權(quán)重衰減機制。

為了探究這一機制，研究者對 Muon 進行了有無權(quán)重衰減的實驗，他們訓(xùn)練了一個包含 800M 參數(shù)和 100B token（約為最優(yōu)訓(xùn)練 token 量的 5 倍）的模型。圖 2 展示了使用 AdamW、無權(quán)重衰減的原始 Muon 以及帶權(quán)重衰減的 Muon 訓(xùn)練的模型的驗證損失曲線。

結(jié)果顯示，雖然原始 Muon 在初期收斂速度更快，但一些模型權(quán)重隨時間的推移增長過大，可能會影響模型的長期性能。

加入權(quán)重衰減后解決了這一問題 —— 結(jié)果表明，帶權(quán)重衰減的 Muon 優(yōu)于原始 Muon 和 AdamW，獲得了更低的驗證損失。公式 3 為表達式，其中 λ 為權(quán)重衰減比率。

一致的 RMS 更新。研究者發(fā)現(xiàn) Adam 和 AdamW 的一個重要特性是，它們將更新的 RMS 維持在 1 左右。然而，月之暗面發(fā)現(xiàn) Muon 更新 RMS 會根據(jù)參數(shù)矩陣形狀的變化而變化，具體如下引理 1 所示：

為了在不同形狀矩陣之間保持一致的 RMS 更新，該研究通過來擴展 Muon 矩陣更新，從而抵消引理 1 中提到的影響。

在實際應(yīng)用中，研究者通常將 AdamW 與 Muon 結(jié)合使用，以處理非矩陣參數(shù)。本文希望優(yōu)化器超參數(shù)（學(xué)習率 η、權(quán)重衰減 λ）能夠在矩陣參數(shù)和非矩陣參數(shù)之間共享。

因此他們提出將 Muon 更新的 RMS 調(diào)整到與 AdamW 相似的范圍。他們通過以下調(diào)整將 Muon 更新 RMS 縮放至這一范圍：

分布式 Muon

月之暗面團隊還提出了一種基于 ZeRO-1 的分布式解決方案，稱為分布式 Muon（Distributed Muon）。分布式 Muon 遵循 ZeRO-1 在數(shù)據(jù)并行（DP）上對優(yōu)化器狀態(tài)進行劃分，并與普通的 ZeRO-1 AdamW 優(yōu)化器相比引入了兩個額外的操作，算法 1 描述了分布式 Muon 的實現(xiàn)。

實驗

RMS 的一致性

為了使所有矩陣參數(shù)更新的 RMS 值與 AdamW 的 RMS 保持一致，研究團隊嘗試了兩種方法來控制參數(shù)更新的 RMS，并將其與只用了 AdamW 的基線的 RMS 進行了對比。

由于大規(guī)模訓(xùn)練模型時，會出現(xiàn)各種意料之外的情況，因此，研究團隊測試了 Muon 對訓(xùn)練早期階段的影響。當矩陣維度差異增大時，更新 RMS 不一致的問題會更加明顯。該團隊對模型架構(gòu)進行了微調(diào)，用標準的 2 層 MLP 替換了 Swiglu MLP，并將其矩陣參數(shù)的形狀從 [H, 2.6H] 改為 [H, 4H]。

團隊評估了模型的損失，并監(jiān)控了關(guān)鍵參數(shù)的 RMS，尤其是形狀為 [H, H] 的注意力查詢權(quán)重和形狀為 [H, 4H] 的 MLP 權(quán)重。

實驗結(jié)果表明（見表 1），Update Norm 和 Adjusted LR 均優(yōu)于基線方法，且 Adjusted LR 的計算成本更低，因此被選用于后續(xù)實驗。

Muon 的 Scaling Law

為了與 AdamW 公平比較，該團隊在一系列基于 Llama 架構(gòu)的模型上對 Muon 進行了拓展。

對于 Muon，由于其 RMS 與 AdamW 匹配，團隊直接復(fù)用了 AdamW 的超參數(shù)。

實驗結(jié)果顯示，擬合的 Scaling Law 曲線表明，在計算最優(yōu)設(shè)置下，Muon 僅需約 52% 的訓(xùn)練 FLOPs 即可達到與 AdamW 相當?shù)男阅?。這進一步說明了 Muon 在大規(guī)模語言模型訓(xùn)練中的高效性。

使用 Muon 進行預(yù)訓(xùn)練

為了評估 Muon 在模型架構(gòu)中的表現(xiàn)，該團隊使用 DeepSeek-V3-Small 架構(gòu)從頭開始預(yù)訓(xùn)練了 Moonlight 模型。

Moonlight 模型總共進行了 5.7 萬億 tokens 的訓(xùn)練，但在訓(xùn)練到 1.2 萬億 tokens 的階段，團隊將其與 DeepSeek-V3-Small（使用 1.33T tokens 訓(xùn)練的 2.4B/16B 參數(shù) MoE 模型）和 Moonlight-A（與 Moonlight 設(shè)置相同，但使用 AdamW 優(yōu)化器）進行了比較。如表 4 所示，Moonlight 在語言、數(shù)學(xué)和編碼等任務(wù)上都顯著優(yōu)于 Moonlight-A，證明了 Muon 的擴展優(yōu)勢。

在完整訓(xùn)練后，Moonlight 與類似規(guī)模的開源模型（如 LLAMA3-3B、Qwen2.5-3B 和 Deepseek-v2-Lite）進行了比較。結(jié)果顯示，Moonlight 在性能上優(yōu)于使用相同數(shù)量 tokens 訓(xùn)練的模型，與更大參數(shù)規(guī)模模型相比，也較有競爭力。

此外，研究團隊還發(fā)現(xiàn)，Muon 可以讓模型的權(quán)重更新更「多樣化」，尤其在 MoE 模型中表現(xiàn)突出。

在微調(diào)階段，在預(yù)訓(xùn)練和微調(diào)階段都使用 Muon，模型表現(xiàn)會比用 AdamW 的組合更好，但如果微調(diào)和預(yù)訓(xùn)練的優(yōu)化器不一致，優(yōu)勢就不明顯了。

更多細節(jié)，請參閱論文原文。