新智元报道
编辑：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 模型进行处理，具体流程如下：
-
问题提取：通过大语言模型识别初始帖文中的数学问题。
-
问题分类：采用大语言模型对每个问题进行多维度分类，并剔除所有选择题、二元判断题及无效问题。
-
问题转化：将证明题转化为需要相似解题技巧的答案导向型问题。
-
答案提取：针对非证明题，从论坛讨论中提取最终答案。
-
基准去污：使用基于 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
