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

機(jī)器之心報(bào)道

機(jī)器之心編輯部

「許多博士（包括過去的我）都陷入了這樣一個(gè)誤區(qū)：認(rèn)為只有在頂級會(huì)議上發(fā)表論文才是終極目標(biāo)。」AI 云服務(wù)商 Hyperbolic CEO Yuchen Jin 如是說。

但現(xiàn)在，發(fā)表論文并不與學(xué)術(shù)影響力直接畫等號(hào)了。

Keller Jordan，OpenAI 深度學(xué)習(xí)團(tuán)隊(duì)主要成員之一，用一篇博客就撬開了 OpenAI 的大門。

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

在這篇博客中，Keller Jordan 提出并構(gòu)建了一種用于神經(jīng)網(wǎng)絡(luò)隱藏層的優(yōu)化器 Muon，其能夠在保證神經(jīng)網(wǎng)絡(luò)（包括 Transformer 和 CNN）的準(zhǔn)確度的前提上大幅提升其訓(xùn)練速度。

為何只發(fā)了博客，而不是發(fā)表一篇正式的 arXiv 論文，Keller Jordan 這樣解釋：能否發(fā)表一篇關(guān)于新優(yōu)化器的論文，且包含大量看起來不錯(cuò)的結(jié)果，和這個(gè)優(yōu)化器是否真的有效之間沒有任何聯(lián)系。「我只相信速通。」

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

接下來，我們看看這篇博客內(nèi)容。

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

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

• 原文地址：https://kellerjordan.github.io/posts/muon/
• Muon 的 PyTorch 實(shí)現(xiàn)：https://github.com/KellerJordan/Muon

Muon：一種用于神經(jīng)網(wǎng)絡(luò)隱藏層的優(yōu)化器

Muon 是一種用于神經(jīng)網(wǎng)絡(luò)隱藏層的優(yōu)化器，可用于快速運(yùn)行 NanoGPT 和 CIFAR-10，并創(chuàng)造了當(dāng)前最快的訓(xùn)練速度紀(jì)錄。

目前，人們已經(jīng)發(fā)布了很多使用 Muon 的實(shí)證研究結(jié)果，所以本文將主要關(guān)注 Muon 的設(shè)計(jì)。

首先，本文將定義 Muon，并概述其迄今為止取得的實(shí)證結(jié)果；然后將詳細(xì)討論其設(shè)計(jì)，包括與先前研究的聯(lián)系以及我們對其工作原理的最佳理解；最后將討論優(yōu)化研究中的證據(jù)標(biāo)準(zhǔn)。

定義

Muon 是一款用戶神經(jīng)網(wǎng)絡(luò)隱藏層的 2D 參數(shù)的優(yōu)化器。其定義如下：

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

使用 Muon 訓(xùn)練神經(jīng)網(wǎng)絡(luò)時(shí)，應(yīng)使用 AdamW 等標(biāo)準(zhǔn)方法優(yōu)化網(wǎng)絡(luò)的標(biāo)量和向量參數(shù)以及輸入層和輸出層。Muon 可用于處理 4D 卷積參數(shù)，方法是將其最后三個(gè)維度展平（如下所示）。

結(jié)果

Muon 已取得以下實(shí)證結(jié)果：

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

圖 1. 按樣本效率比較優(yōu)化器。

圖 2. 按時(shí)間比較優(yōu)化器。

此外，以下是 Muon 和 AdamW 在訓(xùn)練 1.5B 參數(shù)語言模型時(shí)的比較。兩個(gè)優(yōu)化器均已經(jīng)過微調(diào)。

圖 3. Muon 與 AdamW 在 1.5B 短訓(xùn)練上的比較。

Muon 的設(shè)計(jì)

Muon，全稱 MomentUm Orthogonalized by Newton-Schulz，其優(yōu)化 2D 神經(jīng)網(wǎng)絡(luò)參數(shù)的方式是：獲取 SGD - 動(dòng)量生成的更新，對每個(gè)更新應(yīng)用 Newton-Schulz (NS) 迭代作為后處理步驟，然后在將更新應(yīng)用于參數(shù)。

NS 迭代的作用是近似正交化更新矩陣，即應(yīng)用以下運(yùn)算：

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

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

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

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

而在實(shí)驗(yàn)驗(yàn)證中，我們基于人工檢查觀察到，SGD - 動(dòng)量和 Adam 對基于 Transformer 的神經(jīng)網(wǎng)絡(luò)中的 2D 參數(shù)產(chǎn)生的更新通常具有非常高的條件數(shù)。也就是說，它們幾乎都是低秩矩陣，所有神經(jīng)元的更新僅由少數(shù)幾個(gè)方向主導(dǎo)。我們推測，正交化會(huì)有效地增加了其他「稀有方向」的規(guī)模，這些方向在更新中幅度很小，但對學(xué)習(xí)仍然很重要。

消除 NS 迭代的替代方案

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

SVD（即計(jì)算更新的 USV? 分解，然后用 UV? 替換更新）易于理解，但我們不使用它，因?yàn)樗恕?/p>

