英偉達(dá)正式開源了其不久前在 AI 數(shù)學(xué)奧林匹克競(jìng)賽（AIMO，AI Mathematical Olympiad）中斬獲冠軍的核心模型系列。
在本屆 AIMO-2 Kaggle 競(jìng)賽中，超過 2,200 支參賽隊(duì)伍提交了 AI 模型，挑戰(zhàn)在 5 小時(shí)內(nèi)解決 50 道國(guó)家奧林匹克級(jí)別的復(fù)雜數(shù)學(xué)問題。英偉達(dá)的 7 人團(tuán)隊(duì)“NemoSkills”最終正確解答了 34 道題目（相比 2024 年的冠軍提高了 5 道），奪得了冠軍。
現(xiàn)在，英偉達(dá)向全球開放了幫助他們獲勝的核心技術(shù)，包括小參數(shù)的 OpenMath-Nemotron-1.5B、OpenMath-Nemotron-7B 和直接用于競(jìng)賽并優(yōu)化的 OpenMath-Nemotron-14B-Kaggle 模型、性能更為強(qiáng)大的旗艦?zāi)Ｐ?OpenMath-Nemotron-32B，以及訓(xùn)練它們所依賴的 OpenMathReasoning 數(shù)據(jù)集。
基準(zhǔn)測(cè)試的結(jié)果顯示，這幾款模型表現(xiàn)出色，在 AIME 和 HMMT 競(jìng)賽中數(shù)學(xué)問題上的準(zhǔn)確率全面超越了 14B 的 DeepSeek-R1。
英偉達(dá)是如何構(gòu)建 OpenMath-Nemotron 的？
那么，英偉達(dá)是如何構(gòu)建 OpenMath-Nemotron 模型的？
這首先在于一個(gè)大規(guī)模且高質(zhì)量的訓(xùn)練數(shù)據(jù)集。認(rèn)識(shí)到現(xiàn)有資源的不足，英偉達(dá)團(tuán)隊(duì)首先投入了大量的工作來(lái)創(chuàng)建 OpenMathReasoning 數(shù)據(jù)集。
他們先從“Art of Problem Solving（AoPS）”等在線數(shù)學(xué)社區(qū)收集了大量的原始數(shù)學(xué)問題和討論。
隨后，團(tuán)隊(duì)利用 Qwen2.5-32B-Instruct 開發(fā)了一套自動(dòng)化流程，對(duì)這些原始數(shù)據(jù)進(jìn)行細(xì)致處理。這包括從帖子中提取完整的數(shù)學(xué)問題，對(duì)問題進(jìn)行分類（例如，剔除選擇題和是非題），并將一些需要證明過程的問題巧妙地轉(zhuǎn)化為需要具體答案的形式，以便于模型訓(xùn)練和自動(dòng)評(píng)估。同時(shí)，為了保證模型的泛化能力，他們還進(jìn)行了基準(zhǔn)去污染處理，移除了與現(xiàn)有常見數(shù)學(xué)測(cè)試集（如 MATH、GSM8K）中題目過于相似的問題。
最終完成的 OpenMathReasoning 數(shù)據(jù)集，包含了 54 萬(wàn)個(gè)高質(zhì)量數(shù)學(xué)問題，其中涵蓋了從中學(xué)到奧林匹克競(jìng)賽等不同難度級(jí)別。為了讓模型學(xué)會(huì)“思考過程”，團(tuán)隊(duì)更進(jìn)一步地利用 DeepSeek-R1 和 QwQ-32B 等強(qiáng)大的現(xiàn)有模型，為這些問題生成了 320 萬(wàn)條包含詳細(xì)解題步驟的“思維鏈”（CoT，Chain-of-Thought）解決方案。
第二個(gè)核心部分是工具集成推理。
現(xiàn)代 AI 研究的一個(gè)重要趨勢(shì)是讓語(yǔ)言模型學(xué)會(huì)使用外部工具，例如調(diào)用計(jì)算器或執(zhí)行代碼片段，來(lái)輔助解決問題，尤其是在需要精確計(jì)算或模擬的場(chǎng)景下。然而，團(tuán)隊(duì)在實(shí)踐中發(fā)現(xiàn)，即便是當(dāng)時(shí)最強(qiáng)的開源數(shù)學(xué)模型，也難以通過簡(jiǎn)單的提示工程來(lái)引導(dǎo)它們生成高質(zhì)量的、將代碼執(zhí)行與自然語(yǔ)言推理深度融合的解決方案（即 TIR）。這些模型似乎對(duì)其自身固有的純文本推理模式產(chǎn)生了某種“路徑依賴”。
為了克服這一障礙，NemoSkills 團(tuán)隊(duì)設(shè)計(jì)并實(shí)施了一套迭代式開發(fā)流程。他們首先選擇了一個(gè)指令遵循能力較好的基礎(chǔ)模型（LIMO-Qwen-32B），用少量推理數(shù)據(jù)對(duì)其進(jìn)行初步微調(diào)。然后，引導(dǎo)這個(gè)模型生成第一批包含 Python 代碼的 TIR 解決方案。關(guān)鍵的下一步是進(jìn)行嚴(yán)格的質(zhì)量過濾：利用另一個(gè)強(qiáng)大的大模型（ Qwen2.5-32B-Instruct），來(lái)判斷每個(gè)代碼塊的“新穎性”（是產(chǎn)生了新結(jié)果還是僅僅驗(yàn)證已知步驟）和“重要性”（是解決問題的關(guān)鍵環(huán)節(jié)還是可以被幾步簡(jiǎn)單 CoT 取代）。只有那些代碼執(zhí)行提供了顯著推理價(jià)值（而非冗余計(jì)算）的樣本才被保留下來(lái)，形成了約 1.5 萬(wàn)個(gè)樣本的初始 TIR 訓(xùn)練集。
接下來(lái)，他們用這個(gè)高質(zhì)量的初始集去微調(diào)更強(qiáng)大的模型（如 QwQ-32B），使其初步具備生成 TIR 的能力。隨后，利用這個(gè)微調(diào)后的模型生成更多、更高質(zhì)量的 TIR 數(shù)據(jù)，并再次運(yùn)用上述過濾標(biāo)準(zhǔn)進(jìn)行篩選。這個(gè)“生成-過濾-訓(xùn)練”的閉環(huán)被重復(fù)執(zhí)行，每一輪都提升了 TIR 數(shù)據(jù)的規(guī)模和質(zhì)量。最終，團(tuán)隊(duì)構(gòu)建起了一個(gè)包含 170 萬(wàn)條高質(zhì)量 TIR 解決方案的數(shù)據(jù)集?；诖擞?xùn)練出的 OpenMath-Nemotron 模型，能夠熟練地在自然語(yǔ)言推理中嵌入 Python 代碼執(zhí)行，從而攻克那些純文本推理難以解決的復(fù)雜計(jì)算問題。此外，他們還設(shè)計(jì)了一種機(jī)制，使得模型在生成答案時(shí)能夠遵循對(duì)代碼塊使用次數(shù)的限制，這對(duì)于資源受限的推理場(chǎng)景至關(guān)重要。
第三個(gè)核心部分則是團(tuán)隊(duì)提出的生成式解決方案選擇。
在解決困難問題時(shí)，讓模型生成多個(gè)候選答案并從中擇優(yōu)，是提升最終準(zhǔn)確率的常用技巧。傳統(tǒng)的“多數(shù)投票”方法雖然直觀，但往往無(wú)法充分發(fā)掘模型生成的所有答案中的潛在正確信息，其性能通常遠(yuǎn)低于理論上的“pass@k”（即 k 個(gè)答案中至少有一個(gè)正確的概率）上限。
為了彌補(bǔ)這一差距，英偉達(dá)團(tuán)隊(duì)開發(fā)了 GenSelect 技術(shù)。其核心思想不再是簡(jiǎn)單地對(duì)最終答案進(jìn)行投票，而是訓(xùn)練一個(gè)模型，讓它扮演“評(píng)審員”的角色，能夠“閱讀”并“理解”多個(gè)候選解決方案的完整摘要，然后基于對(duì)解題邏輯、步驟合理性等的判斷，選出最可信、最有可能正確的那一個(gè)。
具體來(lái)說，團(tuán)隊(duì)首先利用 Qwen2.5-32B-Instruct 模型為 OpenMathReasoning 數(shù)據(jù)集中所有已生成的 CoT 和 TIR 解決方案重新生成了結(jié)構(gòu)化的、信息更豐富的摘要。然后，他們構(gòu)建了 GenSelect 的訓(xùn)練數(shù)據(jù)：為每個(gè)原始問題，隨機(jī)抽取 2 到 16 個(gè)候選方案的摘要（特別設(shè)計(jì)以確保樣本組中至少包含一個(gè)正確和一個(gè)錯(cuò)誤的解），將這些摘要連同原問題一起輸入給 QwQ-32B 模型，并要求它生成一段詳細(xì)的比較分析文本，最終明確指出哪個(gè)索引號(hào)的解決方案是最佳的。通過篩選掉那些模型判斷錯(cuò)誤（即選擇了錯(cuò)誤答案）的案例，他們構(gòu)建了一個(gè)包含 56.6 萬(wàn)個(gè)樣本的 GenSelect 訓(xùn)練數(shù)據(jù)集。
實(shí)驗(yàn)結(jié)果表明，經(jīng)過 GenSelect 加持的模型，其最終準(zhǔn)確率相比簡(jiǎn)單的多數(shù)投票有了顯著提升，尤其是在候選方案數(shù)量不多時(shí)效果更為明顯。雖然由于 AIMO 競(jìng)賽嚴(yán)格的時(shí)間和計(jì)算限制，GenSelect 未能被納入最終的獲勝提交方案中，但這項(xiàng)技術(shù)已被完全整合到此次發(fā)布的 OpenMath-Nemotron-32B 模型中，構(gòu)成了其支持的三大推理模式之一。
基于上述三大支柱和海量數(shù)據(jù)，英偉達(dá)團(tuán)隊(duì)訓(xùn)練了一系列名為OpenMath-Nemotron的模型，參數(shù)規(guī)模涵蓋 1.5B、7B、14B 和 32B。這些模型均基于強(qiáng)大的 Qwen2.5 基座模型進(jìn)行微調(diào)。對(duì)于 1.5B 和 7B 版本，他們甚至使用了專門為數(shù)學(xué)任務(wù)優(yōu)化的 Qwen2.5-Math 版本作為起點(diǎn)。
訓(xùn)練過程采用了監(jiān)督微調(diào)，混合使用了 CoT、TIR 和 GenSelect 三種任務(wù)的數(shù)據(jù)，總計(jì)達(dá) 550 萬(wàn)個(gè)樣本。這意味著同一個(gè)模型可以通過不同的提示（prompt）在 CoT（純文本推理）、TIR（工具集成推理）和 GenSelect（多方案選擇）模式下工作。
為了處理長(zhǎng)達(dá)數(shù)千甚至上萬(wàn)個(gè) token 的長(zhǎng)序列推理，團(tuán)隊(duì)?wèi)?yīng)用了旋轉(zhuǎn)位置編碼（RoPE，Rotary Position Embedding）擴(kuò)展技術(shù)，并將訓(xùn)練過程中的上下文窗口擴(kuò)展到支持長(zhǎng)序列。訓(xùn)練使用了英偉達(dá)自家的 NeMo-Aligner 工具包，并結(jié)合了序列打包、上下文并行等技術(shù)來(lái)加速長(zhǎng)序列訓(xùn)練。此外，他們還采用了檢查點(diǎn)平均（checkpoint averaging）和在更難問題子集上進(jìn)行第二輪微調(diào)等策略，進(jìn)一步提升模型性能。
多項(xiàng)優(yōu)化推理措施
贏得 AIMO-2 競(jìng)賽不僅需要模型本身強(qiáng)大，還需要在極其苛刻的 5 小時(shí)、4x L4 GPU 限制下高效完成推理。這要求團(tuán)隊(duì)在模型選擇和推理優(yōu)化上做出極致權(quán)衡。
他們的最終提交方案基于 OpenMath-Nemotron-14B 模型的一個(gè)早期版本，該版本在一個(gè)稍小的 CoT 數(shù)據(jù)集（僅 DeepSeek-R1 生成）上訓(xùn)練，并進(jìn)行了輕量級(jí)的 TIR 微調(diào)。值得注意的是，他們采用了模型合并技術(shù)，將純 CoT 訓(xùn)練的檢查點(diǎn)和經(jīng)過 TIR 微調(diào)的檢查點(diǎn)進(jìn)行線性組合。這種簡(jiǎn)單而有效的方法，讓他們能夠在保持 TIR 能力的同時(shí)，部分恢復(fù) CoT 模型的生成流暢性和速度優(yōu)勢(shì)，并減少代碼調(diào)用次數(shù)，從而更好地適應(yīng)競(jìng)賽環(huán)境。
為了在有限的時(shí)間內(nèi)最大化解題數(shù)量和準(zhǔn)確率，團(tuán)隊(duì)實(shí)施了多項(xiàng)推理優(yōu)化措施：
首先，他們使用 TensorRT-LLM 將預(yù)訓(xùn)練模型轉(zhuǎn)換為 TensorRT 引擎。這一工具的動(dòng)態(tài)批處理功能通過動(dòng)態(tài)分組推理請(qǐng)求提高了吞吐量，在樣本完成后即刻釋放，減少延遲并優(yōu)化 GPU 利用率。由于樣本是獨(dú)立處理的，批處理可以無(wú)縫混合不同的提示或推理參數(shù)。TensorRT-LLM 還包括自定義注意力內(nèi)核和分頁(yè) KV 緩存等多種優(yōu)化。
在量化方面，團(tuán)隊(duì)優(yōu)先采用 int8 權(quán)重量化（W8A16）和 FP8 量化，相比 BF16 格式速度提升了 1.5 倍，同時(shí)對(duì)準(zhǔn)確率的影響最小。減小的權(quán)重大小還為更大的鍵值緩存釋放了內(nèi)存，允許處理更長(zhǎng)的序列。團(tuán)隊(duì)還使用了蘋果開發(fā)的 ReDrafter 技術(shù)，這是一種循環(huán)推測(cè)解碼方法，使用基于 RNN 的起草器在每個(gè)解碼步驟提出并驗(yàn)證多個(gè) token。他們訓(xùn)練了一個(gè)能夠在每一步提出最多三個(gè) token 的起草器，在大約 65% 的步驟中成功接受所有三個(gè) token，顯著加速了生成過程。
此外，團(tuán)隊(duì)通過將 CoT 和 TIR 檢查點(diǎn)線性組合創(chuàng)建了最終模型，這種策略允許他們控制每個(gè)微調(diào)階段對(duì)最終模型行為的影響程度。最佳模型是使用 CoT0.3+TIR0.7 的組合創(chuàng)建的，這不僅提高了準(zhǔn)確率，還通過減少解決方案長(zhǎng)度和代碼執(zhí)行次數(shù)加速了生成。團(tuán)隊(duì)實(shí)現(xiàn)了一種緩沖策略，為每個(gè)問題分配 350 秒的基本時(shí)間限制，如果一個(gè)問題提前完成，未使用的時(shí)間會(huì)被添加到共享緩沖區(qū)，供后續(xù)問題使用。
團(tuán)隊(duì)還利用了 NeMo-Skills 的異步生成功能實(shí)現(xiàn)批量處理和早停。例如，在 16 個(gè)樣本的批處理中，如果前 4-5 個(gè)完成的樣本就已經(jīng)對(duì)最終答案達(dá)成一致，則取消剩余的生成并繼續(xù)下一個(gè)問題。這種機(jī)制極大地節(jié)約了在簡(jiǎn)單或中等難度問題上可能浪費(fèi)的時(shí)間，為攻克難題爭(zhēng)取了寶貴的時(shí)間窗口。早停策略增加了響應(yīng)相關(guān)性，因?yàn)檩^短的答案往往質(zhì)量更高。
實(shí)驗(yàn)結(jié)果顯示，在 Comp-Math-24-25 測(cè)試集（包含來(lái)自 AIME 和 HMMT 競(jìng)賽的問題）上，團(tuán)隊(duì)的模型表現(xiàn)出色。1.5B 模型在 CoT 模式下單次通過準(zhǔn)確率為 58.2%，多數(shù)投票準(zhǔn)確率達(dá) 80.0%；在 TIR 模式下，這些數(shù)字分別提高到 64.5% 和 83.3%；使用 GenSelect 技術(shù)后，準(zhǔn)確率進(jìn)一步提升至 83.3%。14B 模型的表現(xiàn)更為出色，在 TIR 模式結(jié)合 GenSelect 使用時(shí)，準(zhǔn)確率高達(dá) 90.0%。最大的 32B 模型在相同條件下甚至達(dá)到了 93.3% 的準(zhǔn)確率。這些結(jié)果也表明，無(wú)論模型大小如何，TIR 模式始終優(yōu)于純 CoT 模式，而 GenSelect 技術(shù)能進(jìn)一步提高準(zhǔn)確率。
目前，英偉達(dá)團(tuán)隊(duì)已將完整的 OpenMathReasoning 數(shù)據(jù)集、訓(xùn)練好的 OpenMath-Nemotron 模型系列以及所有相關(guān)代碼以商業(yè)許可方式發(fā)布到 Hugging Face 和 GitHub 上（項(xiàng)目地址：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
運(yùn)營(yíng)/排版：何晨龍
熱門跟貼