Theoretical guarantees on the best-of-n alignment policy
Authors: Ahmad Beirami
Alekh Agarwal
Jonathan Berant
Alexander Nicholas D'Amour
Jacob Eisenstein
Chirag Nagpal
Ananda Theertha Suresh
Venue: ICML 2025 poster (2025)
Keywords: Alignment
rejection sampling
KL divergence
Primary area: theory
TL;DR: We provide guarantees on KL divergence and win-rate of the best-of-n alignment policy.

Abstract: A simple and effective method for the inference-time alignment of generative models is the best-of-$n$ policy, 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$.

Theoretical Guarantees on the Best-of-n Alignment Policy
Ahmad Beirami' Alekh Agarwal? Jonathan Berant' Alexander D'Amour' Jacob Eisenstein' Ananda Theertha Suresh?
Chirag Nagpal 2
Abstract
too far from the reference model.
A simple and effective method for the inference- time alignment and scaling test-time compute of generative models best-of-n sampling, where n samples are drawn from reference policy, ranked based reward function, and the highest rank- ing one is selected. A commonly used analytical expression in the literature claims that the KL di- vergence 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 the actual KL diver- gence. We also explore the tightness of this upper bound in different regimes, and propose 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 upper bounded by n/(n 1) and derive bounds on the tightness of this characterization. We conclude with analyzing the tradeoff's between win rate and KL divergence of the best-of-n alignment policy, which demon- strate that very good tradeoff's are achievable with n < 1000.
Recently, there has been proliferation of methods for alignment, which include KL-regularized reinforcement learning (Christiano et al., 2017; Ouyang et al., 2022), controlled decoding (Yang & Klein, 2021; Mudgal et al., 2024), SLiC (Zhao et al., 2022), direct preference optimiza- tion (Rafailov et al., 2023), and best-of-n finetuning (Tou- vron et al., 2023). their core, these methods try to solve the following regularized optimization problem:"
max (a)r)())) y) (D [Ta] (1) x(-1æ)
where Tet denotes reference language model; r(xx, y) E R represent sealar reward associated with response y for prompt æ; and the KL divergence (=(()<=)]lp(((<<)) is defined as
DRL. (x))(() == Exna(-la) log ((y
p(y az)
Note that Equation (1) has a closed-form solution (Korbak et al., 2022b;a):
(y(()) x The (ye
which defines exponential family of distributions with nice properties (Yang et al., 2024a). We also define the KL divergence averaged over prompts as D (gllp) = where f is a distribution over prompts. Notice that a small KI. divergence between the aligned policy and the reference policy is desired because it implies that the capabilities of the reference policy largely preserved (Gao et al., 2023; Coste et al., 2024; Eisenstein et al., 2024), which is also theoretically analyzed by Balashankar et al. (2025, Appendix B).
1. Introduction
Generative language models have shown be effective gen- eral purpose tools to solve various problems. While many problems can be solved in zero-shot manner, the output from the so-called reference model may not be outright de- sirable, e.g., it may violate safety rules or may not solve a math problem correctly. Alignment (Christiano et al., 2017: Stiennon et al.. 2020; Ouyang et al., 2022: Bai et al., 2022) and test-time compute scaling (Brown al., 2024; Snell al., 2024) aim at remedying this issue by further nudging the outcome to improve reward function while not drifting
E (4 (<<([=)))p((()))
To compare different alignment techniques, is customary to produce tradeoff curves that measure expected reward (or win rate) function of DRI ( ((ee)rr) for some aligned policy Guarantees on the KL divergence capture the preservation of the core capabilities of the model and tighter estimates on the KL divergence help give guarantees that
'Googee DeepMind 'Google Research. Correspondence comi,
Ahmad Beirami Ananda Theertha Suresh <theertha google.com>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
While theoretically we analyze this optimization problem as function of the prompt æ. in practice we can only solve it by taking another expectation over a set prompts x ~ μ.

---

Theoretical guarantees on the best-of-n alignment policy
the model doesn't lose core capabilities that were present in the reference checkpoint. Thus, it is desirable improve the reward with the least drift measured in KL divergence.
4.0
logini -in - 1// (eg. (41)
Proposed estimator (Eq (251)
Exact KL, (Eq. (5)
af
2.5
Despite all the udvancements in alignment, a simple. popu- lar, and pell-performing method for alignment remains to be the best-of-n policy (Nakano et al., 2021; Stiennon et al., 2020). In fact, Gao et al. (2023); Mudgal et al. (2024); Eisen- stein al. (2024) show that best-of-n consistently achieves compelling win rate vs KL tradeoff' curves, that even dom- inate those of KL-regularized reinforcement learning and other more involved alignment policies. Llama 2 (Touvron et al., 2023) uses best-of-n as a teacher outcomes to further finetune the base model. Mudgal et al. (2024) extended best-of-n through q-learning to block-wise best-of-n de- coding. This has also led to recent research on distilling best-of-n into new models (Gui al., 2024; Amini et al., 2025; Sessa et al., 2025; Qiu et al., 2024). Hughes et al. (2024); Beetham et al. (2024) use best-of-n as an effective method for jailbreaking. Best-of-n is also used as a strong baseline in scaling inference-tine compute (Brown et al., 2024; Snell et al., 2024). This overwhelming empirical suc- cess motivates our theoretical investigation of the best-of-n alignment policy.
2.0
1.5
0.5
0.0
10a
10a
10a
Figure 1. The analytical formula (log(n) (n 1)/n) (Equa- tion (4)), the exact KL divergence (Equation (5)), and the proposed estimator (Equation (25)), for Example 1, illustrating case where the gap between the analytical formula and the exact KI. diver- gence is unbounded.
tradeoffs for the best-of-n policy. Let us further inspect this formula using a toy example.
Example 1. Consider an unprompted model with x = 0 (no input) and binary output, E (0, 13. Let the two outcomes be equiprobable, i.e., (0) = ( = 2- Further, let r(0) = O, and r(1) = 1, i.e. outcome is more desirable than outcome 0. In this example, we can compute x(*) in closed form. Specifically, we can see that (n) (0) 2 and
Subsequent to this work, Yang et al. (2024a) provided theo- retical reasoning for the performance of best-of-n by show- ing it achieves asymptotically optimal Le tradeoffs. Gui et al. (2024) characterized the win rate vs KL gap to be small in the asymptotic regime of a language model whose outcomes have infinitesimally small likelihood. Sun al. (2024) made best-of-n faster through speculative rejection. Mrouch (2024) provided information-theoreti bounds reward vs KL tradeoffs for best-of-n.
T (n)) (1) = 1 +. Thus,
2n
DxL( (1) [[ log(2) - 2
(5)
where h(()) is the binary entropy function. 3 We compare the exact closed-form expression for KL. divergence with the analytical formula in Equation (4). As can be seen in Figure (and is evident from Equation (5)), the true KI. is upper bounded by log(2) for all n, whereas KL.,, grows unbounded n + co. We also report a new estimator for KL divergence that closely mirrors the true KL divergence.
Best-of-n. Let æ be a given input prompt to the language model. Let yy.....WM Un be n i.i.d. samples drawn from Ther( (z). The best-of-n strategy selects?
As we learnt from Example 1, the KL divergence between the best-of-n policy and the reference policy may be quite different from what the analytical formula used in the litera- ture suggests. In the rest of this paper, we shed some light on this formula, derive bounds on the KL divergence, and propose a new estimator for the KL divergence that better captures the behavior of the KL divergence. We also theo- retically reason about the win rate vs KL tradeoff's for the best-of-n policy, and justify its widespread use in language model alignment.
where k" := arg max r(x, yk). ke(n]
(3)
This process inherently leads to sampling from a new policy that is aligned to the reward, denoted by r(a)_ Notice that (1) = Tyels and increasing n increases the reward at the cost of drifting away from the base model.
Our goal in this paper is to better understand the best-of-n alignmnet policy. In particular, we are interested in theo- retical guarantees on DK (zx(*) for different values of n. A commonly used expression in the literature (Stiennon et al., 2020; Hilton & Gao, 2022; Coste et al., 2024; Gao et al., 2023; Go et al., 2023; Scheurer et al., 2023) claims
2. Derivation of the Best-of-n Policy
Our first step is to provide derivation for the best-of-n policy under two simplifying assumptions. Let r(æ, y) E R represent the scalar reward of response y in context æ.
(L (ree))
claim
KL,, := log(n) (n 1)/n. (4)
This formula is commonly used demonstrate reward-KL
*h(z) == x log(z) - (1 - x) log(1 z), for all z E (0. 1), and h()) h(1) = 0. Further, note that all logarithms in this paper are to the hase e.
We define [n] : n1.
2

---

Theoretical guarantees the best-of-n alignment policy
Assumption 2.1. We assume that the reward r(æ, y) is unique for all x. y-
Notice that if n. 1, then (y)ar) Ter(y)a)) For any n, Lemma 2.3 gives closed-form expression for ( (y/z), which we will use subsequently to derive the- oretical guarantees on the KL divergence and win rate of best-of-n.
Assumption 2.2. Lety"= == (y m (y)( 0}. We assume that the language model is such that [3**| < oo, i.e, there are finite possible outcomes (in each context).
We also remark that we can extend this PMF to (()) for real T 2 1. While may not be immediately clear how sample from this extension, it used for best-of-n distillation (Gui et al., 2024; Amini al., 2025; Sessa et al., 2025) and we also use it to give bounds in Section 7.
Note that Assumptions 2.1-2.2 are fairly non-restrictive and make the presentation of the results clearer.
The following result gives the probability mass function (PMF) of the best-of-n policy.
Lemma 2.3. Under Assumptions 2.1-2.2, for all N, the PMF of the best-of-n policy is given by
3. Relations Between the KL Divergence and the Analytical Formula
(h)) (y)
F FF
F (3)l:)"
(6)
Our first result shows that the analytical formula is an up- per bound on the (context-dependent) KL divergence. The proofs for this result and several subsequent results are rele- gated to Appendix A.
wherefor any distribution x,
F+(y(a)= = Pawn()aa(((() 2) rxx y),.
(7)
(u)x) == Pe (e) < r y)]
(8)
Theorem 3.1. For any n. E N, and any a, let KLn be defined in (4). Then,
Proof. Let V be the set of all possible outcomes of the lan- guage model, given prompt i.e., V2 == {y Ter (g) (x) > 0}. Further, let Le == D/ < oo (Assumption 2.2). We order all possible L outcomes as {U: he(r. such that if r(æe, ú,) > rax, ).), then > In other words, D1 is the least desirable outcome associated with the lowest reward, and Vr. is the most desirable outcome associated with the highest reward.
DxL( (1z) (Tg())))(( S KLn log(n) n
Corollary 3.2. For any n, and any prompt distribution p.
(Tt( [ = Emary DKI (n) (-1ee) Tl S KLn-
First notice that sampling from Ner is equivalent to sampling 0100 1], and returning ÿ: such that
Proof. This directly follows from Theorem 3.1.
F y (x)
a.).
(9)
Subsequent to this work, Mroueh (2024) has extended this result to a larger class of stochastic processes (with poten- tially continuous support such as diffusion models) through the application of the strong data processing inequality. In Appendix B, we also extend this result derive bounds the KL divergence of the blockwise best-of-n decod- ing (Mudgal et al., 2024). which generally allows to reach similar reward vs KI. tradeoffs with 10x smaller n. In the rest of this section, we characterize the gap defined as fol- lows:
Similarly, sampling from the best-of-n strategy is akin to sampling u1.... Un ~ 21(0, 1], and returning Wi, such that
Fan Vi- (x) S max Uk < FFar (F.
(10)
kE[n]
On the other hand, we know that the CDF of the maximum iid.
of u1,..... Un U(0, 1), for 7 E [0, 1] is given by
P
max us S -r kE|n]
(11)
(x) KL (x:)
KI, DKI (1(z)) (-rr(=)))) 2 0. (14)
Hence, for all n E N, the PMF of the best-of-n policy, denoted as (1²) given by
3.1. Upper Bounds on the Gap
(n) U. [x) = F (f:lz)) (" F Tue (Ti-1 12))"
Vi
[Ta]]
We need definition to state the upper bound results. Definition 3.3. A model Tee called 5-bound ())()) < 8 for all y E y* and x.
(12)
where F (ool2) x) :- O, and
Frum (is( (zx)
Σ (E(
Tret (tt az).
(13)
In particular, we are interested in characterizing the gap for a 5-bound model for a small 8, which is a model with all outcomes having small likelihoods. Next, we state our main upper bound.
The proof is completed by noticing that Frna (7--1 1/x) Rear y: lz).
3

---

Theoretical guarantees on the best-of-n alignment policy
Theorem 3.4. The gap in Equation (14) is upper bounded by
4. Proposed Estimator for KL Divergence
Motivated by the derivation of the best-of-n policy in Lemma 2.3, we propose new estimator for the KL di- vergence. As a warm-up, first notice the following upper bound:
(x) S 2n(n - 1)e" -a(sw)))
(15)
where Hn x) is the conditional Rényi entropy of order 2 of the language model given context æ, and Ha(a) for any distribution is defined as
Lemma 4.1. For any E N and any a,
(x(z)) = 1 log
Σ (=(ola)))) DEV-
(16)
S Eypaatn) log (1 (1- TTn(()/o))"" )
(B)(()/2)))
Corollary 3.5. Let (x)) S for all y E ,i, Tre is 6-bound. Then, the gap in Equation (14) is upper bounded by
Therefore we may suggest to use the following alternate estimator for KL divergence
Gk2 (x) S 2n(n - 1)5.
(17) DN
Proof. The proof follows by noticing that H2( log(1/5 (16) and invoking Theorem 3.4.
D DKL, (En)) == -log (1-(1-En) En)
(23)
En
Intuitively, if the model outcomes are fairly low probability, making it unlikely to get the same sample twice or more in the n outcomes for best-of-n, the analytical formula KL, in (4) could be relatively accurate, and the gap is bounded above. In other words, if Fee is 6-bound model, and n is sufficiently small such that n²8 < 1, then
where en is defined in Equation (19). Note that the expected value of Dos,,owseeEn) is an upper bound on the KI diver- gence between the best-of-n policy and the reference policy. However, this estimator loose by an additive constant of (n - 1)/n, especially when ne, < 1.
Here we propose a different estimator to close this gap. To derive the estimator, first notice the following result.
DEL (Ir (") (-|zz) (er((( log(n) n
(18)
Corollary 4.2. Let
This assumption is left implicit in the derivation of Hilton & Gao (2022) for the KL divergence best-of-n.
dx(() =0 - 8)" (log + (n 1) logg11111 -8) n - 1
n
3.2. Lower Bounds on the Gap
+ (1-(1=))") log
1 6) 2)"²
(24)
In this section, we characterize cases where the gap may be large. To this end, let us define
Era in Equution (19). Then, Recall the definition on
En T x)),
where y (m)) (-|z))
(19)
DRI (m) (m) [[ (x)) (E())))
Note that En is a random function of æ. In the limit as n + 00, we define
Note that Corollary 4.2 could not be directly used to derive an estimator for the KL divergence because we do not ob- serve c when performing the best-of-n policy. Inspired by this result, and given that can only observe Exs put forth the following practical estimator on the KL divergence.
EO Tere (ymax (z))
(20)
where Vmax (2)) arg maxyeyy r(z, y). Notice that & is a deterministic function of æ.
Theorem 3.6. Let Eco > 0 be defined in Equation (20). For E N, the gap between the analytical formula in Equa- tion (4) and KL divergence is lower bounded by
Definition 4.3. Let be defined in (19). Then, we propose the following estimator for the KI. divergence of the best- of-n policy and the reference policy:
G (x) 21 (1-
==0) (1=Ex))")
(l0g 1 (1 500)
1
DKL(En) := dr, (En).
(25)
-(n - -(n-1)(1 Ex0) log(1 (1- - Eoo) 0. (21)
Note that the estimator proposed in Definition 4.3 is ran- dom variable that depends on E.. We conjecture that in expectation it provides an upper bound on the true KL di- vergence.
Corollary 3.7. As n1 xi, the gap is lower bounded by c (x) > log(ne.>) o,, (logn.)
(22)
Conjecture 4.4. Let en be defined in (19). Then.
In particular, when nF >> 1, then the gap grows un- bounded as we already observed in Example 1.
Dxu(m (" (-)==) S E (1
DKa(en)

---

Theoretical guarantees on the best-of-n alignment policy
Memmstark 105 bast
Alomrle meantits lix 224 Coscr
tarem bonditi 201
Alerra boved EC 123
bac
Du Iterooo
1-10
30
32 12
19 367
101
Aremas femmle MD
ceittatteent bosttils lig 2014 unimite (ig 254 Dact
Do Saateeenln
popsed ennweriia 1250)
Figure 3. The analytical formula (log(n) - (n- - 1)/r), Equa- tion (4), the alternate bound, Equation (23), the proposed estimator, Equation (25). and the exact KL divergence. for two cherry picked examples. In the left panel, the output is supported on alphabet of size 5, where the highest reward outcome has a probability of 10-* and the rest of the probability mass is uniformly distributed over the rest of the outcomes. In the right panel, the output is sup- ported an alphabet of size 200, where the three highest reward outcomes have probabilities 10- 10-3 and 10- respectively. The rest of the probability mass is uniformly distributed over the rest of the outcomes.
107
2=1
Figure 2. The analytical formula (log(n) (ni 1)/n). Equa- tion (4), the alternate bound, Equation (23), the proposed esti- mator, Equation (25), and the exact KL. divergence, for uniform distributions supported alphabets of size L = 10, 10?, 10 10 respectively.
While we don't offer mathematical proof for Conjec- ture 4.4, tens of thousands of randomly generated numerical experiments suggest that it holds true.
n = L. For E < 1, the analytical formula of Equation (4), log(n) (n 1)/n1, has small gap with the actual KL di- vergence (which was also theoretically established in Corol- lary 3.5). On the other hand, when E >>>>>>>>>>> 1, the gap between log(n) - (n - 1)/n and the actual KL divergence becomes large and unbounded (which was also theoretically establish in Corollary 3.7). The alternate bound in Equation (23) cap- (26) tures the behavior of the KL divergence for 7 > 1 and has finite gap. However, it has gap of (n - 1)/n for & < 1 as previously discussed. Finally, we also observe that the proposed estimator in Equation (25) follows the behavior of the true KL divergence closely in all examples.
Let us further inspect the proposed estimator and its variance. We first show that it is strictly upper bounded by KL..
Lemma 4.5. For any realization of En, we have
0 S DKL(En) S KLn,
and hence
Ey (y-o))))1o)))
[Ri(En)
s KLn.
(27)
In the second set of examples, we cherry pick the probability mass function on the outcome create different behaviors of the KI. divergence, shown in Figure 3. In the left panel, the output is supported on an alphabet of size 5, where the highest reward outcome has a probability of 10-4 and the rest of the probability mass is uniformly distributed over the rest of the outcomes. We observe that KL divergence satu- rates early until the highest reward outcome is discovered with = 10 In the right panel, the output supported alphabet of size 200, where the highest reward out- come has probability of 10-5, the second highest reward outcome has a probability of 10 and the third highest reward outcome has a probability of 10-1, The rest of the probability mass is uniformly distributed over the rest of the outcomes. As can be seen, the KL divergence starts to saturate until the next high reward is outcome is discovered around n ss 10 and n 10 As can be seen, the analyti- cal formula in Equation (4) does not capture the behavior of the KL divergence at all whereas the alternate bound in Equation (23) is much better aligned with the actual be- havior. Finally, we observe that the proposed estimator in Equation (25) closely follows the actual KL divergence. We would like to recall that the proposed estimator random
Given this, can immediately bound the variance of the estimator too. Lemma 4.5 implies that standard devia- tion of the estimator is upper bounded by log n, which in turn implies that if the estimator is averaged over M = O (log n log }t) draws from the best-of-n model, the stan- dard deviation is guaranteed to be smaller than 8. Given that we are generally interested in n < 1000. the dependence on is mild Having said that, given each of the M batches contains iid samples (total of M iid samples), one should be able to build bootstrapped estimator for the variance with better guarnatees.
In what follows we numerically inspect the proposed esti- mator in a few scenarios, and compare it with the analytical formula and the exact KL divergence between the best-of-n policy and the reference policy.
The first set of examples, in Figure 2, are uniform dis- tributions over alphabets of varying sizes. Notice that 1 for a uniform distribution, and hence En = Exo = L. the estimator in Equation (25) and Equation (23) are de- terministic. As can be seen KI. divergence saturates around
5

