AI-MO x Numina | 工具集成的数学推理
第二届 AIMO 竞赛将于未来几周在 Kaggle 上正式启动，具体截止日期将在发布时公布。
**这次比赛的难度将提升至 IMO 奥赛级别，**参赛者将面对 100 道全新数学题，要求 AI 展示真正的数学推理能力，而非简单的计算或猜测。本次进步奖的奖金池高达 209.7 万美元，比 2024 年 7 月颁发的首届进步奖翻了一倍。
具体奖项安排：
- 总奖金池为 209.7 万美元。
- Kaggle 平台的私有排行榜将决定 50.79 万美元奖金的分配，奖励给前五支团队。
- 如果某支队伍在公测和私测中均获得 47/50 的分数，奖金池的剩余部分将颁发给该队。如果无人获奖，奖金将滚入下一年度的进步奖。
本篇回顾首届 AI-MO 竞赛的最佳工作 —— Numina，介绍其实现方案和技术细节。
由于获奖队伍必须公开其代码、方法、数据和模型参数。作为第一届 AIMO 的获胜者，Numina 开源了完整的训练代码、数据集及微调后的模型参数，以及最终求解问题的代码，确保工作完全可复现。此外，Numina 还在 HuggingFace 平台上托管了一个在线运行的演示 Demo（不过该服务目前挂了）。
Numina
Numina 于 2023 年底由一群热衷于人工智能与数学的研究者创立，包括 Jia Li、Yann Fleureau、Guillaume Lample、Stan Polu 和 Hélène Evain，得到了 Mistral AI 的初始支持。该项目受到 AIMO 竞赛的启发，专注于开发能够解决高难度数学问题的 AI 模型。
AI-MO
关于 IMO 和 AI-MO 详细的背景介绍，参阅上篇 AI 改写数学竞赛 | 从 IMO Grand 到 AlphaProof，本篇聚焦技术层面。
技术解析
模型微调
简单说，Numina 微调工作包括三部分内容：
-
模型微调：以 DeepSeekMath-Base 7B 为基座模型的进行全量微调，将模型转化为一个“推理代理”，能够通过自然语言推理结合 Python REPL 计算中间结果来解决数学问题。
-
工具集成推理解码算法 (TIR)：团队设计新颖的解码算法，用于工具集成推理，并结合代码执行反馈，在推理过程中生成多个候选解答。
-
内部验证集：为避免模型在公开排行榜上过拟合，Numina 标注了多组内部验证集，帮助指导模型选择。
具体地，微调主要基于 MuMath-Code的研究工作采用两阶段训练，工具集成则依照 ToRA的格式：
-
第一阶段：用 COT 模板构建数据集 NuminaMath-CoT， 并对模型进行初步微调，以增强模型的推理能力。
-
第二阶段：基于 ToRA 方法，使用 GPT-4 生成具有代码执行反馈的 ToRA 格式解答，得到工具集成推理的合成数据集 NuminaMath-TIR。每个数学问题被分解为推理步骤、Python 代码及其输出。用该数据集对模型进一步微调。
一些技术细节：微调过程使用 TRL 的 SFTTrainer 的“打包”功能，将多个样本拼接在 2048 个 token 的块中进行处理。同时，借助梯度检查点与 DeepSpeed ZeRO-3 协议，确保模型的权重、梯度和优化器状态都能在有限的显存中进行有效分片。
实践摸索过程
在最佳解决方案之前，Numina 团队还尝试过其他几种方法，一些关键的实验和结果：
-
纯 CoT 模型 + 多数投票：团队初期尝试了纯 COT 推理，并通过多数投票机制进行评估，取得了 8/50 的成绩。但这种方法未能展现出足够的解题能力。
-
MMOS 数据集的应用：团队进一步尝试了使用 MMOS（Mix of Minimal Optimal Sets）数据集，将问题的解决过程简化为单轮 Python 程序的执行。这种方法提高了模型的表现，达到了 16/50 的成绩，但由于 MMOS 仅生成单轮解答，未能处理复杂的多步推理问题。
-
KTO 优化方法：随后，团队转向了 Kahneman-Tversky Optimization (KTO) 算法。通过从 SFT 模型中每题抽取 4 个解答，与 Ground Truth 比较，生成正负样本，再对 SFT 模型应用 KTO 优化。此方法显著提高了模型表现，达到了 27/50 的成绩。此外，KTO 的一大优势是可以在训练过程中跟踪隐式奖励，帮助调试并优化模型选择。
-
其他模型的尝试：团队测试了多个更大规模的模型，包括 InternLM-20B、CodeLama-33B 和 Mixtral-8x7B，但这些模型在数学问题上表现都不如 DeepSeek 7B，且在推理时存在显著的延迟问题。
-
RLOO 算法的尝试：使用 REINFORCE-leave-one-out (RLOO) 算法与 Proximal Policy Optimization (PPO) 进行强化学习，结合代码执行反馈与奖励机制进行训练。尽管在奖励曲线上看到了些许提升，但最终未能在整体表现上取得显著进步。
-
推理加速与模型合并：在推理加速方面，通过静态 KV 缓存与 torch 编译，可以在 H100 GPU 上将生成速度提升 2-3 倍，但在 Kaggle T4 环境中遇到了各种兼容性问题。团队还实验了多种模型合并技术（如 DARE、TIES 和 WARP），使用 mergekit 工具将 SFT 和 KTO 模型或公共 DeepSeekMath 模型合并。但这些尝试导致内部评估的显著退步，因此未进一步深入探索。
推理算法
除了对模型自身的提升，比如：使用参数更大的模型，构建质量更高的数据集，使用更好的模型基座等等；在模型给定的情况下，推理策略也能显著提升模型性能。比如：
- 通过多轮投票提升模型的自一致性（Self-Consistency)
- 使用外部工具集成 TIR
- 使用策略 + 奖励模型
- 使用 MCTs 树搜索算法，选取合适分支
Numina 的方案组合了前两个策略，SC-TIR 算法：使用 Python 集成推理，并进行多数投票：
算法的具体步骤：
-
生成初始候选解：针对每个问题，复制输入 N 次，作为初始 prompt 批次传递给 vLLM 模型。通过创建多个候选解为后续的多数投票做准备。
-
多轮代码生成：模型通过采样生成 N 个多样化的解答，持续运行直至生成完整的 Python 代码块。此过程可以进行 M 轮，确保代码的完整性。
-
代码执行与反馈：执行每个生成的 Python 代码块，将其输出结果以及可能的错误信息（如追溯错误 traceback）一并返回。这一步可以帮助模型发现并纠正潜在的代码错误。
-
推理轨迹的生成：生成宽度为 N、深度为 M 的推理轨迹，允许模型通过 traceback 的反馈来自我修正代码错误。如果某个候选解未能产生有效结果（如代码块不完整或执行失败），则对其进行修剪。
-
多数投票与后处理：在生成一系列解答后，对候选解进行后处理，并通过多数投票选择最终的答案，确保模型输出结果的可靠性和一致性。
简单说，就是多次代码生成与执行，以及多次采样并进行投票。
这次竞赛的前几名还使用了其他推理策略，后续再开一篇介绍。
数据集和模型
当然，说到底，“Good data is all you need”。
Numina 数据集的来源广泛，包含中国的高中数学练习题、美国和国际数学奥林匹克竞赛（IMO）的问题等。大部分数据来自在线试卷的 PDF 文档和 AoPS 数学论坛。这个论坛包含来自各个国家的数学竞赛题，是 IMO 选手必注册的论坛，我们在上篇 AI 改写数学竞赛 | 从 IMO Grand 到 AlphaProof 也介绍过。
数据的处理过程：对原始 PDF 进行 OCR；细分为问题-解决方案对；翻译成英文；重新调整以产生思想链推理格式；生成最终答案格式。
数据集
Numina 共开源了六个数据集：
| HuggingFace 地址 | 简述 |
|---|---|
| AI-MO/NuminaMath-CoT | 包含 86 万个数学问题及其 CoT 格式的解决方案 |
| AI-MO/NuminaMath-TIR | 从 CoT 数据集中选取 7 万个问题，由 GPT-4 生成的工具集成推理（ToRA）路径和代码 |
| AI-MO/aimo-validation-aime | 包含 AIME 2022-2024 年的 90 个问题及其解决方案 |
| AI-MO/aimo-validation-amc | 包含 AMC12 2022、2023 年的 83 个问题及其解决方案 |
| AI-MO/aimo-validation-math-level-4 | 从 MATH 数据集中提取的 754 个 level 4 难度问题 |
| AI-MO/aimo-validation-math-level-5 | 从 MATH 数据集中提取的 721 个 level 5 难度问题 |
前两个数据集：在搜集问题的基础上，用 GPT-4 和 DeepSeek 进行标注，最后用于模型的两阶段微调。
后四个数据集：用于内部验证集，确保模型的泛化能力。团队使用这四个内部验证集来衡量模型在不同难度数学问题上的表现，并采取措施避免基础模型中的潜在污染，防止模型过度拟合。
除此之外，AIMO 官方提供的资源指南页面也给了一些其他来源。比如英国从小学到高中的竞赛资源：
模型
Numina 团队以 DeepSeek 和 Qwen 为基础模型进行微调，生成了一系列模型用于数学推理任务。这些模型在不同阶段使用 CoT 和工具集成推理（TIR）数据集进行训练，并通过 AutoGPTQ 进行 8-bit 量化，以提升推理效率。
| HuggingFace 地址 | 简述 |
|---|---|
| AI-MO/NuminaMath-7B-CoT | 阶段 1 模型，使用 CoT 数据集对 deepseek-math-7b-base 进行微调 |
| AI-MO/NuminaMath-7B-TIR | 阶段 2 模型，使用 TIR 数据集对上一阶段模型进一步微调 |
| AI-MO/NuminaMath-7B-TIR-GPTQ | 最终模型，使用 AutoGPTQ 进行 8-bit 量化 |
| AI-MO/NuminaMath-72B-CoT | 使用 CoT 数据集对 Qwen/Qwen2-72B 模型进行微调 |
| AI-MO/NuminaMath-72B-TIR | 使用 TIR 数据集对 72B 模型进行进一步微调 |
总结
以上，我们介绍了 Numina 方案的两阶段训练流程，通过全量微调让模型学会链式推理（CoT）和工具集成推理（TIR），并通过代码执行反馈等方式提升模型的推理能力。
在模型微调过程中，最关键的要素是高质量数据集的构建，其次是微调技术和模型基座选择。