Numina 第一名解决方案

感谢 AIMO 和 Kaggle 举办这次非常有趣的竞赛——我们学到了很多关于开源大语言模型 (LLM) 的数学能力，以及在有限计算资源下运行它们所面临的挑战。我们也感谢我们的竞争对手们在竞赛期间分享的各种见解——特别是 Abdur Rafae，他的 公开笔记本 为我们围绕模型训练和推理的许多决策提供了依据。

我们的解决方案由三个主要部分组成：

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

您可以在 Hugging Face Hub 上找到模型、演示和验证集的集合——我们将很快发布训练数据和代码！

我们使用了多种开源库来训练模型， notably TRL, PyTorch, vLLM, 和 DeepSpeed。在一个拥有 8 x H100 GPU 的节点上，我们的模型训练耗时 10 小时。

训练方案

链接：

• 我们的最终模型：AI-MO/NuminaMath-7B-TIR
• 我们在 Kaggle 上使用的 8 位量化模型：AI-MO/NuminaMath-7B-TIR-GPTQ

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

• 阶段 1：在一个大型、多样的自然语言数学问题和解决方案数据集上微调基础模型，其中每个解决方案都模板化为思维链 (CoT) 以促进学习。
• 阶段 2：在阶段 1 的模型基础上，在一个工具集成推理的合成数据集上进行微调，其中每个数学问题被分解为一系列理由、Python 程序及其输出。在这里，我们遵循微软的 ToRA 论文，并提示 GPT-4 以带有代码执行反馈的 ToRA 格式生成解决方案。在此数据上进行微调会产生一个推理代理，可以通过自然语言推理混合使用 Python REPL 来计算中间结果来解决数学问题（见下图）。

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

我们最初的提交仅使用了在阶段 1 上微调的 DeepSeek 7B 模型，但我们发现性能相当有限，使用 maj@32 在公共排行榜上的最佳得分为 8/50。是 Abdur Rafae 的公开笔记本促使我们考虑将代码执行集成到训练方案中，最初，我们专注于 Mix of Minimal Optimal Sets (MMOS) 数据集，因为笔记本标题中描述了这一点。我们发现使用 MMOS 提高了性能，但在公共排行榜上使用 maj@32 仍然 capped 在 16/50，这可能是因为 MMOS 仅由单轮解决方案组成（即模型只生成一个 Python 程序，这对于难题来说是不够的）。我们后来意识到 MMOS 是一个误称，Kaggle 笔记本实际上运行的是能够进行多步推理和代码执行的 DeepSeekMath 7B RL 模型。

就在这时，我们将精力集中在生成类似于 DeepSeekMath Instruct / RL 模型所使用的数据集上，这与 MuMath-Code 方案一起带来了显著的改进。

我们还尝试了在阶段 2 上仅微调整数解决方案，但发现当在所有解决方案类型上进行训练时，模型略微更稳健（可能是因为过滤后留下的样本数量有限）。

带有工具集成推理的自我一致性 (SC-TIR)

正如其他竞争对手所指出的，这次竞赛在模型提交和评估方面提出了几个挑战：

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

最初，我们使用 Abdur Rafae 的公开笔记本进行提交，但发现高方差是个问题。为了解决这个问题，我们采取了一种不同的方法，基于扩展带有工具集成推理的自我一致性解码 (SC-TIR)：

1. 对于每个问题，复制输入 N 次以定义馈送到 vLLM 的初始提示批次。这些有效地定义了用于多数投票的候选数量。
2. 采样 N 个完成项，直到生成完整的 Python 代码块。
3. 执行每个 Python 块并连接输出，包括出现的 traceback。
4. 重复 M 次以生成一组宽度为 N、深度为 M 的推理轨迹。如果轨迹未能产生 sensible 输出（例如，不完整的代码块），我们修剪该轨迹。
5. 后处理解决方案候选项，然后应用多数投票来选择最终答案。

对于我们获胜的提交，我们生成了 N=48 个候选项，深度为 M=4。增加任一参数并未提高性能，我们采取了一种保守的方法以保持在时间限制内。实际上，该算法介于带有 CoT 的自我一致性和思维树 (Tree of Thoughts) 之间（如下所示），但不包括回溯或向前看来规划解码的下一步。

我们发现 SC-TIR 提高了整体性能，并在我们的内部评估和公共排行榜上产生了显著较小的方差。

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

• 上传模型到 Kaggle Hub 非常慢，压缩模型使这一步快了一倍。
• 16 位模型仅加载权重就需要大约 32GB 显存。使用 2xT4，这将需要操作 KV 缓存才能快速运行，我们发现用模型精度换取速度是有益的。

我们使用 AutoGPTQ 以及训练数据集进行校准来量化模型。实际上，这导致准确率略有下降，但主要通过能够在推理期间生成许多候选项来弥补。

好的验证集就是你所需的一切

