PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Audrey Huang, Adam Block, Qinghua Liu, Nan Jiang, Akshay Krishnamurthy, Dylan J. Foster

arXiv:2503.21878v2cs.AIcs.LGstat.ML

TL;DR

The paper asks how imperfect reward models and base-policy coverage limit inference-time alignment as computation increases. It analyzes Best-of-N and introduces InferenceTimePessimism, finding that the latter achieves optimal regret and scaling-monotonic performance through inference-time regularization and rejection sampling.

  • Problem

    The paper studies how reward-model imperfection and base-policy coverage constrain response quality and compute scaling in inference-time alignment.

  • Method

    The paper formulates inference-time alignment statistically and analyzes Best-of-N alongside InferenceTimePessimism, which uses χ2-regularized rejection sampling.

  • Results

    InferenceTimePessimism is regret-optimal, uses near-optimal compute, and is scaling-monotonic, whereas Best-of-N can overoptimize and lacks tight guarantees under weaker coverage.

  • Takeaways & Limitations

    More deliberate use of inference-time computation can mitigate reward overoptimization while preserving performance and compute scaling.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies inference-time alignment as selecting high-quality responses from a base policy using an imperfect reward model, focusing on how coverage and computation affect regret and scaling. It shows that Best-of-N can overoptimize the reward model, while InferenceTimePessimism achieves optimal, scaling-monotonic performance through deliberate regularization and rejection sampling.

  • Framework: Inference-time alignment selects a high-quality response from many base-policy candidates using a reward model as a proxy for task performance.The framework gives black-box access to policy samples, reward labels, and policy log-probabilities.
  • Best-of-N limitations: Best-of-N initially improves quality but can degrade after N passes a critical threshold because reward-model errors produce reward overoptimization.The same tension makes increasing N unsafe: larger candidate pools improve selection opportunities while increasing tail-risk from reward-model errors.
  • Research questions: Imperfect reward models create an information-theoretic barrier to maximizing the unobserved true reward.The framework therefore studies both regret to the true reward and the computational queries required to minimize it.
  • InferenceTimePessimism: InferenceTimePessimism uses rejection sampling to implement χ2-regularization at inference time, mitigating overoptimization through pessimism in the face of uncertainty.The method is designed to use inference-time computation to extract information from the reward model and quantify uncertainty.
  • Best-of-N guarantees: Best-of-N can achieve optimal regret under stringent uniform coverage with tuned N, but lacks tight guarantees under weaker average-case coverage.Its best possible reward can remain below the framework’s optimal skyline regardless of computation.
  • Results: InferenceTimePessimism is regret-optimal, uses near-optimal compute, and remains scaling-monotonic by avoiding overoptimization beyond a regularization-dependent threshold.Experiments across tasks, base policies, and reward models report degradation for Best-of-N as N increases, while InferenceTimePessimism remains monotone and can outperform Best-of-N with untuned budgets.

2 A Statistical Framework for Inference-Time Alignment

The paper formalizes inference-time alignment as selecting a high-quality response from samples of a base policy using an imperfect reward model. Its statistical framework measures performance through regret and computational efficiency through query complexity, highlighting reward-model fidelity and policy coverage as fundamental constraints.

  • Framework: Inference-time alignment selects a response from a base policy using an imperfect reward model, with quality measured by true-reward regret.The base policy supplies candidate responses, while the reward model proxies an unknown true reward.
  • Reward-model quality: The reward model’s expected squared error under the base policy is the framework’s measure of reward quality, although this notion may yield conservative bounds.The paper abstracts reward-model training away and assumes an available imperfect reward model.
  • Computational model: The sample-and-evaluate framework measures computation by the number of sampled responses and reward-model evaluations.The framework provides black-box sampling and evaluation access to the base policy and reward model.
  • Coverage: The best achievable reward is limited by both reward-model error and the base policy’s coverage of high-quality responses.A lower bound formalizes coverage as necessary for minimizing inference-time regret.
  • Coverage: Best-of-N fits this framework by sampling N responses independently and returning the candidate with the highest reward-model score.The framework encompasses Best-of-N and related rejection-sampling schemes.

3 Understanding Best-of-N Alignment: Guarantees and Limitations

