比Adam更有效，POET從譜不變原理出發(fā)，讓LLM訓練又穩(wěn)又快

2025-07-15 09:27:29

研究者們提出了一種基于第一性原理的全新方法——POET（Reparameterized Training via Orthogonal Equivalence Transformation），該方法通過重參數(shù)化優(yōu)化策略，旨在從第一性原理出發(fā)提升訓練效率與穩(wěn)定性。

Zeju Qiu和Tim Z. Xiao是德國馬普所博士生，Simon Buchholz和Maximilian Dax擔任德國馬普所博士后研究員，Bernhard Sch?lkopf是德國馬普所所長，Weiyang Liu是香港中文大學計算機系助理教授。

隨著大型語言模型（LLM）推動人工智能領(lǐng)域取得突破性進展，如何實現(xiàn)高效、穩(wěn)定的超大規(guī)模模型訓練，始終是該領(lǐng)域最富挑戰(zhàn)性的核心議題之一。

針對這一關(guān)鍵問題，研究者們提出了一種基于第一性原理的全新方法——POET（Reparameterized Training via Orthogonal Equivalence Transformation），該方法通過重參數(shù)化優(yōu)化策略，旨在從第一性原理出發(fā)提升訓練效率與穩(wěn)定性。

Paper：Reparameterized LLM Training via Orthogonal Equivalence Transformation

Project page：https://spherelab.ai/poet/

Arxiv：https://www.arxiv.org/abs/2506.08001

POET：基于第一性原理的大型語言模型全新訓練范式

POET 的關(guān)鍵思想是：通過對每個神經(jīng)元進行結(jié)構(gòu)性重參數(shù)化，引入兩個可學習的正交矩陣以及一個固定的隨機權(quán)重矩陣，從而構(gòu)建一個正交等價的變換結(jié)構(gòu)。該方法在訓練過程中嚴格保持權(quán)重的奇異值分布，并天然擁有較低的球面能量，這是 POET 有效性的核心來源。

通過聯(lián)合建模奇異值不變性與最小超球能量，POET為大模型訓練提供了一種兼具物理解釋性與泛化能力的新范式。由于該方法嚴格保持權(quán)重矩陣的譜結(jié)構(gòu)，不僅能穩(wěn)定優(yōu)化過程，還顯著提升了模型的泛化性能。為兼顧計算效率與實用性，研究者還開發(fā)了高效的近似算法，使POET可擴展至超大規(guī)模神經(jīng)網(wǎng)絡(luò)訓練。實驗結(jié)果表明，該方法在大型語言模型訓練中表現(xiàn)出卓越的性能與可擴展性。

圖 POET 的三個學習階段：左—示意圖；中—角度；右—損失值與驗證。

譜性質(zhì)與泛化

當前訓練大型語言模型的事實標準是直接使用Adam優(yōu)化器對權(quán)重矩陣進行更新。盡管這一做法實現(xiàn)簡單，但在計算上往往代價高昂，隨著模型規(guī)模的擴大，其復(fù)雜度迅速增長。此外，該方法對超參數(shù)極為敏感，需精細調(diào)整以保證訓練穩(wěn)定收斂。

更為關(guān)鍵的是，即便訓練損失已經(jīng)被有效最小化，模型的泛化性能仍可能表現(xiàn)不佳。為緩解這一問題，本文提出了多種權(quán)重正則化與歸一化技術(shù)，其核心目標往往可歸結(jié)為：顯式或隱式地改善權(quán)重矩陣的譜結(jié)構(gòu)（即奇異值分布）。

從直觀角度看，權(quán)重矩陣的譜范數(shù)（最大奇異值）描述了其對輸入向量的放大上界，因此與模型的平滑性和泛化能力密切相關(guān)。一般認為，較小的譜范數(shù)（意味著更溫和的變換）往往有助于提升泛化性能。這一觀點促使越來越多研究致力于對譜性質(zhì)進行精細控制。理論研究亦表明，若能有效約束權(quán)重矩陣的譜結(jié)構(gòu)，便可形式化地為模型提供泛化上的保證。

譜保持（Spectrum-preserving）權(quán)重更新

為在避免上述局限的同時實現(xiàn)有效的權(quán)重量譜控制，文章提出一種基于正交等價變換（OrthogonalEquivalenceTransformation）的重參數(shù)化訓練算法，用于間接學習權(quán)重矩陣。具體而言，POET 將權(quán)重矩陣重參數(shù)化為，其中為隨機初始化的權(quán)重矩陣，和為兩個可學習的正交矩陣。在訓練過程中，POET 不直接優(yōu)化權(quán)重矩陣，而是保持隨機初始化矩陣不變，通過學習正交矩陣和來變換。這種重參數(shù)化策略在允許奇異向量靈活調(diào)整的同時，能夠嚴格保持權(quán)重矩陣的奇異值譜，從而實現(xiàn)譜保持的權(quán)重更新。

圖 LLaMA模型中同一權(quán)重矩陣奇異值的訓練動態(tài)。左圖為標準訓練，嚴格遵循大型語言模型的常規(guī)做法（使用AdamW直接優(yōu)化）；右圖為POET，其采用本文提出的近似方法以支持大規(guī)模LLM訓練。POET的奇異值僅出現(xiàn)輕微（幾乎可忽略）的變化，主要歸因于數(shù)值誤差和近似誤差。

奇異值譜的訓練動態(tài)

受 Muon [4]的啟發(fā)，研究者對 AdamW、Muon與 POET 的奇異值譜進行了譜分析。在訓練的不同迭代點，可對訓練后的模型計算 SVD 熵。

