Distributional Process Reward Models: Calibrated Prediction of Future Rewards via Conditional Optimal Transport

Authors: Rachel Ma, Dylan Hadfield-Menell, Kristjan Greenewald

arXiv · PDF

Summary

The authors propose using conditional optimal transport to calibrate process reward models (PRMs) used in inference-time scaling, which currently overestimate success probabilities and are poorly calibrated. They adapt a conditional optimal transport method to learn a monotonic conditional quantile function over PRM scores, giving well-calibrated confidence bounds at any desired level, and plug this into adaptive scaling frameworks that use PRMs to decide how much compute to spend.

Main takeaways:

• Process reward models (PRMs) guide search and scaling at inference time, but they often give overconfident probability estimates.
• The authors use conditional optimal transport to map PRM hidden states to calibrated quantile estimates of success probability, preserving monotonicity.
• This yields valid confidence intervals and integrates into instance-adaptive scaling (spending more compute when the model is uncertain).
• On math reasoning benchmarks (MATH-500, AIME), the method improves calibration substantially and often improves downstream Best-of-N performance over uncalibrated PRMs.

Relevance

Not directly related to my work on persona installation or behavioral fine-tuning. This is about calibrating reward models for inference-time search in reasoning tasks, which is orthogonal to understanding how behaviors are installed or steered during training.

Abstract

arXiv:2605.06785v1 Announce Type: new Abstract: Inference-time scaling methods rely on Process Reward Models (PRMs), which are often poorly calibrated and overestimate success probabilities. We propose, to our knowledge, the first use of conditional optimal transport for calibrating PRMs, modifying conditional OT (CondOT) map learning \cite{bunne2022supervised} to estimate a monotonic conditional quantile function over success probabilities estimated by the PRM, conditioned on PRM hidden states. This yields structurally valid quantile estimates and enables efficient extraction of confidence bounds at arbitrary levels, which we integrate into the instance-adaptive scaling (IAS) framework of \cite{park2025know}. We evaluate on mathematical reasoning benchmarks spanning moderate-difficulty problems (MATH-500) and harder out-of-distribution problems (AIME). For PRMs with reliable ranking signals, our method substantially improves calibration over both uncalibrated PRMs and quantile regression. On downstream Best-of-N IAS performance, our method generally improves over uncalibrated PRMs. These results establish conditional optimal transport as another principled and practical approach to PRM calibration, offering structural guarantees and flexible uncertainty estimation.