BOND (Best-of-N distillation), explained

BOND (Best-of-N Distillation) is an RLHF method that trains a policy to produce, in one sample, the distribution of answers you would get by drawing $N$ samples from the reference model and keeping the one a reward model scores highest. Best-of-N sampling is one of the strongest and simplest ways to spend a reward model at inference time, but it multiplies serving cost by $N$. BOND moves that cost into training. It treats alignment as distribution matching: derive the Best-of-N distribution analytically, then minimise a divergence between the policy and that target, using a mix of supervised learning on Best-of-N samples and a policy-gradient term. Its practical form, J-BOND, distils Best-of-2 of a slowly moving anchor over and over, which compounds into Best-of-$2^M$ without ever choosing $N$ or a KL coefficient. It sits in the post-training stack where PPO-style RLHF or REINFORCE with a leave-one-out baseline would, after SFT and reward modelling.

The problem and where it came from#

Best-of-N has been part of RLHF since Stiennon et al. (2020) used it as a baseline for summarisation, and Gao, Schulman and Hilton (2022) used it to measure reward over-optimisation because it is so well behaved: the policy it defines is a simple reweighting of the reference, and its KL from the reference grows only logarithmically in $N$. Fine-tuning on Best-of-N outputs (the forward-KL direction, as in rejection-sampling fine-tuning) was already common, but it only covers the mode-covering half of the problem.

BOND was introduced by Sessa, Dadashi, Hussenot, Ferret, Vieillard, Ramé and colleagues at Google DeepMind (2024). The paper makes four contributions: an exact expression for the Best-of-N distribution, the observation that this distribution is the solution of a KL-regularised RL problem with a specific reward and coefficient, the use of the Jeffreys divergence (a mix of forward and backward KL) as the matching objective, and an iterative scheme with a moving anchor that reaches large effective $N$ cheaply. They reported a better KL-reward Pareto front than REINFORCE on Gemma 7B and stated that J-BOND was used to fine-tune Gemma 1.1, RecurrentGemma and CodeGemma 1.1. Concurrently, Amini, Vieira, Ash and Cotterell (2024) proposed variational BoN, which matches the same target with the backward KL only.

The objective, with every symbol defined#

The Best-of-N distribution

Fix a prompt and drop it from the notation. ${\pi }_{\text{ref}}$ is the reference policy, $r(y)$ is the reward of a generation $y$, and ties are broken by any fixed strict order. Define the probability that a fresh reference sample is strictly worse than $y$, and the probability that it is not better:

Best-of-N returns $y$ exactly when the best of $N$ independent samples is $y$. The probability that all $N$ samples are at most as good as $y$ is $p_{\le }(y{)}^N$, and that all are strictly worse is $p_{<}(y{)}^N$; the difference is the probability that the maximum is $y$. BOND's Theorem 1 writes the same quantity in a factored form:

Term (A) is an exponential penalty on generations with low reward quantile. Term (B) is a correction for the chance of drawing the same generation more than once; it lies in $[1,N]$ and equals $N$ for continuous distributions, so for long text outputs it is essentially a constant.

Best-of-N is a KL-regularised RL solution

The policy that maximises ${\mathbb{E}}_{\pi }[r^{\prime }(y)]-\beta \mathrm{KL}(\pi \Vert {\pi }_{\text{ref}})$ for any reward $r^{\prime }$ is ${\pi }^{\star }(y)\propto {\pi }_{\text{ref}}(y)\exp (r^{\prime }(y)/\beta )$, the same closed form that DPO builds on. Taking logs of Theorem 1 shows that ${\pi }_{\text{BoN}}$ has exactly this form with

Two consequences follow. $N$ plays the role of the KL coefficient, so choosing $N$ is choosing how far the policy may move. And the effective reward is the log of the reward quantile, which depends only on how $y$ ranks against reference samples: it is invariant to any monotone transformation of the reward model's scores, and because the log is concave it punishes being bad much more than it rewards being excellent. The authors conjecture both properties help against reward hacking. Beirami et al. (2024) showed that the widely quoted value $\log N-(N-1)/N$ for $\mathrm{KL}({\pi }_{\text{BoN}}\Vert {\pi }_{\text{ref}})$ is an upper bound, tight only when individual generations have negligible probability.

