PaperlayerView exact revision on arXiv

Source-linked AI summary

Variational Best-of-N Alignment

Afra Amini, Tim Vieira, Elliott Ash, Ryan Cotterell

arXiv:2407.06057v3cs.CLcs.AIcs.LG

TL;DR

Best-of-N can align language models effectively, but its N-sample inference procedure reduces throughput by a factor of N. This paper derives the BoN-induced distribution and fine-tunes models to approximate it with reverse-KL minimization; experiments report vBoN closer to BoN than other fine-tuning methods across controlled generation and summarization.

  • Problem

    Best-of-N can produce high-reward text but reduces inference throughput by a factor of N, limiting its practicality for large language models.

    Source: §1 · p. 2

  • Method

    The paper derives the distribution induced by BoN and fine-tunes a language model to minimize reverse KL to that distribution.

    Source: §1 · p. 2

  • Results

    Across controlled generation and summarization, vBoN approaches BoN’s performance and compares favorably with models fine-tuned using KL-constrained RL.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Fine-tuned vBoN models retain inference cost comparable to the reference model while approximating BoN’s performance.

    Source: §7 · p. 10

Abstract

from arXiv · show

Best-of-N (BoN) is a popular and effective algorithm for aligning language models to human preferences. The algorithm works as follows: at inference time, N samples are drawn from the language model, and the sample with the highest reward, as judged by a reward model, is returned as the output. Despite its effectiveness, BoN is computationally expensive; it reduces sampling throughput by a factor of N. To make BoN more efficient at inference time, one strategy is to fine-tune the language model to mimic what BoN does during inference. To achieve this, we derive the distribution induced by the BoN algorithm. We then propose to fine-tune the language model to minimize backward KL divergence to the BoN distribution. Our approach is analogous to mean-field variational inference and, thus, we term it variational BoN (vBoN). To the extent this fine-tuning is successful and we end up with a good approximation, we have reduced the inference cost by a factor of N. Our experiments on controlled generation and summarization tasks show that BoN is the most effective alignment method, and our variational approximation to BoN achieves the closest performance to BoN and surpasses models fine-tuned using the standard KL-constrained RL objective. In the controlled generation task, vBoN appears more frequently on the Pareto frontier of reward and KL divergence compared to other alignment methods. In the summarization task, vBoN achieves high reward values across various sampling temperatures.

1 INTRODUCTION

Language models can still produce undesirable content after supervised fine-tuning, motivating alignment to human preferences. This paper converts Best-of-N into a fine-tuning objective and reports that its variational approximation approaches BoN’s performance across controlled generation and summarization.

1 INTRODUCTION

  • BoN’s inference-time throughput cost grows by a factor of N, limiting its practicality for generating text from large language models.BoN often reaches the reward–KL frontier and can surpass RLHF fine-tuning, but requires sampling N outputs.Source: §1 · p. 2
  • vBoN derives the distribution induced by BoN and fine-tunes a language model to minimize reverse KL to that distribution.The objective is insensitive to strictly increasing transformations of reward values, and the model is fine-tuned using PPO.Source: §1 · p. 2
  • In controlled generation, vBoN fine-tuning more often reaches the reward–KL Pareto frontier than other fine-tuning alignment methods.The authors interpret this as a better balance between maximizing reward and staying close to the reference model.Source: §1 · p. 2
  • On summarization, vBoN fine-tuning yields higher average reward values and win rates than models fine-tuned with the KL-constrained RL objective.Source: §1 · p. 3

2 BACKGROUND: REINFORCEMENT LEARNING FROM HUMAN FEEDBACK

The background describes language models as parameterized distributions over strings and outlines the standard RLHF pipeline. Its KL-constrained objective balances reward against divergence from a reference model, while direct sampling from the optimal distribution is intractable.

2 BACKGROUND: REINFORCEMENT LEARNING FROM HUMAN FEEDBACK

  • A language model is a distribution over strings parameterized by θ; prompted models condition response-string distributions on prompt strings.The paper assumes these models are differentiable with respect to θ.Source: §2 · p. 3
  • RLHF first fine-tunes a reference model on task-specific data, trains a reward model, then optimizes a KL-constrained RL objective.The objective encourages high-reward outputs while penalizing excessive divergence from the reference model.Source: §2 · p. 3
  • The KL-constrained objective’s optimal policy reweights the reference model by exponentiated rewards and normalizes by a partition function, but direct sampling is intractable because that function requires an infinite sum.Sampling many reference-model strings and retaining high-reward ones motivates BoN.Source: §2 · p. 3

3 DERIVING THE BEST-OF-N OBJECTIVE

The section derives the distribution induced by Best-of-N and uses reverse-KL minimization to obtain a fine-tuning objective. It also describes the objective’s reward-ranking invariance and a Monte Carlo approach to its computationally expensive CDF term.

