PaperlayerView exact revision on arXiv

Source-linked AI summary

Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems

Christian Walder, Deep Karkhanis

arXiv:2505.15201v5cs.LGcs.AIcs.CLstat.ML

TL;DR

RL with multiple samples traditionally optimizes isolated-sample quality, underusing their collective utility and diversity on difficult problems. PKPO transforms rewards to directly optimize pass@k, and experiments show target-k optimization, improved exploration, and learning gains without sacrificing pass@1 through annealing.

  • Problem

    Independent rewards optimize pass@1 rather than the collective utility and diversity of multiple solution attempts, limiting exploration on challenging task sets.

  • Method

    PKPO derives low-variance unbiased estimators for pass@k and its gradient, then reduces optimization to standard RL with jointly transformed rewards for any k ≤ n.

  • Results

    Experiments show that matching kopt to keval best optimizes pass@keval, while annealing kopt from 8 to 1 improves pass@k without sacrificing pass@1 and PKPO unblocks learning on hard tasks.

  • Takeaways & Limitations

    Optimizing the best reward among sampled sets preserves output diversity, helps solve more problems, and yields stronger policies, especially on harder task sets.

  • Takeaways & Limitations

    Estimator variance increases as kopt approaches n because fewer subsets are available, and the kopt = n special case can struggle to optimize pass@n.

Abstract

from arXiv · show

Reinforcement Learning (RL) algorithms sample multiple n>1 solution attempts for each problem and reward them independently. This optimizes for pass@1 performance and prioritizes the strength of isolated samples at the expense of the diversity and collective utility of sets of samples. This under-utilizes the sampling capacity, limiting exploration and eventual improvement on harder examples. As a fix, we propose Pass-at-k Policy Optimization (PKPO), a transformation on the final rewards which leads to direct optimization of pass@k performance, thus optimizing for sets of samples that maximize reward when considered jointly. Our contribution is to derive novel low variance unbiased estimators for pass@k and its gradient, in both the binary and continuous reward settings. We show optimization with our estimators reduces to standard RL with rewards that have been jointly transformed by a stable and efficient transformation function. While previous efforts are restricted to k=n, ours is the first to enable robust optimization of pass@k for any arbitrary k <= n. Moreover, instead of 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 numbers alongside significant pass@k gains. We validate our reward transformations on toy experiments, which reveal the variance reducing properties of our formulations. We also include real-world examples using the open-source LLM, GEMMA-2. 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 due to better exploration by prioritizing joint utility over the utility of individual samples.

1 Introduction

The paper argues that optimizing single-sample reward can be sub-optimal when multiple samples are used for search, and introduces PKPO to optimize pass@k directly. Its estimators support arbitrary k ≤ n and can be annealed to improve exploration and pass@1 together.

  • Multiple-sample inference couples search with model updates, making naive pass@1 optimization potentially sub-optimal.
  • The method extends pass@k optimization beyond the restricted k = n setting to arbitrary k ≤ n.
  • PKPO constructs robust estimators for pass@k and its gradient by averaging simple estimators over all k-sized subsets.
  • The resulting reward transformations can be applied as drop-in replacements in standard policy-gradient methods.
  • Toy experiments evaluate variance reduction, while LLM experiments test task solving and selective pass@k optimization on real-world problems.

2 Binary Rewards

For binary rewards, pass@k is estimated from n ≥ k samples by averaging subset-based estimators, yielding unbiased estimates of both the metric and its gradient. The construction assigns greater weight to correct samples while retaining reward for incorrect samples to encourage exploration.

  • Pass@k is the probability that at least one of k independent samples is correct.
  • Given n ≥ k samples and c correct samples, ρ(n, c, k) provides an unbiased estimator of pass@k.
  • Averaging over all subsets of size k preserves unbiasedness, with asymptotic variance decreasing at rate 1/n.
  • The gradient estimator assigns more weight to correct samples while also rewarding incorrect samples to encourage exploration.
  • Theorem 2 establishes that the proposed estimator is unbiased for the gradient of pass@k.

3 Continuous Rewards

The paper extends pass@k-style estimation to continuous rewards by sorting sampled rewards and deriving unbiased estimators for maxg@k and its gradient. These estimators support efficient computation through weighted combinations and subset-counting arguments.

  • Continuous rewards generalize the pass@k objective to maxg@k over real-valued reward functions.
  • Sorting rewards g1 ≤ g2 ≤ ··· ≤ gn enables an unbiased maxg@k estimator based on subset contributions.
  • The estimator is computed stably by cancelling factors in binomial coefficients.
  • The gradient estimator b∇(g) is unbiased for the gradient of maxg@k.
  • Under sorted rewards, the estimator coefficients count size-k subsets whose largest element determines the factored reward.
  • The quantities needed for the gradient estimator can be computed in total time O(k + n log n).

