Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Authors: Yunhao Tang
Kunhao Zheng
Gabriel Synnaeve
Remi Munos
Venue: ICML 2025 poster (2025)
Keywords: RL
inference time compute
pass at k
code generation
Primary area: reinforcement_learning->deep_rl
TL;DR: Inference time objectives can be directly optimized via RL

Abstract: In this work, we investigate the merits of explicitly optimizing for inference time algorithmic performance during model training. We show how optimizing for inference time performance can improve overall model efficacy. We consider generic inference time objectives with $k$ samples, with focus on pass@$k$ and majority voting as two main applications. With language model training on reasoning datasets, we showcase the performance trade-off enabled by training with such objectives. When training on code generation tasks, we show that the approach significantly improves pass@$k$ objectives compared to the baseline method.

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Yunhao Tang" Kunhao Zheng?*
Gabriel Synnaeve? 2
Rémi Munos
Abstract
Usually, model training does not explicitly account for the downstream inference time algorithm. It is clear that there is a trade-off: this helps avoid premature specialization, but might also risk not benefiting from the inference time algorithm fully had we known the algorithm in advance.
In this work, we investigate the merits of explic- itly optimizing for inference time algorithmic per- formance during model training. We show how optimizing for inference time performance can im- prove overall model efficacy. We consider generic inference time objectives with i samples, with a focus on pass @k and majority voting as two main applications. With language model training on reasoning datasets, we showcase the performance trade-off enabled by training with such objectives. When training on code generation tasks, we show that the approach significantly improves pass @ k objectives compared the baseline method.
In this work, we investigate the impact of training explicitly for the inference time algorithm. We focus on pass @ k and majority voting, two influential yet relatively simple infer- ence time objectives. The pass k objective arises when the task has verifiers and the system can retry it & times. Major- ity voting is even more broadly applicable: devoting more inference time compute to have different solutions and se- lecting the most voted one. We demonstrate how such objec- tives can be optimized via stochastic gradient descent, and built as part of an online reinforcement learning algorithm (Section 3). We carry out extensive ablations that showcase the trade-off of different objectives, such as an improved inference time performance when the training algorithm is aware of the inference time algorithm (Section 5): we show that when training on mathematical reasoning datasets such as MATH, well challenging code generation datasets such CodeContests, new algorithmic variants achieve significant gains on inference time objectives of interest.
1. Introduction
Traditionally, the performance of machine learning system can be improved via either training time or inference time algorithm. At training time, knowledge is distilled into model weights via gradient descent (Bottou, 2010). At inference time, a fixed model is queried multiple times to deliver a better prediction than single model call.
Though model training and inference time compute im- proves model performance, their scaling properties differ significantly. For example, Lerer et al. (2020) has demon- strated that for the game of Hex, scaling inference time compute can match performance at increased training bud- get, while at a fraction of the cost. More generally, inference compute has proved effective at various machine learning applications, even if the underlying model has been trained extensively. Notable examples include board game (Silver al., 2014: 2018), competitive pokers (Brown & Sandholm, 2018; 2019), general game (Lerer al., 2020; Bakhtin et al., 2022), competitive programming (Li et al., 2022) and lan- guage modeling (Jaech et al., 2024; Guo et al., 2025).
2. Reinforcement learning for language model
A language model can be understood as policy xo in the context of reinforcement learning (RL) (Sutton & Barto, 1998). Given a prompt æ, the policy generates response y. which then gets assessed by human user. Usually, the objective is to optimize such that certain reward function r (x. y) that captures human preference (HF) is maximized (Christiano et al., 2017: Ziegler et al., 2019; Ouyang et al., 2022: Bai et al., 2022). Formally, consider the average reward optimization problem
max Ba (r(x,)))
Equal contribution 'Meta GenAl Meta FAIR. Correspon- dence to: Yunhao Tang <yt2541 @columbia.cdu> Kunhao Zheng <kunhao@ mota.com?>.
where p is a distribution over prompt set. Below, when the context is clear we will omit the dependency on the prompt æ. Given samples y drawn from kg. we can construct stochastic gradient estimates to the policy gradient (Sutton et al., 2000) to iteratively improve the policy. Note that for
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
simplicity, we have omitted the regularization prevalent in the RLHF formulation (Ouyang al., 2022).
3.1. Unbiased stochastic gradient estimate
To optimize for the objective in an unbiased way with stochastic gradient descent, we can sample h generations i..d. (yr) (=() and construct the gradient estimate akin to REINFORCE (Williams, 1992)
Inference time objectives. At inference time when the model is deployed, one might adopt a procedure different from training time to obtain better performance. Depending on the application, different inference time objectives are desired. One important class of inference time objective is pass @ k; where the model is given budget of k gener- ations at test time. Though quite lenient metric, pass @ is especially useful for difficult problems (e.g., (Li et al., 2022)) that allow for trying multiple times. It also approx- imates best-of-k assuming access to a good reward model. Formally,
I (x. y1.. Vo log (V(())))
Importantly, here we sum over the weighted gradient of log probabilities across io samples. This is in contrast to a k-sample policy gradient estimate that averuges over samples. The key difference is that these h samples are coupled through the aggregation function f. As a result, such gradient estimate has high variance on the order of (kk) (Fallah et al., 2020a; Tang, 2022) that increases with the number of samples k as opposed to O(k-!) in k-sample average policy gradient estimate.
E ( [max
rx))],
where r; = r(2, wi) denotes the reward for generation Vi-
Another important example of inference time technique is majority voting. The model is still given budget of k generations and outputs a single solution for assessment. In general, majority voting is applicable if the generation (c, a) can be decomposed into a chain-of-thought and final answer a. The reward scores only the final answer r(,, y) = r(z, a). i.e., for problems with short and verifi- able answers. has proved highly effective at leveraging inference time compute for improved performance in var- ious domains (Wang et al., 2022: Lightman et al., 2023). Concretely, the majority voting objective is
Henceforth for notational simplicity, we let f(y) be a short hand notation for (ar y1-- and let f(y be a short hand notation for leaving out generation i: f(y = fzz 1/.. B1+1.. We also drop the dependency on x when the context is clear.
For variance reduction, we propose the leave-one-out con- trol variate that results in the following gradient estimate
((a( (maj(a1..... a;)))),
(f) - (y- Vo log (yrlz).
(2)
where the operation maj(-) extracts the majority element from set of items. We assume random tie breaks in case more than one element takes the majority.
The effective advantage function A, = f() f (y measures the contribution that has on improving from f (y-:) to (y).
We will focus these two objectives as they self- contained given a language model. There are alternative inference time methods such as best-of-k with access to auxiliary models (c.g., reward models) (Uesato et al., 2022; Lightman et al., 2023) which we do not discuss.
Lemma 1. (Unbiased leave-one-out gradient estimate) The gradient estimate with the leave-one-out control variate in Eqn (2) is unbiased.
3. Optimizing Language Models for Inference Time Objectives using RL
Proof. For any i, [ Va log (3) [x)] since all k samples are independent, hence the proof is concluded.
These inference time objectives share a key feature: they all use i samples. We propose a general formulation of k-sample objectives that be optimized by stochastic gradient descent.
Let us examine a few special cases to make concrete the concept of leave-one-out gradient estimate.
In general, we consider the k-sample objective
E ((o)tt_-))((l))) [f (I, ())),
(1)
defined through a function f : X' x y x y... + R. Here f is an aggregation function that can process an arbitrary number of generations. It is clear that both @ @ and majority voting objectives are special cases.
Average reward: f(y) = Σ Ti. In this case, the form of the policy gradient estimate is equivalent to the leave-one-out control variate proposed in a number of prior work that tackles k-sample average objectives (Kool et al.,
2

---

Let characterize the bias in few special cases. For the average reward case f(y) _/ >x=1 rs there is no bias. This is compatible with the fact that the advantages A, are Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
2019; Mnih al., 2016).
Intuitively, these observations imply that the update has high variance. However in a sense, the leave-one-out control variate (Eqn (2)) marks near optimal trade-off between simplicity and variance reduction for an unbiased gradient estimate. To reduce variance further, we can introduce bias: by subtracting the mean of all advantages, this produces a general gradient estimate
3
1
Σ
Tj
Vo log (v(l))))
k
i=1
j=1
jti
k
Pi
1
Σ
rj
Vo (
h
k
(x) f(y-)) A) Vo log (y(lz))
(3)
@ass@k: k: - m With k the advan- tage function for generation i reduces to A; = max(riss) max(r-4). In contrast to the previous example, the advan- tage A; here is non-negative. In fact, if we assume the reward is ordered r(1) S r(2) S ... S..S ((A) we can rewrite the gradient estimate equivalently as
A
where the additional baseline is 1 = k Σ A, (y- By construction, the new advan- tage A = A, Á is zero-mean and arguably the new esti- mate has even lower variance. However, such an estimate introduces a bias to the objective it optimizes.
x
(ra)) r(k-1)) Ve log (V(1))))))
Lemma 2. (Objective of the biased gradient estimate)
Le.. the advantage is non-zero only for the best generation ((4)- The advantage is effectively the difference between the best and the second best reward. Such rewrite is sensible because, if the second best and the best reward are the same, there is no incentive in updating the policy.
The gradient estimate with zero-mean advantage function in Eqn (3) optimizes for an alternative objective
Σ f
(4)
To impart more intuition on the sparsity of the learning signal, consider a binary reward problem for k we expect the learning signal to be more dense when the average solve rate is at Ol.- If the solve rate is higher, the problem is too easy; otherwise the problem is too hard. In both cases, the learning signal decreases for the gradient estimate.
Proof. The gradient bias comes from the baseline term which evaluates to E E - oV A simple
algebraic manipulation shows that the aggregate gradient expectation
E
fly-,) V Vo log
Majority voting: f(y) = r (maj(a)). Here, we will adopt the notation for a == a1sk and a for the leave-one- out variant. With majority voting, the advantage function for generation is A = r (maj(a)) - (maj(a-))), measuring the impact of answer a; on the majority voted answer a. Assume that all m < k unique answers are sorted in their count la(al as ((1) < ...)(mm) then (maj(a-)) * (maj(a) if and only if a, a(m) and [[m)] = (mm-1)) + 1, since this is the only case where the leave-one-out voted answer changes. In other words, the advantage is all-zero in case a particular answer (a(m) dominates the count such that [mm)) |0(m-1y] + 1 since there is no incentive in updat- ing the policy.
which is the unbiased gradient estimate to the objective
E E E f(y
The objective differs from the initial k-sample objective in Eqn (1) it measures the lcave-one-out k 1-sample objective f (y=:) averaged over all l samples. The bias is hence
B ( (())) 1 -
3.2. Trading-off further variance reduction with bias
already zero-mean. The bias is non-trivial in general.
In general, with the leave-one-out control variate, the advan- tage estimate becomes more sparse. Indeed, the effective advantage A; measures the impact of a individual sample on the global objective. For example, for the pass@ k, A, is only non-zero for the reward maximizing generation; for the majority voting, A, is non nonzzero only when it is possible to alter the majority voted answers.
@ass@k. Recall that the samples are ordered T(1) < r(2; I ... We can write explicitly the new objective for the biased gradient estimate in Eqn (4) as
k 1
E
k T(k) + T(--1)
3

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Mean performance
0.1
0.8
Pass @4
0.4
0.4
mean policy gradient
assagl policy guadient
based passgle policy gradient
0.0
400 Learning steps
000
1000
0.0
0
[L Meet)
KL medd
Figure 1. Comparison of different gradient estimates a bandit case. We set up a bandit problem with y = 100 possible actions and each roward r(y) is deterministic scalar sampled from unit Gaussian. We compare three algorithmic variants: the mean policy gradient, the pass @ k policy gradient and biased variant. All algorithms apply k = samples per update with learning rate n = 1.0. Overall, we that the baseline gradient makes the fastest improvement the mean performance, when graphed against the learning steps (left plot); however, it is generally less KL-efficient than other s-sample alternatives (middle plot). When measuring the pass @ l performance, the k-sample gradient estimates lead significantly faster improvements (right plot).
which is a convex combination of the best and the second best reward. The bias is hence [i T()-1)) which is the gap between the best and the second best reward. Though such bias vanishes as k increases, it also means that the overall gradient estimate has sparser signal. It is less insightful to make explicit the new objective for the majority voting case, which we detail more in Appendix B.
Empirical study on a bandit case. To better illustrate the property of the k-sample gradient variants, inves- tigate the behavior of various policy gradient estimates in the bandit setting. There are a total of | 100 ac- tions and each incurs deterministic reward r(y) sampled from a unit Gaussian. We consider softmax parameteri- zations ( x exp (0(y)) and all algorithmic variants apply k = 4 samples for the update and identical hyper- parameters and initializations. Figure shows the evaluated performance of different variants over time.
3.3. Additional discussions
We provide further discussion on the property of the k- sample gradient estimate.
When measuring the mean performance, the mean policy gradient algorithm clearly maximizes the performance at a faster rate. However, when graphed against the KI. diver- gence KL (To, Tett) against the initial policy Tea we see that the pass k gradient estimate (Eqn (2)) and its biased variant (Eqn (3)) more KL-efficient (Gao et al., 2023). This follows from the fact that the k-sample gradient esti- mates have more conservative updates compared to the mean gradient. In general, maximizing the mean performance can also improve pass@k k performance as we observed in the plot. In this case, the other two algorithmic variants are by design more efficient.
Interpolating mean and general Is-sample objectives. We see that by subtracting the mean of the k-sample advan- tage in Equ (3), we manage to recover a smoother average of leave-one-out objective (Eqn (4)). In fact, by taking this approach this approach further and constructing leave-n-out objectives,
construct an interpolation between the original k-sample objective and the mean objective. See Appendix B for more discussions.
Efffec of increasing number of samples k. For the base- line policy gradient, varying the number of sample & does not impact the expectation of the gradient estimate. The corresponding gradient estimates always approximate the gradient of the mean performance. Increasing the value of k helps reduce variance of the estimate, but has diminishing effect when the batch size is large enough. For the k-sample objectives, varying the value of k changes the objective it- self. For the pass @ i objective, increasing ls means that we care increasingly about the extreme values (i.e., max values) of the reward distribution. For the majority voting objective, as & - 00, we optimize for the margin likelihood that the final answer is matched against a". We will ablate the empirical impact of in Section 5.
4. Related work
We discuss how our approach relates to prior work.
Other multi-sample objectives. Multi-sample objectives have been considered extensively in the variational inference literature (Raiko et al., 2014; Burda et al., 2015; Mnih Rezende, 2016). Using terminologies from this work, they consider k-sample objectives of the following form
1
p(æ; V/i
E log
T:
(a M/.

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Level
Level
Level!
mean perforanicc
bleed pangle odly gredet poodle polk grabet
D055084
30
73
KL.TT
KL
KL
KL, (T.
Figure 2. MATH training pass ls 8B model. We compare three baselines: regular mean policy gradient algorithm and two variants of pass@ k policy gradient algorithms (unbiased and biased). We split the performance across MATH difficulty level and report the mean performance and pass performance over time. We observe that as training progresses, pass & policy gradient algorithms seem to display slight advantage over the baseline algorithm.
where p. q are usually two distributions. Such an objective is by design a lower bound to the marginal likelihood, which is often considered the ultimate objective (Burda et al., 2015). In gradient-basee meta-learning, the k-sample objective takes the following form (Finn et al., 2017; Fallah et al., b020a;b; Tang, 2022)
Optimizing for inference time objectives. Concurrently, Balashankar al. (2024) proposed optimize for pairwise win rate based on best of & sampling. Under certain condi- tions, they showed that the objective is equivalent to the reg- ularized RL problem with transformed reward, which can be approximately optimized. Amini et al. (2024) pursued a similar approach but noticeably took log transformation of the score. Such an objective not possible to estimate in unbiased, instead, a variational approximation is needed. Chow et al. (2024) proposed to optimize best of k perfor- mance for applications where scoring function (or verifier) is available during training. Both work have shown merits in accounting for inference time objectives during training.
E
J
(x, yi)
where g(,, 34) E Rdd is meant be gradient update to the vector 0 E Rd and J Rd + E is a scalar function. In both cases, the effective aggregate function is a transformation of the average function. Intuitively, such objectives are smoother than pass@k k and majority voting. The specitic structure of this objective also produces gradient estimate that trades-off bias with variance (Fallah al., 2020b; Tang development here.
5. Experiments on mathematical reasoning
Throughout, we focus the mathematical reasoning dataset MATH (Hendrycks al., 2021) where the prompt con- sists in asking the model mathematical question with short-form ground truth answer a" available. Given the model generation = (c. a) which typically consists of a step-by-step chain-of-thought reasoning e and a proposed answer a, the reward is computed a". We adopt Sympy (Meurer et al., 2017) to automatically match the answers and assign reward of 1 if there is match and = -1 otherwise. As such, the objective E[r] also measures the average accuracy of the policy. We train on models on the MATH training set with various ob- jective alternatives introduced above. During experiments, we provide the model with a system prompt that asks for step-by-step solution, followed by final answer.
al., 2021), complementary
Leave-one-out control variate. The leave-one-out con- trol variate has been studied in Mnih & Rezende (2016) for variational inference, and in Kool et al. (2019); Ahmadian et al. (2024) for improving the mean objective a REIN- FORCE algorithm. As an alternative, Mnih & Rezende (2016) analyzed the control variate v; Ey-; [f()] which accounts for the credit assigned other lo - sam- ples in expected sense. Such an objective can be espe- cially useful for the k-sample objectives in our case since the advantage A, = fly) - U; is by construction zero-mean and arguably has lower variance. However, estimating such quantities introduces additional overhead practice.
match between a and
Our main experiments are based on the 8B model from the
5

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
MATH/ Level all
gorithm (blue) seems obtain better training performance overall, both in terms of the mean and the majority voting at k performance. Meanwhile, the biased variant seems to perform similarly as the policy gradient baseline.
81
pa5s@4
82
The breakdown also shows that most improvements seem to come from higher level of difficulty compared to the easier categories. This might be explained via a similar argument as the pass k algorithm: the k-sample gradient estimates tend to have sparser advantage and hence less update for the model over fixed number of training steps. This makes them slightly more KL-efficient since they focus more on more challenging problems.
80
mean polccy gradient k-B
plased pass@k policy gradient k-8
pass@k policy gradient
mean policy gredient k=2
alased pasa@k pollcy gradient low
passgk policy aradient k=2
14
2.5
5.0
7.5
10.0
12.5
15.0
17.5
20.0
(TI, rer))
Figure 3. Ablation with number of samples k. We vary the number of samples & for each gradient update for the pass@ k objective. We observe more efficiency gains for the pass@1 k gradient estimates compared to the policy gradient baseline. Importantly, note that as l varies, the pass @ l algorithm changes its objectives.
5.2. Ablations
We carry out ablations along a few dimensions of interest.
Number of samples k. We first ablate with the number of samples & used for constructing the gradient estimates. Throughout, we still apply the same hyper-parameter includ- ing batch size as the base experiments. We focus on the pass@ k objective here and discuss more about the majority voting objective in Appendix A.
Llama 3 model family (Dubey et al., 2024) though we also carry out ablations on models of other sizes. We apply on- line policy gradient algorithmic variants and investigate their performance during training. All variants apply identical pyper-parametee such that they all apply k = 4 samples for gradient estimations, which we detail in Appendix A.
Figure 3 shows the training performance as vary the number of samples k E (2, 8). Overall, we observe that as k: increases, there is no discernible difference in the perfor- mance of the policy gradient estimates. However, increasing k seems to make the pass@ k sample estimates more KL- efficient, i.e., = obtains a better performance than k: = given a fixed KL budget. The performance improvement over the policy gradient estimate is also quite significant. However, though the pass k algorithm is more KL-effiiieet as i increases, they tend to also deviate less from the ref- erence policy Tee with a fixed number of learning steps. This means that to achieve a target level of pass @ perfor- mance, the regular policy gradient baseline, despite being less KL-efficient, can still retain an advantage due to its faster deviation from the reference policy.
5.1. Training performance
@a Figure 2 shows the training performance of few algorithmic variants over time. The MATH dataset contains 5 difficulty levels (the easiest is level and the hardest level 5). We break down performance based on the levels and report both mean performance and pass@ k = 4.
few observations are in order: when measuring the mean performance, the mean policy gradient seems to obtain slight improvement over algorithms that aim for the pass @ k performance. Meanwhile, the dedicated k-sample gradient estimates generally perform better in pass @4 4. The biased variant (green) seems to generally strike a trade-off between the pass@k gradient estimate (blue) and the baseline, which compatible with the theoretical designs. The breakdown shows that for pass @ k. most performance difference comes from the more difficulty categories than easy ones. This might be explained by the fact that since the k-sample gradi- ent estimates update policy more conservatively, they tend to have relatively more updates for the more difficult prompts, compared to the mean policy gradient.
Model size. We also replicate similar experiments across few other model sizes: 3B and 70B, both from the Llama 3 model family. We notice that 3B model tends to dis- play a similar trend of performance trade-off compared to the 8B: the specialized k-sample estimates tend to outper- form baselines on k-sample objective during training, while slightly under-performing for the mean performance. The performance improvement also mostly stems from the more difficult subset of the dataset.
Majority voting. In similar vein, Figure 15 shows the training performance of few algorithmic variants. We break down performance by difficulty level and measure both the mean performance and the majority voting at k = 4 performance. Interestingly, under this set of comparison, the policy gradient algorithm with majority voting at = al-
The performance trade-off for the 70B case is less clear. The policy gradient baseline remains performant even for the k-sample objectives. We suspect that this is because the MATH dataset is a relatively simple task for the 70B model (Yue et al., 2024). Indeed, both the pass@k: k: 4 and
6

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Level
Level
meean performani nic
50.0
45.0
polly protie Ease pasoge polrr adert
pD5s884
35 128 135 KL (hrl
20
10. KL Ri
36. (eed
33
12.0
Figure 4. HARP training pass @ A 70B model. We observe that the regular policy gradient estimate improves over the pass k variants for the mean performance metric, while mnder-performing on the pass@ objective. Such a trade-off is less significant for the MATH dataset. where we speculate that the 70B model is too powerful and learning signals are too sparse to make a difference.
majority voting at A 4 are 90% for the MATH training set. This implies that the k-sample gradient estimates will not produce much signal, given fixed training budget.
training time. We also note that, training progresses, the policy gradient estimate tends to regress pass (iss@) k when k is large (k E {16, 64)). arguably due to impacts from sample diversity.
In light of the previous observation, we examine HARP dataset (Yue et al., 2024), a carefully curated math dataset consisting of hard problems. Though HARP was meant for evaluation, we monitor algorithmic variants' training performance
Training on HARP
6. Experiments on code generation tasks
For code generation, we conduct our experiments Code- Contests (Li et al., 2022), a competitive programming bench- mark containing 13k problems as training set and evaluate on the valid and test set. Similar to previous work (Xu et al., 2024; Gehring et al., 2024), we use a r = +1/ - 1 reward indicating whether the Python code solutions pass all the given tests that come with the dataset. Due to the challeng- ing nature of CodeContests, we use Llama 3.1 70B Instruct model and optimize for pass @8 metries on the CodeCon- tests training set. To measure generalization, we also re- port the performance on another competitive programming benchmark, TACO (L.i et al., 2023) on the easy and hard split, for which we have decontaminated the CodeContests training set against TACO eval set.
sanity check. Fig- ure 4 shows the results for the pass@ @ & objectives. For the pass@k @k k = objective, we observe significant im- provement of k-sample gradient estimates compared to the baseline algorithm. The mean performance of various algo- also observe improvement
rithms is similar, as expected. in the majority voting metric, as well for the 8B model though the improvement appears less significant compared to the 70B case. This might be because the HARP dataset is too difficult for 8B models. See results in Appendix A.
5.3. Evaluation performance
Throughout evaluation the MATH test set, apply temperature sampling with 7 1, top p 1 to be iden- tical to the training setup. In Figure 14 we showcase the evaluation metrics for the @ objectives. We show the majority voting metrics in Appendix A. We noted that training and evaluation performance are not perfectly cor- related- we speculate this might be because since regular policy gradient estimate performs more effective updates during training, it makes up for more generalization gap from training to evaluation.
Figure: 5 shows clear trade-off between optimizing for pass @ and p on CodeContests valid and test set. This trade-off generalizes in-domain to TACO easy and hard split. Across all evaluation set, mean policy gradient achieves the best pass@) performance (or mean perfor- mance). However, also clearly degrades on pass @ as the training progresses. In clear contrast, pass@ k policy gra- dient achieves the best @ performance. The pass @ performance improves since the onset and remains over the course of training. We do not observe a significant differ- ence between different pass @k gradient variants, since they both achieve much better performance than the baseline.
In particular, the policy gradient baseline performs quite well for the evaluation pass @ 4, despite being slightly out- performed by the k-sample pass@k: l gradient estimates at

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
CodeContests/vaidd
CodeContests/tett
TACOleasy
TACO/hard
30
34
26
26
p35s
22
28
10
16
16:
45.0
16
40
42.5
14
pass@8
35
12
10
30
32
25
30
20.0
0.05
0.15
0.10
0.05
0.15 KL(TT, Meed)
(L (ire))
(TT, Rew)
(IT,, Ter)
mean policy gradient
pass@k policy gradient
biased pass@k policy gradient
Figure 5. Code generation task evaluation performance. We observe clear trade-off off between pass I and pass @ CodeContests and TACO using Llama 3.1 70B mean policy gradient achieves the best mean (pass@) 1) performance, while pass gradient variants clearly achieve much better performance fo the pass @ performance.
Table The pass@k performance of Llama 3.1 70B Instruct on CodeContests valid/test up pass 100. Models trained CodeContests training set. The pass@k k objectives optimize for k = 8.
CodeContests / Valid
CodeContests / Test
Method
Variant
pass @
pass 10
pass @ 100
pass
pass 10
pass @ 100
mean
24.6
32.9
38.1
26.2
35,1
39.8
Policy Gradient
pass @ k
21.7
42.2
54.9
20.6
41.7
51.8
biased pass @k
18.0
38.4
54.2
19.3
39.5
51.9
mean pass @k
20.0
27.5
33.5
24.9
34.6
41.1
PPO
19.2
37.3
52.4
21.2
41.4
51.0
biased pass @k
21.6
41.7
55.8
22.6
45.1
56.3
Integration with PPO. We also investigate how the k- sample gradient variants be combined with much more sophisticated algorithmic stack to illustrate its generic practi- cal utility. Note that the @ k objective can be seamlessly integrated into other policy gradient variants, such as PPO (Schulman et al., 7), which adds clipping and an addi- tional value function that serves as the baseline. Figure shows the performance of using PPO. We can observe that PPO with pass @1 objective enters in high-KI. regime the end of the training, due the fact that its update generally leverages much more dense signals. The sparse nature of pass k objective keeps the policy under low-KL regime and maintains higher pass @ performance. For the purpose of training language models with fixation on the reference policy (Ziegler et al., 2019; Ouyang al., 2022), pass @ objectives achieve balance for the performance metries.
Model size. We also include the experiment result with Llama 3.1 8B Instruct model in Appendix A.4. Compared 70B, the performance differences across training objectives are less pronounced. We posit that the model size and the benchmark play a crucial role; for 8B model, the average pass @ on CodeContests modest, resulting in sparse reward signal. Intuitively, pass@ k objective will zero out the advantage for problems solved more than once. Therefore, expect the difference widen as the dataset contains more problems that the model is able to solve multiple times.
6.1. Generalization to different l: for evaluation
We investigate the impact of different training objectives, when combined with increased inference time budget. Con- cretely, we study how pass@k k performance scales with the number of samples k at evaluation time.
8

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
CodeContests/valid
CodeContests/kest
TACO/easy
TACO/hard
24
31
pass@1
30
22
29
22
20
20
27
18
26
10
10
10
10"
10
10
10
10a1
10a
10a
10
10
10
10
10
10
10
10
10P
10
45.0
42.5
44
14
pass@
40.0
12
35
37.5
35.0
30
10
30
32.5
36
30.0
34
10--
10
10
10"
10a
10
10
10
10a
10a
10--
10
10'
10
10-2
10
10
10"
10
[TT, Rest)
(T. Reat)
(T, Reet)
(T,,,, Reet)
mean PPO
pass@k PPO
biased pass@k
Figure 6. PPO code generation evaluation results with Llama 3.1 70B. We combine k-sample gradient estimates with a PPO implementation stack and compare with the default mean policy gradient baseline. We see that biased pass @ gradient estimate especially achieves good performance tradc-off. Details of how the k-sample gradient variants are combined with PPO can be found in Appendix A.4.
Table The @ @ performance of Llama 3.1 70B Instruct on TACO easy and hard split up to pass 100. Models are trained on CodeContests training set. The pass@k objectives optimize for l = 8. When integrated with pass@ and biased pass @ updates, both policy gradients and PPO have improved pass@k k performance, especially for large values of test time
TACO/easy
TACO/hard
Method
Variant
pass @ 1
pass 10
pass @ 100
pass @ 1
pass 10
pass @ 100
mean
29.9
34.5
38.2
8.1
9.9
11.8
Policy Gradient
pass
29.6
43.2
48.9
7.8
14.4
18.0
biased pass @k mass@k
26.7
42.5
51.5
7.0
14.3
21.9
mean pass @k k
27.5
35.3
39.5
7.3
8.8
11.1
PPO
28.0
41.6
49.5
7.6
14.2
18.7
biased pass @k
27.2
44.8
53.8
7.7
14.6
21.3
We show in Table and Table the pass@k performance across different k of Llama 3.1 70B Instruct model trained using different objectives. Models are trained with 6400 gradient update steps. For pass@k k objective, models trained to optimize pass@ & (k: = 8) can generalize to pass rate with different up to = 100. Amazingly, the performance improvements scales with k: for both CodeContest and TACO, the improvements go from 5 10% at k 10 to 10 20% at k 100. In general, improvements hard and test splits are a bit less than improvements easy and validation splits.
eralize to pass rate with different & up k 100. with over +10% solve rate improvement over the baseline. validation and test set.
7. Discussion and limitations
Lots of efforts remain to adapt the training methods in- ference time objectives beyond this work. Applying the methodology here naively would incur large computational burden (Silver et al., 2016; 2018; Brown & Sandholm, 201 8; 2019). Furthermore, many recent inference time algorithms do not treat each sample equally, such as self-correction, reflection and more complicated logic (Yao et al., 2022: Huang al., 2023; Asai et al., 2023). This requires careful algorithmic improvements for better credit assignment.
Generalization to 1: = 100 at test time. We show in Table and Table 2 (in the Appendix) that Llama 3.1 70B Instruct model trained to optimize pass@k: (k 8) can gen-

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Impact Statement
Christiano, P. F., Leike, J., Brown, T., Martic, M., Legg, S., and Amodei, D. Deep reinforcement learning from human preferences. Advances in neural information pro- cessing systems, 30, 2017.
This paper presents work whose goal is to advance funda- mental algorithmic development. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21 2024.
References
Ahmadian, A., Cremer, C., Gallé, M., Fadaee, M., Kreutzer, J., Pietquin, O., Üstün, A., and Hooker, S. Back to basics: Revisiting reinforce style optimization for learning from human feedback in Ilms, arXiv preprint arXiv:2402.14700, 2024.
Fallah, A., Mokhtari, A., and Ozdaglar, A. On the con- vergence theory of gradient-based model-agnostic meta- learning algorithms. In International Conference on Arti- ficial Intelligence and Statistics, pp. 1082-1092. PMLR, 2020a.
Amini, A., Vieira, T., and Cotterell, R. Variational best-of-n alignment. arXiv preprint arXiv:2407.06057 2024.
Fallah, A., Mokhtari, A., and Ozdaglar, A. Provably conver- gent policy gradient methods for model-agnostic meta- reinforcement learning, arXiv preprint urXiv:2002.05133. 2020b.
Asai, A., Wu, Z., Wang, Y., Sil, A., and Hajishirzi, H. Self- rag: Learning retrieve, generate, and critique through self-reflection. arXiv preprint arXiv:2310. 1511,, 2023.
Finn, C., Abbeel, P., and Levine, S. Model-agnostic meta- learning for fast adaptation of deep networks. Interna- tional Conference on Machine Learning, pp. 1126-1135. PMLR, 2017.
Bai, Y., Kadavath, S., Kundu, S., Askell, A., Kernion, J., Jones, A., Chen, A., Goldie, A., Mirhoseini, A., McKin- non, C., al. Constitutional ai: Harmlessness from ai feedback. arXiv preprint arXiv:2212.08077 2022.
Gao, L., Schulman, J., and Hilton, J. Scaling laws for reward model overoptimization. In International Conference on Machine Learning, pp. 10835- 10866. PMLR, 2023.
Bakhtin, A., Brown, N., Dinan, E., Farina, G., Flaherty, C., Fried, D., Goff, A., Gray, J., Hu, H., et al. Human-level play in the game of diplomacy by combining language models with strategic reasoning. Science, 378(6624); 1067-1074, 2022.
Gehring, J., Zheng. K., Copet, J., Mella, V., Cohen, T., and Synnaeve, G. Rlef: Grounding code llms execution feedback with reinforcement learning. arXiv preprint arXiv:2410.0208/ 2024.
Balashankar, A., Sun, Z., Berant, J., Eisenstein, J., Collins, M., Hutter, A., Lee, J., Nagpal, C., Prost, F., Sinha, A., et Infalign: Inference-aware language model align- ment. arXiv preprint arXiv:2412. 2024.
Guo, D., Yang, D., Zhang, H., Song, J., Zhang, R., Xu, R., Zhu, Q., Ma, S., Wang, P., Bi. X., et al. Deepseck-rl: In- centivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.
Bottou, L. Large-scale machine learning with stochastic gradient descent. In Proceedings of 20MPSTAT'010,, pp. 177-186. Springer, 2010.
2025.
Hendrycks, D., Burns, C., Kadavath, S., Arora, A., Basart, S., Tang, E., Song. D.,, and Steinhardt, Measuring math- ematical problem solving with the math dataset. arXiv preprint arXiv:2 103.02300300874. 2021.
Brown, N. and Sandholm, T. Superhuman ai for heads-up no-limit poker: Libratus beats top professionals. Science, 359(6374):418422) 2018
Brown, N. and Sandholm, T. Superhuman ai for multiplayer poker. Science, 365(6455):885-890 2019.
Huang, J., Chen, X., Mishra, S., Zheng, H. S., Yu, A. W., Song, X., and Zhou, D. Large language mod- els cannot self-correct reasoning yet. arXiv preprint arXhv:3310001998, 2023.
Burda, Y., Grosse, R., and Salakhutdinov, R. Importance weighted autoencoders. arXiv preprint 1rxiv:1509.00519, 2015.
Jaech, A., Kalai, A.. Lerer, A., Richardson, A., El-Kishky, A., Low, A., Helyar, A., Madry, A., Beutel, A., Car- ney, A., et al. Openai ol system card. arXiv preprint arXi::2412.16720 2024.
Chow, Y., Tennenholtz, G., Gur, I., Zhuang, V., Dai, B., Thiagarajan, S., Boutilier, C., Agarwal, R., Kumar, A., and Faust, A. Inference-awart fine-tuning for best-of- n sampling in large language models. arXiv preprint MxXi:2241.11288/ 2024.
Kool, W., van Hoof, H., and Welling, M. Buy 4 reinforce samples, get baseline for free! 2019.
10

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Lerer, A., Hu, H., Foerster, J., and Brown, N. Improving policies via search in cooperative partially observable games. In Proceedings of the AAAI conference on artifi- cial intelligence, volume 34, pp. 7187-7194, 2020.
Silver, D., Huang. A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershel V., Lanctot, M., et al. Mastering the game of with deep neural networks and tree search. nature, 529(7587))484 489, 2016.
Li, R., Fu, J., Zhang, B.-W., Huang, T., Sun, Z., Lyu, C., Liu, G., Jin, Z., and Li, G. Taco: Topics in algorithmic code generation dataset, 2023. URL /arxiv. org/abs abs//3122 2212. 14852.
Silver, D., Hubert, T., Schritwieser,, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre. L.,, Kumaran, D., Grae- pel, T., et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Sci- ence, 362(6419): 10:0::111111444 2018.
Li, Y., Choi, D., Chung, J., Kushman, N., Schrittwieser, J., Leblond, R., Eccles, T., Keeling, J., Gimeno, F., Dal Lago, A., et al. Competition-level code generation with alpha- code. Science,
Sutton, R. S. and Barto, A. G. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
2022.
Lightman, H., Kosaraju, V., Burda, Y., Edwards, H., Baker, B., Lee, T., Leike, J., Schulman, J., Sutskever, I., and Cobbe, K. Let's verify step by step. arXiv preprint arXiv:2305.20050 2023.
Sutton, R. S., McAllester, D. A., Singh, P., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. In Advances in neural in- formation processing systems, pp. 1057-1063, 2000.
Meurer, A., Smith, C. P., Paprocki, M., Čertík, O., Kirpichev, S. B., Rocklin, M., Kumar, A., Ivanov, S., Moore, J. K., Singh, S., Rathnayake, T., Vig, S., Granger, B. E., Muller, R. P., Bonazzi, F., Gupta, H., Vats, S., Johansson, F., Pedregosa, F., Curry, M. J., Terrel, A. R., Roučka, v., Saboo, A., Fernando, I,, Kulal, S., Cimrman, R., and Scopatz, A. Sympy: symbolic computing in python. PeerJ Computer Science, 3:e 103, January 2017. ISSN 2376-5992. doi: 10.77 10.7717/peerj cs. 103. URL https: //doi. www.i.oogg10..0 7717 0ttp://0...71//peer-cc.....3.33
Tang, Y. Biased gradient estimate with drastic variance reduction for meta reinforcement learning. In Inter- national Conference Machine Learning, pp. 21050- 21075. PMLR, 2022.
Tang, Y., Kozuno, T., Rowland, M., Munos, R., and Valko, M. Unifying gradient estimators for meta- reinforcement learning via off-policy evaluation. arXiv preprint urXiv:2106. 131.21116611155. 2021.
Uesato, J., Kushman, N., Kumar, R., Song, F., Siegel, N., Wang. L., Creswell, A., Irving, G.. and Higgins, I. Solv- ing math word problems with process-and outcome-based feedback. arXiv preprint ar
Mnih, A. and Rezende, D. Variational inference for monte carlo objectives. In International Conference on Machine Learning, pp. 2188-2196. PMLR, 2016.
2022.
Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., and Kavukcuoglu, K. Asyn- chronous methods for deep reinforcement learning. In International Conference on Machine Learning, pp. 1928- 1937, 2016.
Wang, X., Wei, J., Schuurmans, D., Le, Q., Chi, E., Narang, S., Chowdhery, A., and Zhou, Self-consistency im- proves chain of thought reasoning in language models. arXiv preprint arXiv:2203. 2022.
Williams, R. J. Simple statistical gradient-following al- gorithms for connectionist reinforcement learning. In Reinforcement Learning, pp. 5-32. Springer, 1992.
Ouyang. L., Wu, J.. Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
Xu, S., Fu, W., Gao, J., Ye, W., Liu, W., Mei, Z., Wang, G., Yu, C., and Wu, Y. Is dpo superior to ppo for I1m alignment? a comprehensive study. arXiv preprint arXiv:2404. 10719, 2024.
Raiko, T., Berglund, M., Alain, G., and Dinh, L. Tech- niques for learning binary stochastic feedforward neural networks. arXiv preprint arXiv: 1406. 2989,2999.8 2014.
Yao, S., Zhao, J., Yu, D., Du, N., Shafran, I., Narasimhan, K., and Cao, Y. React: Synergizing reasoning and acting in language models. arXiv preprint arXiv:2210.0362 2022.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint (XXYv:7777.6344 2017.
Yue, A. S., Madaan, L., Moskovitz, T., Strouse, D., and Singh, A. K. Harp: A challenging human-annotated math reasoning benchmark. arXiv preprint arXiv:2412.08811 2024.
Silver, D., Lever, G., Heess, N., Degris, T., Wierstra, D., and Riedmiller, M. Deterministic policy gradient algorithms. In ICML, 2014.

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Zheng, K., Decugis, J., Gehring, J., Cohen, T., Negrevergne, B., and Synnaeve, G. What makes large language models reason in (multi-turn) code generation? arXiv preprint arXiv:2410.08105, 2024.
Ziegler, D. M., Stiennon, N., Wu, J., Brown, T. B., Radford, A., Amodei, D., Christiano, P., and Irving, G. Fine-tuning language models from human preferences. arXiv preprint arXiv: MoXiv:/90000859, 2019.
12

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
APPENDICES: Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
A. Tlyper-paramett and experimental details
We experimented with the Llama 3 model of size 3B, 8B and 70B (Dubey et al., 2024). All experiments are conducted with identical
settings: we always apply batch size of B = 64 prompts per update, and sample k distinct generations per prompt by default. All training and evaluation sampling are conducted at temperature of r = and with top-p- = 1.
We train on the MATH training with 7500 examples and evaluate on the test set with 5000 examples (Hendrycks et al., 2021). A supervised fine-tuning the training is conducted to warm up the RI training. As part of the ablation, also train the HARP dataset, which consists of about 4300 examples of difficult math problems harvested from few public sources (Yue et al., 2024).
For both training and evaluation, we provide system instructions that ask the model to generate a response with step-by-step solution, followed by a final conclusion phrased as the final answer is followed by the answer predicted by the model. This consistent with the prompt structure discussed for Llama models (Dubey et al., 2024; Yue et al., 2024).
A.1. Ablation on number of samples
@. Figure shows the full results for the ablation results. We see as & increases, the s-sample gradient estimates produce performance improvements on the pass k objective. More significant improvements are observed for the more difficult split of the dataset.
MATH / Level
MATH / Level 3
MATH Level
MATH Level all
02
mean performance
50
50
201
bissed pastgik policy gradien knd
patcy a-d
30-
50
passgl policy gradient 1u-5
mean pelicy gradient bo2
treasd pass gli policy cradinnt
10-
paogl palicy gradkent
5.
150
0.0
10.0
22.5
20.0
50
150
22.0
80
92.5
75-
90.0
pass@14
87.5
70-
85.0
60-
13.5
75.2
55
50
125
15.0
17.
20.0
0.0
25
100
12.5
15.0
175
7.5
10.0
12. 27.5 20.0
2.5 50
75
125
150
175
26.0
(TT,
KL TTpe))
KL, (eer)
[TT,,
Figure 7. Ablation with number of samples & for majority voting. We vary number of samples for each gradient update the pass @ & objective. We observe more efficiency gains for the pass@ k gradient estimates compared to the policy gradient baseline. Importantly, note that as & varies, the pass@ i algorithu changes its objectives.
Majority voting. Figure shows the ablation results for the majority voting based results and k-sample gradient based algorithm. We observe a slight improvement of the k-sample based algorithms over policy gradient baseline, though the impact of the number of samples is less signficant compared to the pass @ k case.
A.2. Training on HARP
We train both 8B and 70B on the HARP dataset, with both the pass@k and majority voting based algorithms and objectives.
13

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
MATH / Level 1
MATH / Level 3
MATH / Level 5
MATH / Level all
30-
mean performance
70
60
05
maan pelley cradient len moan polic gradient ke2 pa tcy grasiens Io ased pattcc po foy giadent
55
45
50
12.5
50
60
mV@4
70
65
75
60
0.0
23
12. KL(TT (yar)
10.0
15.0
5.0
10.0
15.0 11.5 70.0
0.0
104 12.
20.0
600
3.0
12.5
13.0
17.3
[T, Mlat)
KL Mird)
KL
Figure 8. Ablation with number of samples & for majority voting. We vary the number of samples for each gradient update for the pass@k & objective. We observe modest efficiency gains for the pass @ le gradient estimates compared the policy gradient bascline but the impact of the number samples k is less significant.
8B model. Figure 9 and Figure 10 show the results for the pass@k and majority voting respectively. Overall, we see that the regular policy gradient algorithm is modestly outperformed by the k-sample based algorithms, in terms of the KL-efficiency trade-off against the training performance.
However, one observation that compatible with prior results, that the regular policy gradient algorithm tends allow the policy to deviate more from the reference policy using fixed number of updates compared the k-sample based algorithms. This presents practical trade-off'- if we want to obtain good level of performance with a fixed number of training steps, the policy gradient algorithm might be better option since it is more compute-efficient.
70B model. Figure 11 shows the HARP training results for the majority voting case. We see that the improvement is less clear for both the mean and the majority voting performance the three algorithmic variants seem to obtain similar performance, with the k-sample policy gradient algorithm obtaining very slight advantage over others.
A.3. Evaluation
Figure 16 shows the evaluation performance for the majority voting on MATH, using the 8B model. Across the various baselines we compare, we note that the biased k-sample policy gradient algorithm seems to slightly outperform the policy gradient algorithm as KL is large, though the policy gradient algorithm itself is clearly a strong baseline.
The unbiased k-sample policy gradient algorithm seems slightly underperform on evaluation, which contrasts the strong training performance in Figure 15. We speculate that this suggests an intriguing interaction between the training and testing performance worthy of further investigation: likely the merits of such s-sample based algorithms depend on the nature and the size of the dataset, which is better tested out in practical applications.
A.4. Experimental details and additional result on CodeContests and TACO
Experimental details. We experimented with Llama 3.1 8B and 70B Instruct model. For both model we use the same hyperparameters. We use a learning rate -c constant learning rate scheduling with 50 warmup steps and weight decay of 0.1. We sample & = generations per prompt. We update the model with a mini batch size 2 with sequence length 8192 and train in total 8k gradient update steps. In training, sampling conducted at a temperature of T = 1 and with top-p = 1. In evaluation, we sample at a temperature r = 1 and with top-p = 0.95. Both pass@1 and pass @ are estimated out of 20 samples. We evaluate the correctness of generated Python code using the official codebase of et (2022). Our
14

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
HARP/A Level 1
HARP/ Level
HARP / Level 5
HARP / Level all
55.0
30.0
25
mean performanc
12.9
27.5
34-
38.0
67.5
12.5
cy
15.0
60.0
14
12.5
33
90
pa55
40 KL. (TT, Reerl
10 20
60
70
18
50 KL(TT, Tieer)
40
60
20
30
40
K2,TT, MTrr)
(L,,,, Tlyet)
Figure Training reward for @ the HARP dataset for 8B model. We can see observe relative performance efficiency of k-sample gradient variants compared to the baseline method. The gains are overall less significant compared to the 70B model, whose model capacity more suitable the challenging HARP dataset.
asynchronous online training decouples the training and the inference, for which include importance sampling term to correct the slightly off-policy nature between current policy and the behavior policy. The original CodeContests training set contains 13328 problems. We follow the decontamination process presented by Zheng et al. (2024) to decontaminate CodeContests training set against TACO evaluation set. We further remove problem instances with less than 5 test cases. This results in total 12275 problems which we use to train our model.
Integrating with PPO. We can integrate advantage clipping and an additional value model Vi to serve as an additional baseline. Given problem æ, we sample multiple responses {yı1 y2. .... Me from Te the training objective aims maximize the following objective:
To (1( Teu (2/(z)
To Te( (r)
e, + A)
(0) = Eynnos
min
A, clip
Teau (3)(z)
&
Te (x(()lz))
To(2)
e, e, 1, Ai).
min
A,, clip
h:
Toaa (v x))
Taa (v1
Å, R4 Vi(x),
where R2 takes the following different forms according to the training objective:
Mcan(pass@@ 1) training objective:
R4 Ti
Pass @ k training objective:
R. JE(I....kk
T3 ry j71
Biased pass @ k training objective:
R₂ max rj max r je(1....k)
max Tj -ax Tj (E(L....kk)
jsk
with being the +1/ binary reward of whether the code is correct. We train the value model Vp to minimize the following clipped value loss:
L (w)) E max (Vv((z) R,) (clip (v), V (x) ac Voiun (x) + + -R,)²)
We set a 0.2 in our experiments.
15

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
HARP Level
1
HARP/ Level 3
HARP Level
HARP Level all
14
326
mean perforaanc
77.1
10
234
175
my oolicy grediett
15.0
Mexxed potks pratent
pelicy grodlant
31.0
$3.5
48.0
mV80
215
35.0
25.0
32.5
22.5
22.5
230
27.5
(TT MTat))
(Trer)
(1T, Meerl
Figure 10. Training reward for majority voting on the HARP datasct for 8B model. Overall, the training performance improvement less significant compared the pass
Llama 3.1 8B Instruct Performance. For Llama 3.1 8B Instruct model, we show Figure 12 the performance of using mean policy gradient, pass@k policy gradient and biased pass@ policy gradient. We also show in Figure 13 the performance when integrating the training objective with PPO.
B. Additional discussion
B.1. Biased objective for majority voting
Following the recipe to derive biased k-sample objective for the k-sample objective, we discuss the case for majority voting. Assume there are unique answers and they are ordered by count a(i), E [m]]. It is not difficult to show that, the resulting objective is a weighted objective the most common and second most common answer
P(y) r (A(m) + (1 - P(v)) (a(m-1))
where P(y) the probability that the leave-one-ous majority voted answer is the second most common answer. Intuitively, the leave-one-out majority voted answer can either be ((m) (when for example [a(m)] [a(m-1)] + 2) the second most common answer ((m-1))
B.2. Interpolation between k-sample objectives and mean objective
We have hinted at the observation that the biased k-sample objective smoother objective than the original k-sample objective. The biased objective also approaches the mean objective Σ which is arguably the most smooth objective one can construct.
Leave-p-out objectives. Extending the idea of leave-one-out further, we discuss leave-p-out, which depicts a spectrum of objectives interpolating the original k-sample objective and the mean objective. We focus the pass @ objective as follows
1 E
Σ SES,
max its
k
where Sp is the set of all subsets of indices from {l.... k). with (x) being its cardinality. Notice that for = 0 we recover the initial problem of the pass@1 objective. For we get the de-meaned objective introduced the previous
16

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
HARP Level
HARP/ / Level 3
HARP / Level 5
HARP Level all
54
mean perforancc
42
msan pubyy gradent pappaity mypallyygdder gradient blseed policy grad em
30
5.0
75
26.0
125
15.0 17.5
0.0
5.0
7.5
10.0
12.5
156
17.5
0.0
2.5
7.5
10.0
12.5
150
37.5
0.0
2.5
7.5
154
12.5
25.0
175
56
90
mV@4
54
52
52
733 36.0 12.3
15.0
175
0.0
75 10.0 125 19.6
10.0 22.5 15.0 37.
0.0
25
125 KL[11, Reg!)
124
25.0
17.5
(M, (Tre))
KL Rear)
KL(M, Heat)
Figure 11. Training reward for majority voting the HARP dataset for 70B model.
section. For = we get the problem of maximizing the expected rewards. For = we get the objective
E
max TI
k(k 1)
practical implementation of gradient estimate of the objective (5) can be written as the de-meaned gradient using the advantages max; T; maxigs T; for Sp- Indeed:
VE
max Ti 165
Sp
E
V log π (U:)
Σ
max Ti iæx'
kE(I....k)
S'ESp
E
V log T(u)
Σ max "48"
p-1
SES,
iEs
S'
E
V T(g)
maxi;
max ids
idss'
SES
iEs
S'ES,
E
Σ
Σ
V
log T(ss)
max T3 max T3
Σ Ag
k-1
k p
#ESp
iEs
S'ES,
For example, the corresponding gradient for 2 is
1
fT
V log F (16) + (1)))
max
max (+.9))
k(k 1)
2
Σ
Ave
=:A...
17

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
CodeContests/valid
CodeContests/test
TACO/easy
TACO/hard
pass@
15
7
6
13
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
18
16
pass@8
12
24
12
10
8-
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
0.0
0.1
0.2
0.3
KL (17, Trre)
(L Nrae))
KL Thr))
KL Ter
mean policy gradient
pass@k policy gradient
biased pass@k policy gradient
Figure 12. Llama 3.1 8B Instruct performance CodeContests and TACO. We compare 3 methods: the policy gradient with leave-one-out control variate that optimizes for pass@1 I, the pass @ policy gradient and the biased pass@ k policy gradient with mean advantage serving the baseline.
Algorithm 1 Online policy optimization
1: INPUT policy Te
2: while t
1, do
3: (i) Sample prompt
4: (ii) Collect & trajectories per prompt (y:) from <)(-1²) for E 21.
5: (iii) Update policy parameter 9 based on one of the gradient variants above (Eqn (2) or Eqn (3)).
6: end while
Softmax objectives. An alternative approach to interpolate between k-sample objective and the mean to through softmax. For example, consider the objective
Σ
Per
where pt x exp(Brr) is the softmax distribution scaled by parameter 2 0. The above objective approaches mean when and @ when / + oo. We find that the softmax objective tends to produce much more learning signal than pass @ k, since its advantage estimates are generally more dense.
In general, might also make of leave-one-out softmax objectives baselines for variance reduction opposed to leave-one-out @ k. This makes sure that the optimization objective is intact while introducing generally smoother signals.
18

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
CodeContests/valid
CodeContests/test
TACO/easy
TACO/hard
9
8
16
pass@1
7
7
15
14
5
5
13
1
10
10
10
10
10-1
20.0
18
28
8
17.5
26
pass@8
15.0
14
22
20
10
18
2
10-3
10
10
10-3
10-2
10
10-2
10-3
10-1
10
(1,,7, Meed)
KL(IT, Tire))
KL(T,, Mred)
KL Medd
mean PPO
pass@k PPO
biased pass@k PPO
Figure 13. Llama 3.1 8B Instruct performance on CodeContests and TACO. We incorporate the PPO training objective and compare
variants: the vanilla PPO objective that optimizes the pass performance, the pass @ l objective and the biased pass @ objective with mean advantage serving as the baseline.
MATH test Level
MATH test Level
MATH best / Level
MATH test Level all
50-
mean performance
46
20
mean pilkcy
25.50
44
69
42
05
D85s04
54.75
67
40
04.25
53.75
KL Rietl
KLAM, (Tar)
KL(TT, mrer)
[L Rrerl
Figure 14. MATH test 8B model. We carry out evaluation the MATH test set, using the same sampling hyper-parameter the training time. We observe that the regular policy gradient baseline obtains a strong performance for pass@k k with [1,4]. However, as increases, it tends to be outperformed by the s-sample gradient variants largely due significant drop in the sample diversity.
19

---

Optimizing Language Models for Inference Time Objectives using Reinforcement Learning
Level
Level
Level
mean perfomance
00
25
3.0
75
100
32.5
150
17.5
332
00
25
50
75
160
325
150
17.5
201
33
30
75
10.0
32.
15.0
11.
20.0
0.0
235
50
13
300
325
15.0
375
300
07.
mvg-
70-
00.
152
25
100 32. 150
175
332
00
25
5.0
100 325 150 KLIR, Rrael
17.5 251
235
3.0
10.0
32.
15.0
135
20.0
360 50
300 125 K2(M. Rhai
1500
175
KL(R. Tiar)
K2
Mhe
Figure 15. MATH training majority voting 8B model. We compare three baselines: regular mean policy gradient algorithm and two variants of majority voting policy gradient algorithms (unbiased and biased). We observe the unbiased majority voting policy gradient algorithm improves slightly over the other two alternatives.
57
63.2
mean performance
15
61.6
mV@4
62.5
mV@64
39.5
59.
52
63.0
$2.0
605
0.
25
12. 15.0 KL(T, Traa)
10.0
17.5
20.0
5.0
73
10.0
125
15.0
37.5
20.0
00
25
75
12.5 KL Beer)
10.0
150
175
20.0
25
50
10.9
12.5
15.0
17.5
200
KL Ryet)
KL(TT. Meet)
Figure MATH test 8B model for majority voting. We carry out evaluation MATH set, using the same sampling hyper- parameter as the training time. We observe that the regular policy gradicnt baseline obtains strong performance overall, though slightly outperformed by the biased k-sample gradient algorithm. The original unbiased k-sample gradient algorith, despite performing better training reward, is slightly underperforming for evaluation.
20


