人工智能（AI）快速發(fā)展，模型訓(xùn)練是核心環(huán)節(jié)，優(yōu)化器扮演著至關(guān)重要的角色，它負(fù)責(zé)調(diào)整模型的參數(shù)，讓模型在數(shù)據(jù)上表現(xiàn)得更好。多年來，AdamW優(yōu)化器一直是優(yōu)化器的標(biāo)桿，因其穩(wěn)定性和高效性深受研究者和工程師的喜愛。然而，隨著AI模型規(guī)模的不斷擴(kuò)大，訓(xùn)練成本和時(shí)間的需求也在激增，這讓人們開始尋找更高效的優(yōu)化方法。近期，一種名為Muon的優(yōu)化器算法悄然出現(xiàn)（源代碼 ????https://github.com/KellerJordan/Muon ????），盡管它還未成為業(yè)界焦點(diǎn)，但其獨(dú)特的設(shè)計(jì)和卓越的性能表明，它可能是AI模型訓(xùn)練領(lǐng)域的一次重大基礎(chǔ)創(chuàng)新。

優(yōu)化器算法：AI訓(xùn)練的“幕后推手”

為什么優(yōu)化器如此重要？

在深度學(xué)習(xí)中，模型訓(xùn)練的目標(biāo)是通過調(diào)整參數(shù)，讓模型的預(yù)測結(jié)果盡可能接近真實(shí)數(shù)據(jù)。這個過程通常是通過定義一個損失函數(shù)來實(shí)現(xiàn)的，損失函數(shù)衡量了模型預(yù)測與真實(shí)值之間的差距。而優(yōu)化器的任務(wù)，就是根據(jù)損失函數(shù)的梯度（gradient），一步步調(diào)整模型的參數(shù)，找到損失最小的“最佳狀態(tài)”。

想象一下，我們在一個崎嶇的山谷中尋找最低點(diǎn)。優(yōu)化器就像我們的導(dǎo)航儀，告訴我們每一步該往哪個方向走、走多遠(yuǎn)。一個好的優(yōu)化器不僅能更快地帶我們到達(dá)谷底（收斂），還能避免在陡峭的坡道上上上下下（訓(xùn)練不穩(wěn)定）。優(yōu)化器的效率直接決定了訓(xùn)練速度、計(jì)算資源需求，以及最終模型的性能。

過去幾年，AdamW（Adam with Weight Decay）一直是訓(xùn)練大型語言模型的首選（如Qwen、DeepSeek、LLaMA等，閉源的大模型不清楚，但大概率也是AdamW）。它結(jié)合了自適應(yīng)學(xué)習(xí)率和權(quán)重衰減（weight decay）的優(yōu)點(diǎn)，能夠在復(fù)雜的參數(shù)空間中穩(wěn)定、高效地工作。然而，隨著模型參數(shù)從幾億增加到幾千億，訓(xùn)練時(shí)間從幾天變成幾周甚至幾個月，AdamW的局限性開始顯現(xiàn)——它在超大規(guī)模場景下的效率開始受到挑戰(zhàn)。進(jìn)一步提升AI能力，我們需要更大的模型和更多的訓(xùn)練資源。但計(jì)算資源的成本高昂，訓(xùn)練時(shí)間過長也會拖慢研究和應(yīng)用的進(jìn)度。因此，開發(fā)更高效的優(yōu)化器，不僅是技術(shù)上的追求，更是經(jīng)濟(jì)和實(shí)踐上的迫切需求。

Muon：從動量到正交化的革新

Muon的基本原理

Muon的全稱是MomentUm Orthogonalized by Newton-Schulz（動量正交化Newton-Schulz），它是一種專為神經(jīng)網(wǎng)絡(luò)隱藏層設(shè)計(jì)的優(yōu)化器。它的核心思想并不復(fù)雜，但卻非常巧妙：先用經(jīng)典的SGD（隨機(jī)梯度下降）動量法生成參數(shù)更新，然后通過一個特殊的“后處理”步驟——Newton-Schulz正交化，讓這些更新更高效。

讓我們一步步拆解這個過程：

1. 動量法生成更新Muon首先沿用了SGD動量的思路。動量法就像給梯度加了一個“慣性”，讓參數(shù)更新不僅依賴當(dāng)前梯度，還參考之前的更新方向。這樣可以加速收斂，避免在參數(shù)空間中“左搖右晃”。在Muon中，這一過程生成一個更新矩陣（記為(B_t)）。
2. 正交化：讓更新更“聰明”接下來，Muon引入了關(guān)鍵創(chuàng)新：通過Newton-Schulz迭代，將更新矩陣(B_t)“正交化”，生成一個新的更新矩陣(O_t)。所謂正交化，簡單來說，就是讓更新矩陣的列（或行）彼此垂直，變成一個“半正交矩陣”（滿足( O^TO = I )或( O O^T= I )）。為什么這樣做？傳統(tǒng)的動量更新往往會被某些“主導(dǎo)方向”牽著走，而其他“稀有方向”（對學(xué)習(xí)也很重要，但幅度小）的貢獻(xiàn)被掩蓋。正交化就像重新分配了這些方向的“發(fā)言權(quán)”，讓更新更全面、更高效地探索參數(shù)空間。
3. 參數(shù)更新最后，Muon用正交化后的更新矩陣( O_t)來調(diào)整參數(shù)：
   
   

