估值840億AI實驗室新成果：給大模型訓(xùn)練加「緊箍咒」
近日，由OpenAI前CTO Mira Murati創(chuàng)辦的Thinking Machines Lab再次發(fā)布研究成果。這是他們繼《克服LLM推理中的不確定性》之后的第二篇研究文章——《模塊流形》。
博客地址：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ù)，將每層的輸出拉回合適范圍，對激活向量進行歸一化。對梯度更新進行歸一化也很常見，例如Muon優(yōu)化器對更新進行譜歸一化處理，使每一步更新的幅度可控。
更進一步，是直接「管住」權(quán)重本體。歸一化權(quán)重矩陣是一個可行方向。文中提出了新視角：將權(quán)重張量約束在某個子流形上，與這些流形約束協(xié)同設(shè)計優(yōu)化算法。這就好比把「救火」變?yōu)椤割A(yù)防」，一開始就將參數(shù)置于健康區(qū)間，讓訓(xùn)練更穩(wěn)定、更具解釋性，使大模型能更高效地訓(xùn)練。
流形優(yōu)化器的形態(tài)
流形是一個局部看起來很平坦的曲面，放大到足夠程度，它就像普通平面。流形上某一點附近的局部平坦空間稱為「切空間」。
如圖1所示，三維球面或更高維度的超球面是流形，圖中紅色部分表示其在某點的切平面。為了讓權(quán)重待在指定流形里，簡單方法是使用普通優(yōu)化器，每步更新后將權(quán)重投影回流形。但如果優(yōu)化步驟偏離流形太多，再強制投影回來，會導(dǎo)致名義學(xué)習(xí)率與參數(shù)在流形上的實際位移不對應(yīng)，削弱我們對「步長—效果」關(guān)系的直覺。
要在流形上設(shè)計訓(xùn)練算法，需先明確在切空間里如何度量「距離」。一個解決思路是直接在切空間中進行優(yōu)化，讓每一步沿著流形「表面」走，使學(xué)習(xí)率更好地對應(yīng)「實際位移」。常見選擇是歐幾里得距離，也可采用其他方式測量距離，如圖2所示。
值得注意的是，距離度量方式的選擇會直接影響最優(yōu)優(yōu)化步驟的方向。
圖3中，粉色箭頭表示原始梯度，即損失函數(shù)對權(quán)重的偏導(dǎo)數(shù)。這意味著我們不一定嚴格按照梯度方向移動。為了用數(shù)學(xué)表達這個過程，可將「在流形約束和特定距離度量下的最優(yōu)更新方向」看作帶約束的優(yōu)化問題，以搭配歐幾里得范數(shù)的超球面為例。
用g表示梯度， w表示超球面上的當前點， a表示更新方向， η表示學(xué)習(xí)率，我們需要解決的問題是：
結(jié)合圖 1、2 和3，這個公式的意思是：綠色箭頭（即a的最優(yōu)解）必須同時滿足兩個條件：一是落在紅色的切平面上，二是在半徑為η的黃色圓圈上。
我們可以用拉格朗日乘數(shù)法求解。
其中λ和μ是拉格朗日乘子。對拉格朗日函數(shù)對a求導(dǎo)并令其為零，結(jié)合兩個約束條件求解λ和μ，可得到最優(yōu)更新方向。
簡單來說，最優(yōu)更新做法是：先從梯度中減去與w同方向的徑向分量，將梯度投影到切空間上，然后將結(jié)果歸一化，再乘以學(xué)習(xí)率。這樣得到的更新方向就在切空間里。
圖4中，這個微小的修正過程被稱為「回縮映射」。完整的流形優(yōu)化算法如下：
總結(jié)來說，一階流形優(yōu)化器包含三個步驟：找到一個單位長度的切向量，在梯度方向上盡可能遠；用學(xué)習(xí)率乘以這個方向，然后從當前權(quán)重中減去；把更新后的權(quán)重通過回縮映射拉回流形上。
執(zhí)行這一流程時，需要決定選擇什么樣的流形作為約束，以及如何定義「長度」的度量方式。根據(jù)這兩個選擇的不同，可得到不同的優(yōu)化算法，具體見下表。
流形Muon
Transformer中的典型權(quán)重矩陣W是「向量變換器」，將輸入向量x轉(zhuǎn)換為輸出向量y=Wx。我們希望設(shè)計一種流形約束和距離函數(shù)，使該矩陣對輸入向量的作用合理，避免輸出值過大或過小，以及更新權(quán)重時輸出向量劇烈變化或幾乎無變化。
使用奇異值分解（SVD）是思考矩陣如何作用于向量的好方法，如圖 5 所示。SVD以分解矩陣的方式顯示矩陣如何沿著不同的軸拉伸輸入向量。
我們希望矩陣的「拉伸效應(yīng)」接近于1，因此選擇所有奇異值均為1的矩陣流形。這種矩陣流形在數(shù)學(xué)上被稱為Stiefel流形，在高矩陣（ m≥n）的假設(shè)下，它可以等價地定義為以下集合：
要為Stiefel流形設(shè)計優(yōu)化器，還需選擇合適的距離函數(shù)。為限制權(quán)重更新對輸入向量的最大拉伸作用，譜范數(shù)是合適的選項。雖然它只約束了最大效應(yīng)，但由于優(yōu)化器會飽和這一上限，也能間接防止最小效應(yīng)過小。這一想法促成了Muon優(yōu)化器的提出。
這一想法與Stiefel流形約束結(jié)合后，形成了「manifold Muon」問題。
文中的一個關(guān)鍵發(fā)現(xiàn)是一個凸優(yōu)化問題，可通過標準方法——對偶上升法來求解。
經(jīng)過推導(dǎo)，對偶函數(shù)的梯度為：
通過一個小實驗，可以驗證算法的可行性，實驗設(shè)置與結(jié)果見圖6。
模塊流形
當我們將多個層組合起來構(gòu)建完整的神經(jīng)網(wǎng)絡(luò)時，會出現(xiàn)什么情況？是否需要關(guān)注層與層之間的交互，并據(jù)此修改優(yōu)化策略？這就需要模塊流形理論，它可以將前文的推導(dǎo)邏輯推廣到整個神經(jīng)網(wǎng)絡(luò)。
該理論的核心思想是構(gòu)建一種抽象機制，指導(dǎo)如何在各層之間合理分配學(xué)習(xí)率。本質(zhì)上，在不同層之間分配學(xué)習(xí)率，或?qū)蝹€層進行縮放，都依賴于我們對網(wǎng)絡(luò)輸出對權(quán)重的Lipschitz敏感性的理解。我們在搭建網(wǎng)絡(luò)的過程中會追蹤這種敏感性，而流形約束有助于我們更精準地把握它。
參考資料：
https://thinkingmachines.ai/blog/modular-manifolds/
本文來自微信公眾號“新智元”，作者：新智元，編輯：元宇，36氪經(jīng)授權(quán)發(fā)布。
本文僅代表作者觀點，版權(quán)歸原創(chuàng)者所有，如需轉(zhuǎn)載請在文中注明來源及作者名字。
免責(zé)聲明：本文系轉(zhuǎn)載編輯文章，僅作分享之用。如分享內(nèi)容、圖片侵犯到您的版權(quán)或非授權(quán)發(fā)布，請及時與我們聯(lián)系進行審核處理或刪除，您可以發(fā)送材料至郵箱：service@tojoy.com