redacted<EMAIL_ADDRESS>
#  BOND: Aligning LLMs with Best-of-N Distillation

Pier Giuseppe Sessa  Google DeepMind  Robert Dadashi  Google DeepMind  Léonard
Hussenot  Google DeepMind  Johan Ferret  Google DeepMind  Nino Vieillard
Google DeepMind  Alexandre Ramé  Google DeepMind  Bobak Shariari  Google
DeepMind  Sarah Perrin  Google DeepMind  Abe Friesen  Google DeepMind
Geoffrey Cideron  Google DeepMind  Sertan Girgin  Google DeepMind  Piotr
Stanczyk  Google DeepMind  Andrea Michi  Google DeepMind  Danila Sinopalnikov
Google DeepMind  Sabela Ramos  Google DeepMind  Amélie Héliou  Google DeepMind
Aliaksei Severyn  Google DeepMind  Matt Hoffman  Google DeepMind  Nikola
Momchev  Google DeepMind  Olivier Bachem  Google DeepMind

###### Abstract

Reinforcement learning from human feedback (RLHF) is a key driver of quality
and safety in state-of-the-art large language models. Yet, a surprisingly
simple and strong inference-time strategy is Best-of-N sampling that selects
the best generation among $N$ candidates. In this paper, we propose Best-of-N
Distillation (BOND), a novel RLHF algorithm that seeks to emulate Best-of-N
but without its significant computational overhead at inference time.
Specifically, BOND is a distribution matching algorithm that forces the
distribution of generations from the policy to get closer to the Best-of-N
distribution. We use the Jeffreys divergence (a linear combination of forward
and backward $\mathrm{KL}$) to balance between mode-covering and mode-seeking
behavior, and derive an iterative formulation that utilizes a moving anchor
for efficiency. We demonstrate the effectiveness of our approach and several
design choices through experiments on abstractive summarization and Gemma
models. Aligning Gemma policies with BOND outperforms other RLHF algorithms by
improving results on several benchmarks.

###### keywords:

LLM, Alignment, RLHF, Best-of-N

##  1 Introduction

State-of-the-art large language models (LLMs) such as Gemini (Gemini Team,
2023; Reid et al., 2024) and GPT-4 (OpenAI, 2023) are generally trained in
three stages. First, LLMs are pre-trained on large corpora of knowledge using
next-token prediction (Radford et al., 2018, 2019). Second, the pre-trained
models are fine-tuned to follow instructions via supervised fine-tuning (SFT)
(Raffel et al., 2020; Wei et al., 2022). Lastly, reinforcement learning from
human feedback (RLHF) (Christiano et al., 2017; Ziegler et al., 2019; Stiennon
et al., 2020) is used to further increase the quality of generations. The RLHF
step generally consists of learning a reward model (RM) (Ouyang et al., 2022)
on human preferences and then optimizing the LLM to maximize predicted rewards
using reinforcement learning algorithms.

RLHF algorithms and their challenges. Fine-tuning LLMs with reinforcement
learning (RL) is challenging (Casper et al., 2023), notably since it can cause
_forgetting_ (French, 1992) of pre-trained knowledge, and since loopholes in
the RM (Clark and Amodei, 2016; Pan et al., 2022) can cause _reward hacking_
(Askell et al., 2021; Skalse et al., 2022). The standard strategy is to use
policy-gradient methods (Williams, 1992) with $\mathrm{KL}$ regularization
towards the SFT policy. Those RL algorithms seek Pareto-optimal policies with
high reward at low $\mathrm{KL}$, to preserve the general capabilities of the
original model and tackle the misalignment (Ngo et al., 2022) concerns.

Best-of-N sampling. In practice, a surprisingly simple inference-time approach
is often used to improve the quality of generations: Best-of-N sampling
(Stiennon et al., 2020). It consists of drawing $N$ candidate generations from
the reference (typically, supervised fine-tuned) model and selecting the one
with the highest reward according to the RM. This strategy empirically
achieves excellent reward-$\mathrm{KL}$ trade-offs (Nakano et al., 2021; Gao
et al., 2023; Touvron et al., 2023) but increases the computational cost by a
factor of $N$.

BOND. In this paper, we propose BOND (Best-of-N Distillation), a novel RLHF
algorithm that _learns_ a policy that achieves the strong performance of Best-
of-N sampling but, crucially, requires only a single sample at inference time,
as depicted in Figure 1. Our key idea is to cast the alignment of the policy
as a distribution matching problem, where we fine-tune the policy to emulate
the Best-of-N distribution. To achieve this, we first derive an analytical
expression for the Best-of-N distribution. This allows us to consider and
optimize different divergence metrics. We first show how to minimize the
_forward_ $\mathrm{KL}$ divergence using samples from the Best-of-N strategy,
leading to a standard imitation learning setup with a mode covering behavior.
We also show how to minimize the _backward_ $\mathrm{KL}$, leading to a new
form of quantile-based advantage, which does not depend on the reward scale,
and corresponds to a mode seeking behavior. Then, we propose to minimize a
linear combination of forward and backward $\mathrm{KL}$, also known as
_Jeffreys divergence_ , which retains the best of both approaches.
Furthermore, to optimize performance while keeping a reduced sample-
complexity, we propose an _iterative_ BOND approach which consists of
iteratively distilling the Best-of-N of a moving anchor policy. Finally, based
on the aforementioned ideas, we propose J-BOND (J for Jeffreys), a novel,
stable, efficient and practical RLHF algorithm to align LLMs.

Experiments. We first demonstrate the effectiveness of BOND and of our design
choices on the abstractive summarization XSum (Narayan et al., 2018) task.
Then, in Section 6, we apply J-BOND to align Gemma (Gemma Team, 2024)
policies. J-BOND not only improves the $\mathrm{KL}$-reward Pareto front
compared to standard RL algorithms, but also enhances performance on academic
benchmarks and side-by-side comparisons against open-source variants such as
Mixtral (Jiang et al., 2024).

Figure 1: Best-of-N is an _inference-time_ strategy that selects the best
generation among $N$ candidates from a reference LLM policy, improving quality
at the cost of a large computational (need to sample and score $N$ times from
the model). In contrast, the proposed BOND approach aims at obtaining a fine-
tuned policy that can directly sample the Best-of-N generation. This would
inherit the quality of Best-of-N sampling, while requiring a single sample at
inference time. We achieve this by distilling the Best-of-N strategy into the
policy via online _distribution matching_.

##  2 Problem Setup

We consider a LLM based on the transformer (Vaswani et al., 2017)
architecture, defining a policy $\pi(x,\cdot)$ by auto-regressively generating
token sequences $y$ from the prompt $x$. Given a pre-trained and typically
supervised fine-tuned reference policy ${\pi_{\text{ref}}}$, we seek to
further align it to human preferences. To achieve this, throughout the rest of
the paper we assume access to a reward model (RM) which we denote as
$r(\cdot)$, trained to reflect human preferences.

Standard RLHF. Most RL algorithms optimize a linear combination of the
expected reward and a $\mathrm{KL}$ divergence between the current and
reference policy:

$\pi_{\text{RL}}=\operatorname{argmax}_{\pi}\mathbb{E}_{\pi}\mathopen{}\mathclose{{}\left[r(y)}\right]-\beta_{\textrm{RL}}\cdot\operatorname{KL}\mathopen{}\mathclose{{}\left(\pi\mid\mid{\pi_{\text{ref}}}}\right),$
(1)

with regularization strength $\beta_{\textrm{RL}}\geq 0$. This $\mathrm{KL}$
regularization forces the policy to remain close to its initialization
${\pi_{\text{ref}}}$ (Geist et al., 2019; Lazaridou et al., 2020), reducing
forgetting (French, 1992) and reward hacking (Skalse et al., 2022). Equation 1
is usually optimized with online algorithms, as they perform better than their
offline counterparts (Tang et al., 2024). Moreover, simple methods have
demonstrated the best results, e.g., REINFORCE (Williams, 1992) with a sampled
baseline for variance reduction (Li et al., 2023; Ahmadian et al., 2024)
outperform PPO (Schulman et al., 2017).

Best-of-N. A complementary alignment strategy is Best-of-N, which is an
inference-time strategy that involves sampling multiple times from
${\pi_{\text{ref}}}$ and selecting the generation with highest reward
according to the RM $r$. In contrast to RLHF strategies, Best-of-N does not
fine-tune the weights of the LLM, but instead modifies the inference
procedure. Best-of-N was empirically shown to be efficient (Touvron et al.,
2023) when looking at reward/$\mathrm{KL}$ trade-offs, and comes with
theoretical guarantees (Qiping Yang et al., 2024) in terms of Pareto-
optimality. Unfortunately, Best-of-N comes at a significantly higher inference
cost which increases linearly with $N$, since producing $N$ generations is (in
general) $N$ times more costly than sampling a single one.

Motivated by the above considerations, we propose a novel alignment method
which we name BOND for Best-of-N Distillation. The goal of BOND is to _distill
the Best-of- N strategy into the policy_. This allows the policy to reach the
strong performance of Best-of-N sampling, while requiring only _a single
sample_ at inference time. We outline our overall approach in the next
section.

##  3 The BOND Approach

We formulate the BOND approach in two main steps. First, we derive an
analytical expression for the Best-of-N distribution (Section 3.1). Second,
using the derived expression, we phrase the problem as a _distribution
matching_ problem (Section 3.2), i.e., we want to steer the policy closer to
the Best-of-N distribution. In Section 3.3, we draw insightful connections
between BOND and standard RLHF.

###  3.1 The Best-of-N distribution

In this section, we derive the exact analytical distribution of Best-of-N
sampling and study its properties. For simplicity, we drop the context $x$
from all notation without loss of generality and assume that the reward $r(y)$
induces a strict ordering on all generations $y$111To distinguish between
generations with the same reward, ties can be broken with any arbitrary strict
ordering.. We can affirm the following main theorem (proof in Section A.1).

