DreamPRM: Domain-Reweighted Process Reward Model for Multimodal Reasoning¶
会议: NeurIPS 2025
arXiv: 2505.20241
代码: GitHub (有)
领域: LLM推理
关键词: process reward model, multimodal reasoning, domain reweighting, bi-level optimization, test-time scaling
一句话总结¶
提出 DreamPRM，通过双层优化自动学习多模态推理数据集的域权重，解决 PRM 训练中的数据质量不均衡问题，在 MathVista 排行榜上以 o4-mini 模型达到 85.2% 的 top-1 准确率。
研究背景与动机¶
- 领域现状: Process Reward Model (PRM) 通过对推理链中间步骤的细粒度评估来引导 LLM 推理，已在文本领域取得成功。将 PRM 扩展到多模态 LLM (MLLM) 是自然的下一步。
- 现有痛点: 多模态推理覆盖更广的任务范围（科学、几何、图表、常识等），导致训练/测试分布偏移更严重，泛化更困难。需要大规模多样化数据集来确保覆盖，但现有多模态推理数据集存在严重的质量不均衡——许多数据集包含无用模态或过于简单的问题。
- 核心矛盾: 需要覆盖多个域的数据来训练泛化性强的 PRM，但简单地混合所有数据集会因低质量数据的噪声而降低 PRM 性能。
- 本文要解决什么？ 自动化地为不同质量的多模态推理数据集分配合适的权重，训练出泛化能力更强的多模态 PRM。
- 切入角度: 借鉴预训练中的域重加权技术（如 DoReMi），将其引入 PRM 训练，并设计适配推理场景的聚合函数损失。
- 核心idea一句话: 用双层优化自动学习数据集权重，让 PRM 聚焦于高质量推理样本，忽略噪声数据。
方法详解¶
整体框架¶
DreamPRM 采用双层优化 (Bi-Level Optimization, BLO) 框架： - 下层优化: 在多个训练域上用加权损失训练 PRM 参数 ϕ - 上层优化: 在独立的元学习数据集上评估 PRM，通过聚合函数损失更新域权重 α
关键设计¶
- 域加权 PRM 训练 (下层): PRM 训练损失为各域损失的加权和 \(\mathcal{L}_{tr} = \sum_{k=1}^K \alpha_k \mathcal{L}_{tr}(\mathcal{D}_k, \phi)\)。每个域的损失使用 Monte Carlo 估计的过程监督信号（通过多次补全并比较答案正确率获得）训练 PRM 预测每步的正确性概率。
- 聚合函数损失 (上层): 设计了与推理时行为一致的元损失——不是直接评估单步预测，而是先将 PRM 对每步的预测通过聚合函数 \(\mathcal{A}(p) = \sum_i \log \frac{p_i}{1-p_i}\) 转换为整条推理链的分数，再与正确性标签计算 MSE。这弥补了训练与推理之间的差异。
- 多阶段推理提示: 借鉴 LLaVA-CoT 的结构化思维，提示 MLLM 按五步推理（复述问题→从图像收集证据→识别背景知识→基于证据推理→总结并得出结论）。
- 数据组织: 15 个训练域覆盖科学、图表、几何、常识四大类；MMMU 作为元学习数据集。
损失函数 / 训练策略¶
- 下层: MSE 损失，PRM 预测值 vs Monte Carlo 估计的过程监督信号
- 上层: MSE 损失，聚合函数输出（经 sigmoid）vs 答案正确性标签
- 下层每 5 步更新一次上层（unroll steps = 5）
- AdamW 优化器，下层 lr=5e-7，上层 lr=0.01
- 总计训练 10000 迭代，单张 A100 约 10 小时
- PRM 基座: Qwen2-VL-2B-Instruct；推理 MLLM: InternVL-2.5-8B-MPO
实验关键数据¶
主实验¶
InternVL-2.5-8B-MPO 作为基座模型，在 5 个多模态推理 benchmark 上的准确率(%)：
| 方法 | WeMath | MathVista | MathVision | MMVet | MMStar |
|---|---|---|---|---|---|
| Base (zero-shot) | 51.7 | 65.4 | 20.4 | 55.9 | 58.9 |
| Self-consistency | 56.4 | 67.1 | 20.7 | 57.4 | 59.6 |
| ORM | 56.9 | 65.3 | 20.5 | 55.9 | 60.1 |
| Vanilla PRM | 54.2 | 67.2 | 20.6 | 58.9 | 60.8 |
| CaR-PRM | 54.7 | 67.5 | 21.0 | 60.6 | 61.1 |
| s1-PRM | 57.1 | 65.8 | 20.2 | 60.1 | 60.4 |
| DreamPRM | 57.4 | 68.9 | 22.1 | 61.4 | 62.3 |
MathVista 排行榜: o4-mini + DreamPRM 达到 85.2% (top-1)，超越 VL-Rethinker、Kimi-k1.6、OpenAI o1 等。
消融实验¶
| 消融配置 | 影响 |
|---|---|
| 移除双层优化 (BLO) | MathVista -3.5%, MMStar -3.4% |
| 移除聚合函数损失 (AFL) | WeMath -1.1%, 一致性下降 1-2% |
| 移除结构化思维 (ST) | MathVision -1.8% |
关键发现¶
- 域权重收敛: M3CoT 和 FigureQA 获得最高权重(~1.5)，AI2D 和 IconQA 获得最低权重(<0.8)——高权重数据集需要更深推理，低权重数据集过于简单
- 可扩展性: CoT 候选数从 2→4→8 时性能单调递增，说明 DreamPRM 能可靠地从更大候选池中筛选
- 跨模型泛化: 用 InternVL 训练的 PRM 在 GPT-4.1-mini 和 o4-mini 上也能持续提升性能
亮点与洞察¶
- 首个多模态 PRM 域加权框架: 将预训练中的域加权思想创新性地引入 PRM 训练
- 聚合函数损失设计精妙: 让上层优化目标与推理时的 PRM 使用方式完全一致，弥合训练-推理 gap
- 实际效果强: MathVista 排行榜 top-1，且方法计算代价低（单卡 10 小时）
- 跨模型泛化: 小模型训练的 PRM 能提升大模型的推理效果，实用价值高
- 学到的域权重可解释: 权重分布与数据集难度/质量直觉一致
局限性 / 可改进方向¶
- PRM 基座较小（Qwen2-VL-2B），更大的 PRM 是否能进一步提升有待验证
- 元学习数据集固定为 MMMU，其选择对结果的敏感性未充分分析
- 域权重是数据集级别的，更细粒度的样本级重加权可能进一步提升效果
- 仅评估了 best-of-N 推理范式，未探索树搜索（MCTS）等更复杂的推理策略
- 推理时需要生成多条 CoT 候选（8 条），计算开销相比 zero-shot 增加约 8 倍
相关工作与启发¶
- DoReMi 和 DOGE 在预训练域加权方面的工作为本文提供了直接灵感
- Math-Shepherd 和 OmegaPRM 的 Monte Carlo 过程监督方法是 PRM 训练的基础
- 与 CaR-PRM、s1-PRM 等启发式数据选择方法相比，自动化的域加权策略更优
- 对 test-time scaling 研究有参考价值：PRM 的质量是 test-time scaling 有效性的关键瓶颈
评分¶
- 新颖性: ⭐⭐⭐⭐ 域加权+双层优化的组合虽非全新，但在多模态 PRM 训练中是首次应用，聚合函数损失的设计有亮点
- 实验充分度: ⭐⭐⭐⭐⭐ 5 个 benchmark、多种 baseline 对比、消融实验、域权重分析、跨模型泛化、排行榜验证，非常全面
- 写作质量: ⭐⭐⭐⭐ 结构清晰，公式推导完整，但部分记号较密集
- 价值: ⭐⭐⭐⭐ 实际效果强，MathVista top-1，方法可直接应用于其他 PRM 训练场景