AMD Athena-PRM 用高质量标签数据赋能多模态复杂推理

作者：Zhenhua Liu, Shuai Wang, Xuanwu Yin, Dong Li, Emad Barsoum

本文介绍Athena-PRM，一个多模态过程奖励模型（PRM），用于为复杂推理问题的每一步推理过程打分。为了高效生成高质量的过程标注数据，我们利用弱 completer 和强 completer 预测结果的一致性，作为筛选可靠过程标签的标准。同时，我们提出了两种有效策略来提升 PRM 的性能：基于 ORM 的初始化以及负样本的上采样。

在多个基准与应用场景中，Athena-PRM 都取得了稳定领先的效果。特别是，当使用 Qwen2.5-VL-7B 作为策略模型时，在Test-Time Scaling的场景下，Athena-PRM 在 WeMath 上提升了 10.2 个百分点，在MathVista 上提升了7.1 个百分点。此外，Athena-PRM 在 VisualProcessBench 上达成了最新的 SoTA（state-of-the-art）结果，相比此前的 SoTA 提升 3.9 F1，说明其在精确判别中间推理步骤正确性方面具有强大的能力。

为什么要使用“弱–强 completer 一致性”？

面向 Test Time Scaling（TTS）的奖励模型主要有两类：结果奖励模型（ORM）和 过程奖励模型（PRM）。

• ORM：对给定问题和完整解答进行整体打分。
• PRM：对推理过程中的每一个中间步骤进行打分，提供更细粒度的反馈。

PRM 通常在性能与跨分布泛化方面更有优势。但高质量的过程标注数据非常难获得：

• PRM800K 通过人工标注收集了 80 万个步骤标签，这不仅耗时，而且需要具备较强数学与多步推理能力的标注者。
• Math-Shepherd 提出了一种基于蒙特卡洛估计的自动标注方法。其核心思想是：用completer（语言模型）从某一中间步骤继续采样大量推理轨迹，用“最终到达正确答案的概率”来衡量该中间步骤的质量。

然而，这种蒙特卡洛式的估计存在两个问题：

1. 计算成本高：需要大量采样才能稳定估计每一步的“到达正确答案的概率”。

2. 标签有噪声：即使采样很多次，中间步骤的估计标签仍然不可避免地会引入噪声。

图1 在相同问题和解题步骤下，不同 completer 的表现示意图

我们希望同时解决以上两个问题：降低计算成本，并减轻步骤标签的噪声。

• 我们首先观察到：MC 方法估计出来的步骤标签质量，强烈依赖 completer 的推理能力：强 completer 即便在中间步骤出错时，也可能“自救”到正确答案（如图 1 所示）；
• 弱 completer 即便当前步骤正确，也可能因为后续能力不足而无法得到正确答案。

因此，基于 completer 的 MC 估计天然带有偏差，估计结果更像是在评估“这一步对特定 completer 的帮助”，而不是这一步本身的逻辑正确性。但从定义上讲，中间步骤是否正确，不应该依赖我们选用哪个 completer。基于这一点，我们提出：

• 同时使用 弱 completer 和 强 completer 进行步骤标签估计；
• 只保留两者估计标签一致的步骤，作为训练数据。

这样可以在不额外放大计算开销的前提下，有效过滤掉受 completer 偏差影响较重的步骤标签，大幅提升步骤标签的质量。

ORM 初始化与负标签上采样

我们进一步提出两种策略，来增强 PRM 的性能：从 ORM 初始化和负样本上采样。

1. 基于 ORM 的初始化

先用大规模样本级标注数据训练 ORM，再用高质量过程标注数据从 ORM 初始化 PRM，可以显著提升 PRM 性能。既有工作表明：ORM 其实在一定程度上也具备判断中间步骤正确性的能力。我们的实验发现：在 PRM 训练前，利用少量高质量样本对 ORM 做“初始化”，非常有助于 PRM 最终效果。

从监督角度看，大规模样本级标注（只有最终结果标签）是一种较弱监督；outcome 级别的监督可以看作是过程监督的一种简化形式；用大规模样本级监督训练 ORM，相当于对 reward model 进行一次“预训练”；在此基础上，用高质量的过程级标注数据训练 PRM，则相当于一种“基于 ORM 的微调”。

2. 负标签上采样

我们还发现，在多数过程标注数据集中（如 PRM800K 和我们合成的数据），标签非常不平衡：正确步骤远多于错误步骤，即“正样本”显著多于“负样本”。

为缓解这一问题，我们对负标签（错误步骤）进行上采样：简单来说，就是在训练批次中提高负样本的出现频率；实验结果表明：负样本上采样能明显提升模型性能，而增加的计算开销非常有限。

应用场景

在训练完 reward model 之后，我们在三个典型场景下评估 ORM 和 PRM 的表现：

1.Test-Time Scaling中的验证

2.推理步骤的直接判定

3.Reward-ranked fine-tuning 中的响应排序

01

Test-time scaling 场景下的验证

我们采用 Best-of-N 的评估范式：

• 对于测试集中的一个问题 (x)，从策略模型 (\pi) 中采样 N 个解答。
• 使用 reward model 对每个解答评分。
• 选取得分最高的解答作为最终输出。

02

推理步骤的直接判定

除了在Best-of-N 场景下用作验证，Athena-PRM 还可以直接用于定位数学推理过程中的错误步骤：

• 给定一个包含 K 步的解题过程作为输入，
• Athena-PRM 在一次前向计算中，给出每一步的正确性得分，
• 用于快速找出有问题的推理环节。

03

Reward-ranked fine-tuning 中的响应排序

我们进一步探索使用 PRM 来为 reward-ranked fine-tuning 生成训练数据：

1.使用当前策略，对每个输入问题生成 (M = 8) 个解答。

2.过滤掉答案错误的解答，并进行去重，去除高度相似的响应，提升剩余解答的多样性。

3.仅保留如下问题：

o 8 个解答中有 2～6 个是正确答案；

o 排除“几乎全错”或“几乎全对”的问题，因为对当前策略 (\pi) 来说，这类问题要么过难，要么过容易。

4.在筛选后的问题上，使用 Athena-PRM 对所有解答打分，为每个问题选出 reward 最高的解答作为“优选响应”标签。

5.使用上述合成数据，对策略 (\pi) 进行 fine-tuning。

通过这种方式，我们在不依赖昂贵人工标注的情况下，构造了高质量的排序监督信号。

结果：在 AMD 硬件上的准确率表现

表 1 展示了 Athena-ORM 与 Athena-PRM 在多种策略模型与多模态推理基准上的表现。总体来看：

• Athena-ORM 和 Athena-PRM 均能普遍提升多模态大模型（MLLM）的推理能力；
• 在以Qwen2.5-VL-7B 作为策略模型的 WeMath 数据集上，Athena-PRM 相比 zero-shot 基线提升 10.2 个百分点；
• ORM 在多数基准和策略模型上也有收益，但部分收益较小，例如在 Qwen2.5-VL-72B + MathVerse 上仅提升 0.4 个百分点；

相比ORM，PRM 在几乎所有场景下的提升幅度都更大，体现出细粒度 reward 在 test-time scaling 中的优势。

延伸阅读

• Athena官方博客原文（英文）：https://rocm.blogs.amd.com/artificial-intelligence/amd-elvm/README.html

---

• ROCm AI Developer Hub：https://www.amd.com/en/developer/resources/rocm-hub/dev-ai.html

---