Theoretical guarantees on the best-of-n alignment policy

Ahmad BeiramiAlekh AgarwalJonathan BerantAlexander Nicholas D'AmourJacob EisensteinChirag NagpalAnanda Theertha Suresh

article2025ICML130 citations

Establishes rigorous theoretical guarantees for best-of-n alignment sampling by disproving a widely used closed-form KL divergence formula, bounding policy drift and win rates, and introducing a tighter KL estimator to guide test-time compute scaling.

Listen

Deploying generative artificial intelligence models safely requires alignment techniques that improve output quality and adhere to safety rules without degrading baseline capabilities. A widely adopted method is best-of-n sampling, where a system generates multiple candidate responses, scores them using a reward function, and selects the highest-scoring candidate. To monitor that the model does not drift too far from its original behavior, practitioners track statistical distance using Kullback-Leibler (KL) divergence, which measures distribution drift. Historically, the field relied on an analytical formula to estimate this drift as sample size grows. The article evaluates the theoretical foundations of best-of-n sampling to determine the mathematical accuracy of this formula and establish rigorous guarantees on distribution drift and win rates against the base model.

To conduct this evaluation, the article develops a closed-form probability mass function for the best-of-n policy under standard assumptions of finite outcomes and unique rewards. The authors derive theoretical upper and lower bounds for distribution drift and win rates, introducing a practical estimator for drift. They also compare best-of-n sampling against alternative rejection sampling mechanisms, such as rewind-and-repeat thresholding, and evaluate their findings through numerical simulations and empirical tests on language tasks using an instruction-tuned model.

The article demonstrates that the standard analytical formula used across the literature is not an exact equality but an upper bound on true distribution drift. When candidate outputs have low probabilities and sample sizes are small, the formula is reasonably close to reality; however, when the sample size is large or specific high-reward outputs have substantial probability mass, the formula drastically overestimates drift by an unbounded margin. Additionally, the article proves that the win rate of best-of-n sampling over the base model is strictly upper-bounded by the ratio of the sample size to the sample size plus one. The proposed drift estimator closely tracks actual behavior across tested regimes, and the analysis confirms that best-of-n sampling achieves near-optimal win rate versus drift tradeoffs at practical sample sizes below 1,000.

These findings have direct operational and governance implications for deploying generative models. Because existing literature relied on an overly conservative upper bound, best-of-n sampling actually preserves base model capabilities and safety guardrails significantly better than previously reported. Decision-makers can achieve competitive performance without the costly retraining required by complex reinforcement learning pipelines. However, this high alignment efficiency also implies that malicious actors can effectively repurpose best-of-n sampling to bypass safety guardrails if reward signals are unconstrained.

Organizations should adopt the article's proposed estimator to track distribution drift more accurately in deployment pipelines and leverage blockwise or standard best-of-n sampling within sample sizes below 1,000 for efficient inference-time alignment. For future research and development, teams should design hybrid approaches that balance compute cost, target reward, and capability preservation, while simultaneously engineering safeguards against adversarial test-time jailbreaking.

The findings carry high confidence due to formal mathematical proofs validated by extensive empirical simulations. Users should note the minor limitation that exact drift estimates can exhibit sample variance across individual draws, which requires averaging across small batches of prompts for reliable point estimates in production monitoring.

Google analytical expressions generative model inference best-of-n alignment policy inference-time alignment preference optimization

arXiv: 2401.01879

Paper: Scaling Laws for Reward Model Overoptimization, Leo Gao et al. (2023). This foundational study derives empirical scaling laws for KL divergence drift and overoptimization in best-of-n sampling, directly providing the analytical baseline formulas and empirical behaviors that the source paper mathematically bounds and refines.
Paper: Direct Preference Optimization: Your Language Model is Secretly a Reward Model, Rafael Rafailov et al. (2023). It introduces standard preference optimization frameworks and analyzes the fundamental trade-offs along reward-KL divergence frontiers that motivate test-time rejection sampling methods.
Paper: Training language models to follow instructions with human feedback, Long Ouyang et al. (2022). It establishes the standard language model alignment paradigm using reward models and KL divergence tracking to constrain drift from the reference policy.
Paper: Let's Sample Step by Step: Adaptive-Consistency for Efficient Reasoning and Coding with LLMs, Pranjal Aggarwal et al. (2023). It explores inference-time multi-candidate sampling and stopping criteria, offering practical context for evaluating compute-efficiency trade-offs in response selection.