链接：
https://huggingface.co/datasets/AI-MO/aimo-validation-aime
https://huggingface.co/datasets/AI-MO/aimo-validation-amc
https://huggingface.co/datasets/AI-MO/aimo-validation-math-level-4
https://huggingface.co/datasets/AI-MO/aimo-validation-math-level-5

为了指导模型选择，我们使用了四个内部验证集来衡量模型在不同难度数学问题上的性能：

• AMC: 我们选取了 2022 年和 2023 年 AMC12 的所有问题，并保留了那些可以转换为整数输出的问题。这产生了 83 个问题的数据集。该验证集旨在模仿 Kaggle 上的私有测试集，我们发现我们的模型可以解决其中约 60-65% 的问题。为了测量方差，我们使用 5-10 个不同的种子运行每个评估，通常看到我们的 TIR 算法有大约 1-3% 的变化。
• AIME: 我们选取了 AIME 22、AIME 23 和 AIME 24 的所有问题，以衡量我们的模型在难题上的表现，以及衡量最常见的失败模式。如上所述，我们使用 5-10 个种子运行每个评估以测量变化。
• MATH 级别 4 & 5: 我们担心中等/困难验证集太小而无法给出可靠的信号，因此我们也从 MATH 测试集中过滤了具有整数解决方案的级别 4 和 5 样本。这给出了每个级别约 750 个问题用于评估，我们对每个模型使用单个种子进行评估。

通过使用这四个验证集，我们能够挑选不同训练阶段最有希望的模型，并缩小超参数的选择范围。我们发现，在这次竞赛中，将小型但具有代表性的验证集与大型验证集结合起来非常有用，因为每次提交都受到采样随机性的影响。

我们尝试过的其他想法

如上所述，我们尝试了一些方法，但最终为了支持 MuMath-Code 方案而放弃了：

• 训练纯 CoT 模型并在评估中使用多数投票
• 训练 MMOS 模型以单步使用 Python 解决问题

我们尝试的另一种技术是将 卡尼曼 - 特沃斯基优化 (KTO) 应用于从 SFT 模型采样的新完成项。这里的方法类似于 OrcaMath，即：

• 使用 SFT 模型对每个问题采样 4 个完成项，使用交错的理由和代码执行。我们使用阶段 2 的 SFT 数据集作为提示源。
• 提取答案并与 ground truth 进行比较。如果正确，将样本标记为正，否则为负。
• 在此数据集上对 SFT 模型应用 KTO。

我们发现这种形式的 on-policy KTO 产生的模型比 SFT 模型更好（在我们的内部评估上提高了几个百分点），并在公共排行榜上得分为 27/50。KTO 的一个 nice 功能是你可以跟踪训练期间的隐式奖励，这真的有助于调试运行——例如，这是我们来自 Weights and Biases 的成功训练日志之一，可以看到所选（即正确）的奖励在训练期间增加，而被拒绝的奖励被抑制。

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

我们还尝试将我们的 SFT 方案应用于更大的模型，如 InternLM-20B、CodeLama-33B 和 Mixtral-8x7B，但发现 (a) DeepSeek 7B 模型由于其在数学上的持续预训练而很难被击败，(b) 推理非常慢，并且我们遇到了一些无法追踪根本原因的神秘超时。

另一个失败的实验包括尝试使用强化学习（具体来说是 REINFORCE-leave-one-out (RLOO) 算法）配合代码执行反馈和正确/错误解决方案的离散奖励。我们将此应用于 DeepSeekMath 7B RL 模型，但没有看到性能的任何显著增益。鉴于像 RLOO 这样的在线方法受到文本生成的瓶颈限制且迭代缓慢，我们放弃了强化学习而 favor KTO。

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

• 使用静态 KV 缓存和 torch 编译。我们发现我们能够在 H100 上的原生 transformers 代码中将生成速度提高 2-3 倍，但在 Kaggle T4 上遇到了各种神秘错误，主要是由于 Hugging Face accelerate 中缺乏对 torch 编译模型分片的支持。
• 各种模型合并技术，如 DARE、TIES 和 WARP。在这里，我们使用 mergekit (https://github.com/arcee-ai/mergekit) 来合并 SFT 和 KTO 模型，或将 SFT 模型与公共 DeepSeekMath 模型合并。总的来说，我们发现这些合并导致我们的内部评估出现显著回归，或者我们没时间更深入地探索这一点。

结束语

我们非常享受参加这次竞赛并从中学到了很多。对于未来的进步奖，我们认为通过启用以下功能，可以进一步提高开源 LLM 的性能：

• 在现代 GPU 如 A100 或 H100 上进行评估。如前所述，最近 LLM 推理优化的许多进展依赖于更新的计算架构和数据类型，使用具有更多显存的 GPU 将使人们能够尝试更强大的模型。
• 将预训练模型截止日期放宽到竞赛开始日期。鉴于 LLM 发展的当前步伐，这将使竞争对手能够使用可用的最强大的开源 LLM，并提供与我们相比人类基线或专有模型有多远的更好衡量标准。