Muon作者僅用一篇博客，就被OpenAI看中了

機器之心報道

機器之心編輯部

「許多博士（包括過去的我）都陷入了這樣一個誤區：認為只有在頂級會議上發表論文才是終極目標?！笰I 云服務商 Hyperbolic CEO Yuchen Jin 如是說。

但現在，發表論文并不與學術影響力直接畫等號了。

Keller Jordan，OpenAI 深度學習團隊主要成員之一，用一篇博客就撬開了 OpenAI 的大門。

這篇名為《Muon: An optimizer for hidden layers in neural networks》的博客發布于 2024 年 12 月，而 Keller Jordan 入職 OpenAI 的時間恰好也在此時。

在這篇博客中，Keller Jordan 提出并構建了一種用于神經網絡隱藏層的優化器 Muon，其能夠在保證神經網絡（包括 Transformer 和 CNN）的準確度的前提上大幅提升其訓練速度。

為何只發了博客，而不是發表一篇正式的 arXiv 論文，Keller Jordan 這樣解釋：能否發表一篇關于新優化器的論文，且包含大量看起來不錯的結果，和這個優化器是否真的有效之間沒有任何聯系?！肝抑幌嘈潘偻?。」

一直以來，研究界的衡量標準過度局限于論文發表，而 Keller Jordan 的案例告訴我們，如果你足夠優秀，一篇博客也能打開頂級 AI 科研機構的大門，甚至是 OpenAI。從中，我們也可以看出，OpenAI 在人才招攬方面更注重能力而非其他外在條件。

接下來，我們看看這篇博客內容。

注意，這篇博客發表于 2024 年 12 月 8 日，因此其中對前沿指標的描述可能會略有過時，比如 NanoGPT 速通結果就已經被多次刷新了，下面展示了 Keller Jordan 托管的 NanoGPT 速通的最新八條世界記錄，其中最新記錄是今年 5 月 25 日創造的，已達到驚人的 2.979 分鐘！當然，如此成績不只靠 Muon，還有 FlexAttention、嵌入優化、架構優化等諸多改進。

https://github.com/KellerJordan/modded-nanogpt

• 原文地址：https://kellerjordan.github.io/posts/muon/
• Muon 的 PyTorch 實現：https://github.com/KellerJordan/Muon

Muon：一種用于神經網絡隱藏層的優化器

Muon 是一種用于神經網絡隱藏層的優化器，可用于快速運行 NanoGPT 和 CIFAR-10，并創造了當前最快的訓練速度紀錄。

目前，人們已經發布了很多使用 Muon 的實證研究結果，所以本文將主要關注 Muon 的設計。

首先，本文將定義 Muon，并概述其迄今為止取得的實證結果；然后將詳細討論其設計，包括與先前研究的聯系以及我們對其工作原理的最佳理解；最后將討論優化研究中的證據標準。

定義

Muon 是一款用戶神經網絡隱藏層的 2D 參數的優化器。其定義如下：

其中，NewtonSchulz5 定義為以下牛頓 - 舒爾茨矩陣迭代：

使用 Muon 訓練神經網絡時，應使用 AdamW 等標準方法優化網絡的標量和向量參數以及輸入層和輸出層。Muon 可用于處理 4D 卷積參數，方法是將其最后三個維度展平（如下所示）。

結果

Muon 已取得以下實證結果：

• 在 CIFAR-10 數據集上，在保證 94% 準確度的前提下，基于 A100 GPU，將訓練速度記錄從 3.3 秒提升至 2.6 秒。
• 在 FineWeb（一項被稱為 NanoGPT 競速的競賽任務）上的訓練速度記錄提升至 3.28 驗證損失，提升了 1.35 倍。
• 在擴展到 774M 個參數和 1.5B 億個參數時，繼續表現出了訓練速度的提升。
• 在 HellaSwag 上用 10 個 8xH100 小時將一個 1.5B 參數的 Transformer 訓練到 GPT-2 XL 級別性能。而使用 AdamW 達到相同結果需要 13.3 小時。

圖 1. 按樣本效率比較優化器。

