OpenAI ο1 模型的發(fā)布掀起了人們對(duì) AI 推理過(guò)程的關(guān)注，甚至讓現(xiàn)在的 AI 行業(yè)開(kāi)始放棄卷越來(lái)越大的模型，而是開(kāi)始針對(duì)推理過(guò)程進(jìn)行優(yōu)化了。今天我們介紹的這項(xiàng)來(lái)自 Meta FAIR 田淵棟團(tuán)隊(duì)的研究也是如此，其從人類(lèi)認(rèn)知理論中獲得了靈感，提出了一種新型 Transformer 架構(gòu)：Dualformer。

根據(jù)人類(lèi)認(rèn)知理論，人類(lèi)的思考受到兩個(gè)系統(tǒng)控制：

• System 1：系統(tǒng) 1，速度快，基于直覺(jué)。
• System 2：系統(tǒng) 2，速度更慢，更加深思熟慮。

近期有研究表明，如果將系統(tǒng) 2 過(guò)程整合進(jìn) Transformer 和大型語(yǔ)言模型中，就能顯著提升它們的推理能力。盡管如此，如果模型只是模仿系統(tǒng) 2 式的思考過(guò)程，那就需要遠(yuǎn)遠(yuǎn)更高的計(jì)算成本才能完成，同時(shí)響應(yīng)速度也會(huì)大幅減慢。

在研究這一難題時(shí)，田淵棟團(tuán)隊(duì)得到了一項(xiàng)驚人發(fā)現(xiàn)：在解決推理任務(wù)時(shí)，一種簡(jiǎn)單的數(shù)據(jù)方案就足以實(shí)現(xiàn)即時(shí)動(dòng)態(tài)的系統(tǒng) 1 和系統(tǒng) 2 配置。

基于此發(fā)現(xiàn)，他們提出了 Dualformer。這是一種可以輕松配置的 Transformer—— 用戶(hù)可以指定在推理過(guò)程中使用快速或慢速模式，在未指定時(shí)模型也可以自行決定。

• 論文標(biāo)題：Dualformer: Controllable Fast and Slow Thinking by Learning with Randomized Reasoning Traces
• 論文地址：https://arxiv.org/pdf/2410.09918

具體而言，為了模仿系統(tǒng) 2 推理過(guò)程，他們讓 Transformer 在包含推理軌跡和最終解答的數(shù)據(jù)上進(jìn)行訓(xùn)練。利用推理步驟的結(jié)構(gòu)，他們?cè)O(shè)計(jì)了特定的軌跡丟棄策略，使得生成的軌跡類(lèi)似于系統(tǒng) 1 在思考過(guò)程中采取的捷徑。在極端情況下，會(huì)丟棄整個(gè)軌跡并鼓勵(lì) Transformer 繞過(guò)所有中間步驟，直接輸出最終解答。在訓(xùn)練時(shí)，他們的策略是隨機(jī)選擇這些結(jié)構(gòu)化的軌跡丟棄策略。

前提準(zhǔn)備

他們的這項(xiàng)研究基于田淵棟團(tuán)隊(duì)之前的另一項(xiàng)研究《Beyond A*: Better planning with transformers via search dynamics bootstrapping》，參閱機(jī)器之心報(bào)道《補(bǔ)齊 Transformer 規(guī)劃短板，田淵棟團(tuán)隊(duì)的 Searchformer 火了》。為了執(zhí)行規(guī)劃，他們要訓(xùn)練一個(gè) Transformer 來(lái)建模一個(gè) token 序列，而該序列則是以順序方式來(lái)表示該規(guī)劃任務(wù)、A* 算法的計(jì)算、由 A* 搜索得到的最優(yōu)解。

圖 3.1 展示了其 token 化方法，其中示例是一個(gè) 3×3 迷宮的導(dǎo)航任務(wù)，目標(biāo)是找到從起點(diǎn)到目標(biāo)單元格的最短路徑。

A* 算法已經(jīng)成功找到了最佳規(guī)劃。這里使用一個(gè) token 序列來(lái)表示該任務(wù)和迷宮結(jié)果，其也被用作 Dualformer 的提示詞。該解答由使用坐標(biāo)描述路徑的規(guī)劃 token 序列描述。A* 算法生成一個(gè)搜索軌跡序列，記錄執(zhí)行的搜索動(dòng)態(tài)，如圖 4.1 所示。

