In-Place Test-Time Training
靜態的「先訓練後部署 (Train then Deploy)」典範從根本上限制了大型語言模型 (Large Language Model, LLM) 無法因應真實世界任務中持續流入的新資訊而動態調整權重。測試時訓練 (Test-Time Training, TTT) 透過在推論時更新模型參數的子集 (稱為快權重 (Fast Weight)) 提供了一個有吸引力的替代方案,但其在當今 LLM 生態系中的潛力受到關鍵障礙阻撓,包括架構不相容、計算效率低下、以及語言建模快權重目標 (Fast Weight Objective) 的錯位。 在這項工作中,我們提出 就地測試時訓練 (In-Place Test-Time Training, In-Place TTT) 框架,能無縫地賦予 LLM 測試時訓練能力。In-Place TTT 將普及的 MLP 區塊 (MLP Block) 中的最終投影矩陣 (Final Projection Matrix) 視為其可調適的快權重,使 LLM 能在不需要昂貴從頭再訓練 (Retraining from Scratch) 的情況下獲得「插即用 (Drop-in)」的增強。此外,我們以一個量身打造、具備理論基礎、明確對齊管理自迴歸語言建模 (Autoregressive Language Modeling) 的下一個 token 預測 (Next-Token Prediction, NTP) 任務的目標,取代 TTT 通用的重建目標 (Reconstruction Objective)。這個原則性目標結合高效的分塊更新機制 (Chunk-wise Update Mechanism),產生一個高度可擴展、相容於上下文平行 (Context Parallelism, CP) 的演算法。 廣泛實驗驗證我們框架的有效性:作為就地增強,它讓 4B 參數的模型在最高 128k 上下文長度的任務上達到優異效能;當從頭開始預訓練時,它持續超越具競爭力的 TTT 相關方法。消融研究進一步提供對我們設計選擇更深入的洞見。整體而言,我們的結果確立 In-Place TTT 是邁向 LLM 持續學習 (Continual Learning) 典範的有前景的一步。 - 代碼:https://github.com/ByteDance-Seed/In-Place-TTT
原文連結:arXiv