BOND: Aligning LLMs with Best-of-N Distillation

## Abstract

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 𝑁 N italic_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 KL KL \mathrm{KL} roman_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:

## 1 Introduction

State-of-the-art large language models (LLMs) such as Gemini (Gemini Team, ; Reid et al., ) and GPT-4 (OpenAI, ) are generally trained in three stages. First, LLMs are pre-trained on large corpora of knowledge using next-token prediction (Radford et al., , ). Second, the pre-trained models are fine-tuned to follow instructions via supervised fine-tuning (SFT) (Raffel et al., ; Wei et al., ). Lastly, reinforcement learning from human feedback (RLHF) (Christiano et al., ; Ziegler et al., ; Stiennon et al., ) is used to further increase the quality of generations. The RLHF step generally consists of learning a reward model (RM) (Ouyang et al., ) 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., ), notably since it can cause *forgetting* (French, ) of pre-trained knowledge, and since loopholes in the RM (Clark and Amodei, ; Pan et al., ) can cause *reward hacking* (Askell et al., ; Skalse et al., ).
The standard strategy is to use policy-gradient methods (Williams, ) 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., ) 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., ). 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., ; Gao et al., ; Touvron et al., ) 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., ) task.
Then, in Section 6, we apply J-BOND to align Gemma (Gemma Team, ) 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., ).

Figure: 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*.
Refer to caption: x1.png

## 2 Problem Setup

We consider a LLM based on the transformer (Vaswani et al., ) 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., ; Lazaridou et al., ), reducing forgetting (French, ) and reward hacking (Skalse et al., ).
Equation 1 is usually optimized with online algorithms, as they perform better than their offline counterparts (Tang et al., ).
Moreover, simple methods have demonstrated the best results, e.g., REINFORCE (Williams, ) with a sampled baseline for variance reduction (Li et al., ; Ahmadian et al., ) outperform PPO (Schulman et al., ).

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., ) when looking at reward/$\mathrm{KL}$ trade-offs, and comes with theoretical guarantees (Qiping Yang et al., ) 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$(^1^11To 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).

###### Theorem 1 .

### 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, ). 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. (); Rafailov et al. ()) 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.
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.
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, ) 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., ). 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, ) 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, ) 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., ) with ${\pi_{\text{ref}}}$ being a T5 supervised fine-tuned policy and $r(\cdot)$ being a T5 NLI reward model (Roit et al., ). 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: 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.
Refer to caption: x2.png

### 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., ). *(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-N^2 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-N^M 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: 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.
Refer to caption: x5.png

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}}}$.

Figure: Algorithm 1 Iterative BOND (meta algorithm)

Figure: 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}}}$.
Refer to caption: x6.png

## 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.

Figure: Algorithm 2 The J-BOND algorithm

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 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., ), 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 extra(^2^22Note 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: 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).
Refer to caption: x9.png

Setup. We consider Gemma (2B and 7B) models (Gemma Team, ) 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, ) 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: 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.
Refer to caption: x12.png

Figure: 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.
Refer to caption: x15.png

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, ) with 2 policy samples per-prompt and a leave-one-out baseline (Ahmadian et al., ) 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, ), RecurrentGemma 2B and 9B (Botev et al., ) as well as CodeGemma 1.1 (CodeGemma Team, ). 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 (, Table 5) for more details).

## 7 Related Work

Best-of-N was introduced in  Stiennon et al. () 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, ; Beirami et al., ), provided an estimator for the average Best-of-N reward (Nakano et al., ), made theoretical connections with $\mathrm{KL}$-constrained RL (Yang et al., ) and provided scaling laws for Best-of-N alignment (Gao et al., ).

Matching Best-of-N for improved alignment is a strategy that was studied in different flavors in the literature. Dong et al. () and Touvron et al. () 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. () 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. () use pairwise AI feedback on online generations to obtain online preferences that are then optimized, while Calandriello et al. () use a dedicated preference reward model instead.
Concurrently and closest to our work, Amini et al. () 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. ().

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., ; Pinto et al., ), for online RLHF (Ahmadian et al., ), offline RLHF (Flet-Berliac et al., ) and preference optimization (Wu et al., ).

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., ; Gorbatovski et al., ; Ramé et al., ), 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., ).
Specifically, BOND incorporates the principles of iterated amplification (Christiano et al., ; Cotra, ), 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., ) and WARP (Ramé et al., ), 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.

## 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, )) 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: 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.
Refer to caption: x18.png

## 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, ):

$$ $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: Figure 9: BOND (with $N=8$) using MC quantile estimates vs. a Learned quantile model.
Refer to caption: x19.png

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., ), 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: 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$.
Refer to caption: x22.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: How does J-BOND performance compare to (Amini et al, 2024) and other concurrent works?

