PAIR: Pairwise-Aware Inclusion Reweighting for Adaptive Rollout Allocation in RLVR

PAIR: Pairwise-Aware Inclusion Reweighting for Adaptive Rollout Allocation in RLVR

Summary

This paper introduces PAIR (Pairwise-Aware Inclusion Reweighting), a method for adaptive rollout allocation in reinforcement learning with verifiable rewards (RLVR). The authors identify a fundamental statistical mismatch in existing adaptive allocation methods: while most approaches treat rollout utility as a pointwise quantity (attached to a prompt, rollout, prefix, or token), the unclipped leave-one-out group-relative score gradient is actually a second-order U-statistic over pairs of rollouts. Specifically, for a prompt q, response reward ri, and sequence score si = ∇θ log πθ(oi q), the gradient equals the average of hij = ½(ri − rj)(si − sj) over all i < j. This pairwise structure means that completing one rollout reveals contrast with every other completed rollout, and adaptive endpoint selection changes which pair terms are observable.

The paper's key insight is that generation cost is paid per rollout endpoint (linear in vertices), while statistical value is quadratic and graph-coupled (edges between completed endpoints). This creates a contrast-graph formulation where short rollout prefixes are vertices and pair-gradient terms are edges. PAIR operates in four stages: (1) generating short independent prefixes for candidate rollouts, (2) using lightweight heads to predict correctness and remaining suffix cost from prefix states, (3) solving a convex program to assign strictly positive continuation probabilities under an expected suffix-token budget, and (4) after randomized continuation, using every edge induced by completed vertices and dividing each kernel by the edge's logged joint inclusion probability.

The theoretical contributions include: proving pairwise equivalence between the leave-one-out estimator and the pair average (Theorem 1); establishing design-unbiasedness of the induced-edge estimator under conditionally independent on-policy rollouts and logged joint inclusion probabilities (Theorem 2); deriving the conditional bias of unweighted adaptive selection and the exact design covariance (Proposition 1); and proving convexity of the allocation surrogate (Proposition 2). The authors explicitly scope their guarantees to the unclipped, unstandardized leave-one-out score gradient, noting that PPO clipping, reward standardization, and weight stabilization are practical approximations studied separately.

Experiments on Qwen3-1.7B and Qwen3-4B compare PAIR against GRPO, DPPO, VIP, HORA, VIGOR, and DUET under compute-matched conditions. On Qwen3-1.7B, PAIR achieves 49.2% average accuracy versus 48.0% for DUET and 45.0% for full-group GRPO, while using 0.49× the GRPO token budget and 42% less wall-clock time. On Qwen3-4B, PAIR attains 57.1% average accuracy at 0.48× tokens, exceeding DUET by +1.4 absolute points. Gains are consistent across mathematics (MATH500, AIME24, AMC23, OlympiadBench) and code (LiveCodeBench) benchmarks.

A frozen finite-population estimator audit with 50,000 randomized design draws confirms the central statistical finding: adaptive unweighted selection achieves low MSE (0.050) but with substantial bias (2.27×10−2), meaning it optimizes a different estimand. Uniform and pointwise vertex HT estimators remain nearly unbiased but leave MSE at 0.154 and 0.174. PAIR recovers near-zero bias (1.90×10−3) with MSE 0.145, improving on both unbiased baselines at matched cost. Mechanism ablations show that removing pair-inclusion correction costs 2.1 accuracy points and raises gradient MSE from 0.145 to 0.218; marginal-only weights recover only part of the gap (MSE 0.195), confirming that joint inclusion—not merely per-rollout propensity—is required.

The paper discusses limitations including: the exact guarantee applies only to the on-policy, sequence-level, unclipped, unstandardized objective; prefix predictors may drift with the policy; severe miscalibration can create high-variance inverse weights; independent Bernoulli continuation controls expected rather than exact token cost; and gains may shrink for short responses or near-deterministic groups. The authors also note that the convex objective is a valid upper bound only under the stated domination condition on edge proxies, and with learned proxies it becomes a surrogate for unknown design variance.