The Jeffreys objective and its two gradients

BOND minimises the generalised Jeffreys divergence with weight $\beta \in [0,1]$ (a different $\beta$ from the RL coefficient):

The forward term is supervised learning on Best-of-N samples, mode-covering:

The backward term is a policy gradient, mode-seeking, with the BOND reward and an implicit KL penalty to the reference:

The quantile is estimated by Monte Carlo: draw $k$ reference samples and set ${\hat{p}}_{\le }(y)=\frac{1}{k}{\sum }_{i}1[r(y_i)\le r(y)]$. The paper's XSum experiments used $k=16$, and the correction term (B) was dropped because it made no measurable difference. Forward KL alone spreads mass and lags on reward; backward KL alone collapses entropy; $\beta =0.5$ reduced both divergences to ${\pi }_{\text{BoN}}$ and raised reward quantiles about as fast as the pure backward objective.

Iterating with a moving anchor, and J-BOND

Best-of-N of a Best-of-N distribution is Best-of-$N^2$ of the original. So instead of picking a large $N$, BOND can repeatedly distil Best-of-$n$ of an anchor policy and then move the anchor to the current policy: after $M$ rounds the target is Best-of-$n^M$ of ${\pi }_{\text{ref}}$. J-BOND fixes $n=2$ and replaces the periodic anchor swap by an exponential moving average of the policy weights, ${\theta }_{\text{anchor}}\leftarrow (1-\eta ){\theta }_{\text{anchor}}+\eta \theta$. With only two anchor samples per prompt, $\log {\hat{p}}_{\le }$ is too noisy, so the backward term uses a calibrated surrogate:

A generation is punished only when it loses to both anchor samples, which mimics the concavity of $\log p_{\le }$. The constant makes the surrogate exact in expectation for a median generation: if $p_{\le }(y)=0.5$, the chance of losing to both anchors is $0.25$ and $0.25\times (-\log 16)=-0.693=\log 0.5$. The full update is $(1-\beta )G_{\text{FW}}+\beta G_{\text{BW}}+\gamma G_{\text{Reg}}$, where $G_{\text{FW}}$ is the negative log-likelihood gradient on the better of the two anchor samples, $G_{\text{BW}}$ is a policy gradient with return $r_{\text{J-BOND}}(y)-(\log {\pi }_{\theta }(y)-\log {\pi }_{\text{anchor}}(y))$ and a leave-one-out baseline, and $G_{\text{Reg}}$ is an extra KL penalty toward the anchor with weight $\gamma \ge 0$.

Worked example: Best-of-2 on four answers#

Suppose a prompt has four possible answers with reference probabilities $0.4,0.3,0.2,0.1$ and rewards $1,2,3,4$, so the best answer is the least likely. The cumulative probabilities are $p_{\le }=(0.4,0.7,0.9,1.0)$ and $p_{<}=(0,0.4,0.7,0.9)$. For $N=2$:

The worst answer drops from 40% to 16% (it must be drawn twice), the best almost doubles, and the expected reward rises from 2.00 to 2.54. The KL from the reference is $\sum {\pi }_{\text{Bo2}}\log ({\pi }_{\text{Bo2}}/{\pi }_{\text{ref}})=0.157$ nats, below the textbook $\log 2-\frac{1}{2}=0.193$ because these answers have large individual probabilities. For $N=4$ the same formula gives $(0.026,0.215,0.416,0.344)$, expected reward 3.08 and KL 0.587 against the bound 0.636. Term (B) matters at this scale: dropping it and renormalising ${\pi }_{\text{ref}}{p}_{\le }^{N-1}$ would give $(0.246,0.323,0.277,0.154)$ for $N=2$, a visibly different distribution. For real text, where no single answer has meaningful probability, (B) is a constant and can be ignored, which is what the paper does.