There are three algorithms discussed in the paper, namely BOND, Iterative-BOND and J-BOND. Is it always preferable to use JBOND or do you recommend using each algorithm in particular situations?

Will the code be made publicly available to serve the research community?
soundness: 3
strengths: The paper makes multiple contributions, namely theoretical derivation for the Best-of-N distribution and a practical RLHF finetuning algorithm that distills the Best-of-N distribution into a policy which is sample efficient and requires just one single sample at inference time

The authors are making a lot of engineering design choices in their proposed model, and carefully analyze the role of each component in the performance of the proposed algorithm

To regularize the model and ensure it is not steering too far from the reference model (the supervised finetuned policy), the authors use a combination of both forward and reverse KL, namely Jeffrey divergence. While the forward KL ensures mode covering behavior, the reverse KL is used for mode seeking behavior; their combination results in better aligned policies that combine the advantages of both divergences

Applying BOND recursively (Iterative BOND) improves the sample efficiency of the BOND algorithm and works for very small values of n (2, 4); its reward/KL tradeoff is comparable to the non-iterative BOND while being more sample efficient

The J-BOND algorithm presents better reward/KL trade-off compared to the REINFORCE algorithm with different values of \beta and does not require using a specific regularization strength

The paper is well written, well-motivated, presents theoretically and experimentally sound insights that would benefit the research community
summary: The paper is focusing on the RLHF alignment problem, in particular on emulating the Best-of-N distribution which is known to perform very well, but is very costly at inference time (for each prompt it requires drawing N candidate generations from the reference model and selecting the one with highest reward according to a reward model). The authors propose the BOND (Best-of-N Distillation) algorithm designed to force the distribution of generations from the finetuned policy to be close to the Best-of-N distribution, requiring the generation of just a single sample (instead of N). To this end, BOND regards the alignment problem as a distribution matching problem and distills the Best-of-N distribution by finetuning the reference policy to imitate the Best-of-N distribution. To stay close to the original reference model, the authors incorporate a KL regularization term that considers both the forward and backward divergence (Jeffrey divergence). In addition, they incorporate Monte-Carlo quantile estimation, and exponential moving anchor, resulting in the J-BOND algorithm. The authors conduct experiments on the abstractive summarization task (XSum dataset) and aligning GEMMA using J-BOND.
weaknesses: The paper combines a lot of distinct ideas already proposed in previous works - it would be good to actually clearly articulate what the novel contribution is. Besides, the comparison with concurrent works is not very clear, in particular the difference with (Amini et al, 2024), WARM, WARP (Rame et al, 2024). 

Figure 4 - It would be interesting to see how the performance of Best-of-N compares to the proposed algorithm J-BOND and REINFORCE

Algorithm 1, line 330 - \pi_t is not defined

Line 329 - \pi \in Capital \pi -Capital \pi  in Algorithm 1 is not defined

Line 456 - “a large previously trained reward model r(.)” - please provide details

Lines 481-482 - there are not many details about the hyperparameters of the REINFORCE algorithm 

The authors are conducting experiments on Gemma 2B and 7B models, while results are convincing it would be good to see if they hold with other models and other tasks than summarization

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: In Section 4.2, the authors utilize 32 Monte Carlo samples to estimate the backward and forward KL divergences between the training policy and the reference distribution. Given the high dimensionality of these distributions, this sample size seems insufficient to reliably capture the divergence and may introduce substantial estimation variance. A sensitivity analysis showing how the estimator's variance changes with an increasing number of Monte Carlo samples would strengthen the results. Alternatively, using a larger sample size for these estimates could enhance the reliability of the reported divergences.

