NuminaMath 如何赢得第一届 AIMO 进步奖

今年，Numina 和 Hugging Face 合作，参加了首届 AI 数学奥林匹克 (AIMO) 进步奖的角逐。这项比赛涉及对开放式 LLM 进行微调，以解决高中生为国际数学奥林匹克训练而使用的难题。我们很高兴地宣布，我们的模型——NuminaMath 7B TIR——荣获冠军，并在私有测试集中成功解决了 50 道问题中的 29 道 🥳！

在这篇博文中，我们将介绍 Numina 倡议以及我们获奖解决方案背后的技术细节。如果您想直接测试模型以解决最难的数学问题，请查看我们的演示。

让我们深入了解！

1. Numina 介绍——一个开放的 AI4Maths 倡议
2. AI 数学奥林匹克 (AIMO) 奖
3. 我们赢得第一个进步奖的解决方案
4. 训练方案
5. 好数据是你的全部所需
6. 通过自我一致性和工具集成推理（SC-TIR）驯服方差
7. 避免过拟合的诅咒
8. 我们尝试过的其他想法
9. Numina 的未来——寻找贡献者和合作伙伴！
10. 致谢

Numina 介绍——一个开放的 AI4Maths 倡议

数学有一些非常特别的地方。

数学是一个人人都能接触的领域，甚至在儿童学会阅读之前就能接触。历史上最伟大的数学家之一 斯里尼瓦瑟·拉马努金 是自学成才的，他于 1887 年出生在印度的一个普通家庭。而对其他人来说，数学从休闲到职业，无所不包。

数学对人类至关重要，它是我们构建商业、iPhone 和核电站等一切事物的支柱。然而，即使是解决关键应用的数学问题，也可能是一种愉快的体验。

纯粹的数学超越了智力，如同浩瀚的海洋，只有心灵才能航行。

这就是为什么当我们启动 Numina 时，选择开源和开放数据集是自然而然的选择。就像人类智能一样，我们相信人工智能在数学领域的进步应该普世。如果计算机是心灵的自行车，那么人工智能就是它的引擎——为我们这个时代的拉马努金们开启了新的视野。

在 Mistral AI 的初步支持下，Numina 于 2023 年末由一群对 AI 和数学充满热情的人士（李佳、Yann Fleureau、Guillaume Lample、Stan Polu 和 Hélène Evain）共同创立，灵感来自于 Alex Gerko 和 XTX Markets 发起的 AI 数学奥林匹克（AIMO）竞赛。

2024年初，Hugging Face 的两位 LLM 微调专家（👋 Lewis Tunstall 和 Ed Beeching）加入 Numina 团队，以应对 2024 年 AIMO 进步奖的挑战。我们还获得了 General Catalyst 和 Answer.ai 的额外支持，截至 2024 年 3 月，Numina 已汇聚了来自 世界各地的顶尖人才。

团队组建完毕，是时候应对 AIMO 挑战了！

AI 数学奥林匹克 (AIMO) 奖

每年，来自世界各地的高中生都会参加国际数学奥林匹克竞赛——这项竞赛旨在解决代数、几何和数论等领域的六个具有挑战性的问题。为了让您了解其难度，以下是去年的问题之一

2023 年 11 月，AIMO 奖启动，旨在推动在数学推理方面表现出色的人工智能模型的开放开发。凡能创造出在 IMO 中获得金牌的人工智能模型者，将获得 500 万美元的巨额奖金。除了大奖，AIMO 还引入了一系列“进步奖”来标记实现这一最终目标的里程碑。第一个进步奖以Kaggle 竞赛的形式举行，其问题比 IMO 中的问题“不那么具有挑战性”，但达到了 IMO 预选赛的水平。以下是一个示例问题，它比上面的 IMO 示例更容易解决，但对 LLM 来说仍然很棘手

设 $k, l > 0$ 为参数。抛物线 $y = kx^2 - 2kx + l$ 与直线 $y = 4$ 在 $A$ 和 $B$ 两点相交，这两点之间距离为 6。那么 $A$ 和 $B$ 到原点距离的平方和是多少？

本次竞赛分为两组各 50 道题目，分别用于公共和私人排行榜，题目对参赛者隐藏。这些问题的难度与 AMC12 和 AIME 考试相当，需要整数输出进行验证。最终排名由私人排行榜决定。参赛者每天可以提交两次解决方案，并且只能使用在 2 月 23 日之前发布的开放权重模型。每次提交分配一块 P100 GPU 或两块 T4 GPU，最长 9 小时用于解决 50 道问题。

考虑到这些限制和规则，战略性选择对于开发我们的获胜解决方案至关重要。

我们赢得第一个进步奖的解决方案