######  Theorem 1.

For any generation $y$, let

${p_{<}}(y)=\mathbb{P}_{y^{\prime}\sim{\pi_{\text{ref}}}}\mathopen{}\mathclose{{}\left[r(y^{\prime})<r(y)}\right]$
(2)

denote the probability that a random generation $y^{\prime}$ from
${\pi_{\text{ref}}}$ is strictly worse than $y$ and let

${p_{\leq}}(y)=\mathbb{P}_{y^{\prime}\sim{\pi_{\text{ref}}}}\mathopen{}\mathclose{{}\left[r(y^{\prime})\leq
r(y)}\right],$ (3)

the probability that $y^{\prime}$ is not better than $y$ (thus including the
equality case). Then, the probability that $y$ is the output of Best-of-N
sampling is given by

${\pi_{\text{BoN}}}(y)={\pi_{\text{ref}}}(y)\times\underbrace{{p_{\leq}}(y)^{N-1}}_{\mathtt{(A)}}\times\underbrace{\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}}_{\mathtt{(B)}}.$
(4)

Interpretation. Theorem 1 provides an intuitive explanation on the behavior of
Best-of-N sampling: it essentially reweights the original sampling
distribution ${\pi_{\text{ref}}}$, by the multiplicative terms $\mathtt{(A)}$
and $\mathtt{(B)}$.

The term $\mathtt{(A)}$ corresponds to a penalty exponential in $N$ based on
the fraction of generations (for the same prompt) that are worse or equal to
the considered generation $y$. Intuitively, this ensures that we sample
exponentially less from bad generations when increasing $N$.

The term $\mathtt{(B)}$ is an additional correction factor due to the
potential of collisions among generations. Importantly, it is at most linear
in $N$ as it is always bounded within $[1,N]$:

$1\leq
1+\sum_{i=2}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}=\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}\leq\sum_{i=1}^{N}1\leq
N$ (5)

It achieves its minimum at $1$ for the worst generation $y_{-}$ since we have
${p_{<}}(y_{-})=0$ by definition. This is not surprising, as we need to sample
$y_{-}$ exactly $N$ times in a row and which corresponds to
${\pi_{\text{BoN}}}(y_{-})={\pi_{\text{ref}}}(y_{-})^{N}$ (note that
${p_{<}}(y_{-})={\pi_{\text{ref}}}(y_{-})$). In contrast, if the likelihood of
individual generations $y$ are low and such generations are good, then
${p_{<}}(y)$ is almost ${p_{\leq}}(y)$ and term $\mathtt{(b)}$ is close to
$N$. Intuitively, this corresponds to the case where sampling a generation $y$
multiple times is unlikely. In the extreme case when ${\pi_{\text{ref}}}$ is a
continuous distribution, term $\mathtt{(B)}$ is constant and equal to $N$ (see
Section A.2).

###  3.2 The BOND objective

The analytical characterization of the Best-of-N distribution allows us to
formulate BOND as a distribution matching problem. That is, we want to solve
the objective:

$\pi_{\texttt{BOND}}=\arg\min_{\pi\in\Pi}\,D(\pi\,\|\,{\pi_{\text{BoN}}}),$
(6)

where $D(\cdot\,\|\,\cdot)$ is a divergence metric steering the training
policy $\pi$ towards ${\pi_{\text{BoN}}}$. For this, a toolbox of possible
divergences exist in the literature including, e.g., forward and backward
$\mathrm{KL}$ (Kullback, 1959). Moreover, we can employ existing distribution
matching techniques to estimate $D$ from online and offline samples. We defer
the choice of suitable divergences and resulting BOND algorithms to Section 4.

###  3.3 Connection with standard RLHF

In this section, we draw important connections between the two seemingly
different objectives of standard RLHF (Equation 1) and BOND (Equation 6).

It is well known (see, e.g., Vieillard et al. (2020); Rafailov et al. (2023))
that the policy maximizing the RLHF objective from Equation 1 is:

$\pi_{\text{RL}}(y)\propto{\pi_{\text{ref}}}(y)\exp\mathopen{}\mathclose{{}\left(\frac{1}{\beta}_{\textrm{RL}}r(y)}\right).$
(7)

From the derived expression of ${\pi_{\text{BoN}}}$ in Theorem 1, we see that
the Best-of-N sampling distribution coincides with the optimal solution of
standard RLHF when using the following specific BOND reward:

$r_{\texttt{BOND}}(y)=\underbrace{\log{p_{\leq}}(y)}_{\texttt{(A)}}+\underbrace{\frac{1}{N-1}\log\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}}_{\texttt{(B)}},$
(8)

and the specific regularization strength $\beta_{\text{BOND}}=\frac{1}{N-1}$.
The term (B) corresponds to the correction factor in Theorem 1, which is
bounded in $\mathopen{}\mathclose{{}\left[0,\frac{\log N}{N-1}}\right]$ for
all generations $y$. Instead term (A) lies in $(-\infty,0]$. This provides two
interesting insights for Best-of-N sampling:

  1. 1.

Best-of-N sampling corresponds to the solution of a standard
$\mathrm{KL}$-regularized RLHF problem where the choice of $N$ determines the
level of KL regularization.

  2. 2.

Best-of-N sampling corresponds to optimizing the expected log reward quantile,
i.e., the log likelihood that the generation has larger reward than a random
sample from the reference distribution. Interestingly, due to the concavity of
the logarithm, $r_{\texttt{BOND}}(y)$ strongly encourages the model to avoid
bad generations rather than encouraging to generate good ones. Moreover,
$r_{\texttt{BOND}}(y)$ is _invariant to monotone transformations of the
reward_ $r(\cdot)$, since it depends only on the rank among the generations.
We conjecture that both these features make the BOND reward
$r_{\texttt{BOND}}(y)$ more robust to reward hacking compared to standard
RLHF.

The connection to RLHF also inspires the proposed approach in this manuscript:
if we can compute the BOND reward or equivalently the Best-of-N distribution
${\pi_{\text{BoN}}}$, then we can steer the policy towards Best-of-N via
distribution matching. In the next section we explore different algorithms to
tackle the main underlying challenges.

##  4 BOND Challenges and Algorithms

Implementing the BOND approach induces the three following challenges: _(1)_
how to estimate the reward quantiles, _(2)_ which is the appropriate
divergence metric to use, and _(3)_ how to choose the hyperparameter $N$
representing the number of sampled generations in Best-of-N. We discuss and
address these challenges in the next three subsections.

###  4.1 Monte-Carlo quantile estimation

One key difficulty in estimating the ${\pi_{\text{BoN}}}$ distribution is that
we need to estimate the quantile

${p_{\leq}}(y)=\mathbb{P}_{y^{\prime}\sim{\pi_{\text{ref}}}}\mathopen{}\mathclose{{}\left[r(y^{\prime})\leq
r(y)}\right],$ (9)

of a given generation $y$. The quantile ${p_{\leq}}(y)$ measures the quality
of $y$ compared to generations from ${\pi_{\text{ref}}}$ when conditioned on
the same prompt (recall that we have suppressed the conditioning on $x$ in our
notation). A very simple but effective quantile estimation method is _Monte-
Carlo sampling_ , sampling $k$ generations from ${\pi_{\text{ref}}}$ and
obtaining the following empirical estimate:

$\hat{p}_{\leq}(y)=\frac{1}{k}\sum_{i=1}^{k}\mathbb{I}\\{r(y_{i})\leq
r(y)\\}.$ (10)

We found this to be a very effective in our experiments, even with a limited
number of samples. In principle, though, one could also use alternative
approaches, e.g., training a learned quantile model (as we explore in Section
B.1).

###  4.2 Jeffreys divergence as a robust objective

The choice of the divergence metric used in BOND is of crucial importance:
different divergences can steer the policy to very different solutions. Here,
we propose the _Jeffreys divergence_ as a robust distribution matching
objective.

The Jeffreys divergence (Jeffreys, 1946) between two distributions is defined
as:

$J_{\text{effreys}}^{\beta}(p\,\|\,q):=(1-\beta)\cdot\underbrace{\text{KL}(q\,\|\,p)}_{\text{Forward
KL}}\,+\,\beta\cdot\underbrace{\text{KL}(p\,\|\,q)}_{\text{Backward KL}}.$
(11)

The (generalized) Jeffreys divergence is a weighted average (with weight
$\beta\in[0,1]$) between the forward and backward $\mathrm{KL}$ divergence.
Notably, when fine-tuning policy $p$, the forward $\text{KL}(q\,\|\,p)$
encourages that generations likely under $q$ are also likely under $p$, thus
encouraging a _mode-covering_ behavior. Instead, the reverse
$\text{KL}(p\,\|\,q)$ is well-known to have a _mode-seeking_ effect, steering
policy $p$ to produce generations that have a high likelihood according to $q$
(Agarwal et al., 2024). While the forward $\mathrm{KL}$ may produce over-
spread distributions, the backward $\mathrm{KL}$ can lead to policy and
entropy collapses. Instead, we empirically show that the Jeffreys divergence
inherits the best of both divergences, producing better aligned policies.

In the context of BOND, this translates into minimizing the divergence
$J_{\text{effreys}}^{\beta}(\pi\,\|\,{\pi_{\text{BoN}}})$ which we can
estimate using samples from the training policy $\pi$ and reference policy
${\pi_{\text{ref}}}$ as follows.

Estimation of the forward KL. The forward KL defined as

$\text{KL}({\pi_{\text{BoN}}}\,\|\,\pi)=\mathbb{E}_{y\sim{\pi_{\text{BoN}}}}\mathopen{}\mathclose{{}\left[\log{\pi_{\text{BoN}}}(y)-\log\pi(y)}\right]$
(12)

