PaperlayerView exact revision on arXiv

Source-linked AI summary

Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models

Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thiagarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, Aleksandra Faust

arXiv:2412.15287v2cs.CLcs.AIcs.LG

TL;DR

Inference-time computation can improve LLM performance, but conventional fine-tuning does not account for the inference strategy used at test time. This paper develops BoN-aware imitation-learning and reinforcement-learning methods that directly optimize for Best-of-N sampling. The methods improve BoN accuracy and pass@N, with reported gains on MATH and HumanEval.

  • Problem

    Conventional fine-tuning is agnostic to test-time inference strategy, although training for multiple attempts may call for diverse responses rather than only the model’s best attempt.

    Source: §1 · p. 2

  • Method

    The paper develops inference-aware supervised imitation-learning and reinforcement-learning methods that directly optimize the Best-of-N policy, using a variational formulation to handle its non-differentiable argmax.

    Source: §1 · pp. 1–2, §2 · p. 3

  • Results

    BoN-aware fine-tuning improves BoN accuracy and pass@N over baselines; Gemma 2B Bo32 performance on Hendrycks MATH rises from 26.8% to 30.8%.

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

  • Takeaways & Limitations

    BoN-aware fine-tuning learns a meta-strategy that interleaves best responses with more diverse responses suited to BoN sampling.

    Source: §7 · p. 13

    The RL formulation requires samples from the BoN distribution; the paper discusses BoN distillation as a way to alleviate this requirement.

    Source: §4 · p. 6

Abstract

from arXiv · show

Recent studies have indicated that effectively utilizing inference-time compute is crucial for attaining better performance from large language models (LLMs). In this work, we propose a novel inference-aware fine-tuning paradigm, in which the model is fine-tuned in a manner that directly optimizes the performance of the inference-time strategy. We study this paradigm using the simple yet effective Best-of-N (BoN) inference strategy, in which a verifier selects the best out of a set of LLM-generated responses. We devise the first imitation learning and reinforcement learning~(RL) methods for BoN-aware fine-tuning, overcoming the challenging, non-differentiable argmax operator within BoN. We empirically demonstrate that our BoN-aware models implicitly learn a meta-strategy that interleaves best responses with more diverse responses that might be better suited to a test-time input -- a process reminiscent of the exploration-exploitation trade-off in RL. Our experiments demonstrate the effectiveness of BoN-aware fine-tuning in terms of improved performance and inference-time compute. In particular, we show that our methods improve the Bo32 performance of Gemma 2B on Hendrycks MATH from 26.8% to 30.8%, and pass@32 from 60.0% to 67.0%, as well as the pass@16 on HumanEval from 61.6% to 67.1%.

1. Introduction

Inference-time computation can improve LLM performance, but it is usually treated separately from training. This work trains models with the Best-of-N (BoN) inference strategy in mind and reports improved BoN and pass@N performance.

1. Introduction

  • Inference-time search, re-ranking, and revision can improve LLM performance, while conventional training usually treats inference strategy as a separate post-hoc choice.The paper argues that training for multiple attempts could favor diverse problem-solving strategies rather than only the model’s best attempt.Source: §1 · p. 1
  • The paper fine-tunes models to directly optimize BoN performance, where a verifier selects the best among multiple generated responses.When the verifier uses ground-truth scoring, BoN is equivalent to pass@N.Source: §1 · pp. 1–2
  • Bo32 performance on Hendrycks MATH rose from 26.8% to 30.8% for Gemma 2B, while pass@32 rose from 60.0% to 67.0%.HumanEval pass@16 rose from 61.6% to 67.1%.Source: §1 · p. 2
  • The authors develop BoN-aware supervised imitation learning and reinforcement learning methods, and identify a temperature–sample-count trade-off for exploration and exploitation.They also describe specialized algorithms inspired by pass@N optimization that promote implicit exploration when environment rewards can verify responses.Source: §1 · p. 2

2. Inference-Aware Fine-Tuning: A Case Study with Best-of-N

The section frames fine-tuning around the inference strategy used at test time, focusing on Best-of-N (BoN). It describes a temperature–sample-count trade-off and a variational route to optimizing BoN despite its non-differentiable selection step.