Paper: DeAL: Decoding-time Alignment for Large Language Models, James Y. Huang et al. (2025). It extends inference-time alignment beyond full-sequence best-of-n candidate selection by developing heuristic-guided search and lookahead alignment dynamically during decoding.
Paper: A General Framework for Inference-time Scaling and Steering of Diffusion Models, Raghav Singhal et al. (2025). It generalizes inference-time reward steering beyond standard sequence-level best-of-n sampling by iteratively filtering and resampling intermediate generation trajectories.
Paper: Reasoning with Sampling: Your Base Model is Smarter Than You Think, Aayush Karan et al. (2026). It builds on test-time sampling dynamics by demonstrating how sampling from base model probability distributions can unlock strong reasoning performance without full reinforcement learning pipelines.
Paper: From RLHF to Direct Alignment: A Theoretical Unification of Preference Learning for Large Language Models, Tarun Raheja et al. (2026). It synthesizes preference learning techniques under a unified theoretical framework, examining policy drift regularization and distribution coverage issues analyzed in the source.

Cover for Theoretical guarantees on the best-of-n alignment policy

Abstract

A simple and effective method for the inference-time alignment and scaling test-time compute of generative models is best-of-n sampling, where n samples are drawn from a reference policy, ranked based on a reward function, and the highest ranking one is selected. A commonly used analytical expression in the literature claims that the KL divergence between the best-of-n policy and the reference policy is equal to log(n) − (n − 1)/n. We disprove the validity of this claim, and show that it is an upper bound on the actual KL divergence. We also explore the tightness of this upper bound in different regimes, and propose a new estimator for the KL divergence and empirically show that it provides a tight approximation. We also show that the win rate of the best-of-n policy against the reference policy is upper bounded by n/(n + 1) and derive bounds on the tightness of this characterization. We conclude with analyzing the tradeoffs between win rate and KL divergence of the best-of-n alignment policy, which demonstrate that very good tradeoffs are achievable with n < 1000.

Topics

analytical expressions generative model inference maximum likelihood estimator probabilistic inference reward function closed-form expression performance guarantees test application time

Keywords

best-of-n alignment policy inference-time alignment preference optimization reinforcement learning rejection sampling reward model KL divergence unbiased estimator

1. Introduction
2. Derivation of the Best-of-n Policy
3. Relations Between the KL Divergence and the Analytical Formula
3.1. Upper Bounds on the Gap
3.2. Lower Bounds on the Gap
4. Proposed Estimator for KL Divergence
5. Win Rate of the Best-of-n Policy
5.1. Upper Bounds on the Win Rate Gap
5.2. Lower Bounds on the Win Rate Gap
6. Rewind-and-Repeat: Rejection Sampling Beyond Best-of-n
7. Win rate vs KL Divergence Tradeoffs
8. Conclusion
Impact Statement
References
A. Proofs of the main results of the paper
A.1. Proofs of Section 3
A.2. Proofs of Section 4
A.3. Proofs of Section 5
A.4. Proofs of Section 6
A.5. Proofs of Section 7
B. KL divergence of blockwise best-of-n
C. Experiments
C.1. Details of experiments on Alpaca dataset
C.2. Computation of tilted min/max
C.3. Experiments with machine translation prompts

Knowls