can be estimated directly drawing samples from the ${\pi_{\text{BoN}}}$ (i.e.,
sampling $N$ times from ${\pi_{\text{ref}}}$ and selecting the best one) and
can be seen as a supervised fine-tuning loss on the Best-of-N samples:

$\nabla_{\pi}\text{KL}({\pi_{\text{BoN}}}\,\|\,\pi)=-\mathbb{E}_{y\sim{\pi_{\text{BoN}}}}\nabla\log\pi(y).$
(13)

Estimation of the backward KL. The backward KL defined as

$\text{KL}(\pi\,\|\,{\pi_{\text{BoN}}})=\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right]$
(14)

can be estimated from the policy samples (note the expectation w.r.t. $\pi$)
and their estimated log-likelihood under ${\pi_{\text{BoN}}}$. In particular,
by the analogies drawn in Section 3.3, we show (in Section A.3) that its
gradient coincides with a policy gradient (e.g., used by REINFORCE (Williams,
1992) in standard RLHF):

$\nabla_{\pi}\text{KL}(\pi\,\|\,{\pi_{\text{BoN}}})=-(N-1)\mathop{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\big{(}r_{\texttt{BOND}}(y)-\beta_{\texttt{BOND}}\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{ref}}}(y)\big{)}}\right)}\right],$
(15)

with the equivalent reward $r_{\texttt{BOND}}$ and regularization
$\beta_{\texttt{BOND}}$ defined in Section 3.3. Note that
$r_{\texttt{BOND}}(y)$ depends on the true unknown quantile ${p_{\leq}}(y)$
and on the correction factor (B) defined in Equation 8. In practice, we
substitute the true quantile by its estimate, while we observed the correction
factor does not play a significant role. Thus, we use
$r_{\texttt{BOND}}(y)=\hat{p}_{\leq}(y)$. Moreover, to reduce the resulting
variance, we use a policy gradient baseline (Sutton and Barto, 1998) which we
compute as the average return for the other generations in the batch.

Thus, the overall $J_{\text{effreys}}^{\beta}$ loss is a linear weighted
combination between a supervised fine-tuning and a policy gradient loss.

Experiments. In Figure 2, we consider the abstractive summarization XSum task
(Narayan et al., 2018) with ${\pi_{\text{ref}}}$ being a T5 supervised fine-
tuned policy and $r(\cdot)$ being a T5 NLI reward model (Roit et al., 2023).
We run BOND with $J_{\text{effreys}}^{\beta}$ objective and
$\beta\in\\{0,0.5,1\\}$. We use 16 MC samples (per prompt) to estimate
quantiles during training. During eval (each 500 training steps), we use 32 MC
samples to estimate the backward and forward KL divergences between the
training policy and the ${\pi_{\text{BoN}}}$ distribution. We set $N=8$ in
Figure 2, and provide similar plots in Section B.2 for $N=4$ and $16$. The
results confirm our intuition: the Jeffreys divergence ($\beta=0.5$) allows to
minimize both divergences from ${\pi_{\text{BoN}}}$ (left and middle plot),
compared to when solely the backward ($\beta=1$) or the forward ($\beta=0$) KL
divergence is minimized. In addition, in the rightmost plot we report the
reward log quantiles (averaged over the eval batches) of the training policy.
Interestingly, BOND with $\beta=0.5$ maximizes the quantiles similarly to the
mode-seeking $\beta=1$ choice, while the mode-covering forward KL ($\beta=0$)
lags behind.

Figure 2: BOND with $N=8$ and different values of $\beta$ for the Jeffreys
divergence objective (_cf._ Equation 11). When using $\beta=0.5$ (Jeffreys
divergence), BOND minimizes both backward (left plot) and forward (middle
plot) KL divergences from ${\pi_{\text{BoN}}}$, achieving best of both
objectives ($\beta=0$ and $\beta=1$). Moreover, it optimizes the reward
quantiles (right plot) significantly more than when using $\beta=0$ (forward
KL minimization). Results for $N=4$ and $N=16$ are reported in Section B.2.

###  4.3 Iterative BOND

Finally, we discuss the choice of the parameter $N$. In practice, choosing $N$
may be difficult for three main reasons: _(1)_ As in standard RLHF, $N$ plays
the role of regularization (see Section 3.3): a large $N$ improves downstream
performance, but if $N$ is too large it will eventually cause reward over
optimization (Gao et al., 2023). _(2)_ The larger the $N$ the more the
estimate of ${\pi_{\text{BoN}}}$ is sensitive to errors in the estimated
quantiles (since ${\pi_{\text{BoN}}}(y)\propto p_{\leq}(y)^{N-1}$). _(3)_
Estimating the forward KL divergence requires sampling from
${\pi_{\text{BoN}}}$ which is prohibitive for large $N$.

To address the above challenges, we propose the _iterative_ BOND approach. The
approach is inspired by the fact that Best-of-N sampling from a Best-of-N
distribution, coincides with Best-of-N2 sampling from the original
distribution. More generally, by informally defining $\text{Bo$N$}(\cdot)$ as
an operator that performs Best-of-N sampling from a base distribution, we
have:

$\underbrace{\text{Bo{N}}(\cdots\text{Bo{N}}(\text{Bo{N}}(}_{M\text{
times}}{\pi_{\text{ref}}})))\,\equiv\,\text{Bo{N}${}^{M}$}({\pi_{\text{ref}}}).$
(16)

This suggests the key idea behind iterative BOND: if we know how to distill
the Best-of-N distribution (i.e., via BOND), then we can apply BOND
recursively (say $M$ times), equivalently to distilling a Best-of-NM of the
initial distribution ${\pi_{\text{ref}}}$. This allows fixing a small $n$
(i.e., $n=2$) and running BOND (with $N=n$) in an _iterative fashion_ , as an
improvement operator. For this, we can introduce an auxiliary _anchor policy_
$\pi_{\text{anchor}}$ initialized as ${\pi_{\text{ref}}}$. We can then run
BOND against $\pi_{\text{anchor}}$ (i.e., we can distill the Best-of-n version
of $\pi_{\text{anchor}}$) and, after a given number of distillation steps,
update $\pi_{\text{anchor}}$ to be the current training policy $\pi_{t}$. The
overall approach is depicted in Figure 3 and summarized in Algorithm 1.

Figure 3: Iterative BOND approach. The policy $\pi_{t}$ is trained to
iteratively distill a Best-of-$N$ (in the figure, $N=2$) version of a _moving
anchor_. This allows to continuously improve the policy performance without
requiring to set a (large) $N$ upfront. Moreover, it leads to better training
stability and a minimal computational complexity, since a small $n$ is used at
each distillation step.

In a nutshell, iterative BOND allows exponential scaling to arbitrary large
$N$ (in fact, it does not require setting $N$ in advance) while keeping a
reduced sample complexity and a stable optimization. The claim is validated in
the results below.

Experiments. In Figure 4, we consider the same experimental setup than in
Figure 2 from Section 4.2, fix the BOND objective to
$J_{\text{effreys}}^{0.5}$ and run iterative BOND with $n\in\\{2,4\\}$ where
the moving anchor is updated every 1000 steps. We report the average reward
(left plot) and average log quantile (middle plot) obtained during training
and compare them to (non-iterative) BOND run with $N\in\\{4,8,16\\}$. As
expected, both reward signals saturate early for non-iterative BOND (the
smaller the $N$ the earlier the reward saturates), while the iterative BOND
approach continuously improves performance (the higher the $n$, the faster).
Moreover, in the rightmost plot we plot the obtained log quantiles against the
$\mathrm{KL}$ from the reference policy. The plot shows that iterative BOND
essentially has the same reward/$\mathrm{KL}$ trade-off of the non-iterative
BOND runs, but crucially allows keeping a small $n$ and to smoothly but
continuously move away from ${\pi_{\text{ref}}}$.

Algorithm 1 Iterative BOND (meta algorithm)

Inputs: ${\pi_{\text{ref}}}$, $n\in\mathbb{N}$.

Initialize $\pi_{0}={\pi_{\text{ref}}}$,
$\pi_{\text{anchor}}^{0}={\pi_{\text{ref}}}$.

for $t=0,\ldots,$ do */ iterations

$\pi_{t+1}=\arg\min_{\pi\in\Pi}D(\pi\,\|\,\text{Best-
of-{n}\@}(\pi_{\text{anchor}}^{t}))$ */ distill the Best-of-n version of
$\pi_{\text{anchor}}^{t}$

$\pi_{\text{anchor}}^{t}=\pi_{t}$

Figure 4: Iterative BOND (with $n=2$ and $n=4$) compared to BOND with
$N=4,8,16$. Iterative BOND continuously improves rewards (left plot) and log
quantiles (middle plot), while they saturate for non-iterative BOND (the
smaller the $N$, the sooner). It allows achieving the same
reward/$\mathrm{KL}$ trade-off (right plot) as non-iterative BOND while
keeping a small $n$ and smoothly moving away from ${\pi_{\text{ref}}}$.

##  5 The J-BOND Algorithm

In this section we present J-BOND, a concrete and practical BOND algorithm
motivated by the results discussed in the previous sections. We describe its
main components below, and summarize it in the pseudo-code of Algorithm 2.

J-BOND follows the template of iterative BOND (Algorithm 1) with $n=2$, i.e.,
it fine-tunes policy $\pi_{t}$ to iteratively distill the Best-of-2 version of
a moving anchor $\pi_{\text{anchor}}^{t}$, initialized as
${\pi_{\text{ref}}}$. The name J-BOND stands for _Jeffreys divergence_ BOND
because it uses the Jeffreys divergence as distribution matching objective,
i.e., it minimizes $J_{\text{effreys}}^{\beta}(\pi\,\|\,\text{Best-
of-2}(\pi_{\text{anchor}}^{t}))$ as defined in Section 4.2.

