第二名解决方案 (Team imagination-research)

感谢 Aimo 和 Kaggle 举办这场精彩的比赛。我们也非常感谢比赛期间所有的竞争对手以及所有分享伟大想法的开源项目。现在是时候与社区分享我们的想法和经验了！

我们的解决方案获得了第 2 名。它在公共排行榜上获得 34/50（排名第 1），在私有排行榜上获得 31/50（排名第 2）。

解决方案总结

本次比赛需要同时优化效率和推理性能。我们的最终解决方案主要由三部分组成：

• 第一部分：面向推理的训练 -- 提高模型的推理能力：阶段 1 - SFT 和阶段 2 - 使用精选数据进行 DPO。
• 第二部分：效率优化 -- 提高推理效率：选择合适的推理引擎、权重量化和 KV Cache 量化。
• 第三部分：推理时策略 -- 优化效率与推理性能的权衡：提示词设计、自一致性聚合、样本级/问题级早停，以及一些启发式超参数调整。

对于本地验证，我们使用了 AIME 2025 测试集（30 个问题）以及参考集（10 个问题），评估了平均样本准确率和聚合准确率（通过自一致性），以便对我们的试验解决方案获得初步判断。

第一部分：面向推理的训练

我们的训练脚本基于 Light-R1 项目，我们非常感谢他们的工作。

阶段 1：SFT

考虑到 DeepSeek-R1-Distill-Qwen-14B 在数学、编码和推理方面的出色表现，我们选择它作为基础模型。
我们将来自 Light-R1 的阶段 2 数据和来自 Limo 的训练数据结合在一起（去除重复项），这两者都是由 deepseek-r1 生成的高难度数学问题的推理轨迹。
我们在单台 8×A800 机器上对基础模型进行了 8 个 epoch 的微调，耗时 11 小时：

准确率提高了，但输出长度也显著增加。

阶段 2：DPO

我们使用 DPO 来减少模型的输出长度。
我们选择 OpenR1-Math-220k 的默认子集来构建我们的数据集。
具体来说，我们尝试使用以下四个标准来构建 DPO 对（y_w, y_l 分别表示被选中的响应和被拒绝的响应）：

• 正确性：y_w 必须正确，y_l 可能正确也可能不正确
• 长度 ratio：len(y_w) < ratio_threshold * len(y_l)
• 最小长度：len(y_w) > min_threshold
• 相似度：sim(y_w,y_l) < sim_threshold
  • 使用 sentence transformer 模型计算 embeddings

应用前三个标准，我们构建了数据集 dpo-1，用于训练我们提交的模型。
应用全部四个标准，我们构建了数据集 dpo-2，用于训练另一个模型，但其性能与我们提交的模型相似。
我们使用 360-LLaMA-Factory，因为他们添加了序列并行 (SP) 技术以支持在有限内存下进行更长上下文的训练。
我们在单台 8×A800 机器上在 dpo-1 数据集（2k 对）上训练了 4 个 epoch，耗时 40 小时。
通过上述过程，我们得到了最终提交的两个模型：
deepseek-14b-sft-dpo2 和 deepseek-14b-sft-dpo4
我们所有的训练数据都可以在这里找到：
训练数据

注意：在上表中，我们对每个问题采样 32 次（直接推理 16 次和代码求解 16 次）。Pass@1 是在 AIME 2025 数据集上计算的。更多关于我们推理方法的细节将在后面讨论。

第二部分：效率优化

推理引擎

我们选择 lmdeploy 作为 LLM 推理框架。与 vllm 相比，带有 TurboMind 引擎的 lmdeploy 框架可以提供更高的吞吐量和更短的模型初始化时间。
我们使用的 lmdeploy 版本是 0.7.0.post1
第一张图片来自 这里

量化

我们应用了 4-bit AWQ 权重量化和 8-bit KV Cache 量化（将配置 main_model.inference_cfg.quant_policy 设置为 8 以使用 lmdeploy 实现的 8-bit KV Cache 量化）。
我们用于提交的两个模型可以在这里找到：模型 1, 模型 2
一些效率结果：