耦合牛頓迭代法 (Coupled Newton iteration) 曾在 Shampoo 的實(shí)現(xiàn)中被用于執(zhí)行逆四次方根，并且可以被輕松地調(diào)整用于執(zhí)行正交化。但我們沒有使用它，因?yàn)槲覀儼l(fā)現(xiàn)它必須至少以 float32 精度運(yùn)行才能避免數(shù)值不穩(wěn)定，而這會(huì)導(dǎo)致它在現(xiàn)代 GPU 上運(yùn)行緩慢。

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

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

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

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

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

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

為了滿足當(dāng) N→∞ 時(shí) φ?(x)→1，會(huì)有一定的自由度，因?yàn)?(a,b,c) 有很多符合此性質(zhì)的可能選擇。稍后我們將優(yōu)化這個(gè)選擇，但現(xiàn)在可在下圖中看到，簡單的基線 (a,b,c)=(2,?1.5,0.5) 已經(jīng)有效。

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

調(diào)整系數(shù)

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

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

這里有一個(gè)令人驚訝的觀察結(jié)果：根據(jù)實(shí)際經(jīng)驗(yàn)，ε 可以高達(dá) 0.3 左右，而不會(huì)損害基于 Muon 的訓(xùn)練的損失曲線。因此，我們的目標(biāo)是最大化 a，使

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

圖 5. 調(diào)整后的牛頓 - 舒爾茨迭代系數(shù)。

在我們的實(shí)驗(yàn)中，當(dāng)使用這些系數(shù)的 Muon 來訓(xùn)練 Transformer 語言模型和小型卷積網(wǎng)絡(luò)時(shí)，只需運(yùn)行 5 步 NS 迭代即可。

我們也考慮過使用三階和七階多項(xiàng)式來進(jìn)行 NS 迭代，但發(fā)現(xiàn)這些方法無法進(jìn)一步改善時(shí)間開銷。

運(yùn)行時(shí)分析

本節(jié)將分析 Muon 的運(yùn)行時(shí)和內(nèi)存需求。

在應(yīng)用 NS 迭代之前，Muon 只是標(biāo)準(zhǔn)的 SGD 動(dòng)量，因此其內(nèi)存需求相同。

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

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

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

下面針對兩個(gè)具體的訓(xùn)練場景計(jì)算了此開銷：NanoGPT 速通和 Llama 405B 訓(xùn)練。

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

由此可以得出結(jié)論，對于典型的語言模型訓(xùn)練場景，無論規(guī)模大小，Muon 的 FLOP 開銷均低于 1%

與先前優(yōu)化器的關(guān)系

Shampoo 優(yōu)化器定義如下：

如果去除預(yù)調(diào)節(jié)器累積，則公式變?yōu)橐韵滦问剑?/p>

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

因此，可以將關(guān)閉動(dòng)量的 Muon 解讀為一種瞬時(shí)或無累積的 Shampoo 優(yōu)化器。

正交 - SGDM

Tuddenham 等人（2022）提出了一種優(yōu)化神經(jīng)網(wǎng)絡(luò)的方法：通過奇異值分解（SVD）對梯度進(jìn)行正交化，對其結(jié)果應(yīng)用動(dòng)量，再將動(dòng)量項(xiàng)作為更新。他們將該優(yōu)化器命名為正交 - SGDM（Orthogonal-SGDM）。這與 Muon 類似，但區(qū)別在于：

• Muon 將動(dòng)量計(jì)算置于正交化之前（實(shí)驗(yàn)表明該設(shè)計(jì)表現(xiàn)更優(yōu)）；
• Muon 采用牛頓 - 舒爾茨迭代代替 SVD，以實(shí)現(xiàn)更高效的正交化。

遺憾的是，Tuddenham 等人（2022）在其最佳實(shí)驗(yàn)配置（表 3）中報(bào)告，他們的方法表現(xiàn)不及精心調(diào)參的標(biāo)準(zhǔn) SGD-Momentum。

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

實(shí)驗(yàn)考量

根據(jù)設(shè)計(jì)，Muon 僅適用于 2D 參數(shù)（對于卷積濾波器則需展平處理），因此網(wǎng)絡(luò)中的其余標(biāo)量和向量參數(shù)仍需使用標(biāo)準(zhǔn)優(yōu)化方法（如 AdamW）。實(shí)驗(yàn)發(fā)現(xiàn)，輸入層和輸出層參數(shù)即使屬于 2D 結(jié)構(gòu)，也需使用 AdamW 優(yōu)化，這對性能至關(guān)重要。具體而言，在訓(xùn)練 Transformer 模型時(shí)，詞嵌入層（embedding）和最終的分類器頭（classifier head）應(yīng)使用 AdamW 才能達(dá)到最佳效果。

嵌入層的優(yōu)化動(dòng)態(tài)應(yīng)與其他層不同，這一結(jié)論符合模范數(shù)理論（modular norm theory）的預(yù)測；但輸出層的優(yōu)化動(dòng)態(tài)差異并未被該理論涵蓋，而是由實(shí)驗(yàn)觀測結(jié)果驅(qū)動(dòng)。

