今天小編分享的科學經驗：月之暗面開源改進版Muon優化器，算力需求比AdamW銳減48%，DeepSeek也适用，歡迎閱讀。
算力需求比 AdamW 直降 48%，OpenAI 技術人員提出的訓練優化算法 Muon，被月之暗面團隊又推進了一步！
團隊發現了 Muon 方法的 Scaling Law，做出改進并證明了 Muon 對更大的模型同樣适用。
在參數量最高 1.5B 的不同 Llama 架構模型上，改進後的 Muon 算力需求僅為 AdamW 的 52%。
同時團隊還基于 DeepSeek 架構訓練出了一個 16B 的 MoE 模型，與改進後的優化算法一同開源。
Muon 技術博客發布當時主要适用于較小的模型和數據集，作者留下了三個懸而未決的問題：
Muon 能否用于更大規模的訓練？
Muon 能否在更大規模的 GPU 集群上使用？
Muon 是否同樣适用于微調和強化學習？
現在月暗團隊用實驗給出了回答——全部都是 Yes。
消息一出，當時 Muon 的作者也都很激動，主要作者Keller Jordan表示這是 Muon 規模化的首個成功報告，為團隊送上了祝賀。
另一名貢獻者，當時負責 Muon 規模化實驗的 Hyperbolic Labs 聯創兼 CTO Yuchen Jin也表示，月暗團隊的這項成果，是 Muon 的一次勝利。
将 AdamW 特點引入 Muon
在介紹月暗團隊的工作之前，先來了解一下 Muon 是個什麼樣的技術。
這是一種神經網絡隐藏層的 2D 參數優化器，主要作者是 OpenAI 深度學習團隊的 Keller Jordan。
這項成果發表于去年的 12 月 8 日，而 Keller 也是去年 12 月加入的 OpenAI。
Muon 的核心思想是通過正交化梯度更新矩陣，避免參數更新陷入局部極小，使模型能夠學習到更加多樣化的特征表示。
在 94% 的精度下，Muon 把 CIFAR-10 在 A100 上的訓練時間從 3.3 秒縮短至 2.6 秒。
不過當時 Muon 團隊只證明了其在小型模型和數據集上的可行性，對于較大的模型能否适用則是個未知數。
現在經過月暗團隊的改進之後，Muon被證明對于更大的模型和數據集同樣适用。
針對模型本身，團隊吸收了 AdamW 中的一些特點，移植到了 Muon 當中，具體包括兩個方面。
一是引入了權重衰減機制，在權重更新公式中添加了一個帶有衰減系數的項。
這樣做的原因是作者發現直接将 Muon 應用到大規模訓練時，模型權重和層輸出的幅度會持續增長，最終超出 bf16 的高精度表示範圍，損害模型性能。
在訓練一個 8 億參數模型至 100B tokens（約 5 倍計算預算最優）的過程中，團隊對比了 AdamW、無權重衰減的 Muon 和帶權重衰減的 Muon。
結果顯示，帶權重衰減的 Muon 在過拟合階段取得了最佳效果，驗證了權重衰減的必要性。
第二項改進，是調整了 Muon 的參數更新尺度，使不同形狀矩陣參數的更新幅度保持一致，并與 AdamW 的更新幅度匹配。
Muon 的一個特點是，對于形狀為 [ A,B ] 的矩陣參數，其理論更新幅度為 sqrt ( 1/max ( A,B ) ) 。
這導致不同形狀矩陣參數的更新幅度差異很大，比如對于 MLP 這種寬矩陣，更新會過小，而将每個 head 看作獨立矩陣時，更新又會過大。
此外，這個幅度也與 AdamW 不一致，給超參數的設定帶來困難。
為了讓不同矩陣參數的更新幅度匹配，并與 AdamW 保持一致，作者嘗試了幾種改進方案，最終選擇直接基于形狀調整每個參數的學習率。
其中 0.2 是通過實驗确定的一個常數，用于将 Muon 的更新尺度與 AdamW 對齊。
除了對 Muon 本身的改進，要想将 Muon 用于更大規模的訓練，還需要将其擴展到分布式訓練環境中。
由于 Muon 需要完整的梯度矩陣來計算正交化的更新量，而現有的分布式訓練框架（如 ZeRO-1、Megatron-LM 等）都假設優化器狀态可以獨立地按元素切分到不同設備上，所以它們無法直接支持 Muon。
為了解決這個問題，論文作者提出了分布式 Muon 的并行化策略。
它在 ZeRO-1 的基礎上引入了兩個額外的操作 :
一是在每個數據并行組内做梯度聚合通信，将分散的梯度切片合并成完整的矩陣；
二是基于聚合後的梯度矩陣并行計算正交化的更新量，然後只保留與本地參數對應的那一部分。
這種實現方式在最小化内存占用和通信開銷的同時，最大限度地保留了原始 Muon 算法的數學性質。
證明 Muon 擴展可行性
基于上述 Muon 改進 , 作者取得了以下成果，作者在 Llama 架構的一系列稠密模型上，進行了 Muon 和 AdamW 的模型縮放對比實驗。
結果表明 , 在計算預算最優的情況下，Muon 的樣本效率是 AdamW 的 1.92 倍，即訓練 FLOPS 只需 AdamW 的 52%，就能達到相當的性能。
這一發現證實了 Muon 在大規模訓練中的效率優勢。
在此基礎之上，作者以 DeepSeek-V3-Small 架構作為基礎，用改進的 Muon 訓練了Moonlight模型。
Moonlight 是一個 MoE 模型，具有 15.29B 的總參數和 2.24B 激活參數，訓練 token 量為 5.7T。
與相同規模和數據量的模型相比，Moonlight 在英語理解與推理（MMLU、TriviaQA、BBH）、代碼生成（HumanEval、MBPP）、數學推理（GSM8K、MATH、CMATH）、中文理解（C-Eval、CMMLU）等各類任務上都取得了明顯更好的性能。
即使與使用更大數據集訓練的稠密模型相比，Moonlight 也展現了極強的競争力。
與多個知名語言模型的對比表明，Moonlight在性能 - 訓練預算平面上推進了帕累托前沿（Pareto Frontier）。
（注：帕累托前沿是一個經濟學和管理學中的概念，描述的是在多目标決策問題中所有可能的最優解的集合，這些解在多個目标之間取得了最佳平衡。在帕累托前沿上的每一個點，都意味着一個目标的改善必然以犧牲另一個目标為代價，因此它代表了在多個目标之間實現的最佳權衡。）
為了進一步分析 Muon 更新矩陣參數的内在機制，作者對比了 Muon 和 AdamW 訓練得到的模型在不同訓練階段的參數矩陣奇異值譜。
結果發現，Muon 優化的矩陣在各層各類參數上，總是比 AdamW 有更高的奇異值熵。這從經驗上驗證了 Muon 通過正交化來學習更多樣化表示的直覺。
最後，在 Moonlight 模型的基礎上，作者還探索了 Muon 在指導微調階段的效果，結果表明，在預訓練和微調階段均使用 Muon 的效果是最佳的。
技術報告：
https://github.com/MoonshotAI/Moonlight/blob/master/Moonlight.pdf
Code：
https://github.com/MoonshotAI/Moonlight
Moonlight 模型：
https://huggingface.co/moonshotai/Moonlight-16B-A3B
