Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems
Authors: Christian Walder
Deep Tejas Karkhanis
Venue: NeurIPS 2025 spotlight ()
Keywords: reinforcement learning
llm
pass at k
inference time compute
monte carlo
gradient estimation
Primary area: reinforcement_learning
TL;DR: We provide a robust method of directly optimizing the pass at k with reinforcement learning, with theory and real world experiments.

Abstract: Reinforcement Learning algorithms commonly sample multiple ($n>1$) solution attempts for each problem and reward them independently. This optimizes for pass@1 performance and prioritizes individual sample performance over the diversity and collective utility of a set of samples. Such algorithms under-utilize the sampling capacity, limiting exploration and eventual improvement on harder examples. As a fix, we propose Pass-at-$k$ Policy Optimization (PKPO), a multivariate transformation on batches of rewards which leads to direct optimization of \passk\ performance, thus optimizing for sets of samples that feature a large maximum reward when considered jointly. Our primary contribution is to derive novel low variance unbiased estimators for the pass@k and its gradient, in both the binary and continuous reward settings. We show that optimizing with these estimators reduces to reinforcement learning with (batches of) rewards that have been jointly transformed by a function that is stable and efficient to compute.

While previous efforts propose transformations for $k=n$, our transformations are the first to enable robust optimization of the pass@k for any arbitrary $k \leq n$. Rather than simply trading off pass@1 performance for pass@k gains, our method allows annealing $k$ during training, optimizing both metrics and often achieving strong pass@1 performance alongside significant pass@k gains.

We validate our transformations on illustrative toy experiments, which reveal the variance reducing properties of our formulations. We also include real-world examples using the open-source models Gemma and Llama . We find that our transformation effectively optimizes for the target $k$. Furthermore, higher $k$ values enable solving more and harder problems, while annealing $k$ boosts both the pass@1 and pass@k. Crucially, for challenging task sets where conventional pass@1 optimization stalls, our pass@k approach unblocks learning, likely by improving exploration through the prioritization of joint utility over the utility of individual samples