Reviews:
[{"review_id": "AamV0Qwy9S", "reviewer": "Reviewer_KzmQ", "summary": "This paper explores the potential of explicitly optimizing language models for inference-time performance objectives, particularly pass@k and majority voting, using reinforcement learning (RL). The authors propose a k-sample objective formulation and derive both unbiased and biased gradient estimators, including a leave-one-out control variate for variance reduction. The paper provides theoretical justification for the estimators, including proofs of unbiasedness and characterizations of induced biases. Empirical evaluations on mathematical reasoning (MATH, HARP) and code generation (CodeContests, TACO) benchmarks demonstrate that training-time optimization of inference-time metrics can yield superior performance on those metrics (especially pass@k), compared to standard mean policy gradient approaches.\n\n**update after rebuttal**:\nThanks for the authors' responses. My concern remains that the generalization sometimes degrades when optimizing inference-time metrics. However, the overall idea and method are interesting, and I encourage the authors to explore more on this.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.\n\n> 1. Generalization of improvements to evaluation time is not always consistent.\n\nThe generalization of improvements to evaluation is indeed less straightforward, as we have explicitly discussed in the paper. The generalization gap depends in a complex way on the training data and model sizes. The cleanest way to assess our algorithmic progress is via the training performance, though we have also included evaluation for completeness. We still see evaluation improvements overall, though they are not as strong as training rewards improvements.\n\n>  2. The method may be less effective for problems where models already achieve high performance (e.g., MATH with LLaMA 70B).\nThis is true and we have elaborated on this observation in the paper. The k-sample objectives are generally less effective for easy problems to the model, since the model does not have strong signals to make improvements. However, for more difficult problems (HARP for Llama 70B), we see more significant improvements, which might be a more useful case for many practitioners as well.\n> Consider emphasizing earlier that these methods do not require auxiliary models (unlike best-of-k with reward models)\n\nWe will make sure to emphasize this point earlier on in this work. We have introduced the inference objectives of interest in Section 2.\n\n> Including a discussion on compute efficiency (vs benefit) can help practitioners decide when to use these techniques.\n\nWe will elaborate on such points more extensively in the revision. In general, we find that the newly proposed algorithms automatically improve inference time objectives with a similar computational budget as the baseline algorithms (which maximizes the mean objectives). The algorithmic change is also almost a few-line code change, which is quite convenient for practitioners.\n\n> Can you elaborate on why training-time improvements on pass@k/majority voting objectives do not always translate to better test-time performance? Could this relate to diversity collapse or overfitting to training prompts?\n\nWe conjecture that this might be more related to the generalization gap between training and test set, which depends in a rather complex way on factors such as the model size, and training data that went into the model pre-training and SFT (the released Llama 3.1 model might have SFT’ed on similar dataset to our RL prompt set, which might reduce the level of generalization we expect). As a result, maybe training set diversity and performance improvements do not transfer 100% to the evaluation set.\n\n> When should practitioners prefer biased vs unbiased gradient estimators? Do you recommend any heuristics for setting k?\n\nWe feel that a good approach is to decide on the estimator adapting to the specific use case. Overall, we find the two estimators to be both quite competitive, and since the unbiased gradient estimator optimizes for the exact objective, it might be the first thing to try.\n\nFor setting k: this depends on the application of interest. We can choose a specific k that will be used at inference time. We can also choose a generic value of k to maximize model performance while diversifying the samples as much as possible - in general, a large value of k indicates more diversity. We find that values of {4,8,16} are good starting points, depending on the application.", "claims_and_evidence": "Claims:  \nInference-time objectives like pass@k and majority voting can be explicitly optimized during training. Doing so leads to improved performance on those inference-time metrics, especially on more difficult problems. The proposed gradient estimators are theoretically justified and practically effective.\n\nWhile the training-time improvements are consistent, test-time generalization improvements are more mixed, especially in the majority voting case and for high-capacity models (e.g., LLaMA 70B on MATH). The authors acknowledge this.", "methods_and_evaluation_criteria": "The methods are clearly described and well-motivated. The paper adapts standard reinforcement learning techniques to k-sample objectives, and proposes novel unbiased and biased gradient estimators for these settings. Evaluation is performed on realistic and challenging benchmarks: MATH and HARP for mathematical reasoning, and CodeContests and TACO for code generation.\n\nMetrics are appropriate for each domain. Evaluation is comprehensive, including training curves, sample-efficiency analysis (via KL-divergence plots), and ablations on model size and number of samples k.", "theoretical_claims": "The paper contains several theoretical claims, mainly regarding gradient estimation:\n1. The leave-one-out gradient estimator is unbiased.\n2. The biased variant optimizes a modified leave-one-out objective.\n\nThese derivations follow standard techniques in policy gradient literature. The proofs are concise, correct, and help contextualize the practical impact of bias vs variance.", "experimental_designs_or_analyses": "The experimental design is sound and robust. Key strengths:\n- Ablations on different model sizes (3B, 8B, 70B).\n- Evaluation across multiple benchmarks and multiple difficulty levels.\n- Consideration of both training performance and test-time generalization.\n- Integration of the proposed gradient estimators with PPO, demonstrating compatibility with standard RL fine-tuning pipelines.\n\nOne weakness is the limited analysis of why generalization sometimes degrades when optimizing inference-time metrics. This could benefit from further exploration.", "relation_to_broader_scientific_literature": "The paper is well-positioned in the context of:\n- RLHF and preference-based fine-tuning.\n- Best-of-k and self-consistency generation\n- Inference-time optimization", "essential_references_not_discussed": "The related work section (and appendix) is comprehensive and cites relevant prior work on:\n- RL with multi-sample objectives,\n- Policy gradient variance reduction,\n- Self-consistency and majority voting.", "other_strengths_and_weaknesses": "Strengths:\n- Good theoretical and empirical contributions.\n- Clearly written, with well-illustrated examples and intuitions.\n\nWeaknesses:\n- Generalization of improvements to evaluation time is not always consistent.\n- The method may be less effective for problems where models already achieve high performance (e.g., MATH with LLaMA 70B).", "other_comments_or_suggestions": "- Consider emphasizing earlier that these methods do not require auxiliary models (unlike best-of-k with reward models)\n- Including a discussion on compute efficiency (vs benefit) can help practitioners decide when to use these techniques.", "questions_for_authors": "- Can you elaborate on why training-time improvements on pass@k/majority voting objectives do not always translate to better test-time performance? Could this relate to diversity collapse or overfitting to training prompts?\n- When should practitioners prefer biased vs unbiased gradient estimators? Do you recommend any heuristics for setting k?", "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": "3947I6aI13", "reviewer": "Reviewer_dbZK", "summary": "This work presents a training objective that directly optimizes pass@k and maj@k performance for LLMs, and optimizes them using standard RL training algorithms (policy gradients with a baseline, PPO) known in practice. They evaluate performance on a synthetic bandit task, math reasoning and coding tasks, and find that they are able to improve performance over algorithms that only optimize pass@1 performance.\n\n## update after rebuttal\nIn the LLM literature, it is well-established that optimizing for pass@1 does not necessarily improve pass@k. Thus, it makes sense to optimize it with an RL training algorithm like Reinforce. The main contribution of the paper is to apply Reinforce on the pass@k objective. The results and discussion are interesting. My main concern is with technical novelty since prior works like Chow et. al. do exactly the same thing (choice of optimizing pass@n for a trained vs. ground-truth reward function does not change the algorithm or setup). **But, I am still willing to raise my score to 2.5 and would not be strongly opposed to accepting the paper if other reviewers are willing to champion it, since the empirical insights are useful to the community**.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.\n\n> The paper lacks key comparisons with InfAlign (Balashankar et. al.), and Chow et. al., that both propose RL training objectives that directly optimize pass@n or BoN performance.\n\nWe have discussed the technical differences between our work and Balashankar et. al, Chow et al, in Section 3.3, with more extended discussion in Appendix C.\n\nAs pointed out in the paper, Balashankar et. al and Chow et al consider structurally different problem formulation as we have in this work. Balashankar et. al considers a pairwise preference setting (vs. point-wise reward setting in our case) while Chow et al considers the case with auxiliary models such as verifiers or scorers (vs. no requirement for auxiliary models in our case). As a result, it is challenging to set up an apple-to-apple comparison against such concurrent work since they tackle very different settings. Further, we also feel that it’s fair not to make direct comparison to unpublished concurrent work.\n\n> The performance on MATH is underwhelming, especially Maj@k in Figure 15, and even pass@4 in Figure 2, is only better in a narrow KL regime, with the proposed policy gradient.\n\nWe acknowledge the fact that the performance improvement on MATH is not huge, though they are still statistically significant. The upper bound of the performance improvement also depends heavily on the base model we started with (which we do not have full control over). MATH as a benchmark task is also relatively easy since the pass@16 has also reached ~86%, at a competitive level as the frontier models such as O1 and Deepseek R1. This also puts a limit on the amount of improvement.\n\nThe narrow KL divergence regime: Fig 2 shows improvement across all KL values we observe during training. For Fig 15, since maj@k signals are much sparser, this makes it more difficult for the model to reach a larger KL divergence deviation during training. However, for the KL divergence regime it reaches, the performance improvements are clear. We have highlighted this point as an algorithmic trade-off in the paper.\n\nAdditionally, depending on practical applications, the “narrow KL” regime might suffice since for general post-training, ideally the RLHF stage does not lead to too much deviation from the reference policy. We leave this to the judgement of practitioners.\n\n> Do the authors think that there is a fundamental tradeoff between optimizing for pass@1 and pass@k performance?\n\nIn principle, pass@1 and pass@k share common goals but there is a conflict between the two objectives, in that pass@k always encourages higher diversity and coverage over the solution space, while pass@1 encourages more determinism. \n\n> Why is it that we cannot learn a single model with good pass@1 performance on easy problems and good pass@k on hard problems? Is this lack of representation capabilities, pre-training biases, or issues with training objectives/algorithms?\n\nIt should be possible in principle: this might require a hybrid objective, and switching between pass@1 vs. pass@k training objectives depending on the problem difficulty. We might even want to devise a novel algorithm that adjusts the objective based on the online sampled pass rate. This definitely is a promising path for future work that integrates the two objectives into a unified algorithm.\n\nThe pretraining biases and representation capabilities certainly plays a key role as well, as suggested by recent investigations to reproduce Deepseek R1 (e.g., Qwen models can reproduce R1 behavior more consistently, and model sizes matter too). That said, training algorithms / objectives should make a difference, as we have demonstrated in the paper, especially because RL has become increasingly important in post-training.", "claims_and_evidence": "Yes, the claims made are clear. They propose an objective to optimize pass@k directly and run experiments that implement this objective and RL algorithm.", "methods_and_evaluation_criteria": "Yes, the math and coding benchmarks used seem standard. The authors only plot performance on the pass@k objective they trained for. It would be interesting to see, if the performance gain on value of k used in training also extends to other values of k. If not, how was this value of k (4 or 8) chosen in practice? This part was a bit unclear.", "theoretical_claims": "The only theoretical claims are made with respect to the bias of their proposed gradient estimates, which seem correct given their readily apparent nature.", "experimental_designs_or_analyses": "Yes, I read through their experimental setup for math reasoning and coding benchmarks, and it looks reasonable to me.", "relation_to_broader_scientific_literature": "The key idea in this paper is to optimize for inference time performance, like pass@k or maj@k. From prior work, we know that optimizing for pass@1 may not necessarily lead to better performance on pass@k. So, the idea is to improve performance at a higher compute budget. In general, the broader goal should be to learn policies that push up the pareto frontier of performance against inference time compute. This work does not do that, but instead makes an attempt to push the performance up at higher test compute budget, i.e., optimize for performance at higher pass@n.", "essential_references_not_discussed": "No, I believe the paper discusses most relevant works that optimize pass@n performance. Though, in the next iteration, it would be useful to include some discussion on training time interventions, that optimize for inference time performance, for example more recent meta-RL training objectives like Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning, Qu. et. al 2025.", "other_strengths_and_weaknesses": "Strengths:\n- The gains on pass@8 performance for coding benchmarks is impressive, compared against a baseline that only optimizes pass@1.\n- The proposed approach is fairly straightforward to implement.\n\nWeaknesses\n- The paper lacks key comparisons with InfAlign (Balashankar et. al.), and Chow et. al., that both propose RL training objectives that directly optimize pass@n or BoN performance. \n- The performance on MATH is underwhelming, especially Maj@k in Figure 15, and even pass@4 in Figure 2, is only better in a narrow KL regime, with the proposed policy gradient.", "other_comments_or_suggestions": "- If the authors can run experiments against suggested baselines that would strengthen the paper. \n- If the authors can also demonstrate that the approach can extend to broader functions like maj@k or even self-verification@k (see Sample, Scrutinize and Scale: Effective Inference-Time Search by Scaling Verification, Zhao et. al.), that would also strengthen this paper.\n- In general if the authors can empirically show that the proposed approach can improve performance on math reasoning, either at a broader range of KL divergence, or for a broader range of k (beyond the pass@k trained for), that would be great.", "questions_for_authors": "- Do the authors think that there is a fundamental tradeoff between optimizing for pass@1 and pass@k performance? Why is it that we cannot learn a single model with good pass@1 performance on easy problems and good pass@k on hard problems? Is this lack of representation capabilities, pre-training biases, or issues with training objectives/algorithms?", "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": "hhQdxXxcE2", "reviewer": "Reviewer_61s6", "summary": "This paper investigates the reinforcement learning algorithms of LLMs in the training time for achieving the test-time objectives. Specifically, it focuses on k-sample policy gradient approaches assuming pass@k and majority vote are the test-time strategies of interest. The authors proposed a leave-one-out like advantage function, which is claimed to be an unbiased estimation of the policy gradient. Furthermore, they showed another variant which further reduces the variance through introducing a bias term by adding an average baseline in the advantage function. Empirical results showing that the biased k-sample method outperforms standard policy gradients. The proposed algorithm is validated through different benchmark datasets (e.g., MATH, CodeContests and TACO) using different sizes of models, demonstrating practical applicability of the proposed methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.\n\n> In the proof of Lemma 1, why is the expectation of two independent variables zero? I didn’t see any assumption made on either variable being zero.\n\nThis is because the expectation of the score function is zero $E[ \\log \\pi(y_i|x)]=0$. More precisely,\n\n$E[f_{-i}\\log \\pi(y_i|x)] = E[f_{-i}] \\cdot E[\\log \\pi(y_i|x)]=0$\n\nWhere the equality makes use of the statistical independence between $i$ and other samples $-i$. We will make this more clear in the revision.\n\n> The proof for Lemma 2 is hard to follow. Not sure which expression the authors refer to. I suggest the authors add a complete proof in the appendix for better understanding of the claims.\n\nThank you for the suggestion, we will add a more detailed explanation in the Appendix.\n\nThe goal of Lemma 2 is mainly to derive the analytic form of the biased gradient estimate, which turns out to optimize a biased objective with an interesting analytic form - it is the average of leave-one-out objectives from $k$ samples.\n\n> The variance reduction is not shown both theoretically and empirically.\n\nThe understanding that the control variate improves variance, either in the form of the leave-one-out control variate (for the unbiased estimator), or subtracting a mean baseline (for the biased estimator), was based on common consensus in the literature on this topic [1]. In fact, we cannot show theoretically that the variance is reduced in all cases, due to the multiplication with the score functions; but we can show that the advantage estimation has lower variance. We can make the theoretical argument more precise in the revision.\n\nFor empirical validation: we never implemented an estimator without control variate for variance reduction, as such estimators are plain REINFORCE gradient estimators and do not work well in practice. We can add such an ablation in the revision in case the reviewer finds it useful.\n\n[1] Kool et al, Buy 4 REINFORCE Samples, Get a Baseline for Free!, 2023\n\n\n\n> why does pass@k policy gradient become much worse than the baseline when PPO is used?\n\nWe have some speculations regarding this observation: the pass@k policy gradient, compared to the mean policy gradient and biased pass@k policy gradient, has the sparsest gradient, i.e., the weight before the policy gradient is non-zero only for positive samples where it is the only 1 positive sample among k samples, while for the other samples the weight is 0. Biased pass@k policy gradient, by subtracting the mean among k samples, makes the policy gradient no-sparse in such cases. When incorporating PPO, adding a learnt value breaks this sparsity due to the imperfectness of the value model. \n\n> What is the performance comparison between the proposed approach (e.g., pass@k policy gradient) and purely SFT trained or pass@1 mean policy gradient method with pass@k test-time strategy? This would be meaningful as it shows how much benefit the proposed method can bring.\n\nWe’ve observed a performance drop when doing SFT on CodeContests training set. Given that Llama 3.1 70B has already been heavily tuned in the post-training, some code\nsolutions in CodeContests training set are of less quality than the data presented in its\npost-training phase. For example, some imports in the Python codes are outdated (e.g., from\nfractions import gcd will throw an ImportError since Python 3.9).\n\nTherefore, we included a baseline for Rejection Sampling Finetuning (RFT) instead, where the model generates 200 rollouts for each problem instance in the training dataset and collect the correct solutions, we downsample them to maximum 50 solutions per problem and conduct SFT on this dataset.\n\nWe show the comparison below, the best performance is bold and the second best is underscored. With these comparisons, the both pass@k (k=8) variants achieves the best performance in terms of the closest metrics pass@10.\n\n| Method                         | pass@1 | pass@10 | pass@100 |\n|---------------------------------|--------|---------|----------|\n| **Llama 3.1 70B Inst.**         | 16.1   | 34.2    | 48.2     |\n| **SFT**                         | 10.0   | 26.7    | 39.2     |\n| **Rejection Sampling Finetuning**| 18.1   | 37.2    | 52.5     |\n| **mean PPO**                    | 24.9   | 34.6    | 41.1     |\n| **pass@k PPO**                  | 21.2   | 41.4    | 51.0     |\n| **biased pass@k PPO**           | 22.6   | 45.1    | 56.3     |", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "- In the proof of Lemma 1, why is the expectation of two independent variables zero? I didn’t see any assumption made on either variable being zero.\n- The proof for Lemma 2 is hard to follow. Not sure which expression the authors refer to. I suggest the authors add a complete proof in the appendix for better understanding of the claims.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The contribution of this paper has shown that aligning the test-time objective with the training-time effort is promising, which is related to RL training and test-time compute of LLMs.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "## Other Weaknesses\n- The variance reduction is not shown both theoretically and empirically.\n- The only baseline the authors considered was mean policy gradient, which is not comprehensive.\n- The comparisons/inconsistent results on PPO integration are not explained adequately. For example, why does pass@k policy gradient become much worse than the baseline when PPO is used?", "other_comments_or_suggestions": "- I suggest the authors add a complete proof for lemmas in the appendix for better understanding.", "questions_for_authors": "- What is the performance comparison between the proposed approach (e.g., pass@k policy gradient) and purely SFT trained or pass@1 mean policy gradient method with pass@k test-time strategy? This would be meaningful as it shows how much benefit the proposed method can bring.", "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": "TvLIOTQXyv", "reviewer": "Reviewer_xLpG", "summary": "This paper explores the potential benefits of explicitly training LLMs for test-time inference. The authors introduce a new RL objective, which explicitly utilizes multiple samples and optimizes LLMs for inference-time objectives, like better pass@k performance or better majority voting performance. The empirical experiments show improvements on various math reasoning tasks and code generation tasks compared with baseline algorithms.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.\n\n> On the other hand, for general RL algorithms, I suspect the performance of algorithms like GRPO, that utilize the same number of samples, will even outperform the objective proposed in this paper.\n\nWe strongly disagree with the reviewer. \n\nIn our manuscript the baseline (mean PG, and mean PPO), as what the reviewer refers to as “general RL algorithm”, already uses the same number of samples as the other variants to enable a fair comparison. And we already provide empirical evidence in our paper that they do not outperform our proposed variants that optimize for pass@k.\n\nIn other words, general RL algorithms such as GRPO only optimize for the mean performance, and would similarly benefit from algorithmic improvements proposed in this work. GRPO is not designed to optimize for pass@k on its own.\n\n> Baselines like GRPO and other single-objective functions with the same sample complexity are necessary for comparison.\n\nPlease refer to our response to Reviewer bYLj for exp. on GRPO; we also include Dr. GRPO as a variant. Note that, our method could be integrated with various policy gradient variants, e.g., the PPO or vanilla Policy Gradient in the manuscript. Please refer to our response to Reviewer 61s6 for results of SFT and Rejection Sampling Finetuning (RFT) for additional comparison.\n\n> To be clear, GRPO [1] doesn't require any prior assumption on the test-time algorithm, but we can still evaluate it on both approaches.\n\nGRPO (or any policy gradient method derived from the objective $\\max E(f(y)))$ **does** have prior assumptions on the test-time algorithm: the one that uses the exact protocol as the $E(f(y))$, i.e., pass@1. The reviewer refers to [2] to argue that a model trained using $\\max E(f(y))$ objective works well on pass@k metrics, i.e., $\\max E(f(y_1, … y_k))$ objective. This could be also true on the other way round.\n\nWe motivate this from the first principle that such an objective is to maximize the k-sample objective $E(f(y_1, … y_k))$, contrary to the common formulation of maximizing cumulative return $E(f(y))$, where the return is a function for individual sample. Our proposed scheme works for settings that fit into this objective $E(f(y_1, … y_k))$. Specifically, we show that in the case of $f = r \\circ \\text{maj}$ and $f = \\max$, this corresponds to the maj@k and pass@k metrics in math reasoning and code generation and gives empirical evidence.\n\nThat said, our method also can be evaluated in multiple sampling budgets in test time (we do not set a constraint that training on k = 8 only works for k = 8 in test time, see our section of generalization to k = 100).\n\nAlso, we share exactly the same hyperparameter as GRPO: the number of samples per prompt. \n\n> As shown in Figure 18 and Figure 19 from [2], the R1-Distill-Qwen-7B model, a single model, outperforms both maj@k and pass@k separately-trained 8B models reported in this paper on MATH. This is not a fair comparison since R1-Distill-Qwen-7B is a stronger base model, but my argument is mainly to question the necessity of separate objectives.\n\nThe reviewer suggests that the access to a stronger model R1-Distill-Qwen-7B questions the necessity of separate training objectives.\n\nWe object to this argument because the R1-Distill-Qwen-7B model is distilled from a capable mode R1, which is specifically trained via RL methods, which is the focus of our paper.\n\nThis does not undermine the necessity of separate objectives but rather reinforces the validity and potential benefits of explicitly targeting inference-time objectives during training, which results in a better model in terms of pass@k performance and could bring the pass@k superiority to other models by following the same distillation process. Also, Figure 18 and Figure 19 from [2] is reported on MATH500, while our numbers are on a full MATH test set (a total 2500 problems).\n\n> There are no implementation details in the paper. How are the hyperparameters selected for different settings in the paper?\n\nPlease find implementation details in the Appendix A (in particular A. and A.4).\n\nFor sampling, we use top-p=1 and temperature $\\tau=1$ sampling with standard parameters, as they are compatible with the RL training, though evaluations can be executed with various sampling scheme. \n\nFor training, the experiments are rather robust to the choice of learning rate and other training parameters: Therefore, we fixed these hyperparameters among different settings (mean/pass@k objective) within math/code experiments respectively, to ensure a fair comparison.\n\nOther hyper-parameters such as $k$ are ablated in the experiments. We bump the $k$ for code experiments since CodeContests is a challenging benchmark.", "claims_and_evidence": "See “Methods And Evaluation Criter” and “Experimental Designs Or Analyses”.", "methods_and_evaluation_criteria": "I am not fully convinced by the motivation of this work. Model training doesn't explicitly account for the downstream inference time algorithm. As presented in work [1], there is no single inference-time algorithm that works best for any scenario, and usually, during model training, it is not practical to assume what inference-time algorithm will be used. It's a user decision, determined by computation budget, latency requirements, etc. However, the objective proposed in this work, if I understand correctly, seems to require separate training for different objectives, which doesn't seem to be well-justified.\n\n[1] Gao, Peizhong, et al. \"Meta reasoning for large language models.\" arXiv preprint arXiv:2406.11698 (2024).", "theoretical_claims": "The proof in section 3 is correct.", "experimental_designs_or_analyses": "As mentioned in the Methods and Evaluation Criteria section, on one hand, I don't think it is a valid assumption to assume which test-time algorithm to use during training. On the other hand, for general RL algorithms, I suspect the performance of algorithms like GRPO, that utilize the same number of samples, will even outperform the objective proposed in this paper. To be clear, GRPO [1] doesn't require any prior assumption on test-time algorithm, but we can still evaluate it on both approaches. As shown in Figure 18 and Figure 19 from [2], the R1-Distill-Qwen-7B model, a single model, outperforms both maj@k and pass@k separately-trained 8B models reported in this paper on MATH. This is not a fair comparison since R1-Distill-Qwen-7B is a stronger base model, but my argument is mainly to question the necessity of separate objectives; a more rigorous comparison is needed. Baselines like GRPO and other single-objective functions with the same sample complexity are necessary for comparison.\n\n[1] Shao, Zhihong, et al. \"Deepseekmath: Pushing the limits of mathematical reasoning in open language models.\" arXiv preprint arXiv:2402.03300 (2024).\n\n[2] Qu, Yuxiao, et al. \"Optimizing Test-Time Compute via Meta Reinforcement Fine-Tuning.\" arXiv preprint arXiv:2503.07572 (2025).", "relation_to_broader_scientific_literature": "This work is related to the claim that improving inference time performance can exceed the benefits of additional training time.", "essential_references_not_discussed": "The key contribution is a multi-sample RL objective that aligns with test-time inference. However, single-sample RL algorithm work is not well-discussed in this paper, such as [1], [2], and [3]. \n\n[1] Kumar, Aviral, et al. \"Training language models to self-correct via reinforcement learning.\" arXiv preprint arXiv:2409.12917 (2024).\n\n[2] Kazemnejad, Amirhossein, et al. \"Vineppo: Unlocking rl potential for llm reasoning through refined credit assignment.\" arXiv preprint arXiv:2410.01679 (2024).\n\n[3] Shao, Zhihong, et al. \"Deepseekmath: Pushing the limits of mathematical reasoning in open language models.\" arXiv preprint arXiv:2402.03300 (2024).", "other_strengths_and_weaknesses": "See “Methods And Evaluation Criter” and “Experimental Designs Or Analyses”.", "other_comments_or_suggestions": "See “Methods And Evaluation Criter” and “Experimental Designs Or Analyses”.\n\n**After Reply Rebuttal Comment**\n\nI maintain my concern about the generalizability of the method. As the authors mentioned in \"Reply Rebuttal Comment\":\n\n> If we train with pass@k and evaluate with maj@k, there is generally no guarantee that we can do better than pass@1.\n\nI think it's a strong assumption that \"inference time objective is known to inform the training time procedure.\" For multiple objectives during inference time, the computational overhead of retraining the model to optimize each target is concerning, and achieving better results by specifically optimizing for each objective seems less appealing. However, the authors have addressed my other concerns; therefore, I have updated my score accordingly.", "questions_for_authors": "There are no implementation details in the paper. How are the hyperparameters selected for different settings in the paper?", "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": "Q26n9r90Qg", "reviewer": "Reviewer_FCtK", "summary": "This paper investigates the impact of reinforcement learning (RL) objectives that optimize multi-sample metrics such as Pass@k and Majority Voting (Maj@k). It contrasts these objectives with standard mean reward RL objectives across mathematical reasoning and code generation tasks. The study highlights the trade-off between optimizing for multi-sample inference objectives and maintaining strong Pass@1 performance.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.\n\n> I think the gains in Pass@k come at a significant cost to Pass@1, particularly in coding tasks.\n\nOne major takeaway message is that there is a trade-off between pass@k and pass@1 in general. For coding tasks, pass@k is generally a more meaningful objective than pass@1 because we focus more on harder tasks, and obtaining one valid solution to a hard problem is enough [1].\n\n[1] Li et al, Competitive level code generation with AlphaCode, 2022\n\n>  Importantly, Mean PG outperforms other multi-sample objectives on the MATH test set, even at Pass@4.\n\nWe are being truthful about the train/test generalization gap here - the performance improvements of our algorithm are better assessed at the training set, since the algorithm does not account for train/test generalization. The generalization gap might depend on multiple factors out of our control, such as the choice of the base model and data used for pretraining the models.\n\nIt is true that the mean PG outperforms multi-sample objectives with MATH at pass@4, but the outperformance is marginal and limited to pass@4. The overall conclusions about the performance improvement from multi-sample gradient estimate still stand, but we are truthful about the fact that it does not necessarily generalize in all cases.\n\n> The approach depends on having a strong verifier, raising doubts about practical improvements from these objectives.\n\nThis is not true. We do not require verifiers at all during training, unlike Chow et al. Can the reviewer clarify on the claim that the approach depends on a strong verifier and where it is used?\n\nFor both training and evaluations, we use the standard numpy package to compute the rewards for training. This is common practice in the reasoning LLM literature.\n\n> The paper lacks a comparison with standard PG + entropy bonus.\nThe entropy bonus in PG algorithms is complementary to the approach we take. In other words, in principle, our method can be combined with an entropy bonus as-is. \n\nAs a result, we do not think adding entropy bonus comparison offers new insights into the current algorithmic designs. Entropy bonus might lead to performance improvements as a result of avoiding premature collapse, but it also introduces a tunable hyper-parameter, whereas our approach is hyper-parameter free compared to the baseline.\n\n> The majority of the related work section is relegated to the appendix, omitting key citations from the main text. Important prior work, such as Variational Best-of-N Alignment, is not discussed.\n\nWe apologize for the inconvenience of moving most related work discussion to the Appendix due to ICML page limit. When the limit is relieved (camera ready), we can move much of the discussion into the main paper.\n\nVariational BoN: this is a related work and we will discuss it. Variational BoN does not optimize for the original objective as-is, but rather requires a variational approximation. This leads to an algorithm akin to the one introduced in Balashankar et al, and in a pairwise comparison setting. In comparison, we focus on the point-wise reward case and require no approximation to the original objective of interest.", "claims_and_evidence": "Improved Pass@k Performance: The paper demonstrates that training with Pass@k policy gradient (PG) and biased Pass@k PG improves Pass@k performance. Supporting evidence is presented in Sections 4 and 5.\nEffect of Increasing k During Training: The paper examines the impact of larger k values on performance (Section 4.2).", "methods_and_evaluation_criteria": "The paper evaluates three RL Objectives: Mean Policy Gradient, Pass@k Policy Gradient, Biased Pass@k Policy Gradient\n\nPerformance is primarily assessed through Pass@1 and Pass@4 scores while comparing these methods against KL with a reference policy. Experiments are conducted on three datasets:\nMATH (Mathematical Reasoning), HARP (More difficult Mathematical Reasoning), CodeContests (Competitive Programming)", "theoretical_claims": "The paper provides analytical insights into optimizing inference-time objectives but does not introduce major theoretical results.", "experimental_designs_or_analyses": "- The paper effectively demonstrates that improving Pass@k often comes at the expense of Pass@1 performance, particularly in code generation tasks.\nPractical Limitations of Pass@k: Pass@k is not a practical metric in real-world inference scenarios as it requires an oracle verifier. While the authors argue that Pass@k represents an upper bound on inference-time scaling, it is unclear if this upper bound is useful given the lack of a reliable verifier (DeepSeek-AI, 2025).\n- The field is shifting towards inference compute scaling within a single trajectory (e.g., longer chains of thought) rather than multi-sample methods. The paper does not address how its proposed techniques integrate with this paradigm.\n- The authors compare their metrics against KL divergence from a reference policy but do not justify why this comparison is meaningful. If the primary objective is to improve Pass@1 and Pass@4, it is unclear why lower KL divergence is relevant.\n- The primary MATH benchmark results are reported for the training set in the main paper, while test set results are relegated to the appendix. This should be reversed.\nBaseline Comparisons: A simple PG with an entropy penalty should be included as a baseline for a fairer comparison.", "relation_to_broader_scientific_literature": "The paper focuses on the generator-verifier framework, where increased inference-time compute is allocated to generating multiple solutions and selecting the best one. However, recent trends have shifted toward single-trajectory inference-time scaling (e.g., longer chain-of-thought reasoning), limiting the practical applicability of methods discussed in the paper.", "essential_references_not_discussed": "Variational Best-of-N Alignment (https://arxiv.org/abs/2407.06057v1)", "other_strengths_and_weaknesses": "Strengths:\nI think the paper is well-organized and easy to follow. It evaluates reasoning and coding tasks across diverse datasets and honestly the analytical insights provide useful perspectives on the optimization of multi-sample objectives.\n\nWeaknesses:\n- I think the gains in Pass@k come at a significant cost to Pass@1, particularly in coding tasks. Importantly, Mean PG outperforms other multi-sample objectives on the MATH test set, even at Pass@4.\n- The approach depends on having a strong verifier, raising doubts about practical improvements from these objectives.\n- The paper lacks a comparison with standard PG + entropy bonus.\n- The majority of the related work section is relegated to the appendix, omitting key citations from the main text.  Important prior work, such as Variational Best-of-N Alignment, is not discussed.", "other_comments_or_suggestions": "I don't have any other comments.", "questions_for_authors": "I think I have discussed my questions in strengths and weaknesses section.", "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": "27k6GK9rvr", "reviewer": "Reviewer_bYLj", "summary": "This paper explores the benefits of explicitly training language models using reinforcement learning to perform well under inference time algorithmic techniques like pass@$k$ and majority voting. The authors argue that directly optimizing for inference time objectives can lead to improved performance on those specific metrics.\nThis paper provides a formulation for $k$-sample objectives and efficient gradient estimation techniques. The authors derive an unbiased gradient estimator using the leave-one-out method, and a biased variant that reduces variance by centering the advantage function.\nThe paper demonstrates these approaches on mathematical reasoning (MATH dataset) and code generation tasks (CodeContests), showing performance benefits when training directly for inference time objectives compared to training for mean performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.\n\n\n> Regarding the computational efficiency issue mentioned in weaknesses, can the authors provide another ablation study on the total number of samples? Namely, allow mean policy gradient to do k\n times of update than the proposed algorithms, and plot the figures.\n\nIn our manuscript the baseline (mean PG, and mean PPO), already uses the same number of samples as the other variants to enable a fair comparison in terms of compute cost. \n\nWe also conducted exp. using k samples for using the PPO loss but remove the value model. For mean objective, this corresponds to the GRPO loss.\nFor the pass@k objective, note that the proposed biased pass@k objective already subtracts the mean out of k samples, which coincides with the Dr.GRPO loss recently proposed https://arxiv.org/abs/2503.20783 (only subtract the mean but not divide by the std. deviation).\n\nFor completeness, we also include:\n1. the mean objective using Dr. GRPO loss \n2. biased pass@k objective using GRPO loss, i.e., further divide by the std. Deviation\n\n\nDue to time and resource constraints we are only able to finish the exp. using llama 3.1 8b model on the coding tasks. We will add the others in the revision.\n\nCheckpoints are all trained for 8000 gradient steps using the same set of hyperparameters. We report the number by avg the perf of the last 5 ckpts, (8k, 7.8k, …, 7.2k) steps, to ensure a non cherry-picked results.\n\npass@1:\n\n| Method                          | CodeContests/Valid | CodeContests/Test | TACO/easy | TACO/hard |\n|----------------------------------|--------------------|-------------------|-----------|-----------|\n| **Mean (Dr. GRPO)**              | 7.17               | 8.88              | 16.47     | 5.40      |\n| **Mean (GRPO)**                  | 9.21               | 8.39              | 17.89     | 4.74      |\n| **Biased pass@k (Dr. GRPO)**     | 6.79               | 7.03              | 15.87     | 4.86      |\n| **Biased pass@k (GRPO)**         | 8.11               | 6.67              | 14.51     | 4.12      |\n\npass@10:\n\n| Method                          | CodeContests/Valid | CodeContests/Test | TACO/easy | TACO/hard |\n|----------------------------------|--------------------|-------------------|-----------|-----------|\n| **Mean (Dr. GRPO)**              | 10.61              | 14.89             | 22.91     | 6.89      |\n| **Mean (GRPO)**                  | 13.88              | 14.41             | 24.11     | 7.06      |\n| **Biased pass@k (Dr. GRPO)**     | 15.71              | 16.44             | 27.01     | 8.55      |\n| **Biased pass@k (GRPO)**         | 16.83              | 15.45             | 26.08     | 7.37      |", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "This paper makes a contribution to language model training by showing a difference between and inference objectives. Most prior work has treated these as separate concerns, but this approach directly incorporates inference techniques into the training process.\nFor practical applications like code generation, where pass@$k$ metrics are commonly used in evaluation, this approach provides a clear path to improving performance.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "### Strengths\n\n- **Technical Contribution:** The theoretical framework for optimizing $k$-sample objectives extends beyond previous work in multi-sample objectives. The proposed estimators are verified to be KL-efficient than directly optimizing mean performance on benchmarks.\n- **Extensive Empirical Validation:** The authors conduct thorough ablation studies across model sizes (3B, 8B, 70B), different values of $k$, and various datasets, strengthening their claims.\n\n### Weaknesses\n\n- **Result Within Expectation:** Since there is gap between mean performance objective and pass@$k$ objective, the results in this paper is normally within expectation. It seems unfair to compare with mean policy gradient, which is not talored for pass@$k$ objective.\n- **Computational Efficiency:** The $k$-sample gradient methods require $k$ times more inference computation during training, which might limit practical applicability for very large models, especially because generating a sample is the most time-consuming process.\n- **Task-Dependent Performance:** As acknowledged by the authors, the benefits vary significantly based on task difficulty - for easier tasks, the standard policy gradient approach remains competitive (and more efficient in samples).", "other_comments_or_suggestions": "### Comments\n\nIn general, I think this is a comprehensive paper, which tells a complete story about the gap between mean performance objective and pass@$k$ objective, showcasing the inference-time RL method better aligns with the latter. However, this result is not surprising, as stated in weaknesses.\n\n### Suggestions\n\nThe authors should use the same colors in different figures. Currently Figure 2 and 3 have different colors for the same algorithms, which is highly confusing in the first sight.", "questions_for_authors": "Regarding the computational efficiency issue mentioned in weaknesses, can the authors provide another ablation study on the **total number of samples**? Namely, allow mean policy gradient to do $k$ times of update than the proposed algorithms, and plot the figures.", "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:

Meta-review: Language models are often optimized for next-token prediction but later evaluated on downstream tasks using metrics such as pass@k. This paper introduces a novel reinforcement learning technique that directly optimizes pass@k. According to reviewers, the paper is well-written and offers both theoretical and experimental contributions. Although there are concerns regarding two concurrent works, the authors address them in the related work section and explain the key differences. I recommend to discuss these reference in the introduction and provide experimental comparisons with at least one of them.
Decision: Accept (poster)