Meta版o1來了!田淵棟團隊整合快慢思考,能走迷宮推箱子

明敏 發自 凹非寺

量子位 | 公眾號 QbitAI

Meta版o1也來了。

田淵棟團隊帶來新作Dualformer,把快慢思考無縫結合,性能提升還成本更低。

能解決迷宮、推箱子等複雜問題。

通過讓模型在推理軌跡和最終答案上進行訓練,再基於特定策略丟掉部分軌跡,Dualformer模型可以在模仿慢思考的同時,像快思考一樣走捷徑。

由此能形成更簡潔的思維鏈(CoT)。

從結果來看,在慢思考模式下,Dualformer的最優解率達到97.6%,推理步驟減少45.5%。

自動切換快慢思考模式下,最優率也達到96.6%,且推理步驟減少59.9%。

搞掂o1玩不來的迷宮遊戲

o1帶火了系統2(慢思考),能讓大模型推理能力大幅提升。

但是隨之而來的計算成本更高。

Dualformer能很好結合快慢思考,從而緩解這一問題。

它建立在Searchformer這項工作的基礎上。Searchformer是一個可以解決複雜推理任務的模型,在A*搜索算法生成的路徑上訓練而來,在路徑規劃任務(如迷宮、推箱子遊戲)上表現良好,可以以更高效率找到最優解。

研究發現,人類會在思考過程中傾向於找捷徑。為了更進一步模擬人類,Dualformer在隨機推理軌跡數據上進行訓練,並在訓練過程中依據定製的丟棄策略丟到部分結構。

比如在處理路徑規劃任務時,根據搜索軌跡中的不同子句(如close子句、子句中的cost tokens、create子句等)設計了四個級別的丟棄策略,從只丟棄close子句到丟棄整個軌跡,並在訓練時隨機選擇應用這些策略。

基於這些策略,Dualformer可以學習更簡潔有效的搜索和推理過程。

在推理階段,Dualformer可配置快速模式(僅輸出解決方案)、慢速模式(輸出推理鏈和最終解決方案)或自動模式(自行決定推理模式)。

這種靈活的推理模式設計使得模型能夠根據不同任務需求和場景進行自適應調整,類似於人類思維在不同情況下的決策方式。

在具體任務上,研究設置了迷宮(Maze)和推箱子遊戲(Sokoban),讓模型進行路徑規劃。以及數學推理任務。

對比來看,在迷宮任務中,o1-preview和o1-mini模型輸出的路徑並不好,會「穿牆」。

快思考模式下,Dualformer的表現如下。

Dualformer以80%的最優率完成這些任務,顯著優於僅基於解決方案數據訓練的Solution-Only模型,後者的最優率僅為 30%。

慢思考模式表現如下。

30×30迷宮任務中,在97.6%的情況下可以達到最優解,同時推理步驟減少45.5%。

自動切換快慢思考模式下,Dualformer的最優率達到 96.6%,與Searchformer相比,推理步驟減少59.9%。

將該方法推廣到Mistral-7B和Llama3-8B上,在Aug-MATH數據集上,模型的表現都有所提升。

比如在Mistral-7B模型上,當p=0.1、0.2和0.3時,Pass@20度量的基線模型,其中絕對正確率增加到61.9%。

最後,來看一下研究團隊陣容。

該研究由田淵棟等人帶來。

田淵棟現在是Meta FAIR的研究科學家主任,領導LLM推理、規劃和決策小組。

Qinqing Zheng是FAIR的工程師,研究方向集中在生成模型和強化學習方面。她本科畢業於浙江大學,在芝加哥大學攻讀博士學位。2017-2019年期間在Facebook擔任研究科學家,幫助Facebook建立了廣告推薦模型的分佈式訓練系統。

Sainbayar Sukhbaatar是FAIR的研究科學家,主要負責大模型推理和記憶方面研究。他曾先後在Google、DeepMind、Meta任職。

Michael Rabbat是FAIR的創始成員之一。加入Meta之前他曾是麥吉爾大學計算機工程系教授。研究領域包括機器學習、分佈式算法、信號處理等。

論文地址:論文地址:

https://arxiv.org/pdf/2410.09918