智東西
編譯 李水青
編輯 心緣

智東西6月16日消息，今天，AI云服務商Hyperbolic的聯合創始人兼CTO Yuchen Jin在社交平臺X上曝料：研究員Keller Jordan僅憑一篇博客文章就加入了OpenAI，并可能正用博客提及的神經網絡隱藏層的優化器Muon訓練GPT-5。

“許多博士（包括以前的我）都陷入了這樣一個誤區：認為在頂級會議上發表論文才是最終目標。但發表論文≠影響力。Muon只作為一篇博客文章發布，它讓Keller加入了OpenAI，他現在可能正在用它訓練GPT-5。”Yuchen Jin說。

▲Yuchen Jin的X推文及Yuchen Jin的自述

Yuchen Jin提及的這篇博客發布于2024年12月，題為《Muon：神經網絡隱藏層的優化器（Muon: An optimizer for hidden layers in neural networks）》。

博客地址：https://kellerjordan.github.io/posts/muon/

從職場社交平臺領英可知，Keller Jordan正是在2024年12月加入OpenAI，由此我們也可以推測他正是憑去年12月發布的一篇博客，成功進入了如日中天的頭部大模型企業。

這篇博客厲害在那兒？Muon憑什么成為OpenAI的敲門磚？讓我們從這篇博客文章內容說起。

一、Muon定義：一個神經網絡隱藏層的優化器

Muon是神經網絡隱藏層的優化器。它被用于NanoGPT和CIFAR-10的快速運行，刷新了當時訓練速度的記錄。

Keller Jordan的博客文章主要關注Muon的設計。首先他定義了Muon并概述其在當時已取得的實證結果；然后他詳細討論了Muon的設計，包括與先前研究的聯系以及對其工作原理的最佳理解；最后他討論了優化研究中的證據標準。

具體來說，Muon是一個針對神經網絡隱藏層二維參數的優化器，其定義如下：

其中“NewtonSchulz5”定義為以下Newton-Schulz矩陣迭代：

使用Muon訓練神經網絡時，應使用AdamW等標準方法優化網絡的標量和矢量參數以及輸入層和輸出層。Muon可用于四維卷積參數，方法是將其最后三個維度展平。

Muon取得了以下實證成果：

1、將CIFAR-10上的訓練速度記錄提高到94%準確率，從3.3秒提高到2.6秒。

2、將FineWeb（一項稱為NanoGPT快速運行的競賽任務）上的訓練速度記錄提高至3.28 val loss，提高了1.35倍。

3、在擴展到774M和1.5B參數的同時，繼續顯示訓練速度的提升。

4、在HellaSwag上用10個8xH100小時訓練了一個1.5B參數轉換器，使其達到GPT-2 XL級別的性能。使用AdamW達到相同結果則需要13.3小時。

以下是針對NanoGPT快速運行的不同強力優化器的比較：

▲按樣本效率比較優化器（可復現日志：https://github.com/KellerJordan/modded-nanogpt/tree/master/records/102924_Optimizers）

▲按掛鐘時間比較優化器

此外，以下是Muon和AdamW在訓練15億參數語言模型時的對比。兩個優化器均已進行調整。

▲Muon與AdamW在15億參數短時間訓練中的對比（可復現日志：https://github.com/KellerJordan/modded-nanogpt/tree/master/records/102024_ScaleUp1B）

二、Muon設計：牛頓-舒爾茨迭代法作為后處理步驟

Muon通過采用SGD-momentum生成的更新來優化二維神經網絡參數，然后在將它們應用于參數之前，對每個更新應用 Newton-Schulz （牛頓-舒爾茨迭代法，簡稱NS）迭代作為后處理步驟。

NS迭代的作用是使更新矩陣近似正交化，即應用下列操作：

換句話說，NS迭代實際上用最接近的半正交矩陣替換了SGD-momentum的更新矩陣。

