剛剛，OpenAI前CTO Mira Murati創(chuàng)辦的Thinking Machines Lab再次發(fā)布成果！

這是他們繼《克服LLM推理中的不確定性》（Defeating Nondeterminism in LLM Inference）之后，第二篇研究文章——《模塊流形》（Modular Manifolds）。

博客地址：https://thinkingmachines.ai/blog/modular-manifolds/

訓(xùn)練大型神經(jīng)網(wǎng)絡(luò)如同「走鋼絲」，必須小心翼翼地維持其內(nèi)部「健康」，防止權(quán)重、激活值或梯度這些關(guān)鍵張量變得過大或過小，以免引發(fā)數(shù)值溢出等一系列問題。

其中一個(gè)重要的思路，是為大模型提供一套統(tǒng)一的量級(jí)管理。

首先是穩(wěn)住基本盤。

使用Layer Norm技術(shù)把每層的輸出拉回合適范圍，對(duì)激活向量進(jìn)行歸一化（normalization），這也是目前一種普遍的做法。

對(duì)梯度更新進(jìn)行歸一化也很常見，例如Muon優(yōu)化器對(duì)更新進(jìn)行譜歸一化處理，使每一步更新的幅度可控。

再進(jìn)一步，是直接「管住」權(quán)重本體。

歸一化權(quán)重矩陣是一個(gè)值得嘗試的方向。

文中提出了一種重新思考優(yōu)化算法提供了新視角：將權(quán)重張量約束在某個(gè)子流形（submanifold）上，以便與這些流形約束協(xié)同設(shè)計(jì)優(yōu)化算法。

這好比把「救火」變「預(yù)防」：

一開始就把參數(shù)放在健康區(qū)間，讓訓(xùn)練更穩(wěn)、更具解釋性，從而使大模型可以更穩(wěn)定、高效地訓(xùn)練起來。

流形優(yōu)化器的形態(tài)

我們知道，流形只是一個(gè)局部看起來很平坦的曲面。

如果放大到足夠多，它看起來就像是一個(gè)普通平面。

流形上某一點(diǎn)附近的局部平坦空間稱為「切空間」（tangent space）。

如圖1所示，三維球面或更高維度的超球面是一個(gè)流形，圖中以紅色部分表示其在某點(diǎn)的切平面。

為了讓權(quán)重能夠「待在」指定的流形里，一個(gè)簡(jiǎn)單的方法是使用普通優(yōu)化器，在每步更新后將權(quán)重投影回流形。

但問題是如果優(yōu)化步驟偏離流形太多，再被強(qiáng)制投影回來，這會(huì)導(dǎo)致名義學(xué)習(xí)率不再對(duì)應(yīng)參數(shù)在流形上的實(shí)際位移，從而削弱我們對(duì)「步長(zhǎng)—效果」關(guān)系的直覺。

想在流形上認(rèn)真設(shè)計(jì)訓(xùn)練算法，必須先想清楚：在切空間里怎么度量「距離」？

一個(gè)解決思路是直接在切空間中進(jìn)行優(yōu)化。這樣，每一步都是沿著流形「表面」走，學(xué)習(xí)率能更好地對(duì)應(yīng)「實(shí)際位移」。

常見的選擇是歐幾里得距離，但也可以選擇以其他方式測(cè)量距離，如圖2所示。

值得注意的是，距離度量方式的選擇會(huì)直接影響最優(yōu)優(yōu)化步驟的方向。

圖3中，粉色箭頭表示原始梯度——即損失函數(shù)對(duì)權(quán)重的偏導(dǎo)數(shù)（partial derivative）。

也就是說，我們不一定非要嚴(yán)格按照梯度方向移動(dòng)。

為了用數(shù)學(xué)表達(dá)這個(gè)過程，我們可以把「在流形約束和特定距離度量下的最優(yōu)更新方向」看作一個(gè)帶約束的優(yōu)化問題，可以用一個(gè)搭配歐幾里得范數(shù)的超球面來舉例。

用g表示梯度， w表示超球面上的當(dāng)前點(diǎn)， a表示更新方向， η表示學(xué)習(xí)率，我們需要解決的問題是：

再回到圖 1、2 和3所展示的可視化語言，這個(gè)公式的意思是：綠色箭頭（也就是a的最優(yōu)解）必須同時(shí)滿足兩個(gè)條件：

一是它要落在紅色的切平面上，二是它必須在半徑為η的黃色圓圈上。

我們可以應(yīng)用拉格朗日乘數(shù)法來求解。

其中λ和μ是拉格朗日乘子。

