PaperlayerView exact revision on arXiv

Source-linked AI summary

Theoretical guarantees on the best-of-n alignment policy

Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D'Amour, Jacob Eisenstein, Chirag Nagpal, Ananda Theertha Suresh

arXiv:2401.01879v3cs.LGcs.CLcs.ITmath.IT

TL;DR

Keeping the aligned policy close to its reference matters, yet a commonly used equality for best-of-n KL divergence is false. The paper derives bounds and proposes an estimator, showing that in its example true KL stays at or below log(2) for all n while the formula grows without bound.

  • Problem

    A commonly used formula treats best-of-n KL divergence as an equality, although small KL is desired because it implies that reference-policy capabilities are largely preserved.

    Source: §1 · pp. 1–2

  • Method

    The paper derives KL bounds, proposes an estimator based on observable εn, and analyzes win-rate versus KL tradeoffs for best-of-n.

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

  • Results

    log(2) upper-bounds true KL for all n in the example, while the literature formula grows unbounded as n increases; the proposed estimator closely mirrors true KL.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The analysis of win-rate versus KL tradeoffs supports the paper's justification for best-of-n's widespread use in language-model alignment.

    Source: §1 · p. 2

    The estimator depends on εn, and its expected upper-bound property is conjectured rather than established.

    Source: §4 · p. 4

Abstract

from arXiv · show

A simple and effective method for the inference-time alignment and scaling test-time compute of generative models is best-of-$n$ sampling, where $n$ samples are drawn from a reference policy, ranked based on a reward function, and the highest ranking one is selected. A commonly used analytical expression in the literature claims that the KL divergence between the best-of-$n$ policy and the reference policy is equal to $\log (n) - (n-1)/n.$ We disprove the validity of this claim, and show that it is an upper bound on the actual KL divergence. We also explore the tightness of this upper bound in different regimes, and propose a new estimator for the KL divergence and empirically show that it provides a tight approximation. We also show that the win rate of the best-of-$n$ policy against the reference policy is upper bounded by $n/(n+1)$ and derive bounds on the tightness of this characterization. We conclude with analyzing the tradeoffs between win rate and KL divergence of the best-of-$n$ alignment policy, which demonstrate that very good tradeoffs are achievable with $n < 1000$.

1. Introduction

The introduction motivates a theoretical study of best-of-n alignment and whether commonly used KL characterizations reflect its actual behavior. It previews KL bounds, a new estimator, and analysis of reward/win-rate–KL tradeoffs.

1. Introduction

  • Small KL divergence is desirable because it indicates that the aligned policy preserves reference-model capabilities, motivating reward improvement with minimal KL drift.The introduction frames reward or win rate against KL divergence as a tradeoff curve for comparing alignment techniques.Source: §1 · pp. 1–2
  • Prior studies report compelling win-rate–KL tradeoffs for best-of-n, motivating its theoretical investigation (Gao et al., 2023; Mudgal et al., 2024; Eisenstein et al., 2024).The introduction also describes best-of-n as a widely used alignment method and a strong baseline for inference-time compute scaling.Source: §1 · p. 2
  • Best-of-n selects the highest-reward response among n reference-policy samples, and increasing n raises reward at the cost of greater drift from the reference model.The strategy samples n independent responses for a prompt before selecting one.Source: §1 · p. 2
  • log(2) upper-bounds the exact KL for all n in the binary example, while the literature formula grows unbounded; the proposed estimator closely tracks the exact KL.Figure 1 compares the formula, exact KL, and estimator for this example.Source: Figure 1 · p. 2, §1 · p. 2
  • The paper investigates KL guarantees for best-of-n across n, derives KL bounds, proposes a new estimator, and analyzes win-rate–KL tradeoffs.The stated goal is to better understand the best-of-n alignment policy and the divergence between it and the reference policy.Source: §1 · p. 2

2. Derivation of the Best-of-n Policy

Under assumptions of unique rewards and finite outcome support, the section derives the best-of-n policy’s PMF by ranking outcomes and connecting best-of-n sampling to the maximum of uniform draws.

2. Derivation of the Best-of-n Policy

  • Unique rewards and finite outcome support are the two assumptions used to derive the best-of-n policy; the authors describe them as fairly non-restrictive and helpful for clearer presentation.Source: §2 · p. 3
  • Lemma 2.3 gives a closed-form PMF under the assumptions, and at n = 1 the best-of-n policy equals the reference policy.The PMF is used to derive subsequent guarantees on KL divergence and win rate.Source: §2 · p. 3
  • The derivation orders possible outcomes by reward and represents reference-policy sampling with a uniform draw, then uses the maximum of n such draws to obtain the best-of-n PMF.Source: §2 · p. 3
  • The PMF extends to real τ ≥ 1, an extension used in best-of-n distillation and to provide bounds in Section 7.The paper notes that sampling from this extension may not be immediately clear.Source: §2 · p. 3