Minimal sample complexity. Compared to the BOND algorithms tested in the
previous section, J-BOND has a minimal sample complexity: for each prompt in
the batch it generates $1$ sample from the policy $\pi_{t}$ and $2$ samples
from the anchor $\pi_{\text{anchor}}^{t}$. While more anchor samples are
generally useful for a better divergence estimation (in Section 4 we used 16
MC samples), autoregressive sampling is the main bottleneck of online RLHF and
we have thus opted for a practical approach working with a small number of
samples.

Algorithm 2 The J-BOND algorithm

Inputs: Prompt dataset $\mathcal{D}$, reference policy ${\pi_{\text{ref}}}$,
reward $r(\cdot)$, $\beta,\eta\in[0,1],\gamma\geq 0$.

Initialize policy and anchor
$\pi_{0}=\pi_{\text{anchor}}^{0}={\pi_{\text{ref}}}$.

for $t=0,\ldots$ do

Sample batch of prompts $\mathcal{D}_{t}\subseteq\mathcal{D}$

For each $x\in\mathcal{D}_{t}$ generate: 1 policy sample $y\sim\pi_{t}(x)$ and
2 anchor samples $y_{1}^{\prime},y_{2}^{\prime}\sim\pi_{\text{anchor}}^{t}(x)$
*/ Forward KL */

Extract Best-of-2 sample:
$y^{\prime}_{\text{Bo2}}=\arg\max_{y^{\prime}\in\\{y_{1}^{\prime},y_{2}^{\prime}\\}}r(y^{\prime})$.

Compute forward KL gradient:
$G_{\text{FW}}(x,\pi_{t})=-\nabla_{\pi_{t}}\log\pi_{t}(x,y_{\text{Bo2}}^{\prime})$
*/ Backward KL */

Compute $r_{\texttt{J-BOND}}(x,y)$ according to Equation 17.

Compute return:
$R(x,y)=r_{\texttt{J-BOND}}(x,y)-(\log\pi_{t}(x,y)-\log\pi_{\text{anchor}}^{t}(x,y))$.

[Optional] Compute baseline $B$, e.g. average return of the other generations
in the batch.

Compute backward KL gradient:
$G_{\text{BW}}(x,\pi_{t})=-\nabla_{\pi_{t}}\log\pi_{t}(x,y)\cdot\mathopen{}\mathclose{{}\left(R(x,y)-B}\right)$.
*/ Additional KL regularization */

KL regularization gradient:
$G_{\text{Reg}}(x,\pi_{t})=-\nabla_{\pi_{t}}\log\pi_{t}(x,y)\cdot(\log\pi_{t}(x,y)-\log\pi_{\text{anchor}}^{t}(x,y))$.
*/ Overall policy update: Jeffreys divergence + KL regularization */

Update policy weights $\theta_{t+1}$ with the overall stochastic gradient:
$\mathbb{E}_{x\sim\mathcal{D}_{t}}\mathopen{}\mathclose{{}\left[(1-\beta)\cdot
G_{\text{FW}}(x,\pi_{t})+\beta\cdot G_{\text{BW}}(x,\pi_{t})+\gamma\cdot
G_{\text{Reg}}(x,\pi_{t})}\right]$ */ Update moving anchor */

Update anchor weights with EMA:
$\theta_{\text{anchor}}^{t+1}\leftarrow(1-\eta)\cdot\theta_{\text{anchor}}^{t}+\eta\cdot\theta_{t+1}$

Crude divergence estimate based on 2 anchor samples. The policy and anchor
samples are used to obtain a crude estimate of the forward and backward KL
components of $J_{\text{effreys}}^{\beta}(\pi\,\|\,\text{Best-
of-2}(\pi_{\text{anchor}}^{t}))$ as described next.

We can minimize the _forward KL_ as described in Section 4.2, by doing
supervised fine-tuning on the best of the 2 anchor samples. To minimize the
_backward KL_ , we utilize the policy gradient-style loss of Equation 15
replacing $r_{\texttt{{BOND}\@}}(y)$ with a different reward which we denote
as $r_{\texttt{J-BOND}}(y)$. The reason for this is that when only 2 anchor
samples are available, the reward $r_{\texttt{BOND}}(y)=\log\hat{p}_{\leq}(y)$
would be quite uninformative due to $\hat{p}_{\leq}(y)$ being a very noisy MC
estimate. Let $y$ be the policy sample and
$\\{y_{1}^{\prime},y_{2}^{\prime}\\}$ be the corresponding anchor samples, we
instead define $r_{\texttt{J-BOND}}(y)$ as

$r_{\texttt{J-BOND}}(y)=\begin{cases}-\log(16)&\text{if}\quad
r(y)<\text{min}\\{r(y_{1}^{\prime}),r(y_{2}^{\prime})\\}\\\
0&\text{otherwise}\end{cases}.$ (17)

That is, generation $y$ receives a negative reward of $-\log(16)$ if it has
worse reward than _both_ anchor samples, while receives $0$ reward otherwise.
The above definition is motivated by the following two main reasons:

  * (i)

We negatively reward $y$ only if it is worse than both the anchor samples, to
mimic the concavity of the ideal (and unknown) reward function
$r_{\texttt{{BOND}\@}}=\log{p_{\leq}}(\cdot)$.

  * (ii)

We choose value $-\log(16)$ to ensure that:
$\mathbb{E}_{y_{1}^{\prime},y_{2}^{\prime}\sim\pi_{\text{anchor}}^{t}}\big{[}r_{\texttt{J-BOND}}(y)\big{]}=\log{p_{\leq}}(y)$
when ${p_{\leq}}(y)=0.5$. The interested reader is referred to Section A.4 for
a derivation and illustration of this fact. In words, the value $-\log(16)$
calibrates the reward function $r_{\texttt{J-BOND}}(\cdot)$ so that, in
expectation under the 2 anchor samples, it matches with the ideal reward
$\log{p_{\leq}}(\cdot)$ for generations $y$ that have median reward (i.e.,
when ${p_{\leq}}(y)=0.5$).

Exponential Moving Average (EMA) anchor. An important component of J-BOND,
which refines the vanilla iterative BOND of Section 4.3, is the use of an
_Exponential Moving Average (EMA)_ anchor. That is, instead of using a
periodically updated anchor, we update the anchor weights
$\theta_{\text{anchor}}^{t}$ _at each fine-tuning step_ as a moving average of
the policy weights $\theta_{t}$:

$\theta_{\text{anchor}}^{t+1}\leftarrow(1-\eta)\cdot\theta_{\text{anchor}}^{t}+\eta\cdot\theta_{t+1}.$
(18)

Consistently with WARP (Ramé et al., 2024), we observed that this weight
averaging procedure has a positive effect on training stability by reducing
variance, and can improve the overall reward/KL trade-off of J-BOND. We
provide an ablation in Section 6.

Additional KL regularization. Finally, we further regularize the policy to
stay closer to the moving anchor via an extra222Note that $\mathrm{KL}$
regularization is already present in the backward $\mathrm{KL}$ component of
the Jeffreys divergence. $\mathrm{KL}$ regularization term, modulated by a
tunable hyperparameter $\gamma\geq 0$. The scope is to further stabilize the
policy updates, viewing the overall operator as a constrained optimization
one:

$\pi_{t+1}=\arg\min_{\pi\in\Pi}J_{\text{effreys}}^{\beta}(\pi\,\|\,\text{Best-
of-2}(\pi_{\text{anchor}}^{t}))\quad+\quad\gamma\cdot\text{KL}(\pi_{t}\,\|\,\pi_{\text{anchor}}^{t}).$
(19)

##  6 Experiments

We test J-BOND on relevant use cases with the following main goals. First, we
ablate and showcase important aspects of J-BOND: the benefits of the EMA
anchor, and the effects of the anchor speed and the additional KL
regularization. Then, we compare J-BOND to classical RLHF baselines using
REINFORCE, demonstrating its efficacy and better performance.

Figure 5: J-BOND with periodic anchor updates (every 50 steps) vs. EMA anchor
($\eta=0.02$) on Gemma 7B. While attaining the same reward (left), using the
EMA anchor displays a significantly lower KL than the reference policy
(middle) and thus a better reward/KL trade-off (right).

Setup. We consider Gemma (2B and 7B) models (Gemma Team, 2024) which we aim to
fine-tune into better conversational agents. For this task, we consider a set
of conversational prompts $\mathcal{D}$, a reference policy
${\pi_{\text{ref}}}$ previously supervised fine-tuned on similar prompts, and
a large previously trained reward model $r(\cdot)$. We use a batch size 128
and the Adam optimizer (Kingma and Ba, 2015) with learning rate $3e^{-6}$ and
100 warm-up steps. For the Jeffreys divergence objective, we set $\beta=0.5$.

EMA vs. hard anchor updates. We ablate the benefits of using an EMA moving
anchor (Equation 18) compared to the periodically updated anchor used in
Section 4.3. For this, we run J-BOND with $\gamma=0$ and EMA coefficient
$\eta=0.02$ on Gemma 7B, and compare it with its variant where the anchor is
only updated every $50$ steps. In Figure 5, we report the average reward of
the policy during training (left plot), the KL from the reference policy
${\pi_{\text{ref}}}$ (middle plot), and the resulting reward/KL trade-off for
the two runs. Both runs produce the same reward increase profile (this is not
surprising since an EMA with $\eta=0.02$ roughly corresponds to an update
period of 50 steps) but, crucially, J-BOND _with an EMA anchor displays a
significantly lower KL increase_ and, as a result, a better reward/KL trade-
off.