经过竞赛中的多次迭代，我们针对第一个进步奖的解决方案包含三个主要部分：

• 一个微调 DeepSeekMath-Base 7B 的方案，使其能够作为一个“推理代理”，通过自然语言推理和使用 Python REPL 计算中间结果来解决数学问题。
• 一种新颖的工具集成推理（TIR）解码算法，带有代码执行反馈，用于在推理过程中生成候选解决方案。
• 我们使用各种内部验证集来指导模型选择，并避免对公共排行榜过拟合。

我们使用了多种开源库来训练模型，特别是 TRL、PyTorch、vLLM 和 DeepSpeed。在 8 块 H100 GPU 的一个节点上，我们的模型训练耗时 10 小时。

训练方案

我们的微调方案主要基于 MuMath-Code 论文，其中包括两个阶段的模型训练：

MuMath-Code 论文中的两阶段训练方法

• 第一阶段：在一个大型、多样化的自然语言数学问题和解决方案数据集上对基础模型进行微调，其中每个解决方案都使用思维链 (CoT) 进行模板化，以方便推理。

• 第二阶段：在第一阶段模型的基础上，使用合成的工具集成推理数据集进行微调。在该数据集中，每个数学问题都被分解为一系列的推理、Python 程序及其输出。在这里，我们遵循微软的 ToRA 论文，并提示 GPT-4 以 ToRA 格式生成带有代码执行反馈的解决方案。对这些数据进行微调，可以生成一个推理代理，该代理可以通过自然语言推理和使用 Python REPL 计算中间结果的组合来解决数学问题（参见下面的截图）。

  

  ToRA 论文中关于我们训练模型所使用的工具集成推理格式的图示。

我们在两个阶段都进行了“完全微调”，即在反向传播过程中更新了所有模型权重。换句话说，我们没有使用 LoRA 或 DoRA 等参数高效技术，因为我们不确定它们能否在不进行大量实验的情况下达到完全微调的性能。我们使用了 TRL 的 SFTTrainer 的“打包”功能，将多个样本连接成一个 2048 个令牌的单一块。所有模型都使用梯度检查点进行训练，并使用 DeepSpeed ZeRO-3 协议进行分片，以确保权重、梯度和优化器状态能够适应可用的 VRAM。下面是我们每个阶段使用的主要超参数：

我们最初的提交只使用了在第一阶段微调过的 DeepSeek 7B 模型，但我们发现其性能相当有限，在公共排行榜上使用 maj@32 的最佳得分仅为 8/50。是 Abdur Rafae 的 公共奖项笔记本 促使我们考虑将代码执行集成到训练方案中。最初，我们专注于 最小最优集混合 (MMOS) 数据集，正如笔记本标题所述。我们发现使用 MMOS 确实提高了性能，但在公共排行榜上使用 maj@32 的得分仍然停留在 16/50，这可能是因为 MMOS 只包含单步解决方案（即模型只生成一个 Python 程序，这对于难题来说是不够的）。我们后来意识到 MMOS 是一个误称，Kaggle 笔记本实际上运行的是 DeepSeekMath 7B RL 模型，该模型能够进行多步推理和代码执行。

此时，我们专注于生成一个类似于 DeepSeekMath Instruct / RL 模型所使用的数据集，再加上 MuMath-Code 方案，这带来了显著的改进。

让我们来看看我们是如何构建这些数据集的。

好数据是你的全部所需

在数据集方面，我们广泛参考了 DeepSeek Math 和其他学者的研究方法，并对其进行了显著扩展。这形成了一个包含数十万个问题-解决方案对的微调数据集，涵盖了从高中数学到竞赛级数学的各种主题。该数据集将在未来几周内完全开源，并可能包含更大的模型，以观察我们的方案如何扩展。有关数据集构建的详细信息，请参阅我们即将发布的数据集技术报告。在本次进步奖竞赛中，我们构建了两个数据集来微调我们的模型。

思维链

该数据集包含数十万个问题，每个问题都以思维链的方式给出解决方案。数据集的来源涵盖了从中国高中数学练习题到美国和国际数学奥林匹克竞赛题。这些数据主要从在线试卷 PDF 和数学讨论论坛收集而来。

处理步骤包括：

1. 从原始 PDF 进行 OCR。
2. 将问题-解决方案对进行分割。
3. 翻译成英文。
4. 重新对齐以生成思维链推理格式。
5. 最终答案格式化。

工具集成推理

工具集成推理（TIR）在本次竞赛中扮演着至关重要的角色。然而，收集和标注此类数据既耗时又昂贵。为了解决这个问题，我们从 Numina 数据集中选择了大约 60,000 个问题，重点关注那些具有数值输出（其中大部分是整数）的问题。