3. Relations Between the KL Divergence and the Analytical Formula

The analytical formula is an upper bound on the context-dependent KL divergence, but the gap can vary substantially: it is bounded under low-probability outcome conditions and can grow unbounded as n increases.

3. Relations Between the KL Divergence and the Analytical Formula

  • The analytical formula upper-bounds the context-dependent KL divergence, and the result also holds for any prompt distribution.The proofs are relegated to Appendix A.Source: §3 · p. 3
  • Mroueh (2024) extended the result to broader stochastic processes, while the appendix derives blockwise best-of-n KL bounds, with similar reward–KL tradeoffs generally possible using 10x smaller n.The cited blockwise result is attributed to Mudgal et al. (2024).Source: §3 · p. 3

3.1. Upper Bounds on the Gap

  • For δ-bound reference models, the gap has an upper bound involving conditional Rényi entropy of order 2; δ-bound means every outcome has probability at most δ.The analysis focuses on small δ, where all outcomes have low likelihood.Source: §3.1 · pp. 3–4
  • When n^2δ ≪ 1, repeated samples are unlikely and the analytical formula can be relatively accurate; this δ-bound condition was implicit in Hilton & Gao (2022).The paper also gives a gap upper bound specifically for δ-bound models.Source: §3.1 · p. 4

3.2. Lower Bounds on the Gap

  • The lower-bound analysis defines a random ε_n(x) and its deterministic limit ε∞, using the reward-maximizing outcome y_max(x).The passage defines y_max(x) as an outcome maximizing the reward for context x.Source: §3.2 · p. 4
  • When nε∞ ≫ 1, the gap between the analytical formula and KL divergence grows unbounded.The theorem provides a lower bound on the gap in terms of ε∞.Source: §3.2 · p. 4

4. Proposed Estimator for KL Divergence

The section proposes a practical KL estimator based on observable εn and examines its bounds, variability, and performance. Across the reported examples, it closely tracks the true KL divergence, while its expected upper-bound property remains a conjecture.

4. Proposed Estimator for KL Divergence

  • The alternate estimator’s expectation upper-bounds KL but is additively loose by (n −1)/n, especially when nεn ≪1.The proposed estimator is introduced to close this gap.Source: §4 · p. 4
  • The practical estimator uses observable εn rather than unobservable ε∞, and the authors conjecture that its expectation upper-bounds the true KL divergence.Tens of thousands of randomly generated numerical experiments suggest the conjectured bound holds, but the authors do not provide a mathematical proof.Source: §4 · pp. 4–5
  • Estimator variance can be high when εn depends on whether an outcome appears; its standard deviation is at most log n, so averaging M = O(log n log 1/δ) draws guarantees it is below δ.The authors note that bootstrapping across the M batches may yield better variance guarantees.Source: §4 · pp. 5–6
  • Across uniform and cherry-picked distributions, the proposed estimator closely follows true KL, while the analytical formula can miss its behavior substantially.For uniform distributions, the analytical formula’s gap is small when n/L ≪ 1 but grows large and unbounded when n/L ≫ 1; the alternate bound captures the latter regime but retains an additive gap of (n −1)/n in the former.Source: §4 · p. 5, Figure 2 · p. 5
  • On four Alpaca/Gemma examples, the proposed estimator substantially outperforms the analytical formula for lower-entropy responses, approaching the better upper bound; negative-length rewards show similar trends.For two examples with large ε∞, including “What is the capital of France?”, the analytical formula is loose even for n ≈ 20.Source: §4 · p. 6, Figure 4 · p. 6

5. Win Rate of the Best-of-n Policy

The section characterizes best-of-n win rates and bounds their gap from n/(n+1). It also shows that win rate depends on reward and the reference policy through a calibrated reward function.

5. Win Rate of the Best-of-n Policy

  • Any policy’s win rate depends on reward and the reference policy through the calibrated reward function.The authors note that this transformation is proposed as the objective in IPO and is used in best-of-n distillation and inference-aware alignment.Source: §5 · pp. 6–7
  • The best-of-n win rate is at most n/(n+1), motivated by its n-to-1 chances against a reference sample.The argument allows for ties when samples have the same reward value.Source: §5 · p. 7

5.1. Upper Bounds on the Win Rate Gap

  • For a δ-bound reference policy, the win rate is fairly close to n/(n+1) when nδ ≪ 1.The section derives an upper bound on the win-rate gap and applies it under this condition.Source: §5.1 · p. 7

