補齊Transformer規劃短板又不放棄快速思考，Dualformer雙重優勢

機器之心報道

編輯：Panda、澤南

OpenAI ο1 模型的發佈掀起了人們對 AI 推理過程的關注，甚至讓現在的 AI 行業開始放棄卷越來越大的模型，而是開始針對推理過程進行優化了。今天我們介紹的這項來自 Meta FAIR 田淵棟團隊的研究也是如此，其從人類認知理論中獲得了靈感，提出了一種新型 Transformer 架構：Dualformer。

根據人類認知理論，人類的思考受到兩個系統控制：

近期有研究表明，如果將系統 2 過程整合進 Transformer 和大型語言模型中，就能顯著提升它們的推理能力。儘管如此，如果模型只是模仿系統 2 式的思考過程，那就需要遠遠更高的計算成本才能完成，同時響應速度也會大幅減慢。

在研究這一難題時，田淵棟團隊得到了一項驚人發現：在解決推理任務時，一種簡單的數據方案就足以實現即時動態的系統 1 和系統 2 配置。

基於此發現，他們提出了 Dualformer。這是一種可以輕鬆配置的 Transformer—— 用戶可以指定在推理過程中使用快速或慢速模式，在未指定時模型也可以自行決定。

具體而言，爲了模仿系統 2 推理過程，他們讓 Transformer 在包含推理軌跡和最終解答的數據上進行訓練。利用推理步驟的結構，他們設計了特定的軌跡丟棄策略，使得生成的軌跡類似於系統 1 在思考過程中採取的捷徑。在極端情況下，會丟棄整個軌跡並鼓勵 Transformer 繞過所有中間步驟，直接輸出最終解答。在訓練時，他們的策略是隨機選擇這些結構化的軌跡丟棄策略。

前提準備

他們的這項研究基於田淵棟團隊之前的另一項研究《Beyond A*: Better planning with transformers via search dynamics bootstrapping》，參閱機器之心報道《補齊 Transformer 規劃短板，田淵棟團隊的 Searchformer 火了》。爲了執行規劃，他們要訓練一個 Transformer 來建模一個 token 序列，而該序列則是以順序方式來表示該規劃任務、A* 算法的計算、由 A* 搜索得到的最優解。

圖 3.1 展示了其 token 化方法，其中示例是一個 3×3 迷宮的導航任務，目標是找到從起點到目標單元格的最短路徑。

A* 算法已經成功找到了最佳規劃。這裡使用一個 token 序列來表示該任務和迷宮結果，其也被用作 Dualformer 的提示詞。該解答由使用座標描述路徑的規劃 token 序列描述。A* 算法生成一個搜索軌跡序列，記錄執行的搜索動態，如圖 4.1 所示。

回想一下，A* 算法是一種在加權圖上的尋路算法。create 子句將節點（由後續座標表示）添加到搜索邊界中，close 子句將節點添加到該閉集。每個子句（create 或 close）後面都跟着 token x、y、c0 和 c1—— 分別表示節點的座標、自開始以來的成本值和啓發值。

結構化軌跡丟棄和隨機訓練

田淵棟團隊之前提出的 Searchformer 已被證明可以有效解決多種複雜的決策任務。但是，它仍有兩個不足。

1. 模型僅能以慢速模式運行並會輸出很長的推理鏈，這會極大延長推理時間。儘管可通過 bootstrapping（一種迭代優化技術，包含 rollout 循環和之後的微調過程）來提速，但這樣的過程會對計算資源產生顯著的額外需求。

2. Searchformer 很難生成多樣化的解答，因爲其經常會採樣相同的 rollout。舉個例子，在他們測試過的 1000 個 30×30 迷宮問題中，Searchformer 的推理鏈平均包含 1500 多個 token，而只能在 64 個響應中找到 7.6 條各不一樣的可行路徑。

爲了解決這些挑戰，他們提出了一個利用隨機化推理軌跡的訓練框架。該方法的靈感來自兩個研究方向：

這些觀察再加上 dropout 技術（在訓練時隨機丟棄神經網絡中的一些單元）的成功，促使該團隊研究了隨機化推理軌跡的作用，並且他們還希望通過利用結構化元素並選擇性地丟棄每個訓練示例的某些部分來簡化 A* 搜索軌跡。該方法的細節如下。