Knowl 1 — Probability Mass Function of the Best-of-n Policy
theoretical result
Let $\pi_{\mathrm{ref}}(\cdot|x)$ be a reference language model distribution over a finite set of discrete responses $\mathcal{Y}_x$ for a given context or prompt $x$ , and let $r(x, y) \in \mathbb{R}$ be a unique reward assigned to each candidate response $y$ . The best-of- $n$ alignment policy $\pi^{(n)}(\cdot|x)$ draws $n$ independent and identically distributed responses $y_1, \dots, y_n \sim \pi_{\mathrm{ref}}(\cdot|x)$ and selects response $y = y_{k^*}$ where $k^* = \arg\max_{k \in \{1, \dots, n\}} r(x, y_k)$ .

The probability mass function (PMF) of the best-of- $n$ policy is given in closed form by:

$\pi^{(n)}(y|x) = F_{\pi_{\mathrm{ref}}}(y|x)^n - F_{\pi_{\mathrm{ref}}}^-(y|x)^n$

where for any policy $\pi$ , the cumulative distribution functions with respect to reward are defined as:

$F_\pi(y|x) := \mathbb{P}_{z \sim \pi(\cdot|x)}[r(x, z) \le r(x, y)]$

$F_\pi^-(y|x) := \mathbb{P}_{z \sim \pi(\cdot|x)}[r(x, z) < r(x, y)]$
Knowl 2 — Upper Bound on the KL Divergence of the Best-of-n Policy
theoretical result
For any sample size $n \in \mathbb{N}$ , context $x$ , and reference policy $\pi_{\mathrm{ref}}$ , the Kullback-Leibler (KL) divergence between the best-of- $n$ policy $\pi^{(n)}$ and the reference policy $\pi_{\mathrm{ref}}$ is upper bounded by:

