Pass@K Policy Gradients for RL & LLM Finetuning

• Pass@K policy gradients are defined as the probability that at least one of K independent samples achieves success, transforming the traditional single-sample reward into a set-based objective.
• They reweight the standard REINFORCE gradient with a factor of K(1-ρ)^(K-1), effectively addressing variance and credit assignment challenges in complex reward settings.
• Applied in RLVR and LLM finetuning, these methods improve exploration, preserve output diversity, and mitigate mode collapse in structured prediction and program synthesis tasks.

Pass@K policy gradients are methods in reinforcement learning that directly optimize for the probability of achieving at least one success in a set of K samples or rollouts from a stochastic policy, rather than optimizing only the expected reward of individual samples. This concept, originally motivated by the evaluation protocols for structured prediction and program synthesis models, is now fundamental in the training of LLMs and generally in RL with verifiable rewards (RLVR) for complex reasoning. The transition from pass@1 (expected success on a single sample) to pass@K as a training objective introduces nontrivial changes in the optimization landscape, estimator construction, and exploration-exploitation balance.

1. Conceptual Overview and Formal Definition

The Pass@K metric is defined as the probability, under the current policy πθ, that at least one sample in a set of K independently drawn outputs achieves the target criterion. For binary rewards, this is: $$\text{pass@}K = \mathbb{P}\left[ \bigvee_{i=1}^K f(x_i) = 1 \right] = \mathbb{E}\left[ 1 - \prod_{i=1}^K (1-f(x_i)) \right]$$ where $x_1,\ldots,x_K \sim p(x|\theta)$ i.i.d., and $f(x) \in \{0,1\}$ is the indicator of success.

For tasks with continuous rewards, the analogous objective is $\max g@K = \mathbb{E}[ \max (g(x_1),...,g(x_K)) ]$.

In the context of RLVR and LLMs, pass@K is empirically measured by generating K model outputs for each input and computing the frequency with which any output is accepted as correct by an automatic checker.

2. Policy Gradient Formulation for Pass@K

The key mathematical property exploited for Pass@K policy gradient methods is that the pass@K objective can be expressed as a monotonic transformation of the expected per-sample success probability $\rho_\theta$:

$$\rho_{K,\theta}(x,a) = 1 - (1 - \rho_\theta(x,a))^K$$

where $$\rho_\theta(x,a) = \mathbb{E}_{y \sim \pi_\theta(\cdot|x)}[r_{0/1}(y, a)]$$.

The gradient of pass@K w.r.t. θ is: $$\nabla_\theta \rho_{K,\theta}(x,a) = K (1 - \rho_\theta(x,a))^{K-1} \nabla_\theta \rho_\theta(x,a)$$ This expresses that the pass@K policy gradient is a weighted version of the standard REINFORCE gradient, where the weight reflects the probability that the first $K-1$ samples fail. Empirically, this leads to greater upweighting of rewards for "hard" (low-$\rho_\theta$) instances as K increases.

Correct unbiased estimators for pass@K and its gradient are provided for both binary and continuous reward settings (Walder et al., 21 May 2025). For the binary case, with $n \geq K$ samples and $c$ correct, the pass@K metric is unbiasedly estimated by: $$\rho(n, c, K) = 1 - \frac{\binom{n-c}{K}}{\binom{n}{K}}$$ The associated policy gradient estimator distributes credit to individual actions using leave-one-out and subset-based weights, enabling direct integration into standard RL algorithms.

For continuous rewards, efficient order-statistic-weighted estimators and corresponding gradients are derived.

3. Variance Reduction Techniques and Efficient Estimators

Variance is a major issue when estimating set-based metrics such as pass@K via Monte Carlo policy gradient. Several advances have led to provably unbiased, low-variance estimators suited for practical RLVR finetuning:

• Leave-One-Out (LOO) Estimator: Baseline-subtracted estimator analogous to the REINFORCE-with-baseline trick, adapted to the combinatorics of pass@K subsets.
• LOO Minus One Estimator: Further variance reduction by leveraging subset constructions omitting not only the focal sample but also one element from the pass@K set (Walder et al., 21 May 2025).

Efficient algorithms allow for complete evaluation of all K-subsets within each batch, supporting arbitrary K ≤ n and computationally scaling as $O(k + n \log n)$.

4. Integration into RL Training Pipelines

Pass@K policy gradient estimators can be seamlessly plugged into general policy gradient frameworks, including REINFORCE, PPO, and GRPO-style objectives:

• The base reward vector over n samples is replaced with its pass@K-transformed counterpart.
• Policy gradient updates are performed as usual, treating the transformed weights as the per-sample reward.
• Estimator computations can be batched over RL microbatches and are compatible with reward shaping, advantage normalization, and entropy/KL regularization.

Annealing k strategies are supported: training can begin with high k to encourage exploration (maximize exploratory pass@K), and then gradually lower k towards 1 to refine exploitation (maximize pass@1), providing smooth interpolation between broad search and precision (Walder et al., 21 May 2025, Chen et al., 14 Aug 2025).