---

Theoretical guarantees on the best-of-n alignment policy
Trandimm the to sking cerien... -u (677
Deecrbe the Sunction of a com (0.00
Traneform the falow ng senter.. lae g.130
Decor be the Runction com. (.000
22-
22-
34-
15-
33
2.
1.-
10-
24
Frgmmnt co dy Ieg)
ta
Papace cstre the (h1)
Prpeoee mil de DNE
Prpprie rntinaaa the ISS
N
6o
15
44
53
10
55
10
40
What of (t. clis)
tres lgo lor staying te... it..0
What the captaal Parnet la c0
Ohe tree lis for slaying Ne.. It. 0.200
100
15-
Gndlphuut bimee do 24 (40
ardgttallerss de By 180
kocal upai bas Ioou ay 22-06-0
60
Betta agar soand Korolay 62 App
6.
Auagiia torva (0g (0)
Anaytia terzaiie (40
30
Figure 4. The analytical formula (log(n) (n 1)/n), Equa- tion (4), better upper bound (Corollary 4.2, Appendix C), the proposed estimator, Equation (25), and the exact KL divergence, for four cherry picked examples from the Alpaca dataset (Taori et al., 2023) using Gemma 9B IT model (Gemma et al., 2024) with reward the lillleiiioood of response under the reference model.
Figure 5. The analytical formula (log(n) - (n - 1)/n), Equa- tion (4), better upper bound (Corollary 4.2, Appendix C), the proposed estimator, Equation (25), and the exact KL divergence, for four cherry picked examples from the Alpaca datasct (Taori ct al., 2023) using Gemma 9B ITT model (Gemma et al., 2024) with reward the negative length of the response.
and we have plotted the expected value of the estimators, whereas in practice both estimators are subject to variance due to the randomness in the value of c, (Equation (19)). To capture the deviation from mean, we compute the tilted min and rilted max (see Appendix C.2 for the definition), which are also plotted in Figure 3. As can be seen, in cases where en could widely vary based on whether a certain outcome appears in the set of n outcomes, the variance could be high.
In other words, w, > z|x²) indicates whether response wins over response z in context æ using the judge r. Then, let the win rate of policy π over the reference policy Tiet for prompt æ be defined as
(1 (-12)=)) Ta
[=a[(- [m]):= E (=)):
(-1a) Wy (y> z|ar).
(28)
It is clear that W- ) || æ:) = 0.5 and the goal of alignment is to improve win rate beyond 0.5 (with the lowest KL divergence between the two models). We further define the following, averaged over prompt distribution μ
In Figure 4, we compare the estimates for four cherry picked examples from the Alpaca dataset (Taori et al., 2023) using Gemma 9B IT model (Gemma et al., 2024) with reward being the bogliikeldhodd of the reference model. Note that for two examples where Soo is large, i.e., the prompts that induce less entropy in the response, such as "What is the capital of France?". the proposed estimator outperforms the analytical formula in Equation (4) considerably and lies very close the better upper bound in Corollary 4.2, whereas KL, (Equation (4)) is loose even for n R 20. The details on the prompts used can be found in Appendix C. In Figure 5, we repeat the same experiment but change the reward to the negative of length to prefer more concise responses and see similar trends. We also include experiments with machine translation in Appendix C.3.
W (T]])Taa) == Egeu (1
(29)
Note that it is clear that if y T (-|2)) and z Teer (-x)) then wr > 2|x)) is an unbiased estimator for win rate of policy against Wref-
We analyze the win rate and derive theoretical guarantees. Let Fa and F be defined in Equation (7) and Equation (8), respectively. We define calibrated reward as:
F ( xx) + F ma (3/x)
CrLA (Ic, y)
(30)
With this definition in place, notice that win rate could be expressed follows.
5. Win Rate of the Best-of-n Policy
Lemma 5.1. The win rate of any policy a. against reference policy Tef could be expressed as
So far, we provided theoretical guarantees on the KL diver- gence of the best-of-n policy with respect to the reference policy. In this section, we extend our study to characterize the win rate of the best-of-n policy against the reference policy. Let win of y against z in context æ be defined as:
([ = Eynrr(la) [Ca.n (x, y)]
(31)
Notice that the above result suggests that the win rate of any policy only depends on reward and the reference policy through the calibrated reward function, Cz,, y). Hence, this notion of calibration may be used as a canonical trans- formation of the reward for preference optimization against
1 (r(ær, y)=r(a, 2))
wr w =|=)= = (r(xx, )x) +))+
2
6

