PaperlayerView exact revision on arXiv

Source-linked AI summary

BoNBoN Alignment for Large Language Models and the Sweetness of Best-of-n Sampling

Lin Gui, Cristina Gârbacea, Victor Veitch

arXiv:2406.00832v3cs.CLcs.LG

TL;DR

The paper asks how best-of-n relates to training-based alignment and whether a trained LLM can mimic BoN without its costly n-sample inference. It develops BoNBoN Alignment and reports that it outperforms IPO and DPO baselines while achieving high win rates with minimal off-target drift.

  • Problem

    The paper asks how BoN relates to training-based alignment and how to avoid BoN’s substantial cost of drawing n samples for each inference.

    Source: §1 · p. 1, §2 · p. 2

  • Method

    BoNBoN trains an LLM to mimic the BoN distribution using best-of-n and worst-of-n samples and an objective derived from BoN’s analytical form.

    Source: §4 · pp. 6–7

  • Results

    BoNBoN vastly outperforms IPO and DPO baselines run on existing preference data, while achieving high win rates with minimal off-target changes.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The empirical results support using on-policy best-of-n samples over original HH or summarization completions and adding an SFT term to fix contrastive methods’ off-target drift.

    Source: §6 · p. 10

    Win rate may be inappropriate when reward magnitudes matter, though it is a natural choice when rewards come purely from binary preferences.

    Source: §6 · p. 11

Abstract

from arXiv · show

This paper concerns the problem of aligning samples from large language models to human preferences using best-of-$n$ sampling, where we draw $n$ samples, rank them, and return the best one. We consider two fundamental problems. First: what is the relationship between best-of-$n$ and approaches to alignment that train LLMs to output samples with a high expected reward (e.g., RLHF or DPO)? To answer this, we embed both the best-of-$n$ distribution and the sampling distributions learned by alignment procedures in a common class of tiltings of the base LLM distribution. We then show that, within this class, best-of-$n$ is essentially optimal in terms of the trade-off between win-rate against the base model vs KL distance from the base model. That is, best-of-$n$ is the best choice of alignment distribution if the goal is to maximize win rate. However, best-of-$n$ requires drawing $n$ samples for each inference, a substantial cost. To avoid this, the second problem we consider is how to fine-tune a LLM to mimic the best-of-$n$ sampling distribution. We derive BoNBoN Alignment to achieve this by exploiting the special structure of the best-of-$n$ distribution. Experiments show that BoNBoN alignment yields substantial improvements in producing a model that is preferred to the base policy while minimally affecting off-target aspects.

1 Introduction

The paper studies best-of-n (BoN) sampling as a way to align LLM outputs with human preferences. BoN draws n responses, ranks them on the target attribute, and returns the highest-ranked response.

1 Introduction

  • BoN aligns outputs by drawing n samples, ranking them on the target attribute, and returning the best one.The paper asks how BoN relates to training-based alignment methods.Source: §1 · p. 1
  • The paper frames the comparison in light of numerous existing approaches to LLM alignment.Examples cited include Ouy+22, Chr+17, Kau+23, Li+24, Raf+23, and Aza+24.Source: §1 · p. 1

2. How can we effectively train a LLM to mimic the BoN sampling distribution?

The paper asks whether BoN can be mimicked by a trained model, given BoN’s substantial inference cost. It introduces BoNBoN Alignment and reports high win rates with minimal off-target changes.

2. How can we effectively train a LLM to mimic the BoN sampling distribution?

  • BoNBoN models achieve high win rates while minimally changing off-target generation aspects, outperforming baselines.Figure 1 describes the comparison on the Anthropic helpful and harmless single-turn dialogue task using n = 8.Source: §2 · pp. 1–2, Figure 1 · p. 2
  • BoN requires drawing n samples per inference, a substantial computational expense that motivates training a model to mimic its distribution.The paper describes the training-versus-BoN comparison as not fully fair because their inference costs differ.Source: §2 · p. 2
  • The paper embeds BoN in a shared class of alignment distributions and shows that it is essentially Pareto-optimal for win rate versus KL distance from the base model.The result concerns the best possible trade-off within that common class.Source: §2 · p. 2
  • BoNBoN trains an LLM to mimic the BoN distribution using best-of-n and worst-of-n samples with an objective derived from BoN’s analytical form.The paper names this procedure BoNBoN Alignment.Source: §2 · p. 2

