來源 | 新智元
編輯 | KingHZ 定慧
AI數學奧賽第一名「答卷」終於公佈！
論文地址：https://arxiv.org/abs/2504.16891
參加本次Kaggle比賽、軟體工程師Chan Kha Vu，則盛讚道：這些模型太不可思議了！從基礎的Qwen模型訓練開始，甚至都不是推理模型。而且沒有利用強化學習！
英偉達團隊參賽的模型叫做OpenMath-Nemotron系列，使用OpenMathReasoning Dataset進行訓練，共釋出了四種引數：
-
OpenMath-Nemotron-1.5B
-
OpenMath-Nemotron-7B
-
OpenMath-Nemotron-14B-Kaggle（AIMO-2 Kaggle競賽中使用的模型）
-
OpenMath-Nemotron-32B
這些模型在流行的數學基準測試中都取得了最好的成績。
甚至1.5B的OpenMath-Nemotron模型，超越14B的DeepSeek-R1蒸餾模型！
圖1：AIME和HMMT競賽中的數學問題準確率
獲勝的關鍵
英偉達能在AIMO-2拔得頭籌，不是沒有理由的。
除了他們有用不完的卡以外
。
團隊在如何復現成果中暗示瞭如果沒有大型GPU叢集，就別試了
英偉達的OpenMath-Nemotron模型能夠獲勝依賴於三個關鍵步驟。
-
高質量的資料集：英偉達建立了一個包含540K個獨特高質量數學問題的大規模資料集，包括奧林匹克級別的問題及其3.2M個長CoT解決方案；
-
TIR(tool-integrated reasoning)工具整合推理：開發了一種新方法，透過迭代訓練、生成和質量過濾將程式碼執行與長CoT整合，從而得到1.7M個高質量的工具整合推理解決方案；
-
GenSelect模式：建立了一個訓練模型的流程，以從多個候選方案中選擇最有希望的解決方案。這種生成式解決方案選擇（GenSelect）顯著優於多數投票基線。
540K來自AoPS論壇的獨特數學問題
首先，英偉達團隊從網際網路上收集了一大批數學問題。
他們從Art of Problem Solving（AoPS）社群論壇收集了大量數學問題資料集。
除「中學數學」（Middle School Math）版塊外，他們收錄了所有論壇討論內容
資料採集後，他們建立系統化流程提取問題和對應答案，使用Qwen2.5-32B-Instruct模型進行處理，具體流程如下：
-
問題提取：透過大語言模型識別初始帖文中的數學問題。
-
問題分類：採用大語言模型對每個問題進行多維度分類，並剔除所有選擇題、二元判斷題及無效問題。
-
問題轉化：將證明題轉化為需要相似解題技巧的答案導向型問題。
-
答案提取：針對非證明題，從論壇討論中提取最終答案。
-
基準去汙：使用基於LLM的相似度比對，剔除與主流數學基準測試高度相似的問題。
基於LLM的問題提取和精煉流程，最終超過構建了包含54萬個問題的資料集，生成了320萬個長推理CoT解決方案。
DeepSeek-R1和QwQ-32B等模型為每個問題生成多個解決方案候選。而較難的問題會獲得更多的候選方案。
錯誤的解決方案透過Qwen2.5-32B-Instruct驗證答案等效性來過濾。如果沒有找到答案，則使用最頻繁的候選答案。
在提交的本次解決方案中，他們使用了由DeepSeek-R1生成的220萬個問題的子集。
TIR：工具整合推理（tool-integrated reasoning）
對於求解數學問題，傳統的LLM單純地預測下一個單詞的機率並不是非常適合。
解決數學問題，更好的做法還是要呼叫專業的計算工具。
對於工具整合推理，模型會在需要的地方提示程式碼進行計算，然後在沙箱中執行程式碼。
英偉達用特殊token <tool_call>和<\tool_call>識別程式碼片段。
然後將程式碼附加到LLM輸出中，位於文字“`和“`output之間。
下面是一個輸出示例片段。
GenSelect選擇最優解
下圖是GenSelect的資料構建流程，主要包含三個步驟：
1. 生成摘要
對於OpenMathReasoning資料集中的每個問題，隨機抽取2到16個候選解答摘要，確保每個樣本組中至少包含一個正確解答和一個錯誤解答。
這個過程會重複進行，直到為每個問題獲得8個不同的比較組。
2. 選擇並過濾答案
然後，使用GenSelect提示詞，將任務交給QwQ-32B，讓它從每個組中選擇最有可能的解答。
GenSelect推理提示詞
這個過程生成了100萬個選擇項，隨後刪除選擇了錯誤解答的例項，將資料量過濾到565K。
3. 總結推理過程（reasoning traces）並輸出
透過Qwen2.5-32B-Instruct總結上一布篩選的正確解答的推理過程，從而形成GenSelect的輸出。
模型訓練
本次提交的Kaggle解決方法 ，使用的訓練方法與論文中詳細描述的略有不同。
參賽團隊發現：這種不同的方法訓練的模型，比公開發布的模型使用的token更少。
新模型表現良好，但由於時間限制，他們沒有在最終模型上進一步實驗減少token。
首先，他們使用SFT在2.2M的CoT解決方案子集上，訓練了一個Qwen2.5-14B-Base模型，共8個epoch。
他們將基礎RoPE改為500k以允許長推理。
該模型的其他訓練引數如下：
使用NVIDIA/Nemo-Skills訓練了8 個epoch，
學習率：1e-4，
最佳化器：AdamW，
權重衰減係數：0.01，
並且有10%的線性預熱衰減到學習率為1e-7，
批大小：1024個樣本。
他們還利用了NVIDIA/NeMo-Aligner中的序列打包和上下文並行化技術，顯著加速了長推理資料的訓練。
論文連結：https://arxiv.org/pdf/2405.01481
在512個H100（是的，512 個！）上，訓練持續了48小時。
在使用20%算力的情況下，他們就已經實現了模型的大部分效能，但他們擴大了訓練規模，觀察學習何時達到飽和。
論文中的圖 3(b)顯示了不同訓練階段的指標。最終權重是從不同階段進行權重平均得到的。
接下來是對15K TIR樣本進行輕量級的TIR微調。
參賽團隊用恆定的學習率1e-5 訓練了TIR 模型400步，並使用最後一個checkpoint而沒有進行平均。
隨後合併CoT和TIR兩個checkpoint，因為這樣做既能提高準確性，又能減少解決方案長度和程式碼執行次數，從而加快生成速度。
評估資料集
在比賽中，他們主要使用2024年的美國邀請數學考試（AIME 24）和哈佛-麻省理工數學錦標賽（HMMT）的題目。
後來增加了兩項測試的2025年度題目。
最終基準Comp-Math-24-25包括256道題目，具體組成如下。
模型推理三步走
模型合併
在這次競賽中，他們探索了多種方法來合併具有CoT和TIR行為的兩個LLM。
主要目標：有效地結合這兩個微調階段的獨特優勢，以提高模型的效能。
他們試驗了mergekit包中的幾種合併技術。
mergekit是專用於合併預訓練語言模型的工具包，採用核外計算（out-of-core）技術
結果出乎意料，令人驚訝：最有效的方法竟然是簡單的線性組合！
也就是在TIR微調之前使用的思維鏈checkpoint以及之後獲得的最佳TIR checkpoint，兩者之間的簡單線性組合。
這種策略，能夠控制每個階段對最終模型行為的影響程度。
對於Comp-Math-24-25資料集，下表展示了合併模型的準確率和生成統計資料。
其中length表示解決方案的平均token數，而code表示解決方案的平均程式碼執行次數。
模型加速
優先考慮了權重為Int8 (W8A16) 和FP8的量化，這比BF16提供了更快的推理速度，且精度損失最小。
減少的權重大小還釋放了記憶體，以便用於更大的鍵值快取。
ReDrafter是由Apple開發的一種推測解碼技術，並在TensorRT-LLM 中實現。
論文地址：https://arxiv.org/abs/2403.09919
在OpenMathReasoning-1資料集的隨機子集上訓練了一個ReDrafter頭。
使用這些問題，用目標模型生成了100k個解決方案。
生成的ReDrafter在每個 LLM 步驟中生成3個token，接受率為65%，實現了大約 1.8 倍的速度提升。
表格中的準確率得分是使用合併模型的maj@12指標，在5次執行中取平均值。
TensorRT-LLM推理
預訓練模型使用TensorRT-LLM轉換為TensorRT引擎。
TensorRT-LLM：專為大語言模型推理最佳化的TensorRT 工具包
TensorRT的動態批處理透過動態組合推理請求來提高吞吐量，每個樣本一旦完成就立即釋放——從而減少延遲並最佳化 GPU 利用率。
vLLM團隊提供的一些最新基準測試， 請參見下圖。
由於樣本處理相互獨立，批次計算可無縫混合不同輸入提示（prompt）或隨機種子。
TensorRT-LLM還集成了多項最佳化技術，包括定製注意力核心（custom attention kernels）和分頁KV快取（paged KV caching）等。
非同步批處理
對於每個新問題，他們使用不同的種子，利用TensorRT中的非同步批處理，啟動12次生成。
每個樣本的流處理會監控程式碼塊、停止語句、最大標記數或超時。
如果LLM生成了程式碼，LLM的生成過程會停止，程式碼塊會在沙箱中執行。
沙箱的輸出（或部分錯誤跟蹤）會被附加到LLM中，生成過程繼續進行。
生成過程會持續，直到遇到另一個程式碼塊。
當沒有遇到其他程式碼塊時，根據最大標記數、超時時間或停止語句之一，LLM會停止。
非同步批處理流程
他們最終提交了基於一種「幾乎」貪心的搜尋策略，因為它在小批次大小下提供了更穩定的結果，並且在猜測解碼的速度上略有提升。
為了提高速度，會監控生成過程是否完成：當初始答案相同時，就會提前停止。
提前停止和快取策略
在監控非同步生成過程中，在12次生成中完成10次，他們會提前停止，避免過度等待任何滯後的生成。
他們還實施了一種緩衝策略。
如果一個問題提前完成，未使用的時間將被加入到共享緩衝區。
下一個問題可以從這個緩衝區中提取最多210秒的額外時間，從而使總時間達到560秒。
推理流程
對於最終選擇的提交，他們選擇了一個14B CoT模型和上述的MIX TIR模型。
MIX TIR模型在交叉驗證資料集上得分明顯更好，在公開排行榜上的得分也得到提高（公開排行榜得分：32, 33, 28）。
最終，私密排行榜的結果更接近交叉驗證資料集的結果，而不是公開排行榜的結果。
由於每次提交的時間限制以及只有50個問題被評分，他們沒有足夠的時間和提交機會來準確縮小交叉驗證資料集和公開排行榜之間的差異，尤其是在每次只能提交一個模型的情況下。
AIMO Progress Prize已經舉辦了兩屆。
在第一屆中，前五名的最高分為29分，最低分只有20分。
在過去一年時間後，前五名中，最高分被英偉達刷到了34分，最低分也和第一屆相同。
AIMO是一個難度非常高的挑戰，在這一屆中，AI解決了50道題目中的34道題。
如果換算成100分，AI在這場考試中已經取得了68分，超過了及格線。
也許明年，或者後面，AI就能在這場測試中獲得「全勝」。
當AI能夠解決所有人類數學家提出的問題，也許數學的邊界也會被重新定義。
參考資料：
https://x.com/jandotai/status/1915345568483991741
https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/574765
https://arxiv.org/pdf/2504.16891
https://huggingface.co/collections/nvidia/openmathreasoning-68072c0154a5099573d2e730
技術交流群邀請函
△長按新增小助手
掃描二維碼新增小助手微信
請備註：姓名-學校/公司-研究方向
（如：小張-哈工大-對話系統）
即可申請加入自然語言處理/Pytorch等技術交流群
關於我們
MLNLP 社群是由國內外機器學習與自然語言處理學者聯合構建的民間學術社群，目前已經發展為國內外知名的機器學習與自然語言處理社群，旨在促進機器學習，自然語言處理學術界、產業界和廣大愛好者之間的進步。
社群可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。
掃描二維碼新增小助手微信
關於我們