---

Theoretical guarantees on the best-of-n alignment policy
given reference policy. In fact, this transformation is theo
Given this upper bound, the win rate of best-of-n would be fairly close to - if Tee is 6-bound model, and is suff-
retically proposed as the objective in IPO (Azar al., 2024)
and is the key to best-of-n distillation (Gui et al., 2024;
ciently small such that na << 1, by combining Theorem 5.3
Amini et al., 2025; Sessa et al., 2025; Yang et al., 2024b)
and Corollary 5.5, we get
and inference-aware alignment (Balashankar et al., 2025).
( ( (+22) [| =
π
(36)
Lemma 5.2. The win rate of best-of-n policy against The is given by
72 n
which is the claim of Gui et al. (2024, Theorem 2) for the win rate of best-of-n
(T( ((m))((T((11122)
= =E Fmn FF - (x,
5.2. Lower Bounds on the Win Rate Gap
2/
Our next result characterizes the cases where the gap bounded away from O.
Proof. This is proved by plugging Lemma 2.3 into Lemma 5.1.
Theorem 5.6. Let Exx: > 0 be defined in Equation (20). Then,
Our next result is an upper bound on the win rate of best- of-n policy. Intuitively, observe that the win rate could be estimated by drawing (n + 1) samples from Te,, associat- ing to the best-of-n model and the remaining one to the reference model. Hence, the best-of-n model gets n-to-1 chances of winning against the reference model, unless there a draw due to samples with the same reward value, Thus, intuitively (a (n) (- (az) [[=))((|(z))
(m) GW
(x)
> n2
n + 1 (1 (1 =t0)
- (1 - (1 - Exx) Exo) (1
Eoc
> 0.
2
Corollary 5.7. As + 00, we have
G (x) ス & (1 + on())
(37)
2
We formalize this as an upper bound on the win rate.
n n+I.
As + 00, G (z)
+ Eze 2
which is bounded from 0.
Theorem 5.3. For all, n, and all æ. the win rate of ofst-offm policy is upper bounded by
6. Rewind-anddRRpent Rejection Sampling Beyond Best-of-n
(1
n S
(32)
The best-of-n policy is a form of rejection sampling. An- other form is called rewind-and-repeal where the process of generating a response and scoring is repeated until cer- tain threshold on reward is met (Kim al., 2025). A more involved blockwise variant of this process is recently used by Li et al. (2024). Formally, let æ be given input prompt to the model, and let { J be sequence of infinite i.i.d. samples drawn from [ -2r). Then, rewind-and -eapee accepts yar such that
n 1
In the rest of this section, we derive bounds on the gap between this upper bound and the actual win rate:
GW
x
n n n+1
(I 2 0. (33)
5.1. Upper Bounds on the Win Rate Gap
Unlike the KI. divergence, is clear that cf (a) could not grow unbounded, and is upper bounded by
r(æ, ym) 2 and Ak M r( yk)) < क (38)
where & R the threshold on reward. In other words, ym is the first draw whose reward reaches a certain threshold t. We also call M the (random) number of trials until the threshold is met, which determines the cost of inference from the model. We denote the resulting policy by π
Theorem 5.4. The win rate gap is upper bounded by
Glyl (xx) S 71 1 -Ha(salle)
(34)
2
where H2(.) denotes the Rényi entropy of order 2 defined in Equation (16).
It is natural to ask: how do the win rate vs KL tradeoffs of rewind-ane anddd-and-raa and-repeat compare with that of best-of-n? To answer this question, first define
Corollary 5.5. Let )nd(y)a)) 5 for all y y", Le., The is 8-bound. Then,
we (x) Ey x) [1(( y) 2 Φ)
(39)
< n1 1-1
as the probability of drawing sample from the reference policy that meets the threshold. Hence, the expected number of trials to output an outcome is E[M] = 1/w (z).
w (x)
8.
(35)
2
Proof. The proof follows by noticing that Ho(T))) log(1/6) and invoking Theorem 5.4.
Z
LI
Next, let us derive the PMF of rewind-and-repeal policy.
7

---

Theoretical guarantees on the best-of-n alignment policy
Lemma 6.1. The probability mass function (PMF) of the policy is given by
0.75
Borly a+) (m)
ifr v) 2 +b ifr(æ, v) < +
TT (y)aa)
(40)
0.65
0.60
The proofs are deferred Appendix A.4. Given its PMF, next we derive KL divergence and win rate of the rewind- and-repeat policy and the reference policy.
0.55
login) (Eq (4))
Proposed estimator 1Eg, (25B)
Exect KL (Ea, (5))
Lemma 6.2. We have
0.50
0.00
0.25
0.50
0.75
1:00
1.25
1:50
1.75
2.00
(Tn((To ( (-[z))][ = log we (a:)
(41)
Figure 6. The win rate against reference policy vs KL divergence tradcoff curve for the best-of-n policy. We have used the analytical formula (log(n) - (n 1)/n), Equation (4), the exact KL diver- gence, Equation (5). and the proposed estimator. Equation (25), for producing the tradeoffss from Example 1, illustrating a case where the actual win rate vs KL divergence tradeoff' curve for the best-of-n policy is more favorable than the one predicted using the upper bound formula, log(n) (n 1)/n.
(1))) (2)) = wi (2)).
(42)
Thus, by sweeping & in R. we will effectively sweep We in [0, 1]. and obtain the respective win rate vs KL divergence tradeoff for the rewind-and-repeat policy. Note that the KL divergence was recently derived by Kim al. (2025, Appendix A.5) and we provide proof for completeness.
win rate vs KL divergence tradeoff' is more favorable than that portrayed by using the formula in Equation (4). In this example, the best-of-n policy in the limit of large n, reaches KI. divergence of log(2) and a win rate of 0.75.
So far, we derived characterization of the KL divergence of the rewind-and-repeat policy. However, when the number of outcomes of model is large, similarly to the case of best-of-n, estimating the KL divergence is intractable.
Definition 7.1. Let W : B+ + [0, 1] be a function that takes in D 2 O and outputs WDD.. Let Dn (.1=)) Trer (ere(((ar)) (az)). We say that W is upper bound on the tradeoff curve of best-of-n if for all n, S (D..... Alternatively, we say
Our main result in this section is an unbiased estimator of the KI. divergence of the rewind-and-repean procedure.
Theorem 6.3. For 2 1, let Hn := Σi=1 be the n-th Harmonic number and Ho = O. Further, let M be the number of trials to achieve an outcome in the rewind-and- repeat policy. Then,
(,(x(") (1 that W is a lower bound on the tradeoff curve of best-of-n if for all n, (x (+12) (- 2 WDD...
(T< (+1x) (1
E[HM-1].
(43)
We can turn our existing bounds into straightforward lower and upper bounds the tradeoff curve for best-of-n. (see Appendix A.5). We conjecture tighter upper bound.
Hence, we propose M- as an (unbiased) estimator for the KL divergence of the rewind-and-repeat procedure with respect to the reference policy.
Conjecture 7.2. For any x and given Teg. let the function Wo(D) = (D) where for all T E [0.5, 1),
7. Win rate vs KL Divergence Tradeoffs
((T) = log 1-7
Thus far, we characterized the KL divergence and win rate of best-of-n. In practice, it is customary to compare differ- ent alignment methods based on their win rate at certain KI. divergence from the reference policy. Note that Theo- rem 3.1 implies that the win rate (or expected reward) vs KL tradeoffs reported in the literature that use the analytical formula in Equation (4) (Gao et al., 2023; Go et al., 2023; Mudgal al., 2024; Scheurer al., 2023) are conservative and the actual tradeoff curve of the best-of-n policy is in fact guaranteed to be no worse than what is reported. To further substantiate this point, let us revisit Example 1 and report the win rate vs KL divergence tradeoff curve (Figure 6), where we used the actual win rate in all cases.* 4 The actual
Wo is an upper bound on the tradeoff curve of best-of-n.
Example 2. We consider ternary language model with alphabet X' = {0, 1, 2). ordered from least preferred to most preferred, with probabilities given by गee = (0.3, 0.6, 0.1). Hence, the calibrated reward in Equation (30) is given by (0.3, 0.75, 0.95). The set of solutions to the KL-regularized RL problem are given by Equation (2). We also compute the best-of-n solutions (for continuous n.) using Lemma 2.3 and and-reaa using Lemma 6. Before discussing the win rate vs KL divergence tradeoffs, we first visualize the set of solutions on the probability simplex in Figure 7 and observe that the solutions could be very different. When we consider the win rate vs KL tradcoffs for this example (Figure 8), we observe that the tradeoffs are strikingly close
'In practice the win rate could be estimated using the unbiased estimator given by the win random variable in Equation (28).
8

---

Theoretical guarantees on the best-of-n alignment policy
10
6h 01---e spirst parco
0.8
bestefr upper scund [Cenjssture 7.2)
bec--the the--hhn upper bound [Canjectuuee 7:31 spcinall
(0.0.1)
11.0.0)
1.0
1.3
2.
Figure 7. The set of solutions for Ex- ample win rate optimal achieved by varying 9 in Equation (2), best-of- n is given by Lemma 2.3, and rewind- and-repeat is given by Lemma 6.1.
Figure 8. Win rate vs KL tradeoff for Exam- ple 2. The tradeoff curve of Best-of-n is close to win rate optimal, and both are better than rewind-and-rrpeat.
Figure Win rate vs KL tradeoff for Exam- ple 3. The tradeoff curve of Best-of-n is close to win rate optimal and the limit behavior, and both are better than rewind-and-repeat.
for best-of-n and win rate optimal models, matching recent findings of Yang et al. (2024a); Gui et al. (2024). Example 3. We consider a language model with low prob- ability outcomes, and calibrated reward, and plot win rate vs KL divergence tradeoff's in Figure 9, and observe that in this asymptotic regime where we observe that the trade- off curve offered by best-of-n is almost optimal and offers better tradeoffs compared to rewind-and-repeat.
offs on reward vs KL divergence (where KL divergence captures preservation of the core model capabílities) com- pared to rewind-and-repeat, the latter is more effective in driving reward up for a given compute budget which is im- portant in test-time compute scaling. Hence, it remains to be seen how to best design a method that achieves Pareto opti- mal tradeoffs between compute, reward, and KL divergence (which captures preservation of capabilities other than what is captured by reward). This might involve combining the future work.
8. Conclusion
with best-of-n and could be an area for
We studied the best-of-n alignment policy and derived its probability mass function (Lemma 2.3). We proved that an analytical formula used in the literature for the KL di- vergence of the best-of-n policy with the reference policy, log(n)- - (n- 1)/n in Equation (4), is false, and only an upper bound on the KL divergence (Theorem 3.1). We de- rived bounds on the gap between this formula and the KL divergence where we roughly showed the following: Let y be a draw from the best-of-n policy. Let E,, be the proba- bility mass of y under the base model. Then, if nen < 1, the gap between the formula in Equation (4) and the exact KL divergence is small (Theorem 3.4); and if ne, >> 1, the gap between the two may be large and unbounded (Theo- rem 3.6). We proposed new estimator for the KL diver- gence (Definition 4.3), which we demonstrated to capture the behavior of the KL divergence on several numerical ex- periments. We showed that the win rate of best-of-n against the reference policy is upper bounded by n/(n + 1) (The- orem 5.3). Similarly to the KL divergence, we provided upper (Theorem 5.4) and lower (Theorem 5.6) bounds on the gap between the actual win rate and the bound. We compared best-of-n with another form of rejection sampling and showed its superiority both
Impact Statement
This paper presents theoretical investigations of best-of-n sampling and other test-time rejection sampling algorithms, which is simple yet effective method for test-time align- ment of generative models. One of the major findings of this paper is that a widely used formula for KL divergence of the best-of-n policy and the reference policy is a theo- retical upper bound on the actual KI. divergence. This may help ensure that the capabilities of the reference model are largely preserved in the aligned model. For example, this may help compliance or risk- management teams preserve safety by guaranteeing that the policy that is served does not drift too far from a safe reference policy.
Our work also showed that best-of-n is an effective (and almost optimal) test-time alignment method which comes with theoretical guarantees on win rate vs KL divergence tradeoffs motivating its use for improving the capabilities and safety of models. On the flip side, this also shows why best-of-n with an adversarial reward may be used to effectively jailbreak generative models. We hope future work can benefit from our findings in making best-of-n more effective, and can also devise safeguards against best- of-n jailbreaks.
through pewind-and-reppe theoretically and empirically. We also extended the bounds to blockwise best-of-n (Mudgal et al., 2024).
While our results showed that best-of-n offers better trade-
9

---

Theoretical guarantees on the best-of-n alignment policy
References
language models a practical size. arXiv preprint 2rXiv:2408.00// 18, 2024.
Amini, A., Vieira, T., and Cotterell, R. Variational best- of-n alignment. International Conference on Learning Representations (ICLR), 2025.
Go, D., Korbak, T., Kruszewski, G., Rozen, J., and Dymet- man, M. Compositional preference models for aligning LMs. arXiv preprini arXiv:2310. 13011, 2023.
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.
Gui, L.,, Gårbacea, C., and Veitch, BoNBoN alignment for large language models and the sweetness of best- of-n sampling. Neural Information Processing Systems (NeurIPS), December 2024.
Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., Das- Sarma, N., Drain, D., Fort, S., Ganguli, D., Henighan, T., et al. Training helpful and harmless assistant with rein- forcement learning from human feedback. arXiv preprint arXiv:2204.05882. 2022.
Hilton, J. and Gao, L. Measuring Goodhart's law, April 2022. URL https: /openai..... com/ resea rch
measuring-goodhrrtes -lw. Accessed: 2024-01- 03.
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.03555 2024.
Balashankar, A., Sun, 7.,, 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 aligmment. International Conference on Machine Learning (ICML), 2025.
Kim, M., Thonet, T., Rozen, J., Lee, H., Jung, K., and Dymetman, M. Guaranteed generation from large lan- guage models. International Conference on Learning Representations (ICLR), 2025.
Beetham, J., Chakraborty, S., Wang, M., Huang, F., Bedi, A. S., and Shah, M. Liar: Leveraging alignment (best-of-n) to jailbreak Ilms in seconds. arXiv preprint arXiv:2412.05232, 2024.
Korbak, T., Elsahar, H., Kruszewski, G., and Dymetman, M. On reinforcement learning and distribution match- ing for fine-tuning language models with no catastrophic forgetting. Advances in Neural Information Processing Systems, 35: 35:16203-16220 2022a.
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 ar arXiv:2407.2178 2024.
Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information pro- cessing systems, 30, 2017.
Korbak, T., Perez, E., and Buckley, C. RL with KL penalties is better viewed Bayesian inference. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 1083-1091, 2022b.
Coste, T., Anwar, U., Kirk, R., and Krueger, D. Reward model ensembles help mitigate overoptimization. Interna- tional Conference on Learning Representations (ICLR), 2024.
Li, T., Beirami, A., Sanjabi, M., and Smith, V. On tilted losses machine learning: Theory and applications. Journal of Machine Learning Research, 24(142):1-79, 2023.
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.
Li Y., Wei, F., Zhao, J., Zhang, C., and Zhang, H. Rain: Your language models can align themselves without fine- tuning. In The Twelfth International Conference on Learn- ing Representations, 2024.
Mroueh, Y. Information theoretic guarantees for policy alignment in large language models. arXiv preprint 005iv::4000088880 2024.
Gao, I.., Schulman, J., and Hilton, Scaling laws for reward model overoptimization. International Conference on Machine Learning, pp. 10865--0066.... PMLR, 2023.
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 decod- ing from language models. International Conference on Machine Learning (ICML), 2024.
Gemma, Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahri- ari, B., Ramé, A., et al. Gemma 2: Improving open
10

---

Theoretical guarantees the best-of-n alignment policy
Nakano, R., Hilton, J., Balaji, S., Wu, J., Ouyang, L.,, Kim, C., Hesse, C., Jain, S., Kosaraju, V., Saunders, W., al. WebGPT: Browser-assisted question-answeriin with hu- man feedback. arXiv preprint arXiv:21 2021.
Yang, Q., Salamatian, S., Sun, Z., Suresh, A. T., and Beirami, A. Asymptotics of language model alignment. International Symposium Information Theory (ISIT), July 2024a.
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 fol- low instructions with human feedback. arXiv preprint arXiv:2203.02155. 2022.
Yang, K. and Klein, D. FUDGE: Controlled text gen- eration with future discriminators. In Proceedings of the 2021 Conference the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 351 1-3535, Online, June 2021. Association for Computational Linguistics. doi: 10. 10.18633w122022.maddmmmiii a aclanthologys www.:/llooooooog..ooooooo wtwp:////....aaaa.a............a www.nnaaccmaaaa 276.
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.
URL https://
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 IIm alignment. arXiv preprint arXh::2410200721 2024b.
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. / /8290, 2023.
Zhao, Y., Khalman, M., Joshi, R., Narayan, S., Saleh, M.. and Liu, P. Calibrating sequence likelihood improves conditional language generation. The Eleventh Inter- national Conference on Learning Representations, 2022.
Scheurer. J., Campos, J. A., Korbak, T., Chan, J. S., Chen, A., Cho, K., and Perez, E. Training language mod- els with language feedback at scale. arXiv preprint arXiv:2303.16755 2023.
Sessa, P. G., Dadashi, R., Hussenot, L., Ferret, J., Vieil- lard, N., Ramé, A., Shariari, B., Perrin, S., Friesen, A., Cideron, G., et al. Bond: Aligning IIms with best-of-n distillation. International Conference on Learning Repre- sentations (ICLR), 2025.
Snell, C., Lee, J., Xu, K., and Kumar, A. Scaling Ilm test- time compute optimally can be more effective than seal- ing model parameters. arXiv preprint arXiv:2408.03314 2024.
Stiennon, N., Ouyang, L., Wu, J., Ziegler, D., Lowe, R., Voss, C., Radford, A., Amodei, D., and Christiano, P. F. Learning summarize with human feedback. Ad- vances in Neural Informatio Processing Systems, 33: 3008-3021, 2020.
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 ehirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Stanford alpaca: An instruction-follooinng llama model, 2023.
Touvron, H., Martin, I.., Stone, K., Albert, P., Almahairi, A., Babaci, 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.

---

Theoretical guarantees the best-of-n alignment policy
A. Proofs of the main results of the paper
A.I. Proofs of Section
We provide the proofs of the main results of the paper. To do so, we need to state two further auxiliary lemmas.
Lemma A.I. For any 0 S < b S 1, and n E N,
b" a"
{
(b
S"' nee log (nw*=1) dv,
(6" - aa) log
nw"11 log (nz= du gn(a, b) &
(44)
b
where g(a, b) 2 0, and is given by
n(b a) In(a, = Dxa = (b" a"" h". a² + (n - 1)(6" log b a² log a) n
(b" - a²), (45)
where S w S b, we define Po(w)=
hm-a" and Pu(w) b-a-
Proof. Notice that
(PelPe))pw) = f(a n Tant-
b n=
non-1(b a) 6"
6 a log du
log
de 2
(46)
b a
a
The inequality is established by algebraic manipulation. The gap In (a, b) is obtained by following Lemma A.2 to derive the right-hand-side of the inequality in closed form.
Lemma A.2. The following identity holds:
[ no log (nw"-1) du (D" aa") + (n 1) (b" logb a" oologa) n1 n-1 (b" a²).
(47)
n
Proof. The proof is completed by noticing that (v" log v) no= log + un
Proof of Theorem 3.1.
Fr (yla)"" F (y)zz)"
(= (-)((=)) = Σ
(F
FT (3)2z)" - -F (z)(z) log
(48)
Tee (t)
DEV-
(ule))
Σ
1w0"1 (nw==1) du
(49)
BEy-
(ul()))
no
log
11) du
(50)
1
log(n)
(51)
n
where Equation (49) follows from Lemma A.1, and Equation (51) follows from Lemma A.2.
Lemma A.3. For any b S and n e N,
gn(a, b) s 2n(n l) - a)2.
(52)
12

---

Theoretical guarantees on the best-of-n alignment policy
Proof. Recall from Equation (45) that where (w) nw"- and Pu (w: 8"-a"
In (a, b) = =6" a" DKL
(53)
I for a S w S b. We can further bound the KL divergence as b-a
DKL (pL(Pllpu)) < max log (w)
(54)
wE(a.b] Pa (w)
max log nw" (b a) we((a,)
(55)
bm a
nb² 1 (b a)
loog b a
(56)
nb"
log n- b2-1-jod j=0
(57)
At this point, we will divide the proof into two cases depending on the value of a, b.
Case We first focus on the case when a < b//2. this case,
non
(pa(Pellpu) S log
(58)
2j ba-1-ja
a
S logn. n.
(59)
Hence,
bn S (b" a²) logn
(60)
< 6" log n
(61)
b logn
(62)
S 4(b a)² log n
(63)
S 2n(n 1) - -
(64)
Case II. For 2 b/2, notice that
nbb
(KL(PollPu) S log
(65)
25= b 1 0
n-
1
ba-1
S log
(66)
8
1)/2 (n- 13/2
b
log
(67)
a
(b a)
< 2
(68)
a
S (n 1) (b
(69)
where the first inequality follows from AM-GM inequality. Hence,
In (a, b) s (b" - -a") (n) 1) bb - a)
(70)
b
(n 1) (b - a)
= (b a)
Σ
(71)
1=0
5 nb 1 1(b - an - 1) (b ②)
(72)
b
2 a) 1)(b a)
(73)
< 2n(n 1)(b a)².
(74)
13

---

Theoretical guarantees the best-of-n alignment policy
The proof is completed by putting together Case and Case II.
Proof of Theorem 3.4.
log(n) 12 12-1
(T =
Σ
In And
(y) , Frnd (y)(()))))
(75)
WEY-
S 2n(n 1) Σ (3)(2) - FF (3)(2))
(76)
VE
2n(n - 1) Σ (Tee)())())
(77)
VE
2n(n 1)e (2(malm)
(78)
where Equation (76) follows from Lemma A.3.
Proof of Theorem 3.6. Notice that the gap least lower bounded by the value of the gap for the highest reward outcome. Hence,
log(n) n
DKL (1))) 2 In Ecos 1),
(79)
where (n(:,)) is defined in Equation (45), and is given by
n2 (1 - (1 1 -xx)))) (80)
gn(1 - 6xc, 1) = (1 (1 - )²) log T- (1 )x0)
(n - 1)(1
- -600 log(1
completing the proof.
Proof of Corollary 3.7. The proof is completed by the following inequalities:
(1 Eoo: 1) = (1 - (1 60o (o log 1
no
(n- - 1)(1 6co) log(1 600)
n- 1 (1 (1 )m)
(81)
(1 Exc )"
n2 1
2 (1 nEr))
log
(n 1) 1 (1 E0o )m
log E0o
(82)
(1 )
n
000 [ 0 )
- 1
2 (1 ncw)
log
+ +(n 1)
(83)
(1 Eoo)
1
E0o
n
(1 me
log
Te
n2 n- 1
ne
600) )01) (1 (1-6x0)
(84)
(1 - E0o) )m
n
(1 Exo )
2 (1 new)
log(ne>x)
(85)
2 e nEW
log(nE -
(86)
Hence, as n - 00,
gn Soo, 1) log(ne.o) + on (log n),
(87)
which completes the proof.
14

---

Theoretical guarantees the best-of-n alignment policy
A.2. Proofs of Section
Proof of Lemma 4.
The proof follows from:
F (n) F (2)(z²)
DKI. (1ar) Pypnria)
log
(88)
Tel (3/ az)
Teat (3/
S Bynniin
log
(89)
Trefl y (3)
ग (00) "/lm"
Proof of Corollary 4.2. Recall that 6 (Yaux 2)) where Vmax
Notice that
S 1- 100"1 log (no**11) du + - (1- (1= 1 (1 Eoo
(90)
Eoo
which follows the same lines in the proof of Theorem 3.1 except that we singled out the highest reward outcome, and bounded the rest of the terms. The proof is completed by invoking Lemma A.2 to express the integral in closed form. 7
Proof of Lemma 4.5. Fist notice that for any 0 S E S 1,
- (1 - log 1 (1 E)R
( S 1-
n="11 log (noa-1) du,
(91)
which is implied by Lemma A.1 and setting 1 and a 1 E. The proof is completed by noticing that
DKL(en)= = for no non-1log (nw do + (1 - (1 - En) log (1 6n)
(92)
Ex
and
1- 1w 1w log (nw du + 1² no log (nw"-1) do.
Kln
(93)
1-en
15

---

Theoretical guarantees on the best-of-n alignment policy
A.3. Proofs of Section
Proof of Lemma 5.1. Recall the definition of Fxx and Fr hence
Eznmal()a) (y > x|xr]] = (ln ( y) > r(a, 2)) + 2 ()(x,))) r(æ,2))
1 2 (1)(z)))))) + F (1)x)))
(94)
which completes the proof.
Lemma A.4. For 0 S a asbsuls 1, ler
(a) == n
n+I (bu+1 (0atl_a"th) 1 - ab" a).
(95)
2 nn ,n( b) 20 0.
We have
(96)
Equivalently,
1 (6" ab)(b + a) S n+1
(bat1 (b"t1-a"th)
(97)
2
Proof. Let consider the function (a, v) for fixed a function of v. Given that (a a) 0, it suffices to show that hm (a, v) 2 0 for all 0 S as v S We have
① hn (a, 0) = no=
n n- (r + a)
1 (v" - a²))
(98)
Ov
2 (nw"~1( (v a) -2=" a"))
1
(99)
2 20
(100)
completing the proof.
Lemma A.S. The following is an upper bound on h (a, b)
(a, b) S
11
(b a)2.
(101)
2
Proof. We have
hn
n
(b a at1)
2 (6" (6"'-a") (bb)
1
(102)
n + 1
2
n-1
= (b a)
n
b'a n-
(a (a+ + b)
Σ
bia n-1-i
(103)
+ 1=0
1=0
1 2( + +b²) + n 1) Σ abb
(b a)
(aa
(104)
12 n++1
2
n - "
$ (b a)
(a" b²)
(105)
2(n
n +
(b a)
71
(6" )")
(106)
2 + 1)
$ (b 0)² (n 1) nn
(107)
2(n + 1)
2 (b a)
(108)
completing the proof.
16

---

Theoretical guarantees the best-of-n alignment policy
Proof of Theorem 5.3. Note that
(T [1 (=)))
= 2
Σ
(F - F (y)(x)") (F (3) + and (y)z))
(109)
NEV*
(F Man (y) Fr (2) (z)=x))
(110)
VEy
(111)
+ 1'
where Equation (109) follows from Lemma 5.2 and Equation (110) follows from Lemma A.4.
Proof of Theorem 5.4. We have
n
=
Σ
hn(F Fid ( Frr (an((y))) ())))
(112)
n + 1
NEY*
1
(F (y)2)) - Frim (2) )))
(113)
2
yEy-
11
1
Σ (y)e))
(114)
2
yEY*
n1
1
(2(aal)))
(115)
2
where Equation (113) follows from Lemma A.5.
Proof of Theorem 5.6. Recall Lemma A.4. Therefore,
GW (x)) 2 (1 Exa, 1) 1)20 0.
(116)
The proof is completed by noticing from Equation (95) that
n
(1 - (1 )oth) - -11-1 -000)))))) (1- E0o
hn - Eoo, 1) n +
(117)
Proof of Corollary 5.7. As - 00.
n + - (1 - (oxt)) (1 (1 - Ex0 (1-
n
-20
Ex:
6.0
(1 + 0n (1)),
(118)
2
2
which completes the proof.
17

---

Theoretical guarantees the best-of-n alignment policy
A.4. Proofs of Section
Proof of Lemma 6.1. The proof is completed by observing that with probability wx)x) good outcome is obtained, and with probability (1 wa (2)) the trial is repeated.
Proof of Lemma 6.2. First consider the KL divergence follows:
(
DKL (-)=)) [Ter) Trer( (x)) Ey (1 log
(119)
Trett y x
1
Ey (11r) log
(120)
ugg. (x)
log
(121)
wg (2))
Next, let us consider the win rate:
W. (a = (1 wx)(x)) x 1 + (xxx) 2
(122)
1
= 1 2 wip u (x),
(123)
completing the proof.
Proof of Theorem 6.3. Let (x) be the probability of selecting a good sequence. Then, we note that
log( (x))
DC
Σ
H& wi (x:)**
Wi (x)
where for k 2 H3 -E== 7 is the k-th Harmonic number and Ho = O. Hence,
(og (x))
H1 (xx) (w) (z)
Let M be the location of the first one (the number of trials). Then,
P(M = k) (1 (x)x))) te
Hence, the following is an unbiased estimator of - -log( (wo (=:)):
H
18

---

Theoretical guarantees the best-of-n alignment policy
A.5. Proofs of Section
Theorem A.6. For any and a given Wees let the function Wi (Di) be implicitly defined for 1:
Wz(T) T+1 2
7 -
(2
1
Dz(T) = Dr(T)=log(T)
WL is a lower bound on the tradeoff curve of best-of-n.
Proof. This is obtained by putting together Theorem 3.1 and Theorem 5.4.
Theorem A.7. For any and given Wef let the function (D(D)) be implicitly defined for T 2 1:
WT(()= T+1²
DU(T) (log()) 1og(T))---
---
2T(T 1)e (a.ll)))))
We an upper bound the tradeoff curve of best-of-n.
Proof. This is obtained by combining Theorem 5.3 and Theorem 3.4.
19

