o1 不是通向大模型推理的唯一路徑!
MIT 的新研究發現,在測試時對大模型進行訓練,可以讓推理水平大幅提升。
在挑戰超難的 ARC 任務時,準确率最高可提升至原來的 5.83 倍。
這樣的表現不僅優于 GPT-4 和 Claude,如果與其他推理方法相結合,還能超越人類的平均水準。
OpenAI o1 團隊成員Noam Brown表示,o1 的大規模計算可能不是最好的方法,很高興看到有學者在提高推理能力上探索新的方法。
在測試中訓練模型
不同于傳統的先訓練後測試模式,測試時訓練(Test-Time Training,TTT)在部署階段面對新的測試樣本時,不直接用訓練好的模型去推理。
在推理之前,測試樣本自身攜帶的信息,會通過快速的訓練過程被用于調整模型參數。
總體來說,TTT 過程中一共有三個關鍵階段——訓練數據生成、模型适應範式設計以及推理階段的策略。
數據生成的核心是将測試任務中蘊含的輸入輸出對關系,通過數據增強的方式最大限度地利用,可具體分爲兩個步驟。
首先是基于 leave-one-out 構造新的任務。
對于包含 K 個輸入輸出對的測試任務,依次将每個樣本留出作爲測試樣本,其餘 K-1 個作爲訓練樣本 , 由此構造出 K 個新的 TTT 訓練任務。
這樣就可以從一個測試任務出發,構造出 K 個結構一緻但内容互補的新任務,從而擴充了 TTT 訓練數據。
在此基礎上,作者還進行了數據增強,主要包括對輸入輸出施加各類幾何變換,以及打亂訓練樣本對的順序。
經過這一步,TTT 訓練集的規模可以得到顯著擴大。
整個 TTT 數據構造過程可高度自動化,不依賴人工标注。
利用構造好的 TTT 數據集,就可以對預訓練好的語言模型進行測試時訓練。
考慮到測試時的資源限制,作者采用了參數高效的 LoRA,爲每個測試任務學習一組獨立的 adapter 參數,附加在預訓練模型的每一層之上,通過一個低秩矩陣與原始權重相乘起到調節作用。
過程中還額外加入了對所有前綴序列的預測,目的是通過在各種長度的演示樣本上都計算損失,鼓勵模型盡早地從少量信息中總結出抽象規律,從而提高魯棒性。
最後,爲了實現 TTT 效果的最大化,作者在推理階段應用了數據增強和集成學習策略。
推理過程中,先利用一系列預定義的幾何變換算子(如旋轉、翻轉等)擴充原始輸入,生成若幹等價視角下的輸入變體。
之後将每個變體輸入并行地送入 LoRA-tuned 模型,獨立完成預測,然後再對齊和還原到原始輸入空間,由此得到一組成對的預測。
在成對預測的基礎上,通過分兩層投票的方式完成集成融合:
第一層在每種變換内部進行投票,選出置信度最高的 Top-3 個預測 ;
第二層在不同變換的 Top-3 預測之間進行全局投票,選出最終的 Top-2 作爲輸出。
這一推理策略,既通過數據增強引入了輸入的多樣性,又用分層投票的方式對不同來源的預測進行了結構化的組合,進一步提升了 TTT 方法的效果。
ARC 任務準确率最高升至 6 倍
爲了評估 TTT 方法的效果,研究團隊以 8B 參數的 GPT-3 作爲基礎模型進行了測試。
如果不使用 TTT 僅進行微調,模型在 ARC 數據集上的準确率隻有 18.3%,加入 TTT 後提升到 47.1%,增長率達到了 157%。
另外,作者還從 ARC 數據集中随機選擇了 80 個任務作爲子集進行了測試。
測試發現,TTT 方法對于 1B 模型的提升效果更加明顯,調整後模型的準确率接近調整前的 6 倍。
并且在調整前後,1B 和 8B 兩個規模的模型之間的相對差距也在縮小。
進一步地,作者還将 TTT 方法與之前在 ARC 任務上取得優異成績的BARC(Bootstrapping Approach for Reward model Construction)方法進行了比較和結合。
具體來說,作者首先獨立運行這兩個系統,得到它們在每個測試任務上的輸出。
如果兩者輸出完全一緻,則直接認爲推理結果是正确的;
如果輸出不一緻,則看 BARC 是否能夠生成确定的、唯一覆蓋所有測試樣本的解題程序,若是則認爲 BARC 的輸出更可靠;
反之,如果 BARC 生成了多個候選程序但無法确定最優解,或者幹脆無法生成任何滿足約束的程序,則認爲 TTT 的輸出更可靠。
兩種方式配合使用後,取得了 61.9% 的 SOTA 成績,已經超過了人類的平均水平。
One More Thing
根據作者在推文中的介紹,在這篇論文發布前,一個叫做 MindsAI 的團隊已經發現使用了相同的技術。
利用 TTT 技術,該團隊已經用 58% 的正确率取得了 ARC 挑戰的第一名。
作者的論文發布之後,MindsAI 團隊領導者 Jack Cole 也發文進行了祝賀:
很高興,我們掀起了這場對 TTT 的興趣風暴。
同時,Jack 還推薦了另一名研究 TTT 的學者——斯坦福大學華人博士後Yu Sun,表示他的研究值得被關注。
Sun 的個人主頁顯示,他針對測試時訓練進行了大量研究,相關成果入選過 ICML、NeurIPS、ICLR 等多個頂級會議。
論文地址:
https://ekinakyurek.github.io/papers/ttt.pdf