3 DERIVING THE BEST-OF-N OBJECTIVE

  • The probability of a string under the BoN distribution is derived under the assumption that the reward function is one-to-one.The paper notes that ties require a fixed tie-breaking order and defines F as the strict cumulative distribution function of reference-model reward values.Source: §3 · pp. 3–4
  • Minimizing reverse KL between the fine-tuned language model and the BoN distribution yields an entropy-regularized objective using the string’s BoN probability as reward.Source: §3 · p. 4
  • The vBoN objective is invariant to strictly increasing transformations of reward because it depends on reward rankings rather than reward magnitudes.The paper suggests this may make it less sensitive to reward outliers and scale.Source: §3 · p. 4
  • Computing log F directly is expensive because it sums probabilities of all strings with lower rewards, so the paper uses a Monte Carlo lower-bound estimate.The estimator uses M reference-model samples and is described as biased but consistent.Source: §3 · p. 4
  • BoN is reported as the most effective approach for positive-sentiment movie reviews, while vBoN is closest to its performance among fine-tuning methods.Source: Figure 2 · p. 5

4 COMPARING THE BON AND RL OBJECTIVES

The section compares a lower bound on the vBoN objective with the KL-constrained RL objective. It identifies similar limiting behavior while distinguishing how the two objectives use reward.

4 COMPARING THE BON AND RL OBJECTIVES

  • As N approaches 1 or infinity, the optimal distributions under the vBoN lower bound and KL-regularized RL objective coincide at the corresponding limits.At N → 1, the distribution approaches the reference model; at N → ∞, it generates only the maximum-reward string.Source: §4 · p. 5
  • N acts as a regularization parameter in vBoN, analogous to β in the KL-constrained RL objective at the stated limits.The paper also reports that models optimized on the lower bound perform competitively with those optimized on the vBoN objective.Source: §4 · p. 5
  • KL-constrained RL maximizes expected reward values, whereas vBoN maximizes the probability that aligned-model samples outrank reference-model samples.Source: §4 · p. 5

5 SENTIMENT CONTROL

The sentiment-control experiments compare fine-tuning methods with BoN on movie reviews, then assess CDF-estimation error and computational costs. vBoN comes closest to BoN among fine-tuning methods, while its fine-tuning approach reduces inference cost.

5 SENTIMENT CONTROL

  • BoN is most effective but incurs inference cost that grows with N; among fine-tuning methods, vBoN comes closest and appears more often on the Pareto front, while DPO beats PPO on win rate but trails on average reward.The comparison measures KL divergence, win rate, and average reward.Source: §5 · pp. 6–7

5.1 ERROR IN ESTIMATING log F(·)

  • With 250 samples, estimated CDFs match the 600-sample estimates for all prompts; across three example prompts, estimates hardly change after 200 samples, or after 100 for the adversarial prompt.The 5,000-prompt comparison reports very few rejections overall; the adversarial prompt has a negatively peaked reward distribution.Source: Table 1 · p. 7, §5.1 · p. 7

5.2 EFFICIENCY ANALYSIS

  • Fine-tuning methods, including vBoN, are N times more efficient at inference than BoN.Estimating log F(·) adds preprocessing; after precomputing F, the vBoN optimization loop takes the same time as PPO and evaluating F is negligible.Source: §5.2 · p. 8
  • Estimating log F(·) requires forward passes through the language model and reward model without gradients, making preprocessing highly parallelizable.The reported experiments use VLLM for efficient inference.Source: §5.2 · p. 8
  • At M = 32 samples, vBoN takes 10 minutes and remains competitive with higher M, while increasing M generally improves rewards and win rates.The analysis fixes N = 10 and averages results across three random seeds on one A100-40GB GPU.Source: §5.2 · p. 8

6 SUMMARIZATION

On summarization, BoN delivers the strongest overall performance, while vBoN comes closest at lower inference cost. Results compare alignment methods across sampling temperatures and in- and out-of-distribution Reddit posts.

6 SUMMARIZATION

  • Figure 4 reports that average reward and win rate improve as the sample size M used to approximate the vBoN loss increases.Table 2 provides example summaries sampled at temperature 0.5 with evaluator reward values.Source: Figure 4 · p. 8, Table 2 · p. 10
  • The evaluation uses a Pythia-2.8B model fine-tuned on human-written Reddit summaries, with a separate Pythia-6.9B reward model for evaluation.The Pythia-2.8B reward model is used for fine-tuning.Source: §6 · p. 8
  • Models are fine-tuned on relationship-subreddit posts and evaluated on equally sized in-distribution and out-of-distribution test groups.Solid traces denote in-distribution posts; dashed traces denote out-of-distribution posts.Source: §6 · pp. 8–9
  • BoN has the highest win rates and average rewards, while vBoN comes closest to BoN at significantly lower inference cost.The comparison spans all sampling temperatures and both in-distribution and out-of-distribution Reddit posts.Source: Figure 5 · p. 9
  • At temperatures 0.25 and 0.5, vBoN significantly outperforms PPO in win rate; at 0.75 and 1, they perform comparably.The reported win rates are measured against samples from the SFT model.Source: §6 · p. 9
  • Average reward declines as temperature rises across methods; DPO and BoNBoN approach or fall below SFT, while vBoN beats PPO at lower temperatures.The reported average-reward trends generally align with the win-rate trends.Source: §6 · p. 10

