Source-linked AI summary
BOND: Aligning LLMs with Best-of-N Distillation
Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amélie Héliou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev, Olivier Bachem
TL;DR
Best-of-N can improve generations but costs N inference-time samples, while RLHF has challenges including forgetting and reward hacking. BOND distills the Best-of-N distribution into a policy using divergence-based matching and a moving anchor. Experiments report improved KL-reward trade-offs and benchmark performance for J-BOND.
Problem
Method
BOND fine-tunes a policy to match the Best-of-N distribution, using a Jeffreys-divergence combination of forward and backward KL and an iterative moving anchor.
Source: §1 · p. 2, §4.2 · p. 6, §5 · p. 10Results
J-BOND improves the KL-reward Pareto front and enhances academic-benchmark and side-by-side performance against open-source variants such as Mixtral.
Source: §1 · p. 2, §8 · p. 13Takeaways & Limitations
BOND aims to retain Best-of-N's strong performance while letting the fine-tuned policy generate with a single inference-time sample.
Source: §1 · p. 2, §2 · p. 3Autoregressive sampling is the main online-RLHF bottleneck, so J-BOND uses a practical approach with a small number of anchor samples.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Reinforcement learning from human feedback (RLHF) is a key driver of quality and safety in state-of-the-art large language models. Yet, a surprisingly simple and strong inference-time strategy is Best-of-N sampling that selects the best generation among N candidates. In this paper, we propose Best-of-N Distillation (BOND), a novel RLHF algorithm that seeks to emulate Best-of-N but without its significant computational overhead at inference time. Specifically, BOND is a distribution matching algorithm that forces the distribution of generations from the policy to get closer to the Best-of-N distribution. We use the Jeffreys divergence (a linear combination of forward and backward 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.
1. Introduction
Best-of-N sampling can improve generation quality but requires many inference-time samples, while RLHF also faces forgetting and reward-hacking challenges. BOND distills the Best-of-N distribution into a policy, and the reported experiments show improvements across reward-quality trade-offs and benchmarks.
RLHF algorithms and their challenges
- RL fine-tuning can cause forgetting and reward hacking, motivating KL-regularized policy-gradient methods that seek high reward while remaining close to the SFT policy.The cited risks include forgetting pre-trained knowledge and exploiting reward-model loopholes.Source: §1 · p. 1
Best-of-N sampling
- Best-of-N selects the highest-reward generation from N candidates, improving reward-KL trade-offs but increasing computational cost by a factor of N.It samples and scores N generations from the reference model.Source: §1 · p. 1
BOND
- BOND matches the policy distribution to the Best-of-N distribution, combining forward and backward KL objectives and iteratively distilling from a moving anchor.Forward KL encourages mode coverage, while backward KL encourages mode seeking; the paper proposes their Jeffreys-divergence combination.Source: §1 · p. 2
Experiments
- J-BOND improves the KL-reward Pareto front and enhances academic-benchmark and side-by-side performance against open-source variants such as Mixtral.The experiments apply J-BOND to Gemma policies.Source: §1 · p. 2
2. Problem Setup
The setup assumes a reference policy and a reward model reflecting human preferences. It contrasts inference-time Best-of-N selection with the proposed goal of distilling that strategy into a policy.
Problem setup
- The setup assumes access to a reward model trained to reflect human preferences and a reference policy, typically supervised fine-tuned.The policy generates token sequences autoregressively from prompts.Source: §2 · p. 2
Best-of-N
- Best-of-N samples repeatedly from the reference policy and selects the highest-reward generation, but its inference cost grows linearly with N.Unlike RLHF, it changes inference rather than fine-tuning model weights.Source: §2 · p. 3
BOND
- BOND aims to distill Best-of-N into the policy so it can attain its strong performance with a single inference-time sample.The method is introduced as an alignment approach motivated by Best-of-N's inference cost.Source: §2 · p. 3
3. The BOND Approach
BOND characterizes the Best-of-N distribution and formulates alignment as matching the policy to that distribution. The paper also relates this formulation to standard KL-regularized RLHF.
3. The BOND Approach
- BOND proceeds by deriving the Best-of-N distribution, matching the policy to it, and connecting the resulting objective to standard RLHF.The approach is organized around these three steps.Source: §3 · p. 3
3.1. The Best-of-N distribution
- Best-of-N reweights the reference distribution: one factor increasingly penalizes generations with many worse-or-equal alternatives, while a collision correction stays between 1 and N.The penalty suppresses bad generations more as N increases; the correction accounts for repeated generations.Source: §3.1 · p. 4
- For the worst generation, the collision correction is 1; for low-likelihood good generations it approaches N, and for continuous reference distributions it equals N.These cases reflect whether repeated sampling of a generation is likely.Source: §3.1 · p. 4
3.2. The BOND objective
- The BOND objective steers the training policy toward the Best-of-N distribution by minimizing a chosen divergence, including forward or backward KL.The paper notes that existing distribution-matching techniques can estimate divergence from online and offline samples.Source: §3.2 · p. 4
3.3. Connection with standard RLHF
- Best-of-N coincides with a KL-regularized RLHF optimum for a specific BOND reward and regularization strength determined by N.The paper further describes the BOND reward as a log-reward quantile, invariant to monotone reward transformations, and conjectures it may be more robust to reward hacking.Source: §3.3 · pp. 4–5
- The connection suggests steering the policy toward Best-of-N through distribution matching when its reward or distribution can be computed.The paper presents this as motivation for the algorithms developed next.Source: §3.3 · p. 5
4. BOND Challenges and Algorithms
This section addresses BOND’s quantile estimation, divergence choice, and selection of N, using Monte Carlo quantiles, a Jeffreys objective, and iterative distillation with a moving anchor.
4.1. Monte-Carlo quantile estimation
- Monte Carlo estimates each generation’s same-prompt reward quantile by sampling k generations from πref; it worked well with limited samples, and learned quantiles are an alternative.The quantile measures a generation’s quality relative to generations from the reference policy conditioned on the same prompt.Source: §4.1 · p. 5
4.2. Jeffreys divergence as a robust objective
- Jeffreys divergence weights forward and backward KL, balancing mode-covering and mode-seeking behavior; the authors report that it avoids forward-KL over-spreading and backward-KL collapse risks.The paper motivates the combined objective as a robust distribution-matching choice for BOND.Source: §4.2 · pp. 5–6
- With N=8, β=0.5 minimizes both KL divergences from πBoN and maximizes reward quantiles similarly to β=1, while β=0 lags.The figure reports backward KL, forward KL, and reward quantiles; β=0.5 combines the two KL directions.Source: §4.2 · p. 6, Figure 2 · p. 7
4.3. Iterative BOND
- Recursive distillation turns repeated Best-of-n improvement into an effective Best-of-n^M target, updating the moving anchor to the current policy after distillation steps.The approach allows a small n without setting a large N upfront, with the authors describing improved stability and reduced sample complexity.Source: §4.3 · pp. 7–8
- Iterative BOND continuously improves rewards and log quantiles as non-iterative BOND saturates, while retaining the same reward/KL trade-off with small n.The comparison uses iterative n=2 and n=4 against non-iterative N=4, 8, and 16.Source: §4.3 · p. 8, Figure 4 · p. 8
5. The J-BOND Algorithm
J-BOND iteratively distills the Best-of-2 distribution of a moving anchor, combining forward and backward KL objectives with practical sampling and update choices. Its EMA anchor and added KL regularization stabilize training and can improve the reward/KL trade-off, while the algorithm uses few samples per prompt.
Iterative distillation and moving anchor
- J-BOND iteratively trains policy π_t to distill the Best-of-2 distribution of a moving anchor, initialized as the reference policy.After a set of distillation steps, the anchor is updated to the current policy; J-BOND uses an EMA update at every fine-tuning step.Source: §5 · pp. 8, 10
Sample and divergence estimation
- Each prompt uses one policy sample and two anchor samples to estimate forward and backward KL components for the Jeffreys-divergence objective.With only two anchor samples, the backward-KL reward is −log(16) only when the policy sample scores below both anchor samples; otherwise it is zero.Source: §5 · pp. 8–9
Sample complexity
- J-BOND uses a small number of samples for practical online RLHF, although more anchor samples can generally improve divergence estimation.The authors choose this trade-off because autoregressive sampling is the main bottleneck of online RLHF.Source: §5 · pp. 8–9
EMA anchor and regularization
- The EMA anchor reduces training variance and can improve the reward/KL trade-off, while additional KL regularization keeps policy updates closer to the moving anchor.The added regularization is controlled by γ and is intended to stabilize updates.Source: §5 · p. 10
6. Experiments
Experiments on Gemma models compare J-BOND with REINFORCE and examine EMA anchor updates, anchor speed, and additional KL regularization. J-BOND achieves a better reward/KL trade-off than the tested REINFORCE baselines, and the reported Gemma fine-tuning results are competitive against open models.
EMA vs. hard anchor updates
- On Gemma 7B, the EMA anchor yields the same reward increase profile as periodic 50-step updates but a lower KL increase and better reward/KL trade-off.The comparison uses EMA coefficient η = 0.02.Source: §6 · p. 11
Anchor speed and KL regularization
- On Gemma 2B, larger EMA mixing η speeds reward increases, while larger γ slows movement away from the reference policy and can improve reward/KL trade-off.The study tests η ∈ {0.01, 0.05, 0.1} with γ = 0, and γ ∈ {0, 0.5, 1, 2} with η = 0.05.Source: §6 · p. 11, Figure 6 · p. 11
Comparison with standard RLHF
- J-BOND produces a better reward/KL trade-off than all tested REINFORCE baselines while continuously improving reward with a stable, linear KL increase.The REINFORCE comparison tests regularization strengths β_RL ∈ {0.001, 0.01, 0.1, 1}; its reward is sensitive to that coefficient.Source: Figure 7 · p. 11, §6 · p. 12
J-BOND for Gemma open models
- Gemma 1.1 IT 7B fine-tuned with J-BOND outperforms Mistral 7B v0.2 Instruct in both safety and instruction following.The paper reports competitive performance for Gemma and other open-weight models fine-tuned with J-BOND.Source: §6 · p. 12
7. Related Work
Related work covers Best-of-N sampling, methods that distill its behavior, and dynamic-anchor approaches. BOND is positioned among these as a distribution-matching method using Jeffreys divergence and an iterative dynamic anchor.
Best-of-N sampling
- Prior Best-of-N work studies its analytical KL and reward properties, theoretical links to KL-constrained RL, and scaling laws.The cited work also develops Best-of-N as a reward-based inference-time method.Source: §7 · p. 12
Matching Best-of-N
- Several methods match Best-of-N through supervised fine-tuning, preference optimization, or distribution matching, while BOND uses Jeffreys divergence and an iterative dynamic anchor.The closest concurrent distribution-matching work uses backward KL alone, whereas BOND uses a different divergence and proposes an iterative anchor procedure.Source: §7 · p. 12
Dynamic anchors
- Dynamic anchors can expand the policy space explored, but updates that are too slow may limit optimization and updates that are too fast may hinder stability.EMA policy references are increasingly used in alignment approaches.Source: §7 · p. 12
Scaling post-training
- BOND frames additional training-time computation as complementary to approaches that scale post-training with multiple reward models or policies.The paper connects its iterative generation-and-reward-comparison process to iterated amplification.Source: §7 · p. 13
8. Conclusion
The paper introduces BOND as online distillation of the Best-of-N distribution and presents J-BOND as a practical implementation. J-BOND combines Monte Carlo quantile estimation, forward and backward KL objectives, and an iterative EMA anchor, with reported improvements to the KL-reward Pareto front.
BOND
- BOND fine-tunes a policy through online distillation of the Best-of-N sampling distribution.The stated goal is to capture Best-of-N performance while requiring only one inference-time sample.Source: §8 · p. 13
J-BOND
- J-BOND combines Monte Carlo quantile estimation, forward and backward KL objectives, and iterative distillation with an EMA anchor.Source: §8 · p. 13
Reported outcome
- J-BOND improves the KL-reward Pareto front and compares favorably with state-of-the-art baselines.Source: §8 · p. 13
A. Supporting results and derivations
The appendix derives the Best-of-N selection probability and connects its discrete formulation to the continuous maximum. It also shows that minimizing backward KL to the Best-of-N distribution is equivalent, up to scaling, to a KL-regularized REINFORCE update.
A.1. Proof of Theorem 1
- The events that a candidate is selected at each sample index are disjoint and together cover selection of that generation by Best-of-N.Summing their likelihoods gives the probability that Best-of-N selects the generation.Source: §A.1 · p. 18
A.2. Link to the continuous case
- For continuous rewards, the Best-of-N correction factor is N, matching the maximum-of-N formula where the CDF corresponds to the discrete cumulative probability.The passage identifies the continuous density as analogous to the reference policy probability in the discrete case.Source: §A.2 · pp. 18–19
A.3. Backward KL and policy gradient equivalence
- Descending the backward-KL gradient is equivalent, up to constant scaling, to REINFORCE on a KL-regularized RLHF objective using rBOND and βBOND.The appendix establishes this connection between the backward-KL gradient and the policy-gradient update.Source: §A.3 · p. 19
A.4. Derivation of J-BOND reward
J-BOND approximates the unknown Best-of-N reward, log p≤(y), using a reward designed around anchor samples. Its expected reward matches the true reward at the median anchor quantile.
A.4. Derivation of J-BOND reward
- J-BOND reward approximates log p≤(y), which is unknown because calculating the quantile generally requires the anchor policy’s reward distribution.The derivation presents the J-BOND reward as an approximation to the true BOND reward.Source: §A.4 · pp. 19–20
- J-BOND assigns a negative reward only when a sample is worse than both anchor samples, reflecting the concavity of the log quantile.The authors report no gains from also rewarding the intermediate case.Source: §A.4 · p. 20
- At the median anchor quantile, p≤(y) = 0.5, the expected J-BOND reward is designed to match the true reward log(0.5).Figure 8 plots the expected J-BOND reward against the true reward.Source: §A.4 · p. 20, Figure 8 · p. 20
B.1. Learned quantile models
Monte-Carlo quantile estimates can be sample-intensive and do not use information across contexts. The appendix explores a learned context-dependent estimator and reports results comparable to Monte-Carlo estimates on the tested task.
B.1. Learned quantile models
- Monte-Carlo quantile estimation may require many samples for accuracy and does not exploit context information, treating new contexts independently.This limits its ability to use estimates from a context with similar reward quantiles.Source: §B.1 · pp. 20–21
- The learned alternative treats a reward quantile as the probability of a binary outcome and trains a context-dependent estimator with binary cross-entropy.The outcome indicates whether a reference-policy sample’s reward is no greater than the candidate’s reward.Source: §B.1 · p. 21
- On abstractive summarization, learned quantiles yield comparable backward and forward KL divergences and average log quantiles to Monte-Carlo estimates.During evaluation, the Best-of-N distribution is approximated with 32 Monte-Carlo samples; Figure 9 compares the learned and Monte-Carlo estimates.Source: §B.1 · p. 21, Figure 9 · p. 21
- The explored learned-quantile approach is described as naive, with alternatives including enforced ordering, quantile regression, and specified reward distributions.The text gives Gaussian reward distributions as one possible specification.Source: §B.1 · p. 21
B.2. Additional plots for BOND with Jeffreys divergence objective
Across the additional settings shown, β = 0.5 minimizes both KL divergences from the Best-of-N distribution. Its log-quantile performance is similar to β = 1.
B.2. Additional plots for BOND with Jeffreys divergence objective
- β = 0.5 minimizes both forward and backward KL divergences from πBoN, while its log quantiles are similar to β = 1.The figure shows results for N = 4 and N = 16; reward quantiles improve significantly over β = 0.Source: §B.2 · p. 21, Figure 10 · p. 22