Now the J-BOND surrogate for the second answer, whose quantile is $p_{\le }=0.7$. It loses to both anchor samples when both are strictly better, probability $(1-0.7{)}^2=0.09$, so its expected surrogate reward is $0.09\times (-2.773)=-0.250$, against the ideal $\log 0.7=-0.357$. The surrogate is exact at the median and somewhat optimistic above it, which is a deliberate trade for using only two anchor samples.

A toy run: exact Best-of-N, J-BOND and REINFORCE#

We built a single-prompt problem with 200 possible responses, reference logits drawn from $\mathcal{N}(0,1.5^2)$ and rewards from $\mathcal{N}(0,1)$ (numpy, seed 0), so the reference has entropy 4.58 nats and expected reward $-0.128$. First we checked Theorem 1: for $N\in \{2,4,8,16,64\}$ the analytic ${\pi }_{\text{BoN}}$ matched 200,000 simulated Best-of-N draws to a total variation distance between 0.009 and 0.0015, which is sampling noise. The exact KLs were 0.193, 0.636, 1.203, 1.826 and 3.071 nats against bounds of 0.193, 0.636, 1.204, 1.835 and 3.175: tight for small $N$ and loose by 0.1 nats at $N=64$, where the best responses start to repeat. Expected reward rose to 0.468, 1.052, 1.603, 2.099 and 2.771.

Then we trained a tabular softmax policy with sampled gradients, batch 64, learning rate 0.02, 6,000 steps, using J-BOND as specified above ($\beta =0.5$, EMA $\eta =0.02$, $\gamma =0$) and, for comparison, REINFORCE on $r-{\beta }_{\text{RL}}\log (\pi /{\pi }_{\text{ref}})$ with a leave-one-out baseline.

Three results are worth stating plainly. First, J-BOND traced the Best-of-N curve: it passed through KL 0.34 at reward 0.65 (step 1,500) and KL 1.67 at reward 1.80 (step 2,000), close to and slightly below the exact Best-of-N points, and both sit below the optimal KL-regularised frontier $\pi \propto {\pi }_{\text{ref}}{e}^{r/\beta }$, which by construction is the best reward for a given KL. Second, it never stopped: by step 6,000 it had reached KL 3.21 and reward 2.535 with an expected log-quantile of $-0.017$, essentially the top response, because a moving anchor keeps compounding the effective $N$. The KL budget is set by when you stop, not by a coefficient. REINFORCE with ${\beta }_{\text{RL}}=1$ instead converged toward its fixed point (KL 0.73, reward 1.25 at step 6,000), and with ${\beta }_{\text{RL}}\in \{0.01,0.1\}$ it went to the top response within about 1,500 steps. Third, the invariance is real and exact: multiplying every reward by 10 produced a bit-identical J-BOND trajectory, while REINFORCE at ${\beta }_{\text{RL}}=0.1$ with the scaled reward jumped to KL 3.17 within 250 steps. In this toy the reward-KL frontier of J-BOND was not better than REINFORCE's; what it bought was robustness to the reward scale and the absence of a coefficient to tune. The paper's reported Pareto improvement on Gemma comes from a setting (a learned reward model, long generations, many prompts) that a tabular toy cannot reproduce.

The algorithm step by step#

