Inference-Aware Fine-Tuning for Best-of-N Sampling in Large Language Models
Authors: Yinlam Chow
Guy Tennenholtz
Izzeddin Gur
Vincent Zhuang
Bo Dai
Aviral Kumar
Rishabh Agarwal
Sridhar Thiagarajan
Craig Boutilier
Aleksandra Faust
Venue: ICLR 2025 Poster (2025)
Keywords: Best-of-N sampling
Reinforcement Learning
Language models
Primary area: foundation or frontier models, including LLMs

Abstract: Recent studies indicate that effectively utilizing inference-time compute is crucial for attaining good performance from large language models (LLMs). Specifically, the Best-of-N (BoN) inference strategy, where an LLM generates multiple responses and a verifier selects the best, has shown strong empirical performance. Motivated by this, we develop a novel inference-aware fine-tuning paradigm, which encompasses the BoN-aware inference framework as a special case. We devise the first imitation learning and reinforcement learning (RL) methods for fine-tuning LLMs using BoN, overcoming the challenging, non-differentiable argmax operator in BoN. We empirically demonstrate that our BoN-aware models implicitly learn a per-example "meta-strategy", which 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 BoN performance of Gemma 2B on Hendrycks MATH from 26.8% to 30.8%, and Pass@K from 60% to 67%.

Published as a conference paper at ICLR 2025
INFERENCE-AWARE TINI-NUNING FOR BEST-OF-N SAMPLING IN LARGE LANGUAGE MODELS
Yinlam Chowl; Guy Tennenholtz', Izzeddin Guri, Vincent Zhuang!, Bo Dai',
Sridhar Thiagarajan', Craig Boutilieri, Rishabh Agarwal", Aviral Kumart, Aleksandra Faust
i Google Deepmind, $ Google Research
ABSTRACT
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 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 verifier selects the best out of set of L.-ggggerarted responses. We devise the first imitation learning and reinforce- ment learning (RL) methods for BoN-aware fine-tuning, overcoming the challeng- ing, 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 test-time input-a process reminiscent of the exploration-exploitaiio exploittaxiion 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. 26.8% to 30.8%, and pass @ 32 from 60.0% to 67. 67.0%, as well the pass @ 16 on HvmanEval from 61.6% to 67.1%.
INTRODUCTION
An effective method for improving the performance of large language models (LLMM)) is leverage additional computation at inference-time: various works (Lightman et al., 2023; Wu et al., 2024; Kumar al., 2024; Hosseini et al., 2024) have shown that by using search, re-ranking, multi-turn revision, and more generally, any approach that makes use of more tokens and inference-time com- pute, the performance of LLMs on various tasks can be significantly improved- so much that in- vesting improving inference-time computation might prove more beneficial than increasing model pre-training compute (Snell et al., 2024).
Despite this promise, existing work largely considers using inference-time computation tional post-hoc design choice, after conventional pre-training and fine-tuning. However, decoupling training and inference-time computation is not optimal: for example, if we knew that LLM is allowed to make multiple attempts to solve a math problem, then it may better to fine-tune it to explore diverse problem-solving strategies, rather than simply generating the candidates that rep- resent the model's best attempt at solving the problem. Within the context of reasoning problems, these performance gains may be sigmificant. as LLMs often fail due to their inability to draw complex inferences about the input and their internal knowledge (Chen et al., 2024).
op-
We argue that the effectiveness of inference-time computation can be substantially increased by explicitly considering the inference procedure during training. We study this inference-awan fine- tuning paradigm using the Best-of-N (BoN) inference strategy, where the LLM generates multiple candidate responses, and verifier selects the best one according to some scoring function (Cobbe al., 2021). When this verifier is the ground-trutt scoring function, BoN is equivalent pass @N, a widely-used method for inference-time compute scaling (Brown et al., 2024). In contrast with traditional fine-tuning methods such supervised fine-tuning (SFT) reinforcement learning (RL),
"Correspondence yin 1 laamhhe@gggggleeee

---