另一個(gè)純粹的經(jīng)驗(yàn)結(jié)果是，在本文測試的所有案例中，使用 Nesterov 式動(dòng)量對 Muon 的效果都比普通的 SGD 動(dòng)量略好。因此，本文在公開的 Muon 實(shí)現(xiàn)中將其設(shè)為默認(rèn)設(shè)置。

第三個(gè)結(jié)果是，如果將 Muon 分別應(yīng)用于 Transformer 的 Q、K、V 參數(shù)，而不是同時(shí)應(yīng)用于 Q、K、V 參數(shù)，則 Muon 在優(yōu)化 Transformer 方面效果更佳，因?yàn)?Transformer 實(shí)現(xiàn)會(huì)將 Q、K、V 參數(shù)設(shè)置為單個(gè)線性層，并將輸出拆分。

當(dāng)前的神經(jīng)網(wǎng)絡(luò)優(yōu)化研究文獻(xiàn)中，充斥著大量宣稱「大幅超越 AdamW」卻最終被社區(qū)棄用的優(yōu)化器。坦白說，這種現(xiàn)象并不令人意外。考慮到行業(yè)每年投入數(shù)十億美元用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練（且迫切希望降低成本），如果這些優(yōu)化器真的有效，理應(yīng)被廣泛采用。因此，問題顯然出在研究層面而非應(yīng)用層面 —— 即現(xiàn)有研究存在系統(tǒng)性缺陷。

通過仔細(xì)分析相關(guān)論文可以發(fā)現(xiàn)，最常見的癥結(jié)在于基線模型（baseline）調(diào)優(yōu)不足：許多研究在將新提出的優(yōu)化器與 AdamW 對比時(shí)，未能對 AdamW 基線進(jìn)行充分調(diào)參。

發(fā)表聲稱有巨大改進(jìn)但無法復(fù)制 / 達(dá)到宣傳效果的新方法并非無害犯罪，因?yàn)樗速M(fèi)了大量個(gè)體研究人員和小型實(shí)驗(yàn)室的時(shí)間、金錢和士氣，他們每天都在為復(fù)制和改進(jìn)這些方法的失敗而感到失望。

為了糾正這種情況，我們應(yīng)該采用以下標(biāo)準(zhǔn)：研究社區(qū)應(yīng)該要求，新的神經(jīng)網(wǎng)絡(luò)訓(xùn)練方法應(yīng)該在競爭性訓(xùn)練任務(wù)中表現(xiàn)出色。

競爭性任務(wù)通過兩種方式解決基線調(diào)優(yōu)不足的問題。

首先，競爭性任務(wù)中的基線是先前的最佳記錄，如果該任務(wù)很熱門，這個(gè)基線很可能已經(jīng)被充分調(diào)優(yōu)。

其次，即使在不太可能的情況下先前記錄未被充分調(diào)優(yōu)，系統(tǒng)也可以通過恢復(fù)到標(biāo)準(zhǔn)訓(xùn)練方法的新記錄實(shí)現(xiàn)自我修正。這種自我修正之所以可行，是因?yàn)闃?biāo)準(zhǔn)方法通常具有經(jīng)過硬件優(yōu)化的高效實(shí)現(xiàn)，而新方法通常會(huì)引入額外的計(jì)算時(shí)間開銷。這樣一來，在熱門的競爭性任務(wù)中，標(biāo)準(zhǔn)方法出現(xiàn)重大但虛假的改進(jìn)并長期保持在記錄歷史中的可能性就很小了。

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

• Muon 能否擴(kuò)展到更大規(guī)模的訓(xùn)練？
• Muon 使用的 Newton-Schulz 迭代能否在大規(guī)模 GPU 集群中合理分布？
• Muon 是否可能僅適用于預(yù)訓(xùn)練，而無法用于微調(diào)或強(qiáng)化學(xué)習(xí)工作負(fù)載？

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

不過，已經(jīng)有研究基于 Muon 優(yōu)化器進(jìn)行了改進(jìn)，比如月之暗面在 Muon 中引入了標(biāo)準(zhǔn)的 AdamW（Loshchilov 等人，2019）權(quán)重衰減機(jī)制。結(jié)果表明，帶權(quán)重衰減的 Muon 優(yōu)于原始 Muon 和 AdamW，獲得了更低的驗(yàn)證損失。

另外，雖然 Muon 誕生于一篇博客，但也已經(jīng)有研究團(tuán)隊(duì) Essential AI 發(fā)布了對該優(yōu)化器的系統(tǒng)性研究論文。感興趣的讀者可擴(kuò)展閱讀：

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

該論文表明，在計(jì)算 - 時(shí)間權(quán)衡方面，Muon 比 AdamW 更能顯著擴(kuò)展帕累托邊界。他們發(fā)現(xiàn)，Muon 在保持大批量（遠(yuǎn)超所謂的臨界批量）數(shù)據(jù)效率的同時(shí)，計(jì)算效率也更高，從而能夠?qū)崿F(xiàn)更經(jīng)濟(jì)的訓(xùn)練。

你嘗試過 Muon 嗎？