Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems

Thanks for your time and effort, and supportive and critical comments. Please see shared comments above as well as those specific to you, here.

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.

Thanks for the pointer to this interesting paper. We will add a reference to it in our own work that highlights the excellent motivating descriptive theory of their section 4.2. For the purposes of this rebuttal, we have some detailed comments on that paper that highlight key features of our own:

• We really enjoyed their descriptive theory in section 4.2, which does indeed help to contextualize our work.
• Generally, computing the pass@1 (which they employ in their theory and hint at in their loss function) is intractable as it requires summing over all correct solutions.
• Unlike our RL scheme, their objective of equation (3) does not optimize the pass@k (nor does it explicitly claim to), but is a heuristic that is suggested by the 1-(1-pass@1)^k form of the pass@k.
• From what we understand they seem to optimize their loss in a supervised learning setting given known ground truth solutions. In contrast, our approach finds correct solutions with RL while optimizing the pass@k. Their scheme may be powerful for distillation type schemes but it is not comparable to ours in its intention.
• Their method boils down to a supervised learning loss function on a single sample, whereas ours boils down to a reinforcement learning reward transformation on a batch of n >= k samples. The use of n >= k samples is important; we can prove that no unbiased estimator of the pass@k exists that uses n < k samples. It is a neat proof that we are happy to provide here add to the appendix of our submission.
• They compare with GRPO, which is not related to pass@k but rather is a basic variance reduction method for RL on pass@1. Our experiments try to show that our method works as intended by comparing to the most closely related methods that address the same problem, and by directly looking at the variance of the different schemes on a toy problem.

[CTJ+21] Evaluating large language models trained on code, OpenAI.

Thanks again for your time and effort. We made a mistake in OpenReview and our rebuttals were not visible to all. Our apologies - please see the rebuttals which should now be visible. We remain very confident in our theory as explained in our feedback to Reviewer Enpe.

In case there is still an issue, here is an exact copy of the response to reviewer Enpe. Please let us know if any other responses are not visible.

(begin copy)

Thanks for your time and effort, and supportive and critical comments. Please see shared comments above as well as those specific to you, here.

Thank you for the detailed technical questions. We believe there may be a misunderstanding on these points, which we hope to clarify below. The unbiasedness of our estimators is a cornerstone of our work, and we appreciate the opportunity to elaborate on the proofs.

Technical clarifications

In Equation (8) ... [the result is incorrect] ... contradicting the claim of r_i being unbiased.

We respectfully believe that the above statement in the review is incorrect. See theorem 2 and the proof of that theorem. We assume that the reviewer missed this theorem given their subsequent comment "Please provide a more detailed justification and proof for ... Equation (8)". We hope this clarification, along with the proof of Theorem 2, fully addresses the reviewer's concern.

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.

We respectfully believe that this statement is incorrect: Conditional independence is guaranteed by the standard sampling process of autoregressive models, it is not an assumption we are making. Given a prompt, each sample is generated independently. Neither the use of a shared KV-cache for the prompt encoding nor standard dropout violates this.

The paper claims that s^{(loo-1)} is an unbiased estimator despite admitting a biased baseline ... undermining the claim of an unbiased gradient estimator.

We respectfully believe this is incorrect. The unbiasedness here is a subtle point and we are glad for the chance to expound on it. It is a cornerstone of our work and also [TZSM25] where it is Lemma 1; if the reviewer comment were correct, that paper would also be deeply flawed (which it is not). The unbiasedness also follows from Corollary 1 in our submission, but some extra work given that the baseline is not constant but rather depends on other samples. In detail:

Assume w.l.o.g. we have two samples x_1 and x_2 and an arbitrary function g(x_2) as the baseline for the estimator for the gradient of the reward r(x_1).

The baselined policy gradient estimator is:

∇_θ J ≈ (r(x_1) - g(x_2)) * ∇_θ log p(x_1 | z, θ)

Consider the expectation over prompts z from the dataset D and conditionally independent samples x_1, x_2 from the policy p(·|z, θ).

E[g(x_2) * ∇_θ log p(x_1 | z, θ)]

Using iterated expectations, we can first condition on a fixed prompt z:

= E_{z∼D} [ E_{x_1, x_2∼p(·|z,θ)} [ g(x_2) * ∇_θ log p(x_1 | z, θ) | z ] ]

The inner expectation of the product can be separated:

