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

arXiv: https://arxiv.org/abs/2605.06785 · HTML: https://mlsanity.com/paper/2605.06785/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 25 of its 33 references that are in the map.
- First posted on arXiv 7 May 2026.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)
- LLM alignment / RLHF / preference optimisation / [Process reward models for reasoning](https://mlsanity.com/alignment/process-reward-models-for-reasoning/)

## Builds on (most cited first, 25 of 25)

- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Language Models (Mostly) Know What They Know](https://mlsanity.com/paper/2207.05221/) (2022)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations](https://mlsanity.com/paper/2312.08935/) (2023)
- [Solving math word problems with process- and outcome-based feedback](https://mlsanity.com/paper/2211.14275/) (2022)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [The Lessons of Developing Process Reward Models in Mathematical Reasoning](https://mlsanity.com/paper/2501.07301/) (2025)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Benchmarking LLMs via Uncertainty Quantification](https://mlsanity.com/paper/2401.12794/) (2024)
- [Beyond Binary Rewards: Training LMs to Reason About Their Uncertainty](https://mlsanity.com/paper/2507.16806/) (2025)
- [PRMBench: A Fine-grained and Challenging Benchmark for Process-Level Reward Models](https://mlsanity.com/paper/2501.03124/) (2025)
- [Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models](https://mlsanity.com/paper/2412.15287/) (2024)
- [Evaluating Mathematical Reasoning Beyond Accuracy](https://mlsanity.com/paper/2404.05692/) (2024)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Know What You Don't Know: Uncertainty Calibration of Process Reward Models](https://mlsanity.com/paper/2506.09338/) (2025)
- [Efficient Neural Network Approaches for Conditional Optimal Transport with Applications in Bayesian Inference](https://mlsanity.com/paper/2310.16975/) (2023)
- [Mapping from Meaning: Addressing the Miscalibration of Prompt-Sensitive Language Models](https://mlsanity.com/paper/2510.17028/) (2025)
- [Neural Conditional Transport Maps](https://mlsanity.com/paper/2505.15808/) (2025)

## Abstract (arXiv)

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 citebunne2022supervised 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 citepark2025know. 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.