inputs: prompts D, SFT policy pi_ref, reward model r, beta=0.5, eta=0.02, gamma>=0
pi = copy(pi_ref); anchor = copy(pi_ref)               # anchor is an EMA of pi's weights
for step in range(T):
    x = sample_batch(D)                                 # e.g. 128 prompts
    y      = pi.generate(x)                             # 1 policy sample per prompt
    a1, a2 = anchor.generate(x), anchor.generate(x)     # 2 anchor samples per prompt
    s, s1, s2 = r(x, y), r(x, a1), r(x, a2)
    best = where(s1 >= s2, a1, a2)
    # forward KL: supervised loss on Best-of-2 of the anchor
    L_fw = -logp(pi, x, best)
    # backward KL: policy gradient with the J-BOND reward and an anchor KL in the return
    r_j = where(s < min(s1, s2), -log(16), 0.0)
    R   = r_j - (logp(pi, x, y) - logp(anchor, x, y)).detach()
    A   = R - leave_one_out_mean(R)
    L_bw = -(A * logp(pi, x, y))
    L_reg = ((logp(pi, x, y) - logp(anchor, x, y)).detach() * logp(pi, x, y))   # score-function KL estimate
    loss = mean((1 - beta) * L_fw + beta * L_bw + gamma * L_reg)
    loss.backward(); opt.step(); opt.zero_grad()
    anchor.weights = (1 - eta) * anchor.weights + eta * pi.weights
    log: reward of y, KL(pi || pi_ref) estimate, win rate of y vs anchor samples

Data needs are those of online RLHF: a prompt set (thousands to hundreds of thousands of prompts), a reward model, and no labelled responses. Every step generates three completions per prompt, one from the policy and two from the anchor, and scores them. There is no critic. The anchor is a second set of weights that must be kept in memory and used for generation and log-probabilities, which is the main structural cost compared with REINFORCE against a fixed reference.

Non-iterative BOND, as in the paper's XSum experiments, is also simple: fix $N$, draw $k$ reference samples per prompt to estimate quantiles (16 in the paper), compute $\log {\hat{p}}_{\le }(y)$ for policy samples, and combine the supervised loss on a Best-of-N sample with the policy gradient. It is useful when you have a specific $N$ you want to match, for example to replace an existing Best-of-16 serving setup.

Hyperparameters that matter#

The one decision J-BOND leaves you is when to stop. Treat it like a KL budget: log the KL to the reference and a held-out reward or win rate at regular checkpoints, and pick the checkpoint on that curve, as you would pick $N$ for Best-of-N serving.

Compute and memory cost#

Per prompt and step, J-BOND generates three sequences, scores three, and runs forward and backward on two (the policy sample and the Best-of-2 anchor sample) plus anchor forward passes for log-ratios. REINFORCE with a leave-one-out baseline and two samples generates two and backpropagates through two; PPO generates fewer but trains a critic. Generation dominates wall-clock time in online RLHF, so expect J-BOND to cost about 1.5 times a two-sample REINFORCE step. Memory holds the policy with optimiser state (about 16 bytes per parameter in mixed precision), an anchor copy in bf16 (2 bytes per parameter), the reward model, and optionally the frozen reference for monitoring. That is one more model copy than REINFORCE against a fixed reference. The amortisation argument is the point of the method: if the alternative is serving Best-of-16, every request costs 16 generations and 16 reward-model calls forever, while BOND pays once in training.

Failure modes and how to detect them#

Running past the useful KL

Because the anchor keeps moving, J-BOND keeps increasing the effective $N$. In our toy it ended on the single top response. With a learned reward model, that end point is where over-optimisation lives (Gao et al., 2022). Watch held-out reward from a different reward model or human or LLM-judge win rates as a function of KL; stop where the curve bends.

Entropy collapse

Too much backward-KL weight or a fast anchor drives per-token entropy down and makes samples repetitive. Track mean token entropy and distinct-n over a fixed prompt set. Lower $\beta$, lower $\eta$ or raise $\gamma$.

Noisy quantiles

For non-iterative BOND with large $N$, errors in ${\hat{p}}_{\le }$ are amplified by the power $N-1$. Symptoms are high-variance policy gradients and a reward curve that does not improve with more steps. Use more reference samples per prompt or switch to the iterative form with small $n$.

Reward model ties and miscalibration

Only rankings matter, so a reward model that gives many responses the same score (or saturates near a maximum) makes $p_{<}$ and $p_{\le }$ diverge and turns comparisons into coin flips. Check the fraction of tied comparisons in a batch. Rank invariance also means a reward model that is well ranked but poorly calibrated is fine for BOND and not for REINFORCE.

When to use it, and what to use instead#

Minimal implementation#