Inference-Aware Fine-Tuning: A Case Study with Best-of-N

  • Standard SFT and RL are agnostic to test-time inference strategy, whereas inference-aware fine-tuning alters their objectives to account for that strategy.The section focuses on BoN, which samples multiple responses at a temperature and selects one using a verifier score.Source: §2 · pp. 2–3
  • BoN applies a verifier to choose among temperature-controlled samples, defining a test-time policy distinct from the learned policy.This distinction is the training–inference gap that the BoN-aware formulation addresses.Source: §2 · p. 3
  • The BoN argmax is non-differentiable, so the authors derive a variational formulation to construct imitation-learning and RL algorithms without top-k operators.The passage notes that differentiable top-k alternatives introduce approximation error and substantially increase computational cost.Source: §2 · p. 3
  • The verifier score may be a learned proxy for true reward; the paper does not assume the two are equal.The footnote says they may be related or even equal, but makes no such assumption.Source: §2 · p. 3
  • For easier MATH problems, small optimal temperature and sample counts can suffice; harder problems favor higher optimal temperatures and sample counts.Marker size in Figure 1 indicates how frequently each temperature–sample-count pair gives the best BoN performance.Source: Figure 1 · p. 3, §2 · p. 4
  • Adjusting temperature and sample count offers a computationally inexpensive model-selection option, while BoN-aware fine-tuning directly trains for the exploration–exploitation trade-off.The text cautions that model selection alone may not capture all nuances of LLM behavior.Source: §2 · p. 4

3. Supervised BoN-Aware Fine-Tuning

BoN-aware SFT adapts likelihood training to the BoN policy when only expert examples, rather than true rewards, are available. A variational formulation avoids differentiating through BoN’s argmax and adds an exploratory win-rate term.

3. Supervised BoN-Aware Fine-Tuning

  • BoN-SFT assumes access to expert examples but not true reward, and the non-differentiable BoN argmax makes its objective difficult to optimize.The authors address this difficulty with a variational approximation.Source: §3 · p. 4
  • The win-rate weight λ_N increases with sample count N and is approximately proportional to N for large N.The passage identifies λ_N as the solution to a one-dimensional search problem.Source: §3 · p. 4
  • The variational objective combines expert-data likelihood with a term that encourages data responses to win against policy samples, accounting for BoN sampling.The win rate can be approximated with a softened sigmoid of verifier-score differences.Source: §3 · pp. 4–5
  • The derived gradient combines expert-response likelihood with a correction based on responses sampled from the BoN policy.This formulation lets BoN-SFT use standard gradient-ascent algorithms despite the non-differentiable distribution.Source: §3 · p. 5
  • The gradient terms clone expert behavior, align verifier rankings, limit likelihood on BoN samples, and reduce the win-rate gap between best and average samples.Together, these terms are described as limiting overfitting while maintaining overall response quality.Source: §3 · p. 5

4. BoN-Aware Fine-Tuning Using Reinforcement Learning

BoN-aware RL trains the policy for the inference-time selection-and-reward pipeline, with a general method and a more efficient variant for binary rewards. Its weighting schemes prioritize difficult prompts and can support exploration using either positive and negative examples or positive examples alone.

4. BoN-Aware Fine-Tuning Using Reinforcement Learning

  • The general BoN-RL gradient is REINFORCE-style, but estimating both its value baseline and policy gradient can be sample-inefficient, especially at large N.The text points to BoN distillation as an approach to alleviate the need for BoN-distribution samples.Source: §4 · pp. 5–6
  • BoN-aware RL draws N responses from the base policy, ranks them with a verifier, and optimizes environment reward on the selected BoN samples.Using environment reward can also let the base policy tolerate potential verifier errors.Source: §4 · p. 5, Figure 2 · p. 6
  • With binary success/failure rewards known to the verifier, a closed-form BoN policy yields an efficient gradient estimator without value estimation.The policy weights base responses according to the prompt's failure probability.Source: §4 · p. 6
  • BoN-RLB asymmetrically weights positive and negative examples, emphasizing successful responses on harder prompts and redistributing likelihood away from errors to support exploration.As prompt failure probability approaches 1, positive samples receive exponentially more influence.Source: §4 · p. 7
  • BoN-RLB(P) learns from positive examples alone and upweights harder BoN-generated positives, addressing settings where collecting negative data may be difficult.The text identifies safety-critical settings as one possible use case.Source: §4 · pp. 7–8
  • BoN-SFT uses expert data and verifier-selected positive outputs, while RL variants differ in whether selection uses a verifier, environment reward, or binary feedback.The methods are summarized in Table 1.Source: §4 · pp. 7–8, Table 1 · p. 8