該指標用于衡量奇異值的多樣性；熵值越高，表示譜分布越均勻、越豐富。[4] 將 Muon 相較于 AdamW 的優(yōu)越性能歸因于其權(quán)重矩陣更新所帶來的更豐富譜分布。正如下圖所示，由于采用正交等價變換，POET 在整個訓練過程中始終保持較高的譜多樣性。

POET方法具備兩項核心優(yōu)勢：

高效的譜控制

由于正交變換并不改變權(quán)重矩陣的奇異值，POET在訓練全程都能保持權(quán)重譜與隨機初始化矩陣一致——即便采用近似實現(xiàn)，這一點也已得到實證驗證。借助恰當?shù)某跏蓟桨?，POET可直接約束奇異值分布，避免標準LLM訓練后權(quán)重出現(xiàn)過大的奇異值。為進一步增強算法效果，研究者們提出了兩種新初始化策略：歸一化高斯初始化（normalizedGaussianinitialization）和 均勻譜初始化（uniformspectruminitialization），均可確保生成的權(quán)重矩陣具有有界奇異值。

高效近似

直接進行POET訓練的計算開銷較高，但方法本身的靈活性為高效、可擴展訓練提供了空間。針對大規(guī)模正交矩陣優(yōu)化這一關(guān)鍵難題，文章提出兩級近似方案：

隨機基元優(yōu)化：將大正交矩陣分解為若干參數(shù)量更少的基元正交矩陣，并結(jié)合“合并再初始化”策略提高效率；

基于Cayley?Neumann參數(shù)化的近似正交性：通過 Neumann 級數(shù)近似 Cayley 正交參數(shù)化，以較低計算成本保持正交性，同樣借助“合并再初始化”策略抑制誤差累積。

LLaMA架構(gòu)的大規(guī)模語言模型預(yù)訓練

本文在多種規(guī)模的LLaMATransformer（60M、130M、350M、1.3B 參數(shù)）上對POET進行了預(yù)訓練實驗。使用的數(shù)據(jù)集為C4——從CommonCrawl清洗得到的網(wǎng)頁語料，已被廣泛用于大型語言模型的預(yù)訓練。下文匯總了實驗結(jié)果，報告了驗證困惑度（perplexity）及可訓練參數(shù)量。

圖 AdamW和POET在模型規(guī)模為350M和1.3B下的可訓練參數(shù)規(guī)模及驗證困惑度（perplexity）。

訓練加速

為突出POET在性能上的顯著改進，文章將AdamW的訓練步數(shù)（即模型實際看到的token數(shù)量）大幅提升至原來的近三倍。即便如此，采用 b=1/2 設(shè)置的POET?FS仍在性能上超越AdamW。

參數(shù)與內(nèi)存復(fù)雜度

通過將超參數(shù) b 作為采樣預(yù)算引入，完全隨機 SPO（StochasticPrimitiveOptimization）成功將參數(shù)復(fù)雜度與權(quán)重矩陣規(guī)模解耦。當 b 取較小值時，POET 的參數(shù)效率顯著提升，但收斂速度有所下降，為使用者提供了效率與速度之間的靈活權(quán)衡。相比之下，塊隨機 SPO 的參數(shù)復(fù)雜度與矩陣尺寸（m+n）成正比，因而較 AdamW（需要 mn 個可訓練參數(shù)）更具可擴展性。在內(nèi)存占用方面，只要采樣預(yù)算 b 設(shè)置得當，兩種 POET 變體均可顯著優(yōu)于 AdamW。下文給出了參數(shù)與內(nèi)存復(fù)雜度的詳細對比。

訓練算法

步驟1：權(quán)重初始化使用歸一化高斯初始化為權(quán)重矩陣賦值：

步驟2：正交矩陣初始化

完全隨機SPO（fullystochasticSPO）：隨機采樣索引集合，并使用CNP（Cayley?NeumannParameterization）對與進行參數(shù)化。與二者均以單位矩陣開始。

塊隨機SPO（block?stochasticSPO）：隨機采樣置換矩陣與，同樣采用CNP對與進行參數(shù)化，并將它們初始化為單位矩陣。同樣，與二者均以單位矩陣開始。

步驟3：高效正交參數(shù)化

對于完全隨機SPO可得：和。

對于塊隨機SPO可得：和。

步驟4：正交矩陣內(nèi)層訓練循環(huán)更新

前向傳播中的等效權(quán)重矩陣為。

反向傳播通過與計算梯度，進而更新完全隨機SPO中的或塊隨機SPO中的；該內(nèi)循環(huán)迭代次數(shù)固定。

步驟5：合并并重新初始化（merge?then?reinitialize）

將已學習的正交矩陣合并進權(quán)重：。

若訓練未結(jié)束，則返回步驟2，重新初始化與，繼續(xù)下一階段訓練。

POET的優(yōu)異表現(xiàn)來自于超球能量與譜保持

神經(jīng)元初始化

鑒于 POET 在訓練過程中會保留初始權(quán)重矩陣的譜特性，初始化策略顯得至關(guān)重要。文章運用了歸一化高斯初始化：先從零均值、固定方差的高斯分布中抽取神經(jīng)元權(quán)重，再對其進行歸一化。下表對多種隨機初始化方案進行了實證比較，結(jié)果顯示歸一化高斯初始化取得了最佳最終性能。研究者推測，這一優(yōu)異表現(xiàn)源于 POET 在該初始化下能夠在訓練過程中同時保持超球能量與譜特性。