A J-BOND loss in PyTorch for a causal LM, given sequence log-probabilities computed with response masks. Generation and reward scoring are left to whatever rollout stack you use; the open-source TRL library provides the sampling and log-probability plumbing for online methods.

import math, torch

def seq_logp(model, input_ids, attn, resp_mask):
    logits = model(input_ids=input_ids, attention_mask=attn).logits[:, :-1]
    lp = torch.log_softmax(logits.float(), -1).gather(-1, input_ids[:, 1:, None]).squeeze(-1)
    return (lp * resp_mask[:, 1:]).sum(-1)

def jbond_loss(policy, anchor, batch_y, batch_best, r_y, r_a1, r_a2, beta=0.5, gamma=0.0):
    # batch_y: the policy sample; batch_best: the better of two anchor samples (tuples of ids, mask, resp_mask)
    lp_best = seq_logp(policy, *batch_best)
    lp_y = seq_logp(policy, *batch_y)
    with torch.no_grad():
        lpa_y = seq_logp(anchor, *batch_y)
        log_ratio = lp_y.detach() - lpa_y
        r_j = torch.where(r_y < torch.minimum(r_a1, r_a2), -math.log(16.0), 0.0)
        R = r_j - log_ratio
        n = R.numel()
        A = R - (R.sum() - R) / (n - 1)                      # leave-one-out baseline
    loss_fw = -lp_best.mean()
    loss_bw = -(A * lp_y).mean()
    loss_reg = (log_ratio * lp_y).mean()                     # gradient of KL(pi || anchor)
    return (1 - beta) * loss_fw + beta * loss_bw + gamma * loss_reg

@torch.no_grad()
def ema_update(anchor, policy, eta=0.02):
    for pa, pp in zip(anchor.parameters(), policy.parameters()):
        pa.mul_(1 - eta).add_(pp.detach().to(pa.dtype), alpha=eta)

Keep the anchor in bf16 and do the EMA in fp32 if you can afford it; at $\eta =0.02$ a bf16 accumulator rounds away much of each update. In POFE, J-BOND is a recipe next to RLOO and GRPO, with KL-to-reference and held-out win rate logged at every checkpoint so the stopping point is a choice you make from a curve.

Further reading#

Sessa et al. (2024) is the BOND paper; Section 3 has the derivation and Section 5 the practical algorithm. Amini et al. (2024) is the concurrent variational BoN, with the same target and backward KL only. Beirami et al. (2024) corrects the KL formula and bounds the win rate of Best-of-N by $N/(N+1)$. Gao, Schulman and Hilton (2022) is the reference for over-optimisation as a function of KL, and Ramé et al. (2024, WARP) motivates the EMA anchor. Ahmadian et al. (2024) describes the REINFORCE leave-one-out baseline that J-BOND's backward term uses and that the paper compares against.

Frequently asked questions#

References#

1. Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, et al. (2024). BOND: Aligning LLMs with Best-of-N Distillation. arxiv.org/abs/2407.14622
2. Afra Amini, Tim Vieira, Elliott Ash, Ryan Cotterell (2024). Variational Best-of-N Alignment. arxiv.org/abs/2407.06057
3. Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D'Amour, Jacob Eisenstein, Chirag Nagpal, Ananda Theertha Suresh (2024). Theoretical guarantees on the best-of-n alignment policy. arxiv.org/abs/2401.01879
4. Leo Gao, John Schulman, Jacob Hilton (2022). Scaling Laws for Reward Model Overoptimization. arxiv.org/abs/2210.10760
5. Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, Paul Christiano (2020). Learning to summarize from human feedback. arxiv.org/abs/2009.01325
6. Alexandre Ramé, Johan Ferret, Nino Vieillard, Robert Dadashi, Léonard Hussenot, et al. (2024). WARP: On the Benefits of Weight Averaged Rewarded Policies. arxiv.org/abs/2406.16768
7. Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, Sara Hooker (2024). Back to Basics: Revisiting REINFORCE Style Optimization for Learning from Human Feedback in LLMs. arxiv.org/abs/2402.14740

Related reading