• 在线测试 (4xL4, batch size=15): 与 W4KV16 相比，W4KV8 将每个输出 token 的时间减少了约 20%，与 FP16 相比减少了 55%。
  
• 本地测试 (2xA800, batch size=32): 与 FP16 相比，W4KV8 和 W4KV16 分别将整体延迟降低了 40% 和 20%-25%。
  一些推理性能结果：
• 本地测试：与 FP16 相比，平均样本准确率（非聚合准确率）下降了 5%~10%；W4KV8 不差于 W4KV16。W4KV4 更差。

第三部分：推理时策略

整体推理工作流

推理工作流如下图所示：提供一个问题作为输入。我们首先准备两种类型的提示词，包括 CoT 提示词和代码提示词（"提示词准备任务"）。然后，我们让 LLM 开始使用 lmdeploy 批量生成多个样本（"LLM 生成任务"）。同时，我们 continuously 尝试从每个样本的流式输出中提取答案，聚合多个样本的答案，并判断是否早停某些生成：

1. 我们在从 stream_infer(...) 调用获得的迭代器的每 N 个 yield 上进行样本级检查，并判断是否早停相应样本的生成。此处使用 Python 代码执行器和答案提取组件。
2. 我们在每个样本结束时进行问题级检查，并判断是否早停当前问题所有剩余样本的生成。此处使用答案聚合组件。
   最后，我们返回聚合的答案。
   注意，对于每个问题，我们会根据剩余时间调整与速度相关的超参数（样本数量、样本级最大时间、问题级早停标准），以便在剩余时间有限时，更平衡地在剩余问题之间分配时间配额。
   

提示词准备

方法：

• 最初，我们对一个问题使用 15 个样本，并通过自一致性聚合它们的答案。
• 注意，我们不一定聚合所有样本的答案，详见 下文讨论。
• 注意，当剩余时间有限时，我们会减少样本数量，详见 下文讨论。
• 我们使用常用的基于代码的推理：(1) 提示模型提供 Python 代码来解决 проблемы；(2) 从输出中提取 Python 代码，创建子进程执行代码；(3) 从执行结果中提取答案。
• 我们使用两种类型的提示词 - CoT 提示词和代码提示词。在 15 个样本中，7 个样本使用 CoT 提示词，8 个样本使用代码提示词：

# CoT 提示词
- system: "你是一个有用的数学助手。请一步步推理并将答案放在 \\boxed{} 中。"
  user_suffix: "\n你擅长推理。\n你必须将最终答案放在 \\boxed{} 中。\n如果最终答案大于 1000，则取模 1000。\n仔细 thoroughly 思考，避免重复。"

# 代码提示词
- system: "你是一个有用的数学助手。请提供 Python 代码来解决数学问题，并将最终答案放在 \\boxed{} 中。"
  user_suffix: "\n你擅长编码\n你必须提供 Python 代码，避免冗余分析。\n如果最终答案大于 1000，则取模 1000。\n答案必须是整数。\n每个问题只有一个答案。\n导入必要的库。"

一些实验：

• 系统提示词选择：我们发现对于推理模型，多样化系统提示词没有帮助。
• 提示词列表选择：在本地测试中，我们发现仅使用我们的代码提示词比使用一半 CoT 和一半代码提示词结果更一致（跨种子和模型）且有小幅提升。然而，当我们提交此方案（仅一次）时，它对公共提交分数没有帮助，因此我们没有进一步测试这一经验选择。
• 样本数量：在本地测试中，我们发现使用 32 个样本比使用 16 个样本取得更好的结果。然而，由于提交平台上的计算能力有限且提交配额有限，我们没有在提交平台上彻底实验更多样本以找到最佳点——我们只是选择了 15 个样本。
• 代码提示词导致代码输出、代码错误、错误答案的频率（32 个样本，16 个 CoT 提示词，16 个代码提示词）：
• 代码运行正确但我们无法从其输出中解析整数的情况很少见，可以忽略。
• 微调前，模型更倾向于输出代码：平均而言，在使用代码提示词的 16 个案例中，只有 1.9 或 3.3 个案例模型没有输出代码。
• 微调后，仅使用数学数据，模型变得不太倾向于输出代码：平均而言，在 16 个案例中约有 11 个案例，代码提示词没有导致模型输出代码。当新模型确实输出代码时，其条件准确率略高于微调前模型（45% 和 55% 对 42%）。

