NuminaMath 7B TIR 榮獲 AIMO 進步獎 – 技術回顧

TL;DR

NuminaMath 7B TIR 透過正確解答 50 個隱藏競賽題目中的 29 題，贏得了首屆 AI Math Olympiad (AIMO) 進步獎。這證明了兩階段微調方案、大規模高品質數學數據，以及結合工具整合推理的自我一致性 (SC-TIR) 推理策略，可以顯著提升開源 LLM 在奧林匹克級別數學上的表現。

介紹 Numina – 一個開放的 AI4Maths 計畫

Numina 是一個於 2023 年底啟動的開源專案，旨在加速 AI 驅動的數學推理。該專案由 Jia Li、Yann Fleureau、Guillaume Lample、Stan Polu 和 Hélène Evain 發起，早期獲得 Mistral AI 的支持。2024 年初，Hugging Face 微調專家 Lewis Tunstall 和 Ed Beeching 加入團隊，並獲得來自 General Catalyst 和 Answer.ai 的額外支持。該合作的首個具體目標是 2024 AIMO 進步獎，這是一項 Kaggle 競賽，測試在 2024 年 2 月 23 日之前發布的開源權重模型在 50 個預選級數學問題（相當於 AMC 12/AIME）上的表現。

AI Math Olympiad (AIMO) 進步獎

AIMO 獎項旨在打造能夠在國際數學奧林匹克競賽中獲得金牌的 AI，並提供 500 萬美元的大獎。首個進步獎要求參與者每天在 Kaggle 平台上提交兩次解答，每次提交可在 P100 或兩張 T4 GPU 上運行最多九小時。題目是取自高中競賽層級的整數輸出問題，設有公開排行榜（50 題）和決定最終排名的隱藏私人排行榜。

獲勝方案架構

獲勝的系統包含三個核心組件：

  1. 微調後的推理代理 (Fine-tuned reasoning agent) – DeepSeekMath-Base 7B 經過兩階段微調，成為一個能將自然語言推理與 Python REPL 調用交織在一起的模型。
  2. SC-TIR 解碼演算法 – 一種新型推理程序，可生成多個候選解答，執行嵌入的 Python 程式碼，並在自我一致性檢查後應用多數投票。
  3. 強大的內部驗證 – 四個精選的驗證集（AMC、AIME 以及 MATH 基準測試的兩個子集）引導模型選擇，並防止對公開排行榜過擬合。

訓練過程在單個節點上使用八張 H100 GPU，使用開源技術棧 (TRL, PyTorch, vLLM, DeepSpeed) 完成，耗時約十小時。

兩階段訓練方案 (MuMath-Code)

該方案遵循 MuMath-Code 論文，包含：

  • 第 1 階段 – 思維鏈 (CoT) 微調 – 基礎模型從數十萬個帶有 CoT 風格解答的自然語言數學問題中學習，鼓勵逐步推理。
  • 第 2 階段 – 工具整合推理 (TIR) 微調 – 使用 GPT-4 以 ToRA 格式生成的合成數據集，每個問題都配有原理說明、Python 程式碼和執行結果。這教會模型為中間計算調用 Python REPL。

這兩個階段都採用全參數微調（非 LoRA/DoRA），並使用 2048-token packing 策略、梯度檢查點 (gradient checkpointing) 和 DeepSpeed ZeRO-3 分片。各階段的關鍵超參數相同：學習率 2e-5、Batch size 32、餘弦退火調度器 (cosine scheduler)，以及啟動比例 (warm-up ratio) 分別為 0.0 (第 1 階段) 或 0.1 (第 2 階段)。

數據就是一切

Numina 構建了兩個大型數據集：

  • 思維鏈數據集 – 數十萬個問題-解答對，來源於中文高中講義、美國考試 PDF 和國際奧林匹克檔案。流水線執行了 OCR、分割、英文翻譯並轉換為 CoT 格式。
  • 工具整合推理數據集 – 約 60,000 個問題（多為整數輸出）透過 GPT-4 流水線進行處理，生成 ToRA 風格的推理、執行程式碼並過濾掉不匹配的結果。每個問題最多重新生成三次以確保正確性。

作者計劃在未來幾週內開源完整數據集。

SC-TIR：自我一致性 + 工具整合推理

在推理過程中，由於 Kaggle 以隨機順序提供題目且 GPU 資源有限，模型會面臨高變異性。SC-TIR 透過以下方式緩解此問題：

  1. 將每個問題複製 N 次（獲勝運行時 N = 48）以創建一批多樣化的提示詞。
  2. 採樣 N 個完成結果，每個結果產生一個 Python 程式碼塊。
  3. 執行每個程式碼塊並擷取輸出或追蹤回溯 (tracebacks)。
  4. 重新採樣最多 M = 4 層深度，允許模型根據先前的錯誤進行自我修正。
  5. 刪除無效樣本並對最終數值答案進行多數投票。

透過 AutoGPTQ 進行 8-bit 量化對於 T4 GPU 至關重要：它將上傳時間減半，避免了 bfloat16 不相容問題，並在僅造成輕微準確率損失的情況下減少了 VRAM 消耗。

防止對公開排行榜過擬合

由於公開測試集僅包含 50 個問題，團隊構建了四個內部驗證套件：

  • AMC (83 個整數輸出問題) – 代表私人測試集；模型解題率約為 60-65%。
  • AIME (90 個問題) – 更高的難度以暴露失敗模式。
  • MATH level 4 (754 個問題) 與 MATH level 5 (721 個問題) – 從 5,000 個問題的 MATH 基準測試中篩選出的整數答案子集。

使用 5-10 個隨機種子進行重複評估以測量變異數（SC-TIR 通常為 1-3%），從而實現可靠的超參數調優。

未能最終入選的實驗

團隊在確定 MuMath-Code 流水線之前探索了幾種替代方案：

  • 帶有多數投票的純 CoT 模型。
  • MMOS (單步 Python) 模型，其表現停滯在 16/50。
  • Kahneman-Tversky 優化 (KTO) 在線微調，雖然將公開分數提高到 27/50，但最終模型未能完成時間限制。
  • 帶有程式碼執行獎勵的強化學習 (PPO 和 REINFORCE-LOO) – 顯示出具前景的獎勵曲線，但沒有可衡量的準確率增益。
  • 擴展到更大的骨幹網路 (InternLM-20B, CodeLlama-33B, Mixtral-8x7B) – 在 T4 GPU 上的推理速度太慢。
  • 模型合併技術 (DARE, TIES, WARP) – 導致內部指標下降。

啟示與未來方向

NuminaMath 7B TIR 的成功證明了 當配備高品質數學數據、兩階段微調機制以及透過 SC-TIR 進行強大推理時，開源 LLM 可以達到具競爭力的奧林匹克級別推理能力。該方法可擴展到更大的模型（團隊計劃為更大的骨幹網路發布數據集），並可應用於其他需要工具使用的領域（例如符號計算、科學編程）。

該專案的開放性質邀請貢獻者和合作夥伴來擴展數據集、改進解碼演算法或探索新的模型架構。持續的社群努力可以讓最終的 AIMO 目標——獲得金牌的 AI——更接近現實。

致謝

作者感謝 Thomas Wolf 和 Leandro von Werra 促進 Numina 與 Hugging Face 的合作，感謝 Hugo Larcher 提供 GPU 資源，感謝 Colin Raffel 提供關於模型合併的建議，以及 Omar Sanseviero 對部落格文章的反饋。額外支持來自 Mistral.ai、General Catalyst、Answer.ai 以及北京大學國際數學研究中心。

獲取模型與演示

Sources