Technical Summary: On Advantage Estimates for Max@K Policy Gradients

Problem Statement

Reinforcement learning with verifiable rewards (RLVR) is a dominant paradigm for post-training large language models (LLMs) on reasoning tasks. However, these settings typically rely on sparse, outcome-based rewards (e.g., a binary correct/incorrect signal), which makes exploration difficult in large combinatorial spaces. While methods like entropy regularization and count-based bonuses exist, they introduce auxiliary objectives that require careful tuning.

A complementary approach is to directly optimize inference-time objectives such as pass@K (the probability that at least one of $K$ sampled completions is correct) and its continuous generalization, max@K (the expected maximum reward among $K$ samples). Several recent works have proposed unbiased policy-gradient (PG) estimators for these objectives. However, the relationships between these estimators remain unclear. Specifically, existing estimators often utilize different signals, baselines, and normalizations. A critical issue identified in this work is that leading estimators (e.g., those based on Expected Improvement) are unbiased as PG estimators but yield non-centered advantages. A non-centered advantage (one with a non-zero mean) corresponds to an underestimated baseline, which can steer optimization toward suboptimal actions and increase gradient variance.

Methodology

The paper proposes MaxPO (Max@K Policy Optimization), a method that introduces a Leave-Two-Out (L2O) baseline to correct the advantage estimation for max@K objectives.

1. Analysis of Existing Estimators

The authors revisit the Expected Improvement (EI) estimator used in prior work (e.g., Walder and Karkhanis [54]). They demonstrate that while the EI-based estimator is an unbiased PG estimator, the resulting advantage signal $s(a)$ is strictly non-negative (due to the ReLU operation in the EI definition). Consequently, $E[s(a)] > 0$, meaning the advantage is not centered. This lack of centering leaves room for variance reduction.

2. The Leave-Two-Out (L2O) Baseline

To achieve an exactly centered advantage while preserving PG unbiasedness, the authors introduce the L2O baseline.

• The Challenge: A standard Leave-One-Out (L1O) baseline fails in the max@K setting because the reward of a sample $i$ can implicitly influence the advantage calculation of other samples $j$ (as $i$ acts as a potential comparator in the max operation). This dependence violates the independence condition required for unbiased PG estimation.
• The Solution: The L2O baseline $b_{-i}^{L2O}$ is computed by averaging the Expected Improvement scores of all other samples $j$, where each score is computed over a subset of comparators that excludes both the target sample $i$ and the sample $j$ itself.
• Properties:
  • Independence: By excluding $i$ from the comparator set for all $j$, the baseline $b_{-i}^{L2O}$ is independent of the reward $r_i$, ensuring the PG estimator remains unbiased.
  • Centering: The authors prove that the realized batch-average advantage is exactly zero for every sampled batch: $\frac{1}{B}\sum_{i=1}^B (s_i - b_{-i}^{L2O}) = 0$.
  • Efficiency: While a naive computation would be $O(B^3)$, the authors derive a vectorized implementation using ReLU-based formulations that computes the L2O baseline in $O(B^2)$ time, making it feasible for GPU-accelerated LLM training.

3. Canonical Finite-Batch Advantage

The paper derives a canonical finite-batch advantage for max@K, defined as $u_i - v_i$:

• $u_i$: The conditional expected max@K value of a group containing response $i$.
• $v_i$: The leave-one-out expected max@K value computed from groups excluding $i$.
  The authors show that the MaxPO estimator (EI + L2O) exactly recovers this canonical signal ($s_i - b_{-i}^{L2O} = u_i - v_i$). This provides a unified view of existing estimators, classifying them into "uncentered" (e.g., EI-only/PKPO) and "(normalized) canonical" (e.g., MaxPO, and certain analytical pass@K variants).

4. Integration into LLM Training

MaxPO is integrated into group-based RL algorithms (like GRPO). For each prompt, a group of $G$ outputs is sampled. The L2O baseline is computed within each group to generate sequence-level advantages, which are then broadcast to token-level losses.

Key Contributions

1. Identification of Non-Centered Advantages: The paper identifies that leading max@K PG estimators (specifically EI-based ones) are unbiased as gradient estimators but fail to produce centered advantages, leading to suboptimal variance properties.
2. Leave-Two-Out (L2O) Baseline: A novel baseline design that preserves PG unbiasedness in the non-linear max@K setting while ensuring the realized batch advantage is exactly zero-mean.
3. Canonical Unification: The derivation of the canonical finite-batch advantage ($u_i - v_i$) and the proof that MaxPO recovers this signal, offering a unified framework to compare and categorize existing pass@K/max@K estimators.
4. Efficient Implementation: An $O(B^2)$ vectorized algorithm for computing the L2O baseline, enabling practical application in LLM post-training.

Experimental Results

The authors validate the method through controlled bandit/maze experiments and LLM reasoning tasks.

• Variance Reduction: In bandit environments, the L2O baseline significantly reduces gradient variance compared to raw EI estimators, with the effect growing as the action space size increases. In LLM training (Llama-3.2-3B-Instruct), the method reduced the gradient variance proxy by 77.4% compared to PKPO (the EI-only baseline).
• Performance Gains: On five math reasoning benchmarks (AIME24, AIME25, AMC23, MATH500, Minerva), MaxPO outperformed strong baselines including PKPO, GRPO, and Entropy-Adv.
  • On Qwen2.5-Math-7B, MaxPO improved task-average pass@256 by 5.2% over PKPO.
  • On Llama-3.2-3B-Instruct, it improved pass@256 by 2.4% over PKPO.
• Inference Compute Scaling: The performance gains of MaxPO were most pronounced at higher inference compute levels (larger $k$ in pass@k), suggesting the method helps the policy better utilize additional test-time samples.

Significance and Claims

The paper claims that MaxPO provides a principled foundation for designing and comparing policy-gradient estimators for inference-time objectives like max@K. By identifying the canonical advantage form and correcting the bias in existing estimators, the work offers a method to improve the efficiency and stability of post-training reasoning models.

The authors modestly note that their analysis focuses on finite-batch, on-policy estimation under i.i.d. sampling. They acknowledge limitations regarding off-policy training, adaptive sampling, and settings with correlated generations as directions for future work. Furthermore, they highlight that while the method improves optimization, better optimization of repeated-sampling reasoning objectives carries broader impacts, potentially strengthening models that could be misused for deceptive assistance or automated problem solving.