---

Theoretical guarantees on the best-of-n alignment policy
B. KL divergence of blockwise best-of-n
While best-of-n could be viewed as a sequence-level rejection sampling, most generative models perform decoding step by step. For example, language models perform generation token- ookenbby-token or diffusion models perform generation through several denoising steps. Best-of-n could be extended to exert control at different levels of granularity. In particular, best-of-n has been extended to provide control through blockwise decoding (Mudgal al., 2024), where block of length B is decoded times and one with highest token-level reward is selected, and decoding is continued as such. Blockwise decoding could also be viewed as a simple form of tree search and also beam search. Our main result in this section characterizes the KL diveregnce of blockwise best-of-n with respect to the reference policy.
For the purpose of this section, we assume that the fully decoded response y == (y1, ,.3r) consist of ' intermediate steps (which are tokens the context of language model and denoising steps in the context of diffusion). For the purposes of this section, we focus on the presentation using generative language models. We let yr EOS be a special token that determines the end of sequence. The autoregressive decoding could be further explained as follows. In the first decoding step, the model иT assigns probability distribution over the first token given by (-1(2)), and one token y1 is drawn from this distribution. Next, the second token, is drawn from T-| V1), and so on. In short, in each step token is drawn from πT (-|ar, 1) where y' = (y)....,,M) 3r). Note that through the chain rule, any language model could be equivalently expressed as a next-token predictor.
Blockwise best-of-n decoding rule was recently proposed by Mudgal al. (2024). Here, the decoder samples n blocks of length B tokens from the reference model and accepts one with highest reward. Formally, if the decoder has already decoded y and the context is x, then
B == arg max r(a,,', x(k))
(124)
where a is the k-th continuation of length B drawn independently from the reference model:
{ (= (k
iid.
Tee((-/æ, (2),)))
(125)
and decoding continues until ECS is reached. We call the resulting distribution T (n) the sequence level best-of-n distribution, (T)).
Note that when B - 0o, aB
(n) becomes
Here, we assume we have access to a reward model can produce scalar values for r(x, for any partially decoded sequence y', which extends the definition of a reward model to also score partial sequences addition to fully decoded sequences. Mudgal al. (2024) argue that for this extension be meaningful the extended function should encode the value function for the sequence-level reward model. However, for the purposes of bounding the KI. divergence, don't care how the token-level reward function is obtained.
In the blockwise best-of-n decoding, control is applied more frequently than the sequence-level best-of-n, and the blockwise decoding enables to effectively score an exponentially large number of fully decoded sequences (similarly to how beam search works). Thus, intuitively the effective n would be exponentially larger and the KL divergence would be expected to grow linearly with the number of times control applied, i.e., by) where |y] the length of the decoded sequence and B the block length. Next, we formalize this intuition by recalling Theorem B.I.
Theorem B.1. Let a decoder, T perform block-wise best-of-n n with blocks of length B steps. Further, let y be draw from this decoder context a: such that by) the length of y, i.e., the total number of decoding steps. Then,
DKL(T (1)) (-1x) (x))) S E By
(m) (-1x)
y B
KLn,
where [.] is the ceiling operator (smallest larger integer), and KLn is defined in Equation (4).
Proof. Let us consider all possible outcomes for decoding a block of length B. The outcomes either finish and we reach EOS withing the block. We call this event E = 1, or leads to a partially decoded sequence = BB and we call this outcome
20

