DeepMath: 輕量級數學推理 Agent，搭配 smolagents

DeepMath 是一個基於 Qwen3-4B Thinking 構建並透過 Group Relative Policy Optimization (GRPO) 微調的數學推理 Agent。它透過在安全沙盒中執行簡潔的 Python 程式碼片段來取代冗長的思維鏈追蹤，從而在提升準確度的同時，將輸出長度減少多達 66%。

技術架構與工作流程

DeepMath 將小型語言模型與 Python 執行器結合，以卸載確定性計算並減少追蹤冗長。該系統透過 smolagents 函式庫實作，並使用 vLLM 作為推理引擎。

核心組件

  • Base Model: Qwen3-4B Thinking.
  • Agent Framework: 使用 smolagents 構建，以允許模型輸出一般標記或包含 Python 程式碼片段的特殊 agent 呼叫。
  • Execution Environment: 一個沙盒環境，具備匯入模組的允許清單及每個片段的超時限制，以確保安全並防止任意程式碼執行（例如，無檔案 I/O 或網路存取）。

推理過程

在推理過程中，模型會為中間計算步驟生成短小的 Python 程式碼片段。這些片段在沙盒中執行，結果會被折回推理追蹤，為後續步驟提供確定性結果給模型。

使用 GRPO 的訓練方法論

DeepMath 透過 TRL 函式庫使用 Group Relative Policy Optimization (GRPO) 進行微調。訓練過程著重於獎勵正確性與簡潔性，而非冗長。

獎勵結構

  • Accuracy Reward: 正確最終答案給予 +1。
  • Code Snippet Reward: 生成程式碼片段給予 +1，相對於準確度獎勵的權重為 10:1。
  • Length Penalty: 透過將 GRPO 完成候選限制為 5,000 個 token，鼓勵較短的輸出。

訓練優化

  • Temperature Scheduling: 使用線性排程 (T=1.2 到 T=0.7) 以平衡初期探索與後期穩定性。
  • In-context Learning: 模型會獲得四個已解決的範例，包含 agent 呼叫與執行器輸出，以學習所需的語法與回應模式。
  • Dataset: 使用 OpenMathReasoning 資料集的 Tool-Integrated Reasoning (TIR) 子集，專注於那些特別受益於外部工具使用的問題。

評估與效能

DeepMath 在四個資料集上進行基準測試：MATH500、AIME、HMMT 和 HLE。評估使用 majority@16 來確保穩健性，並測量平均輸出長度以評估簡潔性。

主要發現

  • Efficiency: 代理方法顯著降低輸出長度，相較於基線可減少多達 66%。
  • Synergy: 消融研究顯示，雖然單獨的代理推理可以降低長度，但當 GRPO 訓練與代理推理同時使用時，才能獲得最佳效果。
  • Accuracy: DeepMath 在具有挑戰性的資料集上展現更高的準確度，同時相比基線 Qwen3-4B-Thinking-2507 模型保持更短的追蹤。

意義與限制

將計算卸載到沙盒可以減少算術錯誤，並由於輸出變短而提升推理速度。然而，研究者指出以下幾項限制：

  • Scope: 目前的焦點僅限於小型模型與數學推理。
  • Generalization: 模型在競賽風格的數學問題上進行評估，可能無法遷移至形式證明或開放式的數學創造。
  • Security: 儘管有嚴格的沙箱與資源限制，執行生成的程式碼仍然是一種固有風險，需要謹慎的部署管理。

Sources