對(duì)這個(gè)拉格朗日函數(shù)對(duì)a求導(dǎo)并令其為零，然后結(jié)合兩個(gè)約束條件求解λ和μ，就可以得到最優(yōu)更新方向。

簡(jiǎn)單來說最優(yōu)更新的做法是：先從梯度中減去與w同方向的徑向分量，即把梯度投影到切空間上，然后將結(jié)果歸一化，再乘以學(xué)習(xí)率。

這樣得到的更新方向就在切空間里了。

圖4中顯示這個(gè)微小的修正過程被稱為「回縮映射」（retraction map）。

完整的流形優(yōu)化算法如下：

總結(jié)來說，一階流形優(yōu)化器包含三個(gè)步驟：

• 找到一個(gè)單位長(zhǎng)度的切向量，在梯度方向上盡可能遠(yuǎn)；
• 用學(xué)習(xí)率乘以這個(gè)方向，然后從當(dāng)前權(quán)重中減去；
• 把更新后的權(quán)重通過回縮映射拉回流形上。

在執(zhí)行這一流程時(shí)，我們需要決定選擇什么樣的流形來作為約束，此外是如何定義「長(zhǎng)度」的度量方式。

根據(jù)這兩個(gè)選擇的不同，我們就能得到不同的優(yōu)化算法，具體見下表。

流形Muon

Transformer中的典型權(quán)重矩陣W是一個(gè)「向量變換器」，即它將輸入向量x
轉(zhuǎn)換為輸出向量y=Wx。

我們希望設(shè)計(jì)一種流形約束和距離函數(shù)，使得該矩陣對(duì)輸入向量的作用合理：既不應(yīng)導(dǎo)致輸出值過大或過小，也不應(yīng)在更新權(quán)重時(shí)引起輸出向量劇烈變化或幾乎無變化。

一個(gè)思考矩陣如何作用于向量的好方法是使用奇異值分解（SVD），如圖 5 所示。

SVD以分解矩陣的方式顯示矩陣如何沿著不同的軸拉伸輸入向量。

我們希望矩陣的「拉伸效應(yīng)」接近于1，因此選擇了一個(gè)所有奇異值均為1的矩陣流形。

這種矩陣流形在數(shù)學(xué)上被稱為Stiefel流形，在高矩陣（ m≥n）的假設(shè)下，它可以等價(jià)地定義為以下集合：

要為Stiefel流形設(shè)計(jì)優(yōu)化器，還需選擇一個(gè)合適的距離函數(shù)。

為限制權(quán)重更新對(duì)輸入向量的最大拉伸作用，譜范數(shù)（spectral norm），即矩陣最大奇異值的度量是一個(gè)合適的選項(xiàng)。

雖然它只約束了最大效應(yīng)，但由于優(yōu)化器會(huì)飽和這一上限，因此也能間接防止最小效應(yīng)過小。

正是這一想法，促成了Muon優(yōu)化器的提出。

這一想法與Stiefel流形約束結(jié)合后，就形成了「manifold Muon」問題。

文中的一個(gè)關(guān)鍵發(fā)現(xiàn)是一個(gè)凸優(yōu)化問題，可以通過標(biāo)準(zhǔn)方法——對(duì)偶上升法（dual ascent）來求解。

經(jīng)過推導(dǎo)，對(duì)偶函數(shù)的梯度為：

通過一個(gè)小實(shí)驗(yàn)，可以驗(yàn)證算法的可行性，實(shí)驗(yàn)設(shè)置與結(jié)果見圖6。

模塊流形

這里還有一個(gè)重要的問題：當(dāng)我們將多個(gè)層組合起來構(gòu)建完整的神經(jīng)網(wǎng)絡(luò)時(shí)，會(huì)發(fā)生什么？

是否需要關(guān)注層與層之間的交互，并據(jù)此修改優(yōu)化策略？

這需要一種可以將前文介紹的推導(dǎo)邏輯推廣到整個(gè)神經(jīng)網(wǎng)絡(luò)的方法——模塊流形（modular manifolds）理論。

該理論的核心思想是：構(gòu)建一種抽象機(jī)制，用來指導(dǎo)如何在各層之間合理分配學(xué)習(xí)率。

在本質(zhì)上，在不同層之間分配學(xué)習(xí)率，或者對(duì)單個(gè)層進(jìn)行縮放，都依賴于我們對(duì)網(wǎng)絡(luò)輸出對(duì)權(quán)重的Lipschitz敏感性的理解。

我們?cè)诖罱ňW(wǎng)絡(luò)的過程中會(huì)追蹤這種敏感性，而流形約束有助于我們更加精準(zhǔn)地把握它。