Figure 6: J-BOND: the role of the EMA mixing parameter $\eta$ and
regularization $\gamma$. The left plot shows J-BOND with $\gamma=0$,
$\eta\in\\{0.01,0.05,0.1\\}$: the larger the $\eta$ the faster the average
reward increases. The middle and right plots show J-BOND with $\eta=0.05$ and
$\gamma\in\\{0,0.5,1,2\\}$: the larger the regularization $\gamma$, the slower
the policy moves away from ${\pi_{\text{ref}}}$, improving the reward/KL
trade-off.

Figure 7: J-BOND ($\eta=0.02$) for Gemma 7B, compared to standard REINFORCE
(with KL-regularized objective of Equation 1) with regularization strength
$\beta_{\textrm{RL}}\in\\{0.001,0.01,0.1,1\\}$. J-BOND does not require
committing to a specific regularization strength, but it continuously improves
the reward displaying a stable and linear $\mathrm{KL}$ increase and a better
reward/$\mathrm{KL}$ trade-off.

Anchor speed and $\mathrm{KL}$ regularization. Here we illustrate the effect
of the anchor mixing parameter $\eta\in[0,1]$ and the benefits of the
additional $\mathrm{KL}$ regularization parameter $\gamma\geq 0$ introduced in
Equation 19. For this, we consider Gemma 2B and first run J-BOND with
$\gamma=0$ and $\eta\in\\{0.01,0.05,0.1\\}$. In the left plot of Figure 6 we
report the average reward of the policy along training. This illustrates that
_the larger the mixing parameter $\eta$_ (i.e., the faster the anchor moves),
_the faster the reward increases_ , as one could intuitively expect. Second,
we fix $\eta=0.05$ and run J-BOND with different regularization strengths
$\gamma\in\\{0,0.5,1,2\\}$. We plot the results in the middle and rightmost
plots of Figure 6. As expected, _the larger the regularization $\gamma$_, the
more constrained are the policy updates and thus, _the slower the policy moves
away from ${\pi_{\text{ref}}}$_ (middle plot). Importantly, the right plot
shows that such a regularization has a positive effect since it can ultimately
improve the reward/$\mathrm{KL}$ trade-off.

Comparison with standard RLHF. We compare J-BOND against standard RLHF
algorithms that aim at maximizing the $\mathrm{KL}$-regularized objective of
Equation 1. To optimize Equation 1, we use REINFORCE (Williams, 1992) with 2
policy samples per-prompt and a leave-one-out baseline (Ahmadian et al., 2024)
for policy gradient advantages. For J-BOND we set the anchor mixing
coefficient to $\eta=0.02$. For REINFORCE, we test possible regularization
strengths $\beta_{\textrm{RL}}\in\\{0.001,0.01,0.1,1\\}$. In Figure 7 we plot
the average reward of the training policy (left plot) and its KL divergence
from the reference ${\pi_{\text{ref}}}$ (middle plot). As presumed, REINFORCE
is quite sensitive to the regularization coefficient $\beta_{\textrm{RL}}$:
the larger the regularization, the lower the reward achieved by REINFORCE (and
the lower the KL from ${\pi_{\text{ref}}}$). This highlight a key advantage of
J-BOND: it does not require committing to a specific regularization level, but
it continuously improves the reward displaying a stable and linear KL
increase. Moreover, in the rightmost plot we plot the corresponding reward/KL
trade-off showing that J-BOND produces a better reward/KL than all of the
REINFORCE baselines.

J-BOND for Gemma open models. J-BOND was also the algorithm used to fine-tune
open-weight models such as Gemma 1.1 2B and 7B (Gemma Team, 2024),
RecurrentGemma 2B and 9B (Botev et al., 2024) as well as CodeGemma 1.1
(CodeGemma Team, 2024). This led to competitive performance. For example, the
Gemma 1.1 IT 7B model outperforms Mistral 7B v0.2 Instruct in both safety and
instruction following (see Gemma Team (2024, Table 5) for more details).

##  7 Related Work

Best-of-N was introduced in Stiennon et al. (2020) as a straightforward but
costly inference method to optimize language generation against a given reward
function. Further works established and refined an analytical form for the KL
divergence against the reference (i.e., Bo$1$) policy (Hilton, 2023; Beirami
et al., 2024), provided an estimator for the average Best-of-N reward (Nakano
et al., 2021), made theoretical connections with $\mathrm{KL}$-constrained RL
(Yang et al., 2024) and provided scaling laws for Best-of-N alignment (Gao et
al., 2023).

Matching Best-of-N for improved alignment is a strategy that was studied in
different flavors in the literature. Dong et al. (2023) and Touvron et al.
(2023) propose to fine-tune LLMs in a supervised fashion on Best-of-N data
actually applying forward $\mathrm{KL}$ minimization. Concurrently to ours,
Gui et al. (2024) proposes to mimic the Best-of-N policy by applying a
combination of supervised fine-tuning on best responses and direct preference
optimization on best-and-worst response pairs. The latter is similar to a
common strategy in online preference optimization methods: Guo et al. (2024)
use pairwise AI feedback on online generations to obtain online preferences
that are then optimized, while Calandriello et al. (2024) use a dedicated
preference reward model instead. Concurrently and closest to our work, Amini
et al. (2024) also apply distribution matching in order to get the benefits of
Best-of-N sampling with amortized cost. While their formalization is
identical, we opt for a different divergence (i.e., Jeffreys) than the one
they use (i.e., only backward KL), and propose an iterative procedure with
dynamic anchor, which we show critical for optimal results. Best-of-N can also
be used for self-improvement in reward modeling, as evidenced in Pace et al.
(2024).

Using a contrastive advantage is an option of J-BOND studied in prior works as
well, which replaced a value estimate by the average Monte Carlo return of
other samples. This was applied in the context of REINFORCE (Kool et al.,
2019; Pinto et al., 2023), for online RLHF (Ahmadian et al., 2024), offline
RLHF (Flet-Berliac et al., 2024) and preference optimization (Wu et al.,
2024).

Exponential moving average (EMA) of policy as reference in regularization,
which we use in J-BOND, is an increasingly popular option. While most
alignment approaches use a static anchor, dynamic anchors bring the benefit of
improving the flexibility of the policy space being explored (Munos et al.,
2023; Gorbatovski et al., 2024; Ramé et al., 2024), with the caveat that too
slow updates limit optimization and too fast updates hinder stability.

Scaling post-training and iterated amplification. BOND hinges on the idea of
investing more resources during training to ensure that computational demands
during inference remain low, a factor often overlooked in traditional scaling
laws (Hoffmann et al., 2022). Specifically, BOND incorporates the principles
of iterated amplification (Christiano et al., 2018; Cotra, 2018), where
amplification in this context consists of producing multiple generations,
comparing their rewards, and using these to iteratively improve the policy
performance. In this regard, BOND is complementary to WARM (Ramé et al., 2024)
and WARP (Ramé et al., 2024), which previously scaled post-training by
training multiple reward models and policies, respectively.

##  8 Conclusion

We introduce BOND, a novel RLHF method that fine-tunes the policy via online
distillation of the Best-of-N sampling distribution. We propose a concrete
algorithm, J-BOND, that integrates multiple components to enhance its
practicality and efficiency; Monte-Carlo quantile estimation, a combination
between forward and backward $\mathrm{KL}$ divergence objectives, and an
iterative procedure with an exponential moving average anchor. J-BOND improves
the $\mathrm{KL}$-reward Pareto front of solutions, and compares favorably
against state-of-the-art baselines. We hope this work can help improve
alignment of AI systems, making them safer and more reliable.

## Acknowledgements

We thank Daniele Calandriello for insightful comments, as well as Gil Shamir,
Bilal Piot, and Remi Munos for helpful discussions.

