英偉達正式開源了其不久前在 AI 數(shù)學奧林匹克競賽（AIMO，AI Mathematical Olympiad）中斬獲冠軍的核心模型系列。

在本屆 AIMO-2 Kaggle 競賽中，超過 2,200 支參賽隊伍提交了 AI 模型，挑戰(zhàn)在 5 小時內解決 50 道國家奧林匹克級別的復雜數(shù)學問題。英偉達的 7 人團隊“NemoSkills”最終正確解答了 34 道題目（相比 2024 年的冠軍提高了 5 道），奪得了冠軍。

圖丨此次比賽的排行榜（來源：Kaggle）

現(xiàn)在，英偉達向全球開放了幫助他們獲勝的核心技術，包括小參數(shù)的 OpenMath-Nemotron-1.5B、OpenMath-Nemotron-7B 和直接用于競賽并優(yōu)化的 OpenMath-Nemotron-14B-Kaggle 模型、性能更為強大的旗艦模型 OpenMath-Nemotron-32B，以及訓練它們所依賴的 OpenMathReasoning 數(shù)據(jù)集。

基準測試的結果顯示，這幾款模型表現(xiàn)出色，在 AIME 和 HMMT 競賽中數(shù)學問題上的準確率全面超越了 14B 的 DeepSeek-R1。

（來源：arXiv）

英偉達是如何構建 OpenMath-Nemotron 的？

那么，英偉達是如何構建 OpenMath-Nemotron 模型的？

這首先在于一個大規(guī)模且高質量的訓練數(shù)據(jù)集。認識到現(xiàn)有資源的不足，英偉達團隊首先投入了大量的工作來創(chuàng)建 OpenMathReasoning 數(shù)據(jù)集。

他們先從“Art of Problem Solving（AoPS）”等在線數(shù)學社區(qū)收集了大量的原始數(shù)學問題和討論。

隨后，團隊利用 Qwen2.5-32B-Instruct 開發(fā)了一套自動化流程，對這些原始數(shù)據(jù)進行細致處理。這包括從帖子中提取完整的數(shù)學問題，對問題進行分類（例如，剔除選擇題和是非題），并將一些需要證明過程的問題巧妙地轉化為需要具體答案的形式，以便于模型訓練和自動評估。同時，為了保證模型的泛化能力，他們還進行了基準去污染處理，移除了與現(xiàn)有常見數(shù)學測試集（如 MATH、GSM8K）中題目過于相似的問題。

最終完成的 OpenMathReasoning 數(shù)據(jù)集，包含了 54 萬個高質量數(shù)學問題，其中涵蓋了從中學到奧林匹克競賽等不同難度級別。為了讓模型學會“思考過程”，團隊更進一步地利用 DeepSeek-R1 和 QwQ-32B 等強大的現(xiàn)有模型，為這些問題生成了 320 萬條包含詳細解題步驟的“思維鏈”（CoT，Chain-of-Thought）解決方案。

圖丨數(shù)據(jù)集組成（來源：arXiv）

第二個核心部分是工具集成推理。

現(xiàn)代 AI 研究的一個重要趨勢是讓語言模型學會使用外部工具，例如調用計算器或執(zhí)行代碼片段，來輔助解決問題，尤其是在需要精確計算或模擬的場景下。然而，團隊在實踐中發(fā)現(xiàn)，即便是當時最強的開源數(shù)學模型，也難以通過簡單的提示工程來引導它們生成高質量的、將代碼執(zhí)行與自然語言推理深度融合的解決方案（即 TIR）。這些模型似乎對其自身固有的純文本推理模式產生了某種“路徑依賴”。

為了克服這一障礙，NemoSkills 團隊設計并實施了一套迭代式開發(fā)流程。他們首先選擇了一個指令遵循能力較好的基礎模型（LIMO-Qwen-32B），用少量推理數(shù)據(jù)對其進行初步微調。然后，引導這個模型生成第一批包含 Python 代碼的 TIR 解決方案。關鍵的下一步是進行嚴格的質量過濾：利用另一個強大的大模型（ Qwen2.5-32B-Instruct），來判斷每個代碼塊的“新穎性”（是產生了新結果還是僅僅驗證已知步驟）和“重要性”（是解決問題的關鍵環(huán)節(jié)還是可以被幾步簡單 CoT 取代）。只有那些代碼執(zhí)行提供了顯著推理價值（而非冗余計算）的樣本才被保留下來，形成了約 1.5 萬個樣本的初始 TIR 訓練集。