為什么正交化更新可行？出于實證研究的動機，作者基于人工檢驗觀察到，SGD-momentum和Adam對基于Transformer的神經網絡中的二維參數產生的更新通常具有非常高的條件數。也就是說，它們幾乎是低秩矩陣，所有神經元的更新僅由少數幾個方向主導。

作者推測，正交化有效地增加了其他“稀有方向”的規模，這些方向在更新中幅度較小，但對學習仍然很重要。

除了NS迭代之外，還有其他幾種方法可以對矩陣進行正交化。但作者沒有使用其中兩種方法，他是如何排除的？

一個是SVD方法，它太慢了，所以作者沒有使用它。另一個是Coupled Newton iteration （耦合牛頓迭代法），它必須至少以float32精度運行才能避免數值不穩定，這導致它在現代GPU上運行速度較慢，所以作者也沒有采用。

相比之下，作者發現NS可以在bfloat16中穩定運行，因此選擇它們作為正交化更新的首選方法。

在Keller Jordan的實驗中，當使用具有調整系數的Muon來訓練Transformer語言模型和小型卷積網絡時，只需運行5步NS迭代就足夠了。

此外，Keller Jordan還分析了Muon的運行時間和內存要求。對于典型的語言訓練場景，無論規模大小，Muon的FLOP開銷都低于1%。

三、Muon實證考慮：批判糟糕的基線，提出新方法

根據設計，Muon僅適用于二維參數，以及通過展平的卷積濾波器，因此網絡中其余的標量和矢量參數必須使用標準方法（例如 AdamW）進行優化。

根據經驗，Keller Jordan發現使用AdamW優化輸入和輸出參數也很重要，即使這些參數通常是二維的。具體來說，在訓練Transformer時，應該將AdamW用于嵌入層和最終分類器頭層，以獲得最佳性能。嵌入層的優化動態應該與其他層不同，這遵循模塊化范數理論。輸出層的這種動態也不同，這似乎并非來自理論，而是由經驗驅動的。

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

第三個結果是，如果將Muon分別應用于變壓器的Q、K、V參數，而不是一起應用于變壓器，則Muon可以更好地優化變壓器，因為對于將QKV參數化為輸出被分割的單個線性層的變壓器實現，默認做法是將它們一起應用。

Keller Jordan認為，神經網絡優化研究文獻目前大多充斥著一堆已死的優化器，它們聲稱能夠擊敗AdamW，而且往往以巨大的優勢獲勝，但卻從未被社區采用。鑒于業界在神經網絡訓練上投入了數十億美元，并渴望降低成本，他們可以推斷，問題出在研究界，而非潛在的采用者。

Keller Jordan犀利地提出：這項研究出了問題。仔細研究每篇論文后，他們發現最常見的罪魁禍首是糟糕的基線：論文在將其與新提出的優化器進行比較之前，往往沒有充分調整AdamW基線。

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

為了糾正這種情況，Keller Jordan建議采用以下證據標準：研究界應該要求，只要有可能，神經網絡訓練的新方法就應該在競爭性訓練任務中取得成功。

競爭性任務通過兩種方式解決了基線欠調問題。首先，競爭性任務的基線是先前的記錄，如果該任務很受歡迎，則很可能已經經過了良好的調整。其次，即使在先前記錄未經過良好調整的不太可能發生的情況下，也可以通過新的記錄進行自我修正，將訓練恢復到標準方法。

結語：全新優化器或成為GPT-5中的重要技術

通過定義、拆解設計及實證研究，Keller Jordan發現了Muon神經網絡隱藏層的優化器具備優于AdamW的效率。通過最新曝料可知，這一技術很有可能成為OpenAI正在研究的GPT-5的重要部分。

Keller Jordan也提出了一些尚未解決的問題。包括：Muon可以擴展到更大規模的訓練嗎？是否有可能在大型GPU集群中正確分布Muon使用的Newton-Schulz迭代？Muon是否僅適用于預訓練，而不適用于微調或強化學習工作負載？或許在GPT-5的研究中，作者已經知道了這些問題的答案。

來源：X平臺、keller Jordan博客