如圖 4.1 所示，A* 搜索軌跡包含 create 和 close 子句，每個子句都包括節點的座標及其到達起始位置和目標位置的（估計）成本。爲了推導得到 Dualformer，他們利用了搜索軌跡的結構，併爲每個訓練示例丟棄軌跡中的某些部分。其有三種自然的丟棄類型：

基於此，他們開發出了四個層級逐層遞進的丟棄策略：

圖 4.1 基於上述迷宮任務演示了這些策略。後面我們會看到，這些策略可有效地引導 Dualformer 學習更簡潔、更高效的搜索和推理過程。

爲了提升訓練數據的多樣性，他們沒有將丟棄作爲一個數據預處理步驟。而是在推理時間，對於一個數據批次中的每個訓練樣本，都從一個分類分佈 Cat (p_0, p_1, p_2, p_3, p_4) 中隨機抽取丟棄策略，其中 p_1, . . . , p_4 是執行 Level 1-4 丟棄的概率，p_0 是保持完整軌跡的概率。這種訓練框架可使 Dualformer 學習多個經過約簡的軌跡，即使對於單個訓練示例也是如此，因爲同一個示例可能出現在多個批次中。

可控式生成

Dualformer 具有一個非常吸引人的特性：在推理時，可以輕鬆地通過提示詞指定以快速或慢速生成模式運行。

該控制機制非常簡單：在標準提示詞之後添加一個 bos 和一個控制 token，其中控制 token 是 plan 或 create 中的一個。

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

而如果僅使用標準提示詞，則 Dualformer 將模仿人類決策的雙重過程 —— 根據情況，它會選擇一種分別對應於系統 1 和系統 2 的推理類型進行響應。

實驗

實驗的目標是解答以下三個問題：

1. Dualformer 在快速、慢速和自動模式下的表現是否優於相應的基線？

2. 在慢速模式下，Dualformer 是否能實現更快的推理，即輸出更短的軌跡？

3. 結構化的軌跡丟棄技術是否適用於在自然語言數據集上訓練的 LLM？

爲了解答問題 1 和 2，該團隊訓練了求解迷宮導航任務和緊密相關的推箱子（Sokoban）任務的 Transformer。爲了解答問題 3，他們微調了 LLama-3.1-8B 和 Mistral-7B 模型來解答數學問題。

導航任務：迷宮和推箱子

迷宮和推箱子任務使用的數據集與 Searchformer 研究的一樣。這裡就不再贅述，我們直接來看結論。

研究表明，Dualformer 可以根據控制指令選擇快速或慢速的運行模式。在快速模式下，它僅輸出最終規劃；在慢速模式下，它還會生成推理軌跡。該團隊在不同的模式下讓 Dualformer 對比了不同的基線。使用的指標包括生成規劃的正確性、最優性和多樣性、推理軌跡的長度等。

表 5.1 分別報告了在迷宮和推箱子任務上，Dualformer 和基線僅解答模型的性能。

可以看到，在生成正確和最優規劃方面，Dualformer 在 1-Solved-64 和 1-Optimal-64 指標上中都明顯優於基線。它在 3-Solved-64 和 3-Optimal-64 指標上也明顯超過了基線，這證明了 Dualformer 在規劃生成方面的穩健性。

尤其需要注意，隨着任務難度提升，Dualformer 的優勢也會增大。對於最大的 30×30 迷宮，Dualformer 的 1-Optimal-64 成功率是僅解答模型的 2.8 倍，在 3-Optimal-64 上是 2.97 倍。

Dualformer 的 SWC 分數也比基線高得多 —— 在每個環境中都高於 0.9。這表明 Dualformer 生成的每個單獨規劃的質量都很高，其成本非常接近最佳規劃。

在實驗考慮的所有問題上，Dualformer 還能穩定地生成更多樣化的規劃。比如在下面這個迷宮示例中，隨着迷宮規模的增加，Dualformer 的多樣性得分（即 64 個響應中不同但正確的規劃的平均數量）會增加。

