新智元報道

編輯：元宇

【新智元導(dǎo)讀】在大模型訓(xùn)練時，如何管理權(quán)重、避免數(shù)值爆炸與丟失？Thinking Machines Lab 的新研究「模塊流形」提出了一種新范式，它將傳統(tǒng)「救火式」的數(shù)值修正，轉(zhuǎn)變?yōu)椤割A(yù)防式」的約束優(yōu)化，為更好地訓(xùn)練大模型提供了全新思路。

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

這是他們繼（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ù)值溢出等一系列問題。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

• 找到一個單位長度的切向量，在梯度方向上盡可能遠(yuǎn)；

• 用學(xué)習(xí)率乘以這個方向，然后從當(dāng)前權(quán)重中減去；

• 把更新后的權(quán)重通過回縮映射拉回流形上。

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

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

流形Muon

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

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

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

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

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

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

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

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

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

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

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

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

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

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

模塊流形

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

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

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

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

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

我們在搭建網(wǎng)絡(luò)的過程中會追蹤這種敏感性，而流形約束有助于我們更加精準(zhǔn)地把握它。

參考資料：

https://thinkingmachines.ai/blog/modular-manifolds/