The analysis shows that Best-of-N improves candidate coverage as N grows but simultaneously increases exposure to reward-model errors. Its guarantees can be optimal under uniform coverage, yet realistic coverage conditions permit suboptimality and reward hacking, motivating more refined regularization.

  • Best-of-N procedure: Best-of-N samples N candidates and returns the one with the largest imperfect reward-model score.Its sole adjustable parameter is the sample size N.
  • Guarantees: The main guarantee balances a coverage-dependent term that decreases with N against an overoptimization term that increases with N.Larger samples improve the chance of finding a response competitive with the comparator, but tail candidates are more vulnerable to reward-model error.
  • Guarantees: BoN-Alignment’s optimal sample size must be large but not too large because N simultaneously improves response quality and regularizes against overoptimization.This opposing dual role produces a slower ε_RM^2/3 rate in the general guarantee.
  • Limitations: For all N, a lower-bound construction shows BoN-Alignment can incur regret that grows as N is scaled, reflecting overfitting to reward-model errors.The construction uses a comparator with C_π⋆ = eO(1), so the issue is not eliminated by favorable ordinary coverage.
  • Uniform coverage: Under bounded uniform coverage, BoN-Alignment achieves a regret scaling matching the skyline’s statistical rate, but this stronger coverage may be unrealistic for softmax policies.Softmax policies can have exponentially larger uniform coverage coefficients than ordinary coverage coefficients.

4 An Optimal Algorithm: Inference-Time Pessimism

InferenceTimePessimism uses inference-time χ2-regularization and rejection sampling to separate computation from statistical tuning. The resulting method achieves optimal regret and remains scaling-monotonic, while its compute requirements are near-optimal.

  • Algorithm: InferenceTimePessimism separates the computational budget N from the regularization parameter β, unlike BoN-Alignment’s single parameter N.This separation enables independent tuning of computation and statistical performance.
  • Algorithm: The χ2-regularizer limits overfitting to low-probability responses with spuriously high estimated rewards.β controls the trade-off between greedier behavior at small values and conservative, heavy-tailed behavior at large values.
  • Guarantees: Choosing β to balance overoptimization and bias yields a regret bound matching the lower bound up to absolute constants.The method is therefore regret-optimal under the framework’s coverage conditions.
  • Guarantees: For fixed β, its regret guarantee holds for all sufficiently large N, preventing performance from dropping below the guaranteed threshold as computation scales.The paper calls this property scaling-monotonicity.
  • Compute: N ≳ 1/εRM(x) queries are necessary to achieve the optimal regret bound, and InferenceTimePessimism matches this requirement.The lower bound applies to sample-and-evaluate algorithms under the stated instance conditions.

5 Experiments

Experiments compare InferenceTimePessimism with BoN-Alignment across multiple tasks, policies, and reward models. BoN-Alignment degrades as N increases, whereas InferenceTimePessimism generally improves or remains robust to increased computation.

  • Setup: The evaluation spans three primary task families, four reward models, and four base policies, with additional AlpacaEval-2.0 results reported preliminarily.Experiments use GSM8K, MMLU, and MATH in the main suite.
  • Parameter effects: As β increases, InferenceTimePessimism requires fewer responses before rejection sampling accepts an answer, while estimated reward decreases.These trends are reported for GSM8K with OASST as the reward model.
  • GSM8K: As N increases on GSM8K, BoN-Alignment’s accuracy exhibits a characteristic dip caused by reward overoptimization.The comparison uses true reward accuracy and estimated reward from OASST.
  • GSM8K: InferenceTimePessimism’s accuracy increases monotonically with N in the GSM8K comparison.This behavior matches the theoretical scaling-monotonicity guarantee.
  • Main comparison: Across GSM8K, MMLU, and MATH, InferenceTimePessimism tends to achieve higher average performance than BoN-Alignment.Many observed differences are not statistically significant.
  • Evaluation scope: Under binary correctness evaluation, the optimal policy is uniform over correct answers, placing these experiments in a regime where BoN is predicted to perform near-optimally.More refined evaluation, such as proof correctness, may produce a different comparison.

6 Conclusion

