在人工智能的世界里，如果說大型神經(jīng)網(wǎng)絡(luò)是那個(gè)能夠?qū)W習(xí)和思考的“大腦”，那么優(yōu)化器（Optimizer），就是指導(dǎo)這個(gè)“大腦”如何學(xué)習(xí)、學(xué)習(xí)多快、以及最終能學(xué)得多好的“隱形之手”。從一個(gè)隨機(jī)初始化的網(wǎng)絡(luò)，到一個(gè)能夠理解語言、生成代碼的強(qiáng)大智能體，優(yōu)化器在其中扮演著至關(guān)重要的角色。

然而，隨著模型規(guī)模的爆炸式增長，我們熟知的經(jīng)典優(yōu)化器正面臨著前所未有的挑戰(zhàn)。近期，一家由前OpenAI核心團(tuán)隊(duì)創(chuàng)立的新公司Thinking Machines Lab（TML），通過其發(fā)布的“模塊流形”研究，為我們揭示了AI優(yōu)化器可能即將迎來的一場深刻的范式革命。

從SGD到AdamW

深度學(xué)習(xí)的早期，最基礎(chǔ)的優(yōu)化器是隨機(jī)梯度下降（SGD）。它的原理很簡單，就像一個(gè)蒙著眼睛的登山者，每走一步都沿著當(dāng)前腳下最陡峭的方向（梯度的反方向）向下走，以期最快到達(dá)谷底（損失函數(shù)的最小值）。

但SGD的問題也很明顯，它容易陷入局部最優(yōu)的“小山谷”里出不來，或者在曲折的山路上反復(fù)震蕩。為了解決這些問題，研究者們引入了“慣性”的概念，催生了Momentum等優(yōu)化器。

隨后，一個(gè)里程碑式的優(yōu)化器Adam誕生了。它不僅考慮了慣性，還能為每個(gè)參數(shù)自適應(yīng)地調(diào)整學(xué)習(xí)率，使其在各種任務(wù)中都表現(xiàn)出色。而我們今天最廣泛使用的AdamW，則是在Adam的基礎(chǔ)上，改進(jìn)了權(quán)重衰減的處理方式，進(jìn)一步提升了其性能和泛化能力。可以說，AdamW是當(dāng)前深度學(xué)習(xí)優(yōu)化器領(lǐng)域的“集大成者”。

新時(shí)代的挑戰(zhàn)與“救火式”修正

當(dāng)模型參數(shù)從百萬級(jí)躍升至千億、萬億級(jí)別時(shí)，AdamW等經(jīng)典優(yōu)化器開始面臨新的挑戰(zhàn)——數(shù)值不穩(wěn)定性。在巨大的網(wǎng)絡(luò)中，參數(shù)的更新很容易出現(xiàn)“爆炸”或“消失”的極端情況，導(dǎo)致訓(xùn)練過程崩潰。

為了應(yīng)對(duì)這個(gè)問題，研究者們發(fā)明了各種“打補(bǔ)丁”式的修正方法。例如，Layer Normalization通過對(duì)每一層網(wǎng)絡(luò)的激活值進(jìn)行歸一化，來保持其穩(wěn)定。譜范數(shù)歸一化等技術(shù)，則通過限制權(quán)重矩陣的“拉伸”程度來控制更新的幅度。這些方法雖然有效，但都屬于“救火式”的被動(dòng)修正。

TML的“流形優(yōu)化”

TML發(fā)布的“模塊流形”研究，提出了一種全新的、從根本上解決問題的思路：不再是在一個(gè)無約束的空間里自由下降，然后再想辦法把跑偏的參數(shù)拉回來；而是一開始，就將參數(shù)約束在一個(gè)永遠(yuǎn)不會(huì)“跑偏”的、被稱為“流形”的特定數(shù)學(xué)空間內(nèi)進(jìn)行優(yōu)化。

為了驗(yàn)證這一理念，他們設(shè)計(jì)了一款名為**“流形Muon”的新型優(yōu)化器。它將神經(jīng)網(wǎng)絡(luò)中的權(quán)重矩陣，約束在一個(gè)名為Stiefel流形**的特殊空間中。在這個(gè)空間里，所有矩陣的奇異值（可以理解為矩陣對(duì)數(shù)據(jù)進(jìn)行拉伸或壓縮的程度）都被嚴(yán)格限制為1。這意味著，無論模型如何更新，其內(nèi)部的參數(shù)始終保持在一個(gè)“健康”、穩(wěn)定的狀態(tài)。

實(shí)驗(yàn)結(jié)果令人振奮。在經(jīng)典的CIFAR-10圖像分類任務(wù)上，使用“流形Muon”訓(xùn)練的神經(jīng)網(wǎng)絡(luò)，其性能優(yōu)于了被廣泛使用的AdamW優(yōu)化器。這初步證明了“流形優(yōu)化”這一新范式的有效性和巨大潛力。

優(yōu)化器的未來展望

“流形Muon”的成功，僅僅是這場優(yōu)化器革命的開端。TML的論文中，還展望了這一新范式未來的多個(gè)研究方向，包括：

• 如何為注意力機(jī)制、嵌入層等更復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)設(shè)計(jì)合適的流形約束。
• 流形約束對(duì)低精度訓(xùn)練的潛在好處。
• 如何設(shè)計(jì)和調(diào)整約束，以提高模型的泛化能力。
• 如何在GPU上高效地實(shí)現(xiàn)相關(guān)的流形數(shù)學(xué)運(yùn)算。

從SGD的簡單梯度下降，到AdamW的自適應(yīng)學(xué)習(xí)，再到如今“流形Muon”所揭示的幾何優(yōu)化新大陸，AI優(yōu)化器的發(fā)展，正不斷向著更穩(wěn)定、更高效、也更接近數(shù)學(xué)本質(zhì)的方向演進(jìn)。

Thinking Machines Lab的這項(xiàng)研究，無疑為我們打開了一扇通往這個(gè)新世界的大門。它預(yù)示著，未來的AI大模型訓(xùn)練，可能將不再是一場充滿不確定性的“煉丹”，而是一門更加嚴(yán)謹(jǐn)、更加可控的科學(xué)。