2 Preliminaries

The preliminaries define alignment in terms of preference quality and similarity to a reference model. They introduce RLHF, contrastive methods, and BoN sampling as different ways to obtain an aligned sampling distribution.

Preliminaries

  • The reference model is the LLM after supervised fine-tuning, and alignment further updates it toward responses with higher reward while retaining similarity to that reference.The paper describes reward as ordering responses according to preference.Source: §2 · p. 3
  • KL divergence formalizes how different the aligned model is from the reference model.The paper defines context-conditional and prompt-set KL divergence.Source: §2 · p. 3
  • Win rate compares an aligned-model response with a reference-model response and is used because preference-based rewards are identified only up to monotone transformation.The comparison asks which response is preferred.Source: §2 · p. 3
  • RLHF estimates a reward model from preference data, then uses KL-regularized reinforcement learning to update the LLM.The deviation from the reference model is controlled by a hyper-parameter.Source: §2 · p. 3
  • Contrastive methods use preferred and dispreferred responses directly in a fine-tuning objective, without explicitly estimating a reward function.The paper gives DPO as an example and says the objective is optimized directly.Source: §2 · p. 3
  • BoN draws n independent responses from the reference model and selects the one with the highest reward.Source: §2 · p. 4

3 Best-of-n is Win-Rate vs KL Optimal

The section compares BoN with alignment distributions using win rate and KL divergence, then explains how explicit KL regularization differs from BoN’s implicit control. It finds that BoN is essentially optimal on the win-rate–KL trade-off.

3.1 A Common Setting for Alignment Policies

  • BoN and training-based alignment distributions fit a shared reward-weighted class, with higher weights assigned to higher-reward responses.The class is defined by a non-decreasing function that may vary across prompts.Source: §3.1 · p. 4

3.2 Optimality: Win Rate versus KL divergence

  • The optimal policy maximizes win rate at a given KL level; its derivation uses a transformed reward and assumes continuous rewards.The paper notes that this continuity assumption ignores the countably infinite set of possible responses and describes the simplification as mild in practice.Source: §3.2 · pp. 4–5

3.3 The best-of-n policy is essentially optimal

  • At n = 2, BoN’s maximum win-rate difference from the optimal policy is less than 1 percentage point, and larger n values approximate it more closely.Figure 2 compares the win-rate-versus-KL curves and their difference across n.Source: Figure 2 · p. 5, §3.3 · p. 6
  • The section states that BoN is essentially optimal in terms of win rate versus KL divergence.Source: §3.3 · p. 6

3.4 Implicit vs Explicit KL regularization

  • Unlike methods that explicitly optimize against an estimated KL drift, BoN controls KL implicitly and avoids estimating that drift.The paper says measured KL can poorly correspond to salient text attributes, including response length.Source: §3.4 · p. 6

4 BoNBoN: Best-of-n fine tuning

BoNBoN trains a language model to mimic the best-of-n distribution, avoiding the inference cost of drawing n samples each time. It combines objectives that target the same policy, using best and worst samples to improve on the data efficiency of SFT-BoN and address IPO-BoN’s tendency to shift probability toward very long responses.

4 BoNBoN: Best-of-n fine tuning

  • Training data come from n reference-model responses ranked by reward, and the objectives are designed to have the best-of-n policy as a minimizer in the infinite-data limit.The method uses the best and worst responses for its pairwise data.Source: §4 · pp. 6–7
  • SFT-BoN maximizes the likelihood of best-of-n samples but is data-inefficient, yielding only marginal improvement over the reference model empirically.It uses only the winning samples.Source: §4 · p. 7
  • IPO-BoN matches the best-to-worst log-likelihood ratio to that of the best-of-n policy, with its β* analytically determined rather than tuned.The optimizer makes the best and worst samples’ log-likelihood ratio equal to the best-of-n policy’s.Source: §4 · p. 7
  • IPO-BoN is more data-efficient than SFT-BoN, but can lower the likelihood of both responses and shift probability mass toward very long outputs.This occurs when the losing response’s likelihood decreases more, increasing the ratio that IPO controls on sampled data.Source: §4 · p. 7
  • BoNBoN combines two objectives that share the best-of-n policy as a minimizer, with α balancing their contributions.The paper chooses α so the SFT and IPO terms contribute approximately equally.Source: §4 · pp. 7–8
  • BoNBoN uses worst-of-n samples to improve SFT-BoN’s data efficiency and forces best-of-n samples’ likelihood to stay high to prevent IPO-BoN’s cheating behavior.Both losses target the same policy rather than conflicting objectives.Source: §4 · pp. 7–8