5. Experiments

The experiments quantify how BoN sample count and temperature interact, then test whether BoN-aware fine-tuning improves BoN and pass@N across math and code benchmarks.

5.1. Co-scaling Analysis of Sample Size 𝑁and Temperature 𝑇in BoN

  • Pass@N rises with N, while BoN accuracy generally declines as N grows, especially at higher temperatures; moderate N balances exploration and exploitation.At very low temperatures, BoN accuracy improves with N, and pass@N’s optimal temperature is positively correlated with N.Source: §5.1 · p. 8, Figure 3 · p. 9

5.2. Inference-aware Fine-tuning with BoN Sampling

  • BoN-SFT with N′=32 achieves the best BoN accuracy and pass@N, whereas base SFT degrades performance relative to the base model.The authors interpret the BoN-SFT result as producing both high-quality and exploratory responses.Source: §5.2 · p. 9, Figure 4 · p. 10
  • 26.8% to 30.8% Bo-32 accuracy and 60.0% to 67.0% pass@32 on MATH, plus 61.6% to 67.1% pass@16 on HumanEval, mark gains from BoN-aware RL.BoN-RL-V improves Bo-32 accuracy; BoN-RL-S improves MATH pass@32, and BoN-RLBP improves HumanEval pass@16.Source: Figure 5 · p. 10, §5.2 · pp. 10–11, Figure 6 · p. 11
  • Training BoN-RL-V with N′=32 improves accuracy across evaluation N=1 to 32, including N=1 accuracy from 22% to 26%.The reported gains are more pronounced at higher evaluation N.Source: §5.2 · p. 10
  • Pass@N-aware RL methods improve pass@N across math and coding, while standard RL with N′=1 lowers pass@32 on MATH and pass@16 on HumanEval.BoN-RL-V does not significantly improve pass@N, which the authors attribute likely to its focus on robustness to a verifier rather than diverse correct solutions.Source: §5.2 · p. 11, Figure 6 · p. 11
  • BoN-RL-V improves on both held-out math benchmarks and outperforms the base model across evaluation temperatures 0.1, 1.0, and 1.5.BoN-RL-S also improves pass@N across temperatures, with larger gains at higher temperatures; BoN-RL-V is more resilient to high-temperature BoN accuracy degradation.Source: §5.2 · p. 11, Figure 7 · p. 12, Figure 8 · p. 12

6. Related Work

Prior work uses inference-time computation through methods such as Best-of-N sampling, self-correction, and search. This paper instead fine-tunes directly for Best-of-N performance and incorporates verifier errors during training.

6. Related Work

  • Best-of-N rejection sampling generates multiple outputs and selects one using oracle information, as in pass@N, or a learned verifier [54].It can also be combined with chain-of-thought prompting [54].Source: §6 · p. 12
  • Prior fine-tuning targets self-correction [Kumar et al., 2024; Snell et al., 2024] or search [Gandhi et al., 2024; Lehnert et al., 2024], whereas this work directly optimizes Best-of-N performance.STaR and ReSTEM use on-policy BoN samples to fine-tune proposal distributions toward higher-quality outputs [Zelikman et al., 2022; Singh et al., 2023].Source: §6 · pp. 12–13
  • For tree search, Dalal et al. (2021) penalize actions leading to states with high Q-value error; this work instead incorporates verifier error into training and frames the verifier as an approximate Q-function.The paper describes its BoN framework as conceptually generalizable to tree search.Source: §6 · p. 13