样本级答案提取与早停

动机：通常，推理模型在早期获得答案后会自我怀疑很多，即使它最终通常给出相同的答案。而且在大多数情况下，在 <think></think> 之间给出答案后，模型会再次重写解决方案（至少两次）。我们能否减少 token 的浪费？
方法：虽然我们实验了来自 "Fu et al., Efficiently Serving LLM Reasoning Programs with Certaindex, arXiv 2412" 的主动探测方法，但我们最终采用了一种更简单的样本级早停技术来简化我们的推理工作流。具体来说，一旦我们检测到第一个成功执行的代码或第一个在 "\boxed{…}" 中的答案，我们就停止该样本的生成过程。
一些实验：一个自然的问题是，这是否会损害后来将最初错误的答案修订为正确答案的潜力。在我们的本地测试中，我们使用 scripts/analyze_early_stop.py 验证了这种情况相对罕见，如下图所示。

问题级答案聚合与早停

我们使用常用的自一致性方法进行答案聚合。我们使用如下问题级早停策略。
动机：不同问题的难度各不相同，因此我们 aim 避免在简单问题上花费太多时间。如下图所示，单个问题的样本之间输出长度差异很大。这表明对于某些问题，我们可能很早就获得了几个正确答案，但仍需要等待最长的样本完成——导致显著的时间浪费（例如 q1, q11, q16–q20 等）。

方法：如果通过检查现有答案达到了足够的确定性，我们可以提前停止问题的生成。具体来说，当大多数输出一致时，我们在问题级终止生成，例如，如果 7 个答案中有 5 个一致。详见 imagination_aimo2/local_eval_kaggle.py 中的 early_stop_strategy.consistency_rules 配置。

速度超参数调整

动机：解决不同难度级别问题所需的时间差异很大，我们设计了一个 adjust_speed 模块来动态调整一些超参数。
方法：随着推理的进行，我们的 adjust_speed 模块计算剩余时间和剩余问题数量，并相应地动态调整模型的采样数量和早停策略。
例如，默认速度是 3(normal)，如果系统检测到每个问题的平均剩余时间少于 5 分钟，它会自动将速度调整为 1(fastest)。这意味着样本数量减少到 10 个，每个问题的最大推理时间也减少了。请参阅我们的代码以获取详细实现。

其他我们尝试过但未纳入最终方案的方法

• GRPO：我们尝试了四次训练运行，但没有观察到准确率的显著提升。
• RAG：我们使用这个 数据集 作为文档库。在解决问题时，我们从库中获取 1~2 个最相似的问题，使用它们及其相应的解决方案作为 few-shot 演示。但我们在 AIME2025 和在线提交上都没有获得提升。
• 基于模型的聚合：我们将 5 个最频繁的答案以及每个答案的一个代表性理由提供给 LLM，要求 LLM 推理出正确答案。但这种聚合方法无法稳定地实现提升。
• 量化相关试验：我们尝试通过在我们的 2k 大小 DPO 数据集上使用 DPO 微调 AWQ 量化模型的 scales 来进行量化感知训练 (QAT)，但这没有带来提升。我们还尝试应用 W_k 和 W_q 的重参数化以获得更通道平衡的 key 以更好地进行 KV cache 量化。这在我们的本地测试中带来了一致但较小的推理性能提升，特别是在 KV4 设置中。然而，由于我们没有足够的提交配额来为 KV4 设置调整其他超参数，我们的最终解决方案没有使用 KV4。
  我们决定简要分享这些初步试验，但请注意，它们不够彻底，无法提供可靠的负面信号。我们对上述试验的定性检查实际上表明有进一步探索的潜力，但我们 simply 缺乏能力进一步探索它们。