一般來說，隨着迷宮規模增大，到達單個目標位置的可能路線也越來越多。這表明 Dualformer 學習了迷宮結構，而僅解答模型可能是記住了最佳規劃，因爲其多樣性得分在所有迷宮規模下都接近 1。

表 5.2 報告了 Dualformer 在慢速模式下運行時的結果。

相應的基線是 Complete-Trace 模型，它使用相同的架構並在具有完整 A* 搜索軌跡的數據上進行了訓練。除了之前報告的指標之外，該研究還報告了在所有 1000 個評估任務中彙總的 64 個響應的推理軌跡平均長度。結果表明，Dualformer 實現了更好的規劃能力和推理速度。它在所有正確性和最優性指標方面都優於 Complete-Trace 模型：包括解決率、最優率和 SWC。

此外，Dualformer 產生的推理軌跡明顯短於基線模型。平均而言，Dualformer 在五個任務中將軌跡長度減少了 49.4%。與以前一樣，與基線相比，Dualformer 還生成了更多不同的規劃。

Complete-Trace 模型是田淵棟團隊的基本 Searchformer 模型。該方法還提出了一種搜索動態引導方法來提高其在推箱子任務上的性能，類似於 Anthony 等人（2017）；Zelikman 等人（2022）的研究。

在訓練 Searchformer 模型後，作者在新創建的自引導數據集上對其進行微調。對於原始數據集中的每個推箱子競賽，此處生成 32 個答案，並將最短的最佳答案納入新數據集。我們可以多次重複此過程。

通過這種方式，Searchformer 學會了生成更短的答案。表 5.4 將 Dualformer 與最多微調 3 步的 Searchformer 模型進行了比較。Dualformer 在大多數指標上與引導模型相當或更好，同時僅使用不到 45.1% 的推理步驟。

該團隊發現，每個引導步驟需要推出 3.2 × 10^6 個總響應和 10^4 次迭代的額外微調。這意味着包括 8 × 10^5 次預訓練迭代。Searchformer 步驟 3 總共需要 8.3 × 10^5 次訓練迭代和 9.6 × 10^6 次 rollout，計算成本很高。相比之下，Dualformer 只需要一個由 8 × 10^5 次迭代組成的訓練階段，沒有額外的 rollout 需求。

自動模式

不僅能通過在 bos 之後注入控制 token 的方式來控制 Dualformer 的推理模式，還可以直接執行採樣，使其自由確定操作模式，類似於人類決策的雙重過程。這種 Dualformer 被稱爲自動模式。表 5.3 報告了結果。對於這裡考慮的所有任務，自動模式 Dualformer 也優於 Complete-Trace 和 Solution-Only 模型。

大模型訓練中的應用：數學推理

作者展示了結構化軌跡丟棄技術在訓練大規模 LLM 解決數學問題方面的有效性。具體來說，作者使用了包含各種數學問題和答案的數據集對 Llama-3-8B 和 Mistral-7B 模型進行微調，其中包含詳細的推理步驟。其中使用了一種軌跡丟棄技術，該技術也利用了數學問題的推理軌跡的特定結構。

最後，作者再對生成的模型與直接在數據集上微調的相應基礎模型進行基準測試。

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

首先來看 Mistral-7B 模型的結果。對於慢速模式推理，使用軌跡丟棄和隨機訓練對模型進行微調可以改進直接在 Aug-MATH 數據集上微調的基線模型。當 p = 0.1 時，絕對 Greedy@1 指標提高了 1.7%（相當於 10% 的相對性能提升），當 p = 0.2 和 0.3 時提高了 0.9%，當 p = 0.4 時提高了 0.1%。當 p = 0.1、0.2 和 0.3 時，新模型也優於 Pass@20 指標的基線模型，其中絕對正確率增加到 61.9%。在兩種評估方案下，推理軌跡的平均長度隨着 p 的增加而下降。

同樣，對於快速模式下的推理，新模型也實現了更高的正確率。Llama-3-8B 模型也具有類似的性能改進趨勢。最後，爲了供讀者參考，作者還列出了在原始 MATH 數據集上微調的 Mistral-7B 和 Llama-3-8B 模型的結果。