Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment

TL;DR

The paper investigates how to best utilize inference-time computation to improve the quality of responses from a pre-trained policy when the reward model is imperfect. It shows that naive Best-of-N (BoN) can suffer reward hacking as compute increases, and formalizes a statistical skyline governed by reward-model error $\varepsilon_{\text{RM}}^2$ and coverage $\mathcal{C}^{\pi^{\star}}(x)$. The authors introduce Inference-Time Pessimism, a $\chi^2$-regularized inference-time algorithm implemented via rejection sampling that achieves regret-optimal and scaling-monotone performance, with near-optimal compute. They provide lower bounds and tight guarantees for BoN, present rigorous analysis of the new algorithm, and validate the approach empirically on multiple tasks and models. The work advances a principled understanding of how to leverage inference-time computation while mitigating reward-model misalignment, with practical implications for scalable alignment of large language models.

Abstract

Inference-time computation offers a powerful axis for scaling the performance of language models. However, naively increasing computation in techniques like Best-of-N sampling can lead to performance degradation due to reward hacking. Toward a theoretical understanding of how to best leverage additional computation, we focus on inference-time alignment, which we formalize as the problem of improving the quality of responses drawn from a pre-trained policy, given a prompt of interest and access to an imperfect reward model. We analyze the performance of inference-time alignment algorithms in terms of (i) response quality, and (ii) compute, and provide new results that highlight the importance of the pre-trained policy's coverage over high-quality responses for performance and compute scaling: 1. We show that Best-of-$N$ alignment with an ideal choice for $N$ can achieve optimal performance under stringent notions of coverage, but provably suffers from reward hacking when $N$ is large, and fails to achieve tight guarantees under more realistic coverage conditions. 2. We introduce $\texttt{InferenceTimePessimism}$, a new algorithm which mitigates reward hacking through deliberate use of inference-time compute, implementing the principle of pessimism in the face of uncertainty via rejection sampling; we prove that its performance is optimal and does not degrade with $N$, meaning it is scaling-monotonic. We complement our theoretical results with an experimental evaluation that demonstrate the benefits of $\texttt{InferenceTimePessimism}$ across a variety of tasks and models.

Figures (11)

• Figure 1: Comparison between performance of BoN-Alignment (dashed lines) and InferenceTimePessimism algorithm (solid lines) on GSM8K with reward model OASST as $\widehat{r}$ and several different choices of $\pi_{\mathsf{ref}}$. Left: $\texttt{BoN-Alignment}\xspace$ initially improves accuracy over $\pi_{\mathsf{ref}}$, but eventually degrades as $N$ increases, while InferenceTimePessimism is monotone, as predicted by our theory. Right: BoN-Alignment overoptimizes the reward model, with high $\widehat{r}$ but lower accuracy as $N$ increases, whereas InferenceTimePessimism stops increasing $\widehat{r}$ with $N$ beyond a certain threshold determined by a regularization parameter.
• Figure 2: Compute-normalized comparison of InferenceTimePessimism (solid lines) and BoN-Alignment (dashed lines) on GSM8K with OASST as $\widehat{r}$, as a function of regularization parameter $\beta$. We run BoN-Alignment with $N=2^{13}$ and run InferenceTimePessimism until rejection sampling accepts (capped to $N=2^{13}$). Left: We see that InferenceTimePessimism can improve significantly over naïve BoN-Alignment for large $N$ due to the reward overoptimization. Center: Number of responses required for InferenceTimePessimism to accept an answer decreases as $\beta$ increases, as predicted by our theory. Right: Estimated reward $\widehat{r}$ for InferenceTimePessimism decreases as $\beta$ increases.
• Figure 3: Comparison of InferenceTimePessimism (solid lines) and BoN-Alignment (dashed lines) in accuracy and estimated reward $\widehat{r}$ for GSM8K for four reward models and choices of $\pi_{\mathsf{ref}}$.
• Figure 4: Comparison of InferenceTimePessimism (solid lines) and BoN-Alignment (dashed lines) in accuracy and estimated reward $\widehat{r}$ for MMLU for four reward models and choices of $\pi_{\mathsf{ref}}$.
• Figure 5: Comparison of InferenceTimePessimism (solid lines) and BoN-Alignment (dashed lines) in accuracy and estimated reward $\widehat{r}$ for MATH for four reward models and choices of $\pi_{\mathsf{ref}}$.

Theorems & Definitions (48)

• remark 1
• definition 1: Sample-and-evaluate framework
• proposition 1: Necessity of coverage
• theorem 1: Guarantee for BoN-Alignment
• theorem 2: Lower bound for BoN-Alignment
• remark 2: Proof technique
• theorem 3: Guarantee for BoN-Alignment under uniform coverage
• theorem 4: Guarantee for InferenceTimePessimism
• theorem 5: Query complexity lower bound
• remark 3: Comparison to KL-regularization