月之暗面和 DeepSeek 這次又「撞車(chē)」了。

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

這次是開(kāi)源。

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

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

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

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

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

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

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

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

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

論文《 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

研究介紹

擴(kuò)展 Muon 

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

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

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

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

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

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

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

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

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

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

分布式 Muon

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

實(shí)驗(yàn)

RMS 的一致性

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

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

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

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

Muon 的 Scaling Law

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

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

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

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

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

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

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

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

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

更多細(xì)節(jié)，請(qǐng)參閱論文原文。