新智元報(bào)道
編輯：KingHZ 定慧
【新智元導(dǎo)讀】AIMO2冠軍「答卷」公布了！英偉達(dá)團(tuán)隊(duì)NemoSkills拔得頭籌，開源了OpenMath-Nemotron系列AI模型，1.5B小模型擊敗14B-DeepSeek「推理大模型」！
AI數(shù)學(xué)奧賽第一名「答卷」終于公布！
論文地址：https://arxiv.org/abs/2504.16891
亞軍隊(duì)成員、清華學(xué)子Yichen You表示冠軍實(shí)至名歸，自己獲益良多。
參加本次Kaggle比賽、軟件工程師Chan Kha Vu，則盛贊道：這些模型太不可思議了！從基礎(chǔ)的Qwen模型訓(xùn)練開始，甚至都不是推理模型。而且沒有利用強(qiáng)化學(xué)習(xí)！
英偉達(dá)團(tuán)隊(duì)參賽的模型叫做OpenMath-Nemotron系列，使用OpenMathReasoning Dataset進(jìn)行訓(xùn)練，共發(fā)布了四種參數(shù)：
OpenMath-Nemotron-1.5B
OpenMath-Nemotron-7B
OpenMath-Nemotron-14B-Kaggle（AIMO-2 Kaggle競(jìng)賽中使用的模型）
OpenMath-Nemotron-32B
這些模型在流行的數(shù)學(xué)基準(zhǔn)測(cè)試中都取得了最好的成績(jī)。
甚至1.5B的OpenMath-Nemotron模型，超越14B的DeepSeek-R1蒸餾模型！
圖1：AIME和HMMT競(jìng)賽中的數(shù)學(xué)問題準(zhǔn)確率
獲勝的關(guān)鍵
英偉達(dá)能在AIMO-2拔得頭籌，不是沒有理由的。
除了他們有用不完的卡以外。
團(tuán)隊(duì)在如何復(fù)現(xiàn)成果中暗示了如果沒有大型GPU集群，就別試了
英偉達(dá)的OpenMath-Nemotron模型能夠獲勝依賴于三個(gè)關(guān)鍵步驟。
高質(zhì)量的數(shù)據(jù)集：英偉達(dá)創(chuàng)建了一個(gè)包含540K個(gè)獨(dú)特高質(zhì)量數(shù)學(xué)問題的大規(guī)模數(shù)據(jù)集，包括奧林匹克級(jí)別的問題及其3.2M個(gè)長(zhǎng)CoT解決方案；
TIR(tool-integrated reasoning)工具集成推理：開發(fā)了一種新方法，通過迭代訓(xùn)練、生成和質(zhì)量過濾將代碼執(zhí)行與長(zhǎng)CoT集成，從而得到1.7M個(gè)高質(zhì)量的工具集成推理解決方案；
GenSelect模式：創(chuàng)建了一個(gè)訓(xùn)練模型的流程，以從多個(gè)候選方案中選擇最有希望的解決方案。這種生成式解決方案選擇（GenSelect）顯著優(yōu)于多數(shù)投票基線。
540K來自AoPS論壇的獨(dú)特?cái)?shù)學(xué)問題
首先，英偉達(dá)團(tuán)隊(duì)從互聯(lián)網(wǎng)上收集了一大批數(shù)學(xué)問題。
他們從Art of Problem Solving（AoPS）社區(qū)論壇收集了大量數(shù)學(xué)問題數(shù)據(jù)集。
除「中學(xué)數(shù)學(xué)」（Middle School Math）版塊外，他們收錄了所有論壇討論內(nèi)容
數(shù)據(jù)采集后，他們建立系統(tǒng)化流程提取問題和對(duì)應(yīng)答案，使用Qwen2.5-32B-Instruct模型進(jìn)行處理，具體流程如下：
問題提?。和ㄟ^大語言模型識(shí)別初始帖文中的數(shù)學(xué)問題。
問題分類：采用大語言模型對(duì)每個(gè)問題進(jìn)行多維度分類，并剔除所有選擇題、二元判斷題及無效問題。
問題轉(zhuǎn)化：將證明題轉(zhuǎn)化為需要相似解題技巧的答案導(dǎo)向型問題。
答案提?。横槍?duì)非證明題，從論壇討論中提取最終答案。
基準(zhǔn)去污：使用基于LLM的相似度比對(duì)，剔除與主流數(shù)學(xué)基準(zhǔn)測(cè)試高度相似的問題。
基于LLM的問題提取和精煉流程，最終超過構(gòu)建了包含54萬個(gè)問題的數(shù)據(jù)集，生成了320萬個(gè)長(zhǎng)推理CoT解決方案。
DeepSeek-R1和QwQ-32B等模型為每個(gè)問題生成多個(gè)解決方案候選。而較難的問題會(huì)獲得更多的候選方案。
錯(cuò)誤的解決方案通過Qwen2.5-32B-Instruct驗(yàn)證答案等效性來過濾。如果沒有找到答案，則使用最頻繁的候選答案。
在提交的本次解決方案中，他們使用了由DeepSeek-R1生成的220萬個(gè)子集。
TIR：工具集成推理（tool-integrated reasoning）
對(duì)于求解數(shù)學(xué)問題，傳統(tǒng)的LLM單純地預(yù)測(cè)下一個(gè)單詞的概率并不是非常適合。
解決數(shù)學(xué)問題，更好的做法還是要調(diào)用專業(yè)的計(jì)算工具。
對(duì)于工具集成推理，模型會(huì)在需要的地方提示代碼進(jìn)行計(jì)算，然后在沙箱中執(zhí)行代碼。
英偉達(dá)用特殊token 和<\tool_call>識(shí)別代碼片段。
然后將代碼附加到LLM輸出中，位于文本```和```output之間。
下面是一個(gè)輸出示例片段。
GenSelect選擇最優(yōu)解
下圖是GenSelect的數(shù)據(jù)構(gòu)建流程，主要包含三個(gè)步驟：
1. 生成摘要
對(duì)于OpenMathReasoning數(shù)據(jù)集中的每個(gè)問題，隨機(jī)抽取2到16個(gè)候選解答摘要，確保每個(gè)樣本組中至少包含一個(gè)正確解答和一個(gè)錯(cuò)誤解答。
這個(gè)過程會(huì)重復(fù)進(jìn)行，直到為每個(gè)問題獲得8個(gè)不同的比較組。
2. 選擇并過濾答案
然后，使用GenSelect提示詞，將任務(wù)交給QwQ-32B，讓它從每個(gè)組中選擇最有可能的解答。
GenSelect推理提示詞
這個(gè)過程生成了100萬個(gè)選擇項(xiàng)，隨后刪除選擇了錯(cuò)誤解答的實(shí)例，將數(shù)據(jù)量過濾到565K。
3. 總結(jié)推理過程（reasoning traces）并輸出
通過Qwen2.5-32B-Instruct總結(jié)上一部篩選的正確解答的推理過程，從而形成GenSelect的輸出。
模型訓(xùn)練
本次提交的Kaggle解決方法 ，使用的訓(xùn)練方法與論文中詳細(xì)描述的略有不同。
參賽團(tuán)隊(duì)發(fā)現(xiàn)：這種不同的方法訓(xùn)練的模型，比公開發(fā)布的模型使用的token更少。
新模型表現(xiàn)良好，但由于時(shí)間限制，他們沒有在最終模型上進(jìn)一步實(shí)驗(yàn)減少token。
首先，他們使用SFT在2.2M的CoT解決方案子集上，訓(xùn)練了一個(gè)Qwen2.5-14B-Base模型，共8個(gè)epoch。
他們將基礎(chǔ)RoPE改為500k以允許長(zhǎng)推理。
該模型的其他訓(xùn)練參數(shù)如下：
使用NVIDIA/Nemo-Skills訓(xùn)練了8 個(gè)epoch，
學(xué)習(xí)率：1e-4，
優(yōu)化器：AdamW，
權(quán)重衰減系數(shù)：0.01，
并且有10%的線性預(yù)熱衰減到學(xué)習(xí)率為1e-7，
批大?。?024個(gè)樣本。
他們還利用了NVIDIA/NeMo-Aligner中的序列打包和上下文并行化技術(shù)，顯著加速了長(zhǎng)推理數(shù)據(jù)的訓(xùn)練。
論文鏈接：https://arxiv.org/pdf/2405.01481
在512個(gè)H100（是的，512 個(gè)！）上，訓(xùn)練持續(xù)了48小時(shí)。
在使用20%算力的情況下，他們就已經(jīng)實(shí)現(xiàn)了模型的大部分性能，但他們擴(kuò)大了訓(xùn)練規(guī)模，觀察學(xué)習(xí)何時(shí)達(dá)到飽和。
論文中的圖 3(b)顯示了不同訓(xùn)練階段的指標(biāo)。最終權(quán)重是從不同階段進(jìn)行權(quán)重平均得到的。
接下來是對(duì)15K TIR樣本進(jìn)行輕量級(jí)的TIR微調(diào)。
參賽團(tuán)隊(duì)用恒定的學(xué)習(xí)率1e-5 訓(xùn)練了TIR 模型400步，并使用最后一個(gè)checkpoint而沒有進(jìn)行平均。
隨后合并CoT和TIR兩個(gè)checkpoint，因?yàn)檫@樣做既能提高準(zhǔn)確性，又能減少解決方案長(zhǎng)度和代碼執(zhí)行次數(shù)，從而加快生成速度。
評(píng)估數(shù)據(jù)集
在比賽中，他們主要使用2024年的美國(guó)邀請(qǐng)數(shù)學(xué)考試（AIME 24）和哈佛-麻省理工數(shù)學(xué)錦標(biāo)賽（HMMT）的題目。
后來增加了兩項(xiàng)測(cè)試的2025年度題目。
最終基準(zhǔn)Comp-Math-24-25包括256道題目，具體組成如下。
模型推理三步走
模型合并
在這次競(jìng)賽中，他們探索了多種方法來合并具有CoT和TIR行為的兩個(gè)LLM。
主要目標(biāo)：有效地結(jié)合這兩個(gè)微調(diào)階段的獨(dú)特優(yōu)勢(shì)，以提高模型的性能。
他們?cè)囼?yàn)了mergekit包中的幾種合并技術(shù)。
mergekit是專用于合并預(yù)訓(xùn)練語言模型的工具包，采用核外計(jì)算（out-of-core）技術(shù)
結(jié)果出乎意料，令人驚訝：最有效的方法竟然是簡(jiǎn)單的線性組合！
也就是在TIR微調(diào)之前使用的思維鏈checkpoint以及之后獲得的最佳TIR checkpoint，兩者之間的簡(jiǎn)單線性組合。
這種策略，能夠控制每個(gè)階段對(duì)最終模型行為的影響程度。
對(duì)于Comp-Math-24-25數(shù)據(jù)集，下表展示了合并模型的準(zhǔn)確率和生成統(tǒng)計(jì)數(shù)據(jù)。
其中l(wèi)ength表示解決方案的平均token數(shù)，而code表示解決方案的平均代碼執(zhí)行次數(shù)。
模型加速
優(yōu)先考慮了權(quán)重為Int8 (W8A16) 和FP8的量化，這比BF16提供了更快的推理速度，且精度損失最小。
減少的權(quán)重大小還釋放了內(nèi)存，以便用于更大的鍵值緩存。
ReDrafter是由Apple開發(fā)的一種推測(cè)解碼技術(shù)，并在TensorRT-LLM 中實(shí)現(xiàn)。
論文地址：https://arxiv.org/abs/2403.09919
在OpenMathReasoning-1數(shù)據(jù)集的隨機(jī)子集上訓(xùn)練了一個(gè)ReDrafter頭。
使用這些問題，用目標(biāo)模型生成了100k個(gè)解決方案。
生成的ReDrafter在每個(gè) LLM 步驟中生成3個(gè)token，接受率為65%，實(shí)現(xiàn)了大約 1.8 倍的速度提升。
表格中的準(zhǔn)確率得分是使用合并模型的maj@12指標(biāo)，在5次運(yùn)行中取平均值。
TensorRT-LLM推理
預(yù)訓(xùn)練模型使用TensorRT-LLM轉(zhuǎn)換為TensorRT引擎。
TensorRT-LLM：專為大語言模型推理優(yōu)化的TensorRT 工具包
TensorRT的動(dòng)態(tài)批處理通過動(dòng)態(tài)組合推理請(qǐng)求來提高吞吐量，每個(gè)樣本一旦完成就立即釋放——從而減少延遲并優(yōu)化 GPU 利用率。
vLLM團(tuán)隊(duì)提供的一些最新基準(zhǔn)測(cè)試， 請(qǐng)參見下圖。
由于樣本處理相互獨(dú)立，批次計(jì)算可無縫混合不同輸入提示（prompt）或隨機(jī)種子。
TensorRT-LLM還集成了多項(xiàng)優(yōu)化技術(shù)，包括定制注意力內(nèi)核（custom attention kernels）和分頁KV緩存（paged KV caching）等。
異步批處理
對(duì)于每個(gè)新問題，他們使用不同的種子，利用TensorRT中的異步批處理，啟動(dòng)12次生成。
每個(gè)樣本的流處理會(huì)監(jiān)控代碼塊、停止語句、最大標(biāo)記數(shù)或超時(shí)。
如果LLM生成了代碼，LLM的生成過程會(huì)停止，代碼塊會(huì)在沙箱中執(zhí)行。
沙箱的輸出（或部分錯(cuò)誤跟蹤）會(huì)被附加到LLM中，生成過程繼續(xù)進(jìn)行。
生成過程會(huì)持續(xù)，直到遇到另一個(gè)代碼塊。
當(dāng)沒有遇到其他代碼塊時(shí)，根據(jù)最大標(biāo)記數(shù)、超時(shí)時(shí)間或停止語句之一，LLM會(huì)停止。
異步批處理流程
他們最終提交了基于一種「幾乎」貪心的搜索策略，因?yàn)樗谛∨看笮∠绿峁┝烁€(wěn)定的結(jié)果，并且在猜測(cè)解碼的速度上略有提升。
為了提高速度，會(huì)監(jiān)控生成過程是否完成：當(dāng)初始答案相同時(shí)，就會(huì)提前停止。
提前停止和緩存策略
在監(jiān)控異步生成過程中，在12次生成中完成10次，他們會(huì)提前停止，避免過度等待任何滯后的生成。
他們還實(shí)施了一種緩沖策略。
如果一個(gè)問題提前完成，未使用的時(shí)間將被加入到共享緩沖區(qū)。
下一個(gè)問題可以從這個(gè)緩沖區(qū)中提取最多210秒的額外時(shí)間，從而使總時(shí)間達(dá)到560秒。
推理流程
對(duì)于最終選擇的提交，他們選擇了一個(gè)14B CoT模型和上述的MIX TIR模型。
MIX TIR模型在交叉驗(yàn)證數(shù)據(jù)集上得分明顯更好，在公開排行榜上的得分也得到提高（公開排行榜得分：32, 33, 28）。
最終，私密排行榜的結(jié)果更接近交叉驗(yàn)證數(shù)據(jù)集的結(jié)果，而不是公開排行榜的結(jié)果。
由于每次提交的時(shí)間限制以及只有50個(gè)問題被評(píng)分，他們沒有足夠的時(shí)間和提交機(jī)會(huì)來準(zhǔn)確縮小交叉驗(yàn)證數(shù)據(jù)集和公開排行榜之間的差異，尤其是在每次只能提交一個(gè)模型的情況下。
AIMO Progress Prize已經(jīng)舉辦了兩屆。
在第一屆中，前五名的最高分為29分，最低分只有20分。
在過去一年時(shí)間后，前五名中，最高分被英偉達(dá)刷到了34分，最低分也和第一屆相同。
AIMO是一個(gè)難度非常高的挑戰(zhàn)，在這一屆中，AI解決了50道題目中的34道題。
如果換算成100分，AI在這場(chǎng)考試中已經(jīng)取得了68分，超過了及格線。
也許明年，或者后面，AI就能在這場(chǎng)測(cè)試中獲得「全勝」。
當(dāng)AI能夠解決所有人類數(shù)學(xué)家提出的問題，也許數(shù)學(xué)的邊界也會(huì)被重新定義。
參考資料：
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
熱門跟貼