5 Experiments

Across dialogue and summarization experiments, BoNBoN achieves a better win-rate/off-target trade-off than the compared alignment approaches. It also outperforms SFT and IPO trained on best-of-n data and generally avoids hyperparameter tuning through its analytic β choice.

5.1 Experimental Setup

  • Experiments use HH dialogue (106,754 training dialogues) and TL;DR summarization (92,831 training posts; 14,764 prompts), with an off-the-shelf reward model as the ground-truth ranker.Pythia-2.8b is the reference model; baselines use n=8 completions, while BoNBoN varies n from 2 to 8.Source: §5.1 · p. 8

5.2 BoNBoN achieves high win rate with little off-target deviation

  • BoNBoN achieves a better win-rate/off-target trade-off than other approaches; comparable-win-rate DPO/IPO settings can nearly double average response length.Figure 3 compares models on win-rate against the base model, estimated KL divergence, and response length for summarization and single-dialogue tasks.Source: §5.2 · p. 8, Figure 3 · p. 9

5.3 BoNBoN mimics the best-of-n policy

  • BoNBoN outperforms best-of-n-trained SFT and IPO at every β and is closer to the optimal BoN distribution; its analytic β* generally avoids hyperparameter tuning.The paper reports an excellent win-rate/off-target trade-off for the default β* choice.Source: §5.3 · pp. 8–9

6 Discussion and Related work

The discussion positions BoN and BoNBoN alongside prior alignment work, arguing that BoN is optimal for the win-rate/KL trade-off and BoNBoN can learn its distribution. It also identifies settings where win rate is not an appropriate alignment goal.

6 Discussion and Related work

  • Prior work analyzes BoN’s discrete-case probability mass function and KL estimator, and compares BoN with a KL-constrained cross-entropy optimum; other work fine-tunes on BoN data.The cited approaches include supervised fine-tuning, iterative DPO, and gradient updates using BoN samples.Source: §6 · p. 9
  • Table 1 reports roughly 85% win rates for BoNBoN, DPO, and IPO, while only BoNBoN leaves off-target attributes unmodified.The compared responses come from the same prompts and models fine-tuned on original HH data without sampling.Source: Table 1 · p. 10
  • The paper embeds BoN in a general alignment framework, establishes its win-rate/KL optimality, and derives BoNBoN to train an LLM to mimic its distribution.Unlike approaches that regularize toward a reference model, BoNBoN does not require estimating that divergence or tuning a final-solution trade-off hyperparameter.Source: §6 · p. 10
  • BoNBoN vastly outperforms IPO and DPO on existing preference data, while whether online methods have a real advantage over it remains open.The discussion notes that online alignment involves greater computational cost and complex, often unstable RL-type optimization.Source: §6 · p. 10
  • Aligning with best-of-n samples substantially outperforms alignment with original HH or summarization completions, and adding an SFT term can fix contrastive methods’ off-target drift.The discussion attributes that drift to contrastive methods pushing down the likelihood of preferred solutions.Source: §6 · p. 10
  • Win rate may be unsuitable when reward magnitudes have absolute meaning, because it treats a small win and a large win alike.The discussion presents win rate as a natural choice when rewards come purely from binary preferences.Source: §6 · pp. 10–11

A Theoretical Results

  • The theoretical appendix defines a general reward-aligned policy class that includes the optimal policy through exponential functions and BoN through power functions.For this policy class, the lemma states that win rate and KL divergence are independent of the prompt distribution.Source: A Theoretical Results · p. 14, §A.1 · p. 14

