比Adam更有效，POET從譜不變原理出發，讓LLM訓練又穩又快
Zeju Qiu和Tim Z. Xiao是德國馬普所博士生，Simon Buchholz和Maximilian Dax擔任德國馬普所博士后研究員，Bernhard Sch?lkopf是德國馬普所所長，Weiyang Liu是香港中文大學計算機系助理教授。
隨著大型語言模型（LLM）推動人工智能領域取得突破性進展，如何實現高效、穩定的超大規模模型訓練，始終是該領域最富挑戰性的核心議題之一。
針對這一關鍵問題，研究者們提出了一種基于第一性原理的全新方法——POET（Reparameterized Training via Orthogonal Equivalence Transformation），該方法通過重參數化優化策略，旨在從第一性原理出發提升訓練效率與穩定性。
Paper：Reparameterized LLM Training via Orthogonal Equivalence Transformation
Project page：https://spherelab.ai/poet/
Arxiv：https://www.arxiv.org/abs/2506.08001
POET：基于第一性原理的大型語言模型全新訓練范式
POET 的關鍵思想是：通過對每個神經元進行結構性重參數化，引入兩個可學習的正交矩陣以及一個固定的隨機權重矩陣，從而構建一個正交等價的變換結構。該方法在訓練過程中嚴格保持權重的奇異值分布，并天然擁有較低的球面能量，這是 POET 有效性的核心來源。
通過聯合建模奇異值不變性與最小超球能量，POET為大模型訓練提供了一種兼具物理解釋性與泛化能力的新范式。由于該方法嚴格保持權重矩陣的譜結構，不僅能穩定優化過程，還顯著提升了模型的泛化性能。為兼顧計算效率與實用性，研究者還開發了高效的近似算法，使POET可擴展至超大規模神經網絡訓練。實驗結果表明，該方法在大型語言模型訓練中表現出卓越的性能與可擴展性。
圖 POET 的三個學習階段：左—示意圖；中—角度；右—損失值與驗證。
譜性質與泛化
當前訓練大型語言模型的事實標準是直接使用Adam優化器對權重矩陣進行更新。盡管這一做法實現簡單，但在計算上往往代價高昂，隨著模型規模的擴大，其復雜度迅速增長。此外，該方法對超參數極為敏感，需精細調整以保證訓練穩定收斂。
更為關鍵的是，即便訓練損失已經被有效最小化，模型的泛化性能仍可能表現不佳。為緩解這一問題，本文提出了多種權重正則化與歸一化技術，其核心目標往往可歸結為：顯式或隱式地改善權重矩陣的譜結構（即奇異值分布）。
從直觀角度看，權重矩陣的譜范數（最大奇異值）描述了其對輸入向量的放大上界，因此與模型的平滑性和泛化能力密切相關。一般認為，較小的譜范數（意味著更溫和的變換）往往有助于提升泛化性能。這一觀點促使越來越多研究致力于對譜性質進行精細控制。理論研究亦表明，若能有效約束權重矩陣的譜結構，便可形式化地為模型提供泛化上的保證。
譜保持（Spectrum-preserving）權重更新
為在避免上述局限的同時實現有效的權重量譜控制，文章提出一種基于正交等價變換（OrthogonalEquivalenceTransformation）的重參數化訓練算法，用于間接學習權重矩陣。具體而言，POET 將權重矩陣
重參數化為
，其中
為隨機初始化的權重矩陣，
和
為兩個可學習的正交矩陣。在訓練過程中，POET 不直接優化權重矩陣，而是保持隨機初始化矩陣
不變，通過學習正交矩陣
和
來變換
。這種重參數化策略在允許奇異向量靈活調整的同時，能夠嚴格保持權重矩陣的奇異值譜，從而實現譜保持的權重更新。
圖 LLaMA模型中同一權重矩陣奇異值的訓練動態。左圖為標準訓練，嚴格遵循大型語言模型的常規做法（使用AdamW直接優化）；右圖為POET，其采用本文提出的近似方法以支持大規模LLM訓練。POET的奇異值僅出現輕微（幾乎可忽略）的變化，主要歸因于數值誤差和近似誤差。
奇異值譜的訓練動態
受 Muon [4]的啟發，研究者對 AdamW、Muon與 POET 的奇異值譜進行了譜分析。在訓練的不同迭代點，可對訓練后的模型計算 SVD 熵。
該指標用于衡量奇異值的多樣性；熵值越高，表示譜分布越均勻、越豐富。[4] 將 Muon 相較于 AdamW 的優越性能歸因于其權重矩陣更新所帶來的更豐富譜分布。正如下圖所示，由于采用正交等價變換，POET 在整個訓練過程中始終保持較高的譜多樣性。
POET方法具備兩項核心優勢：
- 高效的譜控制
由于正交變換并不改變權重矩陣的奇異值，POET在訓練全程都能保持權重譜與隨機初始化矩陣一致——即便采用近似實現，這一點也已得到實證驗證。借助恰當的初始化方案，POET可直接約束奇異值分布，避免標準LLM訓練后權重出現過大的奇異值。為進一步增強算法效果，研究者們提出了兩種新初始化策略：歸一化高斯初始化（normalizedGaussianinitialization）和 均勻譜初始化（uniformspectruminitialization），均可確保生成的權重矩陣具有有界奇異值。
- 高效近似
直接進行POET訓練的計算開銷較高，但方法本身的靈活性為高效、可擴展訓練提供了空間。針對大規模正交矩陣優化這一關鍵難題，文章提出兩級近似方案：
隨機基元優化：將大正交矩陣分解為若干參數量更少的基元正交矩陣，并結合“合并再初始化”策略提高效率；
基于Cayley?Neumann參數化的近似正交性：通過 Neumann 級數近似 Cayley 正交參數化，以較低計算成本保持正交性，同樣借助“合并再初始化”策略抑制誤差累積。
LLaMA架構的大規模語言模型預訓練
本文在多種規模的LLaMATransformer（60M、130M、350M、1.3B 參數）上對POET進行了預訓練實驗。使用的數據集為C4——從CommonCrawl清洗得到的網頁語料，已被廣泛用于大型語言模型的預訓練。下文匯總了實驗結果，報告了驗證困惑度（perplexity）及可訓練參數量。
圖 AdamW和POET在模型規模為350M和1.3B下的可訓練參數規模及驗證困惑度（perplexity）。
訓練加速
為突出POET在性能上的顯著改進，文章將AdamW的訓練步數（即模型實際看到的token數量）大幅提升至原來的近三倍。即便如此，采用 b=1/2 設置的POET?FS仍在性能上超越AdamW。
參數與內存復雜度
通過將超參數 b 作為采樣預算引入，完全隨機 SPO（StochasticPrimitiveOptimization）成功將參數復雜度與權重矩陣規模解耦。當 b 取較小值時，POET 的參數效率顯著提升，但收斂速度有所下降，為使用者提供了效率與速度之間的靈活權衡。相比之下，塊隨機 SPO 的參數復雜度與矩陣尺寸（m+n）成正比，因而較 AdamW（需要 mn 個可訓練參數）更具可擴展性。在內存占用方面，只要采樣預算 b 設置得當，兩種 POET 變體均可顯著優于 AdamW。下文給出了參數與內存復雜度的詳細對比。
訓練算法
步驟1：權重初始化使用歸一化高斯初始化為權重矩陣賦值：
步驟2：正交矩陣初始化
完全隨機SPO（fullystochasticSPO）：隨機采樣索引集合 
，并使用CNP（Cayley?NeumannParameterization）對
與
進行參數化。
與
二者均以單位矩陣開始。
塊隨機SPO（block?stochasticSPO）：隨機采樣置換矩陣
與
，同樣采用CNP對
與進行參數化，并將它們初始化為單位矩陣。同樣，
與
二者均以單位矩陣開始。
步驟3：高效正交參數化
對于完全隨機SPO可得：
和
。
對于塊隨機SPO可得：
和
。
步驟4：正交矩陣內層訓練循環更新
前向傳播中的等效權重矩陣為
。
反向傳播通過
與
計算梯度，進而更新完全隨機SPO中
的或塊隨機SPO中的
；該內循環迭代次數固定。
步驟5：合并并重新初始化（merge?then?reinitialize）
將已學習的正交矩陣合并進權重：
。
若訓練未結束，則返回步驟2，重新初始化
與
，繼續下一階段訓練。
POET的優異表現來自于超球能量與譜保持
神經元初始化
鑒于 POET 在訓練過程中會保留初始權重矩陣的譜特性，初始化策略顯得至關重要。文章運用了歸一化高斯初始化：先從零均值、固定方差的高斯分布中抽取神經元權重，再對其進行歸一化。下表對多種隨機初始化方案進行了實證比較，結果顯示歸一化高斯初始化取得了最佳最終性能。研究者推測，這一優異表現源于 POET 在該初始化下能夠在訓練過程中同時保持超球能量與譜特性。
訓練中的超球能量
超球能量 HE 用于衡量神經元在單位超球面上的均勻分布程度，可作為刻畫各層神經表征的一種度量。文獻[2,3]表明，滿足正交約束的訓練過程可在訓練期間保持這一超球能量不變，從而避免表征退化并提升泛化性能。
歸一化高斯初始化下的POET 可同時保持能量與奇異值分布
在零均值、各向同性的高斯初始化條件下，POET 能夠同時實現譜保持訓練與能量保持訓練。這一特性為歸一化高斯初始化方法的最優性能提供了理論解釋（詳細證明參見附錄 B）。
POET訓練機理解析
為深入理解POET的運行機制，我們用向量探測（vector probing）分析正交矩陣的學習動態。具體做法是：固定隨機生成的單位向量
，計算
，即
與
之間的余弦相似度，以評估正交矩陣
和
的演化。
對七個可學習的正交矩陣在訓練過程中的余弦相似度進行跟蹤后，可以將其學習過程劃分為三個階段（見圖1）：
- 錐殼搜索階段（Conical shell searching）
余弦相似度從1（即
為單位矩陣）逐漸下降并收斂到[0.60,0.65]。這一現象在所有正交矩陣上均一致，說明
將
映射到其原始方向附近的狹窄錐殼上。
- 錐殼上的穩定學習階段（Stable learning on the conical shell）
余弦相似度保持在該區間內不再顯著變化，但模型開始進入穩定學習期；盡管余弦值趨于穩定，驗證困惑度仍在線性下降。
- 最終階段微調（Final adjusting）
隨著學習率逐步衰減至零，學習速度放緩并最終停止。
更為詳盡的討論與實證結果見論文附錄。