4 Variance Reduction

The variance-reduction section develops leave-one-out baselines for maxg@k while preserving unbiasedness. Because naive baselines depend on every sample, the method excludes the differentiated sample and uses efficient subset-based recursions.

  • Subtracting the mean of leave-one-out rewards preserves unbiasedness when the subtracted term does not depend on xi.
  • A naive leave-one-out baseline for si is biased because si depends on all sampled rewards.
  • The corrected baseline excludes element i from the averaged subsets to retain unbiasedness.
  • The leave-one-out quantities use subset averages and can be computed in total time O(k + n log n).
  • For n = k, no subsets remain for constructing the baseline, motivating maxg@(k −1) as a special-case baseline.
  • Averaging smaller, more numerous subsets reduces variance but introduces bias in the baseline rather than in s(loo−1)i.

5 Experiments

Experiments evaluate PKPO on toy estimators and real-world language-model tasks, showing target-k optimization, improved exploration, and benefits on challenging problems.

  • 5.1 One-Dimensional Toy Example: The toy experiment finds that the optimal policy under maxg@k varies with k, while s(loo−1) has the strongest variance performance.Estimator variance is compared in Figure 4.
  • 5.2 RL on Open Source LLMs: Experiments use GEMMA2 and LLAMA3.1 variants on MATH, code generation, and ARC-AGI-1 tasks.The real-world evaluation includes 2B and 9B GEMMA2 variants and an 8B LLAMA3.1 variant.
  • 5.2 RL on Open Source LLMs: Training samples n = 16 completions per prompt, computes rewards for every completion, and applies the s(loo−1) transformation used by PKPO.The kopt = 1 setting is the no-transformation baseline, with leave-one-out mean centering used for stability.
  • 5.2 RL on Open Source LLMs: Higher kopt consistently produces higher cumulative solve rates and entropy on the 12,000-problem MATH training set, indicating greater exploration.The paper reports that optimizing pass@k appears to use the exploration budget more effectively and find more solutions.
  • 5.2.1 Choosing kopt selectively optimizes pass@keval and solves more tasks: The best pass@keval occurs when kopt equals keval or is the closest available value, while pass@1 optimization becomes increasingly sub-optimal as keval grows.The experiments support setting kopt := keval for any keval ≤ n.
  • 5.2.1 Choosing kopt selectively optimizes pass@keval and solves more tasks: As kopt approaches n, estimator variance increases; the reported kopt = 8 gains are stronger for keval ∈ {12, 16} than for keval = 8.The paper associates this with fewer subsets in the estimator and notes that the n = kopt special case struggles with pass@n.
  • 5.2.3 Improving pass@k without sacrificing pass@1: Annealing kopt from 8 to 1 after 1500 steps improves pass@keval for keval > 1 without sacrificing pass@1.The schedule first prioritizes exploration through pass@k optimization, then consolidates the single-sample policy.
  • 5.2.4 PKPO is essential for learning on hard problems: On the challenging ARC-AGI-1 task set, conventional pass@1 optimization stalls, whereas PKPO unblocks learning and yields higher pass@keval for every evaluated keval, including 1.Higher kopt also produces more effective and faster learning in this setting.

6 Conclusions and Outlook

PKPO optimizes the best reward among multiple independent samples rather than their average, preserving diversity and supporting stronger policies. It provides drop-in reward transformations and remains extensible to other search and baseline methods.

  • PKPO maximizes the expected best reward in a sample set instead of the average reward.
  • The authors connect this objective to preserved output diversity, more solved problems, and stronger policies.
  • The method is presented as a drop-in replacement for traditional RL reward transformations.
  • Future extensions include other inference-time search algorithms and more sophisticated baseline techniques.

A.1 Statement and proof that n ≥k samples are required to unbiasedly estimate pass@k

The appendix establishes that unbiased estimation of pass@k requires at least k i.i.d. samples. It derives this through the polynomial characterization of Bernoulli estimability and analyzes estimator variance using U-statistics and Hoeffding theory.

  • A Bernoulli function is unbiasedly estimable from n samples exactly when it is a polynomial in the success probability of degree at most n.
  • Consequently, pass@k is unbiasedly estimable from n i.i.d. samples if and only if n ≥ k.
  • For binary rewards, pass@k equals 1 − (1 − p)^k, a polynomial of degree k in the single-sample success probability.
  • The pass@k estimator is identified as a U-statistic, whose variance is characterized using Hoeffding’s asymptotic theory.
  • The variance analysis conditions on correctness outcomes being i.i.d. Bernoulli variables for a fixed model–problem pair.
  • The estimator’s asymptotic variance decreases at a rate of 1/n.

