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