5.2. Lower Bounds on the Win Rate Gap

  • The section also gives lower bounds identifying cases where the win-rate gap is bounded away from zero.Theorem 5.6 states its result in terms of ε∞, defined earlier in the paper.Source: §5.2 · p. 7

6. Rewind-and-Repeat: Rejection Sampling Beyond Best-of-n

Rewind-and-repeat accepts the first reference-policy sample whose reward reaches a threshold, with the number of trials determining inference cost. The section’s main contribution is an unbiased harmonic-number estimator for its KL divergence.

6. Rewind-and-Repeat: Rejection Sampling Beyond Best-of-n

  • Rewind-and-repeat accepts y_M when its reward reaches threshold Φ and every earlier draw has reward below Φ.M is the random number of trials until the threshold is met, and determines inference cost.Source: §6 · p. 7
  • The expected number of trials is E[M] = 1/w_Φ(x), where w_Φ(x) is the reference-policy probability of meeting the threshold.Source: §6 · p. 7
  • Sweeping Φ sweeps w_Φ over [0, 1], yielding a win-rate-versus-KL tradeoff for rewind-and-repeat.The section derives the policy’s PMF and uses it to characterize KL divergence and win rate against the reference policy.Source: §6 · p. 8
  • H_{M-1} is an unbiased estimator of the rewind-and-repeat policy’s KL divergence from the reference policy.Here M is the number of trials until an outcome is accepted, and H_n is the n-th harmonic number.Source: §6 · p. 8

7. Win rate vs KL Divergence Tradeoffs

The section shows that best-of-n’s actual win-rate–KL tradeoff can be more favorable than the analytical KL upper bound suggests. It also derives tradeoff bounds, conjectures a tighter upper bound, and compares best-of-n with win-rate-optimal and rewind-and-repeat policies in examples.

7. Win rate vs KL Divergence Tradeoffs

  • At large n in Example 1, best-of-n reaches KL divergence log(2) and win rate 0.75, while its actual tradeoff curve is more favorable than the analytical upper-bound formula predicts.The section also states that tradeoffs reported using that formula are conservative, so the actual curve is guaranteed to be no worse.Source: Figure 6 · p. 8, §7 · p. 8
  • The paper turns existing bounds into lower and upper bounds on best-of-n’s tradeoff curve and conjectures that W0(D) = ℓ^-1(D) is an upper bound.The conjecture defines ℓ for τ in [0.5, 1).Source: §7 · p. 8
  • In Example 2, best-of-n’s win-rate–KL tradeoff is close to win-rate optimal and both outperform rewind-and-repeat, although their solutions can differ on the probability simplex.The close tradeoff finding matches recent findings of Yang et al. (2024a) and Gui et al. (2024).Source: §7 · pp. 8–9, Figure 8 · p. 9, Figure 7 · p. 9
  • In Example 3’s low-probability-outcome asymptotic regime, best-of-n is almost optimal and offers better tradeoffs than rewind-and-repeat.Source: §7 · p. 9

8. Conclusion

The paper shows that the commonly used best-of-n KL formula is an upper bound rather than an exact value, and derives bounds for KL divergence and win rate. It also evaluates a KL estimator and compares best-of-n with rewind-and-repeat, identifying unresolved tradeoffs among compute, reward, and KL divergence.

8. Conclusion

  • The literature’s log(n) - (n-1)/n formula is an upper bound, not the exact KL; its gap is small when nε_n ≪ 1 but may be unbounded when nε_n ≫ 1.The gap conditions refer to ε_n, the probability mass under the base model of a draw from the best-of-n policy.Source: §8 · p. 9
  • Best-of-n’s win rate against the reference is at most n/(n + 1), and the paper derives upper and lower bounds on the gap between the actual win rate and this bound.Source: §8 · p. 9
  • The proposed KL estimator captures the behavior of the KL divergence in several numerical experiments.Source: §8 · p. 9
  • The bounds are also extended to blockwise best-of-n (Mudgal et al., 2024).Source: §8 · p. 9
  • Best-of-n has better reward–KL tradeoffs than rewind-and-repeat, but rewind-and-repeat drives reward higher per compute; future work may seek Pareto-optimal compute, reward, and KL methods.The paper suggests combining rewind-and-repeat with best-of-n as one possible direction.Source: §8 · p. 9

Impact Statement … A.4. Proofs of Section 6

The paper argues that best-of-n alignment has theoretical guarantees for preserving closeness to a reference policy and for win-rate/KL tradeoffs, while the appendix supplies proofs and an unbiased estimator used in Section 6.