回想一下，A* 算法是一種在加權(quán)圖上的尋路算法。create 子句將節(jié)點(diǎn)（由后續(xù)坐標(biāo)表示）添加到搜索邊界中，close 子句將節(jié)點(diǎn)添加到該閉集。每個(gè)子句（create 或 close）后面都跟著 token x、y、c0 和 c1—— 分別表示節(jié)點(diǎn)的坐標(biāo)、自開(kāi)始以來(lái)的成本值和啟發(fā)值。

結(jié)構(gòu)化軌跡丟棄和隨機(jī)訓(xùn)練

田淵棟團(tuán)隊(duì)之前提出的 Searchformer 已被證明可以有效解決多種復(fù)雜的決策任務(wù)。但是，它仍有兩個(gè)不足。

1. 模型僅能以慢速模式運(yùn)行并會(huì)輸出很長(zhǎng)的推理鏈，這會(huì)極大延長(zhǎng)推理時(shí)間。盡管可通過(guò) bootstrapping（一種迭代優(yōu)化技術(shù)，包含 rollout 循環(huán)和之后的微調(diào)過(guò)程）來(lái)提速，但這樣的過(guò)程會(huì)對(duì)計(jì)算資源產(chǎn)生顯著的額外需求。

2. Searchformer 很難生成多樣化的解答，因?yàn)槠浣?jīng)常會(huì)采樣相同的 rollout。舉個(gè)例子，在他們測(cè)試過(guò)的 1000 個(gè) 30×30 迷宮問(wèn)題中，Searchformer 的推理鏈平均包含 1500 多個(gè) token，而只能在 64 個(gè)響應(yīng)中找到 7.6 條各不一樣的可行路徑。

為了解決這些挑戰(zhàn)，他們提出了一個(gè)利用隨機(jī)化推理軌跡的訓(xùn)練框架。該方法的靈感來(lái)自?xún)蓚€(gè)研究方向：

• 該團(tuán)隊(duì)注意到，即便 Searchformer 是在完整的 A* 搜索軌跡上訓(xùn)練的，但它也會(huì)生成更短的勾勒搜索過(guò)程的軌跡。
• 研究表明，人類(lèi)在做決策時(shí)往往依賴(lài)捷徑和模式，這一概念被稱(chēng)為系統(tǒng) 1 思維。

這些觀(guān)察再加上 dropout 技術(shù)（在訓(xùn)練時(shí)隨機(jī)丟棄神經(jīng)網(wǎng)絡(luò)中的一些單元）的成功，促使該團(tuán)隊(duì)研究了隨機(jī)化推理軌跡的作用，并且他們還希望通過(guò)利用結(jié)構(gòu)化元素并選擇性地丟棄每個(gè)訓(xùn)練示例的某些部分來(lái)簡(jiǎn)化 A* 搜索軌跡。該方法的細(xì)節(jié)如下。

如圖 4.1 所示，A* 搜索軌跡包含 create 和 close 子句，每個(gè)子句都包括節(jié)點(diǎn)的坐標(biāo)及其到達(dá)起始位置和目標(biāo)位置的（估計(jì)）成本。為了推導(dǎo)得到 Dualformer，他們利用了搜索軌跡的結(jié)構(gòu)，并為每個(gè)訓(xùn)練示例丟棄軌跡中的某些部分。其有三種自然的丟棄類(lèi)型：

• D1：丟棄一個(gè) close 子句；
• D2：丟棄一個(gè)子句中的成本 token；
• D3：丟棄一個(gè) create 子句。

基于此，他們開(kāi)發(fā)出了四個(gè)層級(jí)逐層遞進(jìn)的丟棄策略：

• Level 1：去除搜索軌跡中所有 close 子句。
• Level 2：更進(jìn)一步，額外丟棄所有成本 token。
• Level 3：更加激進(jìn)，進(jìn)一步隨機(jī)丟棄 30% 的 create 子句。
• Level 4：丟棄整條搜索軌跡。

