Technical Summary: Pass@K Policy Optimization (PKPO)

Problem Statement

Reinforcement Learning (RL) algorithms applied to Large Language Models (LLMs) often sample multiple solution attempts ($n > 1$) for a single problem but optimize based on the performance of individual samples (pass@1). This approach prioritizes the utility of a single high-reward sample over the collective utility of the set, effectively under-utilizing the sampling capacity. Consequently, standard policy gradient methods may fail to explore the solution space sufficiently, particularly on harder tasks where the probability of a single correct sample is low. While inference-time search methods exist, the coupling between inference and model updates suggests that naively optimizing the expected single-sample reward is sub-optimal for tasks where correctness can be verified but solutions are not known a priori.

Methodology

The authors propose Pass-at-k Policy Optimization (PKPO), a framework that directly optimizes the pass@k metric (or its continuous analog, maxg@k), defined as the expected maximum reward over $k$ independent samples drawn from the policy.

Core Theoretical Contributions

The primary contribution is the derivation of novel, low-variance, unbiased estimators for both the pass@k objective and its gradient. These estimators function by averaging simple estimators over all $\binom{n}{k}$ subsets of size $k$ from a batch of $n$ samples.

1. Binary Rewards (Section 2):

   • For binary correctness labels ($f \in \{0, 1\}$), the authors derive an unbiased estimator $\rho(n, c, k)$ for pass@k, where $c$ is the count of correct samples.
   • They prove that an unbiased estimator exists only if $n \ge k$.
   • A corresponding unbiased gradient estimator is derived, assigning weights to samples based on their correctness and the combinatorial probability of their inclusion in a successful subset.

2. Continuous Rewards (Section 3):

   • The framework generalizes to continuous rewards $g(x)$ via the maxg@k objective: $E[\max_{i=1}^k g(x_i)]$.
   • By assuming rewards are sorted ($g_1 \le g_2 \le \dots \le g_n$), the authors derive an unbiased estimator $\rho^{(g)}(n, c, k)$ that weights sorted rewards using binomial coefficients.
   • An unbiased gradient estimator is constructed where the weight $s_i$ for a sample is a linear combination of rewards, computed efficiently in $O(k + n \log n)$ time.

3. Variance Reduction (Section 4):

   • To reduce the variance of gradient estimates without introducing bias, the authors propose Leave-One-Out (LOO) baselines.
   • They introduce a specific baseline $s^{(loo-1)}_i$ which subtracts the expected maxg@k of subsets of size $k-1$ (excluding the current sample) from the maxg@k of subsets of size $k$. This generalizes previous work (specifically [TZSM25]) which was restricted to the case $n=k$.
   • This formulation allows for robust optimization even when $n > k$, a common scenario in LLM training where large batches are sampled.

Implementation

The method is designed as a "drop-in replacement" for standard reward transformations in any policy gradient algorithm (e.g., PPO). Given a batch of $n$ samples and their raw rewards, the rewards are transformed using the derived estimators (specifically the sloo_minus_one function in the provided code) before being fed into the policy update step.

Key Results

The authors validate PKPO through toy experiments and real-world benchmarks using GEMMA2 (2B and 9B parameters) and LLAMA3.1 (8B parameters).

1. Toy Experiments:

   • In a 1D Gaussian policy experiment, PKPO successfully shifts the optimal policy mean to be more "risk-tolerant" for higher $k$, allowing the model to explore regions with lower probability but higher potential reward.
   • Variance analysis confirms that the proposed sloo_minus_one estimator significantly reduces gradient variance compared to naive partitioning or estimators without baselines.

2. Real-World Benchmarks (MATH, Coding, ARC-AGI-1):

   • Selective Optimization: Setting the optimization target $k_{opt}$ equal to the evaluation metric $k_{eval}$ yields the best performance. For instance, optimizing for pass@8 significantly improves pass@8 evaluation scores compared to optimizing for pass@1.
   • Solving Harder Problems: On the challenging ARC-AGI-1 benchmark, conventional pass@1 optimization stalls, failing to improve the cumulative solve rate. In contrast, PKPO with $k_{opt} > 1$ "unblocks" learning, achieving higher solve rates and pass@k scores across all $k$.
   • Annealing Strategy: The authors demonstrate that annealing $k$ during training (starting with a high $k_{opt}$ to encourage exploration, then reducing to $k_{opt}=1$ to consolidate) achieves strong pass@1 performance alongside significant pass@k gains, avoiding the typical trade-off where improving pass@k degrades pass@1.
   • Comparison: PKPO outperforms the special case of $n=k$ (as in [TZSM25]) and entropy regularization baselines, particularly on difficult tasks where exploration is critical.

Significance and Claims

The paper claims that PKPO provides a robust, efficient, and theoretically grounded method to optimize for the collective utility of a set of samples rather than just the average or single-best sample.

• Unblocking Exploration: The authors argue that by prioritizing joint utility (the probability that at least one of $k$ samples succeeds) over individual sample utility, the method improves exploration. This is crucial for solving harder problems where the "signal" (a correct solution) is sparse in the search space.
• Flexibility: Unlike previous methods that were restricted to $k=n$ (optimizing for the best of the entire batch) or $k=1$, PKPO enables robust optimization for any arbitrary $k \le n$.
• Practical Impact: The method allows practitioners to tune the exploration-exploitation trade-off via the parameter $k_{opt}$, enabling models to solve more complex reasoning and coding tasks that are otherwise inaccessible to standard pass@1 optimization.

The work positions itself as a generalization and variance reduction of existing order-statistic policy gradient methods, offering a practical pathway to scale RL training for LLMs on tasks requiring deep exploration.