7. Conclusion

The paper presents inference-aware fine-tuning for Best-of-N (BoN), developing algorithms that align training with inference-time sampling. It reports improved BoN accuracy and pass@N on MATH and describes a learned strategy that interleaves strong and diverse responses.

7. Conclusion

  • BoN-aware fine-tuning significantly improves BoN accuracy and pass@N on MATH over state-of-the-art baselines.The paper also highlights robustness and generalizability across BoN configurations.Source: §7 · p. 13
  • The framework develops fine-tuning algorithms for imitation learning and reinforcement learning that account for BoN inference.The paper also reports a BoN co-scaling law to guide temperature and sample-size optimization.Source: §7 · p. 13
  • BoN-aware fine-tuning learns a meta-strategy that interleaves best responses with more diverse responses suited to BoN sampling.The authors identify extending the framework to more complex inference algorithms and larger-scale problems as future work.Source: §7 · p. 13

A. Theoretical Derivations

The appendix derives variational and gradient formulations for BoN-aware training, then outlines implementation practices for its SFT and RL methods.

A.1. Variational Approximation of BoN

  • With unique verifier scores and finitely many base-model outcomes, the BoN policy can be approximated by a variational optimization problem and represented as an exponential twist of the base policy.The twist uses an expected win-rate term scaled by λ_N.Source: §A.1 · p. 18

A.2. Proof of Lemma 1

  • The BoN-aware SFT gradient is expressed as a difference between expectations over expert data and the BoN policy.Its feature-gradient term includes the base-policy log-likelihood gradient and a verifier-score comparison against a sampled response.Source: §A.2 · pp. 18–19

A.3. Proof of Lemma 2

  • Applying REINFORCE to the BoN policy yields an RL policy gradient that weights reward-adjusted gradients by a learned baseline.The derivation includes an additional term involving verifier comparisons and λ_N.Source: §A.3 · p. 19

A.4. Proof of Lemma 3

  • For binary reward, the gradient can be written as an expectation over BoN responses with reward 1.The proof obtains this form using the log-likelihood trick and the BoN distribution.Source: §A.4 · p. 20

A.5. Proof of Corollary 4

  • The Corollary 4 proof derives a gradient for its objective using the log-likelihood trick and the BoN distribution.The supplied passage includes a gradient expression involving the reference indicator and reward-zero responses.Source: §A.5 · p. 20

B. Pseudo-code and Implementation Details

  • Implementation uses an exponential-moving-average anchor policy, BoN-aware losses, and an additional KL-divergence loss to the anchor.The appendix also notes value baselines, batch-normalized advantages, and clipping in RL implementations.Source: B. Pseudo-code and Implementation… · p. 21

B.1. Analysis of Bon-RLB Weights

The appendix analyzes how BoN-RLB weights, entropy-regularized RL formulations, and BoN distillation relate to exploration, optimization, and training efficiency.

B.1. Analysis of Bon-RLB Weights

  • For hard problems and larger N, BoN-RLB’s positive-sample weights rise sharply, emphasizing scarce successes but increasing estimation noise and potentially destabilizing convergence.The positive-only weight rises more moderately because its formula includes a subtraction term.Source: §B.1 · p. 23

C.1. Entropy-regularized RL

  • A pairwise-label formulation can eliminate terms independent of the output and avoid solving for the value function V.The resulting ℓ2 minimization resembles IPO, but its dependence on Qπ prevents a guarantee of a unique solution or an optimal BoN base policy.Source: §C.1 · p. 24
  • An alternative entropy-regularized RL formulation iteratively updates the base policy with importance-weighted BoN likelihood and updates the importance weights themselves.The value function is an auxiliary variable in this weighted-SFT view.Source: §C.1 · pp. 24–25

C.2. Improved Efficiency with BoN Distillation

  • BoN distillation addresses the cost of generating N samples for every gradient update by first matching a BoND policy to the current BoN distribution, then using it for approximate gradients.The text states that convergence to an optimum of the original objective is possible if the approximate gradient remains an ascending direction.Source: §C.2 · p. 25