Impact Statement

  • The widely used best-of-n KL formula is an upper bound on the actual divergence, which may help keep an aligned policy close to a safe reference policy.The impact statement notes this may help compliance or risk-management teams preserve safety.Source: Impact Statement · p. 9
  • Best-of-n is described as an effective, almost optimal alignment method with win-rate/KL guarantees, but adversarial rewards can also enable jailbreaks.The paper calls for future work both to improve best-of-n and to devise safeguards against such jailbreaks.Source: Impact Statement · p. 9

A.1. Proofs of Section 3

  • Appendix A.1 proves Section 3 results using auxiliary lemmas, algebraic manipulation, and a case split on a and b.The proof of Theorem 3.6 also lower-bounds the gap using the gap for the highest-reward outcome.Source: §A.1 · pp. 12–14

A.2. Proofs of Section 4

  • For Section 4, the proof singles out the highest-reward outcome, bounds the remaining terms, and uses a lemma to express an integral in closed form.This approach is stated for the proof of Corollary 4.2.Source: §A.2 · p. 15

A.3. Proofs of Section 5

  • In the Section 5 proof appendix, a displayed inequality includes the factor n/(n+1).The passage gives the inequality as 1/2(b^n-a^n)(b+a) ≤ n/(n+1)(b^(n+1)-a^(n+1)).Source: §A.3 · p. 16
  • Appendix A.3 provides an auxiliary upper bound on h_n(a, b) and inequalities used in the proofs of Section 5 results.The proofs of Theorems 5.3 and 5.4 invoke Lemmas A.4 and A.5, respectively.Source: §A.3 · pp. 16–17

A.4. Proofs of Section 6

  • For Section 6, the appendix defines M as the number of trials until the first success and identifies an unbiased estimator of -log(w_Φ(x)).The proof also describes repeated trials when no good outcome is obtained.Source: §A.4 · p. 18

A.5. Proofs of Section 7

The proofs establish lower and upper bounds on the best-of-n tradeoff curve, deriving them from combinations of earlier theorems.

A.5. Proofs of Section 7

  • Theorem A.6 defines WL(DL) for τ ≥ 1 and establishes it as a lower bound on the best-of-n tradeoff curve.Its proof combines Theorems 3.1 and 5.4.Source: §A.5 · p. 19
  • Theorem A.7 defines WU(DU) for τ ≥ 1 and establishes it as an upper bound on the best-of-n tradeoff curve.Its proof combines Theorems 5.3 and 3.4.Source: §A.5 · p. 19

B. KL divergence of blockwise best-of-n · C. Experiments

The section characterizes blockwise best-of-n decoding and its KL divergence relative to the reference policy. It proposes an expected upper-bound estimator based on the number of decoded blocks.

B. KL divergence of blockwise best-of-n

C.1. Details of experiments on Alpaca dataset · C.2. Computation of tilted min/max · C.3. Experiments with machine translation prompts

The appendix specifies the Alpaca experimental setup, defines tilted min/max as summaries of estimator quantiles, and presents machine-translation experiments comparing KL estimates under two reward choices.

C.1. Details of experiments on Alpaca dataset

  • The experiments select four prompts from the Alpaca dataset, including questions about France’s capital and staying healthy.Source: §C.1 · p. 22
  • The Alpaca experiments use Gemma 2 9B instruction-tuned at temperature 1, with reference-model log-likelihood reward and Corollary 4.2 using 100 samples to bound ε∞.The prompts are formatted according to Gemma instructions.Source: §C.1 · p. 22

C.2. Computation of tilted min/max

  • The appendix computes a tilted average of the estimator defined by Li et al. (2023, Equation (2)).Source: §C.2 · p. 22
  • For t = −1 and t = 1, the estimator’s tilted values are called tilted min and tilted max, capturing low and high quantiles whose difference portrays deviation from the mean.As t approaches −∞ or +∞, the tilted average recovers the minimum or maximum, respectively.Source: §C.2 · p. 22

C.3. Experiments with machine translation prompts

  • Figures 10 and 11 compare the analytical formula, Corollary 4.2 upper bound, proposed estimator, and exact KL across four machine-translation examples.Figure 10 uses reference-model log-likelihood as reward; Figure 11 uses negative response length.Source: Figure 10 · p. 23, Figure 11 · p. 23
  • Figure 10 evaluates the four machine-translation examples using Gemma 9B IT, with reward defined as response log-likelihood under the reference model.It plots the analytical formula, Corollary 4.2 upper bound, proposed estimator, and exact KL.Source: Figure 10 · p. 23
  • Figure 11 evaluates the four machine-translation examples using Gemma 9B IT, with reward defined as negative response length under the reference model.It plots the analytical formula, Corollary 4.2 upper bound, proposed estimator, and exact KL.Source: Figure 11 · p. 23

Related papers

Loading 2401.01879v3…