$D_{\mathrm{KL}}(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \le \overline{\mathrm{KL}}_n := \log(n) - \frac{n - 1}{n}$

Furthermore, averaging over any prompt distribution $\mu$ , the expected KL divergence satisfies:

$D_{\mathrm{KL}}^\mu(\pi^{(n)} \,\|\, \pi_{\mathrm{ref}}) := \mathbb{E}_{x \sim \mu}\left[ D_{\mathrm{KL}}(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \right] \le \log(n) - \frac{n - 1}{n}$

This result disproves the common claim in existing literature that $D_{\mathrm{KL}}(\pi^{(n)} \,\|\, \pi_{\mathrm{ref}}) = \log(n) - (n - 1)/n$ , establishing instead that this analytical formula is an upper bound on the true KL divergence.
Knowl 3 — Bounds on the Best-of-n KL Divergence Approximation Gap
theoretical result
Let the gap between the analytical formula $\overline{\mathrm{KL}}_n = \log(n) - (n - 1)/n$ and the true KL divergence of the best-of- $n$ policy be defined as:

$G_{\mathrm{KL}}^{(n)}(x) := \overline{\mathrm{KL}}_n - D_{\mathrm{KL}}(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \ge 0$

Let $H_2(\pi_{\mathrm{ref}}|x) := -\log \sum_{y \in \mathcal{Y}^*} (\pi_{\mathrm{ref}}(y|x))^2$ denote the conditional Rényi entropy of order 2 of the reference model. The gap satisfies the upper bound:

$G_{\mathrm{KL}}^{(n)}(x) \le 2n(n - 1)e^{-H_2(\pi_{\mathrm{ref}}|x)}$

If the reference model is $\delta$ -bound (i.e., $\pi_{\mathrm{ref}}(y|x) \le \delta$ for all $y$ ), then $G_{\mathrm{KL}}^{(n)}(x) \le 2n(n - 1)\delta$ , showing that the analytical formula is tight when $n^2\delta \ll 1$ .

Conversely, defining $\varepsilon_\infty := \pi_{\mathrm{ref}}(y_{\max}(x)|x)$ where $y_{\max}(x) := \arg\max_y r(x, y)$ , the gap is lower bounded for any $n$ by:

$G_{\mathrm{KL}}^{(n)}(x) \ge \left(1 - (1 - \varepsilon_\infty)^n\right) \left[ \log\left(\frac{n\varepsilon_\infty}{1 - (1 - \varepsilon_\infty)^n}\right) - \frac{n - 1}{n} \right] - (n - 1)(1 - \varepsilon_\infty)^n \log(1 - \varepsilon_\infty) > 0$

In the asymptotic regime $n \to \infty$ , the lower bound scales as $G_{\mathrm{KL}}^{(n)}(x) \ge \log(n\varepsilon_\infty) + o_n(\log n)$ , which grows unbounded whenever $n\varepsilon_\infty \gg 1$ .
Knowl 4 — Estimator for Best-of-n Policy KL Divergence
model/method
Let $y \sim \pi^{(n)}(\cdot|x)$ be an outcome generated by selecting the highest-reward response from $n$ samples drawn from $\pi_{\mathrm{ref}}(\cdot|x)$ , and let $\varepsilon_n := \pi_{\mathrm{ref}}(y|x)$ denote the likelihood of the chosen response under the reference policy.

The single-sample estimator for the KL divergence between $\pi^{(n)}$ and $\pi_{\mathrm{ref}}$ is defined as:

$\widehat{D}_{\mathrm{KL}}(\varepsilon_n) := d_n(\varepsilon_n)$

where the function $d_n(\varepsilon)$ is given by:

$d_n(\varepsilon) := (1 - \varepsilon)^n \left[ \log n + (n - 1)\log(1 - \varepsilon) - \frac{n - 1}{n} \right] + \left(1 - (1 - \varepsilon)^n\right) \log\left( \frac{1 - (1 - \varepsilon)^n}{\varepsilon} \right)$

For any realization of $\varepsilon_n$ , the estimator satisfies deterministic bounds:

$0 \le \widehat{D}_{\mathrm{KL}}(\varepsilon_n) \le \log(n) - \frac{n - 1}{n}$

Because the standard deviation of $\widehat{D}_{\mathrm{KL}}(\varepsilon_n)$ on a single run is at most $\log n$ , averaging over $M = O(\log n \log(1/\delta))$ independent draws from the best-of- $n$ policy guarantees a standard deviation smaller than $\delta$ .
Knowl 5 — Calibrated Reward and General Win Rate Representation
definition
For a context $x$ and reference policy $\pi_{\mathrm{ref}}$ , let $F_{\pi_{\mathrm{ref}}}(y|x) := \mathbb{P}_{z \sim \pi_{\mathrm{ref}}(\cdot|x)}[r(x, z) \le r(x, y)]$ and $F_{\pi_{\mathrm{ref}}}^-(y|x) := \mathbb{P}_{z \sim \pi_{\mathrm{ref}}(\cdot|x)}[r(x, z) < r(x, y)]$ . The calibrated reward function is defined as:

$C_{\pi_{\mathrm{ref}}}(x, y) := \frac{F_{\pi_{\mathrm{ref}}}(y|x) + F_{\pi_{\mathrm{ref}}}^-(y|x)}{2}$

For any policy $\pi$ , its win rate against $\pi_{\mathrm{ref}}$ under reward $r$ , accounting for ties via $w_r(y \succ z | x) := \mathbb{I}(r(x, y) > r(x, z)) + \frac{1}{2}\mathbb{I}(r(x, y) = r(x, z))$ , is defined as:

$W_r(\pi(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) := \mathbb{E}_{y \sim \pi(\cdot|x), z \sim \pi_{\mathrm{ref}}(\cdot|x)}[w_r(y \succ z | x)]$

The win rate of any policy $\pi$ against the reference policy can be equivalently expressed as the expected calibrated reward:

$W_r(\pi(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) = \mathbb{E}_{y \sim \pi(\cdot|x)}\left[ C_{\pi_{\mathrm{ref}}}(x, y) \right]$
Knowl 6 — Theoretical Bounds on the Best-of-n Win Rate
theoretical result
For any sample size $n \in \mathbb{N}$ and context $x$ , the win rate of the best-of- $n$ policy $\pi^{(n)}$ against the reference policy $\pi_{\mathrm{ref}}$ satisfies the upper bound:

$W_r(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \le \frac{n}{n + 1}$

The gap $G_W^{(n)}(x) := \frac{n}{n + 1} - W_r(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \ge 0$ is upper bounded in terms of the conditional Rényi entropy of order 2, $H_2(\pi_{\mathrm{ref}}|x)$ , by:

$G_W^{(n)}(x) \le \frac{n - 1}{2} e^{-H_2(\pi_{\mathrm{ref}}|x)}$

If $\pi_{\mathrm{ref}}$ is $\delta$ -bound ( $\pi_{\mathrm{ref}}(y|x) \le \delta$ for all $y$ ), then $G_W^{(n)}(x) \le \frac{n - 1}{2}\delta$ , establishing that $W_r(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \approx \frac{n}{n + 1}$ when $n\delta \ll 1$ .

Moreover, with $\varepsilon_\infty := \pi_{\mathrm{ref}}(\arg\max_y r(x, y)|x)$ , the gap is lower bounded by:

$G_W^{(n)}(x) \ge \frac{n}{n + 1}\left(1 - (1 - \varepsilon_\infty)^{n+1}\right) - \left(1 - (1 - \varepsilon_\infty)^n\right)\left(1 - \frac{\varepsilon_\infty}{2}\right) > 0$

As $n \to \infty$ , $G_W^{(n)}(x) \ge \frac{\varepsilon_\infty}{2}(1 + o_n(1))$ , demonstrating that whenever $\varepsilon_\infty > 0$ , the gap remains strictly bounded away from zero.
Knowl 7 — Rewind-and-Repeat Policy and its Win Rate vs KL Divergence
theoretical result
The rewind-and-repeat rejection sampling policy $\pi_\Phi$ draws independent samples $y_1, y_2, \dots \sim \pi_{\mathrm{ref}}(\cdot|x)$ sequentially and outputs the first sample $y_M$ whose reward meets a threshold $\Phi \in \mathbb{R}$ (i.e., $r(x, y_M) \ge \Phi$ and $r(x, y_k) < \Phi$ for all $k < M$ ).

Let $w_\Phi(x) := \mathbb{E}_{y \sim \pi_{\mathrm{ref}}(\cdot|x)}[\mathbb{I}(r(x, y) \ge \Phi)]$ denote the per-trial acceptance probability. The probability mass function of $\pi_\Phi$ is:

$\pi_\Phi(y|x) = \begin{cases} \frac{\pi_{\mathrm{ref}}(y|x)}{w_\Phi(x)} & \text{if } r(x, y) \ge \Phi \\ 0 & \text{if } r(x, y) < \Phi \end{cases}$

The exact KL divergence and win rate of the rewind-and-repeat policy against $\pi_{\mathrm{ref}}$ are given in closed form by:

$D_{\mathrm{KL}}(\pi_\Phi(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) = \log\left( \frac{1}{w_\Phi(x)} \right)$

$W_r(\pi_\Phi(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) = 1 - \frac{1}{2} w_\Phi(x)$
Knowl 8 — Unbiased KL Divergence Estimator for Rewind-and-Repeat
theoretical result
Let $M$ be the number of trials until acceptance in the rewind-and-repeat sampling policy with acceptance probability $w_\Phi(x)$ (such that $M$ follows a geometric distribution with parameter $w_\Phi(x)$ , $\mathbb{P}(M = k) = (1 - w_\Phi(x))^{k-1}w_\Phi(x)$ ). Let $H_k := \sum_{i=1}^k \frac{1}{i}$ denote the $k$ -th Harmonic number, with $H_0 := 0$ .

The exact KL divergence between the rewind-and-repeat policy $\pi_\Phi$ and the reference policy $\pi_{\mathrm{ref}}$ satisfies:

$D_{\mathrm{KL}}(\pi_\Phi(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) = \mathbb{E}[H_{M-1}]$

Consequently, $H_{M-1}$ is an unbiased, single-trajectory sample estimator for the KL divergence of the rewind-and-repeat procedure.
Knowl 9 — KL Divergence Upper Bound for Blockwise Best-of-n Decoding
theoretical result
In blockwise best-of- $n$ decoding $\pi_B^{(n)}$ , responses are generated in sequential blocks of length $B$ tokens. Given a partially decoded sequence $y^t$ , $n$ candidate block continuations $z_{(1)}^B, \dots, z_{(n)}^B \sim \pi_{\mathrm{ref}}(\cdot|x, y^t)$ of length $B$ are sampled from the reference model, and the continuation maximizing a reward function $r(x, y^t, z_{(k)}^B)$ is selected. Decoding repeats until the end-of-sequence token is generated.

For a decoded sequence $y$ of total length $|y|$ steps, the KL divergence of the resulting policy $\pi_B^{(n)}$ relative to the reference policy $\pi_{\mathrm{ref}}$ satisfies:

$D_{\mathrm{KL}}(\pi_B^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \le \mathbb{E}_{y \sim \pi_B^{(n)}(\cdot|x)}\left[ \left\lceil \frac{|y|}{B} \right\rceil \right] \left( \log(n) - \frac{n - 1}{n} \right)$

where $\lceil \cdot \rceil$ is the ceiling operator. As block length $B \to \infty$ , this recovers the sequence-level best-of- $n$ upper bound.
Knowl 10 — Conjectured Upper Bound on Best-of-n Win Rate vs KL Divergence Tradeoff
theoretical result
Let $D_n := D_{\mathrm{KL}}(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x))$ denote the KL divergence of the best-of- $n$ policy. The tradeoff between the win rate $W_r(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x))$ and $D_n$ is conjectured to be upper bounded for all $n$ by:

$W_r(\pi^{(n)}(\cdot|x) \,\|\, \pi_{\mathrm{ref}}(\cdot|x)) \le W_0(D_n)$

where $W_0(D) := \ell^{-1}(D)$ is the functional inverse of $\ell : [0.5, 1) \to [0, \infty)$ defined as:

$\ell(\tau) := \log\left(\frac{\tau}{1 - \tau}\right) + \frac{1}{\tau} - 2$

This function characterizes the continuous Pareto frontier bounding the achievable win rate for a given KL divergence under best-of- $n$ sampling.

Coverage note — Omitted material includes the implicit parametric tradeoff curves in Appendix A.5 (Theorems A.6 and A.7, which directly combine the stated KL and win rate bounds) and specific empirical prompt listings from the Alpaca and translation experiments.

References

1.Amini, A., Vieira, T., and Cotterell, R. Variational best-of-n alignment. International Conference on Learning Representations (ICLR), 2025.
2.Azar, M. G., Guo, Z. D., Piot, B., Munos, R., Rowland, M., Valko, M., and Calandriello, D. A general theoretical paradigm to understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics, pp. 4447–4455. PMLR, 2024.
3.Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., DasSarma, N., Drain, D., Fort, S., Ganguli, D., Henighan, T., et al. Training a helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05862, 2022.
4.Balashankar, A., Sun, Z., Berant, J., Eisenstein, J., Collins, M., Hutter, A., Lee, J., Nagpal, C., Prost, F., Sinha, A., Suresh, A. T., and Beirami, A. InfAlign: Inference-aware language model alignment. International Conference on Machine Learning (ICML), 2025.
5.Beetham, J., Chakraborty, S., Wang, M., Huang, F., Bedi, A. S., and Shah, M. Liar: Leveraging alignment (best-of-n) to jailbreak llms in seconds. arXiv preprint arXiv:2412.05232, 2024.
6.Brown, B., Juravsky, J., Ehrlich, R., Clark, R., Le, Q. V., Ré, C., and Mirhoseini, A. Large language monkeys: Scaling inference compute with repeated sampling. arXiv preprint arXiv:2407.21787, 2024.
7.Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
8.Coste, T., Anwar, U., Kirk, R., and Krueger, D. Reward model ensembles help mitigate overoptimization. International Conference on Learning Representations (ICLR), 2024.
9.Eisenstein, J., Nagpal, C., Agarwal, A., Beirami, A., D’Amour, A., Dvijotham, D., Fisch, A., Heller, K., Pfohl, S., Ramachandran, D., Shaw, P., and Berant, J. Helping or herding? reward model ensembles mitigate but do not eliminate reward hacking. Conference on Language Modeling (COLM), 2024.
10.Gao, L., Schulman, J., and Hilton, J. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835–10866. PMLR, 2023.
11.Gemma, Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahriari, B., Ramé, A., et al. Gemma 2: Improving open language models at a practical size. arXiv preprint arXiv:2408.00118, 2024.
12.Go, D., Korbak, T., Kruszewski, G., Rozen, J., and Dymetman, M. Compositional preference models for aligning LMs. arXiv preprint arXiv:2310.13011, 2023.
13.Gui, L., Gârbacea, C., and Veitch, V. BoNBoN alignment for large language models and the sweetness of best-of-n sampling. Neural Information Processing Systems (NeurIPS), December 2024.
14.Hilton, J. and Gao, L. Measuring Goodhart’s law, April 2022. URL https://openai.com/research/measuring-goodharts-law. Accessed: 2024-01-03.
15.Hughes, J., Price, S., Lynch, A., Schaeffer, R., Barez, F., Koyejo, S., Sleight, H., Jones, E., Perez, E., and Sharma, M. Best-of-n jailbreaking. arXiv preprint arXiv:2412.03556, 2024.
16.Kim, M., Thonet, T., Rozen, J., Lee, H., Jung, K., and Dymetman, M. Guaranteed generation from large language models. International Conference on Learning Representations (ICLR), 2025.
17.Korbak, T., Elsahar, H., Kruszewski, G., and Dymetman, M. On reinforcement learning and distribution matching for fine-tuning language models with no catastrophic forgetting. Advances in Neural Information Processing Systems, 35:16203–16220, 2022a.
18.Korbak, T., Perez, E., and Buckley, C. RL with KL penalties is better viewed as Bayesian inference. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 1083–1091, 2022b.
19.Li, T., Beirami, A., Sanjabi, M., and Smith, V. On tilted losses in machine learning: Theory and applications. Journal of Machine Learning Research, 24(142):1–79, 2023.
20.Li, Y., Wei, F., Zhao, J., Zhang, C., and Zhang, H. Rain: Your language models can align themselves without finetuning. In The Twelfth International Conference on Learning Representations, 2024.
21.Mroueh, Y. Information theoretic guarantees for policy alignment in large language models. arXiv preprint arXiv:2406.05883, 2024.
22.Mudgal, S., Lee, J., Ganapathy, H., Li, Y., Wang, T., Huang, Y., Chen, Z., Cheng, H.-T., Collins, M., Strohman, T., Chen, J., Beutel, A., and Beirami, A. Controlled decoding from language models. International Conference on Machine Learning (ICML), 2024.
23.Nakano, R., Hilton, J., Balaji, S., Wu, J., Ouyang, L., Kim, C., Hesse, C., Jain, S., Kosaraju, V., Saunders, W., et al. WebGPT: Browser-assisted question-answering with human feedback. arXiv preprint arXiv:2112.09332, 2021.
24.Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C. L., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.02155, 2022.
25.Qiu, J., Lu, Y., Zeng, Y., Guo, J., Geng, J., Wang, H., Huang, K., Wu, Y., and Wang, M. Treebon: Enhancing inference-time alignment with speculative tree-search and best-of-n sampling. arXiv preprint arXiv:2410.16033, 2024.
26.Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C. D., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. arXiv preprint arXiv:2305.18290, 2023.
27.Scheurer, J., Campos, J. A., Korbak, T., Chan, J. S., Chen, A., Cho, K., and Perez, E. Training language models with language feedback at scale. arXiv preprint arXiv:2303.16755, 2023.
28.Sessa, P. G., Dadashi, R., Hussenot, L., Ferret, J., Vieillard, N., Ramé, A., Shariari, B., Perrin, S., Friesen, A., Cideron, G., et al. Bond: Aligning llms with best-of-n distillation. International Conference on Learning Representations (ICLR), 2025.
29.Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling llm test-time compute optimally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.03314, 2024.
30.Stiennon, N., Ouyang, L., Wu, J., Ziegler, D., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, P. F. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33: 3008–3021, 2020.
31.Sun, H., Haider, M., Zhang, R., Yang, H., Qiu, J., Yin, M., Wang, M., Bartlett, P., and Zanette, A. Fast best-of-n decoding via speculative rejection. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
32.Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Stanford alpaca: An instruction-following llama model, 2023.
33.Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
34.Yang, J. Q., Salamatian, S., Sun, Z., Suresh, A. T., and Beirami, A. Asymptotics of language model alignment. International Symposium on Information Theory (ISIT), July 2024a.
35.Yang, K. and Klein, D. FUDGE: Controlled text generation with future discriminators. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 3511–3535, Online, June 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.naacl-main.276. URL https://aclanthology.org/2021.naacl-main.276.
36.Yang, T., Mei, J., Dai, H., Wen, Z., Cen, S., Schuurmans, D., Chi, Y., and Dai, B. Faster wind: Accelerating iterative best-of-n distillation for llm alignment. arXiv preprint arXiv:2410.20727, 2024b.
37.Zhao, Y., Khalman, M., Joshi, R., Narayan, S., Saleh, M., and Liu, P. J. Calibrating sequence likelihood improves conditional language generation. In The Eleventh International Conference on Learning Representations, 2022.

Citation

MLA

Beirami, A., et al. “Theoretical Guarantees on the Best-of-n Alignment Policy”. arXiv, 2024, http://arxiv.org/abs/2401.01879v3.

APA

Beirami, A., Agarwal, A., Berant, J., D'Amour, A., Eisenstein, J., Nagpal, C., & Suresh, A. T. (2024). Theoretical guarantees on the best-of-n alignment policy. arXiv. http://arxiv.org/abs/2401.01879v3

Chicago

Beirami, A., A. Agarwal, J. Berant, et al. 2024. “Theoretical Guarantees on the Best-of-n Alignment Policy”. arXiv. http://arxiv.org/abs/2401.01879v3.

Harvard

Beirami, A. et al. (2024) “Theoretical guarantees on the best-of-n alignment policy”, arXiv [Preprint]. Available at: http://arxiv.org/abs/2401.01879v3.

Vancouver

1. Beirami A, Agarwal A, Berant J, D'Amour A, Eisenstein J, Nagpal C, Suresh AT (2024) Theoretical guarantees on the best-of-n alignment policy. arXiv

BibTeX

@article{beirami2024theoretical,
  title = {Theoretical guarantees on the best-of-n alignment policy},
  author = {Beirami, Ahmad and Agarwal, Alekh and Berant, Jonathan and D'Amour, Alexander and Eisenstein, Jacob and Nagpal, Chirag and Suresh, Ananda Theertha},
  year = {2024},
  journal = {arXiv},
  url = {http://arxiv.org/abs/2401.01879v3},
  eprint = {2401.01879}
}

Metadata:arXiv

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Theoretical guarantees on the best-of-n alignment policy

Abstract

Topics

Keywords

Table of Contents

Knowls

Knowl 1 — Probability Mass Function of the Best-of-n Policy

Knowl 2 — Upper Bound on the KL Divergence of the Best-of-n Policy

Knowl 3 — Bounds on the Best-of-n KL Divergence Approximation Gap

Knowl 4 — Estimator for Best-of-n Policy KL Divergence

Knowl 5 — Calibrated Reward and General Win Rate Representation

Knowl 6 — Theoretical Bounds on the Best-of-n Win Rate

Knowl 7 — Rewind-and-Repeat Policy and its Win Rate vs KL Divergence

Knowl 8 — Unbiased KL Divergence Estimator for Rewind-and-Repeat

Knowl 9 — KL Divergence Upper Bound for Blockwise Best-of-n Decoding

Knowl 10 — Conjectured Upper Bound on Best-of-n Win Rate vs KL Divergence Tradeoff

References

Citation

BibTeX

Access the Paper