Muon優化器：AI模型訓練算法的下一個里程碑？| 目前還不是業界焦點，但有潛力是重大基礎創新 精華
人工智能（AI）快速發展，模型訓練是核心環節，優化器扮演著至關重要的角色，它負責調整模型的參數，讓模型在數據上表現得更好。多年來，AdamW優化器一直是優化器的標桿，因其穩定性和高效性深受研究者和工程師的喜愛。然而，隨著AI模型規模的不斷擴大，訓練成本和時間的需求也在激增，這讓人們開始尋找更高效的優化方法。近期，一種名為Muon的優化器算法悄然出現（源代碼 ????https://github.com/KellerJordan/Muon ????），盡管它還未成為業界焦點，但其獨特的設計和卓越的性能表明，它可能是AI模型訓練領域的一次重大基礎創新。
優化器算法：AI訓練的“幕后推手”
為什么優化器如此重要？
在深度學習中，模型訓練的目標是通過調整參數，讓模型的預測結果盡可能接近真實數據。這個過程通常是通過定義一個損失函數來實現的，損失函數衡量了模型預測與真實值之間的差距。而優化器的任務，就是根據損失函數的梯度（gradient），一步步調整模型的參數，找到損失最小的“最佳狀態”。
想象一下，我們在一個崎嶇的山谷中尋找最低點。優化器就像我們的導航儀，告訴我們每一步該往哪個方向走、走多遠。一個好的優化器不僅能更快地帶我們到達谷底（收斂），還能避免在陡峭的坡道上上上下下（訓練不穩定）。優化器的效率直接決定了訓練速度、計算資源需求，以及最終模型的性能。
過去幾年，AdamW（Adam with Weight Decay）一直是訓練大型語言模型的首選（如Qwen、DeepSeek、LLaMA等，閉源的大模型不清楚，但大概率也是AdamW）。它結合了自適應學習率和權重衰減（weight decay）的優點，能夠在復雜的參數空間中穩定、高效地工作。然而，隨著模型參數從幾億增加到幾千億，訓練時間從幾天變成幾周甚至幾個月，AdamW的局限性開始顯現——它在超大規模場景下的效率開始受到挑戰。進一步提升AI能力，我們需要更大的模型和更多的訓練資源。但計算資源的成本高昂，訓練時間過長也會拖慢研究和應用的進度。因此，開發更高效的優化器，不僅是技術上的追求，更是經濟和實踐上的迫切需求。
Muon：從動量到正交化的革新
Muon的基本原理
Muon的全稱是MomentUm Orthogonalized by Newton-Schulz（動量正交化Newton-Schulz），它是一種專為神經網絡隱藏層設計的優化器。它的核心思想并不復雜，但卻非常巧妙：先用經典的SGD（隨機梯度下降）動量法生成參數更新，然后通過一個特殊的“后處理”步驟——Newton-Schulz正交化，讓這些更新更高效。
讓我們一步步拆解這個過程：
- 動量法生成更新Muon首先沿用了SGD動量的思路。動量法就像給梯度加了一個“慣性”，讓參數更新不僅依賴當前梯度，還參考之前的更新方向。這樣可以加速收斂，避免在參數空間中“左搖右晃”。在Muon中，這一過程生成一個更新矩陣（記為(Bt)）。
- 正交化：讓更新更“聰明”接下來，Muon引入了關鍵創新：通過Newton-Schulz迭代，將更新矩陣(Bt)“正交化”，生成一個新的更新矩陣(Ot)。所謂正交化，簡單來說，就是讓更新矩陣的列（或行）彼此垂直，變成一個“半正交矩陣”（滿足( OTO = I )或( O OT= I )）。為什么這樣做？傳統的動量更新往往會被某些“主導方向”牽著走，而其他“稀有方向”（對學習也很重要，但幅度小）的貢獻被掩蓋。正交化就像重新分配了這些方向的“發言權”，讓更新更全面、更高效地探索參數空間。
- 參數更新最后，Muon用正交化后的更新矩陣( Ot)來調整參數：
這里，η是學習率，θ是模型參數。
Newton-Schulz迭代：高效的正交化工具
正交化聽起來很高級，計算上也很復雜。如果用傳統的SVD（奇異值分解）來正交化，計算量太大，速度太慢，完全不適合現代GPU加速的訓練環境。而Muon采用的Newton-Schulz迭代，則是一個高效的替代方案。
這個迭代過程的核心是，通過反復計算一個多項式函數（比如五次多項式），逐步將更新矩陣調整到接近正交的狀態。它的優點在于：
- 低計算開銷：每次迭代只需幾次矩陣乘法，遠比SVD快。
- 數值穩定性：可以用bfloat16（一種低精度浮點格式）運行，非常適合現代GPU。
例如，Muon的實現中，Newton-Schulz迭代的公式是：
其中，(a = 3.4445, b = -4.7750, c = 2.0315)是精心調優的系數。經過5次迭代，就能得到一個“足夠正交”的更新矩陣，既高效又實用。
Muon的優勢：效率與潛力的結合
Muon并不是憑空出現的“黑魔法”，它的優勢經過了多次實驗驗證。以下是Muon相比AdamW的幾個突出優點：
- 更快的訓練速度在多個基準測試中，Muon展現了驚人的加速能力。例如：
- 在CIFAR-10圖像分類任務中，Muon將達到94%準確率的時間從3.3 A100-seconds降到2.6 A100-seconds。
- 在NanoGPT speedrunning任務（訓練一個小規模GPT模型）中，Muon將訓練速度提升了1.35倍。 這些結果表明，Muon能顯著縮短訓練時間，尤其是在資源有限的場景下。
- 卓越的可擴展性Muon在大規模模型上同樣表現出色。例如，在訓練一個1.5億參數的語言模型時，Muon僅用10個8xH100-hours就達到了GPT-2 XL的性能水平，而AdamW需要13.3 hours。隨著模型規模繼續擴大，這種差距可能會更明顯。
- 低計算開銷盡管正交化聽起來很復雜，但Muon的額外計算開銷非常小。在典型語言模型訓練中，Muon的FLOP（浮點運算）開銷低于1%。這得益于Newton-Schulz迭代的高效性，以及對現代硬件的優化。
- 與現有框架兼容Muon并非完全替代AdamW，而是與之互補。實際應用中，Muon通常優化網絡的隱藏層參數，而嵌入層（embedding）和分類器頭（head）仍由AdamW處理。這種“分工合作”的方式，讓Muon可以無縫集成到現有訓練流程中。
如果Muon替代AdamW
假設Muon的潛力被充分驗證，并逐步取代AdamW成為新的標準優化器，AI領域可能會迎來以下重大變化：
- 訓練成本大幅降低Muon的高效性意味著，訓練一個大規模模型所需的計算資源和時間將顯著減少。這不僅能節省數千萬RMB的算力成本，還能讓更多中小型研究團隊參與到前沿AI研究中，降低技術門檻。
- 模型規模的進一步突破更高的訓練效率，讓研究者有能力嘗試更大規模的模型。比如，現在訓練一個1000億參數模型可能需要數月，而有了Muon，或許幾周就能完成。這將推動AI能力的上限不斷提升。
- AI應用加速落地更快的訓練速度意味著模型從研究到應用的時間縮短。無論是工業大模型（創新奇智在做的事情），還是其他行業大模型，AI技術的普及速度都可能因此加快。
- 優化器研究的復興Muon的成功可能會重新點燃對優化器算法的興趣。過去幾年，AdamW幾乎“一統江湖”，新優化器的研究相對沉寂。如果Muon證明了創新優化器的價值，研究者可能會投入更多精力，探索其他潛在的突破。
OpenAI：Muon的未來舞臺？
2024年，Muon的開發者Jordan在個人博客中詳細介紹了Muon的設計理念和實驗結果，隨后不久，他宣布加入OpenAI。這一動向并非巧合，很可能OpenAI看中了Muon的潛力。
月之暗面的實踐：Muon的驗證
Moonshot AI（月之暗面）近期在Muon的基礎上進行了大規模實踐，驗證了它在真實場景中的潛力。他們不僅改進了Muon，還訓練了一個名為Moonlight的3B/16B參數Mixture-of-Expert（MoE）模型，用5.7萬億tokens的數據進行了測試。
Moonshot AI的改進
Moonshot AI發現，原始Muon在小規模任務上表現出色，但在超大規模訓練中會遇到問題，比如模型權重增長過大，影響穩定性。為此，他們提出了兩個關鍵改進：
- 引入權重衰減他們將AdamW的權重衰減機制融入Muon，更新公式變為：
- 這有效控制了權重大小，提升了長期訓練的性能。
- 調整更新尺度Muon的更新幅度（RMS）會因參數矩陣的形狀而變化，可能導致訓練不穩定。Moonshot AI提出按矩陣最大維度縮放更新，比如：
這樣可以保持更新幅度一致，并與AdamW兼容。
Moonlight的驚艷表現
基于這些改進，Moonshot AI用Muon訓練了Moonlight模型，并在多個基準測試中取得了優異成績。例如：
- 在MMLU（英語理解）上，Moonlight得分70.0，超越了同規模的Llama3.2-3B（54.7）和Deepseek-v2-Lite（58.3）。
- 在GSM8K（數學推理）上，得分77.4，接近Qwen2.5-3B（79.1），但訓練tokens僅為后者的三分之一。
- 訓練效率上，Moonlight只需約52%的FLOPs，就能達到AdamW的性能水平。
這些結果表明，Muon不僅能加速訓練，還能提升模型性能，尤其在數學和代碼任務上表現突出。
結語：Muon的潛力與未來
Muon作為一種新興的優化器算法，以其獨特的設計和卓越的性能，展現了巨大的潛力。它通過正交化更新矩陣，打破了傳統優化器的局限，在訓練速度、可擴展性和計算效率上都超越了AdamW。盡管目前它還未引起全球關注，但Keller Jordan的加入OpenAI，以及Moonshot AI的成功實踐，都預示著Muon可能成為AI訓練領域的下一個里程碑。
?本文轉載自???后向傳播???，作者： 張發恩