---

Theoretical guarantees on the best-of-n alignment policy
3 = 0. In this case, we write y (= 8).
(n) n
(n) DRL (mT (1(x) Terl [z)) =B ywann (-)o))))))) (-100
(y x+) Tea ( a:)
log
(126)
( ax) [y x))
(m))
"B (y ac)
(x))
E ywanh (+1x)
(11))
1(E 1) log B
E
1(E 0) log
yuunh
Te y x))
(127)
(m) 13
(3) ())
E
yrv B (110)
(n) (-10)
1(8 10g
Teerl y (z)
(n) (2
(n) (sl ar, =B
+ E. 2B w
(m0) (112)
E faw ( 91
11E 1(E=0)
log
"B
Txat (2 (2110) 2x)
log
"B
Teat ( I 2 nn)
(128)
(m) xz)
(m) (2 xx)
E ywu (n) (110)
1(E 1og
AB
E 30 gnn
1(E - 0) log
B
Tirr y a+)
Trat (= B (x)
TB (1) (8) lx, 2B8)
+ E 2l Aawgg (o (120)
E Bowafal (n) (1120,29)
1(E 1(E==))me -0gog
(129)
Teeel @ ar, 2 BB)
<KLn 1B (o) (-12)
1(E =1E==) -0) (1)2,,)))),)))))0 [| (-rret--aee =B1)
(130)
[||
<Knn Ba (n) (12))
1 0) SMa (m) (1)
Kl.,.
(131)
B
y
=E
y (n)
KL
(132)
where Equation (131) follows from recursively applying the same procedure to the subsequent blocks.
This theorem immediately suggests that D KL could be used as an estimator for the KL divergence of block-wise
best-of-n, and in expectation the estimator provides an upper bound on the KL divergence of blockwise best-of-n and the reference model. Sequence-level best-of-n could be viewed as blockwise best-of-n with -> oc, and Theorem B. simplifies to Theorem 3.1 in this asymptotic regime.
21

