o1不是唯一路徑！MIT新研究：在測試時(shí)訓(xùn)練，模型推理能力最高升至5.8倍

克雷西 2024-11-12 15:09:29 來源：量子位

克雷西 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

o1不是通向大模型推理的唯一路徑！

MIT的新研究發(fā)現(xiàn)，在測試時(shí)對大模型進(jìn)行訓(xùn)練，可以讓推理水平大幅提升。

在挑戰(zhàn)超難的ARC任務(wù)時(shí)，準(zhǔn)確率最高可提升至原來的5.83倍。

這樣的表現(xiàn)不僅優(yōu)于GPT-4和Claude，如果與其他推理方法相結(jié)合，還能超越人類的平均水準(zhǔn)。

OpenAI o1團(tuán)隊(duì)成員Noam Brown表示，o1的大規(guī)模計(jì)算可能不是最好的方法，很高興看到有學(xué)者在提高推理能力上探索新的方法。

在測試中訓(xùn)練模型

不同于傳統(tǒng)的先訓(xùn)練后測試模式，測試時(shí)訓(xùn)練（Test-Time Training，TTT）在部署階段面對新的測試樣本時(shí)，不直接用訓(xùn)練好的模型去推理。

在推理之前，測試樣本自身攜帶的信息，會(huì)通過快速的訓(xùn)練過程被用于調(diào)整模型參數(shù)。

總體來說，TTT過程中一共有三個(gè)關(guān)鍵階段——訓(xùn)練數(shù)據(jù)生成、模型適應(yīng)范式設(shè)計(jì)以及推理階段的策略。

數(shù)據(jù)生成的核心是將測試任務(wù)中蘊(yùn)含的輸入輸出對關(guān)系，通過數(shù)據(jù)增強(qiáng)的方式最大限度地利用，可具體分為兩個(gè)步驟。

首先是基于leave-one-out構(gòu)造新的任務(wù)。

對于包含K個(gè)輸入輸出對的測試任務(wù)，依次將每個(gè)樣本留出作為測試樣本，其余K-1個(gè)作為訓(xùn)練樣本,由此構(gòu)造出K個(gè)新的TTT訓(xùn)練任務(wù)。

這樣就可以從一個(gè)測試任務(wù)出發(fā)，構(gòu)造出K個(gè)結(jié)構(gòu)一致但內(nèi)容互補(bǔ)的新任務(wù)，從而擴(kuò)充了TTT訓(xùn)練數(shù)據(jù)。

在此基礎(chǔ)上，作者還進(jìn)行了數(shù)據(jù)增強(qiáng)，主要包括對輸入輸出施加各類幾何變換，以及打亂訓(xùn)練樣本對的順序。

經(jīng)過這一步，TTT訓(xùn)練集的規(guī)?？梢缘玫斤@著擴(kuò)大。

整個(gè)TTT數(shù)據(jù)構(gòu)造過程可高度自動(dòng)化，不依賴人工標(biāo)注。

利用構(gòu)造好的TTT數(shù)據(jù)集，就可以對預(yù)訓(xùn)練好的語言模型進(jìn)行測試時(shí)訓(xùn)練。

考慮到測試時(shí)的資源限制，作者采用了參數(shù)高效的LoRA，為每個(gè)測試任務(wù)學(xué)習(xí)一組獨(dú)立的adapter參數(shù)，附加在預(yù)訓(xùn)練模型的每一層之上，通過一個(gè)低秩矩陣與原始權(quán)重相乘起到調(diào)節(jié)作用。

過程中還額外加入了對所有前綴序列的預(yù)測，目的是通過在各種長度的演示樣本上都計(jì)算損失，鼓勵(lì)模型盡早地從少量信息中總結(jié)出抽象規(guī)律，從而提高魯棒性。

最后，為了實(shí)現(xiàn)TTT效果的最大化，作者在推理階段應(yīng)用了數(shù)據(jù)增強(qiáng)和集成學(xué)習(xí)策略。

