Kaggle | 总奖池超200万美金的数学竞赛AIMO2冠军方案
大家好，我是HxShine
今天分享一篇来自 NVIDIA 的文章（kaggle AIMO2冠军方案，冠军奖金高达26w美金，总奖池超200w美金），Title: AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset。
该方法基于三个关键：1）大规模高质量数学数据集 OpenMathReasoning；2） Tool-Integrated Reasoning (TIR)，将代码执行与长推理链结合；3） Generative Solution Selection (GenSelect)：用于从多个候选解决方案中选择最有希望的一个结果来替代majority vote。
结合这些思想，他们训练了一系列模型，在数学推理基准上取得了最先进的结果，并在 AIMO-2 竞赛中成功解决了 50 个难题中的 34 个（对比deepseek r1满血版大概只能解决20个左右）。
一、概述
- • Title: AIMO-2 Winning Solution: Building State-of-the-Art Mathematical Reasoning Models with OpenMathReasoning dataset
- • URL: https://arxiv.org/abs/2504.16891
- • Authors: Ivan Moshkov, Darragh Hanley, Ivan Sorokin, Shubham Toshniwal, Christof Henkel, Benedikt Schifferer, Wei Du, Igor Gitman (NVIDIA)
- • Code: https://github.com/NVIDIA/NeMo-Skills, https://huggingface.co/collections/nvidia/openmathreasoning-68072c0154a5099573d2e730
1 Motivation
- • 解决像 AI 数学奥林匹亚竞赛（AIMO-2）这样具有挑战性的数学难题，需要模型具备强大的复杂推理能力，现有公开模型在这方面存在局限性。
- • 提高大型语言模型在数学领域的长推理和工具使用能力，现有模型（特别是开源模型）通常难以直接生成高质量的工具集成推理（TIR）解决方案，或难以有效从多个候选方案中选出最优解。
- • 为数学推理领域提供高质量、大规模的训练数据和评估基准，促进社区的进一步研究和进展。
2 Methods
省流版总结:
通过3个步骤来构建 SOTA 数学推理模型：
- 1. 数据构建: 从 AoPS 社区论坛收集大规模数学问题，通过 LLM pipeline处理，构建包含 540K 问题、3.2M CoT 解法和 1.7M TIR 解法的高质量数据集 OpenMathReasoning。
- 2. 工具集成推理 (TIR): 开发迭代式训练、生成、过滤pipeline，使模型能有效结合 Python 代码执行进行长推理。通过指导跟随模型生成初始 TIR 数据，并严格过滤掉非必要或冗余的代码使用，然后用过滤后的数据迭代微调更强大的模型。还设计了控制代码块使用数量的方法。
- 3. 生成式解决方案选择 (GenSelect): 构建pipeline训练模型，使其能够从多个候选解决方案摘要中选择最优方案，以期超越简单的majority vote多数投票。虽然在最终竞赛提交中因时间限制未能完全集成，但在本地评估中显示了潜力。
最终通过在 OpenMath-Nemotron 系列模型（1.5B至32B）上进行多任务 SFT，并针对竞赛环境进行推理优化（如 TensorRT-LLM、投机解码、模型合并、时间管理和沙箱代码执行），赢得了比赛。
详细方法和步骤:
1 OpenMathReasoning 数据集构建 (§2)
从 AoPS 社区论坛收集大规模数学问题，通过 LLM pipeline处理，构建包含 540K 问题、3.2M CoT 解法和 1.7M TIR 解法的高质量数据集 OpenMathReasoning。
| 步骤 | 描述 |
| 问题收集 | 从 Art of Problem Solving (AoPS) 社区论坛收集大量数学问题。 |
| LLM pipeline处理 | 使用 Qwen2.5-32B-Instruct 构建处理pipeline，包括以下子步骤： |
| - 问题提取 | 从论坛帖子中提取所有数学问题。 |
| - 问题分类 | 将问题分为证明题、选择题、二元问题等，并移除选择题、二元问题和无效问题。 |
| - 问题转换 | 将证明题转换为需要数值或表达式答案的问题。 |
| - 答案提取 | 尝试从论坛讨论中提取非证明题的最终答案。 |
| - 基准污染去污 | 使用 LLM 比较移除与流行数学基准高度相似的问题。 |
| 长推理 CoT 解决方案合成 (§2.3) | 使用 DeepSeek-R1 和 QwQ-32B 模型生成多达 32 个 CoT 解决方案候选，并通过过滤保留高质量方案。最终得到 3.2M CoT 解决方案。 |
下面是一个从论坛抽取problem的prompt：
2 Tool-Integrated Reasoning (TIR) (§3)
为了使模型能有效结合 Python 代码执行进行长推理。通过指导instruct model生成初始 TIR 数据，并严格过滤掉非必要或冗余的代码使用，然后用过滤后的数据迭代微调更强大的模型。同时还设计了控制代码块使用数量的方法。
| 步骤 | 描述 |
| 挑战 | DeepSeek-R1 和 QwQ-32B 等现有模型难以直接生成高质量 TIR 解决方案。 |
| 迭代生成和过滤pipeline | 包括以下子步骤： |
| - 初始 TIR 生成 | 使用 LIMO-Qwen-32B 生成初始 TIR 解决方案（stage-0 TIR 数据）。 |
| - 严格过滤 (§3.2) | 使用 Qwen2.5-32B-Instruct 分类代码块并结合规则过滤，得到 15k 高质量 stage-0 TIR 样本。 |
| - 迭代改进 (§3.3) | 使用 stage-0 数据微调 QwQ-32B，生成 700K TIR 解决方案，过滤后得到 260K；再使用 CoT 数据微调过的 14B 模型进行最终轮次生成和过滤，最终得到 1.7M TIR 数据集。 |
| 控制代码块数量 (§3.4) | 在生成过程中明确告知模型剩余可用的代码执行次数，并通过移除超出限制的生成来强化行为。 |
下面是第一步生成使用python工具的prompt example:
下面继续提供一个使用训练后的TIR模型进行推理的prompt：
3 Generative Solution Selection (GenSelect) (§4)
为了使其能够从多个候选解决方案摘要中选择最优方案，以期超越简单的majority vote多数投票。训练了GenSelect方案。我在想他们为啥没用RL方案来讲majority vote的能力转移到pass@1上面来，同时提升速度？
| 步骤 | 描述 |
| 目标 | 从多个候选解决方案摘要中选择最优方案，弥合 pass@k 与 majority@k 之间的差距。 |
| 摘要生成 (§4.1) | 使用 Qwen2.5-32B-Instruct 为每个解决方案生成简洁摘要（最多 2048 tokens），并过滤不忠实的摘要。 |
| 选择候选生成 (§4.2) | 从数据集中随机抽取问题，对每个问题随机采样 2-16 个候选解决方案摘要，确保至少包含一个正确和一个不正确方案。重复此过程为每个问题生成 8 个比较组。 |
| 训练数据生成 | 使用 QwQ-32B 模型评估比较组，选出最有希望的解决方案。将选择正确解决方案的推理轨迹作为 GenSelect 训练数据。生成 1M 选择结果，过滤后得到 565K 样本。 |
| 降低计算成本 (§4.3) | 探索训练模型直接生成比较摘要而不是完整的推理轨迹，提高 GenSelect 推理效率。 |
数据构建流程如下：
GenSelct推理prompt（就是先生成多个solution，然后利用LLM进行汇总得到最终solution，优点是理论上可以比majority vote强）:
4 模型训练与评估 (§5)
| 步骤 | 描述 |
| 模型训练 | 在 Qwen2.5-Base 系列模型（1.5B, 7B, 14B, 32B）上进行 SFT，修改 RoPE 基数到 500K。 |
| 数据集 | 在包含 CoT、TIR 和 GenSelect 三种任务的混合数据集（总计 5.5M 样本）上训练所有模型。 |
| 优化器与学习率 | 使用 AdamW 优化器、cosine 学习率衰减和线性 warmup。 |
| 评估 | 在 Comp-Math-24-25 基准和 HLE-Math 数据集上评估模型性能。 |
5 Kaggle 提交细节 (§6)
| 步骤 | 描述 |
| 模型选择 | 使用 14B 模型的一个早期版本进行提交。 |
| 训练配方 (§6.1) | 在 2.2M CoT 数据子集上训练 8 个 epoch，然后在 15k stage-0 TIR 样本上进行轻量级微调。 |
| 模型合并 (§6.2) | 使用线性组合方式合并 CoT 和 TIR 检查点，以提高准确性并减少生成长度和代码执行次数。 |
| 推理优化 (§6.3) | 实施 TensorRT-LLM 优化（in-flight batching, 异步生成, 量化 W8A16/FP8）和投机解码（ReDrafter）。 |
| 时间管理 (§6.3.4) | 实现缓冲策略，动态分配每个问题的计算时间。 |
| 代码执行集成 (§6.3.5) | 使用 Nemo-Skills 的沙箱环境进行代码执行，并设置严格的安全措施（最大调用次数、超时限制、输出限制）。 |
3 Conclusion
- • 成功赢得 AIMO-2 竞赛： 提出的方法在 AIMO-2 竞赛中取得第一名，在私有测试集上解决了 50 个问题中的 34 个，奖金是美元，听说他们捐了!!!。
- • 构建并发布了大规模高质量数据集 OpenMathReasoning： 包含 540K 问题、3.2M CoT 解决方案、1.7M TIR 解决方案和 566K GenSelect 轨迹，为社区提供了宝贵的资源。
- • 开发了有效的 TIR 和 GenSelect 方法： 证明了通过迭代生成和过滤可以构建强大的 TIR 模型，并通过训练选择模型 (GenSelect) 可以显著提高从多个候选解中选出正确答案的准确率。
- • 发布了一系列 SOTA 开源数学推理模型 OpenMath-Nemotron： 这些模型（1.5B 到 32B）在 CoT、TIR 和 GenSelect 模式下均表现出色，在开源模型中达到了新的 SOTA 水平。
4 Limitation
- • GenSelect 未能在 Kaggle 提交中使用： 由于竞赛严格的时间和计算限制，最终未能将显示出巨大潜力的 GenSelect 方法集成到获胜方案中。
- • TIR pass@1 性能问题： 对于较小的模型，TIR 模式下的 pass@1 准确率可能不如 CoT 模式，这表明小模型在有效和一致地使用工具方面可能存在困难。
- • 评估差异： 本地评估结果与 Kaggle 排行榜得分之间存在差异，可能源于评估集的不同或模型在特定难题上的表现差异。
- • 计算资源需求： 构建数据集、训练模型（尤其是迭代生成和过滤）以及进行 GenSelect 推理都需要大量的计算资源。
- • 摘要生成的依赖： GenSelect 方法依赖于高质量的解决方案摘要生成，这本身也是一个挑战。
二、详细内容
1 GenSelect在采样数量比较少的时候有时明显
说明：图2包含两个子图，分别展示了 14B CoT 模型 (a) 和 14B TIR 模型 (b) 在 Comp-Math-24-25 基准上，pass@k、GenSelect 和 majority voting 三种评估指标的准确率随生成样本数量变化的曲线。
总结：随着生成样本数量的增加，所有指标的准确率基本呈上升趋势。GenSelect 在多数样本数量下优于 majority voting，但随着样本数过多（>32），由于上下文限制，GenSelect 性能可能会变得不稳定，其主要提升体现在较少生成数量时。
2 GenSelect数据构建pipeline
说明：通过采样正确和不正确实例来构建 GenSelect 的输入，确保输入包含至少一个正确和不正确解决方案。这些输入被喂给 QwQ-32B 模型进行评估和选择，选出最优解。
特点：使用 Qwen2.5-32B-Instruct 对选择正确解决方案的推理轨迹进行摘要化，生成 GenSelect 任务的输出数据。
3 不同数据集以及投票方法的性能差异
说明：分别展示了 1.5B 模型 (a) 和 14B 模型 (b) 在 Comp-Math-24-25 基准上，CoT 和 TIR 方法的 majority@64 和 pass@1 准确率随训练 epoch 变化的曲线。
总结1: 随着训练 epoch 数的增加，模型性能普遍提升。
总结2: 较小模型 (1.5B) 的提升更为显著且需要更长的训练时间来获得有意义的改进，而 14B 模型相对训练初期性能更高，后期提升趋缓。
4 OpenMath-Nemotron 系列模型、DeepSeek-R1 和 QwQ-32B 模型在不同基准上的详细评估结果对比
总结：OpenMath-Nemotron 模型在各个基准上均达到了最先进水平(居然最好的效果能领先R1满血版很多)，且 GenSelect 方法能显著提升准确率。
三、总结
结论1：通过结合高质量大规模数据集、工具集成推理和生成式选择方法，可以构建在挑战性数学问题上达到最先进水平的语言模型。 本文的方法在 AIMO-2 竞赛中的胜利以及在多个数学基准上的 SOTA 结果验证了这一策略的有效性。
结论2：发布 OpenMathReasoning 数据集、OpenMath-Nemotron 模型和代码是社区的宝贵财富。 该大规模数据集（540K 问题, 3.2M CoT, 1.7M TIR）为数学推理领域的研究提供了丰富的训练和评估资源，降低了研究门槛，有望加速该领域的进展。
结论3：虽然 TIR 和 GenSelect 显示出巨大潜力，但在实际应用（如竞赛）中仍面临挑战。 TIR 生成的长推理和工具使用更容易产生未完成的解决方案（尤其小模型），GenSelect 虽然理论上能提升性能，但其对样本数量的依赖和推理成本需要在实践中进一步优化以适应严格的计算限制。如何平衡长推理的彻底性与推理效率和可靠性仍是重要课题。