Pass @K Policy Optimization: Solving Harder Reinforcement Learning Problems
Christian Walder & Deep Karkhanis
Google DeepMind
cww.lerr/ddarr//mmnggggggg..... com
Abstract
Reinforcement Learning algorithms commonly sample multiple (n >> 1) so- lution attempts for each problem and reward them independently. This optimizes for pass @1 performance and prioritizes individual sample performance over the diversity and collective utility of a set of samples. Such algorithms under-utilize the sampling capacity, limiting exploration and eventual improvement on harder examples. As a fix, we propose Pass-at-k: Policy Optimization (PKPO), a multi- variate transformation on batches of rewards which leads direct optimization of p performance, thus optimizing for sets of samples that feature large maximum reward when considered jointly. Our primary contribution novel low variance unbiased estimators for the pass@k and its gradient, in both the binary and continuous reward settings. We show that optimizing with these estimators reduces to reinforcement learning with (batches of) rewards that have been jointly transformed by function that is stable and efficient to compute.
derive
While previous efforts propose transformations for k n, transforma- tions are the first to enable robust optimization the pass@k for any arbitrary k Rather than simply trading off pass@1 1 performance for pass@k gains, our method allows annealing l during training, optimizing both metrics and often achieving strong @ performance alongside significant pass@k gains.
We validate our transformations on illustrative toy experiments, which reveal the variance reducing properties of our formulations. We also include real-world examples using the open-source models GEMMA2 and LLAMA3.1. We find that our transformation effectively optimizes for the target k. Furthermore, higher k values enable solving more and harder problems, while annealing & boosts both the pass@1 and pass@k. Crucially, for challenging task sets where conventional pass@1 1 optimization stalls, our pass@k approach unblocks learning, likely by improving exploration through the prioritization of joint utility over the utility individual samples
Introduction
Recent years have seen the rapid rise of large language models (LLMs) trained with internetsccalc pretraining data IRNS+ 18) with post training using both supervised fine-tuning [WBZ+21] and
reinforcement learning (RL) [AAA*23 23] [ea23] Ant GYZ+25). The seminal paradigm of RL with limited by the -erieederived data it is based on and the reward hacking issues that arise from the use of subjective signals more generally [ABC=2) To enable progress toward superhuman capabilities, current work is focusing on grounded reward signals that are free of fine-grad fine-grained human input as in code generation [STTR23 LWG+22 21 YTC+
human feedback (CLB+17
GZC+24) and mathematics 22 AT + + YSG 23
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

moXp@k N&. 0.1) and various
Vamax,@k
Tal nemurr le fun of
swward (al fancion
0.5
mak,pe) 1a
0.4
kmn
K-16
k-4
0.4
K=8
16
0.4
(a) ((x) and max,@k.
(b) g(x) and Vomax,
Figure 1: The effect of k on the optimal policy for one-dimensional toy problem. The policy is mal with mean parameter and fixed standard deviation 0.1. For the max,@k objective (left, defined in Equation (11) ) the optimal corresponds to the horizontal position with maximum max,@k, @k For the derivative (right, the estimation of which is the focus of this paper), the optimal 0 corresponds the location of the zero crossing. For larger & the optimal 0 is more risk tolerant, allowing more samples to exceed one (getting zero reward) in order to increase the chance of obtaining at least one sample close to, but less than one (getting a large reward). See Section 5.1 for more details.
The policy gradients family of RL methods Wi192 has proven effective in language model training 25) To scale to new capabilities, model training needs to tackle challenging RL task sets with no known solutions but for which correctness may verified, as in formal mathematics envi- ronments [AT] In such settings, the RL training loop both updates model parameters and searches for solutions to problems at the continuously advancing frontier of model capabilities.
The specific search algorithm introduces a coupling between inference and model updates, which means that naively optimizing the expected single sample reward, pass@1 may be sub- optimal. While various inference-time search methods are possible [HYM+24] KZA+24, LKB 23 WSL + 22), simply taking multiple independent samples from the model has proven rather effective 23]. Our contribution is to couple this simple search method with model parameter updates by enabling robust optimization of the passūk objective, which is the expected maximum reward over h independent samples.
The pass@k was championed by |CT]+21al who gave a popular unbiased estimator of the metric which derive from a new perspective (to set up our gradient estimators) as Theorem [ generalize to continuous rewards in Theorem 3 and provide additional characteri-
Related Literature
sations Corollaries and 3 Concurrently with our work (TZSM25) offered an elegant variance reduction method for the gradient of the pass@k which corresponds to the special case n k of partial sort. [CVV9) and [XDC 20) present
2
our Equation (33). Interpreting pass@k in terms elegant approximations that are rather general but less efficient in our setting, Others have pro- vided variational approximations for handling the closely related Best-of-N [CTG 24 AVAC24
and other more general [BSB+2 24] inference-time algorithms. 'he contrasting idea of training model approximate the Best-of-N prediction with a single sample was addressed by (4) Our contribution can be interpreted as a generalization and variance reduction of (TZSM25) For a general discussion of gradient estimation, variance reduction, and Monte Carlo, we recommend M Owel3.
Overview and Contributions Our theme is constructing robust estimators of the pass@k and its gradient given
samples by averaging (over all (1) subsets of size k) simple estimators that are functions of k samples. This is straightforward for binary rewards (Section 2). using the counting proof of Theorem 4 We generalize
continuous rewards using the key trick of assuming without loss of generality that the rewards are sorted, as in Section 3 Finally, we give baselining methods that require more involved derivations due to averaging over all subsets that do not include given element (to retain unbiasedness) but which boil down to the same easy-to-apply results in Section yielding our Pass-at-k Policy Optimization (PKPO). We present toy experiments Section which demonstrate the variance reduction afforded by our estimators. Finally, Section 5.2 demonstrates

---

that using our reward transformation solves more tasks and selectively optimizes pass@k through RL experiments on GEMMA2 [T] and LLAMA3.1 [GDJ 24), showcasing real-world impact.
How to Apply this Method It is easy to adapt any policy gradient algorithm use our results. of per-sample rewards for given task. For example, the 2; could be model samples of source code addressing specifie task (which should be the same for all samples), and could provide numeric score that measures how many tests the code passes, or an overall binary pass indicator, or some combination with additional stylistic or brevity terms, etc. Then in order to optimize the pass@k of Equation (I (or the continuous analog max,@k of Equation (II) we simply transform the vector of rewards using either the or the
Assume vector (22)),... ((x)))
sloo wwwummimmmooooo function of Listing | which map R R"
Binary Rewards
Given a binary reward function / - {0, 13 on the action space X, the pass@k for the model p(z[0) is the probability that at least one of k samples drawn i.i.d. is correct:
V lf(x)) =1]
pass P
(1)
E 1 - h
(1 f(x:))
(2)
i=1
where the expectation over i..dd 21, 22, ... 3h p(z
2.1 An Unbiased pass@k Estimator
An estimator for the pass@k given in CTJ-21a given l i.id samples of which correct, the estimator
nn- k
p(n, c,k) =
The following proven in CTT++1aa we give different proof that sets up our gradient estimator. Theorem 1. p(n, c, k) is unbiased estimator of the passäk.
Proof. Let 31, F2, ... Fn ~ p(x without replacement from (1, n). Then
) fa = f(x;), and be a set of & elements sampled uniformly
pass@k Ezy
ET
[I fi) iEz
(4)
Averaging over all subsets of size k recovers p:
Σ
1 [I fi)
=1
Σ
II
f.)
(5)
TET
iET
ZC (1.
IC
(1,2......
(6)
p(n, k),
(7)
where (6) holds because the sum on the r.h.s. (5) is the number of subsets of size & of the (n
incorrect elements. Since averaging in this way retains unbiasedness, this completes the proof.
We show in Corollary asymptotic variance of this estimator decreases at rate of 1/n.
no such unbiased estimator exists for n k. and in Corollary 3 that the

---

Element Weight for Various k
1.0
0.8
MI (normalized)
km
0.6
0.4
0.2
ascending sort index
The effect & has on the effective weight (1(()) of (2) for mini-batch of size n 8. This is the weight of the contribution of each sample assuming that the samples have been sorted ascending order from left to right. The horizontal axis is the sort index. For k n only the largest sample is included; for h I all samples are weighted equally. Intermediate values interpolate these extremes in a precise manner that gives rise to unbiased gradient estimation.
Figure 2:
2.2 An Unbiased pass@k Gradient Estimator
Given a mini-batch of n i.i.d. samples T1, T2, ........... Fn from p(z[0) with corresponding correctness labels f; E (0, 1}, we want to optimize the pass@kk w.r.t. the model parameters Letting c
n fi be the number of correct samples, we will demonstrate unbiasedness of the estimator
if ft =
rs Vo logp(z;(0), where Yi
(8)
p(n - 1, k-ck-
if f: = 0,
that assigns more weight correct samples, while also assigning some reward incorrect samples to encourage exploration. The following well-known results will be used to show that (8) is unbiased.
Lemma (Policy Gradients). For any absolutely continuous distribution p(x|0)
Ex=p(a)o) Vo logp(x(()) = Velem~p(z(0) [r(=)]
(9)
Corollary 1. If is constant w.r.t. both and x then Ep(z(j)) [cV log = O.
Proof. By Lemma 1 Bptajo) (cVo log logp()]0)] [o Voc 0.
can now give our first main result:
Theorem 2. v is unbiased estimator of the gradient of the pass@k:
E Pa R. ~p((()
Vopass@k.
(10)
See Appendix A.3 for proof.
Continuous Rewards
We generalize the pass@k to non-binary rewards g X + as
max,@e 0k E max ((g(::))==))
(11)
While experiments focus sloo.minus_one recommend experimenting with both estimators.

---

3.1 An Unbiased max,0k Estimator
The following estimator for the max, @k a direct analog of p: given 2 ke i.i.d. samples, assuming w.l.o.g. that the rewards = ((x;) are sorted, so that J1 92 S S the estimator is
0 n2, c,
4.g..
(12)
1
where
i 1
#4
(13)
To compute this stably we cancel factors in the binomial coefficients to ge
k
i-j
( =
n- 1
Σ 91
(14)
n-1+1
imk
Theorem 3. pl (n, c, k) is an unbiased estimator of the max @k.
Proof. The proof is similar to Theorem
!
Here we exploit the assumption that the gi are sorted, so
Σ
max
Imax et
(15)
IC (1,
IC 1.....
(16)
= p (0) (n,, c. k),
(17)
since μ4 is the number of subsets of ., at & because all subsets of size include elements that greater than equal to 9k-
i- of size k: 1, which equals ((=)) The sum starts
See Line 17 of Listing for an implementation of pl9).
3.2
An Unbiased max,@k Gradient Estimator
We propose the gradient estimator
(a)
Σ 8 V pogp((;)))
(18)
1=1
where if we assume w.l.o.g. that the gi are sorted, the 8; are a weighted combination of them,
S
miggj,
(19)
j=i
where the diagonals are and the off-diagonals are
(i-1
if --1 otherwise,
ma
(20)
0
(1:2)
if {Ú> 2)
maj
(21)
0
otherwise.
Theorem 4. (g) is an unbiased estimator of the gradient of the @k:
(a)) Vamax, @k.
(22)
We thank Ruixu Zhou of Tsinghua University for correcting in equations 14
31 and
32
5

---

Proof. The proof is analogous to that of Theorem Here we have
= (n,
log ((;
(23)
Σ
max 9j JEI
Vo log p(x;
(24)
=
E
V p(g (0)
Σ Majgi.
(25)
By assumption the are sorted, maxjez 9j = Mmaxzer Therefore maj is the number of subsets
I of {1. 2........ n that
1. are of size k.
2. have 2 as the largest element (so that we can factor out gy).
3. include i (so that 25 holds in expectation by Corollary I
Due to the second condition, the form of maj depends on whether =
The diagonals my zero if l since the largest element of any subset of size k is at least k. If i k then we fix i and are left with
elements from which to choose l - which we can do
ways in line with 20
The maj for i + are obtained by fixing i and j leaving
elements 1, 2, 1.....j - from which to choose & which do (2-2) ways in line with (1).
Theorem 5. $1, S2, be computed in total time O + logn).
See Appendix A.4 for the proof and Line 36 Listing I for an implementation based on it.
4 Variance Reduction
4.1 Leave-One-Out Baseline for the Simple Case
A popular variance reduction method [MRM22) Owel3 GYZ+25 for point-wise rewards g(x) subtracts the mean the leave one out (LOO) rewards within each mini-batch T T2, Zn.
(0o) (x:) = g(x;)
2 g(x)).
(26)
Since the subtracted part does not depend æ,, by Corollary ! this retains unbiasedness.
4.2 Leave-One-Out Baseline for max,@k
Baselining the s; of 19 in this way introduces bias, however, as each s; depends on all 01, ... -n-
We instead apply LOO to the following form of s; that follows from Theorem and the proof thereof:
Σ
max
(27)
JEZ
Z
NET
IC(1,2.....1
S(i, k. n)),
(28)
We can then retain unbiasedness by excluding i from the baseline, by defining
(loo) = S((i, ht {1, (1,
S(j k. (1)
(29)

---

Theorem 6. 81 ,
(1oo)
(loo)
(loo)
can be computed in total time O + logn).
Proof. Given ( it is sufficient to consider computing, for i 2,,,.... .,,.....7..
(
,
.., (1)
(30)
By assuming w.l.o.g, ascending ordering of g(z;), excluding the first index does change the ordering of the remaining indices. The first term therefore
((4)
Σ
mm m4-1,5(i (x:)
(31)
2m
where 31 follows from 49
From 49
obtain for Si < n the left to right recursion
1+1
= b
+
("x²)
(o(2,) g(z++1)) )) mm + m: 2))
(32)
Similar arguments to the proof of Theorem therefore imply the same time complexity.
Listing implements s; (loo) using the recursion in the above proof.
Line 52
4.3
max, 1) Leave-One-Out Baseline for max,
The baseline b/() is an average of the max.0k estimates over sets of size k. For number of samples n equal to k; there are no such subsets to construct the baseline. (TSSM25) recently overcame this issue for the specific case n k by using max,@(k 1) as the baseline statistic. We generalize
their approach to k < n and averaging over all subsets by defining similarly Equation
(27
(lon
Σ
(max JEI
95
max 9b) BEZ
(33)
1EI
IC(1......))
Averaging smaller but more numerous subsets in the baseline reduces variance but introduces bias (in the baseline, not 8; (1oo-1)) Given our previous results it is straightforward to show
(100-1) (loo-1). (1oo-1) cun be computed in total time O(k + n log n).
Theorem 7. $1 822
Proof: By the linearity of the expectation we can split the two terms in the parentheses of Equa- tion (33) into two separate sums. The first summation is by definition simply 8; of Equation (1 The (negation of the) second summation can be computed efficiently using
Σ
max BEJ
Σ
max Ib bEB
k
b
(K-1)
(34)
n(k
JE
B 1
[1.
BC(1.2...,.)1
where the final equality follows with little algebra from Equation 288 and Equation (30p.
Listing implements s! (Ino-1) using (34;; Figure
compares 8,, (lon)
and s (loo-1] 1)
Experiments
5.1 One-Dimensional Toy Example
We start with a policy that Gaussian with a fixed standard deviation and mean parameter wish to learn, so that r. ~ N/(8, 0.1). We set the raw reward to be
05051
g(x) =
(35)
otherwise.
The optimal policy under the max,@k reward varies with & (see Figure estimators is compared in Figure where (100-1) is the strongest.
The variance of our

---

Table 1: Results for GEMMA2-9B on the MATH benchmark.
GEMMA2-9B
k.eval=1
k_eval=2
k_eval=4
k_eval=8
k_eval=16
k_opt=1
22.24 + 0.50
25.35 + 0.55
30.73 = 0.59
37.08 + 0.64
42.59 + 0.68
koopt=2
21.46 0.51
28.61 0.56
32.92 = 0.61
39.59 + 0.66
45.34 = 0.70
k_opt=4
21.25 + 0.53
27.15 + 0.58
34.93 = 0.63
41.71 0.69
47.05 = 0.74
kopt=
20.69 0.56
26.78 0.60
33.68 0.66
42.62 + 0.72
48.37 0.77
TZSM25
19.48 0.61
25.41 0.67
31.17 = 0.73
39.34 + 0.79
44.82 + 0.83
EntropyReg
20.85 + 0.58
26.05 + 0.64
32,48 + 0.70
38.21 + 0.76
43.95 + 0.81
Table 2: Results for LLAMA3.1-8B5 on the MATH benchmark.
LLAMA3.1-8P
keevll=1
k.eval=2
k evvall=4
k keval=
k keeal116
k_opt=1
51.15 + 0.61
51.82
0.64
0.68 53.69
55.41 + 0.72
56.83 + 0.76
k opt=2
49.72 0.62
53.51 + 0.66
55,45 + 0.70
57.23 = 0.74
58.71 = 0.78
k <.opt=4
49.18 + 0.64
52.20 0.68
57.83 + 0.72
58.47 + 0.77
59.28 + 0.81
48.63 0.67
52.14 0.71
56.28 0.75
59.04 = 0.80
61.88 + 0.84
TZSM25
48.21 + 0.70
50.93 0.75
54.38 0.80
57.11 = 0.85
58.55 + 0.90
EntropyReg
48.51 + 0.68
51.95 + 0.73
55.33 + 0.78
56.95 + 0.83
58.18 + 0.88
Table 3: Results for GEMMA2-9B the Coding benchmark.
GEMMA2-9B
k_eval=1
k_eval=2
k_eval=4
k_eval=8
k_eval=16
k opt=1
37.71 0.60
42.03 0.65
48.19
0.69
55.07 = 0.75
60.98 = 0.79
k opt=2
36.84 + 0.61
46.56 + 0.67
52.68 + 0.72
59.73 + 0.78
65.86 + 0.84
k opt=4
36.49 0.63
44.95 0.69
57.09
0.76
63.64 + 0.83
69.51 + 0.88
k_opt=8
35.75 + 0.67
44.41 + 0.73
55.08 + 0.80
65.56 + 0.88
71,91 + 0.94
TZSM25
34.36 - 0.72
42.81 + 0.78
52.36 = 0.86
61.07 + 0.95
66.41 = 1.01
EntropyReg
35.91 + 0.70
43.75 0.76
53.28 = 0.84
60.13 + 0.93
65.29 = 0.99
Table 4: Results for LLAMA3,1-8B on the Coding benchmark.
LLAMA3.1-80
k_eval=1
k_eval=2
k_eval=4
k_eval=8
k_eval=16
k_opt=1
67.38 + 0.72
67.45
0.76
69.22 + 0.80
71.11 + 0.84
72.84 + 0.88
k_opt=2
64.91 + 0.73
69.73 + 0.78
72.03 + 0.82
74.08 + 0.87
75.89 + 0.91
k_opt=4
64.25 + 0.75
68.47 0.80
74.67 + 0.85
75.01 + 0.90
77.75
+ 0.95
k_opt=8
63.57 + 0.78
68.39 + 0.83
72.84 + 0.88
76.82 = 0.94
79.33 + 0.99
TZSM25
62.77 + 0.82
66.86 + 0.88
70.83 + 0.94
73.95 + 1.01
75.47 + 1.06
EntropyReg
63.91 0.80
67.85 0.86
71.78 0.92
72.99
0.98
74.31 + 1.04
5.2 RL on Open Source LLMs
We demonstrate promising RL results with the 2B and 9B parameter variants GEMMA2 and the 8B parameter variant of LLAMA3. 1 real-world problems in MATH HBK 21,. code generation 21] ( and the easy public subset of ARC-AGI-1 [KKL225 The latter
TRP 24
a challenging reasoning task-set even for state-of-the-e the-th--trrttt models much larger than GEMMA2.
For GEMMA2-2B we use v5litepod-128 [Goo] which needs around hours per 1000 train- ing steps. Each RL training run (SWD+17 17] involves sampling fixed n number of completions The re-
{zx ) for given prompt at a given training step. For our experiments, we set wards are computed for every completion using reward function g(.). We transform these rewards ( using our unbiased estimator (loo-1) of (33). which we favour due to Figure 4 and which we refer as PKPD. We repeat the training for a selection of kopt, thus optimizing a different op each time. Since copt =1 leads to no reward transformation, this is our baseline (al-

---

Table 5: Results for GEMMA2-99 on the ARC-AGI- benchmark.
GEMMA2-9B
Cumulative Solve Rate
pass@1
pass@16
k_opt=1
12.00 = 04.33
02.00 = 01.69
08.18 + 04.00
k_opt=4
82.33 = 04.
22.00 = 02.00
38.18 04.67
k.opt=8
84.14 + 04.67
26.67 = 02.50
44.50 + 04.33
[TSMM3]
22.00 = 04.44
06.00 = 02.67
10.16 + 04.57
EntropyReg
24.67 + 04.50
04.00 = 02.33
08.89 04.89
Table 6: Results for LLAMA3.1-8B the ARC-AGI-1 benchmark.
LLAMA3.1-8B
Cumulative Solve Rate
pass@1
pass016
k_opt=1
22.00 + 04.18
03.33 + 02.00
08.00
04.50
k_opt=4
87.17 + 04.14
24.33 + 02.33
42.00 + 04.16
k.opt=8
88.89 = 04.33
29.67 = 02.67
43.13 04.67
TZSM25)
36.00 + 02.50
08.00 + 04.00
18.00 04.89
EntropyReg
28.00 + 04.44
08.00 + 02.50
14.67 + 04.44
though we use basic LOO mean centering of Equation (26). without which the training diverges). For each run, we measure pass@k'val for every koval E 21, 4, 8, 12, 16} at each step. Additionally, also track model entropy and cumulative solve rate during training. The latter is defined as the fraction of tasks from the task-set for which the model has sampled a correct solution at least once; this is critical metric that reflects the success of the model's exploration and measures its ability to find novel solutions.
Entropy regularization baseline In addition to our PKPD and the special case thereof of [TZSM25) we also add the entropy regularization baseline, which is PPO with an additional entropy term the objective. We give this baseline an arguably unfair advantage by performing a small sweep over the values 0.001, 0.005, 0.01, 0.05, 0.1 for the entropy cooefficieenen for each (model, benchmark) pair and only report the best result EntropyReg.
5.2.1 Choosing oopt selectively optimizes pass@keval and solves more tasks
We use the training split of Hendrycks MATH HBK 21] which contains 12,000 problems as our task set. Figure 66 shows that higher &ept in our transformation leads consistently higher cumulative solve rate throughout training, as well as a higher entropy. By optimizing pass@k instead of 1 the model appears to better utilize the exploration budget thus finding more solutions.
In Figure 7 we compare pass@keval across our runs (kopt E (1,4,8) for various keval, We find the best pass@k when hope = geval (or kopt is closest to geval among available kopt). Non- transformed rewards optimize pass@1, leading to sub-optimal @ for boval 1, and the deficit worsens as jeval increases. Thus, our experiments also demonstrate that setting kopE:= Leval in our transformation suffices to optimize @vs for keval n. This generalizes the already powerful result of [ZSM255] by alleviating the coupling that restricts to optimizing either pass@n pass@1. 1. In other words, since RL training of L.MMs typically samples large batch (n >> 1), failing to use our transformation results sub-optimal pass@k performance, especially for modest values of k.
- n, the variance our estimator increases as there are fewer subsets in 33) (see Figure 4h. We presume this is why 1) gains of gopt = 8 over hopt = 4 are more prominent when geval 112, 16} than when geval 8. That is, when eeal is further away from kopt E (4, than when it is closer, and 2) the special case hopt of TZSM25 struggles to optimize the passon.
As hopt
5.2.2 PKPO robustly improves pass@k on held out evaluations
Tables above (and Tables 8 in the appendix) present performance on held-out sets for two
tasks. We report the mean and standard error based on three runs with different random seeds. For math, we train on the train split and evaluate on the test split of Hendryeks MATH [HBK*2 21] To

---

pass@k--v A-va
pass@k geral -
pass@ker) Aeral
16
0.22
T23M22
pasigAm
0.38
T2SM25
pasoggeu
0.44
T2SMES
0.18
0.42
0.16
0.14
0.34
1000
2000
3000
1000
2000
3000
1000
2000
Step
(a) passáe (k 1)
(b) pass (k
(c) pass@11 (k 16)
Figure 3: Annealing kops during PKPO training improves pass@Aeve without sacrificing pass@1.
For a.annealed we train with kopt up to step 1500 and kopt = thereafter.
evaluate coding, we use MBPP AON+21 for training and evaluate on HUMANEVAL CTJ 21b MBPP has multiple unit tests per problem and hence we use this not only proxy for additional benchmarks but also showcase our handling of continuous reward function (% unit tests passed).
5.2.3 Improving pass@k without sacrificing pass@1
Figure demonstrates that PKPO can use any arbitrary kopt < n,, this allows varying kops over the course of training good effect. We show simple annealing procedure which starts training with high kopt = 8 and reduces it to kops = 1 after 1500 steps. This trains the model to initially prioritize exploration (optimize pass@k) and then consolidate the single-sample policy (optimize pass@1). This switch is apparent in Figure 3a at step 1500 where the slope of jaunealed changes. While traditional methods like [TZSM25) suffer from trade-off between pass@k and pass@1, we get final model which has higher pass@keva for all goval> 1 with no sacrifice in pass@1.
5.2.4 PKPO is essential for learning on hard problems
Figure 8 shows the limitation traditional pass optimization through RL on an especially chal- lenging task-set. We the easy subset of ARC-AGI-1 [CKKL25). We observe that conventional pass@ optimization stalls. However, our pass@k approach unblocks learning, and results in higher pass@k across all eevl including -eval= 1. Furthermore, see higher oopt leads to more effec- tive and faster learning. This is likely because the benefits of prioritizing joint utility over individual sample utility more prominent
harder task-set.
Tables 5 [ more extensive experiments on ARC-AGI-1. We make an 80:20 train:test split of the same easy subset before and report the cumulative solve rate on the train set and pass@k rate on the test set. We train to saturation (no change cumulative rate for steps), and again use three random restarts to provide standard errors. By encouraging exploration in direct and stable manner, our method unblocks learning unlike other methods. Entropy Regularization does indeed sacrifice pass@1 1 and slightly improves pass&k by promoting exploration, but it is hard to tune, and significantly outperformed by our method. Moreover, it has no explicit way to optimize for a specific k. eval. [TZSM25] targets the same objective PKPO, but couples the minibatch size k and thereby incurs higher variance than PKPO with l < n.
Conclusions and Outlook
In RL training with multiple independent samples per task, optimizing the pass@k maximizes the expectation of the best reward in the set of samples, rather than the average one. This preserves model output diversity, which leads solving more problems and ultimately yields stronger policies. We provide drop-in replacements for more traditional RL reward transformations that robustly and efficiently optimize the pass@k. This work can be extended in various ways, such as other inference- time search algorithms, and to more sophisticated baseline techniques.
10

---

References
[AAA*23] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anad- kat, et al. Gpt-4 technical report. arXiv preprint Saliv:2303.08774 2023.
[ABC*21] Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant laboratory for alignment. arXiv preprint arXiv:2112.00860
2021. 2021.
[Ant] Anthropic. Claude 3.5 Sonnet.
https://www. anthropic. com/newssclaude=
[AON 21] Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carric Cai, Michael Terry, Quoc Lc, and Charles Sutton. Program synthesis with large language models, 2021.
[AT]
DeepMind AlphaProof and AlphaGeometry Teams. AI achieves silver-medal standard solving international mathematical olympiad problems.
https://tinyurl. com/ alphaproof
[AVAC24] Afra Amini, Tim Vieira, Elliott Ash, and Ryan Cotterell. Variational Best-of-N align- ment. arXiv preprint arXiv:2407.0603. 2024.
[BSB+24] Ananth Balashankar, Ziteng Sun, Jonathan Berant, Jacob Fisenstein, Michael Collins, Adrian Hutter, Jong Lee, Chirag Nagpal, Flavien Prost, Aradhana Sinha, al. Infalign: Inference-aware language model alignment. arXiv preprint 1rX1v:2412.19792, 2024.
[CKKLI25] Francois Chollet, Mike Knoop, Gregory Kamradt, and Bryan Landers. Are prize 2024: Technical report, 2025.
[CLB*17] Paul Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
[CTG+24 24] Yinlam Chow, Guy Tennenholtz, Izzeddin Gur, Vincent Zhuang, Bo Dai, Sridhar Thi-
agarajan, Craig Boutilier, Rishabh Agarwal, Aviral Kumar, and Aleksandra Faust. Inference-aware fine-tuning for Best-of-N sampling in large language models, 2024.
[CTJ 221a] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray. Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Niko- las Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saun- ders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford. Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained on code. arXiv, 2021.
[CT]+21b] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray. Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss William Hebgen Guss, Alex Nichol, Alex Paino, Niko- las Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saun- ders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Mathew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code. 2021.
[CTO+25] Yuri Chervonyi, Trieu H. Trinh, Miroslav Olsak, Xiaomeng Yang, Hoang Nguyen, Marcelo Menegali, Junehyuk Jung, Vikas Verma, Quoe Le, and Thang Luong. Gold- medalist performance in solving olympiad geometry with alphagcometry2, 2025.

---

[CTV19] Marco Cuturi, Olivier Teboul, and Jean- Phin-philippe Vert. Differentiable ranking and sort- ing using optimal transport. In Advances in Neural Information Processing Systems, volume 32, 2019.
[DLJ*24] Shihan Dou, Yan Liu, Haoxiang Jia, Limao Xiong, Enyu Zhou, Wei Shen, Junjie Shan, Caishuang Huang, Xiao Wang. Xiaoran Fan, et al. Stepcoder: Improve code generation with reinforcement learning from compiler feedback. arXiv preprint arXiv:2402.01.399 2024.
[GDJ+24] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Ka- dian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mi- tra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Ma- hajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Fran- cisco Guzmán, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Ander- son, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Ilyann Zarov, Imanol Ar- rieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Al- wala, Karthik Prasad, Kartikeya Upasani. Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala- Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan. Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji. Ning Zhang, Olivier Duchenne, Onur Çelebi. Patrick Alrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng, Prajiwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari. Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hos- seini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Scohyun Sonia Kim, Sergey Edunoy, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do Vish Vogeti, Vitor Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang. Xiaofang Wang, Xiaoqing Ellen Tan, Xide Xia. Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaclle Goldschlag, Yashesh Gaur, Yasmine Babaci, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li. Yuning Mao, Zacharie Delpierre Coud- ert, Zheng Yan, Zhengxing Chen, Zoe Papakipos. Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Fein- stein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf, Arka- bandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yaz-
Albiero,
12

---

dan, Beau James, Ben Maurer, Benjamin Leonhardi. Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape. Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang. Changkyu Kim, Chao Zhou, Chester Hu, Ching- Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Da- vide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Duc Le. Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric- Le, Erik Brinkman, Esteban Ar- caute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Filippos Kokki- nos, Firat Ozgenel, Francesco Caggioni. Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Grant Herman, Grigory Sizov, Guangyi, Zhang. Guna Lakshminarayanan, Hakan Inan, Hamid Sho- janazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, Jean-Baptiste Gaya, Jeff Mar- cus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jes- sica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang. Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraragha- van, Kelly Michelena, Keqian Li, Kiran Jagadeesh, Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov. Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang. Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev. Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Ran- gaprabhu Parthasarathy, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang. Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lind- say, Sheng Feng, Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shugiang Zhang, Shuqiang Zhang, Sinong Wang. Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Su- darshan Govindaprasad, Sumit Gupta, Summer Deng, Sungmin Cho, Sunny Virk, Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler. Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Tim- othy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihaileseu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang, Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang. Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary Vito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, Zhiwei Zhao, and Zhiyu Ma. The llama 3 herd of models, 2024.
[Goo] Google. Google cloud platform. https://cloud, google. com
[GYZ+25] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qi- hao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et Deepseek-rl: Incentivizing reasoning capability in L.LSs via reinforcement learning. arXiv preprint 2501 /2948, 2025.
13

---

[GZC*24] Jonas Gehring, Kunhao Zheng, Jade Copet, Vegard Mella, Quentin Carbonneaux, Taco Cohen, and Gabriel Synnaeve. Rlef: Grounding code 1lms execution feedback with reinforcement learning. arXiv preprint arXiv:2410.02089, 2024.
[HK [21] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math datasct, 2021.
[Hoe48] Wassily Hoeffding. class of statisties with asymptotically normal distribution. Annals of Mathematical Statistics, 19:308-333 -334,,34,,400 1948.
[4YM+24] Arian Hosseini, Xingdi Yuan, Nikolay Malkin, Aaron Courville, Alessandro Sordoni, and Rishabh Agarwal. V-star: Training verifiers for self-taught reasoners. arXiv preprini
arXiv22402.06457, 2024.
[Ko150] A. N. Kolmogorov. Unbiased estimates. Izvestiya Akademii Naak SSSR. Seriya Matem- alicheskaya, 14(4):303-326, 1950.
[KZA*24] Aviral Kumar, Vincent Zhuang, Rishabh Agarwal, Yi Su, John Co-Reyes, Avi Singh, Kate Baumli, Shariq lqbal, Colton Bishop, Rebecca Roelofs, et al. Training language models to self-correct via reinforcement learning. arXiv preprint arXiv:2409. 12917, 2024.
[LCC*22] Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Remi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, al. Competition-level code generation with alphacode. Science, 37866624):0092-1097, 2022.
[Leh98] George Lehmann, E. L. Casella. Theory of Point Estimation. Springer, 2nd edition, 1998.
LKBB+23 Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference Learning Representations, 2023.
[LWG+22] Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Chu Hong Hoi. Coderl: Mastering code generation through pretrained models and deep reinforce- ment learning. Advances in Neural Information Processing Systems, 35:21314-21329 2022.
[MRFM20] Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. Journal of Machine Learning Research, 21(132):1-62, 2020.
[3]W+23] Theo X Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar-Lezama. Is self-repair a silver bullet for code generation? arXiv preprint wxXiv:2306.09866 2023.
[3] Art B. Owen. Monte Carlo theory, methods and examples. https://artoven..5 httpr//oeen.su..ddiiiimmmmmm// 2013.
[PYTG25] Nilay Pande. Sahiti Yerramilli, Jayant Sravan Tamarapalli, and Rynaa Grover. Marvl- qu: A benchmark for mathematical reasoning over visual landscapes. arXiv preprint arXiv:2508. /7/80, 2025.
[RNS 18] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training, 2018.
[SDH+24] 24] Pier Giuseppe Sessa, Robert Dadashi, Leonard Hussenot, Johan Ferret, Nino Vieillard, Alexandre Rame, Bobak Shahriari, Sarah Perrin, Abe Fricsen, Geoffrey Cideron, Sertan Girgin, Piotr Stanczyk, Andrea Michi, Danila Sinopalnikov, Sabela Ramos, Amelic He- liou, Aliaksei Severyn, Matt Hoffman, Nikola Momchev. and Olivier Bachem. BOND: aligning LLMs with Best-of-N distillation. CoRR, abs/2407.14622 146222. 2024.
[SJTR23] Parshin Shojace. Aneesh Jain, Sindhu Tipimeni, and Chandan K Reddy. Execution- based code generation using deep reinforcement learning. arXiv preprint 1rXiv:2301.13816 2023.
[SWD+17] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Prox- imal policy optimization algorithms, 2017.
14

---

[Tea23] Gemini Team. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.1180,, 2023.
(TRP=24) Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhupatiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexan- dre Ramé, Johan Ferret, Peter Liu, Pouya Tafti, Abe Friesen, Michelle Casbon, Sabela Ramos, Ravin Kumar, Charline Le Lan, Sammy Jerome, Anton Tsitsulin, Nino Vieillard, Piotr Stanczyk, Sertan Girgin, Nikola Momchev, Matt Hoffman, Shantanu Thakoor, Jean-Bastien Grill, Behnam Neyshabur, Olivier Bachem, Alanna Walton, Ali- aksei Severyn, Alicia Parrish, Aliya Ahmad, Allen Hutchison, Alvin Abdagic, Amanda Carl, Amy Shen, Andy Brock, Andy Coenen, Anthony Laforge, Antonia Paterson, Ben Bastian, Bilal Piot, Bo Wu, Brandon Royal, Charlic Chen, Chintu Kumar, Chris Perry, Chris Welty, Christopher A. Choquette- Danila Sinopalnikov, David Wein- berger, Dimple Vijaykumar, Dominika Rogozińska, Dustin Herbison, Elisa Bandy, Emma Wang. Eric Noland, Erica Moreira, Evan Senter, Evgenii Eltyshev, Francesco Visin, Gabriel Rasskin, Gary Wei, Glenn Cameron, Gus Martins, Hadi Hashemi, Hanna Climezk-Pluccissie Harleen Batra, Harsh Dhand, Ivan Nardini, Jacinda Mein, Jack Zhou, James Svensson, Jeff Stanway, Jetha Chan, Jin Peng Zhou, Joana Carrasqueira, Joana Iljazi, Jocelyn Becker, Joe Fernandez, Joost van Amersfoort, Josh Gordon, Josh Lipschultz, Josh Newlan, Ju yeong Ji, Kareem Mohamed, Kartikeya Badola, Kat Black, Katie Millican, Keelin McDonell, Kelvin Nguyen, Kiranbir Sodhia, Kish Greene, Lars Lowe Sjoesund, Lauren Usui, Laurent Sifree, Lena Heuermann, Leticia Lago, Lilly McNealus, Livio Baldini Soares, Logan Kilpatrick, Lucas Dixon, Luciano Mar- tins, Machel Reid, Manvinder Singh, Mark Iverson, Martin Görner, Mat Velloso, Ma- teo Wirth, Matt Davidow, Matt Miller, Matthew Rahtz, Matthew Watson, Meg Risdal, Mehran Kazemi, Michael Moynihan, Ming Zhang, Minsuk Kahng, Minwoo Park, Mofi Rahman, Mohit Khatwani, Natalie Dao, Nenshad Bardoliwalla, Nesh Devanathan, Neta Dumai, Nilay Chauhan, Oscar Wahltinez. Pankil Botarda, Parker Barnes, Paul Barham, Paul Michel, Pengchong Jin, Petko Georgiev, Phil Culliton, Pradeep Kuppala, Ramona Comanescu, Ramona Merhej. Reena Jana, Reza Ardeshir Rokni, Rishabh Agarwal, Ryan Mullins, Samanch Saadat, Sara Mc Carthy, Sarah Cogan, Sarah Perrin, Sébastien M. R. Aroold, Sebastian Krause, Shengyang Dai, Shruti Garg, Shruti Sheth, Sue Ron- strom, Susan Chan, Timothy Jordan, Ting Yu, Tom Eccles, Tom Hennigan, Tomas Ko- cisky, Tulsee Doshi, Vihan Jain, Vikas Yadav, Vilobh Meshram, Vishal Dharmadhikari, Warren Barkley, Wei Wei, Wenming Ye, Woohyun Han, Woosuk Kwon, Xiang Xu, Zhe Shen, Zhitao Gong, Zichuan Wei, Victor Cotruta, Phoebe Kirk, Anand Rao, Minh Giang. Ludovic Peran, Tris Warkentin, Eli Collins, Joelle Barral, Zoubin Ghahramani. Raia Hadsell, D. Sculley, Jeanine Banks, Anca Dragan, Slav Petrov, Oriol Vinyals, Jeff Dean, Demis Hassabis. Koray Kavukcuoglu, Clement Farabet, Elena Buchatskaya, Se- bastian Borgcaud, Noah Fiedel, Armand Joulin, Kathleen Kenealy, Robert Dadashi, and Alck Andreev. Gemma 2: Improving open language models at practical size, 2024.
[TZSM25] Yunhao Tang, Kunhao Zheng, Gabriel Synnaeve, and Rémi Munos. Optimizing lan- guage models for inference time objectives using reinforcement learning. arXiv preprint arXiv:2503. 19595, 2025.
[WBZ+21] Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint wrXiv:2109.0166 2021.
[Wil92] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist. reinforcement learning. Machine learning, 8:229-256. 1992.
[WSL+ 24] Yangzhen Wu, Zhiqing Sun, Shanda Li, Sean Welleck, and Yiming Yang. Inference scaling laws: An empirical analysis of compute-optimal inference for problem-solving with language models. arXiv preprint arXiv:2408.0071
2024.
[XDC*20 20] Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. In Advances in Neural Information Processing Systems, volume 33, 2020.
[Y]G+23] Kaiyu Yang, Aidan Swope, Alex Gu, Rahul Chalamala, Peiyang Song, Shixing Yu, Saad Godil, Ryan J Prenger, and Animashree Anandkumar. Leandojo: Theorem prov-
15

---

ing with retrieval-augmented language models. Advances in Neural Information Pro- cessing Systems, 26221573-21612,200 2023.
[YTC+ Zishun Yu, Yunzhe Tao, Liyu Chen, Tao Sun, and Hongxia Yang. B-coder: Value-based deep reinforcement learning for program synthesis. arXiv preprint arXiv:2310.03173, 2023.
16

---

A Additional Theoretical Statements and Proofs
A.1 Statement and proof that n 2 A samples are required to unbiasedly estimate pass Kk
This result is direct consequence of a well-known theorem concerning the unbiased estimability of parametric functions for the Bernoulli distribution.
Yn Theorem 8 (Kolmogorov ) Let Yi....,Y,, be i..d. Bernoulli random variables with suc- cess probability [0, 1]. function p(p) is unbiasedly estimable from this sample and only if can be expressed as polynomial in of degree at most n.
sketch of a proof Theorem can be found in Lehmann and Casella Leh98
Corollary 2. Given a sequence of i..d model samples I,, D2, ... Ena the pass@k k is unbiasedly estimable if and only if 2 k.
Proof. It is sufficient to consider a single a fixed and observed correctness function f. so that the independence of the z; implies the independence of the correctness events [f(x)) = 1]. Let p [[[f(x;) 1]] be the prohability that any single sample is correct. The pass@k is defined as the complement of the probability that all k samples are incorrect, which for the specific assumptions adopted in this proof is 1 - (1 p).. Because this expression is polynomial in of degree k, the result follows immediately from Theorem g
A.2 Characterization of the Variance
Our proof of Theorem iidntitiess the pass@k estimator p(n, c, k) U-statistic. To characterize its variance, we apply Hoeffding's 's asymptotic theory.
Theorem 9 (Hoeffding [Hoe48)) ). Let XI...., Xn be independent and identically distributed random variables with distribution F. Let h(x1....., rk) be a symmetric kernel with E[M(X,..... x+) < co. Define the parameter p. Ex(h(X1..... Xx)] and the U-statistic:
Un
12
Σ .(X........ X.
(36)
k
1 ChE
Let - proved that if Ci Var(h, (X,)) ))) 0, then
= E[h(r. ....... X1 1 be the projection of the kernel onto a single variable. Hoeffding
0o:
vn(U, - n1 4 N(00
(37)
In the standard application of pass@k we evaluate the estimator on specific problem defined by a prompt and a correctness oracle. While the true pass rate v is unknown to the observer, it is fixed property of the model-problem pair. Consequently, the correctness outcomes of the generated samples are i.i.d. conditioned on the problem.
The following lemma derives the variance parameter under this conditioning. We abuse the notation by allowing the X, to denote correctness.
Lemma 2 (Conditional Variance of the Projection). Fix problem instance such that the correct- nesses X, are i.i.d. Bernoulliy). For the pass@k kernel h(r1,...,, max(x1,..... 24), the variance of the first-order projection is:
(i(u,k) v(1 22)
(38)
Proof. The projection hy (x) the expected value of the kernel given the first sample is fixed x, while X2,..., X remain random variates drawn from Bernoulli(y).
h (x) E[max(#, X Xx)].
(39)
We evaluate this for the two possible realizations of z:
17

---

1. Case z (Success): The maximum is 1 regardless of the remaining samples.
hu (1) 1.
2. Case 0 (Failure): The maximum is O if and only if all remaining k samples fail. Since the remaining samples are i..d. with failure probability (1 1/),
h (0) 1 レ)
The projection h1 (X1) is thus binary random variable taking value h1 (1) with probability v and h, (0) with probability 1
v, so that
Ci = v(1 v) (h1(1) h0))))²
=v(1 -v)
1/) (1 - (1-11-1)~-1)
101 - 12) ((1 (1-1)*11)²
= v01 z)ak-1
We can substitute this explicit form back into Hoeffding's general result.
Corollary 3. For a fixed problem with pass rate v, as n + oc, the asymptotic variance of the estimator p(n, e, k) is:
Var(p) 1 [R²u(1 2)
(40)
n
A.3 Proof of Theorem 2
Although Theorem 2 is special case of Theorem we include both because the following proof uses different approach from that of the more general statement, and is arguably the easier of the two.
Proof. By Lemma 1 the gradient V Vopass@k has the unbiased estimator
p(n, k) V
log p(x;10)
(41)
"
Σ -|== IC
1 II f.)
Σ
V a log p(x,10)
(42)
FT
(1,2..)))))))
my log
(43)
where 42 substitutes the l.h.s. of 5 m, is the number of subsets I of {,,2,... nn that
1. are of size k;
2. contain at least one correct element, so that - (1 = 1.
3. contain i, so that (43) holds in expectation by Corollary I
Due to the second condition, m, therefore equals one of values, which we denote by (1) and (0) depending whether f,
f. 0, respectively.
fa = 1 then all subsets that include also include at least one correct element (i itself), so that m (1) just the number of subsets of size & of 21, n.
that include i, which equals the number
of subsets size of (1, 21,2..... 1}:
m 11)
n
(44)
1
18

---

If ft O then we assume w.l.o.g. that i = so that m0) is the number of subsets of size k 1 of (1,2. ........ n } with least one correct element.
(l))
E
I- 13) JEJ
=
n2 11-1 p(n 1,c,- 1).
(45)
JS IJ|=k-1
where we again used (5) this time to get expression in terms of p. Using m and ml)) can compute y(()) and r(1) using 43
( m (1)
k
(46)
and
(0)
(9)
n- k-
1
p(n
c. 1-
1)
p(n 1, 1),
(47)
1
n
in line with
8).
A.4 Proof of Theorem
Proof. The vector (81, ................................ 8n) can be written Mg where have introduced
(g(x1) ((. 9(x,)) well as the matrix M with
1. diagonal elements ma given by (20
2. upper diagonals mij for < j given by 21 which is independent of i,
3. lower diagonals may for > j equal to zero
Because of the structure of M, have that
8n
Man g(En);
(48)
and, for 1 S < ni the right to left recursion
(g(z))m + g(2:4)) (m4t---)))
S; 8i+1
(49)
The ratios of maa 11,++1 and M4+144+1 divided by (2) can be simplified by cancelling factors in the binomial coefficients and writing the remaining factors as a product ratios similarly to (14). for total cost of (nk)); this computation can be further simplified by noting that the required ratios can be lazily computed in sequence (for example to obtain M,41,i+1 from m;)) at a cost of O(1) after computing the first at cost of O(k), giving total cost of O(k+n) The additional O(n log comes from assuming the i are sorted in increasing order of g(x:).
19

---

B Implementation
def _m. normed (N: int, K: int, int, j: int) -> float:
1f 1 -- J and >- K-1:
return
K (N-K+1)
m up-prod mp arange (i-K+2, i+1) up. arange (N-K+2, x+1))
elif j i and j >= K-1 and K >= 2:
return
K (N-K+1) (K-1) N
10
up. prod arange (j j j) / up. arange (N-K+2, R))
11
return 0
def diagona (N: int, int) np ndarray:
15
return mp. array n[_m_ (N, K, i, i) for in range (N)
16
def rho (g: np. ndarray, K: int) -> float:
See Equation (12) "..
return (ortsort(g)
_m_diagonal (g), K)) sum
20
21
def dellaaa int, K: int, i: int) -> float: return (M_normed K, i, i+1) _m_normed K, it1, i+1)
def _ellts (N: int, int) up ndarray:
25
return np array (N-1, K, i) for i in range (N 2)))
def _orteed poported_apllyyooe apply (func: Callable) -> Callable:
def inner (x: пp adaraay *args, **vargs)) np. ndarray:
29
i_sort np. arggort(( func_x = up zepos
30
[in_xxi_sort]] = func sxnosx((_ *args, **kvargs)
func_x
return inner
34
e_sorted_aplly
def (g: np ndarray, K: int):
See Equation 19)
38
(g)(g)
29
c g __dianoonallammmmaa (N, K)
cN: += g[i:] _deltas (N+1, K)
np. cumsum ( (:::
43
8 s_sorted_appl
def b(g: ndarray, K: int) -> ndarray: N len(g)
mp. arange B)) astype aatype(float) a [1:] _eelta (N, K) np arange (1, N-1)
(m__diaoonal (N-1,
48
c1 ap.array ([(M .g [(:]). sum]) ()))
c2 (g[:-1) [[:]]
return np
(np. concatenate ((ci, c2)))
52
def sloo (g: np ndarray, K: int) -> np ndarray:
*** See Equation 29) aen
return (g, K) (g. K) (len (g) 1)
56 def sloo_minus www.s min
one (g: : mp ndarray, K: int) -> np. ndarray:
57
*** See Equation (33) uen
return ( K) (g. K-1) * K / (K-1) / len (g)
I Listing Python reward batch transformations, Functions with names that begin with an underscore helpers, while the remaining four functions rho, s, sloo and sloo minus one implement pو9),
8i, 8, (100) and 8, (loo-1), respectively. For simplicity this implementation costs O(nk: n log n)
reducing this to Olk + n log nz) would require optimizing _deltas and m.diagonal.
20

---

C
Additional Figures
Empirical Variance the Gradient Estimator
Variance of Estimated Gradient of pass, for k
+ all subsets baselined sllel subsets baseline
102
nalve partitionea baselines ed
nalve partitioned baceline loo minus one partitioned minus
101
10"
10
15
20
25
30
35
Number of model samples
max,@k Figure 4: The variance of different estimators of the gradient of max, @k with k = 4 for the one- dimensional problem depicted in Figure [ the location x 1. Each data-point is the sample variance of 10,000 independent unbiased gradient estimates (lower is better). The horizontal axis denotes the number of samples used to construct each of the 10,000 estimates. We compare the
following methods:
all subsets baselined: ((oo)) our novel estimator of Equation 29 that analytically sums over all subsets of size of the samples with unbiased baseline method that subtracts for each element i the mean of the estimator over all subsets of size k that do not include i.
novel estimator of Equation (19) that analytically sums over all subsets of size & of the samples but that does not include a varrance-reducing baseline.
all subsets
baseline:
naive partitioned baselined a naive transformation that sets all & transformed rewards in subset of & samples equal to the largest raw reward in that subset. To extend this method to n >> k we partition the n samples (for integer multiples of k) into disjoint subsets of size k and average the estimated gradient obtained from each. Furthermore, as simple variance reduction method, for each such set of lt samples we subtract the mean of the transformed rewards from the other sets of k samples (thereby averaging over (n - k) samples and subtracting the result from the samples and repeating n/k times leave-oneoout fashion over the subsets of size k). to randomly sample increasing number of partitions of the samples and average over all of them, then intuitively the resulting estimator would approach the variance of (lonn) but this would be expensive and indeed the limiting case of considering all partitions is intractable for general n and k. Our estimators have the key property of summing over all such partitions while nonetheless being efficient to compute.
were
naive partitioned no baseline similar method to the previous one, but without the naive mean subtraction based variance reduction step. loo minus one partitioned but instead of using the naive estimate (which sets every transformed reward to simple max of the raw reward given set of i: samples) uses the gloo--1) method applied separately to each disjoint set of & samples, and averages that over all such subsets. In this way, this is a trivial generalization of (ZSM25) which extends to n > by applying the basic method disjoint subsets and averaging the results. We do not subtract a baseline across sets as this did not improve the variance, possibly because the method within each &: already includes variance reduction baseline.
method that uses the same partitioning approach the previous two,
loo minus one all subsets: glooo-1) over all subsets of size i of the n samples and uses all appropriate subsets of size k: - 1 to form the verucing-reducing baseline that retains unbiasedness, thereby non-trivially generalizing [TSM225) to all n > k with strong variance reduction.
our novel estimator of Equation (33) that analytically sums
21

---

Raw vs Effective pa Rewards
Raw vs Effective passy@k Rewards
2
6
F
0.4
A=
R- 16
Effective
0.0
2-
-0.
-04
(aW rewardwddddo gbo)
[3W reward g(A)
(a) All subsets without baseline:
(b) subsets with 1.00 baseline: (loo)
Raw vs Effective pass.@k Rewards
1.0
Effective Raaword 60--1
0.8
0.2
0.4
0.2
0.0
0.2
0.4
raw reward glx))
(c) All sized subsets with sized subsets baseline: y (loo-1)
Figure 5: The effect of the LOO baseline on the effective rewards derived from n 32 raw re- wards g(x;) sampled uniformly from 1-1/2,11/2] The non baselined effective rewards (a) from (19) include vertical offset that grows with & despite being function of raw rewards (horizontal axis) that are centered around zero. The baselined effective rewards (b) and (c) from (29) and (33) respectively are more centered, and give rise to reduced gradient estimator variance. To construct the figure we grouped reward values into regularly spaced bins and averaged the transformed reward for each bin to construct curves. Note; because our transformations are from IR + R not possible to directly inspect one-dimensional transformation.
Cumulative solve rate
Entropy when optimizing pass@kon
0.6
89
Entropy
0.2
-91
-92-
0.0
93
2000
4000
6000
8000
10000
0
500
1000
1500
2000
2500
3000
Step
Step
(a) Cumulative salve (N 12000 problems)
(b) Effect of changing entropy during training
Figure 6: (a): Increasing hopt in PKPO training solves more problems during GEMMA2 RL. (b): higher hopt makes the model learn to have higher entropy during RL. Thus, by optimizing for
pass@k with A > instead of pass@1, the model tends to have higher entropy leading to better exploration and solving more problems. Note that the size of one epoch, which 750 steps, is evident in (a), where
the slope decrease at each epoch boundary.
22

---

passig ond p
pass gend
0.20
0.26
0.32
T2SM25
TZSM25
T25M25
0.16
0.22
0.28
500
1000
1560
2000
2500
3000
530
1000
1500
2000
2500
3000
530
1000
1560
2000
2500
3000
Stap
50ep
S0ap
(a) passá1 (h 1)
(h) pass02 (Aevaal 2)
(c) passld (h 4)
Pa5s
Pass@gA
Aond
0.44
1:9
panggee
T25
0.42
TZ5M25
2-42
0.38
0.34
(d) pass@8 (keval 8)
(c) pass@112 (k 12)
pass@16 (kmval 16)
Figure 7: Effect of kopt (used in our PKPO training) on the rolling pass@keval in GEMMA2 RL. Setting kope leval usually achieves the best pass@keval, Prior work [TSM255] (which is equiva- lent to the specific case of hp here presumably due to the larger estimator variance and unreliable gradient (see also Figure 4).
16 in our notation) is also shown for comparison, and suffers
Cumulative
pass@a Keval geval
pass@k eval k
0.020
1:07
0.10
Salve rate
pasagee
0.015
k =
passggee
0.08
0.06
0.4
0.2
0.02
0.0
800
1000
800
1000
Step
Stop
Stop
(a) Cumulative solve
(b) pass@1 (v.vval 1)
(c) pass@16 (k 16)
Figure 8: Our PKPO (kopt > 1) dramatically improves progress the challenging
Table 7: Results for GEMMA2-2B on the MATH benchmark.
GEMMA2-2B
k.eval=1
k.eval=2
keeval=4
k.eval=8
k.eval=16
koop==1
15.91 + 0.40
18.12 + 0.43
23.37 + 0.48
29.58 + 0.53
35.02 + 0.58
koop==2
15.15 0.41
20.73 0.45
25.81 0.50
31.96 + 0.55
37.75 1 0.60
k opt=4
14.86 + 0.43
19.86 + 0.47
0.52
34.27 + 0.58
38.91 = 0.62
k opt==8
14.19 0.46
19.33 0.50
26.45 0.55
35.49 + 0.60
40.73 = 0.65
TZSM25)
13.11 0.50
18.09 0.54
24.58
0.60
31.81 + 0.66
37.24 + 0.71
EntropyReg
14.51 + 0.48
18.95 + 0.52
25.33 + 0.58
30.95 + 0.64
36.18 + 0.69
Table 8: Results for GEMMA2-2B the Coding benchmark.
GEMMA2-2B
k_eval=1
k_eval=2
k_eval=4
k_eval=8
keeval116
k_opt=1
19.82 0.53
23.81 0.57
29.75
0.62
36.33 + 0.66
42.04 + 0.71
k_opt=2
18.70 + 0.54
26.94 + 0.59
33.82 + 0.64
40.95 + 0.69
47.03 + 0.74
koopt=4
18.69 0.56
26.43 0.61
36.81 0.67
44.81 + 0.73
50.55 = 0.78
k_op==8
17.94 + 0.59
25.86 + 0.64
35.88 + 0.70
46.45 + 0.77
52.83 + 0.83
TZSM25
16.81 + 0.65
24.27 + 0.69
33.11 = 0.76
41.98 + 0.84
47.26 0.89
EntropyReg
18.05 + 0.62
25.13 + 0.67
34.01 + 0.74
40.88 + 0.81
46.15 + 0.86
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: the abstract and introduction are typical and represent the paper's contribution and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The paper is largely theoretical, and the theorems include appropriate quali-
fiers and assumptions.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairnsss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't it acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
24

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: have made every effort ensure that the results precise and rigorous.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof's in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: the main contribution is reward transformation method that need only mod- ify standard RL algorithm by mapping batches of scalar rewards to their transformed values. We have provided the code for this transformation in Listing
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way accomplish this, but reproducibility detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
25

