Dualformer：透過隨機推理軌跡學習實現可控的快速與慢速思考
https://arxiv.org/pdf/2410.09918v2
在人類認知理論中，人類思維受兩個系統主導：快速且直覺的系統1，以及較慢但更具推理性的系統2。近期研究表明，將系統2的思維過程融入Transformer模型（包括大型語言模型LLMs）能夠顯著提升其推理能力。然而，單純模擬系統2思維的模型需要更高的計算成本，反應速度也較慢。為了解決這項挑戰，我們提出了Dualformer，這是一種單一的Transformer模型，能夠無縫整合快速和慢速兩種推理模式。
Dualformer是透過在包含隨機推理軌跡的資料上進行訓練而獲得的，訓練過程中會隨機捨棄推理軌跡的不同部分。這些捨棄策略是根據推理軌跡的結構專門設計的，類似於分析我們的思維過程，並透過模式形成思維捷徑。在推理階段，我們的模型可以配置為僅輸出解決方案（快速模式），或同時輸出推理鏈和最終答案（慢速模式），也可以自動決定啟用哪種模式（自動模式）。
在所有情況下，Dualformer在性能和計算效率方面均優於相應的基準模型：(1) 在慢速模式下，Dualformer對未見過的30×30迷宮導航任務的最佳解決率達到97.6%，超越了在完整推理軌跡資料上訓練的Searchformer基準模型93.3%的表現，同時推理步驟減少了45.5%；(2) 在快速模式下，Dualformer完成這些任務的最佳率達到80%，顯著優於僅在答案資料上訓練的Solution-Only模型（其最佳率僅為30%）；(3) 在自動模式下，Dualformer實現了96.6%的最佳率，同時相比Searchformer減少了59.9%的推理步驟。
在數學問題上，我們的方法透過大型語言模型微調也取得了性能提升，表明這項技術具有超越特定任務模型的廣泛適用性。
1 導論
在心理學中，雙歷程理論（Wason 和 Evans，1974）解釋了思維可以透過兩種不同的歷程產生。通常，一種歷程是內隱的（自動且無意識的），另一種歷程是外顯的（受控且有意識的）。
著名的書籍《快思慢想》（Kahneman，2017）深入探討了雙歷程理論的概念。書中描述了兩種不同的思維方式：系統1對應於內隱歷程，它快速、直覺且帶有情感色彩；系統2對應於外顯歷程，它較慢、更具推理性，也更加邏輯化。Kahneman 將這項理論廣泛應用於各種認知與行為活動中，以說明這兩個系統如何影響我們的決策和推理歷程。如書中所述，系統1容易導致認知偏差，由於缺乏主動控制，容易產生系統性錯誤；而作為平衡，系統2更適合需要仔細分析、權衡和規劃的任務。
Transformer（Vaswani 等，2017）是一種序列建模工具，已成為包括大型語言模型（LLMs）在內的各類基礎模型的核心技術（Dosovitskiy，2020；Baevski 等，2020；Radford 等，2021；Touvron 等，2021；Hsu 等，2021；Touvron 等，2023；Dubey 等，2024）。在許多研究中，Transformer 被廣泛用於解決推理和規劃問題，例如 Zhou 等（2022）；Kojima 等（2022）；Pallagani 等（2022）；Valmeekam 等（2023a）；Chen 等（2024）；Gundawar 等（2024）；Wang 和 Zhou（2024）。有趣的是，雙歷程理論同樣適用於人工智慧領域，我們可以將 Transformer 的推理模式劃分為「快速」和「慢速」兩種。在快速模式下，Transformer 會直接輸出最終解決方案，不包含任何推理步驟；而在慢速模式下，模型會生成思維的中間步驟，例如尋找最短路徑時的搜尋軌跡，並與最終計畫一併輸出。這兩種推理模式與人類本身具備的兩種思維系統的優缺點非常相似。如先前研究（Wei 等，2022；Valmeekam 等，2023b；Lehnert 等，2024；Gandhi 等，2024；Saha 等，2024）所討論的，運行在快速模式下的模型計算成本較低，反應速度更快，但在準確性和最佳性方面不如慢速模式的模型。
這就引出一個問題：我們能否將快速和慢速兩種模式整合到基於Transformer的推理代理中，就像人類擁有兩種不同的思維系統一樣，讓它們相互補充？
已有許多方法被提出。一種流行的典範是首先建構純粹的系統2模型，然後透過微調使其像系統1一樣更有效率，例如將系統2的輸出「蒸餾」為系統1的形式（Yu 等，2024；Deng 等，2024），或透過知識蒸餾（Wang 等，2023；Shridhar 等，2023）或從符號系統中自舉（如 Searchformer（Lehnert 等，2024）、Stream of Search（Gandhi 等，2024））來提升現有系統2的推理效率。然而，在這些方法中，往往需要額外的微調，計算成本高昂，且難以在運行時靈活地讓最終系統在系統1和系統2之間切換。為了解決這項問題，Saha 等（2024）設計了一個顯式的元控制器，用於在兩個不同系統之間切換。
在本研究中，我們展示了一個令人驚訝的發現：一個簡單的資料處理方法就足以實現在解決推理任務時對系統1和系統2的即時配置。由此得到的模型 Dualformer 在推理過程中可輕鬆配置為運行在快速或慢速模式下；若未指定模式，模型還能自行決定使用哪種模式。具體而言，為了模擬系統2的推理歷程，我們的Transformer在包含推理軌跡和最終解決方案的資料上進行訓練。利用推理步驟的結構特徵，我們設計了特定的軌跡捨棄策略，使得生成的軌跡類似於系統1在思維歷程中所採取的「捷徑」。在極端情況下，我們完全捨棄整個推理軌跡，鼓勵Transformer直接輸出最終解決方案，跳過所有中間步驟。在訓練過程中，我們隨機選擇這些結構化的軌跡捨棄策略。詳見第4節。
我們首先將該框架應用於訓練一個編碼器-解碼器結構的Transformer模型，以解決路徑搜尋問題，其中推理軌跡由A*搜尋演算法生成。我們考慮兩個任務領域：迷宮導航和推箱子遊戲（Sokoban），與Lehnert等（2024）中的設定一致，並採用相同的標記化方案。有趣的是，我們發現這類問題對目前最先進的大型語言模型（如o1-preview和o1-mini）具有挑戰性，其輸出路徑常常「穿牆」（見附錄G中的範例）。
在每種推理模式下，Dualformer均優於既定的基準模型，在問題解決率和最佳率方面均取得更強的表現。此外，Dualformer顯著提升了生成計畫的多樣性，能夠找到更多通往目標的獨特路徑。值得注意的是，即使在慢速模式下，Dualformer也表現出很高的效率，生成的推理軌跡遠短於基準模型。接下來，我們將該框架應用於大型語言模型的微調，以回答數學問題。參照Yu等（2023）的方法，訓練樣本來自MATH資料集（Hendrycks等，2021），其答案由Llama-3.1-70B-Instruct模型重寫，以包含詳細的中間推理步驟。同樣，獲得的LLMs也展現出更高的效能和效率。
2 相關工作
學習規劃與推理為提升基於Transformer的模型在長週期任務中的規劃與推理能力，研究者已付出巨大努力。目前主要發展出兩類方法。
第一類方法利用現有的大型語言模型（LLMs）。例如，研究者訓練大型語言模型呼叫外部已有的符號求解器，如Ahn等（2022）、Besta等（2024）、Sel等（2023）、He-Yueya等（2023）、Liu等（2023）、Silver等（2024）的工作。Pallagani等（2022, 2024）研究了對大型語言模型進行微調以使用符號求解器；Schick等、Hao等（2024）對大型語言模型進行微調以使用外部工具；Hao等提出對大型語言模型進行微調，使其在世界模型（world model）內部進行推理與規劃。其中，一些研究嘗試將系統1與系統2的思維模式融入大型語言模型的推理歷程。Weston和Sukhbaatar（2023）提出了「系統2注意力」（System 2 Attention, S2A），這是一種更具深思熟慮性的注意力機制，旨在透過減少模型對上下文中虛假相關性的依賴，提升大型語言模型的推理能力。Yu等（2024）則將系統2的輸出「蒸餾」到系統1模型中，目的是將系統2方法生成的高品質結果重新編譯回大型語言模型的生成歷程中，而無需中間的推理標記序列。
第二類方法旨在從零開始訓練Transformer模型，使其能夠獨立進行規劃與推理（Lehnert等，2024；Saha等，2024；Gandhi等，2024），通常使用特定任務的語言表示。Lehnert等（2024）和Gandhi等（2024）均嘗試透過語言化地表示搜尋歷程來訓練大型語言模型進行搜尋：Lehnert等（2024）開發了Searchformer，該模型被訓練以模仿A*搜尋演算法在路徑規劃問題中的歷程；Gandhi等（2024）則將其模型應用於「倒數遊戲」（Countdown game）。在同期工作中，Saha等（2024）訓練了兩個獨立的模型，並透過一個外部元控制器進行管理：其中一個模型生成快速但無推理軌跡的響應，另一個則響應較慢但包含完整推理鏈。類似地，Lin等（2023）利用多個模型，透過代理式工作流實現慢速與快速思維。我們的研究與上述工作密切相關，但存在關鍵區別：與Lehnert等（2024）、Gandhi等（2024）在訓練資料中不修改推理軌跡的做法不同，我們的方法在訓練過程中對推理軌跡進行隨機化處理；與Saha等（2024）、Lin等（2023）不同，我們不使用任何顯式的控制器，也不為每種模式分別訓練兩個網路，而是將快速和慢速兩種模式的功能整合到單一模型之中。
利用大型語言模型生成合成資料大型語言模型（LLMs）已被廣泛用於多個領域的合成資料生成。例如，Wei等（2021）、Longpre等（2023）、Taori等（2023）透過從包含自然語言指令的多樣化模板中取樣，建構了合成指令資料集。這項方法也被應用於視覺領域（Liu等，2024b,a；Zhu等，2023；Brooks等，2023；Peng等，2023）。為了提升大型語言模型解答數學問題的能力，Yu等（2023）開發了一種方法，使用專門設計的提示詞對原始MATH資料集（Hendrycks等，2021）進行重寫、驗證和擴充。類似的方法也在其他研究中得到進一步探索，包括Yuan等（2023）、Luo等（2023）、Lee等（2023）、Yue等（2023）、Tong等（2025）的研究。
3 預賽
我們的工作建立在 Lehnert 等（2024）的研究基礎之上。為了實現規劃，我們訓練一個 Transformer 模型，用於建模一個由若干標記（tokens）組成的序列，該序列依序表示規劃任務、A* 演算法的計算歷程，以及從 A* 搜尋中得出的最佳解。標記化方法如圖 3.1 所示。
以一個簡單的範例說明：我們考慮一個 3×3 迷宮中的導航任務，目標是從起始單元格到目標單元格找到一條最短路徑，且不能穿過牆單元格。A* 演算法已成功確定了一條最佳路徑。我們使用一串標記來同時表達該任務和迷宮的結構，這些標記也作為 Dualformer 的輸入提示（prompt）。解則透過一個表示路徑座標的「計畫標記序列」（plan token sequence）來展示。
A* 演算法生成一個搜尋軌跡序列，記錄了搜尋歷程中的動態行為，如圖 4.1 所示。需要回顧的是，A* 演算法是一種在加權圖上的路徑搜尋演算法。「create」子句將節點（由後續座標表示）加入搜尋前沿（frontier），「close」子句則將節點加入已關閉集合（closed set）。每個子句，無論是 create 還是 close，後面都跟著 x、y、c0 和 c1 四個標記，分別表示該節點的座標、從起點出發的實際代價（cost-since-start）、以及啟發式估計值。關於 A* 演算法和標記化方法的更多細節，我們建議讀者參考 Russell 和 Norvig（2021）以及 Lehnert 等（2024）的研究。
4 結構化軌跡捨棄與隨機化訓練
Lehnert 等（2024）提出的 Searchformer 已被證明在解決多種複雜決策任務方面具有有效性。然而，它仍存在兩個重要局限。
首先，該模型僅在慢速模式下運行，輸出冗長的推理鏈，顯著增加了推理時間。雖然可以透過自舉（bootstrapping）方法（Lehnert 等，2024）來緩解這項問題——這是一種包含多次 rollout 與微調循環的迭代最佳化技術——但該歷程會帶來顯著增加的額外計算資源需求。
其次，Searchformer 在生成多樣化解決方案方面存在困難，經常重複取樣相同的推理路徑。例如，在我們測試的 1000 個 30×30 迷宮問題中，Searchformer 的推理鏈平均包含超過 1500 個標記，在 64 次響應中僅能找到 7.6 條不同的可行路徑（見第 5.1.2 節）。
為應對這些挑戰，我們提出了一種利用隨機推理軌跡的訓練框架。我們的方法受到兩方面研究的啟發：第一，我們注意到，儘管 Searchformer 是在完整的 A* 搜尋軌跡上訓練的，但它生成的推理鏈更短，僅粗略勾勒出搜尋歷程；第二，研究表明，人類在決策時常依賴「捷徑」和模式，這項概念被稱為系統1思維（Kahneman，2017）。這些觀察，結合 dropout 技術（Hinton，2012；Srivastava 等，2014）的成功經驗——即在訓練過程中隨機捨棄類神經網路中的部分單元——促使我們探索在框架中使用隨機推理軌跡的可能性。我們旨在利用 A* 搜尋軌跡的結構特徵，針對每個訓練樣本有選擇性地捨棄部分軌跡內容。具體方法如下所述。
如圖 4.1 所示，A* 搜尋軌跡包含「create」和「close」兩個子句，每個子句都包括節點的座標及其（估計的）從起點出發的實際代價和到目標的啟發式代價。為了得到 Dualformer，我們利用搜尋軌跡的結構，在每個訓練樣本中捨棄其部分資訊。共有三種自然的捨棄類型：
• D1：捨棄一個 close 子句
• D2：捨棄子句中的代價標記（cost tokens）
• D3：捨棄一個 create 子句
基於這些基本類型，我們設計了四個層級的捨棄策略，每一級都在前一級的基礎上進一步增強：
• 一級策略（Level 1）：從搜尋軌跡中刪除所有的 close 子句。
• 二級策略（Level 2）：在一級策略基礎上，進一步刪除所有子句中的代價標記。
• 三級策略（Level 3）：更為激進，額外隨機捨棄 30% 的 create 子句。
• 四級策略（Level 4）：最終策略，直接捨棄整個推理軌跡。
透過在訓練過程中隨機應用這些結構化的捨棄策略，模型被鼓勵學習從不完整或高度簡化的推理歷程中直接推導出正確答案，從而模擬系統1的快速直覺思維，同時保留系統2的完整推理能力。
圖4.1以之前提到的迷宮任務為例，展示了我們的捨棄策略。直觀來看，一級捨棄（Level 1）指示Dualformer有效跳過A*搜尋中的「關閉集」（close-set）計算；二級捨棄（Level 2）促使Dualformer同時跳過「關閉集」和代價計算；三級和四級捨棄則鼓勵Dualformer省略部分或全部的搜尋步驟。我們將在第5節中展示，這些策略有效地引導Dualformer學習到更簡潔、高效的搜尋與推理歷程。
為了提升訓練資料的多樣性，我們並未將捨棄操作作為資料預處理步驟。相反，在訓練過程中，對於每個批次中的每一條訓練樣本，我們從一個分類分布 Cat(p₀, p₁, p₂, p₃, p₄) 中隨機取樣捨棄策略，其中 p₁ 到 p₄ 分別表示採用一級到四級捨棄策略的概率，p₀ 表示保留完整推理軌跡的概率。這種訓練框架使得Dualformer即使對同一條訓練樣本，也能從多種簡化後的軌跡中進行學習，因為同一樣本可能在不同批次中以不同捨棄方式反覆出現。
與標記遮罩技術的比較感興趣的讀者可能會想到，我們的訓練框架是否類似於BERT等著名大型語言模型所使用的標記遮罩（token masking）技術（Devlin, 2018；Liu, 2019；Song, 2019；Gauthier 和 Levy, 2019；Sinha 等, 2021；Kitouni 等, 2024）。然而，我們的方法與這些遮罩技術存在顯著差異。首先，標準的遮罩技術通常是對序列中的輸入標記進行均勻的隨機遮罩，而我們的捨棄策略僅應用於搜尋軌跡部分。其次，被遮罩的大型語言模型通常採用雙向注意力機制，並以預測被遮罩的標記為目標；而Dualformer使用的是因果注意力機制（causal attention），其訓練目標僅為下一標記預測，整體目標是提升模型的推理與規劃能力。在計算效率方面，我們的訓練歷程也更具優勢：捨棄標記會縮短輸入序列，從而節省計算資源。例如，在8塊Tesla V100 32GB GPU上訓練Dualformer完成30×30迷宮任務需要30小時，而若使用完整的推理軌跡則需要36小時。更多訓練細節見第5節及附錄A。
4.1 可控生成
Dualformer的一個吸引人的特性是：在推理階段，它可以輕鬆透過提示詞控制，選擇以「快速」或「慢速」模式生成結果。控制機制極為簡單：我們在標準提示（包含環境和任務描述）的基礎上，添加起始標記 bos 和一個控制標記（control token），該控制標記為 plan 或 create。若使用 plan，Dualformer將進入快速模式，直接輸出最終路徑規劃，跳過所有中間推理步驟；反之，若在bos後插入 create，Dualformer將進入慢速模式，生成完整的推理軌跡及最終方案。具體範例請參見附錄B。如果僅使用標準提示（不加控制標記），Dualformer將模擬人類決策的雙歷程機制——根據具體情境自主決定生成何種類型的響應，分別對應系統1和系統2的推理模式。
5 實驗
我們的實驗旨在回答以下問題：
1. Dualformer 在快速模式、慢速模式和自動模式下是否優於相應的基準模型？它是否能生成更多樣化的規劃方案？
2. 在慢速模式下，Dualformer 是否能實現更快的推理，即輸出更短的推理軌跡？
3. 結構化軌跡捨棄技術是否可推廣到在自然語言資料集上訓練的大型語言模型（LLMs）？
我們在第5.1節中回答問題1和2，其中我們訓練Transformer模型來解決迷宮導航任務以及相關的推箱子（Sokoban）遊戲，方法類似於Searchformer（Lehnert等，2024）。為回答問題3，我們在第5.2節中對Llama-3.1-8B和Mistral-7B模型進行微調，使其能夠解決數學問題。
5.1 導航任務：迷宮與推箱子
我們沿用Lehnert等（2024）的方法，考慮迷宮和推箱子任務，並使用相同的資料集。兩個資料集均包含10⁵個帶有完整A*搜尋軌跡的訓練樣本。A*演算法的實現具有非確定性：當路徑代價相同時，演算法會隨機選擇；子節點的擴展順序也是隨機的。迷宮的尺寸從15×15變化到30×30。對於所有迷宮任務，我們隨機生成30%–50%的牆單元格作為障礙物，並隨機取樣起點和終點位置。推箱子地圖大小為7×7，包含兩個隨機放置的碼頭（docks）、兩個箱子和一個工人位置。我們還在地圖內部隨機添加兩個額外的牆單元格。關於地圖生成的具體流程和範例圖像，詳見附錄A.1。
首先，我們在第4.1節中展示了Dualformer可以被明確控制，以在快速或慢速模式下運行：在快速模式下僅輸出最終方案，在慢速模式下則生成完整的推理軌跡。在第5.1.1和5.1.2節中，我們將Dualformer在每種模式下分別與相應的基準模型進行比較。我們採用多種指標系統性地評估模型性能，包括生成方案的正確性、最佳性、多樣性，以及推理軌跡長度等。最後，在第5.1.5節中我們對設計選擇進行消融實驗。
超參數設定我們採用與Lehnert等（2024）相同的編碼器-解碼器架構實例化Dualformer。編碼器基於T5架構（Raffel等，2020）並引入旋轉位置編碼（rotary embeddings），解碼器則採用GPT風格的架構。迷宮環境使用1500萬參數（15M）的模型，推箱子環境使用4600萬參數（46M）的模型。所有模型均在10萬個訓練樣本上進行訓練。迷宮環境訓練4×10⁵輪次，推箱子環境訓練8×10⁵輪次。模型架構及其他超參數的詳細資訊見附錄A。
我們使用第4節所述的結構化軌跡捨棄策略訓練Dualformer。對於每個訓練樣本所採用的捨棄策略，我們嘗試了三組概率設定：(1) {p₀ = 0.45, p₁ = p₂ = p₃ = 1/6, p₄ = 0.05}，(2) {p₀ = 0.8, p₁ = p₂ = p₃ = p₄ = 0.05}，(3) {p₀ = 0.7, p₁ = 0.05, p₂ = p₃ = 0.1, p₄ = 0.05}，並選擇驗證誤差最低的一組。最終選擇結果為：迷宮任務採用第(1)組，推箱子任務採用第(3)組。
基準模型對於快速模式，我們的基準是「僅解模型」（Solution-Only model），該模型與Dualformer架構相同，但僅在包含最佳最終解、不含任何推理軌跡的序列資料上進行訓練。對於慢速模式，基準是「完整軌跡模型」（Complete-Trace model），即在包含完整A*搜尋軌跡的資料上訓練的模型。該模型在Lehnert等（2024）中被稱為「搜尋增強模型」（search augmented model），也是未進行搜尋動態自舉（bootstrapping）的原始Searchformer模型。我們還將在第5.1.2節中將Dualformer與經過自舉的模型進行比較。所有這些模型的參數量相同：迷宮問題為1500萬，推箱子問題為4600萬。
評估指標我們使用兩個指標評估模型是否生成正確且最佳的方案：1-Solved-64 和 1-Optimal-64。具體而言，對於每個評估任務（例如一個迷宮或一個推箱子遊戲），我們從訓練好的模型中隨機取樣64個響應。無論生成的推理軌跡如何，我們僅解析並評估最終方案部分。如果64個方案中至少有一個是正確的（即可行且能到達目標位置），則該任務在1-Solved-64指標下標記為成功；如果至少有一個方案是最佳的，則在1-Optimal-64指標下標記為成功。我們對1000個未見過的評估任務重複此歷程，並報告平均成功率。
為考察各方法的穩健性，我們還報告了3-Solved-64和3-Optimal-64指標：當至少有3個方案正確或最佳時，任務才被標記為成功。
此外，我們採用「成本加權成功率」（Success Weighted by Cost, SWC）（Wu等，2019）來衡量生成方案的整體品質，基於其代價進行聚合：SWC = (1/64) Σ₆₄ᵢ₌₁ I(方案i正確) · (c∗ / cᵢ) 其中I為指示函數，c∗為最佳方案的代價，cᵢ為第i個方案的代價。顯然，SWC得分越高，生成的方案越接近最佳；若所有生成方案均為最佳，則SWC達到最大值1。
最後，為量化生成方案的多樣性，我們統計每個任務在64個響應中產生的唯一正確方案的數量，並報告在1000個評估任務上的平均值。
5.1.1 快速模式
表5.1報告了Dualformer與基準模型「僅解模型」（Solution-Only model）在迷宮和推箱子任務上的性能表現。在生成正確和最佳方案方面，Dualformer在1-Solved-64和1-Optimal-64兩項指標上均顯著優於基準模型。同時，在3-Solved-64和3-Optimal-64指標上也明顯超越基準，表明Dualformer在方案生成方面具有更強的穩健性。
特別是，隨著任務難度的增加，性能差距也隨之擴大。對於最大的30×30迷宮任務，Dualformer的1-Optimal-64率是「僅解模型」的2.8倍，3-Optimal-64率更是達到其2.97倍。Dualformer的SWC得分也顯著高於基準模型，在所有環境中的得分均超過0.9，表明Dualformer生成的每一個方案品質都很高，其代價非常接近最佳方案。
此外，對於所有考慮的問題，Dualformer始終能生成更多樣化的方案。在附錄C中，我們展示了一個迷宮範例，並繪製了Dualformer與基準模型生成的唯一正確路徑。一個有趣的觀察是，Dualformer的多樣性得分（即在64次響應中生成的唯一正確方案的平均數量）隨著迷宮尺寸的增大而增加。直觀來看，迷宮越大，通往同一目標位置的可能路徑就越多。
這表明Dualformer學習到了迷宮的結構特徵，而「僅解模型」可能只是在記憶最佳路徑，因為其在所有迷宮尺寸下的多樣性得分均接近1。
5.1.2 慢速模式
表5.2報告了Dualformer在慢速模式下運行時的結果。對應的基準模型是「完整軌跡模型」（Complete-Trace model），該模型使用相同的架構，並在包含完整A*搜尋軌跡的資料上進行訓練。除了此前報告的各項指標外，我們還統計了在所有1000個評估任務中，64次響應所生成推理軌跡的平均長度。
結果表明，Dualformer在規劃能力和推理速度兩方面均實現了提升。在所有正確性和最佳性指標上——包括問題解決率、最佳率以及SWC得分——Dualformer均優於完整軌跡模型。此外，Dualformer生成的推理軌跡明顯短於基準模型。平均而言，在五項任務中，Dualformer將推理軌跡長度減少了49.4%。與之前一致，Dualformer生成的方案也比基準模型更具多樣性。具體範例請參見附錄C。
與搜尋動態自舉方法的比較
完整軌跡模型是Lehnert等（2024）中提出的Searchformer基礎模型。該研究還提出了一種「搜尋動態自舉」（search dynamics bootstrapping）方法，用於提升模型在推箱子（Sokoban）任務上的表現，類似於Anthony等（2017）和Zelikman等（2022）中的方法。在訓練完Searchformer模型後，會在一個新建構的自舉資料集上對其進行微調。具體而言，對於原始資料集中的每個推箱子關卡，生成32個響應，並將其中最短的最佳響應加入新的訓練資料集。該歷程可重複多次，從而使Searchformer逐步學會生成更短的響應。
表5.4將Dualformer與經過最多3步自舉微調的Searchformer模型進行了比較。結果顯示，Dualformer在大多數指標上達到或優於經過自舉的模型，同時使用的推理步驟減少了超過45.1%。
值得注意的是，每一步自舉歷程都需要生成總計3.2×10⁶個響應樣本，並額外進行10⁴輪次的微調。這意味著，包含最初的8×10⁵輪次預訓練，Searchformer在第3步自舉時總共需要8.3×10⁵輪次的訓練和9.6×10⁶次rollout，計算成本非常高。相比之下，Dualformer僅需一個訓練階段，共8×10⁵輪次，且無需任何額外的rollout歷程，顯著降低了計算開銷。
5.1.3 自動模式：雙歷程
除了透過在 bos 後注入控制標記來控制 Dualformer 的推理模式外，我們也可以直接對其取樣，使其自由決定運行模式，類似於人類決策中的雙歷程機制。我們將這種模式稱為 Dualformer 的自動模式（auto mode）。表 5.3 報告了相關結果。在我們考慮的所有任務中，自動模式下的 Dualformer 同樣優於完整軌跡模型和僅解模型。
一個值得探討的有趣問題是：Dualformer 是否能夠根據問題的難度自動調整其運行模式？為探究這點，我們對1000個未見過的迷宮（牆密度在0.3至0.5之間變化）各自生成了64個自動模式下的響應，然後分析了所有可行解中屬於慢速模式路徑的佔比。結果如圖 5.1 所示。隨著牆密度的增加（意味著迷宮可能更具挑戰性），慢速模式路徑的比例也隨之上升。這可能由兩個因素導致：(1) 解決更困難的問題本身需要更慢的思考歷程；(2) Dualformer 更頻繁地選擇了慢速模式。類似地，我們還觀察到，隨著迷宮尺寸增大、問題難度提升，Dualformer 也持續地更多地採用慢速思考模式。
5.1.4 泛化性能
到目前為止，我們所有實驗的測試均在與訓練集相同尺寸和牆密度的迷宮上進行。接下來，我們進一步考察Dualformer在分佈外（OOD）場景下的泛化能力。在本節中，我們考慮一個在20×20尺寸迷宮上訓練的Dualformer模型，訓練時的牆密度在0.3至0.5之間均勻隨機取樣。在測試階段，我們將牆密度從0.1變化到0.6。表5.5展示了慢速模式下的測試結果，每種情況測試50個未見過的迷宮實例。
隨著牆密度的增加，迷宮難度上升，同時輸入提示（prompt）的長度也隨之增加。因此，對於分佈內測試情況（牆密度為0.3、0.4、0.5），Dualformer（慢速模式）實現了接近100%的最佳率。然而，當牆密度增加到0.6（分佈外情況）時，性能有所下降。有趣的是，在較低的牆密度（如0.1和0.2）下，Dualformer（慢速模式）未能成功解決任何一個迷宮。我們的推測是：當牆密度過低時，提示資訊過短，導致模型難以泛化。
除了改變牆密度，我們還測試了Dualformer在矩形迷宮（例如高度為20，寬度為10）上的慢速模式表現。模型仍是在20×20的方形迷宮上訓練的。表5.6展示了測試結果。令人驚訝的是，Dualformer在這些情況下表現出良好的泛化能力。
5.1.5 消融研究
如第4節所述，用於訓練Dualformer的隨機化軌跡來源於多種軌跡捨棄策略的組合，而這些策略的組合方式有多種可能。本節中，我們對所做設計選擇進行消融分析。
首先，為了實現快速和慢速兩種推理模式的運行，一種簡單的替代方法是：在訓練時混合使用「僅解資料」和「完整軌跡資料」，即在我們的隨機化策略中令 p₁ = p₂ = p₃ = 0。我們將此類變體稱為 Mix-p 模型，其中 p 表示訓練資料中「僅解資料」所佔的比例。需要注意的是，「僅解模型」本質上是 Mix-1 模型，而「完整軌跡模型」則是 Mix-0 模型。接下來，我們將Dualformer與不同p值的Mix-p模型在兩種推理模式下進行比較。
其次，我們的捨棄策略是分層級、結構化設計的。例如，二級捨棄是在一級策略的基礎上進一步擴展的。我們進一步探究：當在某個特定層級停止捨棄策略時，模型性能將如何變化。
與 Mix-p 模型的比較圖D.1對比了Dualformer與Mix-p模型的1-Optimal-64率。我們測試了8個p值：0（等同於完整軌跡模型）、0.05、0.1、0.2、0.4、0.6、0.8 和 1.0（等同於僅解模型）。在兩種推理模式下，Dualformer在我們考慮的五項任務中均優於所有Mix-p模型。事實上，在其他各項指標上，Dualformer也全面優於Mix-p模型，詳見附錄D。特別地，在慢速模式下，Dualformer是生成推理軌跡最短的模型，即推理效率最高。
隨機化策略的組合分析我們還將Dualformer與其變體進行比較，這些變體在特定層級停止捨棄策略的應用。對於迷宮環境，我們固定每個訓練樣本使用完整推理軌跡的概率為 p₀ = 0.5，並調整其他捨棄策略的概率；對於推箱子環境，我們固定一級捨棄的概率為 p₀ = 0.05。表E.1列出了所有模型所使用的不同捨棄策略的概率分佈，表5.7展示了實驗結果。
需要注意的是，這些變體無法在快速模式下運行，因為它們未在任何四級捨棄（即完全捨棄軌跡）的資料上進行訓練。因此，我們僅報告其在慢速模式下的性能。隨著捨棄策略層級的提升，生成的推理軌跡長度逐漸縮短。在其他指標方面，採用一級+二級+三級捨棄的模型與Dualformer性能相當。然而，Dualformer在推理軌跡更短的同時，還具備在快速模式下運行的能力，展現出更全面的優勢。
5.2 應用於大型語言模型訓練：數學推理
在本節中，我們展示結構化軌跡捨棄技術在訓練大規模大型語言模型（LLMs）解決數學問題方面的有效性。具體而言，我們使用一個包含多種數學問題及附帶詳細推理步驟的答案的資料集，對 Llama-3-8B 和 Mistral-7B 模型進行微調，並同樣利用數學問題推理軌跡的特定結構設計軌跡捨棄方法。我們將微調後的模型與直接在該資料集上微調得到的基準模型進行對比。
資料集我們使用一個名為 Aug-MATH 的資料集評估所有模型，該資料集源自 MATH 資料集（Hendrycks 等，2021），原始資料集包含 7500 個訓練樣本和 5000 個測試樣本的數學問題與解答。沿用 Yu 等（2023）的方法，我們呼叫 Llama-3.1-70B-Instruct 模型，按照指定格式將原始答案重寫，使其包含更詳細的中間推理步驟。為了提升推理軌跡的多樣性，我們對每個問題以溫度 0.7 和 top-p = 0.9 取樣 4 個不同的大型語言模型響應。最終得到的資料集包含 30,000 個訓練樣本和 5,000 個測試樣本。附錄 F.1 展示了我們用於答案重寫的提示模板，以及一個重寫前後的具體訓練範例。
結構化軌跡捨棄與隨機化由 Llama-3.1-70B-Instruct 重寫後的數學問題答案平均包含 6 到 10 個中間推理步驟，每個步驟可能包含多個句子。我們採用與第4節中提出的框架類似的隨機化訓練方法：在每個訓練批次中，對每個取樣的訓練樣本，我們以概率 p 隨機捨棄其每一個中間推理步驟。下文將報告在不同 p 值下的實驗結果。
超參數設定我們對兩個基礎模型——Mistral-7B 和 Llama-3-8B——均進行兩輪微調，批次大小（batch size）為 32。最佳化器採用 AdamW（Loshchilov 和 Hutter，2019），Mistral 模型的學習率為 8e−6，Llama-3 模型為 5e−6。完整的訓練細節，包括學習率及其他超參數的選擇方法，詳見附錄 A.2。微調過程中使用的提示模板與 Yu 等（2023）一致，具體見附錄 F.2。
評估方法我們採用以下兩個評估指標：
Greedy@1（Dubey 等，2024；Yu 等，2023）：對每個問題，使用溫度為 0 生成 1 個響應，並判斷其答案是否正確。
pass@20（Chen 等，2021）：對每個問題，以溫度 0.5 隨機取樣 20 個響應，統計其中至少有一個正確的比例。
作為參考，我們也報告了在原始 MATH 資料集上直接微調這些模型的結果。
結果結果如表 5.8 所示。我們測試了四個 p 值：0.1、0.2、0.3 和 0.4。實驗結果表明，所提出的訓練策略使兩個大型語言模型在效能和效率方面均得到提升。
我們首先分析 Mistral-7B 模型的結果。在慢速模式推理下，使用軌跡捨棄和隨機化訓練進行微調的模型，優於直接在 Aug-MATH 資料集上微調的基準模型。當 p = 0.1 時，Greedy@1 指標絕對提升了 1.7%（相當於相對性能提升 10%）；當 p = 0.2 和 p = 0.3 時，提升為 0.9%；當 p = 0.4 時，提升為 0.1%。
在 Pass@20 指標上，當 p = 0.1、0.2 和 0.3 時，我們的模型也優於基準模型，最高絕對正確率達到 61.9%。在兩種評估方式下，隨著 p 值增大，推理軌跡的平均長度逐漸縮短。
同樣，在快速模式下的推理中，我們的模型也實現了更高的正確率。Llama-3-8B 模型也表現出類似的性能提升趨勢。
最後，為便於讀者參考，我們列出了 Mistral-7B 和 Llama-3-8B 模型在原始 MATH 資料集上微調的結果。這些模型的性能明顯落後於在 Aug-MATH 上微調的模型。
6 結論
我們提出了一種簡單且易於實現的框架，用於訓練Transformer模型解決推理與規劃任務。我們深入分析了推理軌跡的結構，並設計了相應的捨棄策略，以模擬人類思維歷程中的「捷徑」行為。透過對訓練樣本隨機應用這些捨棄策略，所得到的模型Dualformer可在推理時被控制以快速模式、慢速模式或自動模式運行——在自動模式下，模型自行決定採用何種推理方式。Dualformer在迷宮導航任務和推箱子遊戲中均實現了性能提升，同時減少了所需的推理步驟數量。
值得注意的是，我們的方法不僅限於從零開始訓練特定任務的模型。我們將相同思想的技術應用於大型語言模型（LLMs）的微調，以解答數學問題，並同樣獲得了性能提升。此外，所提出的框架還能降低計算開銷，因為在捨棄推理軌跡後，輸入序列長度得以縮短。
未來的研究可以進一步探究該方法是否有助於模型的擴展性，並嘗試結合課程學習（curriculum learning）和分層規劃（hierarchical planning）等方法，使Dualformer能夠適應更複雜的任務場景。
附錄A：網路架構和超參數
我們使用與Lehnert等（2024）相同的編碼器-解碼器Transformer架構用於Dualformer。它首先將每個標記轉換為一個獨熱向量，然後透過嵌入層將其轉換為一組向量。嵌入向量隨後透過Lehnert等（2024，圖4）中展示的後續層。我們使用RoPE嵌入進行位置編碼，並且在我們的架構中沒有使用dropout。
對於模型大小、架構參數和批次大小，我們使用與Lehnert等（2024）相同的設定。具體來說，對於迷宮任務，我們使用一個包含1500萬參數的模型，該模型包含3個注意力頭和6層，隱藏層大小為64。我們使用AdamW最佳化器（Loshchilov和Hutter，2019）進行模型最佳化，學習率為2.5e-4，批次大小為16，其中β₀和β₁分別設定為0.9和0.99。在前2000次梯度更新中，我們採用了線性預熱策略。此後，我們使用餘弦學習率排程器（Loshchilov和Hutter，2016）。
對於Sokoban任務，我們使用一個包含4600萬參數的模型，該模型包含4個注意力頭和8層，隱藏層大小為96。我們使用64的批次大小和7.5e-5的學習率，而其他超參數與上述相同。
A.1 迷宮和推箱子
我們使用與Lehnert等（2024）相同的資料集，該資料集可在https://github.com/facebookresearch/searchformer中找到。迷宮資料集包含1000萬個範例，推箱子資料集也包含1000萬個範例，我們按照mongodb中的「id」字段對它們進行排序，並使用前10萬個範例（遵循Lehnert等（2024）的相同方法）。為了讀者的參考，資料集是按以下方式生成的。對於迷宮任務，首先隨機指定30-50%的單元格為牆壁。接下來，隨機選擇一個起點和一個目標位置。然後，應用A*演算法生成最佳計畫。對於推箱子任務，我們使用一個7×7的網格地圖，其中隨機插入兩個牆壁單元格作為障礙物。此外，還隨機放置了兩個碼頭、兩個箱子和兩個工人位置。一旦生成了一個遊戲，只有當它能被A*演算法解決時，才會將其添加到資料集中。
A.2 數學推理
我們使用https://github.com/meta-llama/llama-recipes提供的實現來微調模型。我們對所有模型進行2個周期的訓練，使用32的批次大小。對於Llama模型，我們使用AdamW最佳化器，學習率為5×10⁻⁶；對於Mistral模型，學習率為8×10⁻⁶。學習率的選擇歷程如下：我們在2×10⁻⁶、5×10⁻⁶、8×10⁻⁶這3個值中進行搜尋，並選擇使驗證損失最低的學習率。然後，我們使用選定的學習率在完整的訓練資料集上重新訓練模型，並報告結果。對於其他超參數，我們使用預設值。具體來說，我們沒有使用線性學習率預熱、權重衰減，也沒有使用多步梯度累積。我們為AdamW使用了以下參數：betas=(0.9, 0.999)，eps=1e⁻⁸，γ=0.85（乘法逐步學習率衰減），並採用「packing」作為批次策略。
C 生成計畫的多樣性
Dualformer在發現獨特可行解方面優於基準模型。為了直觀展示這點，我們選取了一個迷宮任務範例，並使用快速模式的Dualformer生成了64個響應。圖C.1繪製了快速模式Dualformer發現的所有獨特可行路徑，以及僅解決方案基準（64個響應）找到的路徑。快速模式的Dualformer識別出了42條獨特可行路徑，而僅解決方案模型僅找到3條。同樣，圖C.2比較了慢速模式的Dualformer和完整軌跡（Searchformer）基準。慢速模式的Dualformer發現了39條獨特可行路徑，而完整軌跡模型僅找到17條。