Published as a conference paper at ICLR 2025
which agnostic to the inference strategy used at test-time, our inference-aware (BoN-aware) methods directly optimize the performance of the BoN policy, and lead to significantly improved BoN performance at test time.
Our work makes several key contributions to the understanding and optimization of batch-of- neighbors (BoN) inference. (1) We formally define the inference-aware and BoN-aware problem setting, recognizing the erucial role of the inference strategy during training.
(2) We establish co- scaling behavior for BoN, quantifying the inherent trade-off between exploration and exploitation governed by the temperature (T") and the number of samples (N). This analysis directly informs the design and optimization of our BoN-aware algorithms. By understanding how these interactions influence performance, we can effectively tune these parameters to achieve an optimal balance be- tween exploration and exploitation in both our supervised and reinforcement learning settings. (3) We develop a BoN-aware supervised fine-tuning algorithm that aligns the target distribution with the BoN policy distribution. (4) We extend our method general Now-a RL framework, en- abling the policy to learn and solve downstream tasks under the BoN inference strategy. To further enhance BoN-aware fine-tuning, pass N accuracy, which promote implicit exploration and connect with established self-supervised learning techniques, particularly in scenarios where the environment reward can be used for verifi- cation. (5) Empirically, our experiments demonstrate the effectiveness of BoN-aware fine-tuning in terms of improved performance and inference riference-time compute. In particular, we show that our meth- ods 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. 67.1%.
devise specialized algorithms inspired by methods optimizing
2 INFERENCE-AWARE
A CASE STUDY WITH BEST-OF-N
Standard fine-tuning methods typically train LLMs to produce the best response for a given prompt. In LLM ane-tuning. a model (or policy) is trained via supervised fine-tuning (SFT), by maximizing the likelihood w.r.t. ground-truth data. Formally, we search for a policy я : X' + Ay that maxi- mizes the likelihood Ex--Par " (ullz) [log æ (y]zz)], where here, X' and y are the space of prompts and outputs of an LLM, P is the prompt distribution, and nt' is a distribution of expert responses. Alternatively, the policy can be fine-tuned via reinforcement learning (RL) (Schulman al., 2017): max Exp-Payattt) [R(2, v)], to align the LLM's behaviors with the reward function R(x, y)- While popular, these methods have not taken the LLM's inference-time strategies into the account.
Inference-Award Aware Fine-Tuning, To address the gap between how LLMs are trained and how they are used at inference time, we develop inference-aware fine-tuning. During inference, the learned policy π is often not directly used; rather some inference strategy I II X + Ay applied to it. For example. I can be the BoN strategy, which samples multiple candidate responses, and selects the best using the score function of some verifier; or might be a search mechanism (Lightman et al., 2023) or self-correction (Kumar et al., 2024). To account for this inference strategy 7, we alter the objective SFT and RL objectives to be "aware" of the inference strategy:
max
TEll Bg-r.gg (aln) log (TT, 1²) (v)], and
(Inference- Aware SFT)
max J J(TT) == EP- [Rx,,)) y)]. TEll
(Inference-Aware RL)
Indeed, Inference-Awane SFT and Inference-Aware RL are aware of the strategy 7. In what follows, we focus on the case where the inference strategy is BoN (i.e., = BoN), in both the SFT and RL. setups. As we will later see, this brings about new algorithms for training the policy.
BoN-Aware Problem Formulation. We begin by defining the BoN strategy. This inference strat- egy samples N resposnes from a model with some temperature T, and then selects the best one, based on some verifier score. Formally, the BoN inference policy can be written as:
((T = bon (y) R. N, = arg max V'E uNY
((zgg") s.t. ys I (.1=), (1)
where I is a sample with temperature T. and r X R is a verifier score'. what follows, when N, T are clear from context, we write boon (y) Iz; T). We see that the above strategy defines class of BoN policies that is different from the learned policy म, demonstrating the gap
'The verifier score and the true reward it can be related, even equal, yet do not make that assumption here. Usually, r model trained to predict R, and therefore serves a proxy of the true reward.
2

---

Published as a conference paper at ICLR 2025
between training and inference. We inject this class of BoN polices into the Inference- SFT and Inference-Aware RL frameworks to derive the instantiation of inference-aware fine-tuning.
Besides closing the gap between training and inference and mitigating potential discrepancies be- tween the verifier score r and the true reward R. BoN policies provide further benefits. The BoN mechanism introduces implicit exploration during training, bypassing the computational burden of explicit exploration (Cen al., 2024). Selecting the best of N samples allows the base policy to explore output variations, inducing a controlled exploration that can lead to more robust and gener- alizable policies, particularly in scaling behavior w.r.t. temperature T and number of samples N.
Optimizing the BoN policy class notoriously difficult due to the non-differentiabiiiy of the arg max operator. Although several differentiable top-k operators (Cuturi et al., 2019; Xie et al., 2020) might be exploited Thon. they induce approximation error, and more importantly, increase the computational cost dramatically. In our work, we derive a variational formulation of the learning problem w.r.t. Tbo without top-k operators. allowing to construct novel algorithms for inference- aware BoN, using both standard supervised imitation learning (Section 3) and RL (Section 4).
Exploration- Exxpotttotinloooittiiinn with BoN Sampling. BoN sampling offers a natural mechanism for balancing explo- ration and exploitation in response generation. Adding BoN inference to the base model effectively explore the diverse possibilities within the model's s ouiput space
275
390
while also exploiting its knowledge to generate high- trade-
quality candidates. This iilloraatio-eeploittaiie off is crucial for solving various tasks and improving gen- eralizability. To quantify such trade-off, we empirically verify the implicit exploration and exploitation properties of BoN. We do this by revealing optimal co-scaling w.r.t. temperature T and number of samples N. Specifically, for a fixed base policy æ, at any prompt E X there is an optimal temperature T* (1) and optimal number of sam- ples No (z) which maximize performance of BoN:
The relationship between the op- timal number of samples (N") and optimal temperature (I")) in BoN. The size of cach marker
Figure 1:
N (z; r). (x)))) e arg max N,T
[R(x, v)-
given (T, coordinate indi- which the (T, N) pair resulted in the hest BoN per- formance. The plot reveals a trade-off: "eas- ier" problems have small T and N', while "harder" problems require larger T" for exploration and consequently larger N
To understand the connection between NO (x)) and T (z).
cates the frequency of problems
we assess the performance of Gemma 2B (Team et al.,
2024) on the MATH benchmark (Hendrycks et al., 2021), when applying the BoN inference strategy. Figure shows empirical frequencies of problems, when varying (x) and N (x) (larger marker size signifies higher fie- quency). The figure depicts tradeoff' between 7 and N, reminiscent of the particularly also a small N). These "easier" problems do not require heavy exploration (small T*) and can therefore be more exploitative (smail N*). On the other hand, T increases, the base policy π becomes more stochastic, resulting in more diversity and exploration. These more "diffi- cult" problems, require more exploration (larger T*), hence less exploitation (larger N*). Indeed, in such cases, the distribution of N shifts to high values. Our results suggest a tradeoff between exploration and exploitation, and further motivates the BoN-aware setup, which can account for this tradeoff uniformly across all samples.
trade-off. When Tx)) is small, any N is optimal (and
Figure also uncovers a cost-effective recipe for adjusting T' and N for optimal BoN performance: learn how to fine-tune the model for better inference by simply adjusting these accessible parameters. However, it is important to note that relying solely on model selection has limitations. While this approach offers computationally inexpensive way to improve BoN's 's inference-tin per- formance, it may not fully capture the nuances the LLM's behavior. With sufficient computational resources, general BoN-aware fine-tuning can further unlock performance gains by directly training the LLM to optimize for the ceporaatoon-pplittaaiio trade-off of the BoN inference process.
SUPERVISED BoN- AWARE -UNNNG
We begin by developing the BoN-aware SET framework. Under this setting we assume we do not have access to the true reward, and only wish to maximize the likelihood of dataset of expert
3

---

Published as a conference paper at ICLR 2025
examples. Recall the definition of the BoN policy Teea in Equation (1). The Inference-Aware SFT version of BoN becomes:
max (((x.)))<D [log Thom (y x;; m)];
(2)
TIll
A major difficulty in solving Equation (2) is the non-differentiabiiity of the arg max operator in the BoN procedure. To address this, we use the variational approximation of Thon (see Section A.1)
Tooa (y/zz) xc [T(U|=)) exp (x,
(3)
where Q- (I. = Exy-a((la)) ((aa))rraaa) is the expected win-rate over base policy x, charac- terizing the probability for which response outperforms the responses generated by the base over the verifier score r. The constant AN > is solution of a ID-search problem (Gui al., 2024) (see details in Appendix A.I). It can be shown that AN is monotonically increasing in N, and Ay o N approximately for large N. Plugging the variational form of Equation (3) into Equation (2) yields:
max E ((z,))) [log [b][() == B( log ग (ulz) AN Qx (z. v) log Z (x)
(4)
FEll
Likelibood
Inference- Awarensss
where (x) = ( (yla)) jexp (AN Q (ar, v))] is the partition function.
The above optimization problem reveals two term. While the first term tries to push the base pol- icy π into maximizing the likelihood of the data, the second term regularizes the policy to be more exploratory by increasing the data win rate over the policy. This in turn accounts for the sampling in BoN. For data efficiency when estimating the win rate Q~ (z, y) we leverage a common practice in modeling pairwise preferences (Rafailov et al., 2023) to approximate the win rate with its "soft- ened" counterpart: Q= (x, y) Rs Egy>>((-a)) [ar (r(x, y) - r(æ, y y))], where a is the sigmoid function. Next, we exploit properties of policy gradient (Sutton et al., 1999) and the gradient of energy- -ased policies (Rafailov et al., 2024) derive the gradient for Equation (4) (see Appendix A.2 for proof):
Lemma (BoN-SFT). The gradient of Equation (4) w.r.f. LLM parameters E O of is given by ((ay))D [o (x, y: 0)] - EEwD, D.gr-mool--- J (x,, wi 0)}. where
Vof (z. y: 0) = Ve log Te (y (r) AA -y'' [V log me (v'lz)-1(r(x; < r(2,//) (5)
Our formulation circumvents the di-edeefeeeetibbiittyy of the BoN distribution, allowing solution of BoN-SFT via standard gradient-ascent algorithms. The individual terms of the gradient imply the following: (1) clones the expert behavior by maximizing its likelihood over D; (2) it aligns with the verifier score ranking, which assigns a high win-rate to the expert over the base; (3) it avoids fitt-fitting by limiting its likelihood over the BoN sample; and (4) it maintains overall response quality by reducing the win rate between its best and average samples.
4 BON-AWARE FINE-TUNING USING REMENTTCEMENT LEARNING
Training LLMs that are amenable to BoN sampling can be framed within the RL paradigm, which trains an agent (LLM) that optimizes its actions within an environment. In this context, the LLM generates N responses (candidates actions) for given prompt (contexts). A separate macro agent (verifier) selects the candidate deemed most suitable according to a predefined criterion (e.g., prob- ability of success). This action then deployed the environment, yielding a reward (e.g., task completion). The key challenge in training this agent lies in achieving two objectives simultane- ously: (i) Enhancing agent's 's exploration capabilities generate diverse candidates that cover the space of potential solutions and align with the verifier's preferences; (ii) Maximizing the environ- ment reward of the final response. Motivated by this observation, we utilize RL for BoN-aware fine-tuning, enabling the development of more robust and adaptable LLMs. A schematic of the BoN-Awart RL framework is shown in Figure 2.
The BoN-Aware RL problem takes the following form:
(6) We train the BoN policy Thon (paramterized by л) to attain a high environment reward. Apart from enabling better exploration, using the environment reward R(x, y) in BoN-RL allows the base policy to tolerate potential errors in the verifier r(z, y). We first develop general algorithm for solving the BoN-aware RL problem. We then study an important subclass which assumes a binary reward, common feature of many reasoning problems (e.g., math, code).
max J (T) == EEg- m (1.7) [R(z, y)].
HEIl
4

---

Published as a conference paper at ICLR 2025
Thom
Environment Reward
Verifier
aggnac
22
π
3/
R
(k.st))
Update
Figure 2: BoN-Aware RL fine-tuning where N independent samples are first drawn from base и and ranked by the verifier r. The BoN policy then optimizes the environment reward of the BoN samples (Lemma 2).
We begin with deriving a gradient estimator to the objective in Equation (6). Exploiting the con- nection between the BoN policy and its energy-based policy counterpart in Equation (10), and using derivations analogous to those in Lemma 1, we compute the gradient of J(Ø), which leads REItFORCE-style algorithm (Williams, 1992) (see Appendix A.3 for proof):
Lemma (BoN-RL). The gradient of Equation (6) w.F. LLM parameters
of is given by
(0) EDwD D.g.D go (-12)) [V log TT (I, 9) (R(x, y- b(x))]
AN Er N m p V log Te (v/[z) (r((,, y) < r(x,y')3 (R(x, y) - b(z))}, (7) where b(z) = E y (m) R(x y) a baseline for variance reduction (Schulman et al., 2015).
This formulation resembles the standard REINFORCE gradient with the main difference of drawing samples from the BoN policy (instead from the base policy л). This allows one to solve BoN-RL via actor-critic methods (Sutton et al., 2009a). In practice, one can replace b(x) with a learned value baseline (z)) parameterized by th, for which & is updated by gradient descent w.r.t. the critic value loss. While BoN-RL inherits the benefits verifier alignment from BoN-SFT, and can be viewed as reware-weighted variant of the popular STaR method (Zelikman et al., 2022), generally it can be rather sample inefficient (especially when N large), as estimating both the value function b(x:) and the policy gradient in BoN-RL require samples from the BoN distribution. See Appendix C for discussion on alleviation using BoN distillation (Sessa al., 2024).
BoN-RL with Binary Reward and Verifier. While Lemma 2 provides a general method for BoN- aware RL, the policy gradient estimator in Equation (7) is sample ineffiecient for general rewards and verifiers. However, many domains admit binary success/failure metrics (e.g., reasoning tasks, math, coding) which allow an efficient gradient estimator, obviating the need for value estimation. Specifically, with binary reward known to the verifier, i.e., R(x, y) r(x, y) E {0, 1), Theorem of Sessa et al. (2024) implics the following closed-form solution the BoN policy Abon:
( [z) Prai (x)
if R(I, y) =0
Thon [x)
(8)
7 m) (1 M)) if R(x, y) =1"
where Pan Ey--)(()l=) [1 [(==g(===)) is the fraction of problems on which the base policy is incorrect, Under the binary assumption, Thon is a weighted distribution of the base policy x, whose importance sampling ratio depends the its failure probability (z). Introducing this closed form of Toon to Lemma 2, we obtain the following policy gradient (see Appendix A.4 for proof):
Lemma 3 (BoN-RLB). Assume R(x, v) E {0. The gradient of Equation (6) w.E.t LLM param- elers E 4 T
π given by
EznD
B Byy-oon (o log [og))(y(z)] g (Pra((1)) + By--tan, R=o (vo log [o(y|>)) g (Poulx))
where the positive and negative sample-dependent weights are given by
N N. N
N.p. 1-
g (v)
N 9
Lemma 3 not only reveals an efficient policy gradient estimator for binary reward, but more impor- tantly demonstrates how BoN-RLB balances positive and negative examples in its gradient update. It proposes novel way to re-weigh BoN-RLB's training examples, which prioritizes harder examples (as (x) - 1) by giving their positive samples exponentially more influence and aggressively redistributing log likelihood away from incorrect responses. The significance of this asymmetric
5

---

Published as a conference paper at ICLR 2025
Table 1: Summary of Our Best-of-N-Aware Fine-Tuning Methods
Offline Data
Reward (R)
Verifier (R+r)
Binary (R=r)
Positive Only
Closed
Method
Form
BoN-SFT
X
BoN-RL-V
X
X
BON-RL-S
x
BoN-RLB
X
X
X
BoN-RLB(P)
weighting scheme is that it infuses implicit exploration capabilities to the base policy. As Tajwar et al. (2024) observed, when the model reduces the likelihood of negative responses, it shifts that probability mass towards mode of the learned policy essentially reinforcing existing successful strategies (exploitation). However, these high-likelihood regions later produce errors, the resulting negative gradient redistributes this mass again, pushing the model to explore other potential solu- tions. This iterative process of concentrating probability mass and subsequent redistribution through negative gradients drives a form of exploration: as long as the model can produce correct solutions, it need not devote most of its sampling budget N on that but can also explore more diverse solutions.
Positive-only Weighting. Although we have illustrated the benefits of an asymmetric weighting scheme in BoN-RLB for exploration, training with both positive and negative examples may be infeasible (e.g., in data-limited online RL system that only records positive examples). To tackle this, we apply change of measure to Lemma 3 with the BoN distribution to derive policy gradient that only involves positive examples (see Appendix A.5 for proof):
Corollary 4 (BoN-RLB(P). Assume R(xx v) E (0, The gradient of Equation (6) W.E. LLM parameters 0 E ② of π is given by EEnD EygoThen Rt1 Vo log no To(y)x)] (x(((uu()))] where (p) N.p (1-p) (1-p")
Notice that the weighting gt (p. N) is monotonically increasing in p E [0,1] and lies within [0, 1] for any N. Using this gradient update, BoN-RLB(P) resembles weighted version of BoN-STaR (see Remark C.3 in the appendix). where it clones positive examples generated by the current BoN policy and up- weights the more difficult ones, where (r is close to 1.
Table summarizes different BoN-aware fine-tuning methods including: (1) BON-SFT from Lemma 1. (2) BoN-RL from Lemma 2 with a verifier, BoN-RL-V, (3) BoN-RL with environment reward as verifier. BoN-RL-S, (4) BON-RLB from Lemma 3, and (5) BoN-RLB(P) from Corollary 4. These methods are designed to leverage the BoN selection strategy during model fine-tuming, under different settings mentioned in Sections 3 and 4. BoN-SFT uses expert data but doesn't involve explicit reward learning. It relies on a verifier to select the best output and only trains with positive examples. BoN-RL-V and BoN-RL-S employ RL optimize BoN performance, with the former using a verifier for response selection while the latter relying solely on a reward signal. BoN-RLB and (ON-RIB((P) also utilize RI. on the special case of binary reward feedback (R They both have closed-form solutions, which can lead more efficient learning. BoN-RLB(P)) only learns from positive examples, which might be beneficial for (e.g., safety-critical situations where collecting negative data is difficult.
(0. 1))-
5 EXPERIMENTS
In this section we address the following questions: (1) Can we quantify the co-scaling relationship between the BoN number of samples N and temperature T, enabling joint optimization of these parameters? (2) Do inference-aware fine-tuning methods (SFT and RL) enhance the effectiveness of BoN sampling? (3) Do these improvements generalize across problem domains and BoN settings?
5.1 CO-SCALING ANALYSIS OF SAMPLE SIZE N AND TEMPERATURE T IN BON
We study the co-scaling behaviors of BoN and pass@ N performance over varying N and T by experimenting pre-trained Gemma 2B and 9B models on the Hendrycks MATH bench- mark (Hendrycks et al., 2021). Our experiments showed consistent co-scaling behaviors across different model sizes, therefore we summarize the results of 2B model below and include the 9B findings in Appendix D. As illustrated in Figure (Figure 9, Figure 10 in Appendix D.1), pass @ N consistently increases with higher N, as commonly observed (Brown et al., 2024). Our analy- sis suggests that this relationship can be captured by a power-law function of the following form:

---

Published as a conference paper at ICLR 2025
Temperature
0.275
0.01
0.250
PasseN Acturacy
0.1
BoN Accuracy
0.7
0.225
1.0
0.200
0.5
0.175
0.4
0.150
0.3
0.125
0.2
0.100
10"
10 10P
10
10
Figure 3: BoN and pass@ N performance of Gemma 2B policy and reward models w.r.t varying N and T. pass @ N monotonically improves with N; BoN shows inflection points as N increases. Colored dashed lines denote predictions of scaling functions; black dashed lines in BoN plot denote the last inflection points.
PT))@N(T) R2 exp(a(T)NN*T) where the parameters (T) and b(T) are lemperature-dependent and derived by fitting the model to data at a specific temperature T. Further analysis of this scaling behavior (detailed in Appendix 1.1) indicates that there is a strong positive correlation between the optimal temperature and N, which aligns with the intuition that larger N benefits from broader ploration (higher 7',, while smaller N favors focused exploitation (constant T). This relationship is straightforward, as there are no verifier errors confounding the selection of best responses.
Our experiments demonstrate an intriguing relationship between BoN accuracy, T, and N. We find that lower temperatures generally yield better BoN accuracy. Furthermore, BoN accuracy generally decreases as N increases, but degrades more rapidly with higher temperatures. With larger N and T. the increased randomness in the base policy inherently generates more "bad" samples (with poor accuracy). This phenomenon suggests that the verifier is sensitive to noise and may mistakenly select random outputs generated at higher temperatures as the best responses due to misalignment with the true reward (Type II error). Conversely, at very low temperatures, BoN accuracy improves with N, indicating the algorithm remains in an exploitation phase. Optimal performance is observed at moderate N values, striking a balance between exploration and exploitation.
5.2 INFERENCE-AWANE FINE-TUNING WITH BON SAMPLING
Experimental setup. We study 2B and 9B Gemma 2 models (Team et al., 2024) on canonical mathematical problem-solving and code generation benchmarks. For math, we train and evaluate on Hendrycks MATH benchmark (Hendrycks et al., 1021), and additionally evaluate on two held-out math benchmarks: Functional MATH (Srivastava al., 2024) and MathÖdyssey (Fang al., 2024). For coding, we train on MBPP (Austin et al., 2021), 1), and evaluate on HumanEval (Chen et al., 2021). More details of our experiments can be found in Appendix D.2.
Our main evaluation metries are BoN and pass@N N accuracies, i.e. the accuracies of the policies defined in Equation (1) with a learned verifier and ground-truth reward respectively. For math, we consider both metrics, using learned verifier for BoN, and for coding, we consider only pass @ N because test-case feedback is usually available.
We benchmark our proposed methods BoN-SFT, BoN-RL-V, BON-RL-S, BoN-RLB, and BoN- RLB(P) against several baselines: (1) STaR from Remark C.3, which uses self-training over cor- rectly generated responses; (2) RL (Lee et al., 2023) with verifier feedback (RL-V); (3) RL with environment feedback (RL-S); (4) standard SFT of the base policy (SFT, N' 1); and (5) BoN with the preitrained model (Base model). We denote by (N', and (N, T) the number of samples and temperature used in training and evaluation, respectively. We use T = T" 1.0 unless speci- fied otherwise. Similar to co-scaling, our experiments show similar trends with 2B and 9B models, therefore we summarize that of the 2B model below and defer the 9B results to Appendix D.3.
BoN-aware supervised fine-tuning. We first evaluate the BoN and pass@) N performance of of- fline SFT methods, including BoN-SFT with various N', and base SFT (N" 1), with results shown in Figure 4. We find that base SFT significantly degrades upon the base model, indicating that it causes overfitting or lack of generalization. BoN-SFT able to improve the BoN accuracy significantly, especially with increasing N'. We find that BON-SFT with N' 32 achieves the best performance for both BoN accuracy and pass @ N, suggesting that it is able produce both high-
7

---

Published as a conference paper at ICLR 2025
0.28
0.26
Pass Accuracy
BoN Accuracy
Base
0.24
SFT N==1
0.4
BoN-SFT
0.22
BoN-SFT N'
BoN SFT =2=
0.3
0.18
10
15
20
25
30
5
15
25
Figure 4: BoN Accuracy and pass@) N accuracy for BoN-SFT with Gemma 2B models.
0.30
BoN Accuracy
0.29
BoN Accuracy
0.26
0.28
Base model
BoN-RL-S M1
BoN-RL-V
0.26
BoN-RLD #1=32
AV Nw1
0.24
Base model
BoN BoN-RL-V N-
BON-RL-V N'=4
RLV N'
10
20
30
10
15
20
25
N
(a) Comparison of different methods.
(b) RoN-RL-V with various N' values.
Figure 5: BoN accuracy on BoN-RL. methods and baselines, with Gemma 2B model on MATH.
quality and exploratory responses. To improve substantively over the base model, we next turn to RL, which should be more effective by virtue of being on-policy,
BoN-RL-V improves BoN Accuracy. Our BoN RL algorithm, BoN-RL-V, with N' 32, sig- nificantly outperforms several baselines (Figure 5a). Specifically, it boosts the Bo-32 accuracy of our base model from 26.8% to 30.8%. As expected, the inference-unaware RL-V method performs poorly, likely due to common reward hacking issues (Jinnai al., 2024). We also find that training with the same verifier used as test-time is critical our other proposed inference-aware methods that use the environment reward instead of the verifier (BoN-RL-S and BoN-RLB) show improvement over the base model but are substantially worse than BoN-RL-V.
We plot the performance of BoN-RL-V with different training N' in Figure 5b, and observed the best performance when training with 32 samples (N' 32). Interestingly, although the gains are more pronounced at higher N values, training with large N' leads to consistent improvements across all N 1 to 32. This suggests that RL-BoN-V not only optimizes the specific BoN setting was trained on, but also generalizes to other BoN configurations, including large improvements on direct N accuracy (from 22% to 26%). We hypothesize that the RL-BoN-V method significantly enhances BoN accuracy by effectively exploring larger sample space during training. Training with large N' allows the base policy to explore wider range of responses to generate higher-quality responses, similar to how effective exploration in RL leads to better performance and generalization across different scenarios.
BoN-RL-S, BoN-RLB, and BoN-RLB(P) improve pass@ N. We next analyze the pass@ N per- formance of various methods on math and coding benchmarks, aiming to understand how differ- training approaches impact the models' ability generate diverse and correct solutions. As shown in Figure 6, our inference-awan aware methods designed to explicitly optimize pass during training (BoN-RL-S, BoN-RLB, or BoN-RLB(P)) can lead to better test-time @ across both domains, highlighting the importance of directly considering the desired evaluation metric during training. Notably, in the MATH domain (Figure 6a), BoN-RL-S significantly improves the pass@32 of Gemma 2B from 60.0% 67.0%. This suggests that by encouraging the model to explore wider range of solution strategies during training, we can substantially enhance its ability to gener- correct answers. Conversely, standard RL with N' 1 slightly degrades pass @2, indicating
8

---

Published as a conference paper at ICLR 2025
0.6
Pass@N Accuracy
Passi Accuracy
0.5
0.55
0.4
0.50
0.3
BoN-RLB N''32
BoN-RL-S N2-32
RL-S M-=1
RL-S N'=1
BON.RLB N
BoN-RL-V M'-32
0.40
0.2
N
N
(a) MATH.
(h) HumanEval.
Figure 6: pass @ N of BoN-RL methods with binary reward as verifier, with Gemma 2B models.
0,200
0.7
Pass Accuracy
0.6
PN Accuracy
0.175
0.150
0.5
0.125
0.4
Base-model
0.100
Base-model BON RLV N-16
BoN-ALVV N-16
0.075
0.3
BoN N-16
BoN-RLS
STaR_ 16
0.050
S1aR_16
0.2
BoN RLBP N-16
0.025
BoN-RLBP N-16
BoN-ALB N=16
BoN RLB N=16
10
15
20
25
30
(a) Functional MATH.
(b) MathOdyssey.
Figure 7: @ N of BoN-RL and baselines on held-out benchmarks with Gemma 2B models.
that over-optimizing for immediate performance can actually hinder the model's exploration capa- bilities and limit its overall problem-solving ability. Interestingly, BoN-RL-V does not significantly improve pass @ N, likely because its training objective focuses on robustness to verifier, which may not perfectly align with generating diverse correct solutions. Similar trends observed coding benchmarks (Figure 6b), where pass @ N -aware methods, particularly BoN-RLBP significantly out- perform the base model, increasing the pass 16 performance of Gemma 2B from 61. 61.6% to 67. 1 1%. Again, standard RI. fine-tuning (N" 1) negatively impacts performance, reducing pass @ 16 to 59.8%, demonstrating the downsides of solely focusing on reward maximization during training
Generalization of BoN-RL to held-out benchmarks and other temperatures. To further assess the generalization capabilities of our BoN aware fine-tuned models, we evaluate their performance two challenging, held-out benchmarks: Functional Math (Srivastava et al., 2024) and Math- Odyssey (Fang et al., 2024). As shown in Figure 7, our fine-tuned models demonstrate clear im- provements on both benchmarks. Our BoN-aware policies also generalize well to different sampling settings. As shown in Figure 8, BoN-RL-V V (trained with T' 1.0) consistently outperform the base model across all evaluation temperatures (I' E {0.1, 1.0, 1.5)). Similarly, BoN-RL-S show superior performance for pass @ N all temperatures, with increasing gains higher temperatures. This highlights the benefits of broader exploration, even after BoN-aware training. For BoN-accuracy, lower temperatures favored both models, but BoN-RL-V consistently excelled, particularly at lower regimes.
temperatures, demonstrating its generalizability across different apploration-exploitaeeo Furthermore, BoN-RL-V show greater resilience to accuracy degradation at higher temperatures, suggesting an enhanced ability to adapt to verifier failure modes and mitigate Type-Il errors.
6 RELATED WORK
Large language models (LLMs) can leverage inference- time computation to improve the quality of their generated outputs (Welleck et al., 2024), particularly on reasoning tasks. One common approach is to use chain-of-thought (Wei et al., 2022), where the model generates a step-by-step rationale before generating the final output. Another useful approach that be combined with

---

Published as a conference paper at ICLR 2025
0.325
Pass@ Accuracy
BoN Accuracy
0.300
0.275
0.5
0.250
0.4
0.225
T=0.1
Base Gemma 2b v2
0.3
T-0.1
Base Gemma 2b v2
T=100
BoN-RL-V 33232
T-10
BoN N'=32
T=1.5
0.2
5
10
15
20
25
30
10
20
25
30
N
(a) BoN Accuracy
(b) pass N
Figure 8: BoN and pass @ N over various temperatures. BoN-RI. with verifier and exact reward (solid lines) trained with T 1.0. Dashed and solid lines show the base (Gemma 2B) and finetuned models respectively.
Shain-of-thought is Best-of-N rejection sampling (Charniak & Johnson, 2005; Stiennon et al., 2020), which is our focus in this work. In Best-of-N, generate multiple candidate outputs from LLM and select the best output. BoN re-ranking can be done either using oracle information, such as checking final answers for solving math problems, which is also known pass @ N (Chen et al., 2021), or learned verifiers (Cobbe al., 2021; Lightman al., 2023; Hosseini et al., 2024; Zhang et al., 2024). Recent work also empirically analyzes strategies that optimally trade off additional test-time compute for improved performance (Wu al., 2024; Snell et al., 2024).
Closely related to our approach is prior work that fine-tunes LLMs to improve their self-correction capabilities (Kumar et al., 2024; Snell al., 2024) or search capabilities on planning tasks (Gandhi al., 2024; Lehnert et al., 2024), which allows for more efficient scaling with test-time compute. By contrast, our work focuses on inference-aware fine-tuning that directly optimizes for Best-of-N performance, instead of an intermediate capability that be used at test-time.
To make an LLM amenable to test-time scaling, techniques like STaR (Zelikman et al., 2022) or ReSTEM (Singh et al., 2023) have been employed fine-tune the model using on-policy data. This process leverages BoN sampling to iteratively generate better responses, and fine-tunes on this cu- rated data, for which the LLM learns to improve its proposal distribution, effectively increasing the likelihood of generating high-quality outputs during inference.
Finally, our work is related to recent work on leveraging tree search to enhance decision decision-making in RL (Dalal al., 2021). A key challenge in both BoN sampling and tree search lies in mitigating the impact of imperfect value estimation. Dalal et al. (2021) address this in tree search by penalizing actions leading to states with high Q-value emrrr, making inference more pessimistic for out-of- distribution samples. In contrast, we address BoN verifier error during training, learning responses robust to these errors and aligning training with inference. Our framework extends to tree search, using the verifier as an approximate Q-function and optimizing policy robustness to errors.
7 CONCLUSION
We introduced inference-aware fine-tuning, a novel paradigm that bridges the gap between training and inference for I.MMs. Specifically for the Best-of-N inference strategy, discovered scaling law for BoN that guides the optimization of temperature and sample size, developed a gamut of fine-tuning algorithms that handle various imitation learning and reinforcement learning settings, training L.MMs to generate diverse and high-quality outputs tailored for BoN inference, demonstrated the efficacy of these methods by significantly improving on BoN accuracy and pass@ N on the standard MATH reasoning benchmark over state-of-the baselines, highlighting the robustness and generalizability of our approaches across various BoN configurations.
Our work exemplified how BoN-aware fine-tuning learns meta mtratestrategy, which interleaves best sponses with more diverse responses that might be better suited for BoN sampling. These findings underscore the potential of inference-aware fine-tuning to unlock previously undiscovered capa- bilities LLMs through aligning training methodologies with inference-time compute strategies. Future work includes extending this framework to incorporate more complex, inference algorithms (e.g., reasoning, critique-ant reviise MCTS), developing contextual BoN-aware algorithms that can generalize to various tasks, investigating the interplay between the co-scaling of temperature, sample size, and BoN-aware fnn-nn and applying algorithms more larger-scale problems.
10

---

Published as a conference paper at ICLR 2025
Reproducibility Statement. We utilize the publicly available Gemma 2B and 9B language mod- els, the Hendrycks MATH benchmark, and the HumanEval coding benchmark - all accessible the research community. Our experimental setup is described in detail in Section 5. Furthermore, the appendix provides comprehensive pseudo-code (Algorithms 1 to 4) and implementation details for our BoN-aware fine-tuning algorithms BB.N-SF.. BoN-RL, BoN-RLB, and BoN-RLB(P)) We also delve into the theoretical underpinnings of our methods in the main text and the appendix, enabling a thorough understanding of our approach.
REFERENCES
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang. Carrie Cai, Michael Terry. Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.077:: 2021.
Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learn- ing from human preferences. In International Conference on Artificial Intelligence and Statistics. pp. 4447-4452 4455. PMLR, 2024.
Ahmad Beirami, Alekh Agarwal, Jonathan Berant, Alexander D' Amour, Jacob Eisenstein, Chirag Nagpal, and Ananda Theertha Suresh. Theoretical guarantees on the best-of-n alignment policy. arXiv preprint vrXiv:2401.01873 2024.
Bradley Brown, Jordan Juravsky, Ryan Ehrlich, Ronald Clark, Quoe V. Le, Christopher Ré, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2024. URL https//arra.oorjaas//47777.227
21777
Miguel A Carreira- Perpinan and Geoffrey Hinton. On contrastive divergence learning. Interna- tional workshop on artificial intelligence and statistics, pp. 33 40. PMLR, 2005.
Carriara-Perpinan
Shicong Cen, Jincheng Mei, Katayoon Goshvadi, Hanjun Dai, Tong Yang, Sherry Yang, Dale Schu- urmans, Yuejie Chi, and Bo Dai. Value-incentivized preference optimization: A unified approach to online and offline rlhf. arXiv preprint arXiv: 19320, 2024.
Eugene Charniak and Mark Johnson. Coarse-to-fine n-best parsing and maxent discriminative reranking. In Proceedings of the 43rd Annual Meeting of the Association for Computational Linguistics (ACL '05), pp. 173-180, 2005.
Lingjiao Chen, Jared Quincy Davis, Boris Hanin, Peter Bailis, Ion Stoica, Matei Zaharia, and James Zou. Are more Ilm calls all you need? towards scaling laws of compound inference systems. arXiv preprint 24X14:2403.02419, 2024.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374 2021.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint 10. 14116 2021.
Marco Cuturi, Olivier Teboul, and Jean-Philippp Vert. Differentiable ranking and sorting using optimal transport. Advances in neural information processing systems, 32, 2019.
Bo Dai, Niao He, Hanjun Dai, and Le Song. Provable bayesian inference via particle mirror descent. In Artificial Intelligence and Statistics, pp. 985-994. PMLR, 2016.
Gal Dalal, Assaf Hallak, Steven Dalton, Shie Mannor, Gal Chechik, et al. Improve agents with- out retraining: Parallel tree search with off- oof-policy correction. Advances in Neural Information Processing Systems, 34:5518-555 2021.
Meng Fang, Xiangpeng Wan, Fei Lu, Fei Xing, and Kai Zou. Mathodyssey: Benchmarking mathe- matical problem-solving skills in large language models using odyssey math data, arXiv preprint arXiv:2406. 18321, 2024.
11

---

Published as a conference paper at ICLR 2025
Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, and Noah D Goodman. Stream search (sos); Learning to search in language. arXiv preprint (rXiv:2404.03683 2024.
Hans U Gerber, Elias SW Shiu, et al. Option pricing by Esscher transforms. HEC Ecole des hautes études commerciales, 1993.
Lin Gui, Cristina Gårbacea, and Victor Veitch. Bonbon alignment for large language models and the sweetness best-of-n sampling. arXiv preprint arXiv:2406. 00832, 2024.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021.
Arian Hosscini, Xingdi Yuan, Nikolay Malkin, Aaron Courville, Alessandro Sordoni, and Rishabh Agarwal. V-star: Training verifiers for talf-t reasoners. arXiv preprint arXiv:2402.06457, 2024.
Yuu Jinnai, Tetsuro Morimura, Kaito Ariu, and Kenshi Abe. Regularized best-of-n sampling to mitigate reward hacking for language model alignment. arXiv preprint arXiv:2404.01054, 2024.
Aviral Kumar, Vincent Zhuang. Rishabh Agarwal, Yi Su, John D Co-Reyes, Avi Singh, Kate Baumli, Shariq Iqbal, Colton Bishop, Rebecca Roelofs, et al. Training language models to self-correct via reinforcement learning. arXiv preprint arXiv:24099 12917, 2024.
Harrison Lee, Samrat Phatale, Hassan Mansoor, Kellie Lu, Thomas Mesnard, Colton Bishop, Victor Carbune, and Abhinav Rastogi. Rlaif: Scaling reinforcement learning from human feedback with ai feedback. arXiv preprint arXiv:2309.000 002670 2023.
Lucas Lehnert, Sainbayar Sukhbaatar, Paul Movay, Michael Rabbat, and Yuandong Tian. Be- yond a#: Better planning with transformers via search dynamics bootstrapping. arXiv preprint arXiv:2402. 14083, 2024.
Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards. Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. 2023.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. arXiv preprint arXiv:2305. 2.t.::2955...00 18290, 2023.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gem- ini 1.5: Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv:2403.0555 055300 2024.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint Stti::55060223333 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: arXiv:1707.06:: 1707.06347, 2017.
Pier Giuseppe Sessa, Robert Dadashi, Léonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Ramé, Bobak Shariari, Sarah Perrin, Abe Friesen, Geoffrey Cideron, et al. Bond: Aligning 1lms with best-of-n distillation. arXiv preprint 2rXiv:2407.1462. 2024.
Avi Singh, John D Co-Reyes, Rishabh Agarwal, Ankesh Anand, Piyush Patil, Peter Liu, James Harrison, Jaehoon Lee, Kelvin Xu, Aaron Parisi, et al. Beyond human data: Scaling self telf-training for problem-solving with language models. arXiv preprint arXiv:23/2.06580 2023.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ulm test-time compute optimally can be more effective than scaling model parameters. arXiv preprint arXiv:2408.0331 14 2024.
12

---

Published as a conference paper at ICLR 2025
Saurabh Srivastava, Anto PV, Shashank Menon, Ajay Sukumar, Alan Philipose, Stevin Prince, Sooraj Thomas, et al. Functional benchmarks for robust evaluation of reasoning performance. and the reasoning gap. arXiv preprint arXiv:2402. 19450, 2024.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-30211 2020.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Richard S. Sutton, Hamid Reza Maei, Doina Precup, Shalabh Bhatnagar, David Silver, Csaba Szepesvári, and Eric Wiewiora. Fast gradient-descent methods for temporal-difference with linear function approximation. In International Conference on Machine Learning (ICML), 2009a.
difference learning
Richard S. Sutton, Hamid Reza Maci, and Csaba Szepesvári. A convergent onn temporal-different algorithm for off-policy learning with linear function approximation. In Advances in Neural Information Processing Systems (NeurIPS), 2009b.
Fahim Tajwar, Anikait Singh, Archit Sharma, Rafael Rafailov, Jeff Schneider, Tengyang Xie, Ste- fano Ermon, Chelsea Finn, and Aviral Kumar. Preference Fine-Tuning of LLMs Should Leverage Suboptimal, On-Policy Data, 2024.
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, Johan Fer- ret, Peter Liu, Pouya Tafti, Abe Friesen, Michelle Casbon, Sabela Ramos, Ravin Kumar, Char- line Le Lan, Sammy Jerome, Anton Tsitsulin, Nino Vieillard, Piotr Stanczyk, Sertan Girgin, Nikola Momchev, Matt Hoffman, Shantanu Thakoor, Jcan-Bastien Grill, Behnam Neyshabur, Olivier Bachem, Alanna Walton, Aliaksei Severyn, Alicia Parrish, Aliya Ahmad, Allen Hutchi- son, Alvin Abdagic, Amanda Carl. Amy Shen, Andy Brock, Andy Coenen, Anthony Laforge, Antonia Paterson, Ben Bastian, Bilal Piot, Bo Wu, Brandon Royal, Charlie Chen, Chintu Kumar, Chris Perry, Chris Welty, Christopher A. Choquette- berger, Dimple Vijaykumar, Dominika Rogozińska, Dustin Herbison, Elisa Bandy, Emma Wang, Eric Noland, Erica Moreira, Evan Senter, Evgenii Eltyshev, Francesco Visin, Gabriel Rasskin, Gary Wei, Glenn Cameron, Gus Martins, Hadi Hashemi, Hanna Alimezak-Plucifssk. Harleen Batra, Harsh Dhand, Ivan Nardini, Jacinda Mein, Jack Zhou, James Svensson, Jeff Stanway, Jetha Chan, Jin Peng Zhou, Joana Carrasqueira, Joana Iljazi, Jocelyn Becker, Joe Fernandez, Joost van Amersfoort, Josh Gordon, Josh Lipschultz. Josh Newlan, Ju yeong Ji, Kareem Mo- hamed, Kartikeya Badola, Kat Black, Katie Millican, Keelin McDonell, Kelvin Nguyen, Kiranbir Sodhia, Kish Greene, Lars Lowe Sjoesund. Lauren Usui, Laurent Sifre, Lena Heuermann, Leti- cia Lago, Lilly McNealus, Livio Baldini Soares, Logan Kilpatrick, Lucas Dixon, Luciano Mar- tins, Machel Reid, Manvinder Singh, Mark Iverson. Martin Görner, Mat Velloso, Mateo Wirth, Matt Davidow. Matt Miller, Matthew Rahtz, Matthew Watson, Meg Risdal, Mehran Kazemi, Michael Moynihan, Ming Zhang, Minsuk Kahng, Minwoo Park, Mofi Rahman, Mohit Khat- wani, Natalie Dao, Nenshad Bardoliwalla, Nesh Devanathan, Neta Dumai, Nilay Chauhan, Os- car Wahltinez, Pankil Botarda, Parker Barnes, Paul Barham, Paul Michel, Pengehong Jin, Petko Georgiev, Phil Culliton. Pradeep Kuppala, Ramona Comanescu, Ramona Merhej. Reena Jana, Reza Ardeshir Rokni, Rishabh Agarwal, Ryan Mullins, Samanch Saadat, Sara Mc Carthy, Sarah Perrin, Sébastien M. R. Arnold, Sebastian Krause, Shengyang Dai, Shruti Garg, Shruti Sheth, Sue Ronstrom, Susan Chan, Timothy Jordan, Ting Yu, Tom Eccles, Tom Hennigan, Tomas Ko- cisky, Tulsee Doshi, Vihan Jain, Vikas Yadav, Vilobh Meshram, Vishal Dharmadhikari, Warren Barkley, Wei Wei, Wenming Ye, Woohyun Han, Woosuk Kwon, Xiang Xu, Zhe Shen, Zhitao Gong, Zichuan Wei, Victor Cotruta, Phoebe Kirk, Anand Rao, Minh Giang, Ludovic Peran, Tris Warkentin, Eli Collins, Joelle Barral, Zoubin Ghahramani, Raia Hadsell, D. Sculley, Jeanine Banks, Anca Dragan, Slav Petrov, Oriol Vinyals, Jeff Dean, Demis Hassabis, Koray Kavukcuoglu, Clement Farabet, Elena Buchatskaya, Sebastian Borgeaud, Noah Fiedel, Armand Joulin, Kathleen Kenealy, Robert Dadashi, and Alek Andreev. Gemma Improving open language models at a practical size, 2024. URL https://arii.org/abs//400.8 0tbbs/2408.000188
Danila Sinopalnikov, David Wein-
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi. Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837, 2022.
13

---

Published as a conference paper at ICLR 2025
Sean Welleck, Amanda Bertsch, Matthew Finlayson, Hailey Schoelkopf. Alex Xie, Graham Neubig, Ilia Kulikov, and Zaid Harchaoui. From decoding to meta-generation: Inference-time algorithms for large language models. arXiv preprint arXiv:2406. 16838, 2024.
Peter M Williams. Bayesian conditionalisation and the principle of minimum information. The British Journal for the Philosophy of Science, 31(2):131-144, 1980.
Ronald J Williams. Simple statistical gradient algorithms for connectionist reinforcement learning, Machine learning, 8(3-4):229-25) 256, 1992.
Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. An empirical anal- ysis of compute-optimal inference for problem-solving with language models. arXiv preprint arXiv:2408.00724 2024.
Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. Advances in Neural Information Processing Systems, 33:20520-20531, 2020.
Eric Zelikman, Yuhuai Wu, Jesse Mu, and Noah Goodman. Star: Bootstrapping reasoning with reasoning. Advances in Neural Information Processing Systems, 35:15476-15444 154888 2022.
Arnold Zellner. Optimal information processing and bayes's theorem. The American Statistician, 42(4):278-280, 1988.
Lunjun Zhang, Arian Hosseini, Hritik Bansal, Mehran Kazemi, Aviral Kumar, and Rishabh Agarwal. Generative verifiers: Reward modeling as next-token prediction. arXiv preprint DrXiv:2408.15240. 2024.
Jun Zhu, Ning Chen, and Eric Xing. Bayesian inference with posterior regularization and appli- cations to infinite latent syms. The Journal of Machine Learning Research, 15(1):1799-188 2014.
14

---

Published as a conference paper at ICLR 2025
A THEORETICAL DERIVATIONS
A. VARIATIONAL APPROXIMATION OF BON
We assume that the verifier score r(ar, y) is unique for all x. y, and the base model has finite set of possible outcomes for each context (Beirami et al., 2024).
Proposition 5 (Theorem 2 in Gui et al. (2024)), With negligible error, one may effectively approxi- mate bon as the solution to the following optimization problem:
1
Thon (yla)) E arg max Eynp 1Q (2, M(la)eey
KL (UllTT) (z),
(10)
Av
where Q (x, 1) Ey By(-a((a)) (((x))))(a)))
the expected win-rate over n, and
(AN - 1) exp (AN + 1) log
ex = log N
N
(11)
exp (AN - 1)
An
N
through AN scaling sub-linearly with BoN number of samples N.
We can show the optimal solution to Equation (10) has a closed form Than X [T exp (x)Qx))((le) (yll)) This can also be revealed by viewing Equation (10) as the variational form of Bayes' rule (Williams, 1980; Zellner, 1988; Zhu al., 2014 4; Dai et al., 2016). whose optimal solution is the posterior. This implies Thon can be represented by an exponentialltwistee icy ग with energy function Ax Q2 (y, x²). partition function Z (I) Ex(()= [exp (AN Qx (z, B)], and an appropriate xN from Equation (11).
policy (Gerber et al., 1993) over base pol-
In this section we will provide proofs for the technical results in this paper.
A.2 PROOF OF LEMMA 1
Therorem 2 of Gui et al. (2024) shows that, with negligible error, one may effectively approximate Thon as the solution to the following optimization problem:
Thon ()) E arg max Eysu 1Q (z, v)1
KL (x)l))))
(12)
AN
where 2 (vc, y) = [v-x(()))) [1 (( is the expected win-rate over n, and this yields the variational form Thon x [ir exp (AN (y)). Plugging the variational form of Thon into (2) yields the learning problem for m:
max E (xx [log (boo(y)z) = B(x log v (ylz)+AN Qx (I, (og Zn (x)
(13)
Elll
Likelihood
Inference-Awareness
Taking gradient of this objective function over 0 θ implics
VolB(a,))
[log no (ulx) +AN QT, (x, y) log (x)]
V log no (y/=)] + AN VolE 12., (z, v)l 1 VOEAND olo Zmo (x)]
- V log no (vlx)) + AN E(za)) [EgrrNT (V Tolog <o(l/>>) (r(x. v) r(x, y')
ENDD
V o log (x)()))) [exp (AN QT, (I, U)]]
E (
V e log ne
AN B( [Eg'rto (V log To [z) (y/12>)
(r(z.)) - r(x y )ll
my [V log No (v)xr) exp (AN Qxo (z, v))l + En)(yia) [V exp (AN Qno (x,)))]
E
Bro(u(z) (exp (AN QTy (x,
This further implies that V (D [log To (y(z) +AN Qn (x, y)
log ZTo (x)] exp (AN Q v))
-E E [Vo (z, yı 0)] ELnD E Eno(sin)
Ero(ylz) [exp (AN Qro (x, y))l Vof (z. y: 0)
through collecting terms from the above expression and recalling the definition of V (a, y; 0) as Vol (a, y; 0): : Va log To (w/x) + AN Byy [Vo log no (2/lx)) (r(x, v) r(x, s/))]
This further implies that
V ( [log To + Qxo (x, y) - log Zno (x) ( [Vol (ar, (;:0) EDDDD Boo.DDyymmmm.. Tow. [V (x, y: 0)],
completing the proof of this lemma.
15

---

Published as a conference paper at ICLR 2025
A.3 PROOF OF LEMMA 2
Recall the RL objective function
max [ := ((
v)l.
(14)
TElf
Applying the REINFORCE trick (Sutton et al., 2009b) to this problem over the BoN policy class
and using the analgous argument from the proof of Lemma 1, we have the following expression for the policy gradient:
E Ey (l))
Vo log Moon (y)2)) R(a, y)]
-E m (aar))
V fef (a,, y; 0) R(x, y)] - B.n..y Wof (I. y: Ey a)aaD [R(x, y)
-EWD B.wD.ymmen yomm (ll)
V Vaja (zr, y) (R(x, y) - b(x))]
= -REND, -DawD.yytae syotttee (1l))
V log (wlx)) (R(æ, v) - b(x))]
+ AN D.y.
Vo log no (1/12)-11((z.g)) r(x,//))))))))))) (R(r,y) b(x))]
[Va log me ()) (R(z, y) - b(x))
D.y.m (x)
Ax Bn (1)z))))) (m Va log na (1/1z) <r(x)))))))((((()))))) (R(r, - b(x))
This completes the proof of this lemma.
A.4 PROOF OF LEMMA 3
Using the log-likelihood trick, and plugging in the BoN distribution from Equation (8), the gradient of Equation (6) can be computed as
E By-amu
[Va log Tbon R(x, y)] mon (1)) Ve log Toon [z]]
1 By / (+ ((2,)/== N
-E
V (og B (10))
+ Ey ( R( Va log 1 Eyywaal (-1x) R(2,80 )=)
1 (x)
1 IR(aN'
-EwDD
(yymma((lo) V log no (y/=²]
By'na Vo log 1[(2,))==1)
Ey (x(+)) [(()))))=))
+ Ey'nr Volog -= N By)nx(((ar) [1 [[(ar)))=
(yTTma((1)) V log To (ul=)) Nhe((2) N (1 Inor(x)))
1 - [e)(x))
N [ze(
EBnD
Vo log ()()))) (xr)
- ((-1) Byy
1 Iner (x)
A.5 PROOF OF COROLLARY 4
Using the log-likelihood trick, and plugging in the BoN distribution from Equation (8), the gradient of problem Equation (6) can be computed as
N N=
Ez-D
B ( (11).
Vo [y))))
1 [
N Inet
+
E
log
[r)]
Lot (x)
(x)
N4 - (oe((()))
= =RyaD (1))
Vo Volg (y)
1 - (z)
16

