Know What You Don't Know: Uncertainty Calibration of Process Reward Models¶

会议: NeurIPS 2025
arXiv: 2506.09338
代码: http://young-j-park.github.io/know-what-you-dont-know
领域: LLM推理
关键词: 过程奖励模型, 校准, 分位数回归, 推理时缩放, 自适应采样

一句话总结¶

本文提出了一种基于分位数回归的PRM校准方法，使PRM输出的分数更准确地反映LLM实际推理成功概率，并基于校准后的PRM设计了实例自适应推理时缩放（IAS）策略，在保持准确率的同时显著降低推理成本。

研究背景与动机¶

领域现状：Process Reward Model (PRM) 在推理时缩放中扮演关键角色，用于引导Best-of-N采样和Beam Search等方法。PRM为每个中间推理步骤打分，评估当前推理路径最终得到正确答案的概率。

现有痛点：即使是SOTA的PRM（如Qwen-PRM、Shepherd-PRM）也存在严重的校准问题——它们倾向于高估成功概率，特别是在使用较弱的LLM或面对困难问题时。这是因为PRM是在特定策略模型上训练的，换用不同能力的LLM时会产生分布不匹配。

核心矛盾：PRM训练时绑定了特定策略模型的生成分布 \(\pi_\theta\)，但部署时可能使用不同能力的模型。一个在72B模型上训练的PRM会系统性高估1B模型的成功概率。

本文目标 (1) 如何校准现成PRM使其输出可靠的成功概率估计？(2) 如何利用校准后的概率实现自适应计算资源分配？

切入角度：传统的温度缩放等校准方法不适用于PRM（因为成功概率本身就是中间概率而非二元标签），因此提出用分位数回归来预测成功概率的分布，并利用保守的下分位数估计来指导资源分配。

核心 idea：用分位数回归微调PRM，预测成功概率的下界，从而实现保守但可靠的实例自适应推理预算分配。

方法详解¶

整体框架¶

给定一个现成PRM和目标LLM，通过三阶段数据收集（初始轨迹生成→前缀提取+Monte Carlo rollout→成功概率估计）构建校准数据集，然后用分位数回归微调PRM的预测头，最后基于校准后的PRM实现自适应采样。

关键设计¶

校准数据收集三阶段:
- 功能：为每个问题生成多条推理轨迹，对每条轨迹的每个前缀进行Monte Carlo rollout估计真实成功概率
- 核心思路：从MATH训练集采样500个问题，每个问题生成 \(N_{\text{val}}=8\) 条轨迹，对每个前缀再生成 \(N_{\text{MC}}=8\) 条续写，统计正确率作为真实成功概率 \(\tilde{p}^{(i,t)}\)
- 设计动机：Monte Carlo采样提供了与具体LLM绑定的经验成功概率，解决了PRM与策略模型不匹配的问题
分位数回归校准:
- 功能：修改PRM的预测头，输出多个分位数（如10%、50%、90%）的预测
- 核心思路：使用加权分位数损失 \(\text{wQL}(\hat{r}, \tilde{p}) = \frac{1}{N_q} \sum_{n=1}^{N_q} [\beta_n \cdot \max(0, \tilde{p} - \hat{r}^{(\beta_n)}) + (1-\beta_n) \cdot \max(0, \hat{r}^{(\beta_n)} - \tilde{p})]\)
- 设计动机：预测条件均值会导致约50%的情况高估成功概率，预测低分位数（如10th percentile）可以提供保守的下界估计，避免计算资源分配不足
实例自适应缩放（IAS）:
- 功能：根据校准后PRM估计的成功概率，动态调整每个问题的采样数量
- 核心思路：对于成功概率为 \(p\) 的问题，要达到目标正确率 \(C\)，所需最小采样数为 \(N_{\text{IAS}}(p, C) = \lceil \frac{\log(1-C)}{\log(1-p)} \rceil\)
- 设计动机：简单问题（高 \(p\)）需要很少的样本，困难问题（低 \(p\)）需要更多样本——这模拟了人类"简单题快做，难题多花时间"的策略

训练策略¶

仅微调PRM的预测头（扩展输出维度以支持多分位数预测），保持主干模型冻结
使用Qwen-PRM-7B和Shepherd-PRM-7B作为基础模型进行校准

实验关键数据¶

主实验 - PRM校准效果（Qwen-PRM-7B, Brier Score↓）¶

数据集	目标LLM	未校准	校准后	降幅
MATH500	Llama-3.2-1B	0.2414	0.0692	-71%
MATH500	Qwen-2.5-7B	0.1008	0.0818	-19%
MATH500	R1-Qwen-7B	0.1480	0.0828	-44%
AIME24-25	Llama-3.2-1B	0.1936	0.0029	-98%
AIME24-25	R1-Qwen-7B	0.4144	0.0694	-83%

消融实验 - BoN+IAS计算节约¶

数据集	模型	BoN(N=64)	BoN+IAS(校准)	Budget Ratio
MATH500	Qwen-2.5-7B	0.854	0.837	23.4%
MATH500	R1-Qwen-7B	0.864	0.857	31.3%
MATH500	Llama-3.2-1B	0.476	0.462	63.8%
AIME24-25	R1-Qwen-7B	0.267	0.180	96.4%

关键发现¶

PRM系统性高估成功概率，尤其对弱模型和OOD问题（偏差分布右偏，峰值接近1.0）
校准是IAS有效的前提——未校准的PRM无法可靠地指导资源分配
使用低分位数（如10%）比条件均值更适合IAS场景，因为它提供了保守估计
IAS在简单问题上节约大量计算，在难问题上分配更多资源，整体效率显著提升

亮点与洞察¶

分位数回归的巧妙应用：不同于传统校准方法只预测均值，分位数回归可以给出成功概率的分布信息（上界、中位数、下界）。低分位数提供保守估计，这对资源分配决策非常关键——宁可多采样也不要漏掉正确答案。
PRM不校准的根本原因分析：PRM的训练是策略依赖的，这意味着同一个PRM在不同LLM上的校准性完全不同。本文清楚地揭示了这一根本原因，并提供了优雅的解决方案。
IAS的采样复杂度公式 \(N_{\text{IAS}} = \lceil \log(1-C)/\log(1-p) \rceil\) 非常简洁实用，可以直接迁移到任何需要自适应采样的场景。

局限与展望¶

校准数据收集需要对目标LLM做大量rollout，成本不低
仅在数学推理任务上验证，未测试代码生成、开放问答等场景
Monte Carlo rollout的样本数（\(N_{\text{MC}}=8\)）相对较少，可能导致经验概率估计不准
改进方向：可以探索跨模型的校准迁移；减少rollout数量的高效校准方法；与更复杂的搜索策略（如MCTS）结合

评分¶

新颖性: ⭐⭐⭐⭐ 首次系统研究PRM校准问题，分位数回归的应用很巧妙
实验充分度: ⭐⭐⭐⭐ 多种PRM、多种LLM、多个benchmark，但可以更多样化
写作质量: ⭐⭐⭐⭐⭐ 问题定义清晰，理论和实验结合紧密
价值: ⭐⭐⭐⭐⭐ PRM校准是推理时缩放的关键瓶颈，本文的解决方案简洁实用