在大模型訓(xùn)練領(lǐng)域，如何有效管理權(quán)重、避免數(shù)值異常成為關(guān)鍵挑戰(zhàn)。Thinking Machines Lab提出的“模塊流形”理論為這一難題提供了創(chuàng)新性解決方案，將傳統(tǒng)的事后數(shù)值修正轉(zhuǎn)變?yōu)槭虑暗募s束優(yōu)化，為大模型訓(xùn)練開(kāi)辟了新路徑。

訓(xùn)練大型神經(jīng)網(wǎng)絡(luò)如同精密操作，需嚴(yán)格控制權(quán)重、激活值和梯度等關(guān)鍵參數(shù)的數(shù)值范圍。任何參數(shù)的過(guò)度膨脹或萎縮都可能引發(fā)數(shù)值溢出，導(dǎo)致訓(xùn)練失敗。當(dāng)前主流方法包括使用Layer Norm技術(shù)歸一化層輸出，以及通過(guò)Muon優(yōu)化器等工具對(duì)梯度更新進(jìn)行譜歸一化處理，這些手段旨在維持參數(shù)的數(shù)值穩(wěn)定。

研究團(tuán)隊(duì)提出更深入的解決方案：將權(quán)重張量約束在特定子流形上，并據(jù)此設(shè)計(jì)優(yōu)化算法。這種“預(yù)防式”方法通過(guò)初始參數(shù)設(shè)置確保訓(xùn)練穩(wěn)定性，使模型具有更強(qiáng)的可解釋性。流形優(yōu)化的核心在于將參數(shù)限制在局部平坦的曲面上，通過(guò)切空間優(yōu)化實(shí)現(xiàn)精確控制。

切空間作為流形上某點(diǎn)的局部平坦區(qū)域，是實(shí)施優(yōu)化的關(guān)鍵場(chǎng)所。以三維球面為例，其切平面構(gòu)成優(yōu)化操作的基準(zhǔn)面。傳統(tǒng)方法在每步優(yōu)化后強(qiáng)制投影參數(shù)回流形，但可能導(dǎo)致實(shí)際位移與學(xué)習(xí)率脫節(jié)。研究團(tuán)隊(duì)提出直接在切空間進(jìn)行優(yōu)化，使學(xué)習(xí)率能準(zhǔn)確反映參數(shù)的實(shí)際移動(dòng)距離。

距離度量方式的選擇直接影響優(yōu)化方向。研究采用歐幾里得距離作為基準(zhǔn)，同時(shí)探索其他度量標(biāo)準(zhǔn)。通過(guò)數(shù)學(xué)建模，將流形約束下的最優(yōu)更新方向轉(zhuǎn)化為帶約束的優(yōu)化問(wèn)題。以超球面為例，最優(yōu)更新方向需同時(shí)滿(mǎn)足切平面約束和半徑約束，這一過(guò)程可通過(guò)拉格朗日乘數(shù)法精確求解。

具體操作中，最優(yōu)更新包含三個(gè)步驟：首先將梯度投影到切空間，去除與當(dāng)前點(diǎn)同方向的徑向分量；其次對(duì)投影結(jié)果進(jìn)行歸一化處理；最后乘以學(xué)習(xí)率得到更新方向。這種“回縮映射”機(jī)制確保參數(shù)始終保持在流形約束范圍內(nèi)。

不同流形選擇和距離度量組合可衍生出多種優(yōu)化算法。研究團(tuán)隊(duì)特別關(guān)注Transformer權(quán)重矩陣的優(yōu)化，提出基于Stiefel流形的約束方法。該流形要求所有奇異值均為1，通過(guò)奇異值分解可直觀理解矩陣對(duì)輸入向量的拉伸效應(yīng)。配合譜范數(shù)作為距離函數(shù)，有效限制權(quán)重更新的最大和最小效應(yīng)。

結(jié)合Stiefel流形約束和譜范數(shù)度量，研究團(tuán)隊(duì)開(kāi)發(fā)出Muon優(yōu)化器的流形版本。通過(guò)凸優(yōu)化問(wèn)題建模和對(duì)偶上升法求解，實(shí)現(xiàn)了參數(shù)更新的精確控制。實(shí)驗(yàn)驗(yàn)證表明，該算法能有效維持參數(shù)穩(wěn)定性，防止數(shù)值異常。

當(dāng)擴(kuò)展至多層神經(jīng)網(wǎng)絡(luò)時(shí)，研究提出“模塊流形”理論。該理論通過(guò)追蹤網(wǎng)絡(luò)輸出的Lipschitz敏感性，指導(dǎo)各層學(xué)習(xí)率的合理分配。流形約束為這種敏感性分析提供了精確框架，使不同層之間的參數(shù)更新保持協(xié)調(diào)，避免因?qū)娱g交互導(dǎo)致的訓(xùn)練不穩(wěn)定。