Pass@K Policy Optimization: Multi-Sample RL

• Pass@K Policy Optimization is a reinforcement learning framework that maximizes the chance of success by optimizing the probability of at least one correct output among k samples.
• It employs reward transformations and analytical advantage shaping to reduce gradient variance and mitigate issues like vanishing gradients and mode collapse.
• Empirical results indicate improved multi-sample performance and diversity, though careful tuning is needed to balance single-sample accuracy with broader exploration.

Pass@K Policy Optimization (PKPO) refers to a class of reinforcement learning algorithms and reward transformations designed to maximize the chance that at least one of $k$ independent samples from a learned policy succeeds—where “success” is measured by an external verifier or a binary reward signal. PKPO directly optimizes for multi-sample success (the pass@$k$ metric) as opposed to the traditional focus on single-shot accuracy (pass@$1$), addressing the exploration-exploitation limitations of standard RL with verifiable rewards, and supporting reasoning tasks such as mathematical problem solving and code synthesis with LLMs.

1. The Pass@$k$ Objective and Its Policy Gradient

The pass@$k$ metric is defined for a policy $\pi_\theta$ and input $x$ as the probability that at least one of $k$ i.i.d. samples yields a correct output: $\mathrm{Pass@}k(x) = 1 - (1 - \rho_x)^k$ where $\rho_x = \sum_{o\in O(x)} \pi_\theta(o|x)$ is the probability mass assigned to correct answers for $k$0. This formulation generalizes to continuous rewards and subsumes pass@$k$1 as a special case.

Direct optimization of pass@$k$2 via expected reward encourages probability concentration—policy mass collapses onto the maximal-reward mode—which results in poor diversity and under-utilized sampling capacity for larger $k$3 (Walder et al., 21 May 2025). The pass@$k$4 objective, in contrast, is non-linear and saturating: it is optimized not just by high mean reward, but by maintaining sufficient probability mass on multiple correct or useful outputs.

The policy-gradient for pass@$k$5 can be written as (Thrampoulidis et al., 27 Oct 2025, Chen et al., 14 Aug 2025, Le et al., 30 Jan 2026): $k$6 where $k$7 is the 0/1 correctness of the output $k$8. Thus, policy-gradient updates under pass@$k$9 act as positive reweightings of the base pass@$1$0 gradient, particularly upweighting “hard” prompts where the base policy's success probability is low (Yu, 20 Nov 2025).

2. Motivations, Limits, and Failure Modes of Naive Pass@$1$1 Optimization

A naively applied pass@$1$2 objective is subject to several notable shortcomings:

• Vanishing Gradients at Extremes: When $1$3 (hard prompts), empirical gradients disappear because correct outputs are almost never seen; as $1$4 (easy prompts), the multiplicative factor $1$5, so little updating occurs (Yu, 20 Nov 2025). Thus, the learning signal vanishes exactly in the regimes where exploration or marginal refinement are most needed.
• Exploration Collapse: Focus on maximizing pass@$1$6 with a standard policy-gradient causes excessive mode concentration (diversity collapse), especially as training proceeds. Multi-sample performance (pass@$1$7) converges to single-sample performance (pass@$1$8) as the probability mass collapses onto a single solution (Le et al., 30 Jan 2026, Yu, 20 Nov 2025, Walder et al., 21 May 2025).
• Gradient Conflict with Pass@1: PKPO's implicit reweighting toward “hard” prompts can result in a gradient that is anti-aligned with the pass@$1$9 gradient, especially when those prompts exhibit negative parameter interference (i.e., improvements for hard prompts degrade performance on easy ones) (Barakat et al., 24 Feb 2026).

In summary, naive pass@$k$0 policy-gradient optimization risks inactivity on hard cases, over-concentration on a dominant solution, and can degrade single-sample accuracy (Yu, 20 Nov 2025, Barakat et al., 24 Feb 2026). Addressing these issues has motivated a cascade of refined PKPO algorithms.