While BOND’s benefits, such as improved KL/reward trade-offs and dynamic regularization, are discussed in the body of the paper, they are not clearly summarized in the introduction or abstract. A brief overview in these sections would effectively communicate BOND’s main advantages over traditional RLHF approaches, aiding readers in understanding its unique contributions and practical value.
soundness: 4
strengths: The paper is well-structured and clearly presents its methodology, with detailed explanations and algorithms that allow readers to follow the progression. From iterative BOND to the addition of KL regularization in Sections 4 and 5, the additional experimental results effectively support these methodological advancements. 
BOND is notable for its originality, offering a practical and computationally efficient alternative to traditional RLHF that achieves a superior KL-reward balance without requiring commitment to a specific regularization level. The work is significant in its potential impact on RLHF practices, as it provides a scalable solution for optimizing performance and efficiency while minimizing trade-offs between KL divergence from the reference distribution and reward maximization.
summary: The paper introduces Best-of-N Distillation (BOND), a novel alignment tuning algorithm designed to emulate the Best-of-N sampling method in a more computationally efficient manner. BOND aims to achieve the same high-quality output as Best-of-N sampling without the inference-time computational overhead, by aligning model outputs to match the distribution of the Best-of-N candidates. 
In addition, to ensure stability and scalability, the authors introduce an iterative approximation strategy that operates effectively even with a minimal sample size (e.g., 2 or 4). 
Further, based on the two types of loss function derivation aiming forward and reverse KL respectively, the author leverages Jeffreys divergence proposing J-BOND, an enhanced algorithm incorporating iterative distribution matching with an exponential moving average (EMA) anchor. J-BOND demonstrates effectiveness in maintaining a stable training and superior KL-reward trade-off through experiments.
weaknesses: The paper relies heavily on the Jeffreys divergence without sufficient comparative analysis against alternative divergence metrics. The mode-covering and mode-seeking behavior property paper mentioned about are only observed in lower dimension such as multimodal distribution in 1-dimension. An inclusion of other divergence types, especially in the iterative stages, could offer clearer insights into the unique advantages of Jeffreys divergence. Further, relevant literature on divergence measures in alignment tuning, should also be cited to contextualize this choice.
- Go, Dongyoung, et al. "Aligning language models with preferences through f-divergence minimization." Proceedings of the 40th International Conference on Machine Learning. 2023.

As the paper discusses the method’s efficiency, the paper would benefit from explicit comparison of the computational cost saved by BOND relative to traditional Best-of-N sampling, or comparisons with sampling approaches used in RLHF. This would clarify BOND’s potential advantages in real-world applications.

Additionally, while the paper addresses the challenge of sampling size N through iterative approximation, showing practical advantages like non-saturation compared to non-iterative BOND, this helpful randomness raised in iterative BOND is solely introduced by approximation randomness, which lacks controls or specific directions. This calls into question whether the proposed algorithm genuinely achieves a distilled $\text{Bo}N^M$ distribution. 
The substantial difference in $r(y)$ between iterative and non-iterative BOND  in Figure3 suggests a potential vulnerability to reward hacking, as discussed in Gao et al.
- Gao, Leo, John Schulman, and Jacob Hilton. "Scaling laws for reward model overoptimization." International Conference on Machine Learning. PMLR, 2023.


The introduction combines related works and problem setup, which could be structured more effectively. Detailed discussions on RLHF and Best-of-N would be more suitable in a separate related works section or could be incorporated into the problem setup. In the introduction, it would be clearer to emphasize the limitations of existing methods and highlight the advantages of the proposed approach over current methods.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. Can the authors compare with more fundamental baseline methods, such as BoN or other BoN distillation algorithms?

2. Can the authors supplement additional experiments, including downstream validation and more ablation studies as discussed in Weakness 3?

3. Can the authors prove more clearly the advantages of the BOND algorithm over other Alignment algorithms, in terms of both performance and efficiency, to make the argument more convincing?
soundness: 3
strengths: 1. Rigorous Theoretical Analysis: This work rigorously analyzes the distribution characteristics under Best-of-N sampling and establishes its connection with standard RLHF, as well as the specific reward value $r_{BOND}$ under this correlation. This provides a reliable theoretical foundation for the work, rather than being based on naive assumptions.

2. Some Degree of Novelty: Although there is some concurrent work, the idea of distilling distributions from Best-of-N is fairly novel and important.

3. Consideration of Practical Efficiency: I appreciate the authors' consideration of the practical efficiency of the algorithm. The proposed J-BOND algorithm theoretically has lower sampling complexity, which should increase the efficiency of RLHF.
summary: This paper introduces a distribution matching-based Best-of-N distillation method that simulates the Best-of-N distribution space, while reducing the time overhead of N inferences to just one. Starting from the theoretical distribution of BoN, the authors construct the Iterative BOND algorithm based on Quantile estimation and the choice of Jeffreys Divergence, and further propose the more practically meaningful J-BOND algorithm.
weaknesses: 1. Lack of Important Baselines: Given that the main purpose of the paper is to distill Best-of-N sampling, BoN performance should straightforwardly serve as an important baseline to analyze pros and cons in terms of performance and efficiency. Moreover, other concurrent BoN distillation algorithms [1] should also be considered.

2. Lack of Downstream Validation: The main metrics in the paper, such as reward value and KL divergence, cannot be directly equated to the model's performance on downstream tasks. For an RLHF method, it is necessary to conduct experiments on downstream tasks and present more intuitive metrics to demonstrate the model's alignment performance.

3. Insufficient Experimental Setup: The paper lacks exploration of several issues. For instance, BoN sampling heavily depends on the Reward Model, and the influence of different RMs on the BOND algorithm is not investigated. Additionally, a more nuanced exploration of Jeffreys Divergence with smoother β variations could be included; and the comparison between J-BOND and standard Iterative BOND lacks investigation.