---

Theoretical guarantees on the best-of-n alignment policy
C. Experiments
C.1. Details of experiments on Alpaca dataset
We select the following four prompts from the Alpaca dataset (Taori et al., 2023).
P1 Transform the following sentence into yes/no question. It is going to rain tomorrow.
P2 Describe the function of a computer motherboard
P3 What is the capital of France?
P4 Give three tips for staying healthy.
We plot results based on Gemma 2 9B parameter instruction tuned model (Gemma et al., 2024) with temperature one. We further modify the prompts as per the instructions in ht httpp:/ //ai. google dev/ / /gemma/ does/ formatting, We use log-likelihood of the reference model as the reward. To get the better upper bound we use Corollary 4.2, where we bound Coo using 100 samples.
C.2. Computation of tilted min/max
We compute the tilted average of the estimator defined by (Li et al., 2023, Equation (2)):
log E ())l)) [
(133)
Note that Equation (133) recovers ( (dm (sn) for t - - 0o, and maxy~-(()((1e [c.)) for -> co (Li et al., 2023). To capture the variance, we call the value for & 1 the tilted min, and the value for t the rilted max. The tilted min/max capture the low/high quantiles of the value of the estimator and their difference portrays the deviation from the mean.
C.3. Experiments with machine translation prompts
In this section, we demonstrate the value of the new estimator using machine translation prompts.
P1 Translate the next sentences to German. I want to buy bread.
P2 Translate the next sentences to French. want to buy bread.
P3 Translate the next sentences to German. A simple and effective method for the inference- time alignment and scaling test-time compute of generative models is the best-of-n policy, where samples are drawn from reference policy, ranked based reward function, and the highest ranking one is selected. A commonly used analytical expression in the literature claims that the KI. divergence between the best-of-n policy and the reference policy is equal to log(n) (n - 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 KI. divergence and empirically show that provides a tight approximation. We also show that the win rate 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 tradeoff's between win rate and KL divergence of the best-of-n alignment policy, which demonstrate that very good tradeoffss are achievable with
1000.
P4 Translate the next sentences to French. A simple and effective method for the inference-time alignment and scaling test-time compute of generative models is the best-of-n policy, where samples are drawn from 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-m alignment policy, which demonstrate that very good tradeoff's are achievable with n < 1000.
22

---

Theoretical guarantees the best-of-n alignment policy
Translate the next sentences t... (E. 2 0.12)
Translate the next sentences t... (E. a 0.09)
De (best-of-n || [[d
2.5
Dq. bbest-of-n || Mead
2.5
2.0
2.0
1.5
1.5
1.0
1.0
Analytical formula (Eq. (4))
Analycical formula (Eq. (4))
Proposed estimator (Eq. (251)
Proposed estimator (Eq. (25))
Better upper bound (Coroliary 4.2, App.
Better upper bound (Corollary 4.2, App.
0.0
0.0
50
Translate the next sentences t... (e. 0.00)
Translate the next sentences t... (E. 0.00)
3.0
3.0
DxL (best-of-n Teed
2.5
Del (best-of-n Rted
2.5
2.0
2.0
1.5
1.5
1.0
1.0
Analytical formula (Eq,
Analytical fomula (Eq.
0.5
0.5
Proposed estimator (Eq. (25))
Proposed estimator (Eq. (251)
Better upper bound (Corollary 4.2, App. C)
Better upper bound (Corollaayy 4.2. App. C)
0.0
10
20
30
40
50
10
20
30
40
50
Figure 10. The analytical formula (log(n) (n 1)/n). Equation (4), better upper bound (Corollary 4.2, Appendix C), the proposed estimator, Equation (25), and the exact KL divergence, for four machine translation examples using Gemma 9B IT model (Gemma et al., 2024) with reward the log- ligeeiidd of response under the reference model.
Translate the next sentences t... (e. = 0.12)
Translate the next sentences t... (E. z 0.09)
Del (best-of-n || ntd
2.5
DeL (best- of-n || (eee))
2.5
2.0
2.0
1.5
1.5
1.0
1.0
0.5
0.5
Analytical formula (Eq. (4))
Analytical farmula (Eq, (4))
Proposed estimator (Eq, (25))
Proposed estimator (Eq, (25))
0.0
0.0
20
20
50
Translate the next sentences t... (s_ 0.00)
Translate the next sentences t... (E_ 0.00)
3.0
3.0
DEL (best-of-n Red
DeL (best-of-n ll Rad
2.0
2.0
1.5
1.0
1.0
0.5
0.5
Analytical formula (Eq. (4)
Analytical formula (Eq, (4))
Propesed estimator (Eq. (25))
Proposed estimator (Eq, (25))
0.0
0.0
10
20
30
40
50
10
20
30
40
50
Figure 11. The analytical formula (log(n) (n 1)/n), Equation (4), better upper bound (Corollary 4.2, Appendix C), the proposed estimator, Equation (25), and exact KI. divergence, for four machine translation examples using Gemma 9B model (Gemma al., 2024) with reward the negative length of response under the reference model.
23


Reviews:
[{"review_id": "RibwZ1KBAY", "reviewer": "Reviewer_vBEz", "summary": "The paper provides theoretical analyses for the widely used Best-of-N (BoN) policy, especially focusing on the KL divergence from the reference model and its win-rate. They first point out that the conventionally used formula for the KL divergence does not hold and actually gives only an upper bound. They evaluate the gap between the true KL divergence and the upper bound, showing the non-triviality of the gap under when N is sufficiently large (Section 3). They also propose an alternative estimator for the KL divergence, which is still a conjecture but experimental results indicate the plausibility of the conjecture (Section 4). Next, they analyze the win-rate of the BoN against the reference model, and provide an upper bound and its tightness evaluation (Section 5). Meanwhile, they also analyze the rewind-and-repeat model $\\pi_{\\Phi}$ with a threshold $\\Phi$, as a variant of BoN model. In such cases, both the win-rate and KL divergence can be exactly calculated (Section 6). Finally, they discuss on the conventional evaluation of the win-rate vs KL divergence tradeoff, and they conclude that the previous results are too pessimistic and can be improved with their results in Section 3 and Section 5 (Section 7).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful review and encouraging comments.\n\n> How the equation (34) can be derived?\n\nEq. (34) is derived by combining Theorem 5.3, i.e., Eq. (30), and Corollary 5.5, i.e., Eq. (33). We will clarify this.\n\n> The proposed estimator for the KL divergence has some evidences both theoretically and empirically, but the actual guarantee is still a conjecture. However, this is not a fault of the paper since it is explicitly stated in Conjecture 4.4.\n\nWhile we acknowledge this weakness of our work. We would like to mention two points: \n* We have tested this conjecture in tens of thousands of randomly generated numerical examples with varying support sizes and n and have not found a counter example.\n* If we replaced $\\epsilon_n$ with $\\epsilon_\\infty$, then the estimator would lead to an actual upper bound on the KL divergence (Corollary 4.2), which at least suggests that for large $n$ the estimator gives an upper bound.", "claims_and_evidence": "- Claim 1: The conventionally used formula for the KL divergence does not hold and actually gives only an upper bound, especially when N is sufficiently large.\n    - Theorem 3.1, 3.4 analyze the tightness of the formula, and Theorem 3.6 shows the looseness when N is sufficiently large.\n- Claim 2: The proposed estimator for the KL divergence might give a tighter upper bound than the conventionally used formula.\n    - empirically checked in Figure 2-4.\n    - However, the claim is theoretically still a conjecture.\n- Claim 3: They analyzed the win-rate and show that it can be approximated by $\\frac{N}{N+1}$ as expected in previous literature when N is sufficiently small.\n    - proved by Theorem 5.3, 5.4.\n    - Question: How the equation (34) can be derived?\n- Claim 4: In the case of the rewind-and-repeat model $\\pi_{\\Phi}$ with a threshold $\\Phi$, which can be seen as a variant of BoN model, both the KL divergence and win-rate can be explicitly evaluated in terms of the probability distribution of the reference model.\n    - provided by Lemma 6.2 and Theorem 6.3\n- Claim 5: The previous (empirical) results on the tradeoff curve between the win-rate and KL divergence are too pessimistic and can be improved by their results.\n    - theoretically obtained as corollaries of Theorem 3.1, 3.4, 5.3, 5.4. (Theorem A.6, A.7) \n    - empirically checked in Figure 7-8.\n\nOverall, their claims are solid and well-supported by theoretical and empirical results.", "methods_and_evaluation_criteria": "The proposed estimator for the KL divergence has some evidences both theoretically and empirically, but the actual guarantee is still a conjecture. However, this is not a fault of the paper since it is explicitly stated in Conjecture 4.4.", "theoretical_claims": "I roughly checked all proofs in Appendix.", "experimental_designs_or_analyses": "Their experiments seem to be well-designed.", "relation_to_broader_scientific_literature": "One of the key contributions of this paper is that it fixes the conventional misunderstanding on KL divergence between the BoN and reference models. The other analyses are also valuable to the community of inference-time alignment.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "See Claims and Evidences.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See Claims and Evidences.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zPflGnIrkI", "reviewer": "Reviewer_jbod", "summary": "This paper revisits the best-of-n alignment policy: given n samples from a reference language model, pick the sample that scores highest under a reward (alignment) function. A long-used formula in prior papers, $D_{kl}=log(n)-(n-1)/n$, has been cited as the KL divergence of the best-of-n policy from the reference. The authors show that although this expression frequently appears, it is only an upper bound—not an exact value—in many realistic scenarios. They derive exact or tighter bounds, provide a new practical KL estimator, and show that in some regimes (particularly if a small number of high-probability completions exist under the reference), the conventional formula can substantially overestimate the real KL drift.\n\nIn addition to these theoretical clarifications, they also: (1) Analyze blockwise best-of-n (where re-ranking occurs multiple times during generation). (2) Compare best-of-n to “rewind-and-repeat,” another rejection-sampling-like policy. (3) Show how their refined analysis can help practitioners accurately measure or cap KL divergence in alignment pipelines—important for settings where controlling distribution shift is crucial (e.g., compliance, RLHF with KL constraints).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful review and encouraging comments\n\n> The real data experiments are limited and do not show large-scale or broad tasks.\n\nThe goal of the experiments was to show cases where the could be a gap between the analytical formula and the true KL. We have now experimented with other prompts for machine translation and other rewards, namely output length, and still show that the KL estimates could be loose in more tasks. See https://anonymous.4open.science/r/bon\n\n> The variance of the proposed KL estimator in practical, large n sampling not deeply explored. Would you recommend running it repeatedly and averaging over many sequences, or do you envision an online, per-sample approach?\n\nThanks for this very important question. The proposed estimator is upper bounded by the analytical formula, and hence by log(n). It is also non-negative and lower bounded by 0. Hence, a crude upper bound on its standard deviation is $\\log(n)$. Thus, if the estimator is averaged out over $M = O(\\log(n) \\log\\frac{1}{\\delta})$ responses the standard deviation could be driven down to below $\\delta$. Given that we are generally interested in $n < 1000$, the dependence on $n$ is manageable. Having said that, given each of the M batches has $n$ iid samples (total of $M \\times n$ iid samples), we can use a bootstrapped estimator and should be able to remove the dependence on $\\log(n)$. We will think more about the variance estimation and include it in the next iteration.\n\nWe would like to also mention that the variance of the estimator in Figure 2 is exactly zero given all outcomes are the same likelihood. We have now also computed and plotted the standard deviation of the estimator in Figure 3 in the new PDF.  See https://anonymous.4open.science/r/bon.\n\n> The theorems assume finite support or at least that the model probabilities for outcomes can be well-defined in a somewhat discrete sense. Real LMs often have large vocabularies ...\n\nThe finite support assumption could be relaxed in practice by limit arguments given that PMF is bounded, and we will include this relaxation in the next version. We would like to also mention that subsequent to our work, Mroueh (2024) in fact has proved a variant of Theorem 3.1 under much weaker assumptions that even applies to continuous distributions such as diffusion process.\n\nMroueh, Youssef. \"Information theoretic guarantees for policy alignment in large language models.\" arXiv preprint arXiv:2406.05883 (2024).\n\n> Showing this in a broader study would underscore the conditions under which the new analysis provides a “significant” improvement. A concrete scenario (e.g. in compliance or enterprise alignment) where the old formula’s overestimation of KL meaningfully hampered performance? \n\nWe would like to emphasize that the message of this paper is not just about providing a new estimator for KL divergence. We also prove that the existing formula is an upper bound and hence the works that use it still give guarantees on the win rate vs KL divergence of best-of-n. \n\n> Because the paper is about “how far the policy can deviate,” it might be worth an explicit mention that tighter estimates can help compliance or risk-management teams ensure that the model does not drift too far from a safe baseline.\n\nThanks for the suggestion. Will explicitly mention it.\n\n> The authors’ theoretical perspective might help design an adaptive best-of-n procedure that halts once an approximate KL threshold is reached, or once a certain minimum reward is achieved.\n\nBest-of-infinity with halting when a certain threshold on reward is reached is akin to rewind-and-repeat which we analyze in Section 6, where we show that the resulting tradeoff between win rate and KL is less favorable compared to best-of-n. However, the resulting tradeoff between win rate and cost (number of decoding trials) is more favorable than best-of-n, which is important in test-time scaling. Hence, it remains to be seen how to best design a method that achieves Pareto optimal tradeoffs between compute, reward, and KL divergence (which captures preservation of capabilities other than what is captured by reward), which might involve combining the rewind-and-repeat with best-of-n for some finite n as the reviewer suggests. We will add a discussion to this end in the concluding remarks as an area for future work.  \n\n> Have you observed in practice (beyond the toy examples) that blockwise best-of-n greatly outperforms single-step best-of-n in reward–KL tradeoffs, or do you expect diminishing returns?\n\nBlockwise best-of-n has been studied comprehensively by (Mudgal et al., 2024). While the reward vs KL tradeoffs are generally not better than best-of-n (due to the fact that best-of-n is already almost optimal), blockwise best-of-n generally allows to achieve similar reward vs KL tradeoffs with ~10x smaller n, which means a 10x reduction in test-time compute. \n\nMudgal, Sidharth, et al. \"Controlled Decoding from Language Models.\" ICML, 2024.", "claims_and_evidence": "The paper’s key claims are relatively well substantiated by proofs and by numeric examples carefully chosen to highlight the regimes where the old formula fails.", "methods_and_evaluation_criteria": "The real data experiments are limited and do not show large-scale or broad tasks.\nThe variance of the proposed KL estimator in practical, large n sampling is not deeply explored.", "theoretical_claims": "The main theorems (e.g. Theorem 3.1, Theorem 3.4, Theorem 5.3) use standard tools from information theory (KL divergence definitions, integrals, combinatorial arguments) and appear consistent. The proofs in the appendices logically match the statements in the main text.\n\n**Potential concerns**: The theorems assume finite support or at least that the model probabilities for outcomes can be well-defined in a somewhat discrete sense. Real LMs often have large vocabularies, but the authors argue that for bounding or approximate computation, this finite-support assumption can be relaxed in practice.", "experimental_designs_or_analyses": "The chosen synthetic examples clearly demonstrate the paper’s main theoretical points.The real prompts confirm that skewed probability distributions are common in real queries.\n\n**Potential Weaknesses**:\nThe real data experiments are limited and do not show large-scale or broad tasks.\nThe variance of the proposed KL estimator in practical, large n sampling is not deeply explored.", "relation_to_broader_scientific_literature": "- **Alignment & KL Regularization**: The paper ties into RLHF and reward-based alignment (Christiano et al., Ouyang et al., etc.). Many works track or constrain $KL(\\pi||\\pi_{ref} )$ as a measure of “distribution drift.”\n\n- **Controlled Decoding**: Ties well to prior re-ranking or decoding approaches like beam search, blockwise best-of-n, or other re-ranking methods.\n\n- **Relation to Rejection Sampling**: They connect best-of-n and “rewind-and-repeat” as forms of rejection sampling, referencing literature on approximate generation, safe decoding, or iterative expansions.\n\n- **Practical Gains**: The new results clarify that in a “$\\delta$-bound” setting (where the reference rarely repeats outcomes) the old formula is indeed tight, but in the presence of a few large-prob outcomes, the real KL can be drastically smaller.", "essential_references_not_discussed": "The authors already cite many highly relevant alignment and preference-based optimization works (including those that have used or mentioned the log(n) formula for best-of-n). No glaring omission is evident.\n\nOverall, the bibliography covers standard RLHF, preference optimization, and controlled decoding references.", "other_strengths_and_weaknesses": "**Strengths**\n- Clarification of a widely cited formula: The old rule was occasionally treated as exact, which could be misleading in some alignment or compliance scenarios.\n- Practical new estimator: Offers a run-time method to measure “how much drift best-of-n actually used,” potentially letting practitioners adapt n on the fly or confirm they have not exceeded a KL budget.\n- Extensions: The coverage of blockwise decoding and comparisons to alternative rejection sampling (rewind-and-repeat) broadens applicability.\n\n**Weaknesses**\n- Empirical scope: The paper mostly uses small or contrived examples. Large-scale, high-n tasks in more varied domains could strengthen claims about real-world use.\n- Variance analysis: The new estimator’s variance or confidence intervals in complex distributions is left for future work.", "other_comments_or_suggestions": "- (1) While it is likely that in many real tasks with big n and small probabilities per token, the difference is small, showing this in a broader study would underscore the conditions under which the new analysis provides a “significant” improvement.\n- (2)  Because the paper is about “how far the policy can deviate,” it might be worth an explicit mention that tighter estimates can help compliance or risk-management teams ensure that the model does not drift too far from a safe baseline.\n- (3) The authors’ theoretical perspective might help design an adaptive best-of-n procedure that halts once an approximate KL threshold is reached, or once a certain minimum reward is achieved.", "questions_for_authors": "**Estimator Variance**\n- How large can the variance of the proposed KL estimator get in real usage (especially if $\\epsilon_n$  is extremely small)?\n- Would you recommend running it repeatedly and averaging over many sequences, or do you envision an online, per-sample approach?\n\n**Blockwise vs. Full-Sequence**\n- Have you observed in practice (beyond the toy examples) that blockwise best-of-n greatly outperforms single-step best-of-n in reward–KL tradeoffs, or do you expect diminishing returns?\n\n**Real-World Deployments**\n- Can you give a concrete scenario (e.g. in compliance or enterprise alignment) where the old formula’s overestimation of KL meaningfully hampered performance or forced overly conservative constraints?\nA real case study would clarify the direct practical payoff.\n\n**Infinite/Very Large Vocabularies**\n- Your proofs rely on finite support or a finite set of “possible outcomes” in each context. For large-vocabulary LMs, do you see a direct extension or an approximate argument? Would it still hold to treat large but finite vocab sizes similarly?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oRmVtV20yu", "reviewer": "Reviewer_3Tjm", "summary": "This paper discusses the best-of-n policy, a method used for inference-time alignment of generative models. The key idea is to draw n samples from a reference policy, rank them based on a reward function, and select the highest-ranking sample. The passage critiques a commonly used analytical expression in the literature, which claims that the KL divergence between the best-of-n policy and the reference policy is equal to log(n)−(n−1)/n. The authors disprove this claim, showing that it is instead an upper bound on the actual KL divergence. They also explore the tightness of this bound and propose a new estimator for the KL divergence that provides a tighter approximation.\nAdditionally, the authors analyze the win rate of the best-of-n policy against the reference policy, showing that it is upper bounded by n/(n+1) . They derive bounds on the tightness of this characterization and conclude by analyzing the tradeoffs between win rate and KL divergence. Their findings suggest that very good tradeoffs can be achieved with n<1000", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful review and encouraging comments.\n\n> This assumption simplifies the theoretical analysis; however, in practical applications, the output space of generative models is typically infinite (as in natural language generation tasks), which may limit the universality of the theory.\n\nWe would like to mention that (1) in practice, language models have a finite output tokens and will eventually stop hence practically the number of possible generated sequences is finite (even though it could be very large), which is covered by our theory; (2) this assumption could be relaxed using limit arguments, and we are happy to include a relaxation of it; (3) this result has already been extended to models with potentially continuous outputs (which would not even be countable) by Mroueh (2024) (subsequent to our work) so the extension beyond what is proved here is certainly possible.\n\nMroueh, Youssef. \"Information theoretic guarantees for policy alignment in large language models.\" arXiv preprint arXiv:2406.05883 (2024).\n\n\n> The experiment did not involve validation with real generative models, resulting in uncertainty regarding the practical application effectiveness of the theoretical results. In particular, the output space of real generative models is typically large and complex, necessitating further verification of the applicability of the theoretical results in these scenarios.\n\nWe would like to mention that we already showed a certain case with prompts from AlpacaEval dataset and Gemma2 IT 9B in Figure 4. We have now expanded the scope of experiments with more tasks and rewards in https://anonymous.4open.science/r/bon\n\n> This paper does not provide detailed guidance on the specific implementation and optimization of the best-of-n strategy in practical applications.\n\nWhile we agree with the reviewer that we do not provide an optimization of best-of-n, our results imply that with n~100-1000, win rate against base model will be already saturated without the need for excessively large n which also implies that KL divergence of ~5 is enough to reach a good policy. This implies that RLHF practitioners can aim to keep the KL divergence of their aligned models <10 and achieve good policies, which also significantly helps mitigate reward overoptimization.", "claims_and_evidence": "The claims made in the submission supported by clear and convincing evidence. Assumption 2.2: It is posited that the output space of the language model is finite. This assumption simplifies the theoretical analysis; however, in practical applications, the output space of generative models is typically infinite (as in natural language generation tasks), which may limit the universality of the theory.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make good sense for the problem or application at hand.", "theoretical_claims": "I have examined all the theoretical proofs to ensure their accuracy.", "experimental_designs_or_analyses": "The experiment did not involve validation with real generative models, resulting in uncertainty regarding the practical application effectiveness of the theoretical results. In particular, the output space of real generative models is typically large and complex, necessitating further verification of the applicability of the theoretical results in these scenarios.", "relation_to_broader_scientific_literature": "This paper refutes the commonly used expression for KL divergence in the literature, demonstrating that it merely serves as an upper bound to the actual KL divergence. Furthermore, a novel estimator for KL divergence is proposed.", "essential_references_not_discussed": "Most of the essential references have already been cited within this paper.", "other_strengths_and_weaknesses": "Strengths:\n1. This paper not only analyzes the KL divergence but also delves into the win rate of the best-of-n strategy, proving an upper bound of the win rate to be n/(n+1). Furthermore, the authors derive a trade-off relationship between the win rate and KL divergence, demonstrating how adjusting n can achieve a better balance between performance and model alignment in various scenarios.\n2. This paper introduces a novel estimator for KL divergence and validates its effectiveness through numerical experiments. This estimator provides a more accurate reflection of the KL divergence between the best-of-n strategy and the reference strategy.\n\nWeaknesses:\n1. This paper does not provide detailed guidance on the specific implementation and optimization of the best-of-n strategy in practical applications.", "other_comments_or_suggestions": "I do not have any other comments or suggestions.", "questions_for_authors": "I do not have any important questions for the authors.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t5PPKYSFYf", "reviewer": "Reviewer_nvEQ", "summary": "This paper provides a theoretical analysis of the best-of-$n$ policy $\\pi^{(n)}$, which is a simple inference-time method for aligning language models where $n$ samples are drawn from a reference policy $\\pi_{ref}$, the highest ranking one based on a reward function is selected.\n\nThe authors first disprove a commonly used analytical expression for $KL (\\pi^{(n)} || \\pi_{ref})$. They demonstrate that the formula $\\widetilde{KL}\\_n := \\log(n) - \\frac{n-1}{n}$, which is widely cited in the literature, is actually just an upper bound on the true KL divergence. They theoretically characterize when this bound is tight and when it can be arbitrarily loose.\nAdditionally, the paper develops a new estimator for $KL (\\pi^{(n)} || \\pi_{ref})$ that more accurately tracks the true value, as demonstrated by experiments.\nThe authors also analyze the win rate of $\\pi^{(n)}$ against $\\pi_{ref}$. This analysis shows that the win rate is upper bounded by $\\frac{n}{n+1}$ and characterizes when this bound is tight.\n\nThe authors also compare best-of-$n$ with another rejection sampling approach called rewind-and-repeat, ultimately showing the superiority of best-of-$n$ in terms of win rate vs. KL divergence tradeoffs.\nThe paper concludes that the actual win rate vs. KL divergence tradeoffs for best-of-$n$ are better than what has been reported in the literature when using the incorrect formula, and that very good tradeoffs are achievable with $n < 1000$.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your insightful review and encouraging comments.\n\n\n> My concern is with the experiment in Figure 4, where the authors test language models on benchmarks. The reward is set as the log-likelihood of the reference model, which already correlates with the policies, thus reducing generality. I suggest the authors provide more explanation about this choice, justify this approach, or discuss a more general reward function.\n\nWe have now experimented with other tasks and other rewards (e.g., generation length) to showcase that this is not just an issue that pertains to the log-likelihood reward. See https://anonymous.4open.science/r/bon.\nWe would like to reiterate that anytime there is a non-trivial chance of collision between the outcomes in a set of size $n$, the analytical formula for the KL divergence overestimates the true KL.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No, but they seems to be correct by intuition.", "experimental_designs_or_analyses": "I reviewed all experiments mentioned in the main text. My concern is with the experiment in Figure 4, where the authors test language models on benchmarks. The reward is set as the log-likelihood of the reference model, which already correlates with the policies, thus reducing generality. I suggest the authors provide more explanation about this choice, justify this approach, or discuss a more general reward function.", "relation_to_broader_scientific_literature": "This paper is significant to the field of language model alignment for several reasons:\n\n- It corrects a fundamental misunderstanding about the KL divergence of best-of-$n$. It also provides theoretical justification for why best-of-$n$ performs so well empirically, showing that its win rate vs. KL divergence tradeoffs are actually better than previously thought.\n- It offers theoretical bounds and a more accurate estimator for KL divergence that can be used to better evaluate alignment methods. These results can give hints on choosing the best $n$ in practice.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "### Strengths\n- **Theoretical guarantees:** The paper provides mathematical analysis, establishing bounds on both KL divergence and win rates. It corrects a misunderstanding in the literature by showing the formula $\\widetilde{KL}\\_n := \\log(n) - \\frac{n-1}{n}$ is only an upper bound for $KL (\\pi^{(n)} || \\pi_{ref})$.\n- **Practical estimator:** The proposed KL divergence estimator has practical value for researchers evaluating alignment methods, as it more accurately captures the true KL divergence, as demonstrated by experiments.\n- **Comparative analysis:** The comparison with rewind-and-repeat provides context on why best-of-$n$ performs well in practice through the lens of tradeoffs between KL divergence and win rates.\n\n### Weaknesses\n- **Limited empirical validation:** While the paper includes examples, it lacks extensive empirical validation on real language models and datasets. More comprehensive experiments would strengthen the claims. Additionally, in the experiment in Figure 4, the reward is highly correlated with the reference policy, which represents a constrained case.", "other_comments_or_suggestions": "No", "questions_for_authors": "Regarding the empirical validation on benchmarks, what will the results be if the reward model is no longer the log likelihood of $\\pi_{ref}$?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper considers finetuning alignment of generative models via best-of-n policy. In particular, the paper shed insights into the commonly-referred formula $\tilde{KL}_n$ for the regularization term just serves as an upper-bound of the actual value of it and provided a more accurate estimator. The paper also provided a new rejection sampling that outperforms best-of-n theoretically and empirically. The paper is easy to read. All reviewers are positive and the reviews are informative enough to recommend an acceptance. 

I recommend the authors to revise the paper based on the reviews and responses. In particular, additional experimental results with real data (conducted/promised to be conducted) during the author response will increase the value of this paper.
Decision: Accept (poster)