圖 2. 按時間比較優化器。

此外，以下是 Muon 和 AdamW 在訓練 1.5B 參數語言模型時的比較。兩個優化器均已經過微調。

圖 3. Muon 與 AdamW 在 1.5B 短訓練上的比較。

Muon 的設計

Muon，全稱 MomentUm Orthogonalized by Newton-Schulz，其優化 2D 神經網絡參數的方式是：獲取 SGD - 動量生成的更新，對每個更新應用 Newton-Schulz (NS) 迭代作為后處理步驟，然后在將更新應用于參數。

NS 迭代的作用是近似正交化更新矩陣，即應用以下運算：

也就是說，NS 迭代實際上會將 SGD - 動量的更新矩陣替換為與其最接近的半正交矩陣。這相當于用 UV? 替換更新，其中 USV? 是其奇異值分解 (SVD)。

為什么正交化更新是有益的？

我們首先想指出一個有效的答案是：這樣就沒問題了嗎？

但是，對于源自 Bernstein & Newhouse (2024) 對 Shampoo (Gupta et al. 2018) 分析的理論動機，請參閱后文。

而在實驗驗證中，我們基于人工檢查觀察到，SGD - 動量和 Adam 對基于 Transformer 的神經網絡中的 2D 參數產生的更新通常具有非常高的條件數。也就是說，它們幾乎都是低秩矩陣，所有神經元的更新僅由少數幾個方向主導。我們推測，正交化會有效地增加了其他「稀有方向」的規模，這些方向在更新中幅度很小，但對學習仍然很重要。

消除 NS 迭代的替代方案

除了 NS 迭代之外，還有其他幾種矩陣正交化的選項。本小節將解釋為什么沒有使用其中兩種方法。請參閱 Bernstein & Newhouse (2024) 的附錄 A，獲取更完整的可用方法列表。

SVD（即計算更新的 USV? 分解，然后用 UV? 替換更新）易于理解，但我們不使用它，因為它太慢了。

耦合牛頓迭代法 (Coupled Newton iteration) 曾在 Shampoo 的實現中被用于執行逆四次方根，并且可以被輕松地調整用于執行正交化。但我們沒有使用它，因為我們發現它必須至少以 float32 精度運行才能避免數值不穩定，而這會導致它在現代 GPU 上運行緩慢。

相比之下，我們發現牛頓 - 舒爾茨迭代可以在 bfloat16 精度下穩定運行。因此，我們選擇它們作為正交化更新的首選方法。

證明 NS 迭代能夠讓更新正交化

為了理解 NS 迭代使更新正交化的原因，令 G=USV? 為 SGD - 動量生成的更新矩陣的 SVD。然后，對系數 (a,b,c) 運行一步 NS 迭代，輸出結果如下：

一般來說，如果定義五次多項式 φ(x)=ax+bx3+cx?，那么對系數 (a,b,c) 進行 N 步 NS 迭代會輸出結果 Uφ?(S) V?，其中 φ?(S) 表示對構成 S 對角線的奇異值 N 次逐元素應用 φ。

因此，為了保證 NS 迭代收斂于 Ortho (G)=UV?，需要做的就是 (1) 確保 S 的初始元素在 [0,1] 范圍內；(2) 選擇系數，使得當 N→∞ 時，φ?(x)→1。

為了滿足第一個條件，只需在開始 NS 迭代之前將 G 替換為 G/‖G‖F。這種重新縮放是有益的，因為 Ortho (cG)=Ortho (G)。

為了滿足當 N→∞ 時 φ?(x)→1，會有一定的自由度，因為 (a,b,c) 有很多符合此性質的可能選擇。稍后我們將優化這個選擇，但現在可在下圖中看到，簡單的基線 (a,b,c)=(2,?1.5,0.5) 已經有效。

圖 4. 牛頓 - 舒爾茨迭代的基線系數。

調整系數

雖然 NS 系數 (a,b,c)=(2,?1.5,0.5) 已經能夠完美地實現更新的正交化，但我們可以進一步調整它們，以減少需要運行的 NS 迭代步數。

