Notes on OpenMath-Nemotron

Wed, 15 Apr 2026 09:50:39 +0800

NVIDIA 在 AIMO-2 比赛中的 winning solution.

Introduction

AIMO 是一个专门针对 AI 模型的数学奥林匹克竞赛，参赛选手需要训练模型在有限的计算资源下解决 50 个高难度数学问题。

NVIDIA 主要进行了三点改进：

构建了包含 540K 高质量的数学推理数据集以及 3.2M long reasoning solutions
构建了一个集成 code execution 的方法来合成 1.7M TIR solutions
构建了一个 pipeline 来训练模来从多个 candidates 中选出最好的答案

最终，NVIDIA 的方案解决了 34 个数学问题。

Methods

Data Preparation

首先作者从 AoPS 上收集得到大量数学问题，然后剔除掉初中级别的数学题目，接下来，作者构建了一个 pipeline 用于提取问题以及对应的答案，这里作者使用了 Qwen2.5-32B-Instruct 模型来完成这个任务，处理流程有：

使用 LLM 来识别并提取所有的问题
对问题进行分类：proof, MCQ, binary, valid
transformation: 将证明题展缓为 answer based question
answer extraction: 提取最终答案
decontamination: 移除与 math benchmark 高度相似的数据

经过数据处理之后，每个阶段的数据量如下表所示

Pipeline Stage	Data Size
Original forum discussions	620K
Extracted problems	580K
Removing “bad” problems	550K
Benchmark decontamination	540K

最终数据集的统计如下

Subset	Size
Converted proofs	260K
With extracted answer	190K
No extracted answer	90K
Total	540K

接下来，作者基于 AIME 和 HMMT 构建了一个验证集，验证集包含了 256 个问题。

为了生成 cot solution, 作者首先使用 Qwen2.5-72B-Math-Instruct 来给数据进行难度分级，接下来作者使用 DeepSeek-R1 和 QwQ-32B 来生成 32 个 candidate solution. 最后作者使用 Qwen2.5-32B-Instruct 来 judge 生成的答案，最终的数据汇总如下所示

Model	all	after filtering
QwQ-32B	1.0M	0.5M
DeepSeek-R1	4.2M	2.7M
Total	5.2M	3.2M

Tool Integrated Reasoning

作者认为对于 instruction model, 我们可以在 long reasoning 过程中保持其 instruction following 能力，然后作者使用 LIMO-Qwen-32B 来生成带有 python code execution 的 solution, 最终得到了 1.2M solution.

然而，作者发现合成的数据里有一部分工具调用是无效的，为了过滤数据，作者使用了 Qwen2.5-32B-Instruct 来对 code block 进行分类，分类使用了两个标准：

novel calculation/verification: 代码执行是否产生了比较好的结果或者简化了之前的步骤
significant/modrate/trivial. 代码是否完成了 solution 比较重要的一部分

作者还移除了结果不对的数据以及包含 2 个以上 code block 的数据，经过这个过程后，得到了 stage-0 TIR data, 包含 15k samples.

然后，作者使用 QwQ-32B 来在 stage-0 TIR data 上微调 7 个 epoch, 再生成 700K 样本并过滤得到 260K 样本。

作者最后使用了一个 14B 的模型重复上面的过程，这个模型训练使用了 QwQ-32B 生成的答案。

最终，TIR 数据集包含了 1.7M 样本。为了控制模型工具调用的次数，作者还在 prompt 里说明了可用的工具调用次数

Generative Solution Selection

在本节中，作者构建了一个从 candidate solution summaries 中挑选最优答案的方法。作者发现从多个 solution 中进行对比然后挑选出最优答案比给单独答案进行打分效果更好。

首先，作者使用 Qwen-2.5-32B-Instruct 来给 solution 生成 summary, summary 的 max length 为 2028 token. 如果 Summary 的答案和预测答案不一致，样本会被丢掉。

为了实现 GenSelect inference, 作者首先随机采样多个 candidate solution summary 组成 sample group, 然后使用 QwQ-32B 来进行挑选过滤得到 565K 数据集

为了进一步减少计算消耗，作者使用 summary 来进行训练，结果发现模型的准确率下降不到 $1-2\%$. 作者通过实验发现 GenSelect 可以有效提高模型 CoT 以及 TIR 场景下的准确率。

Experiments

作者基于 Qwen-2.5-Base 系列模型进行训练，包括 1.5B, 7B, 14B 和 32B.

为了有效提高模型上下文，作者将 RoPE 的 base frequency 提高到 500K.

训练使用了三个任务：

COT solution generation
TIR solution generation
GenSelect

最终数据集包含了 5.5M 样本（3.2M CoT, 1.7M TIR, 566K GenSelect)

使用了 NeMo-Aligner 来加速训练，使用 model merging 来得到更好的效果。

试验结果如下图所示

作者还构建了一个更难问题的数据集，包含了 2.2M 样本，作者使用了 32B 的模型进行训练，结果如下图所示

Model	First SFT	Second SFT
1.5B CoT	55.1	58.2
1.5B TIR	64.1	64.5
7B CoT	61.3	62.5
7B TIR	71.1	70.7
14B CoT	62.9	65.2
14B TIR	74.6	73.4

Submission

针对 Kaggle 竞赛，作者进一步进行了改进

首先，作者使用 Qwen2.5-14B-Base 在 2.2M CoT solutions 上进行训练 8 epoch. 然后作者在 15K Stage-0 TIR samples 上轻量化微调，最后作者对 CoT checkpoint 以及 TIR checkpoint 进行 model merging. 作者并没有使用 GenSelect.

Model	maj@16	pass@16	length	code
CoT	62.9	76.2	11203	-
TIR	66.8	80.1	15834	2.73
CoT_0.3 + TIR_0.7	69.1	81.3	12489	0.85

在推理上，作者使用了 TensorRT-LLM 将模型转换为 TensorRT 格式

作者还使用了 ReDrafter 来实现 speculative decoding.

Conclusion

作者在本文中构建了一个大规模的数学推理数据集 OpenMathReasoning, 包含了 540K 问题，3.2M long CoT solutions, 1.7M TIR solutiosn 以及 566K GenSelect traces.

作者还训练得到了 OpenMath-Nemotron 模型系列，该模型系列在 AIMO-2 上取得了 SOTA 的成绩。

AIMO on Mao Song(毛松)'s Homepage