## References

  * Agarwal et al. (2024) R. Agarwal, N. Vieillard, Y. Zhou, P. Stanczyk, S. R. Garea, M. Geist, and O. Bachem.  On-policy distillation of language models: Learning from self-generated mistakes.  In _ICLR_ , 2024.
  * Ahmadian et al. (2024) A. Ahmadian, C. Cremer, M. Gallé, M. Fadaee, J. Kreutzer, A. Üstün, and S. Hooker.  Back to basics: Revisiting REINFORCE style optimization for learning from human feedback in LLMs.  _arXiv preprint_ , 2024.
  * Amini et al. (2024) A. Amini, T. Vieira, and R. Cotterell.  Variational Best-of-N alignment.  _arXiv preprint_ , 2024.
  * Askell et al. (2021) A. Askell, Y. Bai, A. Chen, D. Drain, D. Ganguli, T. Henighan, A. Jones, N. Joseph, B. Mann, N. DasSarma, N. Elhage, Z. Hatfield-Dodds, D. Hernandez, J. Kernion, K. Ndousse, C. Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, and J. Kaplan.  A general language assistant as a laboratory for alignment.  _arXiv preprint_ , 2021.
  * Beirami et al. (2024) A. Beirami, A. Agarwal, J. Berant, A. D’Amour, J. Eisenstein, C. Nagpal, and A. T. Suresh.  Theoretical guarantees on the Best-of-N alignment policy.  _arXiv preprint_ , 2024.
  * Botev et al. (2024) A. Botev, S. De, S. L. Smith, A. Fernando, G.-C. Muraru, R. Haroun, L. Berrada, R. Pascanu, P. G. Sessa, R. Dadashi, L. Hussenot, J. Ferret, S. Girgin, O. Bachem, A. Andreev, K. Kenealy, T. Mesnard, C. Hardin, S. Bhupatiraju, S. Pathak, L. Sifre, M. Rivière, M. S. Kale, J. Love, P. Tafti, A. Joulin, N. Fiedel, E. Senter, Y. Chen, S. Srinivasan, G. Desjardins, D. Budden, A. Doucet, S. Vikram, A. Paszke, T. Gale, S. Borgeaud, C. Chen, A. Brock, A. Paterson, J. Brennan, M. Risdal, R. Gundluru, N. Devanathan, P. Mooney, N. Chauhan, P. Culliton, L. G. Martins, E. Bandy, D. Huntsperger, G. Cameron, A. Zucker, T. Warkentin, L. Peran, M. Giang, Z. Ghahramani, C. Farabet, K. Kavukcuoglu, D. Hassabis, R. Hadsell, Y. W. Teh, and N. de Frietas.  Recurrentgemma: Moving past transformers for efficient open language models.  _arXiv preprint_ , 2024.
  * Calandriello et al. (2024) D. Calandriello, D. Guo, R. Munos, M. Rowland, Y. Tang, B. A. Pires, P. H. Richemond, C. L. Lan, M. Valko, T. Liu, et al.  Human alignment of large language models through online preference optimisation.  In _ICML_ , 2024.
  * Casper et al. (2023) S. Casper, X. Davies, C. Shi, T. K. Gilbert, J. Scheurer, J. Rando, R. Freedman, T. Korbak, D. Lindner, P. Freire, et al.  Open problems and fundamental limitations of reinforcement learning from human feedback.  _TMLR_ , 2023.
  * Christiano et al. (2018) P. Christiano, B. Shlegeris, and D. Amodei.  Supervising strong learners by amplifying weak experts.  _arXiv preprint_ , 2018.
  * Christiano et al. (2017) P. F. Christiano, J. Leike, T. Brown, M. Martic, S. Legg, and D. Amodei.  Deep reinforcement learning from human preferences.  In _NeurIPS_ , 2017.
  * Clark and Amodei (2016) J. Clark and D. Amodei.  Faulty Reward Functions in the Wild.  https://openai.com/research/faulty-reward-functions, 2016.
  * CodeGemma Team (2024) CodeGemma Team.  Codegemma: Open code models based on gemma.  _arXiv preprint_ , 2024.
  * Cotra (2018) A. Cotra.  Iterated distillation and amplification.  https://ai-alignment.com/iterated-distillation-and-amplification-157debfd1616, 2018\.
  * Cover (1999) T. M. Cover.  _Elements of information theory_.  1999\.
  * Dabney et al. (2017) W. Dabney, M. Rowland, M. G. Bellemare, and R. Munos.  Distributional reinforcement learning with quantile regression.  In _AAAI Conference on Artificial Intelligence_ , 2017.
  * Dong et al. (2023) H. Dong, W. Xiong, D. Goyal, Y. Zhang, W. Chow, R. Pan, S. Diao, J. Zhang, K. SHUM, and T. Zhang.  RAFT: Reward ranked finetuning for generative foundation model alignment.  _TMLR_ , 2023.
  * Flet-Berliac et al. (2024) Y. Flet-Berliac, N. Grinsztajn, F. Strub, E. Choi, C. Cremer, A. Ahmadian, Y. Chandak, M. G. Azar, O. Pietquin, and M. Geist.  Contrastive policy gradient: Aligning llms on sequence-level scores in a supervised-friendly fashion.  _arXiv preprint_ , 2024.
  * French (1992) R. M. French.  Semi-distributed representations and catastrophic forgetting in connectionist networks.  _Connection Science_ , 1992.
  * Gao et al. (2023) L. Gao, J. Schulman, and J. Hilton.  Scaling laws for reward model overoptimization.  In _ICML_ , 2023.
  * Geist et al. (2019) M. Geist, B. Scherrer, and O. Pietquin.  A theory of regularized markov decision processes.  In _ICML_ , 2019.
  * Gemini Team (2023) Gemini Team.  Gemini: A family of highly capable multimodal models.  2023\.
  * Gemma Team (2024) Gemma Team.  Gemma: Open models based on gemini research and technology.  _arXiv preprint_ , 2024.
  * Gorbatovski et al. (2024) A. Gorbatovski, B. Shaposhnikov, A. Malakhov, N. Surnachev, Y. Aksenov, I. Maksimov, N. Balagansky, and D. Gavrilov.  Learn your reference model for real good alignment.  _arXiv preprint_ , 2024.
  * Gui et al. (2024) L. Gui, C. Gârbacea, and V. Veitch.  Bonbon alignment for large language models and the sweetness of Best-of-N sampling.  _arXiv preprint_ , 2024.
  * Guo et al. (2024) S. Guo, B. Zhang, T. Liu, T. Liu, M. Khalman, F. Llinares, A. Rame, T. Mesnard, Y. Zhao, B. Piot, J. Ferret, and M. Blondel.  Direct language model alignment from online ai feedback.  _arXiv preprint_ , 2024.
  * Hilton (2023) J. Hilton.  KL divergence of max-of-n, 2023.
  * Hoffmann et al. (2022) J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al.  Training compute-optimal large language models.  In _NeurIPS_ , 2022.
  * Jeffreys (1946) H. Jeffreys.  An invariant form for the prior probability in estimation problems.  _Proceedings of the Royal Society of London. Series A, Mathematical and Physical Sciences_ , 186(1007):453–461, 1946\.
  * Jiang et al. (2024) A. Q. Jiang, A. Sablayrolles, A. Roux, A. Mensch, B. Savary, C. Bamford, D. S. Chaplot, D. d. l. Casas, E. B. Hanna, F. Bressand, et al.  Mixtral of experts.  _arXiv preprint arXiv:2401.04088_ , 2024.
  * Kingma and Ba (2015) D. P. Kingma and J. Ba.  Adam: A method for stochastic optimization.  In _ICLR_ , 2015.
  * Kool et al. (2019) W. Kool, H. van Hoof, and M. Welling.  Buy 4 REINFORCE samples, get a baseline for free!  2019\.
  * Kullback (1959) S. Kullback.  _Information Theory and Statistics_.  New York, 1959.
  * Lazaridou et al. (2020) A. Lazaridou, A. Potapenko, and O. Tieleman.  Multi-agent communication meets natural language: Synergies between functional and structural language learning.  In _ACL_ , 2020.
  * Li et al. (2023) Z. Li, T. Xu, Y. Zhang, Y. Yu, R. Sun, and Z.-Q. Luo.  Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models.  _arXiv preprint_ , 2023.
  * Munos et al. (2023) R. Munos, M. Valko, D. Calandriello, M. G. Azar, M. Rowland, Z. D. Guo, Y. Tang, M. Geist, T. Mesnard, A. Michi, et al.  Nash learning from human feedback.  _arXiv preprint_ , 2023.
  * Nakano et al. (2021) R. Nakano, J. Hilton, S. Balaji, J. Wu, L. Ouyang, C. Kim, C. Hesse, S. Jain, V. Kosaraju, W. Saunders, et al.  Webgpt: Browser-assisted question-answering with human feedback.  _arXiv preprint_ , 2021.
  * Narayan et al. (2018) S. Narayan, S. B. Cohen, and M. Lapata.  Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization.  In E. Riloff, D. Chiang, J. Hockenmaier, and J. Tsujii, editors, _Conference on Empirical Methods in Natural Language Processing_ , 2018.
  * Ngo et al. (2022) R. Ngo, L. Chan, and S. Mindermann.  The alignment problem from a deep learning perspective.  _arXiv preprint_ , 2022.
  * OpenAI (2023) OpenAI.  Gpt-4 technical report.  2023\.
  * Ouyang et al. (2022) L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al.  Training language models to follow instructions with human feedback.  _NeurIPS_ , 2022.
  * Pace et al. (2024) A. Pace, J. Mallinson, E. Malmi, S. Krause, and A. Severyn.  West-of-n: Synthetic preference generation for improved reward modeling.  _arXiv preprint_ , 2024.
  * Pan et al. (2022) A. Pan, K. Bhatia, and J. Steinhardt.  The effects of reward misspecification: Mapping and mitigating misaligned models.  In _ICLR_ , 2022.
  * Pinto et al. (2023) A. S. Pinto, A. Kolesnikov, Y. Shi, L. Beyer, and X. Zhai.  Tuning computer vision models with task rewards.  In _ICML_ , 2023.
  * Qiping Yang et al. (2024) J. Qiping Yang, S. Salamatian, Z. Sun, A. Theertha Suresh, and A. Beirami.  Asymptotics of language model alignment.  _arXiv_ , 2024.
  * Radford et al. (2018) A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever.  Improving language understanding by generative pre-training.  2018\.
  * Radford et al. (2019) A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Rafailov et al. (2023) R. Rafailov, A. Sharma, E. Mitchell, S. Ermon, C. D. Manning, and C. Finn.  Direct preference optimization: Your language model is secretly a reward model.  _arXiv preprint_ , 2023.
  * Raffel et al. (2020) C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _JMLR_ , 2020.
  * Ramé et al. (2024) A. Ramé, N. Vieillard, L. Hussenot, R. Dadashi, G. Cideron, O. Bachem, and J. Ferret.  WARM: On the benefits of weight averaged reward models.  In _ICML_ , 2024.
  * Ramé et al. (2024) A. Ramé, J. Ferret, N. Vieillard, R. Dadashi, L. Hussenot, P.-L. Cedoz, P. G. Sessa, S. Girgin, A. Douillard, and O. Bachem.  WARP: On the benefits of weight averaged rewarded policies.  _arXiv preprint_ , 2024.
  * Reid et al. (2024) M. Reid, N. Savinov, D. Teplyashin, D. Lepikhin, T. Lillicrap, J.-b. Alayrac, R. Soricut, A. Lazaridou, O. Firat, J. Schrittwieser, et al.  Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context.  _arXiv preprint_ , 2024.
  * Roit et al. (2023) P. Roit, J. Ferret, L. Shani, R. Aharoni, G. Cideron, R. Dadashi, M. Geist, S. Girgin, L. Hussenot, O. Keller, et al.  Factually consistent summarization via reinforcement learning with textual entailment feedback.  In _ACL_ , 2023.
  * Schulman et al. (2017) J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov.  Proximal policy optimization algorithms.  _arXiv preprint_ , 2017.
  * Skalse et al. (2022) J. M. V. Skalse, N. H. R. Howe, D. Krasheninnikov, and D. Krueger.  Defining and characterizing reward gaming.  In _NeurIPS_ , 2022.
  * Stiennon et al. (2020) N. Stiennon, L. Ouyang, J. Wu, D. Ziegler, R. Lowe, C. Voss, A. Radford, D. Amodei, and P. F. Christiano.  Learning to summarize with human feedback.  _NeurIPS_ , 2020.
  * Sutton and Barto (1998) R. S. Sutton and A. G. Barto.  _Reinforcement Learning: An Introduction_.  MIT Press, 1998.
  * Tang et al. (2024) Y. Tang, D. Z. Guo, Z. Zheng, D. Calandriello, Y. Cao, E. Tarassov, R. Munos, B. Ávila Pires, M. Valko, Y. Cheng, and W. Dabney.  Understanding the performance gap between online and offline alignment algorithms.  _arXiv preprint_ , 2024.
  * Touvron et al. (2023) H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, et al.  Llama 2: Open foundation and fine-tuned chat models.  _arXiv preprint_ , 2023.
  * Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin.  Attention is all you need.  In _NeurIPS_ , 2017.
  * Vieillard et al. (2020) N. Vieillard, T. Kozuno, B. Scherrer, O. Pietquin, R. Munos, and M. Geist.  Leverage the average: an analysis of kl regularization in reinforcement learning.  _NeurIPS_ , 2020.
  * Wei et al. (2022) J. Wei, M. Bosma, V. Zhao, K. Guu, A. W. Yu, B. Lester, N. Du, A. M. Dai, and Q. V. Le.  Finetuned language models are zero-shot learners.  In _ICLR_ , 2022.
  * Williams (1992) R. J. Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Reinforcement learning_ , 1992.
  * Wu et al. (2024) Y. Wu, Z. Sun, H. Yuan, K. Ji, Y. Yang, and Q. Gu.  Self-play preference optimization for language model alignment.  _arXiv preprint_ , 2024.
  * Yang et al. (2024) J. Q. Yang, S. Salamatian, Z. Sun, A. T. Suresh, and A. Beirami.  Asymptotics of language model alignment.  _arXiv preprint_ , 2024.
  * Ziegler et al. (2019) D. M. Ziegler, N. Stiennon, J. Wu, T. B. Brown, A. Radford, D. Amodei, P. Christiano, and G. Irving.  Fine-tuning language models from human preferences.  _arXiv preprint_ , 2019.