圖 4.1 基于上述迷宮任務(wù)演示了這些策略。后面我們會(huì)看到，這些策略可有效地引導(dǎo) Dualformer 學(xué)習(xí)更簡(jiǎn)潔、更高效的搜索和推理過(guò)程。

為了提升訓(xùn)練數(shù)據(jù)的多樣性，他們沒(méi)有將丟棄作為一個(gè)數(shù)據(jù)預(yù)處理步驟。而是在推理時(shí)間，對(duì)于一個(gè)數(shù)據(jù)批次中的每個(gè)訓(xùn)練樣本，都從一個(gè)分類(lèi)分布 Cat (p_0, p_1, p_2, p_3, p_4) 中隨機(jī)抽取丟棄策略，其中 p_1, . . . , p_4 是執(zhí)行 Level 1-4 丟棄的概率，p_0 是保持完整軌跡的概率。這種訓(xùn)練框架可使 Dualformer 學(xué)習(xí)多個(gè)經(jīng)過(guò)約簡(jiǎn)的軌跡，即使對(duì)于單個(gè)訓(xùn)練示例也是如此，因?yàn)橥粋€(gè)示例可能出現(xiàn)在多個(gè)批次中。

可控式生成

Dualformer 具有一個(gè)非常吸引人的特性：在推理時(shí)，可以輕松地通過(guò)提示詞指定以快速或慢速生成模式運(yùn)行。

該控制機(jī)制非常簡(jiǎn)單：在標(biāo)準(zhǔn)提示詞之后添加一個(gè) bos 和一個(gè)控制 token，其中控制 token 是 plan 或 create 中的一個(gè)。

如果使用 plan，則 Dualformer 將以快速模式運(yùn)行，繞過(guò)推理步驟并直接輸出規(guī)劃。另一方面，如果在 bos 之后注入 create，則 Dualformer 將以慢速模式工作并生成推理軌跡和最終規(guī)劃。下面基于迷宮任務(wù)展示了這兩種模式的示意圖。

而如果僅使用標(biāo)準(zhǔn)提示詞，則 Dualformer 將模仿人類(lèi)決策的雙重過(guò)程 —— 根據(jù)情況，它會(huì)選擇一種分別對(duì)應(yīng)于系統(tǒng) 1 和系統(tǒng) 2 的推理類(lèi)型進(jìn)行響應(yīng)。

實(shí)驗(yàn)

實(shí)驗(yàn)的目標(biāo)是解答以下三個(gè)問(wèn)題：

1. Dualformer 在快速、慢速和自動(dòng)模式下的表現(xiàn)是否優(yōu)于相應(yīng)的基線(xiàn)？

2. 在慢速模式下，Dualformer 是否能實(shí)現(xiàn)更快的推理，即輸出更短的軌跡？

3. 結(jié)構(gòu)化的軌跡丟棄技術(shù)是否適用于在自然語(yǔ)言數(shù)據(jù)集上訓(xùn)練的 LLM？

為了解答問(wèn)題 1 和 2，該團(tuán)隊(duì)訓(xùn)練了求解迷宮導(dǎo)航任務(wù)和緊密相關(guān)的推箱子（Sokoban）任務(wù)的 Transformer。為了解答問(wèn)題 3，他們微調(diào)了 LLama-3.1-8B 和 Mistral-7B 模型來(lái)解答數(shù)學(xué)問(wèn)題。

導(dǎo)航任務(wù)：迷宮和推箱子

迷宮和推箱子任務(wù)使用的數(shù)據(jù)集與 Searchformer 研究的一樣。這里就不再贅述，我們直接來(lái)看結(jié)論。

研究表明，Dualformer 可以根據(jù)控制指令選擇快速或慢速的運(yùn)行模式。在快速模式下，它僅輸出最終規(guī)劃；在慢速模式下，它還會(huì)生成推理軌跡。該團(tuán)隊(duì)在不同的模式下讓 Dualformer 對(duì)比了不同的基線(xiàn)。使用的指標(biāo)包括生成規(guī)劃的正確性、最優(yōu)性和多樣性、推理軌跡的長(zhǎng)度等。

• 快速模式

