Optimal Policy Minimum Bayesian Risk
Ramón Fernandez Astudillo, Md Arafat Sultan, Aashka Trivedi, Yousef El-Kurdi, Tahira Naseem, Radu Florian, Salim Roukos
TL;DR
This work tackles robust and scalable inference-time decoding for LLMs by reframing minimum Bayesian risk decoding (MBRD) through a KL-controlled optimal policy. It defines an optimal posterior $p^*(y|x) \propto p_R(y|x) \exp(R(y,x)/\beta)$ and derives Rao-Blackwellized, self-normalized importance sampling estimators to compute expectations with respect to $p^*$. The paper introduces OP-MBRD and its efficient variant OPE-MBRD, which leverage process reward models and a sample-budgeted stopping rule to balance reward, similarity, and reference signals while maintaining asymptotic guarantees. Empirical results on math (MATH-500) and coding (HumanEval) show OP-MBRD generally outperforms standard MBRD and BoN across multiple generator-PRM configurations, with substantial throughput gains when calibration is favorable. The work offers a principled, modular framework for combining rewards and similarity signals at inference time with well-characterized behavior and practical efficiency improvements.
Abstract
Inference scaling helps LLMs solve complex reasoning problems through extended runtime computation. On top of long chain-of-thought (long-CoT) models, purely inference-time techniques such as best-of-N (BoN) sampling, majority voting, or more generally, minimum Bayes risk decoding (MBRD), can further improve LLM accuracy by generating multiple candidate solutions and aggregating over them. These methods typically leverage additional signals in the form of reward models and risk/similarity functions that compare generated samples, e.g., exact match in some normalized space or standard similarity metrics such as Rouge. Here we present a novel method for incorporating reward and risk/similarity signals into MBRD. Based on the concept of optimal policy in KL-controlled reinforcement learning, our framework provides a simple and well-defined mechanism for leveraging such signals, offering several advantages over traditional inference-time methods: higher robustness, improved accuracy, and well-understood asymptotic behavior. In addition, it allows for the development of a sample-efficient variant of MBRD that can adjust the number of samples to generate according to the difficulty of the problem, without relying on majority vote counts. We empirically demonstrate the advantages of our approach on math (MATH-$500$) and coding (HumanEval) tasks using recent open-source models. We also present a comprehensive analysis of its accuracy-compute trade-offs.