3. Practical PKPO Algorithms: Reward Transformations and Advantage Shaping

Contemporary PKPO methods introduce explicit reward transformations and advantage-shaping procedures to directly and stably optimize the non-linear pass@$k$1 objective.

Reward Transformations

Modern PKPO constructs groupwise, low-variance, unbiased reward transformations that admit standard policy-gradients (PPO, GRPO), directly maximizing expected “best-of-$k$2” reward (Walder et al., 21 May 2025, Chen et al., 14 Aug 2025). For $k$3 samples per prompt, an unbiased estimator for the pass@$k$4 group reward is: $k$5 where $k$6 is the count of correct samples.

To minimize variance, leave-one-out (LOO) and LOO-minus-one baselines are employed. These reward vectors are fed into the policy-gradient update in lieu of per-sample rewards, making plug-and-play integration with existing RL frameworks possible.

Analytical Advantage Functions

Efficient, closed-form expressions for per-sample advantages are derived via combinatorial analysis of group outcomes (Chen et al., 14 Aug 2025): $k$7 Here, $k$8 and $k$9 are the group-level mean and std of rewards, and the formula ensures correct credit assignment for both positive and negative samples without explicit enumeration of groupings.

Variance Reduction and Sample Efficiency

Adoption of these analytical and LOO-based baselines constrains gradient estimator variance, reducing wasted updates and stabilizing RLVR training even at high $k$0 (Walder et al., 21 May 2025).

4. Extensions: Exploration, Exploitation, and Hybrid Policies

PKPO frameworks admit customization for balancing exploration (diversity) and exploitation (greedy accuracy):

• Annealing $k$1: Progressive scheduling, e.g., training with high $k$2 early for exploration followed by low $k$3 (or $k$4) for exploitation, improves both pass@$k$5 and pass@$k$6 (Walder et al., 21 May 2025).
• Advantage Shaping: Analytical or hand-crafted advantage functions allow shifting the “peak gradient attention” onto problem difficulty regimes of interest (Chen et al., 14 Aug 2025).
• Surrogate Reward Maximization: Advantage-shaping heuristics (e.g., hard-example upweighting, entropy or uncertainty regularization) are interpretable as regularization at the reward level and can be derived systematically via forward- or reverse-engineering of surrogate objectives (Thrampoulidis et al., 27 Oct 2025).
• Hybrid/Interpolated Objectives: Leveraging combinations of pass@$k$7 and pass@$k$8-style advantages or risk-sensitive objectives to maintain pass@$k$9 while improving diversity (Chen et al., 14 Aug 2025, Barakat et al., 24 Feb 2026).
• SimKO and Transform-Augmentation: Methods such as SimKO redistribute positive mass among top-$\pi_\theta$0 alternatives for correct tokens and apply asymmetric penalties to top-1 candidates on incorrect tokens, discouraging over-concentration and preserving answer diversity (Peng et al., 16 Oct 2025). Transform-augmented approaches like TA-GRPO pool advantages over semantically equivalent question variants, counteracting diversity collapse and gradient-diminishing regimes while also providing robustness to phrasing shift (Le et al., 30 Jan 2026).

5. Theoretical Properties, Gradient Analysis, and Trade-Offs

PKPO methodologies are grounded in theoretical guarantees, unbiasedness, and variance minimization, but also reveal potential trade-offs:

• Unbiasedness and Variance Reduction: All major PKPO reward transformations are provably unbiased estimators of the gradient of the pass@$\pi_\theta$1 population objective (Walder et al., 21 May 2025). LOO/LOO-1 baselines achieve minimal variance among known strategies.
• Gradient Alignment and Conflict: For pass@$\pi_\theta$2, the per-example gradient is a scaled version of the pass@$\pi_\theta$3 gradient, but with a scaling factor that vanishes at the regime’s extremities (Yu, 20 Nov 2025). When hard prompts exhibit negative prompt interference (parameter gradients for them are negatively correlated with gradients for the rest), PKPO can degrade pass@$\pi_\theta$4 (Barakat et al., 24 Feb 2026). Recommendations include monitoring gradient alignment and designing hybrid or tempered reweighting schemes.
• Exploration–Exploitation Curve: The “attention” curve (update magnitude as a function of prompt difficulty) for standard RL peaks at 50% accuracy, focusing updates on mid-difficulty prompts, while analytical PKPO can shift this peak according to $\pi_\theta$5, focusing updates where extra exploration is most impactful (Chen et al., 14 Aug 2025).

6. Empirical Results and Observed Performance Advantages

Experiments on reasoning and code-generation benchmarks demonstrate consistent gains in multi-sample success, improved model entropy, and diversity when employing PKPO and its variants (Walder et al., 21 May 2025, Le et al., 30 Jan 2026, Peng et al., 16 Oct 2025, Chen et al., 14 Aug 2025). Notable findings include:

Method	Pass@1 (%)	Pass@k (k=8/16/32) (%)	Diversity/Entropy Trends
Standard RLVR	Baseline	Limited at high k	Probability mass collapse; low diversity
PKPO	+5–7 pts	+10–20 pts	Higher entropy, success on harder prompts
SimKO	≈+1–2 pts	+2–5 pts	Blocks over-concentration; maintains modes
TA-GRPO (N=3)	+9.84	+8.69 (AIME24), +5.05 (GPQA)	Maintains solution strategies; reduced zero-gradient probability
APO	+0.8 Pass@1	+2.3 at Pass@16	Breaks accuracy-diversity trade-off

On competitive math and scientific reasoning benchmarks, e.g., Qwen3-1.7B on MATH with TA-GRPO, Pass@32 improves by +9.84 points over GRPO and +5.05 on out-of-distribution science tasks (Le et al., 30 Jan 2026). SimKO and APO demonstrate that tailored support coverage and selective mass re-inflation can further mitigate the diversity collapse seen in vanilla RLVR (Wang et al., 5 Feb 2026, Peng et al., 16 Oct 2025).

7. Limitations, Recommendations, and Open Directions

The principal challenges for PKPO entail managing the trade-off between pass@$\pi_\theta$6 and pass@$\pi_\theta$7—in particular, avoiding degradation of single-sample accuracy due to negative prompt interference or over-emphasis on hard but idiosyncratic samples (Barakat et al., 24 Feb 2026). Practically, monitoring gradient alignment and adjusting $\pi_\theta$8 or advantage shaping is recommended (Barakat et al., 24 Feb 2026).

Other research frontiers include:

• Automated validation of semantically-equivalent transformations in TA-GRPO (Le et al., 30 Jan 2026).
• Adaptive selection of the number of transformations or top-$\pi_\theta$9 for support coverage (Wang et al., 5 Feb 2026, Le et al., 30 Jan 2026).
• Extensions to non-binary or multi-answer tasks, such as code generation using continuous reward proxies (e.g. pass rates on unit tests) (Walder et al., 21 May 2025).
• Integrating PKPO with entropy regularization, determinantal point process diversity bonuses, or risk-sensitive interpolated objectives for robust multi-objective optimization (Yu, 20 Nov 2025, Thrampoulidis et al., 27 Oct 2025).
• Scaling investigations on larger models and more complex domains, including proof synthesis and multi-agent dialogue (Le et al., 30 Jan 2026).

PKPO constitutes a principled framework for multi-sample success maximization in reinforcement learning with verifiable rewards, grounded in unbiased gradient estimators and supported by a growing toolkit of advantage shaping and data augmentation strategies. It provides state-of-the-art improvements in exploration, diversity, and generalized solution coverage, while spotlighting the nuanced interplay between groupwise optimization and single-sample reliability.