C.3. Connection to STaR

  • BoN-RL adds an importance-sampling gradient term to the reward-weighted SFT-like term shared with BoN-STaR.The added term accounts for the distribution shift between the base policy and BoN policy and emphasizes reward improvement.Source: §C.3 · p. 25
  • Figure 9 plots the pass@N exponent and BoN optimal sample count against temperature, marking in-training predictions with dashed curves and extrapolations with stars.The caption identifies the plotted quantities and marker conventions but does not state an outcome.Source: Figure 9 · p. 26

D.1. Additional Scaling Results with Gemma 9B verifier and policy models

Additional results examine how scaling predictions generalize across models and inference algorithms and compare Gemma-9B with Gemma-2B.

D.1. Additional Scaling Results with Gemma 9B verifier and policy models

  • The same power-law function with a linear trend term applies to BoN and MajorityVoting across different LLMs and inference algorithms, while MajorityVoting’s MC estimates yield less smooth curves and slightly lower prediction performance.Table 3 reports R-squared values for the language models and inference algorithms.Source: §D.1 · p. 26
  • Gemma-9B improves both pass@N and BoN significantly compared with Gemma-2B, and reward-model overoptimization is lesser for Gemma-9B.Figure 10 presents the pass@N and BoN performance curves for Gemma-9B.Source: §D.1 · p. 27, Figure 10 · p. 27
  • For Gemma-9B, the performance gap between temperatures 0.7 or 1.0 and temperature 0.1 also increases.The comparison is reported for the Gemma-9B policy and reward models.Source: §D.1 · p. 27

D.2. Details of BoN-aware Fine-tuning Experiments

This material details MATH and code-generation evaluation setups and reports Gemma 9B findings on MATH and held-out math benchmarks.

D.2. Details of BoN-aware Fine-tuning Experiments

  • For MATH, 7,500 original training problems were augmented with 4,500 test problems, and the remaining 500 were used for evaluation.Source: §D.2 · p. 27
  • SFT distills correct Gemini 1.5 Flash solutions, while RL uses binary answer correctness and a separate pretrained Gemma 2B verifier trained on Gemini-generated data.Source: §D.2 · p. 27
  • Code-generation models train on MBPP and are evaluated on HumanEval using standard procedures (Austin et al., 2021; Chen et al., 2021; Kumar et al., 2024).Source: §D.2 · p. 27

D.3.1. Hendrycks MATH with Gemma 9B

  • BoN-RL-V achieves the best BoN performance and BoN-RL-S the best pass@N performance on Gemma 9B MATH, with both substantially improving over the base model.Source: Figure 11 · p. 28, §D.3.1 · p. 28

D.3.2. Held-out math benchmarks with Gemma 9B

D.3.3. pass@N Accuracy with Different Training Samples, with Gemma 2B

These experiments compare pass@N training variants and BoN distillation baselines, with BoN-RL-V outperforming the distillation baselines in BoN accuracy.

D.3.3. pass@N Accuracy with Different Training Samples, with Gemma 2B

  • STaR's weaker pass@32 result is attributed to omitting negative samples, difficulty-based reweighting, and importance sampling between the base and BoN policies.Source: §D.3.3 · pp. 28–29
  • At N′=16, BoN-RL-S, BoN-RLB, and BoN-RLB(P) outperform baselines in pass@N; STaR reaches 55% pass@32 versus 60% for BoN-RLB.Source: §D.3.3 · p. 28, Figure 13 · p. 29
  • BoN-RLB and BoN-RLB(P) outperform BoN-RL-S at N′=16 but underperform it at N′=32, suggesting instability as N′ grows.Source: §D.3.3 · p. 29

D.3.4. Comparing BoN-RL with Base-BoN Distillation Baselines, with Gemma 2B

  • Four distillation baselines target the best sample, all samples, verifier-weighted resamples, or a majority-voted sample from N=16 responses.Source: §D.3.4 · p. 29
  • BoN-RL-V outperforms four distillation baselines in BoN accuracy, although those baselines also improve over base Gemma 2B.Source: §D.3.4 · p. 29, Figure 14 · p. 30
Loading 2412.15287v2…