A.2 Proof of Theorem 1

These appendices derive the optimal-policy and BoN comparisons, then discuss how discrete response distributions relate to the continuous analysis. The discrete CDF approaches uniformity when response probabilities are small.

A.2 Proof of Theorem 1

  • Theorem 1 derives an exponential-form optimal policy and its context-conditional win-rate and KL expressions.The policy uses a positive constant and a normalizer.Source: §A.2 · p. 15

A.3 Proof of Theorem 2

  • The best-of-n policy has context-conditional win rate n/(n+1) against the base policy.The theorem also gives a context-conditional KL expression, but its value is not legible in the supplied passage.Source: §A.3 · p. 16

A.4 Proof of Theorem 3

  • In the discrete case, the BoN policy has a different PMF from the continuous expression, and its actual KL is smaller than the continuous-case upper bound.The passage says the actual KL is almost that bound when the corpus is large and each response has low probability.Source: §A.4 · p. 16

B Discrete versus Continuous Uniform Distribution

  • The discrete setup assumes distinct rewards for the finite response set, ordered from lowest to highest.The appendix treats the response set as finite because the vocabulary and possible word combinations are limited.Source: B Discrete versus Continuous Uniform… · p. 17

B.1 Qx normalization in discrete case

  • The discrete transformed-reward CDF is a staircase, but is nearly uniform when the response set is large and each response has low probability.Figure 4 contrasts a substantial difference for small L with a negligible difference for large L.Source: §B.1 · pp. 17–18, Figure 4 · p. 17

B.2 KL Divergence and Win Rate

The appendix bounds differences between the discrete and continuous cases using the area between their CDFs. As individual response probabilities shrink, the discrete quantities approach their continuous counterparts.

B.2 KL Divergence and Win Rate

  • The discrete-case KL divergence is upper bounded by the continuous-case KL divergence.The stated bound follows from convexity of x log(x) and Jensen’s inequality.Source: §B.2 · pp. 18–19
  • The continuous-case win rate lies between the discrete-case win rates computed with ties and without ties.The appendix also bounds the win-rate difference by a factor involving f(1), F(1)−F(0), and the CDF-area difference.Source: §B.2 · pp. 19–20
  • As the maximum response probability tends to zero, discrete and continuous KL and win-rate differences diminish, matching the practical low-probability-per-response setting.The appendix relates these differences to the area between the uniform CDF and the discrete transformed-reward CDF.Source: §B.2 · p. 20

C Experimental Details

The experimental-details appendix describes preference-label preparation for DPO and IPO, BoN hyperparameter choices, and BoNBoN training and evaluation settings.

C Experimental Details

  • DPO and IPO training uses preference labels only after checking that the reward model agrees with the chosen and rejected responses.The chosen and rejected responses are scored with the reward model before the resulting data are used for training.Source: C Experimental Details · p. 21
  • The appendix lists five β values for IPO BoN and uses α = 0.005 with β*8 = 0.0275482094 as the default BoNBoN setting.It also tests a β five times smaller or larger while holding α fixed.Source: C Experimental Details · p. 22
  • Models are trained with RMSprop at learning rate 5e−7, sampled at the 20k checkpoint, and evaluated against the SFT baseline using a reward model as judge.The appendix reports approximately three hours to reach 20k timesteps using three A100s with 140 GB of memory.Source: C Experimental Details · p. 22

D Additional results

The theoretically derived β gives consistently strong results, while other settings either reduce win rate or increase average-response-length change. Tables 2–6 provide additional model responses using best-and-worst-of-8 training samples for BoNBoN, DPO BoN, and IPO BoN.

D Additional results

  • The theoretically derived β gives consistently strong results; other settings have substantially worse win rate or substantially higher average-response-length change.Source: Figure 5 · p. 23
  • Tables 2–6 provide additional responses from different models, with BoNBoN, DPO BoN, and IPO BoN all trained on best-and-worst-of-8 samples.Source: Table 2 · p. 24, Table 3 · p. 25, Table 4 · p. 26
Loading 2406.00832v3…