[1] Variational Best-of-N Alignment, https://arxiv.org/pdf/2407.06057

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: Would the curves in Figure 4(a) converge to similar spots if you run the algorithms long enough?

Do the authors expect that the conclusions would be similar for much larger models?

Writing: 
- pi(x,y) makes it look like a joint distribution; do the authors mean pi(y|x)?
- 3e-6 means 0.000003; 3e^(-6) means 0.007. Which one are you referring to?
soundness: 3
strengths: The motivation of distillation is great and this direction should be deeply explored. The paper is written carefully. The algorithm is clearly written. I checked the derivations and they seem correct so far.

REINFORCE is a natural baseline and the authors have attempted multiple beta values for the baseline.
summary: The paper is essentially distilling inference-time best-of-N sampling into training time. Specifically, the authors propose to train the policy to match the best-of-N distribution (which is an analytical form derived by the authors). The distribution matching is done through minimizing a combination of forward and backward KL. The behavior of J-BOND appears better for reward vs. num steps as well as KL(policy || ref policy) vs. num steps, compared to REINFORCE with various beta values.
weaknesses: Have the authors tried reward shaping techniques for RLHF baseline, e.g., making the reward values more peaky – either very large or very small? 

I’d appreciate a more comprehensive discussion on how much the authors expect this technique to benefit downstream tasks.

It’ll be great if the authors can include more discussion on whether the baseline B is important or not in the algorithm.

What ablations on beta and gamma in Algorithms 2 (balancing among forward KL, backward KL, additional regularization) would likely benefit downstream tasks more? It's still unclear to me why we want to put an equal/large weight on backward KL. More motivation would be nice.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 2
questions: #### Question 1:
Line 298, why is the proposed approach "equivalently to distilling a Best-of-$N^M$ of the initial distribution $\pi_{ref}$"? Is this a qualitative or quantitative statement?
#### Question 2:
Figure 4, seems like a linear relationship between # steps and r(y) for iterative BOND.
For Figure 5, it seems like another kind of trending between $r(y)$ and # steps or $KL(\pi||\pi_{ref})$.
Figures 6 and 7, present a log-like trend between the KL and reward similar to the REINFORCE algorithm.
A similar trend is also observed in http://arxiv.org/abs/2406.00832 and http://arxiv.org/abs/2407.06057.
As discussed in http://arxiv.org/abs/2204.05862 and my experience, there can be an approximately linear relationship between $\sqrt{D_{KL}}$  and $r$ for BoN sampling. 
It would be interesting if the author could provide some empirical or theoretical intuition about such a relationship.
Does this indicate that even though performing a BoN distribution match, it is still more similar to the general policy-gradient RL algorithm (which may try to match another distribution)?
#### Question 3:
Why is there an approximately linear relationship between #steps and KL in, as presented in Figures 5 and 6 for the BOND algorithm?
For the REINFORCE algorithm, it seems a quite different trend.
#### Question 4:
Figure 4 presents consistent trending between KL and $\log_{p\le(y)}$ across varying $N$ and algorithm, which is quite interesting.
Is there any explanation for this phenomenon?
#### Question 5:
Is there any analysis or comparison of the reward overoptimization of this algorithm?
soundness: 4
strengths: 1. This work implements the LLM alignment problem through Best-of-N distillation, which can be a sound direction for the development of algorithms in this field.
2. This work formulates and discusses the BoN distribution and its relationship with the general RLHF target. 
3. This work proposes to utilize Jeffreys divergence to balance the mode-covering and mode-seeking behavior introduced by forward- and backward-KL- KL optimization.
4. This work further integrates their method with EMA techniques and proposes an iterative version.
summary: This article models the alignment of LLMs as a distribution-matching problem. Specifically, it considers the distribution induced by Best-of-N sampling as the target distribution and optimizes the Jeffreys divergence with respect to it for balancing the characteristics of forward & backward KL-based optimization. Additionally, this work derives an iterative form, updating the reference model during training.
weaknesses: 1. My biggest concern is the absence of direct comparison with close-related works http://arxiv.org/abs/2407.06057, http://arxiv.org/abs/2406.00832. It would be much more convincing if there is a clear comparison like Figure 1 in http://arxiv.org/abs/2407.06057, rather than a simple baseline REINFORCE presented in Figure 7 in this work.
2. More discussion about why BoN distribution as the target can be included.
3. Other possible additional analyses are discussed in the Questions.

Meta-review: This paper proposes to align an LM based on the best-of-N distribution. This paper is backed up by formal mathematical characterizations (which may be useful for others working with these types of distributions), as well as solid empirical results. There were some concerns with regard to justifying particular choices (e.g., Jeffrey's divergences), but I think this paper is a clear accept.

Decision: Accept (Poster)