##  Appendix A Supporting results and derivations

###  A.1 Proof of Theorem 1

Consider $N$ random generations $y_{1}$, $y_{2}$, …, $y_{N}$ from
${\pi_{\text{ref}}}$ and an arbitrary generation $y$ among them. Let
$A_{i}(y)$ denote the event that $y$ is the best sample (i.e., $r(y)\geq
r(y_{i})$ for all $i$) and that $i$ is the lowest index for which $y_{i}=y$.
It is trivial to see that the the events $\\{A_{i}(y)\\}_{i=1,2,\dots,N}$ are
disjoint and that their union corresponds to $y$ being selected by Best-of-N
sampling.

The event $A_{i}(y)$ occurs if and only if three conditions are met:
$r(y_{j})<r(y)$ for all $j<i$, $y_{i}=y$, and $r(y_{j})<r(y)$ for all $j<i$.
This allows us to derive the likelihood of the event $A_{i}(y)$:

$\displaystyle\mathbb{P}\mathopen{}\mathclose{{}\left[A_{i}(y)}\right]$
$\displaystyle=\mathopen{}\mathclose{{}\left(\prod_{j=1}^{i-1}\mathbb{P}\mathopen{}\mathclose{{}\left[r(y_{j})<r(y)}\right]}\right)\times{\pi_{\text{ref}}}(y)\times\mathopen{}\mathclose{{}\left(\prod_{j=i+1}^{N}\mathbb{P}\mathopen{}\mathclose{{}\left[r(y_{j})\leq
r(y)}\right]}\right)$
$\displaystyle={p_{<}}(y)^{i-1}\times{\pi_{\text{ref}}}(y)\times{p_{\leq}}(y)^{N-i-1}.$

The likelihood that Best-of-N sampling selects the generation $y$ is then
given by

$\displaystyle{\pi_{\text{BoN}}}(y)$
$\displaystyle=\sum_{i=1}^{N}\mathbb{P}\mathopen{}\mathclose{{}\left[A_{i}(y)}\right]$
$\displaystyle=\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[{p_{<}}(y)^{i-1}\times{\pi_{\text{ref}}}(y)\times{p_{\leq}}(y)^{N-i}}\right]$
$\displaystyle={\pi_{\text{ref}}}(y)\times\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[{p_{<}}(y)^{i-1}\times{p_{\leq}}(y)^{N-i}}\right]$
$\displaystyle={{\pi_{\text{ref}}}(y)}\times{p_{\leq}}(y)^{N-1}\times\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}.\hfill\qed$

###  A.2 Link to the continuous case

A noteworthy observation is that we can relate the Best-of-N expression to the
case of a continuous distribution, in which case the term $\mathtt{(B)}$ in
Equation 4 is constant and equal to $N$ (which is intuitively natural as
${p_{<}}(y)$ and ${p_{\leq}}(y)$ have the same value in this case).

Indeed, recall that the probability for a sequence $y$ to be drawn from the
Best-of-N distribution is

${\pi_{\text{BoN}}}(y)={\pi_{\text{ref}}}(y)\times\underbrace{{p_{\leq}}(y)^{N-1}}_{\mathtt{(A)}}\times\underbrace{\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}}_{\mathtt{(B)}}.$
(20)

Here, $y$ is a discrete variable, as it lives in $\llbracket
1;T\rrbracket^{L}$ where $T$ is the number of tokens and $L$ is the maximum
length of a sequence.

Now, we show why Equation 20 matches the classic formula for the max of $N$
continuous variables. Formally, let $X$ be a real valued random variable with
density $f_{X}$ and a cumulative distribution function $F_{X}$. Taking
$Y_{1},...Y_{N}$ i.i.d. variables with the same density, define
$X_{N}=\max\\{Y_{1},...Y_{N}\\}$ as the maximum over the $N$ variables. Then,
we have that

$F_{X_{N}}(y)=\mathbb{P}(Y_{1}\leq y,\dots Y_{N}\leq y)=F_{X}(y)^{N},$ (21)

and thus

$f_{X_{N}}(y)=f_{X}(y)F_{X}(y)^{N-1}N.$ (22)

In Equation 22, we recognize the Best-of-N formula in the case where the
correction factor $\mathtt{(B)}$ is $N$. For the term $\mathtt{(A)}$,
$F_{X}(y)$ plays the role of ${p_{\leq}}(y)$, as by definition
$F_{X}(y)=\mathbb{P}(X\leq y)$. Finally, $f_{X}(y)$ is the density of $X$,
which is analogous to the probability ${\pi_{\text{ref}}}(y)$ in the discrete
case.

###  A.3 Backward KL and policy gradient equivalence

We formally show the analogy between the gradient of the backward KL
divergence of Equation 14 and the standard (e.g., REINFORCE (Williams, 1992))
policy gradient of a KL-regularized RLHF problem with equivalent reward
$r_{\texttt{BOND}}$ and regularization $\beta_{\texttt{BOND}}$.

The exact backward KL gradient can be derived as:

$\displaystyle\nabla_{\pi}\operatorname{KL}\mathopen{}\mathclose{{}\left(\pi\mid\mid{\pi_{\text{BoN}}}}\right)$
$\displaystyle=\nabla_{\pi}\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right]$
$\displaystyle=\nabla_{\pi}\sum_{y}\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right)$
$\displaystyle=\sum_{y}\nabla_{\pi}\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right)+\pi(y)\nabla_{\pi}\log\pi(y)$
$\displaystyle=\sum_{y}\pi(y)\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right)+\pi(y)\nabla_{\pi}\log\pi(y)$
$\displaystyle=\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right)+\nabla_{\pi}\log\pi(y)}\right]$
$\displaystyle=\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right)}\right].$

Above, we have used the product rule of gradient, the rule
$\nabla_{\pi}\pi(y)=\pi(y)\nabla_{\pi}\log\pi(y)$ and the fact that
$\mathbb{E}_{y\sim\pi}\nabla_{\pi}\log\pi(y)=0$.

_Equivalence with Policy Gradient RL._ As anticipated, one can verify that
descending the above gradient is equivalent – up to a constant scaling – to
running the RL policy gradient REINFORCE algorithm on the RL objective of
Equation 1 with $r=r_{\texttt{BOND}}$ and
$\beta_{\textrm{RL}}=\beta_{\texttt{BOND}}$. Indeed, we can use the expression
for ${\pi_{\text{BoN}}}$ to break down the above gradient into:

$\displaystyle\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{BoN}}}(y)}\right)}\right]$
$\displaystyle=\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{ref}}}(y)-(N-1)\log{p_{\leq}}(y)-\log\sum_{i=1}^{N}\mathopen{}\mathclose{{}\left[\frac{{p_{<}}(y)}{{p_{\leq}}(y)}}\right]^{i-1}}\right)}\right]$
$\displaystyle=\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{ref}}}(y)-\frac{r_{\texttt{BOND}}(y)}{N-1}}\right)}\right]$
$\displaystyle=-(N-1)\underbrace{\mathbb{E}_{y\sim\pi}\mathopen{}\mathclose{{}\left[\nabla_{\pi}\log\pi(y)\mathopen{}\mathclose{{}\left(r_{\texttt{BOND}}(y)-\beta_{\texttt{BOND}}\mathopen{}\mathclose{{}\left(\log\pi(y)-\log{\pi_{\text{ref}}}(y)}\right)}\right)}\right]}_{\text{gradient
used by REINFORCE}}.$