Based on the paper, here are the specific improvements I can make to AI systems:

1. Implement pairwise-aware gradient estimation in RL fine-tuning

• Replace pointwise reward-weighting (e.g., per-prompt or per-token) with edge-based contrast terms h ij = ½(r i − r j)(s i − s j) over all completed rollout pairs.

• This corrects the statistical mismatch that biases current adaptive allocation methods (GRPO, DPPO, VIP, etc.) toward a different estimand, reducing gradient bias from 2.27×10−2 to 1.90×10−3 in finite-population audits.

2. Add joint-inclusion probability weighting for adaptive rollout selection

• When adaptively terminating or continuing rollouts, log the joint inclusion probability of each pair of completed rollouts and divide each pair-gradient kernel by that logged probability.

• This removes the conditional bias introduced by adaptive endpoint selection, which unweighted or marginal-only propensity methods fail to correct (MSE improves from 0.218 to 0.145 in ablations).

3. Build a contrast-graph-based rollout allocator

• Generate short independent prefixes first, then use lightweight correctness/suffix-cost predictors to solve a convex program that assigns strictly positive continuation probabilities under a token budget.

• This yields a 0.49× token reduction vs. full-group GRPO while improving accuracy (e.g., 49.2% vs. 45.0% on Qwen3-1.7B), and reduces wall-clock time by 42%.

4. Enable unbiased gradient estimation under compute constraints

• Use the induced-edge estimator with logged joint inclusion probabilities to guarantee design-unbiasedness (Theorem 2) for the unclipped leave-one-out score gradient, even when only a subset of rollouts is completed.

• This allows the AI system to train on fewer rollouts per prompt without sacrificing gradient fidelity, making RLVR feasible for larger models or tighter latency budgets.

5. Provide a diagnostic tool for allocation bias

• Expose the conditional bias formula (Proposition 1) as a monitoring metric during training. The system can flag when adaptive selection is drifting from the true estimand, enabling early correction or fallback to uniform allocation.

The improved AI system can:

• Fine-tune reasoning and code-generation models with verifiable rewards (math, code) using 50% fewer tokens and 40% less time, while achieving higher accuracy than current state-of-the-art adaptive methods (e.g., +1.4 points over DUET on Qwen3-4B).

• Maintain unbiased gradient estimates even when rollouts are adaptively pruned or extended, preventing silent degradation from selection bias.

• Scale to larger models or longer response sequences without proportional compute growth, by allocating rollout effort based on predicted value rather than uniform sampling.

Reinforcement learning with verifiable rewards (RLVR) spends most of its compute generating groups of long reasoning trajectories. Recent allocators reduce this cost by assigning budgets to prompts, rollouts, or tokens according to a pointwise notion of difficulty or utility. We identify a statistical mismatch: the unclipped leave-one-out group-relative score gradient is not a sum of independent point contributions, but a second-order U-statistic over pairs of rollouts. Completing one rollout therefore reveals contrast with every other completed rollout, and adaptive endpoint selection changes which pair terms are observable. We introduce PAIR (Pairwise-Aware Inclusion Reweighting), which treats short rollout prefixes as vertices and pair-gradient terms as edges of a contrast graph. A prefix-only predictor estimates correctness and remaining token cost; a convex design chooses positive continuation probabilities under an expected suffix-token budget; and each edge induced by completed vertices is inverse-weighted by its logged joint inclusion probability. Under conditionally independent on-policy rollouts and an unclipped, unstandardized objective, the resulting estimator is design-unbiased for the complete candidate-pair gradient. Across compute-matched RLVR runs on Qwen3-1.7B/4B, PAIR improves average accuracy by +1.2 and +1.4 over the strongest pointwise allocator while using 51% and 52% fewer generated tokens than full-group GRPO. A frozen-population estimator audit confirms that unweighted adaptive selection is biased, whereas pair-inclusion correction recovers the complete-pair target at matched suffix cost.