---

Published as a conference paper at ICLR 2025
Algorithm BoN-SFT
1: Input: Verifier r, environment reward R, expert dataset D
2: for t- 1,2,.. ..2....do do
3: Sample a batch of prompts and solutions (z,, yi from the expert data D.
4: for = ,,...,, B do
5: Sample responses [bt.))==1 from (=(-))))
6: Select the BoN response y." arg max, r(z,, Mt.j).
7: Compute the gradient Vefo(z; Vs) using Equation (5).
8: end for
9:
Update 0 by following the gradient in Theorem at learning rate a > 0, i.e,
N
0-0+a
VoJ (zi, Us: 0)) - [Vol (xi, vi: :
N
10: end for
B
PSEUDO-CODE AND IMPLEMENTATION DETAILS
Pseudo-code for all our SFT and RL methods is presented in Algorithms to 4. Our implementation follows the standard use of an anchor policy, updated using exponential moving average. The policy trained via BoN-aware losses, with additional KL divergence loss to the anchor policy. Table shows the hyper-parameters used for all of our experiments.
We use linear annealing for the KL-coefficient. For all our RL experiments, we use value baselines to reduce variance of our reward estimates. We normalize our advantage estimates w.r.t. the batch. For BoN-RLB the value network estimates Pran (1)). We add additional clipping of the coefficients . 9N by clipping the value estimates for Pran-
Table 2: Hyperparameters used in experiments.
Hyperparameter
Value
Base model
Gemma 2b v2
Optimizer
AdamW
Learning rate policy
3e-6
Policy warmup steps
100
Learning rate value
le-5
Anchor EMA
0.01
Training steps
2500
Batch size
32
Sampling temperature
1.0
KL coefficient anneal steps
2500
KI. coefficient anneal range
1.0 - 0.075
KI. coefficient anneal delay
10
Clipping values for Prail
(0.01,0.99)
B. ANALYSIS OF BON-RLB WEIGHTS
The shifting balance between (p) and g (p) with varying p directly reflects the exploration- exploitation trade-off. As p approaches 1, signifying very difficult problems, both g (p) and g (p) increase, but (p) rises more dramatically, especially for larger values of N. This sharp increase in (p) highlights the algorithm's 's increasing emphasis on learning from the few correct responses that are available in challenging scenarios. The effect is amplified by larger sample sizes: the more attempts are made, the more valuable the scarce successes become. The g (p) weight, used when only positive feedback is available, exhibits similar upward trend with p but with less pronounced increase. This more moderate behavior can be attributed to the subtraction term in its formula, which tempers the influence of the positive samples and promotes a more balanced learning approach.
Finally, the potentially very large values of x (p) for hard problems and larger N introduce chal- lenges for estimation. These high weights amplify the impact of individual positive samples, making
17

---

Published as a conference paper at ICLR 2025
Algorithm 2 BoN-RLB(P)
1: Input: Environment reward R, dataset D
2: for 1, 1...... do
3: Sample batch of prompts { from D.
4: for = ,,..., B do
5: Sample responses [ from TT(-1z;))
6: Sample rewards for all candidate responses { (R(2, y4) y:) from environment.
7:
Select the BoN response y = arg max, R(2,, Vi.j).
8:
Empirically estimate the base failure probability for each Ty, i E {l,..., B),
NV.
Pain = N
Σ
1 R(#ann))==-
9:
end for
10:
Update by following the gradient in Corollary 4 at learning rate a >> 0, i.e.
B
0 0
Σ
Vo (3 (x.) g" (P)
vo
B
1=1
where y/, + represents the BoN sample that achieves a reward of 1. 11: end for
Algorithm 3 BoN-RLB
Input: Environment reward R, dataset !
2: for t = 1, 2, do
3: Sample batch of prompts {= from D.
4: for = 1,..., B do
,...,, Sample N responses {yr. 1 1-= from To(-)x;).
5:
6:
Sample rewards for all candidate responses (:)) from environment.
7:
Select the BoN response y = arg max, R( Vi.j). Empirically estimate the base failure probability for each IG, i E (I,..., B).
8:
(I,) ==
1 m=zaan.))
N
9:
end for
10:
Update 0 by following the gradient in Lemma 3 at learning rate > 0, i.e.,
0 - 0+a
a log #o( + lanz) g (Paan) Vo log no(y" (as)) of (Pan(
where y
represent the BoN samples that achieve rewards of and respectively.
end for
the training more vulnerable to noise and potentially hindering convergence to a stable optimal pol- icy. This underscores the need for techniques like gradient clipping or regularization to mitigate the destabilizing effects of high weight values and ensure robust learning, or alternatively. using T (p) with Corollary 4 to ensure boundness of the gradient weights.
18

---

Published as a conference paper at ICLR 2025
Algorithm 4 BoN-RL
1: Input: Verifier
r, environment reward R, dataset 'D
2: for f 1,22, 1,2.... .. do
3: Sample a batch of prompts { from D.
4: for = ,,..., B do
5:
Sample responses [ from TT(-1x;)-
6:
Select the BoN response y. arg max, ri.i,
Mi.j).
7:
(If environment reward It is available to the BoN algorithm, we replace verifier with
that.)
8:
Sample the reward R(z,, vi from environment.
9:
Compute the gradient Vefo (Zi, y/) using Equation (5).
10:
end for
11:
Update by following the gradient in Lemma at learning rate a >> i.e.,
-- 0 +
B VoJo (Ii v() (R(as, u.) - b(xx))
where by (xi), B is a learned baseline value function of Tbbun i.e.,
/ E arg min LIR(E, 2/) (xx))
Update value estimate & using the current environment reward target and BoN policy tra- jectories.
13: end for
C ALGORITHMIC EXTENSIONS
C.I ENTROPY-REGULLARIZE RL
We would like to study an entropy-regularizer RL problem for the Thon policy. Recall that generally
in antropy-regllarize RL, we solve
max EEnD (1y [R(x, v)] B (T(()|(>))())))
(15)
T(-1Z)E
where R(x, y) is the environment reward (that is not necessarily identical to the verifier score model), Ty is baseline policy, and B > 0 is the weight for the KI regularization term. Using the consistency condition Ki.-regularized MDP, solving for the optimal policy of this problem equivalent to finding solution pair of V and π E A of the following equation:
(16) Now, we further parameterize the policy variable with the BoN policy bomn then with the suffi ciency part of the consistency condition one can show that Toca is an optimal RL policy of Equa- tion (15) there exists pair of V and π that satisfies the following equation
((x) = R(x,8) y) + 3 log nH (T at) log π ( x), Vr E D. Vy
V(x) R(x, v) + 3 log Ta ()((ar))
3 (log π (yl=r + AN Qx (u.z) x) log (x)), Vo E D, Vy (17)
There are two ways to approximately find the solution in Equation (17). The first way is to refor- mulate the above equation with condition that equates the values between any pairwise states and outputs (x, V. v')::
Ts(u|o) R(æ, v/ +3 log To(y'>)) BAN (V'xx) = R(x, x)+ log T (V)l)) + 3AN (y,4), Vo E DD Vy. <({(=)
(18)
Suppose one have access to pairwise labels in the data-set, then this formulation eliminates any terms that are independent to y and circumvents the need of solving for the value function V. One may approximately solve Equation (18) by minimizing the following & loss:
minE, (.D)))) [[9(x, y: T) (g(2, +)) FEA
T T
((. : R(x, y) log क y x))
+ BAN Qu (y, x)
This formulation is similar to that in IPO (Azar et al., 2024). However, unlike IPO, where the term g(,, y; ) is linear in the logits of π and therefore one can show that its Ea minimization problem has unique solution, in this case g(,, y: also depends on Q.. which is function of (and
19

---

Published as a conference paper at ICLR 2025
thus nonlinear function of its logits), preventing us from drawing similar conclusions that the l2 minimization problem has unique solution. Therefore, even if one can exactly solve this & minimization problem (and make the loss zero), there is no guarantee that the solution policy x corresponds to the base policy of an optimal Thon policy to the KL-regularized RL problem.
For the second approach, consider the following linear programming reformulation of Equation (17): min [x)D[V(<r]
VAEA
s.t.
V(x) > R(z.y) (y)l)())>)) - 3 (log क (y/a) AN Qx (y. log Z (x)) Vo E D, Vy
(19)
Since the inequality constraint is convex function in and an affine function in V. by strong duality it has the following equivalent Lagrangian-dus dual formulation:
R(xx, v) + _ 3 (Ax Q (v: x)) log 7x(x) - V(x))
max min (awweeD (x)) + m(x,y)
<(-))20 VISEA
max min IDD [(1 x(z,v)) ((x) + s(a,y) (R(x, + 3 []((]=))] max ED [<r( y) llog Than (2/ (a;; xT)
x(:120 V
REA
(20) This formulation can be viewed as an weighted-SFF approach that iteratively updates (i) the base policy that maximizes the likelihood of Thea over data D, weighted with importance weights k(x, y), and (ii) the importance weight function itself. Here, the value function (x) is simply an auxiliary variable.
C.2 IMPROVED EFFICIENCY WITH BON DISTILLATION
While Lemma provides a recipe for training base policy to adapt to the BoN inference strategy, a key challenge lies in the computational cost and data inefficiency associated with BoN sampling, especially when N is large. Particularly, each gradient update requires generating N samples from the current base policy, which can be prohibitively expensive. Furthermore, using these samples solely for single gradient update may deem wasteful.
To alleviate this issue, leveraging the recent advances in BoN Distillation (BoND) (Sessa et al., 2024), an RLHF algorithm that distills BoN behaviors into a standard LLM, we approximate the BoN distribution of the current π. This results in an iterative, two-step procedure. First, estimate a BoND policy TROND (parameterized by weights Φ) of π by solving the distribution- metrii
matching problem: minş Ez-D [L (oo) (z)]. where the backward- Wwardd-LLaa metric induces quantile- based advantage and mode-seeking behaviors TBoND- Utilizing the variational form Toon (y) (z) xx
exp (AvQ=) (yl=r), this problem can be further reformulated as
(oaN((/(>) E arg max à Second, equipped with the BoND policy, we change the gradient of Lemma with the approximate gradient [a.,y)~ e (z. v. Er (=))m Vf (z, 0)]. In general, this approach is also well-connected with Contrastive Divergence (Carreira-Perpie Perpinan & Hinton, 2005) in energy- based learning, which promotes the idea of approximately sample from the current target distribution (Toon in case). It shows that the learning algorithm can still converge to optimum w.r.t. the original objective function as long as the gradient estimated by the approximate samples still points at an ascending direction.
EEED E ( (1, AN
(L(TollT)(=))
(21)
C.3 CONNECTION TO STAR
Consider the popular STaR method (Zelikman et al., 2022) applied for training Toon, which updates á by following the weward-weghhedd gradient:
Egg. D. (-lm
[Vo log (v(lx) R(a, v)]
(22)
Notice that the policy gradient of BoN-RL is a sum of two terms: VoJ(0) = 91 (0) + g0 (0). where 9i (0) EEpDD .yw..e-- Ve log To (v/lz) R(x, y)] is equivalent to that of BoN-STaR, updating via weighted supervised tine-tuning over the responses and the rewards obtained by the cur- rent BoN policy, and (f) Em E.-D)yr=ma(((1)) (Ve(Av Qn - lxgE,exp(Ax (z) R(x, accounts for the gradient effect of the importance sampling term (exp AN (r. y) between π and Thon emphasizing how much it can improve the reward. The additional g0) component
20

---

Published as a conference paper at ICLR 2025
04
180
03
10
Data Prodiction
sprinal Weaic3an
02
0.
00
109
10
(a) pass N
(b) BoN
Figure 9: Scaling of exponent w.r.t temperature in pass N and optimal N w.r. temperature in BoN. Dashed eurves denote in-training predictions, stars denote extrapolation values for the corre- sponding temperatures.
pass@
BoN
Majority Voting
Gemma-98 9B
0.986
0.989
0.89
Gemma-2B 0.998
0.998
0.784
Table R-squared values for different language models and inference algorithms.
makes BoN-RL amenable to the distributional shifts introduced by the BoN procedure, enabling the base policy to be adept at utilizing the BoN exploration mechanism optimize the reward.
D EXPERIMENTAL DETAILS
D.1 ADDITIONAL SCALING RESULTS WITH GEMMA 9B VERIFIER AND POLICY MODELS
Similar to Gemma2B co-scaling experiments, for Gemma 9B co-scaling, we present additional re- sults in Figure We analyze the optimal exponent b" (T) w.r.t different temperatures (see co-scaling in Section 5) and find that power law functional form can explain the relationship very accurately, achieving very low extrapolation error for pass @ and BoN, i.c., 2.75e 05 and 2.87, respectively. Results for pass@) suggest that exponent can be accurately predicted from just temperature. We also inspect how optimal N scales with T in BoN. We fit power law function plus linear term which accurately predicts optimal N for unseen temperatures. Predictions of the fitted model can be used to achieve close to optimal performance, achieving less than 0.001 point drop in BoN perfor- mance. This suggests that our predictive model makes accurate predictions that keeps the optimal performance.
Generalization of scaling predictions. In Table 3, we compare various inference algorithms and LLMs of different sizes. For Majority algorithm, use MC estimation simulate different sample sizes. We use the same functional form used for co-scaling experiments in Section for Votini Our results remark the strong generalization performance of our predictive models across different L.MSs (both policy and reward models) and inference algorithms. The same power- law function with a linear trend term applies well to both BoN and Majority Voting. While we use unbiased estimates for pass @ N and BoN, we use MC estimation for Majority MajorityVoting which leads to less smooth curves and slightly lower prediction performance.
Gemma-9B Results. Figure 10, we present results for Gemma-9B policy and reward models. Using Gemma-9B improves both pass@1 N and BoN significantly compared to Gemma-2B. We ob- serve that the gap between using large temperatures (0.7 or 1.0) and very small temperatures (0.1) also increased. While Gemma-2B showed very strong reward model overoptimization for larger N and temperatures,
lesser overoptimization for Gemma-9B models.
21

---

Published as a conference paper at ICLR 2025
Temperature
0.9
0.01
Pass.gN Accuracy
BoN Accuracy
1,0
0.7
1.5
0.3
0.5
10
10
10a
16"
10:
10
N
N
Figure 10: pass N (left) and BoN (right) performance for Gemma-9B. While curves show similar shape Gemma-2B models, overall performance is globally improved and overoptimization is reduced.
0.80
0.56
0.75
FSsSNN Accuracy
BoN Accurac)
0.54
0.52
0.65
0.50
0.60
0.48
0.55
0.46
Base-model
AL. N-1
N-1
BoN-RLV N-8
BON-RLV
RoN-RLS
0.42
BON-RLS N-B
BoN-RLBP N-8
5
10
15
20
25
30
10
15
20
25
30
N
N
(a) BoN
(b) pass N
Figure 11: BoN and pass @ Accuracy Results on Hendrycks MATH with Gemma 9B.
D.2 DETAILS BON- FINE-TUNING EXPERIMENTS
For the MATH benchmark, we trained the Gemma 2B and 9B models with the Hendrycks MATH dataset. Following Lightman et al. (2023), we augment the original 7500 MATH training problems with 4500 problems from the test set, evaluating performance on the remaining 500 problems. In the supervised setting, we leverage larger Gemini 1.5 Flash model (Reid et al., 2024) generate MATH solutions with answers and steps (32 candidates for each the MATH problems), sub- sampling only the correct responses and distilling knowledge into the Gemma 2B model. In the RL setting, we use a binary environment reward denoting whether the model's answer matches the ground truth answer. The verifier used in all BoN experiments is separate pre- pre-trained Gemma 2B model that predicts the probability of correct response given the prompt. The verifier trained with the data collected from the Gemini 1.5 Flash model.
Alternatively, to benchmark and evaluate on the HumanEval benchmark (Chen et al., 2021), following the standard procedures delineated in Kumar et al. (2024).
models code generation, train MBPP (Austin et al., 2021)
D.3 ADDITIONAL BON-AWARE FINE-TUNING RESULTS
We now present additional results on BoN-aware fine-tuning (both SFT and RL).
D.3.1 HENDRYCKS MATH WITH GEMMA 9B
We additionally benchmark a larger model, GemmaV2 V2 9B, on Hendryeks MATH, with results shown Figure 11. We observe that, similar the trends of the experiments run with the Gemma 2B coun- terpart, BoN-RLV achieves the best BoN performance, while BoN-RL-S achieves the best pass @ N performance, with both substantially improving the base model.
22

---

Published as a conference paper at ICLR 2025
0.95
0.75
0.90
0250
Pass@NN Accuracy
Accuracy
0.225
0.80
0.175
0.75
0.150
Base-model
RL. Nm1
0:25
RL Nm1
DoN-RLV RLV N
DoN-RLV
0.65
BoN-RLS Nm
0.100
BoN-RLS NeB
BON-RLBP BP N-8
0.075
DoN RLBP N
10
15
20
25
30
5
15
20
30
N
(a) Functional MATH.
(b) MathOdyssey.
Figure 12: pass@N of BoN-RL and baselines on held-out benchmarks with Gemma 9B models.
0.6
PassigNN Accuracy
0.5
0.4
0.3
Bsse model
BoN-RL-S Nw16 RLB #1
M/=16
Figure 13: pass @ N of various methods with N' = 16
D.3.2 HELD-OUT MATI BENCIMARKS WITH GEMMA 9B
To evaluate the generalization capabilities of our BoN-aware finetuned models, we additionally eval- uate on two completely held-out and challenging benchmarks, Functional Math (Srivastava et al., 2024) and MathOdyssey (Fang et al., 2024). We present the results of these held-out benchmarks with 9B models in Figure 12, and observe that our fine-tuned models improve on both BoN and pass @ N for these held-out benchmarks similarly as in the case of Gemma 2B.
D.3.3 PASS N ACCURACY WITH DIFFERENT TRAINING SAMPLES, WITH GEMMA 2B
Similar to the experiments in Figure 6a, our BoN-RL-S, BON-RLB, and BoN-RLB(P) models also demonstrate superior performance over the baseline methods in pass @ N evaluations with N' = 16 (Figure 13). In this case, STaR performs the worst (55% in pass @ 32) as it fails to (i) utilize negative samples in training for implicit exploration (unlike BoN-RLB, 60% pass @2). (ii) re-weight sam- ples based difficulty, prioritizing learning from challenging problems and avoiding overfitting simpler ones (unlike BoN-RLB(P). 60% pass @ 32), and (iii) account for the importance sampling factor between the base policy and the BoN policy (unlike BoN-RL-S, 58% pass @ 32). BoN-RLB and BoN-RLB(P) are superior to BoN-RL-S on N' 16 (Figure 13), but worse on N' ure 6a), suggesting that they suffer from instability with increasing N'. This is potentially due to the following observations: (i) The asymmetry between the positive (g+) and negative (g weights in BoN-RLB increases with N', destabilizing its learning at larger N' values; (ii) RL-BoN-S uti- lizes the variational approximation of Thon in its gradient update, introducing approximation errors that may cause its sub-optimal performance (relative to a stable instance of BoN-RLB trained at N' 16); BoN-RLB(P only uses positive samples, which inherits the shortcomings of STaR (lack of implicit exploration), yet re-balances the examples with the difficulty the problems. Overall, it leads to consistent yet mild performance degradation over BoN-RLB.
(Fig-
D.3.4 COMPARING BON-RL WITH BASE-BON DISTILLATION BASELINES, WITH GEMMA 2B
We consider various alternative methods to improve Gemma 2B BoN accuracy through various data generation methods. We distill the Gemma 2B model using these datasets and compare the
23

---

Published as a conference paper at ICLR 2025
BoN Accuracy
0.25
0.20
BoN-RL-V
0.15
Base-BoN-Shi
Base-All-SFT
Base Base-WWighhtdd.o SFT
Base-Maj-SFT
N
Figure 14: BoN Accuracy on MATH comparing Base Gemma 2B and BoN RL-V with other fine- tuning techniques: (I) BoN-SFT: Distillation of BoN sample for N=16; (2) All-SFT: Distillation of all N=16 samples (i.e., average sample); (3) Weighted-SFT: Distillation of all N=16 samples by average re-sampling w.r.t. verifier scores; and (4) Maj-SFT: Distillation of majority voting strategy.
0.30
BoN Accuracy
0.25
0.20
Base-model
0.15
BoN-RL-V
BoN-RL-S
0.10
BoN-RLBP
BoN-RLB
5
10
20
25
30
N
Figure 15: Verifier Mismatch. Plots show BoN accuracy under verifier-reward mismatch using Gemma 2B MATH. During training verifier was used for BoN. On test, environment reward was used as verifier of the BoN strategy, inducing a mismatch in verifiers.
base Gemma 2B model and BoN-RL-V method. We consider the following four distillation benchmarks (all run over Hendrycks MATH):
Base-BoN-SFTT In this method we generate a dataset of the best of 16 samples for each example in the dataset. We use the best sample as target to distill Gemma 2B.
2. Base-All-SFT: We use the full range N samples as targets. This dataset is used to distill Gemma 2B to the average effective sample of the base model.
3. Base- Weighted-SWi SFT: Similar to Pase-All-SFT we sample N = 16 samples for each exam- ple. We then re-sample N 16 examples (from these samples, with repetition), weighted according to verifier scores. This dataset used to distill Gemma 2B the average effec- tive sample, weighted by verifier scores.
4. Base-Maj Base-Maj-SFFT SFT: We use majority voting over N 16 samples to select a target. We distill Gemma 2B to predict the majority voted target.
We show the BoN accuracy results of these methods in Figure 10. While the aforementioned base- lines do improve BoN performance the Base Gemma 2B model, they are still out out-performed by our BoN-RL-V method, indicating the value of utilizing the inference BoN strategy explicitly during training.
24


Reviews:
[{"review_id": "EdZQ2qf85O", "reviewer": "Reviewer_4VnU", "summary": "The paper proposes a novel inference-aware fine-tuning paradigm, aiming to enhance the inference performance when scaling the compute. \n\nTo overcome the non-differential argmax operator within best-of-N (BoN), the paper proposes the first imitation learning and reinforcement learning methods for fine-tuning language models with BoN. \n\nThe experimental results show the BoN improvement of Gemma 2B on Hendrycks MATH from 26.8% to 30.8%, and Pass@K from 60% to 67%.", "questions": "1. How is the accuracy of the used reward model? \n\n2. How does best-of-N sampling relate to majority vote in math reasoning?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The studied problem is interesting and important: how to improve best-of-N sampling in the test time through advanced training methods applied to fine-tune the models. \n\n2. The proposed method is reasonable and could be sound and useful if the method can be verified carefully and thoroughly.", "weaknesses": "- The paper writing can be improved: it is not easy to follow the core idea of the proposed methods. A visual Figure showing the core mechanism of the method is strongly suggested. \n\n- The experiments are not sound enough to demonstrate the effectiveness of the proposed method, specifically:\n\n**Regarding the co-scaling behavior of sample number N and temperature**: this analysis was conducted using only a single policy model and a single reward model.\n\nHow might the curves in Figure 2 appear if different language models or stronger reward models were used?\n\nThe conclusions drawn here cannot be readily generalized to other language models or reward models.\n\n\n**The results presented in Figure 4 make it challenging to conclude that the proposed inference-aware fine-tuning significantly improves performance.**\n\nOne potential issue lies in the weakness of the reward model: the solution selected as having the highest reward is often inaccurate, and fine-tuning solely on this data may result in suboptimal model performance.\n\nSelecting a broader range of solutions could benefit fine-tuning, as a correct answer might be found within these additional options.\n\nTherefore, the observed improvement may stem from identifying a correct solution through an increased sample selection rather than from the inference-aware fine-tuning itself.\n\n\n- I would also suggest conducting experiments on the alignment tasks and testing the models on AlpacaEval or Arena-Hard test sets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YHTrGpQ5Ow", "reviewer": "Reviewer_VBde", "summary": "This paper presents a finetuning approach for large language models (LLMs) that incorporates the inference strategy directly into the model’s training process, with a particular focus on Best-of-N (BoN) inference. The authors argue that there is a disparity between the best-of-N policy typically used at inference time and the policy learned through standard supervised fine-tuning (SFT). To address this, they formalize an inference-aware fine-tuning problem that integrates the inference strategy into the training objective. Specifically, they develop a supervised BoN-aware fine-tuning method and a BoN-aware reinforcement learning (RL) with binary rewards, using policy gradients derived only from positive examples to mitigate data inefficiency. The proposed methods are evaluated on the Gemma 2B model, primarily using the Hendrycks MATH benchmark, with an emphasis on understanding the relationship between the number of samples and temperature.", "questions": "1. Will inference-aware SFT and inference-aware RL reduce the model’s generalizability? How does the model perform when using beam search alone, without a verifier?\n2. For inference-aware SFT experiments, what type of verifier is used? How does the method perform if there’s a mismatch between the training and testing verifiers?\n3. In Figure 1, how is the empirical frequency calculated? What defines the best BoN performance? Also, what are the definitions of \"easy\" and \"difficult\" problems, given that the figure is based on the same set of MATH benchmark problems?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper addresses an important and timely problem, as best-of-N inference with verifiers has become increasingly popular for reasoning tasks.\n- The paper provides a solid theoretical framework supporting the proposed methods with a clear formalization of inference-aware training objectives.", "weaknesses": "- The method is only evaluated on a single model and a single task, which limits insights into its broader applicability.\n- Although BoN-aware fine-tuning is a key contribution, the experiments predominantly examine the relationship between the number of samples and temperature. This leans more toward an analysis of the proposed methods than a comparative assessment against strong baselines, leaving questions about its competitive advantages unanswered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QQZEV5Wnf8", "reviewer": "Reviewer_gx5H", "summary": "This paper proposes an inference aware fine-tuning strategy for Best-of-N (BoN) to overcome the non-differentiable argmax operator for BoN. The proposed strategy incentivizes the model to balance the best and diverse responses for fine-tuning, achieves better performance compared to the base Gemma 2B model on the MATH dataset.", "questions": "See \"Weaknesses\".\n\nI would be happy to discuss further with the authors and reassess my score based on the rebuttal stage.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The incorporating of inference strategy into training is a compelling and well-founded idea.\n2. The proposed BoN-aware fine-tuning to balance the trade-off between the best and diverse responses is well-motivated.", "weaknesses": "1. The experiments are conducted on a base model (Gemma 2B) and an evaluation dataset (MATH). More evaluation results on different base models and evaluation datasets are needed.\n2. What is the underlying rationale for the BoN-SFT loss function presented in Equation (7)? The function appears to penalize y', even in scenarios where the sampled y' demonstrates high quality, as evidenced by r(x, y') approaching r(x, y). This raises a question: If the sampled y' exhibits desirable characteristics, wouldn't it be more beneficial for the policy model to learn from these high-quality responses?\n3. It would be valuable to conduct a comparative analysis of BoN-SFT against several straightforward baseline approaches. These could include:\n\n(1) Fine-tuning the base model using the highest-quality response selected from N generated samples.\n\n(2) Fine-tuning the base model utilizing the ground truth responses.\n\n(3) Fine-tuning the base model with all the sampled N responses.\n\n(4) Fine-tuning the base model with the combination of all the sampled N responses and the ground truth responses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)