1. 我們希望 a 盡可能大，因為 φ′(0)=a 意味著該系數控制著初始奇異值較小時的收斂速度。
2. 對于每個 x∈[0,1]，我們希望 φ?(x) 在 N→∞ 時收斂到 [1?ε,1+ε] 范圍內的一個值，使得 NS 迭代的結果與 Ortho (G) 相差不大。

這里有一個令人驚訝的觀察結果：根據實際經驗，ε 可以高達 0.3 左右，而不會損害基于 Muon 的訓練的損失曲線。因此，我們的目標是最大化 a，使

有很多方法可以解決這個約束優化問題。而這里使用一種基于梯度的臨時方法，最終得到系數 (3.4445,4.7750,2.0315)，這也是最終設計 Muon 時所使用的稀疏。這些系數的變化如下圖所示。請注意 x=0 附近的陡然增長。

圖 5. 調整后的牛頓 - 舒爾茨迭代系數。

在我們的實驗中，當使用這些系數的 Muon 來訓練 Transformer 語言模型和小型卷積網絡時，只需運行 5 步 NS 迭代即可。

我們也考慮過使用三階和七階多項式來進行 NS 迭代，但發現這些方法無法進一步改善時間開銷。

運行時分析

本節將分析 Muon 的運行時和內存需求。

在應用 NS 迭代之前，Muon 只是標準的 SGD 動量，因此其內存需求相同。

對于網絡中的每個 n×m 矩陣參數（例如，設 m≤n），NS 迭代的每一步都需要 2 (2nm2+m3) 個矩陣乘法 FLOP，對于平方參數，最多為 6nm2。因此，與 SGD 相比，Muon 所需的額外 FLOP 最多為 6Tnm2，其中 T 是 NS 迭代次數（通常我們取 T=5）。

如果該參數參數化了一個線性層，那么執行一個訓練步驟（即前向和后向傳遞）所需的 FLOP 基準量為 6nm2，其中 B 是該步驟中通過該層的輸入數量。

因此，Muon 的 FLOP 開銷最多為 Tm/B，其中 m 為模型維度，B 為以 token 為單位的批量大小，T 為 NS 迭代步數（通常 T=5）

下面針對兩個具體的訓練場景計算了此開銷：NanoGPT 速通和 Llama 405B 訓練。

1. 對于當前的 NanoGPT 速通記錄，模型維度為 m=768，每批次的 token 數量為 B=524288。因此，開銷為 5?768/524288=0.7%。
2. 對于 Llama 405B 訓練，模型維度為 m=16384，每批次的 token 數量為 B=16000000（Dubey et al. 2024）。因此，使用 Muon 進行此訓練的開銷為 5?16384/16000000=0.5%。

由此可以得出結論，對于典型的語言模型訓練場景，無論規模大小，Muon 的 FLOP 開銷均低于 1%

與先前優化器的關系

Shampoo 優化器定義如下：

如果去除預調節器累積，則公式變為以下形式：

這就是正交化梯度。如果我們在正交化之前添加動量，就能恢復 Muon 更新，但由于使用了四次方根求逆而不是牛頓 - 舒爾茨迭代，因此時間和 FLOP 開銷會更高。

因此，可以將關閉動量的 Muon 解讀為一種瞬時或無累積的 Shampoo 優化器。

正交 - SGDM

Tuddenham 等人（2022）提出了一種優化神經網絡的方法：通過奇異值分解（SVD）對梯度進行正交化，對其結果應用動量，再將動量項作為更新。他們將該優化器命名為正交 - SGDM（Orthogonal-SGDM）。這與 Muon 類似，但區別在于：

• Muon 將動量計算置于正交化之前（實驗表明該設計表現更優）；
• Muon 采用牛頓 - 舒爾茨迭代代替 SVD，以實現更高效的正交化。

遺憾的是，Tuddenham 等人（2022）在其最佳實驗配置（表 3）中報告，他們的方法表現不及精心調參的標準 SGD-Momentum。

圖源：https://arxiv.org/pdf/2202.07052

實驗考量

