NuminaMath 7B TIR 获得 AIMO 进步奖 – 技术回顾

TL;DR

NuminaMath 7B TIR 在首届 AI Math Olympiad (AIMO) 进步奖中通过正确解答 50 道隐藏赛题中的 29 道，展示了两阶段微调方案、大规模高质量数学数据以及自洽+工具集成推理（SC‑TIR）推理策略能够显著提升开源大模型在奥林匹克级数学上的表现。

介绍 Numina – 一个开放的 AI4Maths 项目

Numina 是一个于 2023 年底启动的开源项目，旨在加速 AI 驱动的数学推理。项目由 Jia Li、Yann Fleureau、Guillaume Lample、Stan Polu 和 Hélène Evain 发起，早期得到 Mistral AI 的支持。2024 年初，Hugging Face 的微调专家 Lewis Tunstall 和 Ed Beeching 加入团队，随后获得 General Catalyst 和 Answer.ai 的进一步支持。该合作的首个具体目标是 2024 年 AIMO 进步奖——一场在 Kaggle 平台上对 50 道预选级数学题（相当于 AMC 12/AIME）进行测试的比赛，仅使用 2024 年 2 月 23 日之前发布的开源权重模型。

AI Math Olympiad (AIMO) 进步奖

AIMO 奖项旨在培育能够在国际数学奥林匹克上夺得金牌的 AI，设有 500 万美元的大奖。首届进步奖要求参赛者每日在 Kaggle 平台提交两次解答，每次提交在 P100 或两块 T4 GPU 上运行，最长可达九小时。题目为整数输出的高中竞赛水平题目，设有公开排行榜（50 题）和隐藏的私有排行榜，后者决定最终排名。

获胜方案架构

获胜系统由三大核心组件构成：

  1. 微调推理代理 – DeepSeekMath‑Base 7B 经过两阶段微调，成为能够在自然语言推理与 Python REPL 调用之间交叉的模型。
  2. SC‑TIR 解码算法 – 一种新颖的推理流程，生成大量候选解答，执行嵌入的 Python 代码，并在自洽检查后进行多数投票。
  3. 稳健的内部验证 – 四套精心策划的验证集（AMC、AIME 以及 MATH 基准的两个子集）用于模型选择，防止对公开排行榜的过拟合。

训练使用开源技术栈（TRL、PyTorch、vLLM、DeepSpeed）在单节点八块 H100 GPU 上完成，耗时约十小时。

两阶段训练方案（MuMath‑Code）

该方案遵循 MuMath‑Code 论文，包含：

  • 阶段 1 – 思路链（CoT）微调 – 基础模型学习数十万条带有 CoT 风格解答的自然语言数学题，鼓励逐步推理。
  • 阶段 2 – 工具集成推理（TIR）微调 – 使用 GPT‑4 生成的 ToRA 格式合成数据，每道题配有推理过程、Python 程序及执行结果，教会模型在中间计算时调用 Python REPL。

两阶段均采用全参数微调（不使用 LoRA/DoRA），使用 2048‑token 打包策略、梯度检查点和 DeepSpeed ZeRO‑3 分片。关键超参数在两阶段保持一致：学习率 2e‑5、批大小 32、余弦调度器，热身比例为 0.0（阶段 1）或 0.1（阶段 2）。

数据即一切

Numina 构建了两大数据集：

  • 思路链数据集 – 数十万对来源于中国高中练习册、美国考试 PDF 与国际奥林匹克题库的题目‑解答对。流水线包括 OCR、分割、英文翻译以及转化为 CoT 格式。
  • 工具集成推理数据集 – 约 60 000 道（多数为整数输出）题目经 GPT‑4 流水线处理，生成 ToRA 风格推理、执行代码并过滤不匹配项。每道题最多重生成三次以确保正确性。

作者计划在近期开源完整数据集。

SC‑TIR：自洽 + 工具集成推理

推理阶段模型会出现高方差，因为 Kaggle 随机提供题目且 GPU 资源受限。SC‑TIR 通过以下方式缓解：

  1. 将每道题 N 次复制（获胜运行中 N = 48），形成多样化提示批次。
  2. 采样 N 次完成，每次生成一段 Python 代码。
  3. 执行所有代码块并捕获输出或回溯信息。
  4. 重新采样至 M = 4 层深度，允许模型基于前一次错误自行纠正。
  5. 剔除无效样本，对最终数值答案进行多数投票。

在 T4 GPU 上使用 AutoGPTQ 将模型量化至 8‑bit 至关重要：它将上传时间减半，避免了 bfloat16 不兼容问题，并降低了显存占用，仅带来轻微的精度下降。

防止对公开排行榜的过拟合

由于公开测试集仅含 50 题，团队构建了四套内部验证套件：

  • AMC（83 道整数输出题） – 代表私有测试集，模型解答率约 60‑65%。
  • AIME（90 题） – 难度更高，用于暴露模型失效模式。
  • MATH level 4（754 题） 与 MATH level 5（721 题） – 从 5 000 题的 MATH 基准中筛选出整数答案子集。

通过 5‑10 次随机种子重复评估，测量方差（SC‑TIR 通常为 1‑3%），从而实现可靠的超参数调优。

未进入最终方案的实验

在确定 MuMath‑Code 流水线前，团队尝试了多种备选方案：

  • 纯 CoT 模型配多数投票。
  • MMOS（单步 Python）模型，最高止步 16/50。
  • Kahneman‑Tversky 优化（KTO）策略微调，公开得分提升至 27/50，但未能在最终模型上完成时间限制。
  • 强化学习（PPO 与 REINFORCE‑LOO）结合代码执行奖励——奖励曲线表现良好，却未带来可测量的准确率提升。
  • 扩展至更大骨干模型（InternLM‑20B、CodeLlama‑33B、Mixtral‑8x7B）——在 T4 GPU 上推理速度过慢。
  • 模型合并技术（DARE、TIES、WARP）——导致内部指标回退。

含义与未来方向

NuminaMath 7B TIR 的成功表明，在高质量数学数据、两阶段微调方案以及 SC‑TIR 稳健推理的加持下，开源大模型完全可以在奥林匹克级别的推理中具备竞争力。该方法可扩展至更大模型（团队计划为更大骨干模型发布数据集），并可迁移至其他需要工具使用的领域（如符号计算、科学编程）。

项目的开放性邀请社区贡献者和合作伙伴扩展数据集、改进解码算法或探索新模型架构。持续的社区努力有望让最终目标——能够在国际数学奥林匹克上夺金的 AI——更进一步。

致谢

作者感谢 Thomas Wolf 与 Leandro von Werra 为 Numina–Hugging Face 合作提供的便利，感谢 Hugo Larcher 的 GPU 供应，感谢 Colin Raffel 在模型合并方面的建议，以及 Omar Sanseviero 对博客文章的反馈。额外支持来自 Mistral.ai、General Catalyst、Answer.ai 与北京大学数学研究国际中心。

获取模型与演示

Sources