推理過程中，先利用一系列預(yù)定義的幾何變換算子（如旋轉(zhuǎn)、翻轉(zhuǎn)等）擴(kuò)充原始輸入，生成若干等價(jià)視角下的輸入變體。

之后將每個(gè)變體輸入并行地送入LoRA-tuned模型，獨(dú)立完成預(yù)測，然后再對齊和還原到原始輸入空間，由此得到一組成對的預(yù)測。

在成對預(yù)測的基礎(chǔ)上，通過分兩層投票的方式完成集成融合：

• 第一層在每種變換內(nèi)部進(jìn)行投票，選出置信度最高的Top-3個(gè)預(yù)測;
• 第二層在不同變換的Top-3預(yù)測之間進(jìn)行全局投票，選出最終的Top-2作為輸出。

這一推理策略，既通過數(shù)據(jù)增強(qiáng)引入了輸入的多樣性，又用分層投票的方式對不同來源的預(yù)測進(jìn)行了結(jié)構(gòu)化的組合，進(jìn)一步提升了TTT方法的效果。

ARC任務(wù)準(zhǔn)確率最高升至6倍

為了評估TTT方法的效果，研究團(tuán)隊(duì)以8B參數(shù)的GPT-3作為基礎(chǔ)模型進(jìn)行了測試。

如果不使用TTT僅進(jìn)行微調(diào)，模型在ARC數(shù)據(jù)集上的準(zhǔn)確率只有18.3%，加入TTT后提升到47.1%，增長率達(dá)到了157%。

另外，作者還從ARC數(shù)據(jù)集中隨機(jī)選擇了80個(gè)任務(wù)作為子集進(jìn)行了測試。

測試發(fā)現(xiàn)，TTT方法對于1B模型的提升效果更加明顯，調(diào)整后模型的準(zhǔn)確率接近調(diào)整前的6倍。

并且在調(diào)整前后，1B和8B兩個(gè)規(guī)模的模型之間的相對差距也在縮小。

進(jìn)一步地，作者還將TTT方法與之前在ARC任務(wù)上取得優(yōu)異成績的BARC（Bootstrapping Approach for Reward model Construction）方法進(jìn)行了比較和結(jié)合。

具體來說，作者首先獨(dú)立運(yùn)行這兩個(gè)系統(tǒng)，得到它們在每個(gè)測試任務(wù)上的輸出。

如果兩者輸出完全一致，則直接認(rèn)為推理結(jié)果是正確的；

如果輸出不一致，則看BARC是否能夠生成確定的、唯一覆蓋所有測試樣本的解題程序，若是則認(rèn)為BARC的輸出更可靠；

反之，如果BARC生成了多個(gè)候選程序但無法確定最優(yōu)解，或者干脆無法生成任何滿足約束的程序，則認(rèn)為TTT的輸出更可靠。

兩種方式配合使用后，取得了61.9%的SOTA成績，已經(jīng)超過了人類的平均水平。

One More Thing

根據(jù)作者在推文中的介紹，在這篇論文發(fā)布前，一個(gè)叫做MindsAI的團(tuán)隊(duì)已經(jīng)發(fā)現(xiàn)使用了相同的技術(shù)。

利用TTT技術(shù)，該團(tuán)隊(duì)已經(jīng)用58%的正確率取得了ARC挑戰(zhàn)的第一名。

作者的論文發(fā)布之后，MindsAI團(tuán)隊(duì)領(lǐng)導(dǎo)者Jack Cole也發(fā)文進(jìn)行了祝賀：

很高興，我們掀起了這場對TTT的興趣風(fēng)暴。

同時(shí)，Jack還推薦了另一名研究TTT的學(xué)者——斯坦福大學(xué)華人博士后Yu Sun，表示他的研究值得被關(guān)注。

Sun的個(gè)人主頁顯示，他針對測試時(shí)訓(xùn)練進(jìn)行了大量研究，相關(guān)成果入選過ICML、NeurIPS、ICLR等多個(gè)頂級會(huì)議。

論文地址：
https://ekinakyurek.github.io/papers/ttt.pdf