首届AI奥林匹克（ AI Math Olympiad (AIMO).），第一名由 Numina 和 Hugging Face 的佬们合作拿下，解决了private test中的29个问题！
排名
原文地址：
https://huggingface.co/blog/winning-aimo-progress-prize
以前，来自世界各地的高中生每年都会参加国际数学奥林匹克竞赛，解决代数、几何和数论等领域的六个具有挑战性的问题。
2023 年 11 月，AIMO 奖设立，目的是推动擅长数学推理的AI模型的开发。谁能创建能够赢得 IMO 金牌的 AI 模型，谁就将获得 500 万美元的大奖!
除了大奖之外，AIMO 还推出了一系列进步奖，以标志着实现这一最终目标的里程碑。一等奖以 Kaggle 竞赛形式举行，题目难度低于 IMO，但属于 IMO 预选水平。
下面这是一个示例问题，相比于 IMO 问题更容易解决，但对于 LLMs 来说仍然很棘手：
让 𝑘,𝑙 >0 为参数。抛物线 𝑦=𝑘𝑥^2−2𝑘𝑥+𝑙与线 𝑦=4相交于两个点 𝐴和𝐵，2个点相距为6，那么从𝐴、𝐵到原点的距离的平方和是多少？
参赛者每天可以提交两次解决方案，仅使用 2 月 23 日之前发布的开放权重模型。每次提交都分配一个 P100 GPU 或 2xT4 GPU，最多 9 小时来解决 50 个问题。
由3个部分
- 微调 DeepSeekMath-Base 7B ，使其充当“推理 agent”，可以通过自然语言推理和使用 Python REPL 计算中间结果的混合来解决数学问题。
- 一种用于工具集成推理 (TIR) 的新颖解码算法，具有代码执行反馈，可在推理过程中生成候选解决方案。
- 用来指导模型选择并避免过度拟合公共排行榜的各种内部验证集。
微调方法
微调方法主要基于 MuMath-Code 论文，分两个阶段训练模型：
MuMath-Code 论文中的两阶段训练方法
- 第 1 阶段：在包含自然语言数学问题和解决方案的大型、多样化数据集上微调基本模型，其中每个解决方案均使用思维链 (CoT) 进行模板化，以促进推理。
- 第 2 阶段：在工具集成推理的合成数据集上微调第 1 阶段的模型，其中每个数学问题都分解为一系列基本原理、Python 程序及其输出。在这里，遵循微软的 ToRA 论文，使用 GPT-4 生成带有代码执行反馈的 ToRA 格式的解决方案。对此数据进行微调会产生推理 agent，它可以通过自然语言推理和使用 Python REPL 计算中间结果的组合来解决数学问题（请参见下面的截图）。
图来自 ToRA 论文，介绍了用来训练模型的工具集成推理格式。
在两个阶段都进行了“全参微调”，所有模型权重在反向传播期间都得到了更新。使用 TRL 的 SFTTrainer 中的“packing”功能将多个样本拼接成 2048 个token的单个块中。所有模型都经过梯度检查点（gradient checkpointing）训练，并使用 DeepSpeed ZeRO-3 协，以确保权重、梯度和优化器状态能够适合可用的 VRAM。在每个阶段使用的主要超参数如下：
| Stage 1 | Stage 2 | |
|---|---|---|
| learning rate | 2.0 E-5 | 2.0 E-5 |
| total batch size | 32 | 32 |
| block size | 2048 | 1024 |
| num epochs | 3 | 4 |
| lr scheduler | cosine | cosine |
| warmup ratio | 0.1 | 0.1 |
最初提交的作品使用了 DeepSeek 7B 模型，该模型仅在第 1 阶段进行了微调，但发现性能相当有限，8/50 是在使用 maj@32 的公共排行榜上的最佳成绩。
受到公开notebook的启发，集中精力生成一个与 DeepSeekMath Instruct / RL 模型使用的数据集类似的数据集，这与 MuMath-Code 的流程一起带来了明显的改进。
数据集构建方法（Good data is all you need）
参考了DeepSeek Math和其他工作的方法，并对其进行了大幅扩展。产生了包含数十万问题-解决方案对的微调数据集，涵盖从高中数学到竞赛级别数学的主题。该数据集将在接下来的几周内完全开源。有关数据集构建的详细信息，请参阅即将开源的数据集技术报告。我们构建了两个数据集来微调我们的模型。
- Chain of Thought
该数据集由数十万个问题组成，每个问题都有以思维链的方式编写的解决方案。数据集的来源范围从中国高中数学练习到美国和国际数学奥林匹克竞赛问题。数据主要来自在线试卷 PDF 和数学讨论论坛。
处理步骤包括：
- 对原始 PDF 进行 OCR。
- 细分为问题-解决方案对。
- 翻译成英文。
- 重新调整以产生思想链推理格式。
- 最终答案格式。
- Tool-integrated reasoning
工具集成推理（TIR）在本次比赛中发挥着至关重要的作用。然而，收集和标注此类数据既昂贵又耗时。为了解决这个问题，从 Numina 数据集中选择了大约 60,000 个问题，重点关注那些具有数字输出的问题。
然后，利用 GPT-4 生成类似 TORA 的推理路径，执行代码并生成结果，直到解决方案完成。筛选出最终答案与参考不匹配的解决方案，并重复此过程三次，以确保准确性和一致性。这种迭代方法使我们能够高效地生成高质量的 TORA 数据。
作为参考，以下是我们的第 1 阶段模型 NuminaMath-7B-CoT 和最终第 2 阶段模型 NuminaMath-7B-TIR 在 MATH 基准上与其他开放和专有模型相比的性能：
| Model | MATH (%) |
|---|---|
| Chain of Thought Reasoning | |
| GPT-4 (2023) | 42.5 |
| GPT-4o | 76.6 |
| Claude 3.5 Sonnet | 71.1 |
| DeepSeekMath-7B-Instruct | 46.8 |
| DeepSeekMath-7B-RL | 51.7 |
| NuminaMath-7B-CoT | 56.3 |
| Tool-Integrated Reasoning | |
| DeepSeekMath-7B-Instruct | 57.4 |
| DeepSeekMath-7B-RL | 58.8 |
| NuminaMath-7B-TIR | 68.2 |
结合Self-Consistency和Tool-integrated reasoning 来避免随机性 - (SC-TIR)
如其他参赛者指出的那样，本次比赛在模型提交和评估方面提出了一些挑战：
- 评估 API 以随机顺序提供问题，因此提前停止等策略会产生较高的方差。因为一次运行一开始可能会遇到更多困难问题，从而为剩余部分留下更少的时间（反之亦然）
- LLM 推理中的大多数创新， 如 Flash Attention 2 或 torch.compile 等标准方法不适用于 T4 GPU。同样，不支持 bfloat16 ，这促使我们探索 AWQ 和 GPTQ 等训练后量化方法。
- 对于每个问题，将输入复制N次，作为一个batch。
- 对 N 个不同的补全进行采样，直到生成完整的 Python 代码块。
- 执行每个 Python 块并拼接输出，包括回溯（如果出现）。
- 重复 M 次以生成一批大小为 N、深度为 M 的输出结果，允许模型使用回溯自我纠正代码错误。如果样本无法产生合理的输出（例如，不完整的代码块），请修剪该结果。
- 对候选解决方案进行后处理，然后应用多数投票来选择最终答案
对于获胜的提交，生成了 N=48 个候选，深度为 M=4。增加任一参数都不会提高性能，因此采取保守的方法来保持在时间限制内。实际上，该算法通过工具集成推理增强了 CoT（如下所示）的自我一致性。
SC-TIR 算法产生了更稳健的结果，并且在内部评估和公共排行榜上的方差显著减少。
值得一提的一个技术细节是，以 8 位精度量化模型很有帮助。这是出于三个原因：
- 将模型上传到 Kaggle Hub 的速度非常慢，压缩模型使这一步速度提高了一倍。
- T4 GPU 不支持 bfloat16，转换为 float16 会导致模型性能下降。无法转换为 float32，因为它超出了可用的 GPU 内存。
- 此外，16 位模型仅用于加载权重就消耗约 32GB VRAM。对于 2xT4，这需要操纵 KV 缓存才能快速运行，权衡模型精度和速度是有益的。
使用 AutoGPTQ 以及用于校准的训练数据集来量化我们的模型。在实践中，这会导致准确性小幅下降，但提供了最佳折衷方案，以适应 Kaggle 平台上评估所施加的限制。
避免过度拟合的诅咒
过度拟合公共排行榜是 Kaggle 竞赛中的常见风险，当测试集只有 50 个问题时更是如此。此外，规则允许每天最多提交两次，这使得强大的内部验证数据集对于我们的开发节奏至关重要。根据AIMO团队的规定，测试问题的难度为中等，介于AMC12和AIME级别之间，具有整数输出。
为了选择模型，使用了四个内部验证集来衡量模型在不同难度的数学问题上的性能。从 AMC12（2022、2023）和 AIME（2022、2023、2024）中选择问题来创建两个内部验证数据集：
- AMC（83 个问题）：从 AMC12 22、AMC12 23 中挑选了所有问题，并保留了那些可以转换为整数输出的问题。这产生了包含 83 个问题的数据集。该验证集代表 Kaggle 上的private test集，因为从竞赛描述中知道问题属于这个级别或更难。测试发现我们的模型可以解决大约 60-65% 的问题。为了测量方差，使用 5-10 个不同的种子进行每次评估，并且使用 SC-TIR 算法通常会看到大约 1-3% 的变化。
- AIME（90 个问题）：从 AIME 22、AIME 23 和 AIME 24 中挑选了所有问题来衡量模型在困难问题上的表现，以及衡量最常见的故障模式。如上所述，我们使用 5-10 个种子进行每次评估以测量变化。
由于 AMC/AIME 验证集规模较小，这些数据集上的模型性能容易受到噪声的影响，与公共排行榜类似。为了更好地评估模型的性能，还使用 MATH 测试集的子集（包含 5,000 个问题）对其进行了评估。仅保留整数输出的问题，以简化多数投票并模拟竞争评估。这导致了两个额外的验证集：
- 数学 4 级（754 个问题）
- 数学 5 级（721 个问题）
通过使用这四个验证集，能够在不同的训练阶段选择最好的模型，并缩小超参数的选择范围。在这个特定的竞赛中，将小型但具有代表性的验证集与较大的验证集相结合是有用的，其中每个提交都受到采样的一些随机性的影响。
如上所述，我们尝试了一些方法，但最终被放弃，转而采用 MuMath 代码配方：
- 训练纯 CoT 模型并使用多数投票进行评估
- 训练 MMOS 模型以通过 Python 一步解决问题
尝试的另一种技术是将 Kahneman-Tversky Optimization (KTO) 应用于从 SFT 模型采样的结果。这里的方法类似于 OrcaMath，即：
- 使用第 2 阶段的 SFT 数据集作为提示数据集，生成4个补全结果。
- 提取答案并将其与正确答案进行比较。如果正确，则将样本标记为positive，否则标记为negative。
- 将 KTO 应用于此数据集上的 SFT 模型。
这种形式的同策略 KTO 生成的模型比 SFT 模型稍好（内部评估好几个百分点），并且在公共排行榜上得分为 27/50。
KTO 的一个很好的功能是，您可以在训练期间跟踪隐式奖励，这确实有助于调试运行 - 例如，下图是很成功的训练日志之一，其中人们可以看到所选（即正确的解决方案）奖励随着训练而增加，而被拒绝的则被抑制。
但是，我们没有时间将此方法应用于最终的 SFT 模型，因此我们可能还能解决 1-2 个以上的问题！
我们还尝试将我们的sft方法应用于 InternLM-20B、CodeLama-33B 和 Mixtral-8x7B 等更大的模型，但发现 (a) DeepSeek 7B 模型由于持续的数学预训练而很难被击败，并且 (b) 更大模型在 2xT4 GPU 上推理速度非常慢，并且经历了许多莫名其妙，无法追踪原因的超时。
另一个失败的实验包括尝试使用强化学习（特别是近端策略优化算法和 REINFORCE 留一法 (RLOO) 算法）和代码执行反馈，并为编写代码和获得正确/不正确的解决方案形成奖励。将其应用于 DeepSeekMath 7B RL 模型。虽然我们看到了一些有希望的奖励曲线，但没有看到结果有任何显著的提升。鉴于像 RLOO 这样的在线方法受到文本生成的瓶颈并且迭代缓慢，我们放弃了强化学习，转而尝试 KTO。
在推理方面，也进行了实验：
- 使用静态 KV 缓存和 torch 编译。发现能够在 H100 上将原生 Transformer 代码的生成速度加快 2-3 倍，但在 Kaggle T4 上遇到各种神秘错误，主要是由于加速中的 torch 编译缺乏对模型分片的支持。
多种模型合并技术，例如 DARE、TIES 和 WARP。这里我们使用 mergekit 来合并 SFT 和 KTO 模型，或者将 SFT 模型与公共 DeepSeekMath 模型合并。总的来说，我们发现这些合并导致我们的内部评估出现明显的退化，并且没有时间更深入地探索这一点。
最后，感谢AIMO团队发起了如此激动人心、鼓舞人心的比赛！
PS：给公众号添加【星标⭐️】不迷路！您的点赞、在看、关注 是我坚持的最大动力！
欢迎多多关注公众号「NLP前沿」，加入交流群，交个朋友吧，一起学习，一起进步！
最新文章推荐阅读