表 5.1 分別報(bào)告了在迷宮和推箱子任務(wù)上，Dualformer 和基線(xiàn)僅解答模型的性能。

可以看到，在生成正確和最優(yōu)規(guī)劃方面，Dualformer 在 1-Solved-64 和 1-Optimal-64 指標(biāo)上中都明顯優(yōu)于基線(xiàn)。它在 3-Solved-64 和 3-Optimal-64 指標(biāo)上也明顯超過(guò)了基線(xiàn)，這證明了 Dualformer 在規(guī)劃生成方面的穩(wěn)健性。

尤其需要注意，隨著任務(wù)難度提升，Dualformer 的優(yōu)勢(shì)也會(huì)增大。對(duì)于最大的 30×30 迷宮，Dualformer 的 1-Optimal-64 成功率是僅解答模型的 2.8 倍，在 3-Optimal-64 上是 2.97 倍。

Dualformer 的 SWC 分?jǐn)?shù)也比基線(xiàn)高得多 —— 在每個(gè)環(huán)境中都高于 0.9。這表明 Dualformer 生成的每個(gè)單獨(dú)規(guī)劃的質(zhì)量都很高，其成本非常接近最佳規(guī)劃。

在實(shí)驗(yàn)考慮的所有問(wèn)題上，Dualformer 還能穩(wěn)定地生成更多樣化的規(guī)劃。比如在下面這個(gè)迷宮示例中，隨著迷宮規(guī)模的增加，Dualformer 的多樣性得分（即 64 個(gè)響應(yīng)中不同但正確的規(guī)劃的平均數(shù)量）會(huì)增加。

一般來(lái)說(shuō)，隨著迷宮規(guī)模增大，到達(dá)單個(gè)目標(biāo)位置的可能路線(xiàn)也越來(lái)越多。這表明 Dualformer 學(xué)習(xí)了迷宮結(jié)構(gòu)，而僅解答模型可能是記住了最佳規(guī)劃，因?yàn)槠涠鄻有缘梅衷谒忻詫m規(guī)模下都接近 1。

• 慢速模式

表 5.2 報(bào)告了 Dualformer 在慢速模式下運(yùn)行時(shí)的結(jié)果。

相應(yīng)的基線(xiàn)是 Complete-Trace 模型，它使用相同的架構(gòu)并在具有完整 A* 搜索軌跡的數(shù)據(jù)上進(jìn)行了訓(xùn)練。除了之前報(bào)告的指標(biāo)之外，該研究還報(bào)告了在所有 1000 個(gè)評(píng)估任務(wù)中匯總的 64 個(gè)響應(yīng)的推理軌跡平均長(zhǎng)度。結(jié)果表明，Dualformer 實(shí)現(xiàn)了更好的規(guī)劃能力和推理速度。它在所有正確性和最優(yōu)性指標(biāo)方面都優(yōu)于 Complete-Trace 模型：包括解決率、最優(yōu)率和 SWC。

此外，Dualformer 產(chǎn)生的推理軌跡明顯短于基線(xiàn)模型。平均而言，Dualformer 在五個(gè)任務(wù)中將軌跡長(zhǎng)度減少了 49.4%。與以前一樣，與基線(xiàn)相比，Dualformer 還生成了更多不同的規(guī)劃。

• 與搜索動(dòng)態(tài)引導(dǎo)的比較

Complete-Trace 模型是田淵棟團(tuán)隊(duì)的基本 Searchformer 模型。該方法還提出了一種搜索動(dòng)態(tài)引導(dǎo)方法來(lái)提高其在推箱子任務(wù)上的性能，類(lèi)似于 Anthony 等人（2017）；Zelikman 等人（2022）的研究。

在訓(xùn)練 Searchformer 模型后，作者在新創(chuàng)建的自引導(dǎo)數(shù)據(jù)集上對(duì)其進(jìn)行微調(diào)。對(duì)于原始數(shù)據(jù)集中的每個(gè)推箱子競(jìng)賽，此處生成 32 個(gè)答案，并將最短的最佳答案納入新數(shù)據(jù)集。我們可以多次重復(fù)此過(guò)程。