The paper argues that inference-time alignment depends on coverage, scaling, and optimality, and that deliberate computation can mitigate reward overoptimization. It also identifies extensions involving richer policy information, joint training and inference design, and online exploration.

  • Conclusion: InferenceTimePessimism and its analysis show that sophisticated inference-time computation can achieve optimal regret and compute scaling.The contribution includes efficient sampling from regularized distributions, query-complexity lower bounds, and approximate rejection-sampling analysis.
  • Future directions: Because the base policy is treated as a black box beyond coverage, future work could use learned representations or other refined policy properties.The paper frames this as an open direction rather than a demonstrated improvement.
  • Future directions: The paper proposes co-designing training-time and inference-time alignment because inference-time BoN can be suboptimal in general.This direction concerns procedures that optimize downstream inference behavior jointly with training.
  • Future directions: An additional open direction is inference-time exploration with true rewards or high-quality verifiers in online settings.This would differ from the paper’s imperfect-reward-model formulation.

A.1 Connection to Offline (Training-Time) Alignment

The paper formulates inference-time alignment as a per-instance variant of offline alignment, abstracting away reward-model training while retaining imperfect reward evaluation. Its experiments instantiate this framework across diverse tasks, policies, and reward models.

  • Connection to offline alignment: Inference-time alignment is treated as a variant of offline alignment that abstracts away how the reward model was trained.The objective is per-instance regret for an arbitrary fixed prompt rather than an average over i.i.d. prompts.
  • Connection to offline alignment: The algorithms and analyses can be combined with reward-estimation procedures to derive end-to-end sample-complexity guarantees for offline alignment.This connects the per-instance framework back to reward-model training from preference data.
  • Connection to online alignment: Unlike online alignment, the framework makes decisions using an imperfect reward model rather than feedback from the true reward model.The paper identifies online feedback as a complementary line of theoretical research.
  • Experimental setting: The empirical suite covers GSM8K, MMLU, MATH, and a preliminary AlpacaEval-2.0 study.Correctness is measured by final answers for the math tasks, while AlpacaEval uses win rate against an evaluator language model.
  • Experimental setting: Experiments use multiple reward models and base policies, with approximately 20K temperature-1 responses generated per prompt, task, and policy.The main reward-model set includes OASST, GEMMA-RM, LLAMA-RM, and ARMO-RM.
  • Experimental setting: Bootstrap replicates estimate accuracy across sampled subsets, and fixed-N comparisons reuse the same samples for normalization and rejection sampling.The reported policy accuracy averages over prompts, with standard errors marginalized over prompts.

B.2 Results for Further Policies and Tasks

Across additional policies and tasks, InferenceTimePessimism remains essentially monotonic in compute, while properly tuned BoN-Alignment achieves similar performance in several comparisons.

  • InferenceTimePessimism is essentially monotonic in compute budget across GSM8K, MMLU, and MATH.
  • BoN-Alignment can also be monotonic when task-policy pairs are more in-distribution for a reward model such as ARMO-RM.
  • BoN-Alignment with properly tuned N and InferenceTimePessimism have similar performance across Phi-3-Small, Mistral-7B, and LLAMA-3-3B evaluations.LLAMA-3-3B is evaluated only on GSM8K and MMLU.
  • Under binary true rewards and uniform coverage over correct answers, BoN-Alignment performs near-optimally.

B.3 Further Experiments

Further experiments examine reward distributions, task and policy comparisons, compute normalization, and robustness to β. They broadly support monotonic scaling for InferenceTimePessimism while identifying evaluation and small-budget caveats.

  • Reward distributions indicate ARMO-RM separates correct from incorrect GEMMA-2-2B responses best for a representative GSM8K prompt.Its incorrect-response support is empirically below the correct-response support, so sufficiently large N selects the correct answer for that prompt.
  • InferenceTimePessimism shows less over-optimization across β, while BoN-Alignment typically requires β tuning to avoid over-optimization entirely.
  • Small β can produce a small InferenceTimePessimism performance dip when rejection sampling accepts no responses and the method defaults to BoN-Alignment.This is more likely when N is small relative to 1/β.
  • AlpacaEval-2.0 results cannot statistically separate BoN-Alignment from InferenceTimePessimism because evaluation is noisy and uses fewer replicates and prompts.The experiment uses 5 replicates per prompt instead of 50.
  • The normalization-constant procedure sorts and bins sampled responses by estimated reward before iterating toward a threshold.

C.2 Guarantee for ComputeNormConstant

