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.