5. Advantage Shaping and Surrogate Reward Connections

There is a unification between direct pass@K REINFORCE optimization and advantage shaping methods that implicitly maximize smooth surrogates or regularized versions of pass@K (Thrampoulidis et al., 27 Oct 2025). For example, by applying a variance-stabilizing transformation $$F(\rho_{K,\theta}) = \frac{2}{K}\arcsin(\sqrt{\rho_{K,\theta}})$$, practical GRPO-style advantage shaping is shown to optimize a smoothed surrogate of the strict pass@K metric, leading to robust and more stable credit assignment, especially for rare-event instances.

Regularization at the reward level (e.g., upweighting hard examples or including entropy terms) can be recast as modifying the effective pass@K objective, balancing exploration and exploitation while maintaining interpretable gradual transitions across task difficulty (Thrampoulidis et al., 27 Oct 2025).

A formal "recipe" for deriving new advantage shaping rules from surrogate reward objectives is established (select a smooth monotonic transformation F, differentiate with respect to θ, substitute empirical batch means, use variance-reduced estimators) (Thrampoulidis et al., 27 Oct 2025).

6. Exploration–Exploitation Tradeoff in Pass@K Optimization

Pass@K training provably encourages and sustains greater exploration. Empirical evidence confirms several phenomena (Chen et al., 14 Aug 2025, Peng et al., 16 Oct 2025):

• Standard pass@1-focused RLVR methods rapidly collapse output entropy, causing sharp probability concentration and reduced diversity (mode collapse).
• Pass@K training, either by direct reward transformation or token-level adjustments (e.g., SimKO (Peng et al., 16 Oct 2025)), maintains higher entropy and diversity specifically at "semantic forks" in generation, critical for solving harder reasoning problems.
• Analytical advantage functions can be constructed to explicitly maximize exploration on unsolved problems, suppressing overfitting to "easy" cases and facilitating transfer to previously unsolved or out-of-domain tasks.
• Correctly tuned pass@K optimization achieves strong or no-loss pass@1 while significantly boosting pass@K, in contrast to prior methods that incur trade-offs.

7. Empirical Results and Practical Recommendations

• Pass@K policy gradient methods (PKPO, Pass@K Training, SimKO) consistently and robustly outperform independent-sample RLVR baselines under challenging reasoning and code-generation benchmarks, both in cumulative solve-rate and diversity metrics (Walder et al., 21 May 2025, Chen et al., 14 Aug 2025, Peng et al., 16 Oct 2025).
• Annealing or dynamically selecting k can yield models that match the pass@1 of vanilla RLVR/GRPO while exceeding them by large margins in pass@K for K > 1 (Walder et al., 21 May 2025).
• Fine-grained RLVR adjustments at the token or action level (e.g., SimKO) can solve the mode-collapse problem without harming exploitation and are more robust to OOD shifts (Peng et al., 16 Oct 2025).
• All transformation recipes and implementations can be realized as drop-in replacements or minor code augmentations within typical RL pipelines.

Estimator Type	Bias	Variance	K arbitrary?	Works with RL
Naive MC (independent)	Unbiased	High	Yes	Standard PG
PKPO (LOO-l1)	Unbiased	Low	Yes	Any PG variant
GRPO-K/Bydance	Surrogate	Lower	Yes	GRPO-style
SimKO	Surrogate/Empirical	Lowest for LLM RL	Yes	Token-level RLVR

8. Theoretical Foundations and Generalization

The mathematical backbone of pass@K policy gradients is the classical policy gradient theorem (Kämmerer, 2019), bootstrapped to set-based objectives via transformations of the policy-induced marginal success probability. All related convergence, monotonic improvement, and regularization results for expected-reward RL carry over to pass@K maximization when the metric is a monotonic function of the expected reward (as is the case for pass@K) (Papini et al., 2019). Analytical advantage functions for pass@K, derived explicitly from group statistics, ensure stable training and open a path for generalizable advantage/credit-shaping in RLVR (Chen et al., 14 Aug 2025, Thrampoulidis et al., 27 Oct 2025).

9. Conclusion

Pass@K policy gradients represent a principled and tractable generalization of policy gradient RL for optimizing non-decomposable, set-based success metrics. Recent algorithmic advances provide efficient, unbiased, and low-variance estimators, theoretically grounded surrogate and advantage shaping frameworks, and robust empirical recipes for integrating these methods in RLVR and LLM finetuning scenarios. The approach bridges the gap between practical evaluation in structured prediction and theoretical RL optimization, ensuring both diversity and high performance on complex, hard-to-solve tasks that cannot be addressed by pass@1 training alone (Walder et al., 21 May 2025, Chen et al., 14 Aug 2025, Thrampoulidis et al., 27 Oct 2025, Peng et al., 16 Oct 2025).