A.3 Proof of Theorem 2

The proof counts subsets containing each sample and distinguishes whether that sample is correct, then represents the resulting reward transformation as a structured matrix operation. The implementation exploits this structure for efficient computation.

  • The proof counts size-k subsets that contain a given element and include at least one correct element.
  • The subset count takes one of two values depending on whether the selected sample is correct or incorrect.
  • When the selected sample is correct, every qualifying subset containing it is successful, yielding a count based on size-(k−1) subsets.
  • When the selected sample is incorrect, the count instead requires the remaining subset to contain at least one correct element.
  • The transformed reward vector is written as s = Mg, with M defined by diagonal and upper-diagonal entries and zero lower-diagonal entries.
  • After sorting, the required computation costs O(k+n) plus O(n log n) for sorting.

B Implementation

The implementation supplies Python batch transformations for the estimator and several baselines, including leave-one-out variants. It uses sorting and cumulative computations, with a simple implementation cost of O(nk + n log n).

  • The code computes diagonal coefficients through repeated normalized matrix evaluations.
  • Delta terms are formed from differences between neighboring normalized matrix evaluations.
  • The implementation uses cumulative sums and sorted application helpers in its batch transformations.
  • The implementation exposes functions for ρ(g), s_i, s^(loo), and s^(loo−1) reward transformations.
  • The straightforward implementation costs O(nk + n log n), while reducing it to O(k + n log n) requires optimizing delta and diagonal computations.

C Additional Figures

The additional figures examine estimator variance, reward baselines, entropy, and pass@k performance across PKPO settings and benchmarks. Together, they show how k and reward transformation affect exploration, learning progress, and evaluation outcomes.

  • Estimator variance: Figure 4 compares the sample variance of unbiased maxg@k gradient estimators as the number of samples n increases, with lower variance preferred.The figure uses k = 4 and 10,000 independent gradient estimates for a one-dimensional problem.
  • Reward transformation: LOO baselines center effective rewards more closely and reduce gradient-estimator variance compared with non-baselined rewards.The non-baselined rewards exhibit a vertical offset that grows with k, whereas the baselined rewards are more centered.
  • GEMMA2 training: Higher kopt increases entropy during GEMMA2 RL and solves more problems, consistent with improved exploration.The figure notes that epoch boundaries occur every 750 steps.
  • GEMMA2 evaluation: Setting kopt = keval usually achieves the best rolling pass@keval, while prior work using kopt = n = 16 suffers from larger estimator variance and unreliable gradients.The comparison includes kopt values used in PKPO training and several evaluation values keval.
  • ARC-AGI-1: PKPO with kopt > 1 dramatically improves progress on the challenging ARC-AGI-1 task set.The supplied figure caption identifies this result without reporting a numerical value.

NeurIPS Paper Checklist

The checklist states that the paper’s abstract and introduction represent its contribution and scope, while the paper reports implementation and experimental details to support reproducibility. It also identifies limited confidence-interval reporting and calls for explicit assumptions and limitations.

  • Claims and scope: The checklist justification says the abstract and introduction are typical and represent the paper’s contribution and scope.The checklist guidance links this assessment to whether claims match the theoretical and experimental results and their expected generality.
  • Theory assumptions and proofs: The checklist emphasizes that theoretical results should state assumptions clearly and provide complete, correct proofs.The supplied justification says the authors made efforts to ensure precision and rigor.
  • Reproducibility: The main reward-transformation code is provided in Listing 1, enabling adaptation of a standard RL algorithm by mapping scalar rewards to transformed values.The checklist describes this transformation as the main non-trivial code needed for the contribution.
  • Experimental details: The checklist justification says the paper makes efforts to report experimental information at an appropriate level of detail, including training and test details.The supplied checklist text also states that the compute requirements are strongly indicated by training the open-source 2B GEMMA2 model, with further expansion possible for the camera-ready version.
  • Statistical reporting: Because of computational cost, the paper does not provide detailed confidence intervals, although it reports making efforts to describe the scope and significance of its results.This is the clearest limitation stated in the supplied checklist material.
Loading 2505.15201v5…