接下來，他們用這個高質量的初始集去微調更強大的模型（如 QwQ-32B），使其初步具備生成 TIR 的能力。隨后，利用這個微調后的模型生成更多、更高質量的 TIR 數(shù)據(jù)，并再次運用上述過濾標準進行篩選。這個“生成-過濾-訓練”的閉環(huán)被重復執(zhí)行，每一輪都提升了 TIR 數(shù)據(jù)的規(guī)模和質量。最終，團隊構建起了一個包含 170 萬條高質量 TIR 解決方案的數(shù)據(jù)集。基于此訓練出的 OpenMath-Nemotron 模型，能夠熟練地在自然語言推理中嵌入 Python 代碼執(zhí)行，從而攻克那些純文本推理難以解決的復雜計算問題。此外，他們還設計了一種機制，使得模型在生成答案時能夠遵循對代碼塊使用次數(shù)的限制，這對于資源受限的推理場景至關重要。

第三個核心部分則是團隊提出的生成式解決方案選擇。

在解決困難問題時，讓模型生成多個候選答案并從中擇優(yōu)，是提升最終準確率的常用技巧。傳統(tǒng)的“多數(shù)投票”方法雖然直觀，但往往無法充分發(fā)掘模型生成的所有答案中的潛在正確信息，其性能通常遠低于理論上的“pass@k”（即 k 個答案中至少有一個正確的概率）上限。

為了彌補這一差距，英偉達團隊開發(fā)了 GenSelect 技術。其核心思想不再是簡單地對最終答案進行投票，而是訓練一個模型，讓它扮演“評審員”的角色，能夠“閱讀”并“理解”多個候選解決方案的完整摘要，然后基于對解題邏輯、步驟合理性等的判斷，選出最可信、最有可能正確的那一個。

圖丨 GenSelect 的數(shù)據(jù)構建流程（來源：arXiv）

具體來說，團隊首先利用 Qwen2.5-32B-Instruct 模型為 OpenMathReasoning 數(shù)據(jù)集中所有已生成的 CoT 和 TIR 解決方案重新生成了結構化的、信息更豐富的摘要。然后，他們構建了 GenSelect 的訓練數(shù)據(jù)：為每個原始問題，隨機抽取 2 到 16 個候選方案的摘要（特別設計以確保樣本組中至少包含一個正確和一個錯誤的解），將這些摘要連同原問題一起輸入給 QwQ-32B 模型，并要求它生成一段詳細的比較分析文本，最終明確指出哪個索引號的解決方案是最佳的。通過篩選掉那些模型判斷錯誤（即選擇了錯誤答案）的案例，他們構建了一個包含 56.6 萬個樣本的 GenSelect 訓練數(shù)據(jù)集。

實驗結果表明，經過 GenSelect 加持的模型，其最終準確率相比簡單的多數(shù)投票有了顯著提升，尤其是在候選方案數(shù)量不多時效果更為明顯。雖然由于 AIMO 競賽嚴格的時間和計算限制，GenSelect 未能被納入最終的獲勝提交方案中，但這項技術已被完全整合到此次發(fā)布的 OpenMath-Nemotron-32B 模型中，構成了其支持的三大推理模式之一。

基于上述三大支柱和海量數(shù)據(jù)，英偉達團隊訓練了一系列名為OpenMath-Nemotron的模型，參數(shù)規(guī)模涵蓋 1.5B、7B、14B 和 32B。這些模型均基于強大的 Qwen2.5 基座模型進行微調。對于 1.5B 和 7B 版本，他們甚至使用了專門為數(shù)學任務優(yōu)化的 Qwen2.5-Math 版本作為起點。

訓練過程采用了監(jiān)督微調，混合使用了 CoT、TIR 和 GenSelect 三種任務的數(shù)據(jù)，總計達 550 萬個樣本。這意味著同一個模型可以通過不同的提示（prompt）在 CoT（純文本推理）、TIR（工具集成推理）和 GenSelect（多方案選擇）模式下工作。

圖丨訓練過程中準確率的提升（來源：arXiv）

為了處理長達數(shù)千甚至上萬個 token 的長序列推理，團隊應用了旋轉位置編碼（RoPE，Rotary Position Embedding）擴展技術，并將訓練過程中的上下文窗口擴展到支持長序列。訓練使用了英偉達自家的 NeMo-Aligner 工具包，并結合了序列打包、上下文并行等技術來加速長序列訓練。此外，他們還采用了檢查點平均（checkpoint averaging）和在更難問題子集上進行第二輪微調等策略，進一步提升模型性能。

多項優(yōu)化推理措施