這里，η是學(xué)習(xí)率，θ是模型參數(shù)。

Newton-Schulz迭代：高效的正交化工具

正交化聽起來很高級，計(jì)算上也很復(fù)雜。如果用傳統(tǒng)的SVD（奇異值分解）來正交化，計(jì)算量太大，速度太慢，完全不適合現(xiàn)代GPU加速的訓(xùn)練環(huán)境。而Muon采用的Newton-Schulz迭代，則是一個高效的替代方案。

這個迭代過程的核心是，通過反復(fù)計(jì)算一個多項(xiàng)式函數(shù)（比如五次多項(xiàng)式），逐步將更新矩陣調(diào)整到接近正交的狀態(tài)。它的優(yōu)點(diǎn)在于：

• 低計(jì)算開銷：每次迭代只需幾次矩陣乘法，遠(yuǎn)比SVD快。
• 數(shù)值穩(wěn)定性：可以用bfloat16（一種低精度浮點(diǎn)格式）運(yùn)行，非常適合現(xiàn)代GPU。

例如，Muon的實(shí)現(xiàn)中，Newton-Schulz迭代的公式是：

其中，(a = 3.4445, b = -4.7750, c = 2.0315)是精心調(diào)優(yōu)的系數(shù)。經(jīng)過5次迭代，就能得到一個“足夠正交”的更新矩陣，既高效又實(shí)用。

Muon的優(yōu)勢：效率與潛力的結(jié)合

Muon并不是憑空出現(xiàn)的“黑魔法”，它的優(yōu)勢經(jīng)過了多次實(shí)驗(yàn)驗(yàn)證。以下是Muon相比AdamW的幾個突出優(yōu)點(diǎn)：

1. 更快的訓(xùn)練速度在多個基準(zhǔn)測試中，Muon展現(xiàn)了驚人的加速能力。例如：

• 在CIFAR-10圖像分類任務(wù)中，Muon將達(dá)到94%準(zhǔn)確率的時(shí)間從3.3 A100-seconds降到2.6 A100-seconds。
• 在NanoGPT speedrunning任務(wù)（訓(xùn)練一個小規(guī)模GPT模型）中，Muon將訓(xùn)練速度提升了1.35倍。 這些結(jié)果表明，Muon能顯著縮短訓(xùn)練時(shí)間，尤其是在資源有限的場景下。

2. 卓越的可擴(kuò)展性Muon在大規(guī)模模型上同樣表現(xiàn)出色。例如，在訓(xùn)練一個1.5億參數(shù)的語言模型時(shí)，Muon僅用10個8xH100-hours就達(dá)到了GPT-2 XL的性能水平，而AdamW需要13.3 hours。隨著模型規(guī)模繼續(xù)擴(kuò)大，這種差距可能會更明顯。
3. 低計(jì)算開銷盡管正交化聽起來很復(fù)雜，但Muon的額外計(jì)算開銷非常小。在典型語言模型訓(xùn)練中，Muon的FLOP（浮點(diǎn)運(yùn)算）開銷低于1%。這得益于Newton-Schulz迭代的高效性，以及對現(xiàn)代硬件的優(yōu)化。
4. 與現(xiàn)有框架兼容Muon并非完全替代AdamW，而是與之互補(bǔ)。實(shí)際應(yīng)用中，Muon通常優(yōu)化網(wǎng)絡(luò)的隱藏層參數(shù)，而嵌入層（embedding）和分類器頭（head）仍由AdamW處理。這種“分工合作”的方式，讓Muon可以無縫集成到現(xiàn)有訓(xùn)練流程中。

如果Muon替代AdamW

假設(shè)Muon的潛力被充分驗(yàn)證，并逐步取代AdamW成為新的標(biāo)準(zhǔn)優(yōu)化器，AI領(lǐng)域可能會迎來以下重大變化：