###  A.4 Derivation of J-BOND reward

Here we provide a theoretical explanation behind the design of the J-BOND
reward function discussed in Section 5:

$r_{\texttt{J-BOND}}(y)=\begin{cases}-\log(16)&\text{if}\quad
r(y)\leq\text{min}\\{r(y_{1}^{\prime}),r(y_{2}^{\prime})\\}\\\
0&\text{otherwise}\end{cases}.$

Recall that $r_{\texttt{J-BOND}}(\cdot)$ is meant to approximate the the true
reward function $r_{\texttt{BOND}}(y)=\log{p_{\leq}}(\cdot)$ which is unknown
since ${p_{\leq}}(\cdot)$ in general requires knowing the reward distribution
of $\pi_{\text{anchor}}^{t}$ (in J-BOND, we only take 2 samples
$y_{1}^{\prime},y_{2}^{\prime}\sim\pi_{\text{anchor}}^{t}$).

As mentioned in Section 5, we designed $r_{\texttt{J-BOND}}(\cdot)$ to assign
a negative reward only if sample $y$ is worse than both the anchor samples, to
mimic the concavity of the log quantile $\log{p_{\leq}}(\cdot)$. In practice,
we did not observe gains when rewarding also the intermediate case. The
particular choice of value $-\log(16)$ is motivated by the following main
reason.

We want that, when sample $y$ has median reward compared to the anchor
rewards’ distribution (i.e., ${p_{\leq}}(y)=0.5$), then – in expectation –
$r_{\texttt{J-BOND}}(y)$ coincides with the true reward
$r_{\texttt{BOND}}(y)=\log{p_{\leq}}(\cdot)=\log(0.5)$. For this purpose, let
us consider the parametrized function:

$r_{\texttt{J-BOND}}^{\alpha}(y)=\begin{cases}\alpha&\text{if}\quad
r(y)<\text{min}\\{r(y_{1}^{\prime}),r(y_{2}^{\prime})\\}\\\
0&\text{otherwise}\end{cases}.$

Note that the stochasticity of $r_{\texttt{J-BOND}}^{\alpha}(y)$ is due to the
2 random anchor samples $y_{1}^{\prime},y_{2}^{\prime}$ and its expectation
can be computed as:

$\displaystyle\mathbb{E}_{y_{1}^{\prime},y_{2}^{\prime}\sim\pi_{\text{anchor}}^{t}}\mathopen{}\mathclose{{}\left[r_{\texttt{J-BOND}}^{\alpha}(y)}\right]$
$\displaystyle=\alpha\cdot\mathbb{P}\mathopen{}\mathclose{{}\left[\\{r(y_{1}^{\prime})>r(y)\\}\cap\\{r(y_{2}^{\prime})>r(y)\\}}\right]+0\cdot\mathbb{P}\mathopen{}\mathclose{{}\left[\text{"otherwise"}}\right]$
(23) $\displaystyle=\alpha\cdot(1-{p_{\leq}}(y))^{2},$ (24)

where we have used the definition of
${p_{\leq}}(y)=\mathbb{P}_{y^{\prime}\sim\pi_{\text{anchor}}^{t}}\mathopen{}\mathclose{{}\left[r(y^{\prime})\leq
r(y)}\right]$. Using the expression above, we can find the $\alpha$ for which
$\mathbb{E}_{y_{1}^{\prime},y_{2}^{\prime}\sim\pi_{\text{anchor}}^{t}}\mathopen{}\mathclose{{}\left[r_{\texttt{J-BOND}}^{\alpha}(y)}\right]=r_{\texttt{{BOND}\@}}(y)$
when ${p_{\leq}}(y)=0.5$:

$\alpha\cdot(1-0.5)^{2}=\log(0.5)\quad\rightarrow\alpha=-\log(16).$

We illustrate this in Figure 8, where we plot the expected
$r_{\texttt{J-BOND}}(y)$ reward and the true reward $r_{\texttt{{BOND}\@}}(y)$
as a function of ${p_{\leq}}(\cdot)$.

Figure 8: Expected value of the J-BOND reward, i.e.,
$\mathbb{E}_{y_{1}^{\prime},y_{2}^{\prime}\sim\pi_{\text{anchor}}^{t}}\mathopen{}\mathclose{{}\left[r_{\texttt{J-BOND}}(y)}\right]$
(see Equation 23), compared to the true (unknown) reward $\log{p_{\leq}}(y)$,
as a function of the quantile ${p_{\leq}}(y)$. By design of
$r_{\texttt{J-BOND}}$, the two curves coincide when ${p_{\leq}}(y)=0.5$, i.e.,
when $y$ has median reward w.r.t. the anchor distribution.

##  Appendix B Additional Experiments

###  B.1 Learned quantile models

Monte-Carlo quantile estimation (Section 4.1) approximates the reward
quantiles by sampling multiple times from the reference policy
${\pi_{\text{ref}}}$, for each observed context. While we found it to be very
simple and effective, it may require many sample for an accurate quantile
estimation and, in addition, it does not exploit any information about the
given context. For instance, assuming we have a good quantile estimation for
context $x$ and are presented a new context $x^{\prime}$. MC quantile
estimates treat $x^{\prime}$ independently from $x$, although they may have
very similar reward quantiles.

Motivated by this, in this section we explore an alternative approach that
aims at _learning_ a context-dependent quantile estimator
$\widehat{{p_{\leq}}}_{\theta}(\cdot)$, parametrized by parameter $\theta$.
The idea is to view quantile ${p_{\leq}}(y)$ as the parameter of a Binomial
random variable $Z$ where $Z=\mathbb{I}\\{r(y_{i})\leq r(y)\\}$ for
$y_{i}\sim{\pi_{\text{ref}}}$. Under such a view, we can interpret
$\widehat{{p_{\leq}}}_{\theta}(\cdot)$ as the output of a binary classifier
and train it via maximum likelihood estimation using the standard binary
cross-entropy loss (Cover, 1999):

$L(\theta,x,y)=-\mathbb{E}_{y^{\prime}\sim{\pi_{\text{ref}}}(x)}\mathopen{}\mathclose{{}\left[\log\widehat{{p_{\leq}}}_{\theta}(x,y)\mathbb{I}_{\\{r(x,y^{\prime})\leq
r(x,y)\\}}+\log(1-\widehat{{p_{\leq}}}_{\theta}(x,y))\mathbb{I}_{\\{r(x,y^{\prime})>r(x,y)\\}}}\right].$
(25)

We test such an approach in the abtractive summarization task considered in
Section 4. We parametrize $\widehat{{p_{\leq}}}_{\theta}(\cdot)$ with a LLM
initialized as ${\pi_{\text{ref}}}$ and fine-tuned using the loss of Equation
25. Simultaneously, the policy $\pi_{t}$ is fine-tuned using BOND and
utilizing $\widehat{{p_{\leq}}}_{\theta_{t}}(\cdot)$ as quantile estimator at
each step. Notably, we approximate the expectation in Equation 25 via a
_single_ sample from ${\pi_{\text{ref}}}$ for each prompt in the batch. In
Figure 9 we report the backward and forward $\mathrm{KL}$ divergences between
the training policy and the ${\pi_{\text{BoN}}}$ distribution as well as the
average log quantiles, and compare them to the ones obtained when running BOND
with MC quantile estimation. Note that in both cases, ${\pi_{\text{BoN}}}$ is
approximated by 32 MC samples during evaluation. When using the learned
quantile model $\widehat{{p_{\leq}}}_{\theta}(\cdot)$, we observe BOND
achieves very comparable KL divergences and log quantiles compared to using MC
quantile estimaton. This illustrates that the use of learned quantiles is
valid and promising, potentially offering interesting computational advantages
in situations where, e.g., $\widehat{{p_{\leq}}}_{\theta}(\cdot)$ can be re-
used or learned offline with a fixed sample budget.

Figure 9: BOND (with $N=8$) using MC quantile estimates vs. a Learned quantile
model.

Finally, we remark that the explored approach is quite naive, and alternative
learned quantile models can definitely be derived, e.g., further enforcing
ordering in the predicted quantiles, using quantile regression (Dabney et al.,
2017), or assuming a pre-specified (e.g., Gaussian) rewards’ distributions.

###  B.2 Additional plots for BOND with Jeffreys divergence objective

We provide additional experiments that complement the ablation of Section 4.2
when running BOND with a Jeffreys divergence objective
$J_{\text{effreys}}^{\beta}$ for different values of $\beta$. In particular,
in Figure 10 we show BOND results when using $N=4$ (top plots) and $N=16$
(bottom plots). Results are consistent with what discussed in Section 4.2:
When using $\beta=0.5$, BOND minimizes both backward and forward KL
divergences from ${\pi_{\text{BoN}}}$ compared to when solely one of them is
used as BOND objective. Moreover, it optimizes the log quantiles similarly to
$\beta=1$.

Figure 10: BOND with $N=4$ (top plots) and $N=16$ (bottom plots), and
different values of $\beta$ for the Jeffreys divergence objective (_cf._
Equation 11). When using $\beta=0.5$ (Jeffreys divergence), BOND minimizes
both backward (left plots) and forward (middle plots) KL divergences from
${\pi_{\text{BoN}}}$, achieving best of both objectives ($\beta=0$ and
$\beta=1$). Moreover, it optimizes the reward quantiles (right plots)
significantly more than when using $\beta=0$.