= E_{z∼D} [ E_{x_2∼p(·|z,θ)}[g(x_2)] * E_{x_1∼p(·|z,θ)}[∇_θ log p(x_1 | z, θ)] ]

The second term is always zero by Lemma 2 from our submission, i.e.:

E_{x_1∼p(·|z,θ)}[∇_θ log p(x_1 | z, θ)] = ∫ p(x_1|z,θ) * (∇_θ p(x_1|z,θ) / p(x_1|z,θ)) dx_1 = ∇_θ ∫ p(x_1|z,θ) dx_1 = ∇_θ(1) = 0

so the inner expectation becomes zero for any prompt z:

= E_{z∼D} [ E_{x_2∼p(·|z,θ)}[g(x_2)] * 0 ] = E_{z∼D}[0] = 0

The introduced bias is therefore equal to zero.

Other responses

the provided Python implementation (Listing 1) explicitly notes that it costs O(nk+n log(n)) for simplicity ... provide a clear path towards achieving an implementation with implementation with complexity as claimed in Theorem 5

Please see lines 414-419 of the supplementary material.

O(n log n) ... significantly impacts scalability ...

As per lines 72-79, n is the number of samples for a given task (such as a single math question) in a single minibatch. As such this cost was negligible compared to the gradient and sampling costs in our experiments (which for a typical quadratic attention transformer costs O(f(|theta|) n t^2) where t is the sequence length and f(|theta|) is a function of the number of parameters). We were surprised that such an efficient algorithm exists given that the naive cost is O(n choose k), and we hope the reviewer appreciates the effort it took to find these efficient algorithms.

(end copy)

Thanks for your time and effort, and supportive and critical comments. Please see shared comments above as well as those specific to you, here.

Statistical rigor ... if only a single seed is used throughout, this raises concerns about the statistical robustness of the findings.

This is a fair and important point. Due to time and computational constraints, we were unable to run all experiments over multiple seeds for this rebuttal. We acknowledge this as a limitation and will add a note to the paper. We will make it a priority to add results with multiple seeds to the camera-ready version.

"[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.

We consider all subsets of size k of 1,2,...,n. Naively applying their generic methods would have cost at least O(n choose k).

Baselines (concurrent methods): I am surprised by the lack of baselines in the paper. Is there only one concurrent method to compare to?

Yes! There is only one (very recent) method that derives unbiased estimators for the pass@k, [TZSM25]. However it only considers number of samples n=k. Or work goes much further in theoretical depth and is more generally applicable (by decoupling the minibatch size and k). However, their insight of using subsets of size k-1 in the baseline was new to us and dramatically improved our estimators, leading to our s^{loo-1} (which was also by far the most challenging to derive in our setup, but to our suprise can be computed efficiently).

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.

We have run this for [TZSM25] and it indeed performs poorly. We will include this plot in our paper as it is indeed good evidence indicative of the strength of our method, thanks for the suggestion.

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). ... To be fully fair, the claims should be qualified to the LLM context, where those baselines have not (to my knowledge) been evaluated.

Indeed, these are very interesting but very different to our drop in reward transformations for RL that guarantee pass@k optimization. We agree that lots of great work could be done combining these with RL and LLMs.

Thanks for your reply, and for sharing these additional elements. I am curious to hear the other reviewers' opinion to your rebuttal. If the other reviewers do not submit their updated recommendations, I will nevertheless raise my score to a weak accept (4), trusting that the authors will address the points discussed.

Thanks for your time and effort, and supportive and critical comments. Please see shared comments above as well as those specific to you, here.

Thank you for the detailed technical questions. We believe there may be a misunderstanding on these points, which we hope to clarify below. The unbiasedness of our estimators is a cornerstone of our work, and we appreciate the opportunity to elaborate on the proofs.

Technical clarifications

In Equation (8) ... [the result is incorrect] ... contradicting the claim of r_i being unbiased.

We respectfully believe that the above statement in the review is incorrect. See theorem 2 and the proof of that theorem. We assume that the reviewer missed this theorem given their subsequent comment "Please provide a more detailed justification and proof for ... Equation (8)". We hope this clarification, along with the proof of Theorem 2, fully addresses the reviewer's concern.

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.

We respectfully believe that this statement is incorrect: Conditional independence is guaranteed by the standard sampling process of autoregressive models, it is not an assumption we are making. Given a prompt, each sample is generated independently. Neither the use of a shared KV-cache for the prompt encoding nor standard dropout violates this.

