新智元報道

編輯：KingHZ 定慧

【新智元導讀】AIMO2冠軍「答卷」公布了！英偉達團隊NemoSkills拔得頭籌，開源了OpenMath-Nemotron系列AI模型，1.5B小模型擊敗14B-DeepSeek「推理大模型」！

AI數學奧賽第一名「答卷」終于公布！

論文地址：https://arxiv.org/abs/2504.16891

亞軍隊成員、清華學子Yichen You表示冠軍實至名歸，自己獲益良多。

參加本次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模型進行處理，具體流程如下：

1. 問題提取：通過大語言模型識別初始帖文中的數學問題。

2. 問題分類：采用大語言模型對每個問題進行多維度分類，并剔除所有選擇題、二元判斷題及無效問題。

3. 問題轉化：將證明題轉化為需要相似解題技巧的答案導向型問題。

4. 答案提取：針對非證明題，從論壇討論中提取最終答案。

5. 基準去污：使用基于LLM的相似度比對，剔除與主流數學基準測試高度相似的問題。

基于LLM的問題提取和精煉流程，最終超過構建了包含54萬個問題的數據集，生成了320萬個長推理CoT解決方案。

DeepSeek-R1和QwQ-32B等模型為每個問題生成多個解決方案候選。而較難的問題會獲得更多的候選方案。

錯誤的解決方案通過Qwen2.5-32B-Instruct驗證答案等效性來過濾。如果沒有找到答案，則使用最頻繁的候選答案。

在提交的本次解決方案中，他們使用了由DeepSeek-R1生成的220萬個子集。

TIR：工具集成推理（tool-integrated reasoning）

對于求解數學問題，傳統的LLM單純地預測下一個單詞的概率并不是非常適合。

解決數學問題，更好的做法還是要調用專業的計算工具。

對于工具集成推理，模型會在需要的地方提示代碼進行計算，然后在沙箱中執行代碼。

英偉達用特殊token 和<\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