根據設計，Muon 僅適用于 2D 參數（對于卷積濾波器則需展平處理），因此網絡中的其余標量和向量參數仍需使用標準優化方法（如 AdamW）。實驗發現，輸入層和輸出層參數即使屬于 2D 結構，也需使用 AdamW 優化，這對性能至關重要。具體而言，在訓練 Transformer 模型時，詞嵌入層（embedding）和最終的分類器頭（classifier head）應使用 AdamW 才能達到最佳效果。

嵌入層的優化動態應與其他層不同，這一結論符合模范數理論（modular norm theory）的預測；但輸出層的優化動態差異并未被該理論涵蓋，而是由實驗觀測結果驅動。

另一個純粹的經驗結果是，在本文測試的所有案例中，使用 Nesterov 式動量對 Muon 的效果都比普通的 SGD 動量略好。因此，本文在公開的 Muon 實現中將其設為默認設置。

第三個結果是，如果將 Muon 分別應用于 Transformer 的 Q、K、V 參數，而不是同時應用于 Q、K、V 參數，則 Muon 在優化 Transformer 方面效果更佳，因為 Transformer 實現會將 Q、K、V 參數設置為單個線性層，并將輸出拆分。

當前的神經網絡優化研究文獻中，充斥著大量宣稱「大幅超越 AdamW」卻最終被社區棄用的優化器。坦白說，這種現象并不令人意外?？紤]到行業每年投入數十億美元用于神經網絡訓練（且迫切希望降低成本），如果這些優化器真的有效，理應被廣泛采用。因此，問題顯然出在研究層面而非應用層面 —— 即現有研究存在系統性缺陷。

通過仔細分析相關論文可以發現，最常見的癥結在于基線模型（baseline）調優不足：許多研究在將新提出的優化器與 AdamW 對比時，未能對 AdamW 基線進行充分調參。

發表聲稱有巨大改進但無法復制 / 達到宣傳效果的新方法并非無害犯罪，因為它浪費了大量個體研究人員和小型實驗室的時間、金錢和士氣，他們每天都在為復制和改進這些方法的失敗而感到失望。

為了糾正這種情況，我們應該采用以下標準：研究社區應該要求，新的神經網絡訓練方法應該在競爭性訓練任務中表現出色。

競爭性任務通過兩種方式解決基線調優不足的問題。

首先，競爭性任務中的基線是先前的最佳記錄，如果該任務很熱門，這個基線很可能已經被充分調優。

其次，即使在不太可能的情況下先前記錄未被充分調優，系統也可以通過恢復到標準訓練方法的新記錄實現自我修正。這種自我修正之所以可行，是因為標準方法通常具有經過硬件優化的高效實現，而新方法通常會引入額外的計算時間開銷。這樣一來，在熱門的競爭性任務中，標準方法出現重大但虛假的改進并長期保持在記錄歷史中的可能性就很小了。

這篇博客還列舉了一些待解決問題：

• Muon 能否擴展到更大規模的訓練？
• Muon 使用的 Newton-Schulz 迭代能否在大規模 GPU 集群中合理分布？
• Muon 是否可能僅適用于預訓練，而無法用于微調或強化學習工作負載？

在撰寫本文時，Keller Jordan 還不知道這些問題的答案。

不過，已經有研究基于 Muon 優化器進行了改進，比如月之暗面在 Muon 中引入了標準的 AdamW（Loshchilov 等人，2019）權重衰減機制。結果表明，帶權重衰減的 Muon 優于原始 Muon 和 AdamW，獲得了更低的驗證損失。

另外，雖然 Muon 誕生于一篇博客，但也已經有研究團隊 Essential AI 發布了對該優化器的系統性研究論文。感興趣的讀者可擴展閱讀：

• 論文標題：Practical Efficiency of Muon for Pretraining
• 論文地址：https://arxiv.org/pdf/2505.02222

該論文表明，在計算 - 時間權衡方面，Muon 比 AdamW 更能顯著擴展帕累托邊界。他們發現，Muon 在保持大批量（遠超所謂的臨界批量）數據效率的同時，計算效率也更高，從而能夠實現更經濟的訓練。

你嘗試過 Muon 嗎？