7 CONCLUSION

The paper derives a fine-tuning approximation to the BoN distribution that retains high rewards and proximity to the reference model at reference-model inference cost. Related work covers BoN’s alignment performance and cost, its applications, and concurrent fine-tuning approaches.

7 CONCLUSION

  • The variational BoN approximation attains high rewards while maintaining proximity to reference models, and its inference cost matches that of the original reference model.The paper derives the approximation by fine-tuning language models to approximate the distribution induced by BoN.Source: §7 · p. 10

A RELATED WORK

B PROOF OF PROP. 1

The proposition derives the probability of a string under the BoN distribution using events defined by reward thresholds. For non-injective rewards, the proof extends the ordering with a fixed tie-break rule.

B PROOF OF PROP. 1

  • The proof follows Casella and Berger’s stated theorem to compute πbon(y) from the two reward-threshold events.The passage introduces the events as all N samples having rewards less than or equal to r(y), and all having rewards less than r(y).Source: B PROOF OF PROP. 1 · pp. 14–15
  • Under the proposition’s one-to-one reward assumption, the proof obtains πbon(y) by subtracting the probability that all N samples have lower rewards from the probability that all have rewards no greater than r(y).The proof interprets the difference as the event that at least one sample has reward r(y) and the remaining samples have rewards no greater than r(y).Source: B PROOF OF PROP. 1 · p. 15

C STRATEGIES FOR NON-INJECTIVE REWARD FUNCTIONS

  • For non-injective rewards, the proof orders strings first by reward and breaks equal-reward ties using an arbitrary but fixed total order, such as lexicographic order.It then defines the relevant events and probabilities under that total order; the rest of the proof is unchanged.Source: C STRATEGIES FOR NON-INJECTIVE REWARD… · p. 15

D PROOF OF THM. 2

The section describes lower-bound derivations, the vBoN training procedure and experimental settings, and a comparison showing performance close to the vBoN objective.

D PROOF OF THM. 2

  • Dropping positive summation terms and retaining only the first term yields a lower bound for the objective.The section also presents a Jensen’s-inequality approach to deriving a lower bound.Source: D PROOF OF THM. 2 · p. 16
  • The Jensen-derived L1 bound is looser than L, which is established as the tighter lower bound.The proof applies Jensen’s inequality using the convexity of −log(x) for x ≥ 0.Source: D PROOF OF THM. 2 · p. 17

F EXPERIMENTAL DETAILS

  • The experimental details list PPO settings and describe varying hyperparameters to visualize expected-reward and KL-divergence trade-offs.Listed settings include 10,000 episodes, batch size 32, and sampling temperature 0.7.Source: D PROOF OF THM. 2 · p. 18, F EXPERIMENTAL DETAILS · p. 18

E VBON PSEUDOCODE

  • The vBoN algorithm initializes πθ from πref, samples responses, computes their rewards, and optimizes with Eq. (5) or Eq. (8) using PPO.The pseudocode specifies one response per prompt in each batch.Source: E VBON PSEUDOCODE · p. 18

G COMPARING THE VBON OBJECTIVE AND L LOWER BOUND

H ADDITIONAL EXPERIMENTS WITH BON-SFT

The section reports a BoN-SFT experiment using maximum-likelihood training on BoN generations while varying N, and finds that BoN diverges more from the reference model than other fine-tuning methods.

H ADDITIONAL EXPERIMENTS WITH BON-SFT

  • BoN diverges too much from the reference model compared with other fine-tuning methods for alignment.The experiment trains with maximum likelihood on BoN generations while varying N; results are shown in Figure 7.Source: H ADDITIONAL EXPERIMENTS WITH BON-SFT · p. 19

I QUALITATIVE RESULTS

Tables 4–7 provide example summaries and evaluator-reward values at sampling temperatures 0.25, 0.5, 0.75, and 1.

I QUALITATIVE RESULTS

  • Table 4 gives example summaries sampled at temperature 0.25 and their evaluator rewards.Source: Table 4 · p. 20
  • Tables 5 and 6 give example summaries and evaluator rewards at temperatures 0.5 and 0.75.Source: Table 5 · p. 20, Table 6 · p. 21
  • Table 7 gives example summaries sampled at temperature 1 and their evaluator rewards.Source: Table 7 · p. 21

Related papers

Loading 2407.06057v3…