然后，我们利用 GPT-4 生成 TORA 样式的推理路径，执行代码并生成结果，直到解决方案完成。我们过滤掉最终答案与参考不匹配的解决方案，并重复此过程三次，以确保准确性和一致性。这种迭代方法使我们能够高效地生成高质量的 TORA 数据。

作为参考，以下是我们的第一阶段模型 NuminaMath-7B-CoT 和最终第二阶段模型 NuminaMath-7B-TIR 在 MATH 基准测试上的表现，与其他开放和专有模型进行对比：

MATH 基准测试上的表现。除非另有明确说明，所有数据均通过零样本贪婪解码获得。

通过自我一致性和工具集成推理（SC-TIR）驯服方差

正如其他参赛者所指出的，本次比赛在模型提交和评估方面存在若干挑战：

• 评估 API 以随机顺序提供问题，因此早期停止等策略会产生高方差，因为一次运行可能在一开始就有更多难题，从而导致剩余时间减少（反之亦然）
• LLM 推理中的大多数创新都需要访问现代 GPU，因此 Flash Attention 2 或 torch.compile 等标准方法无法在 T4 GPU 上运行。同样，bfloat16 等现代数据类型也不受支持，这促使我们探索 AWQ 和 GPTQ 等后训练量化方法。

最初，我们使用 Abdur Rafae 的 公共笔记本 进行提交，但发现高方差是个问题。为了解决这个问题，我们采用了基于工具集成推理的不同方法：

1. 对于每个问题，将输入复制 N 次，以定义要馈送 vLLM 的初始提示批次。这有效地定义了用于多数投票的候选数量。
2. 采样 N 个不同的补全，直到生成完整的 Python 代码块。
3. 执行每个 Python 代码块并连接输出，包括追溯（如果出现）。
4. 重复 M 次，生成大小为 N 且深度为 M 的批次，允许模型使用追溯自纠代码错误。如果样本未能产生合理输出（例如，不完整的代码块），则剪枝该结果。
5. 对候选解决方案进行后处理，然后应用多数投票选择最终答案。

对于我们的获胜提交，我们生成了 N=48 个候选，深度为 M=4。增加任何一个参数都不会提高性能，因此我们采取了保守的方法，以确保在时间限制内完成。实际上，此算法通过工具集成推理增强了与 CoT 的自我一致性（如下所示）。

我们发现我们的 SC-TIR 算法在我们的内部评估和公共排行榜上都产生了更稳健的结果，方差显著降低。

值得一提的一个技术细节是，我们发现将模型量化为 8 位精度很有帮助。原因有三：

• 将模型上传到 Kaggle Hub 非常慢，压缩模型使这一步的速度提高了一倍。
• T4 GPU 不支持 bfloat16，转换为 float16 会导致模型性能下降。转换为 float32 不可能，因为那会超出可用的 GPU 内存。
• 此外，一个 16 位模型仅加载权重就需要消耗大约 32GB 的显存。对于 2xT4s，这将需要操作 KV 缓存才能快速运行，我们发现权衡模型精度以换取速度是有益的。

我们使用 AutoGPTQ 和训练数据集进行校准，对模型进行了量化。实际上，这导致了准确性的小幅下降，但为适应 Kaggle 平台上的评估所施加的约束提供了最佳折衷方案。

避免过拟合的诅咒

过拟合公共排行榜是 Kaggle 竞赛中的常见风险，尤其是在测试集只有 50 个问题时更是如此。此外，规则规定每天最多提交两次，这使得鲁棒的内部验证数据集对于我们控制开发进度至关重要。根据 AIMO 团队的规定，测试题难度中等，介于 AMC12 和 AIME 之间，并且需要整数输出。

为了指导模型选择，我们使用了四个内部验证集来评估模型在不同难度数学问题上的性能。为了避免基础模型中可能存在的污染，我们从 AMC12（2022、2023）和 AIME（2022、2023、2024）中选择了问题，创建了两个内部验证数据集：

• AMC（83 个问题）：我们选择了 AMC12 22、AMC12 23 中的所有问题，并保留了那些可以转换为整数输出的问题。这形成了一个包含 83 个问题的数据集。这个验证集旨在代表 Kaggle 上的私有测试集，因为我们从比赛描述中得知这些问题的难度与此级别或更高。我们发现我们的模型可以解决这些问题的大约 60-65%。为了衡量方差，我们每次评估都使用 5-10 个不同的种子进行，通常使用我们的 SC-TIR 算法可以看到大约 1-3% 的变化。
• AIME（90 个问题）：我们选择了 AIME 22、AIME 23 和 AIME 24 中的所有问题，以衡量我们的模型在难题上的表现，并评估最常见的失败模式。与上文一样，我们每次评估都使用 5-10 个种子来衡量变化。