The ComputeNormConstant analysis establishes that the algorithm finds a normalization threshold for the sampled reward distribution, with O(N log N) computational complexity.

  • A continuous, piecewise-linear normalization function that decreases from infinity to zero guarantees existence of the required λ.
  • Algorithm 3 finds λ such that the estimated normalization function satisfies bΦ(λ) = 1.
  • The generalized ComputeNormConstant procedure also finds λ satisfying bΦ(λ) = 1 in O(N log N) time.
  • The proof replaces reward-equivalent response bins with surrogate responses and shows their normalization functions coincide.
  • The algorithm locates the interval containing λ by scanning sorted reward bins and returns the corresponding threshold.
  • The rejection-sampling framework uses N base-policy samples and accepts each according to a truncated importance weight divided by M.
  • The threshold M trades approximation bias against termination cost: larger M reduces truncation bias but requires more samples.

D.3 Lower Bounds

The lower-bound results show that any selection strategy drawing only from N responses of the base policy requires sufficient coverage and sample complexity to approximate a target policy.

  • Any selection strategy receives only responses sampled from πref and must return one of those N candidates.
  • The lower bound applies to arbitrary selection algorithms in the sample-and-evaluate framework, not only to BoN-Alignment.
  • If N is below the target-policy rejection threshold Mπ,ε, the induced response distribution cannot be ε-close to the target policy.
  • The proof analyzes the regime N < M and omits prompt dependence while relating the threshold to approximate rejection-sampling coverage.

E Proofs from Section 2

The appendix develops technical guarantees and lower bounds for Best-of-N alignment using rejection sampling, EM-divergence, and coverage coefficients. It shows both information-theoretic limits and reward-model-driven overoptimization as computation increases.

  • General BoN analysis: Lemma F.1 bounds Best-of-N regret through EM-divergence between a comparator policy and the reference policy.The proof uses a rejection-sampling policy as an intermediate comparator and translates estimated-reward performance into true-reward regret.
  • General BoN analysis: The analysis parameter M trades off lower achievable regret against increasing sample complexity.The bound becomes useful when EM-divergence is translated into coverage coefficients.
  • Lower bounds: For insufficient computation, information-theoretic lower bounds apply to any inference-time selection algorithm, with thresholds governed by coverage quantities.The appendix also derives BoN-specific lower bounds once N is large enough.
  • Lower bounds: The lower-bound constructions separate L1- and L∞-coverage, showing that realistic coverage conditions can make approximation difficult and permit persistent BoN regret.One construction has Cπ⋆∞ exponentially larger than Cπ⋆.
  • Lower bounds: When the bad response appears, its modeled reward is maximal, so Best-of-N selects it even though the construction makes it task-suboptimal.This exposes overfitting to reward-model errors in the tail of the reference distribution.

F.6 Additional Results

These additional results establish information-theoretic guarantees for Best-of-N and characterize the optimization and estimation procedures used in the paper’s inference-time alignment analysis. They also verify the convex optimization structure underlying the policy construction.

  • Additional guarantees: The appendix gives an information-theoretic regret guarantee for Best-of-N and presents the main theorems as coverage-based relaxations.The interpretable theorems isolate dependencies on coverage coefficients and N.
  • Proof conventions: The proofs condition on a single prompt and suppress prompt dependence in notation while retaining prompt-specific quantities in the underlying guarantees.Expected returns are defined under either the reward model or the true reward function.
  • Policy construction: Strong duality and strong convexity make the KKT conditions sufficient to establish optimality and uniqueness of the proposed policy.A strictly feasible point establishes Slater’s condition for the convex program.
  • Policy construction: The optimal policy reweights the reference policy by relu(β^-1(r(y) − λ)), assigning zero mass when r(y) − λ ≤ 0.The normalization condition determines λ so the resulting policy has the required total mass.

G.2 Proof of Theorem 4.2

The proof of Theorem 4.2 applies the lower-bound construction to a logarithmic L1-coverage regime and transfers its information-theoretic regret bound to the target sample size.

  • Proof of Theorem 4.2: The construction has Cπ⋆(x) = O(log C), while the theorem considers computation on the order of 1/εRM.The proof invokes the first part of Theorem F.4 with p = 1/(2Cπ⋆(x)).
  • Proof of Theorem 4.2: Any algorithm using the specified smaller sample budget must suffer the theorem’s lower-bounded regret.The argument sets the lower-bound sample size to the target N and concludes the result directly.
Loading 2503.21878v2…