贏得 AIMO-2 競賽不僅需要模型本身強大，還需要在極其苛刻的 5 小時、4x L4 GPU 限制下高效完成推理。這要求團隊在模型選擇和推理優(yōu)化上做出極致權衡。

他們的最終提交方案基于 OpenMath-Nemotron-14B 模型的一個早期版本，該版本在一個稍小的 CoT 數(shù)據(jù)集（僅 DeepSeek-R1 生成）上訓練，并進行了輕量級的 TIR 微調。值得注意的是，他們采用了模型合并技術，將純 CoT 訓練的檢查點和經過 TIR 微調的檢查點進行線性組合。這種簡單而有效的方法，讓他們能夠在保持 TIR 能力的同時，部分恢復 CoT 模型的生成流暢性和速度優(yōu)勢，并減少代碼調用次數(shù)，從而更好地適應競賽環(huán)境。

為了在有限的時間內最大化解題數(shù)量和準確率，團隊實施了多項推理優(yōu)化措施：

首先，他們使用 TensorRT-LLM 將預訓練模型轉換為 TensorRT 引擎。這一工具的動態(tài)批處理功能通過動態(tài)分組推理請求提高了吞吐量，在樣本完成后即刻釋放，減少延遲并優(yōu)化 GPU 利用率。由于樣本是獨立處理的，批處理可以無縫混合不同的提示或推理參數(shù)。TensorRT-LLM 還包括自定義注意力內核和分頁 KV 緩存等多種優(yōu)化。

在量化方面，團隊優(yōu)先采用 int8 權重量化（W8A16）和 FP8 量化，相比 BF16 格式速度提升了 1.5 倍，同時對準確率的影響最小。減小的權重大小還為更大的鍵值緩存釋放了內存，允許處理更長的序列。團隊還使用了蘋果開發(fā)的 ReDrafter 技術，這是一種循環(huán)推測解碼方法，使用基于 RNN 的起草器在每個解碼步驟提出并驗證多個 token。他們訓練了一個能夠在每一步提出最多三個 token 的起草器，在大約 65% 的步驟中成功接受所有三個 token，顯著加速了生成過程。

（來源：arXiv）

此外，團隊通過將 CoT 和 TIR 檢查點線性組合創(chuàng)建了最終模型，這種策略允許他們控制每個微調階段對最終模型行為的影響程度。最佳模型是使用 CoT0.3+TIR0.7 的組合創(chuàng)建的，這不僅提高了準確率，還通過減少解決方案長度和代碼執(zhí)行次數(shù)加速了生成。團隊實現(xiàn)了一種緩沖策略，為每個問題分配 350 秒的基本時間限制，如果一個問題提前完成，未使用的時間會被添加到共享緩沖區(qū)，供后續(xù)問題使用。

團隊還利用了 NeMo-Skills 的異步生成功能實現(xiàn)批量處理和早停。例如，在 16 個樣本的批處理中，如果前 4-5 個完成的樣本就已經對最終答案達成一致，則取消剩余的生成并繼續(xù)下一個問題。這種機制極大地節(jié)約了在簡單或中等難度問題上可能浪費的時間，為攻克難題爭取了寶貴的時間窗口。早停策略增加了響應相關性，因為較短的答案往往質量更高。

圖丨異步批處理流程（來源：Kaggle）

實驗結果顯示，在 Comp-Math-24-25 測試集（包含來自 AIME 和 HMMT 競賽的問題）上，團隊的模型表現(xiàn)出色。1.5B 模型在 CoT 模式下單次通過準確率為 58.2%，多數(shù)投票準確率達 80.0%；在 TIR 模式下，這些數(shù)字分別提高到 64.5% 和 83.3%；使用 GenSelect 技術后，準確率進一步提升至 83.3%。14B 模型的表現(xiàn)更為出色，在 TIR 模式結合 GenSelect 使用時，準確率高達 90.0%。最大的 32B 模型在相同條件下甚至達到了 93.3% 的準確率。這些結果也表明，無論模型大小如何，TIR 模式始終優(yōu)于純 CoT 模式，而 GenSelect 技術能進一步提高準確率。

圖｜數(shù)學基準測試的評估結果（來源：arXiv）

目前，英偉達團隊已將完整的 OpenMathReasoning 數(shù)據(jù)集、訓練好的 OpenMath-Nemotron 模型系列以及所有相關代碼以商業(yè)許可方式發(fā)布到 Hugging Face 和 GitHub 上（項目地址：https://huggingface.co/collections/nvidia/openmathreasoning-68072c0154a5099573d2e730）。

參考資料：

1.https://arxiv.org/abs/2504.16891

2.https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/574765

運營/排版：何晨龍