The paper claims that s^{(loo-1)} is an unbiased estimator despite admitting a biased baseline ... undermining the claim of an unbiased gradient estimator.

We respectfully believe this is incorrect. The unbiasedness here is a subtle point and we are glad for the chance to expound on it. It is a cornerstone of our work and also [TZSM25] where it is Lemma 1; if the reviewer comment were correct, that paper would also be deeply flawed (which it is not). The unbiasedness also follows from Corollary 1 in our submission, but some extra work given that the baseline is not constant but rather depends on other samples. In detail:

Assume w.l.o.g. we have two samples x_1 and x_2 and an arbitrary function g(x_2) as the baseline for the estimator for the gradient of the reward r(x_1).

The baselined policy gradient estimator is:

∇_θ J ≈ (r(x_1) - g(x_2)) * ∇_θ log p(x_1 | z, θ)

Consider the expectation over prompts z from the dataset D and conditionally independent samples x_1, x_2 from the policy p(·|z, θ).

E[g(x_2) * ∇_θ log p(x_1 | z, θ)]

Using iterated expectations, we can first condition on a fixed prompt z:

= E_{z∼D} [ E_{x_1, x_2∼p(·|z,θ)} [ g(x_2) * ∇_θ log p(x_1 | z, θ) | z ] ]

The inner expectation of the product can be separated:

= E_{z∼D} [ E_{x_2∼p(·|z,θ)}[g(x_2)] * E_{x_1∼p(·|z,θ)}[∇_θ log p(x_1 | z, θ)] ]

The second term is always zero by Lemma 2 from our submission, i.e.:

E_{x_1∼p(·|z,θ)}[∇_θ log p(x_1 | z, θ)] = ∫ p(x_1|z,θ) * (∇_θ p(x_1|z,θ) / p(x_1|z,θ)) dx_1 = ∇_θ ∫ p(x_1|z,θ) dx_1 = ∇_θ(1) = 0

so the inner expectation becomes zero for any prompt z:

= E_{z∼D} [ E_{x_2∼p(·|z,θ)}[g(x_2)] * 0 ] = E_{z∼D}[0] = 0

The introduced bias is therefore equal to zero.

Other responses

the provided Python implementation (Listing 1) explicitly notes that it costs O(nk+n log(n)) for simplicity ... provide a clear path towards achieving an implementation with implementation with complexity as claimed in Theorem 5

Please see lines 414-419 of the supplementary material.

O(n log n) ... significantly impacts scalability ...

As per lines 72-79, n is the number of samples for a given task (such as a single math question) in a single minibatch. As such this cost was negligible compared to the gradient and sampling costs in our experiments (which for a typical quadratic attention transformer costs O(f(|theta|) n t^2) where t is the sequence length and f(|theta|) is a function of the number of parameters). We were surprised that such an efficient algorithm exists given that the naive cost is O(n choose k), and we hope the reviewer appreciates the effort it took to find these efficient algorithms.

Thanks again for your time and effort, and for diving deeply into the paper. We hope that your theoretical concerns have been addressed in our rebuttal. If not, we are happy to explain further, just let us know.

Thanks for your time and effort, and supportive and critical comments. Please see shared comments above as well as those specific to you, here.

[use larger models, more baselines, more datasets, more random seeds, out of sample accuracy] ... [remove the cumulative solve rate plot]

We appreciate the comments, but respectfully point the reviewer to common review comments, particularly comment G1. Our goal is a proof of concept for our theory, and a demonstration that we can find more solutions to an RL training set with verifiable rewards but no ground truth solutions. We apologize for not bringing this out more clearly in the paper, and will remedy this with a new discussion section.

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!)

Diversity is not intrinsically important. Rather, only diversity that yields more correct solutions matters. We therefore consider our cumulative solve rate plot 2b the ideal way of presenting diversity for this work. We respectfully remind that your review suggested to actually remove that cumulative solve rate plot; we wonder if you would reconsider that point in light of our rebuttal, both here and general comment G1 in the shared part of the rebuttal, above. Furthermore, entropy is very widely used and well better understood measure of the diversity of model samples as compared to ad hoc analyses such as clustering.

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)

