Stanford CS329A: Planning and Multi-Step Reasoning in AI Agents

概述

AI Agent 的多步推理需要將推理（規劃要做什麼）、行動（透過工具或搜尋來執行步驟）以及搜尋（根據回饋優化軌跡）進行協調整合。核心挑戰在於如何超越簡單的序列生成，進而鼓勵模型在尋求正確解法時展現出多樣性、探索性與效率。

Language Agent Tree Search (LATS)

LATS 透過將 Monte Carlo Tree Search (MCTS) 應用於語言模型軌跡，將推理、行動與規劃進行統一。它將線性的思維鏈（chain-of-thought）過程轉化為基於樹狀結構的搜尋，以優化通往最終答案的路徑。

LATS 框架

LATS 透過六個不同的階段進行運作：

  1. 選擇 (Selection)：使用 Upper Confidence Bound applied to Trees (UCT) 選擇一個節點進行擴展，以平衡新路徑的探索與已知高價值路徑的利用。
  2. 擴展 (Expansion)：模型從所選節點中採樣多個潛在的行動。
  3. 評估 (Evaluation)：每個產生的狀態會使用 LLM-as-a-Judge（提示模型給出 0-1 分）與 self-consistency score（特定行動被採樣的頻率）的加權平均值進行評分。
  4. 模擬 (Simulation)：模型會貪婪地擴展得分最高的狀態，直到達到終端狀態（成功或失敗）或達到計算預算。
  5. 反向傳播 (Backpropagation)：軌跡的結果會被用來更新通往該結果的所有狀態的值。
  6. 反思 (Reflection)：模型會分析軌跡以確定成功或失敗的原因，隨後將其附加到上下文（context）中以改進未來的搜尋。

效能與權衡

LATS 在 HotPotQA（多跳推理）與 WebShop（實際的電子商務任務）上展現了強大的結果，在無需微調的情況下達到了接近人類專家的效能。然而，由於樹狀結構的重複擴展與反向傳播，該方法會引入顯著的推理成本。此外，LATS 目前尚未解決不可逆的行動（例如：金融交易）問題，在這些情況下，試錯搜尋是不可行的。

SPRINT: Parallelizing Reasoning Traces

SPRINT 解決了前沿模型（如 DeepSeek-R1 或 Gemini）中序列化「思考」的低效率問題，在這些模型中，較長的推理鏈通常與較高的準確度相關，但會增加延遲與成本。

平行規劃與執行

SPRINT 基於以下觀察：許多推理步驟是相互獨立的，可以平行執行。該框架使用一種後訓練/微調方法，教導模型同時擔任 planner 與 executor：

  • 數據建立：來自 DeepSeek-R1 等模型的推理軌跡會由 GPT-4o 進行標註，以識別規劃步驟與執行步驟。接著，這些步驟會被組織成一個有向無環圖 (DAG) 以識別哪些步驟可以同時運行。
  • 微調 (Fine-tuning)：模型會在這些重新格式化的軌跡上進行監督式微調 (SFT)，教導它為平行規劃（例如：「Plan 1」與「Plan 2」）及其對應的執行輸出特定的標籤。
  • 推理 (Inference)：在測試時，模型會生成多個獨立的規劃；接著，這些規劃會被平行地執行（例如：透過多個 Python tool calls），並在模型進入下一個規劃階段之前，將結果同步回上下文 (context) 中。

關鍵發現

  • 效率：SPRINT 顯著減少了達到答案所需的序列化 token 數量，在某些任務上將序列化 token 數減少了約 40%。
  • 準確度：與僅追求效率的目標不同，結構化的平行思考過程實際上提升了在 MATH, Countdown, 與 GPQA Diamond 等基準測試上的準確度。
  • 泛化能力：在 MATH 數據上訓練的模型，其平行思考能力可以泛化到領域外 (out-of-domain) 的任務，無需進一步訓練。

SWiRL: Synthetic Multi-Step RL

SWiRL 專注於教導模型使用工具並進行多步推理，而無需在強化學習 (RL) 過程中執行工具所帶來的開銷與不穩定性。

SWiRL 方法論

SWiRL 將數據收集與模型優化分離：

  1. 離線合成生成 (Offline Synthetic Generation)：模型會迭代地生成多步軌跡（推理 $\rightarrow$ 工具呼叫 $\rightarrow$ 環境回應）。
  2. 過程標註 (Process Labeling)：LLM-as-a-Judge 會根據提議的行動/查詢的品質，為每個個別步驟分配一個獎勵值，而不論工具的實際輸出為何。
  3. 多步強化學習 (Multi-Step RL)：模型使用 RL 來優化在給定先前上下文 (context) 的情況下，單一行動的預期獎勵。關鍵在於，訓練期間並不執行工具；模型僅因根據預先收集的離線數據提出「好」的行動而獲得獎勵。

數據過濾與泛化

  • 過程 vs. 結果過濾 (Process vs. Outcome Filtering)：研究人員發現，過程過濾數據（保留被 LLM 判斷為「好」的步驟）比 結果過濾數據（僅保留具有正確最終答案的軌跡）在訓練上更有效。這顯示出學習正確的過程推理比單純模仿正確的最終結果更具價值。
  • 跨領域轉移 (Cross-Domain Transfer)：SWiRL 展現了卓越的泛化能力。在 GSM8K（使用計算機工具）上訓練的模型，在 HotPotQA（使用搜尋工具）上的表現有所提升，且反之亦然。這顯示模型正在學習多步工具呼叫與推理的一般性元能力 (meta-skill)，而非僅僅學習特定工具的 API。
  • RL vs. SFT：多步 RL 的表現顯著優於監督式微調 (SFT)，因為 RL 允許模型探索並從模仿數據中存在的錯誤軌跡中脫離出來。

Sources