---

(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: the main non-trivial code that is required to transform the rewards is provided in the document as Listing I.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cr cc/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: ctttlsssccppublccc//c http://pbbigggidds/s/ppp SccoodSubisssiinPP...
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Every effort has been made to report this information to an appropriate level of detail.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
26

---

the paper is typical in that due to the computational cost of the experiments, 7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification:
we do not provide detailed condifence intervals, etc. However, we made every effort to provide suitable level of detail on the scope and significance of the experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
we have provided much information as we are permitted to by our organi- zation at this stage. The compute requirements are already strongly indicated by the fact that specify that we are training with the open source 2B GEMMA2 model. We can expand on this for the camera ready.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics http:////w.pppppcppp//pp////ctta
27

---

Answer: [Yes]
Justification:
The paper conforms rather comfortably, it is mainly theoretical paper
with standard experimental results on existing publicly available data.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: this is a methodological work that is not tied to any specific applications.
There no new and direct path from this paper to specific societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: this paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
28

---

Released models that have high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (c.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: have made every effort include these details.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, paaerssithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should provided.
this information not available online, the authors are encouraged to reach out to the asset's ss creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documenta- tion provided alongside the assets?
Answer: [NA]
Justification: the paper does not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/mmodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsoureing experiments and research with human subjects, does the pa- per include the full text instructions given participants and sereenshots, applicable, well as details about compensation (if any)?
29

---

Answer: [NA]
Justification: the paper does not use human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: see above.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMSs if it is important, original, or non-standard component of the core methods in this research? Note that if the L.LM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality the research, declaration is not required.
Question:
Answer: [NA]
Justification: the LLM training our experiments is standard and on standard publicly available tasks.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (httpss//neurips. ct/ http::/acc.aeeeeeeee//////////// for what should or should not be described.
30


Reviews:
[{"review_id": "ex7vMv48cu", "reviewer": "Reviewer_uuR2", "summary": "The authors derive a set of new algorithms for direct pass@k optimization. Starting from the unbiased pass@k /max@k estimators in the case of binary / continuous rewards, the authors derive a gradient estimator (both without and with variance reduction via LOO baseline).\n\nThe empirical evaluation involves training Gemma-2-2B on MATH and ARC-AGI-1 (easy subset), and baseline against another pass@k estimator (Tang et al, 2025). These experiments show improved pass@k performance when directly optimizing for the same k, as well as improvement over pass@k estimator from (Tang et al, 2025). \n\nThe authors also propose a method in which k is annealed during training, leading to models that balance better pass@1 and pass@k performance.", "questions": "1. Can you give an intuitive explanation for the resulting adjustments to the gradient estimator? (Equations 8 and 18-21) \n\n2. Do you have evidence that optimizing for pass@k leads to higher diversity of solutions beyond simple metrics like entropy?", "limitations": "None.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Summary:** the theoretical contribution of this work is clearly strong; the main crux is whether the much weaker empirical section is sufficient to justify acceptance. In this reviewer's opinion, this paper does not meet the bar for the strength of empirical support in RL-for-LMs theoretical papers published at NeurIPS, hence Borderline Reject. However, I'd happily change my rating to Accept with the addition of several new experiments (see the Weaknesses section).\n\n\n## Strengths\n\nThe paper is very well-written and addresses an important problem: improving the effectiveness of multiple trials by directly optimizing for pass@k / max@k. The derived estimators are (to my knowledge) all novel contributions, and provide a clear improvement over the prior work by generalizing to continuous rewards and cases where sample size $n \\neq k$, as well as reducing the variance of the estimator. I didn't find any issues with the proofs, although it is possible I missed some. \n\nThe method is also very straightforward to apply, since it just transforms the rewards (python implementation in Listing 1 is a nice addition).\n\nThe demonstrated results are quite promising: pass@k performance is generally better when optimizing for the same (or the at least closest) k and consistently outperforms another estimator (Tang et al, 2025). The results are particularly promising for ARC-AGI-1, where the advantages of directly optimizing for k>1 unstucks the training. I also appreciated the illustrative example in Figure 1. \n\n## Weaknesses\n\nThe evidence of the practical usefulness of the derived method is quite preliminary. The following is a list of weaknesses of the Experimental section, together with proposed additions/changes:\n\n1. only one small model (Gemma-2-2B) and two dataset (MATH, easy subset of ARC-AGI-1) (fix: add different and larger base models; possibly also more datasets)\n2. reporting train set performance only (fix: report test set performance instead, which is the more relevant metric for practical usefulness of the method)\n3. no experiments with continuous/max@k estimator (fix: add experiment on some continuous reward domain; could be a coding task, where g(x_j) = % unit tests passed)\n4. all results have only one random seed / no error bars, and the training curves are far from convergence -- this is all non-standard for a RL paper (fix: include $\\geq 3$ random seeds, train longer)\n5. lack of baselines: currently include ony pass@$k^{eval}$ comparison to the same algorithm with pass@$k^{opt}$ estimator (for different values of $k^{opt}$) or the pass@k estimator from (Tang et al, 2025) (fix: include more baselines, inc standard RLVR algorithms like GRPO, a best-of-N approximation method (Amini et al, 2024; Chow et al, 2024; Sessa et al, 2024), training with entropy bonus (which would show the proposed algorithm is incentivising something more interesting that increasing entropy))\n6. the claims of empirical diversity / higher exploration are unsupported: the only provided support is in Fig 2b, showing slightly higher entropy increase when optimizing for higher k -- but this is a poor proxy for diversity of attempted approaches. (fix: report some semantic measure of solutions obtained with pass@1 and pass@k optimization -- distribution of answers, semantic clustering of reasoning traces / programs / proofs --> greater diversity here would be a very exciting result!)\n7. reporting cumulative solve rate -- which seems less important and makes comparison to other RLVR work where this isn't used harder (fix: skip, report pass@k only)\n\n\n## Minor \n\n- the citation format is wrong for this venue\n- the illustrative 1-D example seems misplaced in the Experiments section? Perhaps move into its own section on intuitive explanation of the proposed method\n- specify the dataset in the Figure 2-5 captions\n- Fig 3-5: why not show the same $k^{eval}$ and $k^{opt}$ sets? \n- Fig 3-5: the learning curves matter less than final performance (fine in Fig 4 to show annealing point), perhaps summarize the results with bar plots and the learning curves could be in the appendix\n- in Fig 5: the y-axis cuts off very low -- we don't see the rate of improvement past 0.02/0.1\n- include a full list of training hyperparameters in the appendix\n- there's a tendency to provide interpretation of the results in the Experiments instead of a dedicated Discussion section\n- no proper related work section", "was_revised": true, "final_justification": "I increased the rating as the authors addressed my main concerns about the lack of empirical support (points 1-5 on the list of Weaknesses). \n\nI've been also convinced that BoN approximation baseline makes less sense (point 5) and the cumulative solve rate is a good indication of diversity here (point 7). However, I'd still argue for the inclusion of the GRPO baseline (point 5) and the inclusion of semantic measures of diversity (point 6) (see discussion for more details).", "rebuttal": ""}, {"review_id": "SJ7LDeXyu5", "reviewer": "Reviewer_Enpe", "summary": "This paper introduces Pass-at-k Policy Optimization (PKPO), a method for optimizing the pass@k metric in Reinforcement Learning rather than the conventional pass@1 approach. The key insight is that when sampling multiple solutions per problem, optimizing for the probability that at least one of k samples succeeds can improve exploration and performance on challenging tasks. The authors provide mathematical foundations, including unbiased estimators for pass@k and its gradient in both binary and continuous reward settings. The method is validated on toy problems and a small subset of real-world experiments using Gemma-2 on mathematical reasoning tasks.", "questions": "1. Given that practical RL fine-tuning often involves non-i.i.d. samples (e.g., due to shared dropout, KV-cache), how do these violations of the i.i.d. assumption affect the theoretical guarantees of unbiasedness and variance reduction? Can the authors provide theoretical analysis or empirical validation under these realistic, non-i.i.d. conditions to demonstrate the robustness of PKPO?\n2.   The theoretical claim of an unbiased gradient estimator is crucial. Can the authors provide a rigorous proof that the overall gradient estimator remains unbiased when using $s^{(loo-1)}$, specifically addressing how the potential admitted bias in the baseline term within the score-function product does not propagate to make the entire gradient estimate biased?\n3. * Can the authors provide a clear path towards achieving an implementation with $O(k + n \\log n)$ complexity as claimed in Theorem 5, rather than the currently noted $O(nk + n \\log n)$?\n    * How do you propose to address the super-linear scaling in wall-clock time across multiple tasks that require sorting, particularly for large-scale applications with many tasks?\n4.  Please provide a more detailed justification and proof for the specific form of $r_i$ in Equation (8), especially for the $f_i=0$ case and its interaction with the count `c`. How does this formulation ensure unbiasedness for all $k \\le n$, and specifically, address concerns about potential bias for small $k$ values?\n5. To strengthen the theoretical contributions and their applicability, can the authors provide experiments on domains or tasks where the i.i.d. assumption (or approximations thereof) is better satisfied or can be explicitly controlled? This would help validate the method's performance under its theoretical premises.", "limitations": "The authors acknowledge some limitations but miss several critical issues. These unaddressed points fundamentally impact the method's theoretical soundness and practical applicability. Specifically, the paper does not adequately discuss: (1) the potential bias in the leave-one-out formulation for $r_i$ in Equation (8), (2) the violation of i.i.d. assumptions in practical RL settings for LLMs, (3) the computational scaling issues across multiple tasks due to sorting, and (4) the potential for gradient bias originating from the admitted bias in the baseline terms within score-function estimators. Addressing these fundamental problems is crucial for the work.", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "This paper addresses an important problem in reinforcement learning: how to effectively optimize when multiple solution attempts per problem are sampled. The core idea of directly optimizing for pass@k is valuable. However, the proposed method suffers from several fundamental methodological and computational issues that undermine its theoretical claims and practical scalability.\n\n**Strengths:**\n*  The core insight of directly optimizing pass@k rather than pass@1 is valuable, as it addresses a real limitation in current RL practice where multiple samples are commonly used but optimization focuses only on individual sample quality[cite: 2, 4].\n* The approach of deriving unbiased estimators for pass@k and its gradient, applicable to both binary and continuous reward settings, is sophisticated and aims to provide a complete theoretical treatment.\n\n**Weaknesses:**\n*  The paper contains several mathematical flaws that compromise its theoretical rigor:\n    *  In Equation (8), the `r_i` definition for $f_i=0$ (incorrect samples) uses $\\rho(n-1,c,k-1)$. This formulation, particularly how the count `c` is used for samples that are assumed incorrect, appears to include the incorrect sample being evaluated, which violates the leave-one-out principle. This can introduce bias, especially for small `k`, directly contradicting the claim of `r_i` being unbiased.\n    * All theoretical proofs (e.g., Theorem 1, Theorem 2) assume `k` independent and identically distributed (i.i.d.) samples. However, in practical LLM fine-tuning settings, samples from the same mini-batch often share components like dropout masks and KV-cache. This makes them non-independent, directly violating a core assumption of the theoretical guarantees and potentially invalidating the unbiasedness claims in the empirical context.\n    * The paper claims that $s^{(loo-1)}$ is an unbiased estimator despite admitting a biased baseline. However, when a biased baseline appears directly inside the score-function product (as implied by the general form of policy gradient estimators), the overall gradient estimate itself becomes biased. This specific technical issue is not adequately addressed, undermining the claim of an unbiased gradient estimator.\n* **some computational complexity issues**\n    * Theorem 5 claims an $O(k + n \\log n)$ complexity for computing the reward transformations. However, the provided Python implementation (Listing 1) explicitly notes that it \"costs $O(nk + n \\log n)$ for simplicity\". This discrepancy is significant for scalability.\n    * For continuous rewards, the method requires $O(n \\log n)$ sorting per task (due to the assumption that rewards $g_i$ are sorted in Theorem 3 and its application to `s_i` computation), which would make the wall-clock time super-linear in `n` across multiple tasks. This significantly impacts scalability for large `k` or when processing many tasks.\n* The experimental evaluation is limited to one specific model architecture (Gemma-2 2B variant) and primarily two domains (mathematical reasoning and logic puzzles). This narrow scope limits the empirical contribution and does not validate the method's effectiveness on diverse RL applications or in settings where the theoretical assumptions might hold better.\n* The comparison largely focuses on vanilla pass@1 optimization. The paper lacks comprehensive comparisons with other contemporary inference-time optimization methods, advanced exploration techniques, or alternative approaches for handling multiple samples in RL literature.", "was_revised": true, "final_justification": "I am not an expert in the field. I am trying to evaluate this paper based on the mathematical rigor and empirical results.", "rebuttal": ""}, {"review_id": "gSNwVk7qrQ", "reviewer": "Reviewer_vPUj", "summary": "This paper focuses on optimising policies jointly over multiple samples, rather than following the usual single-sample optimisation paradigm. Anticipating the fact that multiple attempts will be made at inference time, it enables training policies that have a higher chance of producing at least one very good output. The paper proposes optimising over the best k out of n samples, thereby generalising prior work which fixes k = n. Additionally, it introduces an unbiased and low-variance gradient estimator for this setting, supported by mathematical proofs, and provides experiments to evaluate its effect on performance during training.", "questions": "What is the statistical significance of the empirical results? Specifically, how many seeds were used per experiment?\n\nCan we reasonably expect this method to outperform a standard entropy-regularized baseline? It would be helpful to understand whether the benefits go beyond what entropy regularisation alone can achieve.", "limitations": "There does not seem to be any discussion of the limitations. Could the authors add one?", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n- The pass@k optimization paradigm is relevant and increasingly important in RL for generative models. This paper contributes to it by introducing a general formulation (arbitrary k≤n), supporting it with rigorous mathematical derivations, and providing empirical validation.\n- The paper is well written. The mathematical content is clearly presented, and the proofs are systematically included, either in the main text or in the appendix.\n- The experimental section includes a toy problem to validate and illustrate the key claims (i.e., improved sample performance and reduced gradient variance).\n- The paper also presents large-scale experiments using Gemma-2 on challenging benchmarks (e.g., MATH, ARC-AGI), demonstrating that the proposed method can be applied to realistic and difficult problems.\n\n\n\n**Weaknesses:**\n- The title feels too broad relative to the empirical scope of the paper. While it refers to “solving harder RL tasks,” the experiments are limited to two tasks in the LLM domain. The claim would be more accurate if it were restricted to LLMs.\n- There already exist methods designed to jointly optimize over multiple attempts in hard RL tasks, such as Grinsztajn et al. (2022), Chalumeau et al. (2023), and follow-ups like Hottung et al. (2024) and Bonnet et al. (2024). It is unclear whether the approach proposed here would outperform these methods in other domains (e.g., combinatorial optimization). To be fully fair, the claims should be qualified to the LLM context, where those baselines have not (to my knowledge) been evaluated.\n- Statistical rigor: It is unclear whether the experimental results are averaged over multiple seeds. If only a single seed is used throughout, this raises concerns about the statistical robustness of the findings.\n- Line 54/55 claims that “[CTV19] and [XDC+20] present elegant approximations that are rather general but less efficient in our setting.” However, no empirical or theoretical justification is provided for this efficiency claim.\n- Baselines (concurrent methods): I am surprised by the lack of baselines in the paper. Is there only one concurrent method to compare to? \n- I’m unclear on what the loss function contributes beyond encouraging high-entropy solutions. - It would be useful to compare against a baseline that optimizes pass@1 (or pass@n) with an added entropy regularization term. This would help disentangle the benefits of the proposed method from standard exploration techniques in RL, such as entropy bonuses, which are widely used and easy to implement.\n- In Section 5.2.3, the authors demonstrate promising results on a difficult benchmark, but do not compare against other methods. The results would be compelling if competing approaches were also shown to fail in this setting, similar to the k=1 baseline. It remains unclear how well [TZSM25] or other baselines would perform on this task.\n\nReferences:\n- Nathan Grinsztajn, Daniel Furelos-Blanco, Shikha Surana, Clément Bonnet, and Thomas D. Barrett. Winner takes it all: Training performant rl populations for combinatorial optimization. Neurips 2023.\n- Felix Chalumeau, Shikha Surana, Clément Bonnet, Nathan Grinsztajn, Arnu Pretorius, Alexandre Laterre, and Thomas D Barrett. Combinatorial optimization with policy adaptation using latent space search. Neurips 2023\n- Andre Hottung, Mridul Mahajan, and Kevin Tierney. Polynet: Learning diverse solution strategies for neural combinatorial optimization. ICLR 2024\n- Clement Bonnet, Matthew Macfarlane. Searching Latent Program Spaces.", "was_revised": true, "final_justification": "I am willing to support the paper, although with a low confidence.\n\nThe paper seems to make important theoretical contributions. And support them with empirical evidence. The initial experiments lack statistical reports, but from the rebuttal, I trust that the authors will address this in the revised manuscript. I also trust that the authors will update their title and reduce the breadth of their claim to their mathematical contribution.\n\nI regret that the others reviewers did not engage more in the discussions, which would have helped the assessment of this paper. From the information I have, I am leaning towards a weak accept.", "rebuttal": ""}, {"review_id": "8AIf5G8chM", "reviewer": "Reviewer_jhr4", "summary": "In this paper, \"Solving Harder RL Tasks with Direct Pass@k Optimization\", the authors introduce a specific transformation on the final rewards which leads to a direct optimization of the pass@k objective.\n\nThe primary contribution is the theoretical motivation for the proposed transformation. Specifically the authors introduce unbiased estimators for the pass@k objective and gradient in the case of binary rewards. These estimators are then extended to the case of real rewards. The authors further propose various baselines for variance reduction for both binary and real valued rewards, and show that these can be efficient computed.\n\nFinally the authors show some empirical results of training with the transformed rewards on some toy and real domains.", "questions": "1. Compute restrictions aside were there any considerations for running $k^{eval}$ up to just 16? I would expect some sort of asymptotic behavior here and it would be interesting to observe this on, say, MATH\n2. The ability to implement an annealing procedure for k<n feels like it has a lot of promise for practitioners. Extending these results might provide just the small-scale but robust empirical justification for this reward scaling suitable to a theoretical paper.\n3. The step-function improvement on ARC-AGI-1 is pretty surprising (especially for the pass@1 eval). What's going on there?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper is clearly presented and easy to follow without extraneous content\n2. The paper presents a timely and important contribution. The implicit misalignment between the pass@1 training and pass@k usage of models leaves much on the table. The paper's contribution, to provide a more general (and practical) reward transformation to support the direct pass@k, is, to my mind, a clearly valuable contribution with immediate application.\n\nWeaknesses:\n1. There is likely some missing literature. Specifically, there are results like (https://arxiv.org/abs/2502.07154) which would both help to further motivate the problem and better contextualize the current contribution.\n2. The Introduction's positioning of this work as \"...unblock(ing) the path to artificial super intelligence...\" seems unnecessary\n3. The paper title itself appears somewhat misleading. While the authors show _some_ results which, expectedly, demonstrate improvement on \"harder\" tasks, it is clear that the paper's primary contribution is a theoretical advancement. I would prefer to have seen the title more closely match the content. \n4. ^^ Relatedly, the empirical validation, while suitable for a theoretically focused paper, is weak when making the more general claim suggested by the title (small models, few datasets).", "was_revised": true, "final_justification": "I retain my rating of 5-Accept. The authors appear to have address most concerns. I would still like to see the title changed (as author's have sort of agreed to).", "rebuttal": ""}]

Author rebuttal:

Meta-review: The paper consider optimizing the pass@k metric rather than the more common pass@1. They propose a variance reduced estimator that works for any n >= k, and several baselines. The technical contribution appears solid, and the motivation of the paper is also solid. In fact, pass@k is a very important objective, because it is connected to diversity of reasoning, and it is also directly useful as some application (such as writing a code that runs as fast as possible, or formal math proofs) can use test-time verifiers, making pass@k a critical metric to optimize if one expects to use higher test-time compute. Thus, I am glad to recommend the paper for acceptance.
Decision: Accept (spotlight)