Thanks for the suggestions. We checked each one in detail and offer the following comments. [GRPO] optimizes pass@1, not pass@k. It uses a relative advantage transformation that is very similar to our leave one out reward transformation for the pass@1 results reported in our paper, but with additional innovation to handle multi step RL, which we do not address and which is orthogonal to our contribution. [Amini et al] and [Sessa et al] learn to approximate the BoN result (that comes from ranking multiple samples and taking the best) in a single sample; this is in some sense the exact opposite of our goal, which is to obtain more diverse individual samples such that at least one our of k samples are correct at the expense of individual sample strength. [Chow et al] is closer to our work but a comparison here is rather a demanding request since it is a very new and non trivial method for which no implementation is available. We can make some comparisons by eyeballing their results however. For example, a number of their results (such as Fig 5b) show that when optimizing for larger N they dominate models optimized for smaller N on the small N evaluations. This is not indicative that their objective works as intended; rather, while their methods somehow provide some benefit in some settings, they do not appear to trade pass@1 and pass@k as intended. This is could be due the approximations they make in deriving their tractable objective, which is ultimately not guaranteed to optimize pass@k. In contrast we have an unbiased gradient which guarantees pass@k optimization.

Can you give an intuitive explanation for the resulting adjustments to the gradient estimator? (Equations 8 and 18-21)

We tried to do this with Figures 6 and 7 in the appendix. These visualize the contribution of individual terms to our estimators. Roughly speaking, the larger k, the closer we go toward rewarding all n samples with the same reward value, namely the max (continuous case) or any (binary case) reward. As k is reduced to one, we revert to assigning to each sample the unchanged reward associated with it. Values in between interpolate between these extremes.

Thanks again for your time and effort. We made a mistake in OpenReview and our rebuttals were not visible to all. Our apologies - please see the rebuttals which should now be visible. We remain very confident in our theory as explained in our feedback to Reviewer Enpe.

In case there is still an issue, here is an exact copy of the response to reviewer Enpe. Please let us know if any other responses are not visible.

(begin copy)

Thanks for your time and effort, and supportive and critical comments. Please see shared comments above as well as those specific to you, here.

Thank you for the detailed technical questions. We believe there may be a misunderstanding on these points, which we hope to clarify below. The unbiasedness of our estimators is a cornerstone of our work, and we appreciate the opportunity to elaborate on the proofs.

Technical clarifications

In Equation (8) ... [the result is incorrect] ... contradicting the claim of r_i being unbiased.

We respectfully believe that the above statement in the review is incorrect. See theorem 2 and the proof of that theorem. We assume that the reviewer missed this theorem given their subsequent comment "Please provide a more detailed justification and proof for ... Equation (8)". We hope this clarification, along with the proof of Theorem 2, fully addresses the reviewer's concern.

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.

We respectfully believe that this statement is incorrect: Conditional independence is guaranteed by the standard sampling process of autoregressive models, it is not an assumption we are making. Given a prompt, each sample is generated independently. Neither the use of a shared KV-cache for the prompt encoding nor standard dropout violates this.

The paper claims that s^{(loo-1)} is an unbiased estimator despite admitting a biased baseline ... undermining the claim of an unbiased gradient estimator.

We respectfully believe this is incorrect. The unbiasedness here is a subtle point and we are glad for the chance to expound on it. It is a cornerstone of our work and also [TZSM25] where it is Lemma 1; if the reviewer comment were correct, that paper would also be deeply flawed (which it is not). The unbiasedness also follows from Corollary 1 in our submission, but some extra work given that the baseline is not constant but rather depends on other samples. In detail:

Assume w.l.o.g. we have two samples x_1 and x_2 and an arbitrary function g(x_2) as the baseline for the estimator for the gradient of the reward r(x_1).

The baselined policy gradient estimator is:

∇_θ J ≈ (r(x_1) - g(x_2)) * ∇_θ log p(x_1 | z, θ)

Consider the expectation over prompts z from the dataset D and conditionally independent samples x_1, x_2 from the policy p(·|z, θ).

E[g(x_2) * ∇_θ log p(x_1 | z, θ)]

Using iterated expectations, we can first condition on a fixed prompt z:

= E_{z∼D} [ E_{x_1, x_2∼p(·|z,θ)} [ g(x_2) * ∇_θ log p(x_1 | z, θ) | z ] ]

The inner expectation of the product can be separated:

= E_{z∼D} [ E_{x_2∼p(·|z,θ)}[g(x_2)] * E_{x_1∼p(·|z,θ)}[∇_θ log p(x_1 | z, θ)] ]

