分布过程奖励模型：通过条件最优传输校准未来奖励的预测

# 分布过程奖励模型：通过条件最优传输校准未来奖励预测

来源：https://arxiv.org/html/2605.06785
Rachel Ma MIT CSAIL rachelm8@mit\.edu &Dylan Hadfield\-Menell MIT CSAIL &Kristjan Greenewald IBM Research

###### 摘要

推理时扩展方法依赖于过程奖励模型（PRM），但 PRM 往往校准不佳且高估成功概率。据我们所知，我们首次提出利用条件最优传输（Conditional Optimal Transport, CondOT）来校准 PRM，通过修改 CondOT 映射学习方法Bunne等人(2022 (https://arxiv.org/html/2605.06785#bib.bib2))，以估计在给定 PRM 隐藏状态下、由 PRM 预测的成功概率上的单调条件分位数函数。这种方法产生了结构有效的分位数估计，并支持在任意置信水平下高效提取置信区间，我们将其集成到Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))提出的实例自适应扩展（IAS）框架中。我们在涵盖中等难度问题（MATH-500）和更难分布外问题（AIME）的数学推理基准上进行了评估。对于具有可靠排序信号的 PRM，我们的方法在校准性能上显著优于未校准的 PRM 和分位数回归。在下游 Best-of-N IAS 性能方面，我们的方法通常优于未校准的 PRM。这些结果确立了条件最优传输作为 PRM 校准的另一种合理且实用的方法，提供了结构保证和灵活的不确定性估计。

## 1 引言

扩展推理时计算已成为提升大型语言模型（LLM）在推理任务上性能的强力范式Snell等人(2024 (https://arxiv.org/html/2605.06785#bib.bib8)); Brown等人(2024 (https://arxiv.org/html/2605.06785#bib.bib16))。推理时扩展方法不仅仅依赖固定模型，而是生成多个候选推理轨迹，并使用评分模型从中选择。过程奖励模型（PRM）对中间推理步骤相对于任务进行评分Cobbe等人(2021 (https://arxiv.org/html/2605.06785#bib.bib5)); Uesato等人(2022 (https://arxiv.org/html/2605.06785#bib.bib6)); Lightman等人(2023 (https://arxiv.org/html/2605.06785#bib.bib7))，可提供逐步信号以指导搜索、选择和预算分配。这些决策的质量直接取决于 PRM 分数反映真实成功概率的程度。

![Figure 1](https://arxiv.org/html/2605.06785#fig.1)
图 1：数学推理轨迹的估计成功概率：未校准的基础 PRM 通常高估，分位数回归提供的灵活性有限且可能出现交叉违规。我们的条件 OT 方法保证了单调的分位数函数，并支持在任意置信水平下进行灵活的不确定性估计。

然而，在实际应用中，最先进的 PRM 往往校准不佳且过于乐观Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))。这在推理时扩展中尤为有害，因为诸如 Best-of-N Brown等人(2024 (https://arxiv.org/html/2605.06785#bib.bib16))采样和实例自适应扩展（IAS）Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))之类的预算分配策略将 PRM 分数视为成功概率的代理。过度自信的 PRM 会对错误轨迹赋予过高的分数，导致算法对困难问题采样不足并错误地提交错误解决方案。改善 PRM 校准是使推理时扩展按预期工作的先决条件。

Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))近期的工作通过分位数回归解决了这一问题，拟合模型以在给定 PRM 表示的情况下预测成功概率分布的一组固定分位数水平。虽然有效，但这种方法存在结构局限性：分位数水平必须在训练时固定，因此模型无法在不重新训练的情况下查询任意置信水平。此外，分位数回归独立对待每个分位数，因此较高分位数可能产生比低分位数更低的预测值，违反了有效分位数函数的基本属性。在推理时扩展决策依赖于不确定性分布形状的场景中，这些局限性限制了灵活性和可靠性。

我们提出使用条件最优传输（OT）来解决这些局限性。基于Bunne等人(2022 (https://arxiv.org/html/2605.06785#bib.bib2))的双网络架构，我们修改 CondOT 架构以条件化于 PRM 隐藏状态，学习从 PRM 表示到经验成功结果分布的校准最优传输映射。校准后的不确定性估计可以直接通过Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))的 IAS 框架指导每道题目的计算分配。与其将校准后的后验坍缩为点估计，我们的方法将完整的预测分布传播到分配决策中，在成功概率不确定时增加样本预算，在置信度高时减少预算。

我们总结具体贡献如下：

1. 1\.一种基于条件最优传输的 PRM 校准方法，从 PRM 隐藏状态学习完整的单调条件分位数函数，无需重新训练即可从单一模型中在任意置信水平下提供灵活的不确定性估计。
2. 2\.经验证据表明，对于定义良好的 PRM，OT 校准在分布内（MATH-500）和分布外（AIME24-25）基准测试中，其 Brier 分数、ECE 和加权分位数损失均显著优于未校准的 PRM 和分位数回归。
3. 3\.对下游 Best-of-N IAS Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))性能的分析显示，OT 灵活的预测分布使得准确率优于未校准的基础 PRM。

## 2 相关工作

过程奖励模型（PRM）：过程奖励模型（PRM）对中间推理步骤进行评分，估计其对产生正确最终解决方案的贡献Cobbe等人(2021 (https://arxiv.org/html/2605.06785#bib.bib5)); Lightman等人(2023 (https://arxiv.org/html/2605.06785#bib.bib7)); Uesato等人(2022 (https://arxiv.org/html/2605.06785#bib.bib6))。它们广泛应用于推理时扩展算法，后者通常依赖于 Best-of-N 采样：生成多个候选响应并使用奖励模型选择得分最高的输出Chow等人(2024 (https://arxiv.org/html/2605.06785#bib.bib15)); Cobbe等人(2021 (https://arxiv.org/html/2605.06785#bib.bib5)); Brown等人(2024 (https://arxiv.org/html/2605.06785#bib.bib16))。最近的 PRM 如 Qwen-PRM Zhang等人(2025 (https://arxiv.org/html/2605.06785#bib.bib17))、Shepherd-PRM Wang等人(2024 (https://arxiv.org/html/2605.06785#bib.bib14))和 ReasonEval Xia等人(2025 (https://arxiv.org/html/2605.06785#bib.bib18))在推理基准测试中表现出强劲的性能。Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))也提出了用于 LLM 推理时扩展的实例自适应采样策略，我们在本工作中也采用了这些策略。

然而，PRM 往往校准不佳，产生过度自信的成功估计，可能导致次优的搜索决策。先前的工作使用分位数回归对 PRM 输出的不确定性进行建模来解决这一问题Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))，但此类方法需要在训练期间固定分位数水平，限制了推理时的灵活性。我们提出通过条件最优传输学习完整的条件分位数函数来校准 PRM，这允许进行灵活高效的不确定性估计，而无需针对不同分位数水平进行重新训练。

LLM 的不确定性量化：大型语言模型的不确定性量化（UQ）已在多种设置下得到研究，包括词元级预测、序列级生成以及对结构化输出的决策。方法包括基于预测似然的度量Kadavath等人(2022 (https://arxiv.org/html/2605.06785#bib.bib22))、自我一致性Wang等人(2022 (https://arxiv.org/html/2605.06785#bib.bib23))和共形预测Ye等人(2024 (https://arxiv.org/html/2605.06785#bib.bib25))。这些方法主要关注最终输出或词元预测的不确定性估计，而我们关注中间步骤。已提出各种校准 LLM 不确定性预测的方法，如温度缩放Guo等人(2017 (https://arxiv.org/html/2605.06785#bib.bib26))、带有奖励的强化学习Damani等人(2025 (https://arxiv.org/html/2605.06785#bib.bib27))、学习从语义意义到置信度分数的映射Cox等人(2025 (https://arxiv.org/html/2605.06785#bib.bib24))，或通过探针技术Liu等人(2024 (https://arxiv.org/html/2605.06785#bib.bib28))。与这些通常产生点对点或特定任务不确定性估计的方法不同，我们的方法学习 PRM 输出上结构化的、表示条件化的不确定性模型，从而能够在跨置信水平的一致估计。

神经网络的条件最优传输：最优传输为概率分布之间的映射提供了合理的框架（Villani 等人 (2009 (https://arxiv.org/html/2605.06785#bib.bib9)); Peyré 和 Cuturi (2019 (https://arxiv.org/html/2605.06785#bib.bib10))），并已应用于生成建模Arjovsky等人(2017 (https://arxiv.org/html/2605.06785#bib.bib11))和域自适应Courty等人(2016 (https://arxiv.org/html/2605.06785#bib.bib12), 2017 (https://arxiv.org/html/2605.06785#bib.bib13))。近期工作探索了使用神经网络学习条件于上下文的最优传输映射Rodriguez-Pardo等人(2025 (https://arxiv.org/html/2605.06785#bib.bib3)); Bunne等人(2022 (https://arxiv.org/html/2605.06785#bib.bib2)); Wang等人(2025 (https://arxiv.org/html/2605.06785#bib.bib4))。在本工作中，我们将Bunne等人(2022 (https://arxiv.org/html/2605.06785#bib.bib2))的双网络适配以条件化于大型语言模型（LLM）隐藏状态。虽然先前工作主要使用条件最优传输进行分布建模和生成建模，但我们将其用于学习预测分数与结果之间的校准映射，旨在为不确定性估计学习一致的条件分位数函数。

## 3 预备知识

我们在大型语言模型（LLM）的推理时扩展Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))中使用过程奖励模型（PRM）生成和评估多个推理轨迹。我们引入该设置中使用的轨迹、成功概率和实例自适应推理时扩展的符号。

Best-of-N Brown等人(2024 (https://arxiv.org/html/2605.06785#bib.bib16)):

$\mathbf{x}^{(i)} = (x_1^{(i)}, x_2^{(i)}, \dots, x_{T^{(i)}}^{(i)}) \sim \text{LLM}(q), \quad \text{for } i=1, \dots, N.$

其中 $q$ 是查询，$\mathbf{x}$ 是由 LLM 生成的推理轨迹（总共生成 $N$ 条完整轨迹），$x_i$ 是第 $i$ 个推理步骤，$T$ 是轨迹的总长度。

每条轨迹的分数由 PRM 分配 ($r^{(i)} = \text{PRM}(q, \mathbf{x}^{(i)})$)，最终输出是获得最高奖励的轨迹。

**成功概率**：推理时扩展的一个关键量是部分轨迹的成功概率。

$p \triangleq \Pr(x_{t+1:T} \text{ generated by LLM yields a correct answer} \mid q, \mathbf{x}_{0:t})$

其中 $q$ 是查询，$\mathbf{x}_{0:t}$ 是从步骤 1 到步骤 $t$ 到目前为止生成的轨迹部分。注意 $\mathbf{x}_{0:0}$ 是一个空序列。该量捕捉了继续当前推理路径导致正确解决方案的可能性。在实践中，$p$ 是未知的且必须估计，这引入了直接影响下游决策的不确定性。

**实例自适应推理时扩展（IAS）**：根据估计的成功概率分配计算预算。给定目标置信度 $C \in (0,1)$，达到此置信度所需的样本数量为：

$N^{\star}(p,C) \overset{\triangle}{=} \min \{n \in \mathbb{N} : \Pr(\text{at least one out of } n \text{ trajectories is correct}) \geq C \}.$

在实践中，这近似为Park等人(2025 (https://arxiv.org/html/2605.06785#bib.bib1))使用 PRM 计算，其中 $\hat{r}^{(\beta)}$ 是 PRM 在分位数水平 $\beta$ 处的估计成功概率，$N_{\text{max}}$ 是最大预算约束。

$N_{\text{IAS}}(p,C) \overset{\triangle}{=} \frac{\log(1-C)}{\log(1-p)}$

$N_{\text{IAS}} \overset{\triangle}{=} \min \{\lceil N_{\text{IAS}}(\hat{r}^{(\beta)},C) \rceil, N_{\text{max}}\} \quad (1)$

对于 IAS 的有效性，拥有 $p$ 的良好校准估计至关重要。

## 4 通过条件最优传输校准 PRM

我们的目标是学习给定 PRM 的成功概率 $p$ 的校准估计。与其预测单个标量或固定的一组分位数（如分位数回归校准所做的那样），我们旨在学习结果的完整条件分布，从而在所有分位数水平上实现一致的不确定性估计。

**用于校准的最优传输**：最优传输（OT）为概率分布之间的映射提供了合理的框架。在 Monge 表述中，OT 寻求一个映射 $T^{\star}$，将源分布 $\mu$ 推送到目标分布 $\nu$，同时最小化传输成本：

$T^{\star} := \arg \inf_{T_{\#}\mu=\nu} \int_{\mathbb{R}^{d}} \|x-T(x)\|^{2} \,d\mu(x).$

在我们的设置中，$\mu$ 对应于条件于 PRM 隐藏状态的 PRM 预测分布，而 $\nu$ 对应于成功结果的实证分布。学习到的传输映射将预测分数与校准后的结果分布对齐。

**条件最优传输（CondOT）**：我们对 CondOT Bunne等人(2022 (https://arxiv.org/html/2605.06785#bib.bib2))进行了修改。CondOT 通过最优传输的对偶表述学习条件于上下文的最优传输映射。给定上下文 $h$，CondOT 学习一个将源分布传输到目标分布的映射。在对偶表述中，最优传输用两个标量势函数 $f$ 和 $g$ 表示。CondOT 使用两个部分输入凸神经网络（PICNNs）参数化这些对偶势函数，$g: \text{PICNN}_{\theta_g}(\cdot, h)$ 和 $f: \text{PICNN}_{\theta_f}(\cdot, h)$，在传输变量上强制凸性的同时实现灵活的函数逼近。对偶势函数 $f$ 和 $g$ 可以通过以下最小-最大目标学习：

$\ell^{f}_{\