1. 訓(xùn)練成本大幅降低Muon的高效性意味著，訓(xùn)練一個大規(guī)模模型所需的計(jì)算資源和時(shí)間將顯著減少。這不僅能節(jié)省數(shù)千萬RMB的算力成本，還能讓更多中小型研究團(tuán)隊(duì)參與到前沿AI研究中，降低技術(shù)門檻。
2. 模型規(guī)模的進(jìn)一步突破更高的訓(xùn)練效率，讓研究者有能力嘗試更大規(guī)模的模型。比如，現(xiàn)在訓(xùn)練一個1000億參數(shù)模型可能需要數(shù)月，而有了Muon，或許幾周就能完成。這將推動AI能力的上限不斷提升。
3. AI應(yīng)用加速落地更快的訓(xùn)練速度意味著模型從研究到應(yīng)用的時(shí)間縮短。無論是工業(yè)大模型（創(chuàng)新奇智在做的事情），還是其他行業(yè)大模型，AI技術(shù)的普及速度都可能因此加快。
4. 優(yōu)化器研究的復(fù)興Muon的成功可能會重新點(diǎn)燃對優(yōu)化器算法的興趣。過去幾年，AdamW幾乎“一統(tǒng)江湖”，新優(yōu)化器的研究相對沉寂。如果Muon證明了創(chuàng)新優(yōu)化器的價(jià)值，研究者可能會投入更多精力，探索其他潛在的突破。

OpenAI：Muon的未來舞臺？

2024年，Muon的開發(fā)者Jordan在個人博客中詳細(xì)介紹了Muon的設(shè)計(jì)理念和實(shí)驗(yàn)結(jié)果，隨后不久，他宣布加入OpenAI。這一動向并非巧合，很可能OpenAI看中了Muon的潛力。

月之暗面的實(shí)踐：Muon的驗(yàn)證

Moonshot AI（月之暗面）近期在Muon的基礎(chǔ)上進(jìn)行了大規(guī)模實(shí)踐，驗(yàn)證了它在真實(shí)場景中的潛力。他們不僅改進(jìn)了Muon，還訓(xùn)練了一個名為Moonlight的3B/16B參數(shù)Mixture-of-Expert（MoE）模型，用5.7萬億tokens的數(shù)據(jù)進(jìn)行了測試。

Moonshot AI的改進(jìn)

Moonshot AI發(fā)現(xiàn)，原始Muon在小規(guī)模任務(wù)上表現(xiàn)出色，但在超大規(guī)模訓(xùn)練中會遇到問題，比如模型權(quán)重增長過大，影響穩(wěn)定性。為此，他們提出了兩個關(guān)鍵改進(jìn)：

1. 引入權(quán)重衰減他們將AdamW的權(quán)重衰減機(jī)制融入Muon，更新公式變?yōu)椋?img src="https://s2.51cto.com/oss/202503/25/a23c96b252c95212734675cbb6f9aa77fc850b.png" alt='Muon優(yōu)化器：AI模型訓(xùn)練算法的下一個里程碑？| 目前還不是業(yè)界焦點(diǎn)，但有潛力是重大基礎(chǔ)創(chuàng)新-AI.x社區(qū)' title='Muon優(yōu)化器：AI模型訓(xùn)練算法的下一個里程碑？| 目前還不是業(yè)界焦點(diǎn)，但有潛力是重大基礎(chǔ)創(chuàng)新-AI.x社區(qū)' style="width: 484px; visibility: visible;" data-type="inline">這有效控制了權(quán)重大小，提升了長期訓(xùn)練的性能。
2. 調(diào)整更新尺度Muon的更新幅度（RMS）會因參數(shù)矩陣的形狀而變化，可能導(dǎo)致訓(xùn)練不穩(wěn)定。Moonshot AI提出按矩陣最大維度縮放更新，比如：

 這樣可以保持更新幅度一致，并與AdamW兼容。

Moonlight的驚艷表現(xiàn)

基于這些改進(jìn)，Moonshot AI用Muon訓(xùn)練了Moonlight模型，并在多個基準(zhǔn)測試中取得了優(yōu)異成績。例如：

• 在MMLU（英語理解）上，Moonlight得分70.0，超越了同規(guī)模的Llama3.2-3B（54.7）和Deepseek-v2-Lite（58.3）。
• 在GSM8K（數(shù)學(xué)推理）上，得分77.4，接近Qwen2.5-3B（79.1），但訓(xùn)練tokens僅為后者的三分之一。
• 訓(xùn)練效率上，Moonlight只需約52%的FLOPs，就能達(dá)到AdamW的性能水平。

這些結(jié)果表明，Muon不僅能加速訓(xùn)練，還能提升模型性能，尤其在數(shù)學(xué)和代碼任務(wù)上表現(xiàn)突出。

結(jié)語：Muon的潛力與未來

Muon作為一種新興的優(yōu)化器算法，以其獨(dú)特的設(shè)計(jì)和卓越的性能，展現(xiàn)了巨大的潛力。它通過正交化更新矩陣，打破了傳統(tǒng)優(yōu)化器的局限，在訓(xùn)練速度、可擴(kuò)展性和計(jì)算效率上都超越了AdamW。盡管目前它還未引起全球關(guān)注，但Keller Jordan的加入OpenAI，以及Moonshot AI的成功實(shí)踐，都預(yù)示著Muon可能成為AI訓(xùn)練領(lǐng)域的下一個里程碑。

? 圖片 ?

本文轉(zhuǎn)載自??后向傳播??，作者：張發(fā)恩