The second term is always zero by Lemma 2 from our submission, i.e.:

E_{x_1∼p(·|z,θ)}[∇_θ log p(x_1 | z, θ)] = ∫ p(x_1|z,θ) * (∇_θ p(x_1|z,θ) / p(x_1|z,θ)) dx_1 = ∇_θ ∫ p(x_1|z,θ) dx_1 = ∇_θ(1) = 0

so the inner expectation becomes zero for any prompt z:

= E_{z∼D} [ E_{x_2∼p(·|z,θ)}[g(x_2)] * 0 ] = E_{z∼D}[0] = 0

The introduced bias is therefore equal to zero.

Other responses

the provided Python implementation (Listing 1) explicitly notes that it costs O(nk+n log(n)) for simplicity ... provide a clear path towards achieving an implementation with implementation with complexity as claimed in Theorem 5

Please see lines 414-419 of the supplementary material.

O(n log n) ... significantly impacts scalability ...

As per lines 72-79, n is the number of samples for a given task (such as a single math question) in a single minibatch. As such this cost was negligible compared to the gradient and sampling costs in our experiments (which for a typical quadratic attention transformer costs O(f(|theta|) n t^2) where t is the sequence length and f(|theta|) is a function of the number of parameters). We were surprised that such an efficient algorithm exists given that the naive cost is O(n choose k), and we hope the reviewer appreciates the effort it took to find these efficient algorithms.

(end copy)

Details of experiments

We thank the reviewers for suggesting additional experiments which could support our claims. In our new experiments, we do the following:

1. Run our procedure on Gemma-2 2B,9B and Llama3.1 8B
2. Held-out test-sets: To evaluate the final trained model to strengthen our claims,
3. Train until loss curves have saturated. As a consequence, the new experiments are now trained for 10k steps as opposed to 3k steps
4. Added error bars: Each run is repeated thrice with different seeds. The eval numbers have been reported with error bars thus accounting for 3 runs for every (model, dataset, k_opt).
5. Entropy regularization baseline: Apart from TZSM25, we also add the entropy regularization baseline. We use PPO with entropy bonus. To give this baseline the best shot, we ran a small sweep [0.001, 0.005, 0.01, 0.05, 0.1] of the entropy_coefficient for each (model,benchmark), and only report the best performance each time as Entropy Reg (tuned). For every entropy_coefficient, runs are also seeded to get error bars as in 4 before choosing the best one.

MATH benchmark:

We use the same MATH dataset for training. In addition, we now use the held out eval set for model evals below. We evaluate the final checkpoint in each run

Coding Benchmark:

We use MBPP for RL training and HumanEval for held-out evaluation. MBPP has multiple unit tests per problem and hence we use this not only as a proxy for additional benchmarks but also to show RL training with continuous reward function (% unit tests passed)

ARC-AGI

We make a 80/20 train-test split of the same easy set as in our paper. We report the cumulative solve rate on the train set and pass@k rate on the test set. We train till saturation (no change in cumulative rate past 1k steps)

Observations

1. The tables below show strong trend such that for a given k_eval, pass@k is highest when k=k_opt and clearly outperforms baselines. Thus generalizing our observations to multiple models, datasets and held out evals
2. As we posit in the paper, the existing pass@k optimization method TZSM25 (k_opt=16) is sub-optimal since it directly optimizing for k=batch_size and has high noise in the estimator. In fact, a lower k_opt=8 out-performs TZSM25 for all pass@k_eval. Showing the clear value of our method, and the flexibility it provides, supported by its theoretical grounding.
3. Entropy Regularization does indeed sacrifice pass@1 and slightly improves pass@k by promoting exploration. But is hard to tune, and is significantly outperformed by our method. Moreover, it has no explicit way to optimize for a specific k_eval.
4. Error bars: We also point out that all of the experiments below had stable RL training with little variation observed across different seeds, as is indicated by the low value of the standard error. This observation is consistent with previous experiments and also why we originally only reported single runs
5. Both baselines don't unblock ARC-AGI training, but our method for high k (k_opt=4,8) does. This is indicated in the high cumulative solve rate (+60%) on train and significant performance gain (+25%) in test set

Results ARC-AGI

Results MATH Benchmark

Gemma-2B

Gemma-9B

Llama 3.1 8B

Results Coding

Gemma-2B

Gemma-9B

Llama 3.1 8B