通過(guò)這種方式，Searchformer 學(xué)會(huì)了生成更短的答案。表 5.4 將 Dualformer 與最多微調(diào) 3 步的 Searchformer 模型進(jìn)行了比較。Dualformer 在大多數(shù)指標(biāo)上與引導(dǎo)模型相當(dāng)或更好，同時(shí)僅使用不到 45.1% 的推理步驟。

該團(tuán)隊(duì)發(fā)現(xiàn)，每個(gè)引導(dǎo)步驟需要推出 3.2 × 10^6 個(gè)總響應(yīng)和 10^4 次迭代的額外微調(diào)。這意味著包括 8 × 10^5 次預(yù)訓(xùn)練迭代。Searchformer 步驟 3 總共需要 8.3 × 10^5 次訓(xùn)練迭代和 9.6 × 10^6 次 rollout，計(jì)算成本很高。相比之下，Dualformer 只需要一個(gè)由 8 × 10^5 次迭代組成的訓(xùn)練階段，沒(méi)有額外的 rollout 需求。

自動(dòng)模式

不僅能通過(guò)在 bos 之后注入控制 token 的方式來(lái)控制 Dualformer 的推理模式，還可以直接執(zhí)行采樣，使其自由確定操作模式，類(lèi)似于人類(lèi)決策的雙重過(guò)程。這種 Dualformer 被稱(chēng)為自動(dòng)模式。表 5.3 報(bào)告了結(jié)果。對(duì)于這里考慮的所有任務(wù)，自動(dòng)模式 Dualformer 也優(yōu)于 Complete-Trace 和 Solution-Only 模型。

大模型訓(xùn)練中的應(yīng)用：數(shù)學(xué)推理

作者展示了結(jié)構(gòu)化軌跡丟棄技術(shù)在訓(xùn)練大規(guī)模 LLM 解決數(shù)學(xué)問(wèn)題方面的有效性。具體來(lái)說(shuō)，作者使用了包含各種數(shù)學(xué)問(wèn)題和答案的數(shù)據(jù)集對(duì) Llama-3-8B 和 Mistral-7B 模型進(jìn)行微調(diào)，其中包含詳細(xì)的推理步驟。其中使用了一種軌跡丟棄技術(shù)，該技術(shù)也利用了數(shù)學(xué)問(wèn)題的推理軌跡的特定結(jié)構(gòu)。

最后，作者再對(duì)生成的模型與直接在數(shù)據(jù)集上微調(diào)的相應(yīng)基礎(chǔ)模型進(jìn)行基準(zhǔn)測(cè)試。

結(jié)果見(jiàn)表 5.6。作者共測(cè)試了 p 的四個(gè)值：0.1、0.2、0.3 和 0.4。結(jié)果表明，新研究所提出的訓(xùn)練策略使這兩個(gè) LLM 更加有效和高效。

首先來(lái)看 Mistral-7B 模型的結(jié)果。對(duì)于慢速模式推理，使用軌跡丟棄和隨機(jī)訓(xùn)練對(duì)模型進(jìn)行微調(diào)可以改進(jìn)直接在 Aug-MATH 數(shù)據(jù)集上微調(diào)的基線(xiàn)模型。當(dāng) p = 0.1 時(shí)，絕對(duì) Greedy@1 指標(biāo)提高了 1.7%（相當(dāng)于 10% 的相對(duì)性能提升），當(dāng) p = 0.2 和 0.3 時(shí)提高了 0.9%，當(dāng) p = 0.4 時(shí)提高了 0.1%。當(dāng) p = 0.1、0.2 和 0.3 時(shí)，新模型也優(yōu)于 Pass@20 指標(biāo)的基線(xiàn)模型，其中絕對(duì)正確率增加到 61.9%。在兩種評(píng)估方案下，推理軌跡的平均長(zhǎng)度隨著 p 的增加而下降。

同樣，對(duì)于快速模式下的推理，新模型也實(shí)現(xiàn)了更高的正確率。Llama-3-8B 模型也具有類(lèi)似的性能改進(jìn)趨勢(shì)。最后，為了供讀者參考，作者還列出了在原始 MATH 數(shù)據(jù)集上微調(diào)的 Mistral-7B 和 Llama-3-8B 模型的結(jié)果。