由于 AMC/AIME 验证集规模较小，模型在这些数据集上的性能容易受到噪声影响，这与公共排行榜类似。为了更好地评估我们模型的性能，我们还使用 MATH 测试集的一个子集对其进行了评估，该测试集包含 5,000 个问题。我们只保留了具有整数输出的问题，以简化多数投票并模仿竞赛评估。这产生了两个额外的验证集：

• MATH 等级 4（754 个问题）
• MATH 等级 5（721 个问题）

通过使用这四个验证集，我们能够从不同的训练阶段中挑选出最有希望的模型，并缩小超参数的选择范围。我们发现，在这种特殊的比赛中，结合小型但具有代表性的验证集和大型验证集很有用，因为每次提交都会受到抽样带来的一些随机性影响。

我们尝试过的其他想法

如上所述，我们尝试了几种方法，但最终都放弃了，转而使用 MuMath-Code 配方：

• 训练纯 CoT 模型并使用多数投票进行评估。
• 训练 MMOS 模型以一步解决 Python 问题。

我们尝试的另一种技术是对 SFT 模型采样的新补全应用 Kahneman-Tversky 优化 (KTO)。这里的方法与 OrcaMath 类似，即：

• 使用 SFT 模型为每个问题采样 4 个补全，交错推理和代码执行。我们使用第二阶段的 SFT 数据集作为提示源。
• 提取答案并与真实值进行比较。如果正确，将样本标记为正，否则标记为负。
• 将 KTO 应用于该数据集上的 SFT 模型。

我们发现这种形式的在线 KTO 产生了一个比 SFT 模型略好的模型（在我们的内部评估中提高了几个百分点），并在公共排行榜上获得了 27/50 的分数。

KTO 的一个优点是可以在训练期间跟踪隐式奖励，这确实有助于调试运行——例如，这是我们成功的训练日志之一，可以看到所选（即正确解决方案）奖励在训练过程中增加，而被拒绝的奖励则受到抑制。

不幸的是，我们没有足够的时间将这种方法应用于我们的最终 SFT 模型，所以我们可能能够多解决 1-2 个问题！

我们还尝试将 SFT 方案应用于更大的模型，如 InternLM-20B、CodeLlama-33B 和 Mixtral-8x7B，但发现 (a) DeepSeek 7B 模型由于在数学上的持续预训练而难以超越，(b) 在 2xT4 GPU 上的推理速度非常慢，并且我们遇到了许多无法追踪根源的神秘超时问题。

另一个失败的实验包括尝试使用强化学习（特别是近端策略优化算法和 REINFORCE-leave-one-out (RLOO) 算法），通过代码执行反馈和针对编写代码以及获得正确/不正确解决方案的形状奖励。我们将其应用于 DeepSeekMath 7B RL 模型。虽然我们看到了一些有希望的奖励曲线，但我们没有看到性能的显著提升。鉴于 RLOO 等在线方法受到文本生成和迭代缓慢的瓶颈，我们放弃了强化学习，转而尝试 KTO。

在推理方面，我们还尝试了：

• 使用静态 KV 缓存和 torch 编译。我们发现在 H100 上可以将原生 transformers 代码的生成速度提高 2-3 倍，但在 Kaggle T4s 上遇到了各种神秘错误，主要是由于 Accelerate 中 torch 编译不支持模型分片。

各种模型合并技术，如 DARE、TIES 和 WARP。在这里，我们使用 mergekit 来合并 SFT 和 KTO 模型，或者 SFT 模型与公共 DeepSeekMath 模型。总的来说，我们发现这些合并导致了我们在内部评估中出现了显著的性能下降，并且我们没有足够的时间更深入地探索这一点。

Numina 的未来——寻找贡献者和合作伙伴！

在 Numina 成功赢得 AIMO 2024 进步奖之后，我们现在旨在继续我们的使命，即促进人工智能和人类智能在数学领域的发展。您可以访问我们的网站了解更多项目信息，并随时通过 contact@projectnumina.ai 与我们联系。

Numina，如同数学本身，旨在向所有愿意通过人工智能推动数学进步的全球人才和支持者开放！

致谢

我们感谢 Thomas Wolf 和 Leandro von Werra 促成了 Numina 和 Hugging Face 的合作。我们还要感谢 Hugo Larcher 帮助 Hugging Face 集群上的 GPU 全速运行，Colin Raffel 在模型合并方法方面的建议，以及 Omar Sanseviero 对博文的反馈。

我们还要向 Mistral.ai、General Catalyst、Answer.AI 以及北京大学北京国际数学研究中心致以诚挚的谢意，他们从项目伊始就给予了支持。

最后，我们感谢 AIMO 奖团队举办了如此激动人心和富有启发性的竞赛！