Dataset Mismatch Matters in Group Relative Policy Optimization for Reinforcement Learning from Verifiable Rewards

AI summary

Overview

Research area: Reinforcement learning for large language model reasoning — specifically reinforcement learning from verifiable rewards (RLVR) and the GRPO training algorithm.

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods, GRPO, bi-level optimization, and online-learning regret analysis.

Scope: The paper introduces learning-to-reweight rollout (L2RR), a method that uses limited, aggregated feedback from a small held-out target dataset to learn which source prompts to sample during GRPO training, aiming to improve target-domain accuracy while cutting rollout cost.

What This Paper Is About

GRPO trains a language model by generating groups of responses per prompt and normalizing rewards within each group; when every response in a group gets the same verifier score, the advantages collapse to near zero and the rollout is wasted. Existing selective-rollout methods decide which prompts to skip using only source-side training statistics, but the source training set is rarely distributed like the target test set, so those decisions can be misaligned with target performance. The authors ask whether a small amount of aggregated target-side feedback — obtained without sharing target prompts or rollout traces — can be used to learn a source-prompt sampling distribution that both improves target accuracy and reduces wasted rollouts.

Key Contributions

1. A target-guided selective rollout formulation for GRPO. L2RR learns a time-varying discrete sampling distribution over source prompts, posed as a bi-level objective (Eq. 4) that maximizes accuracy on a held-out target set while penalizing rollouts that yield near-zero group-normalized advantages.

2. A privacy-restricted feedback interface. The sampler update needs only directional derivatives of the form ⟨g^tgt_{t+1}, g_{t,i} − ḡ_t⟩, which the target owner can compute locally without revealing target prompts or rollout traces; formal privacy guarantees are explicitly left to future work.

3. A rollout-allocation rule derived from the theory. Since the target-loss bound tightens as the variance of the target-guided gradient estimate falls, the authors allocate a per-round rollout budget RO_t = BG·clip(s̄_t, ρ_min, 1) across prompts in proportion to w_{i,t}^{3/2}, with a floor G_min and cap G (Eq. 8).

4. A regret analysis. Theorem 1 gives an average sampler regret bound of order O(T^{-1/2}) (specifically (log N_S/a + aH²/2)/√T) under a bounded-second-moment assumption and η_q = a/√T, with an additive λ_T term when the target-feedback losses are used; choosing λ_T ≤ c_λ/√T preserves the rate.

Main Findings

• Highest average accuracy in every reported configuration. Across six benchmarks (Math500, AIME24, AMC, Gaokao, Minerva Math, Olympiad Bench), three models, and two source sets, L2RR has the best average. On Qwen2.5-Math-1.5B with DAPO+MATH, average rises from 47.7% (DAPO) to 50.3% while rollouts fall from 7.6M to 3.0M.

• Better accuracy-per-rollout than source-only selective rollout. On Qwen2.5-Math-7B with DAPO+MATH, L2RR reaches 58.8% average with 6.0M rollouts, versus 57.0% / 6.3M for GRESO and 57.3% / 6.6M for AR3PO. On DeepSeek-R1-Distill-Qwen-1.5B with Open-R1, L2RR reaches 58.1% at 1.5M rollouts versus 55.7% at 2.4M for DAPO.

• The learned sampler is the most important component. Ablating to a uniform sampling distribution gives the largest and most consistent accuracy drop (50.9% → 48.9% average); removing the rollout-waste term drops accuracy to 49.5%; removing rollout allocation drops the RO ratio from 2.5× to 1.8×.

• Small held-out target sets already help. Average accuracy rises monotonically with the held-out target fraction: 49.3% at 4%, 49.7% at 6%, 49.9% at 8%, and 50.9% at 10%, with no saturation indicated in the studied range.

• Efficiency gains show up in training dynamics. After roughly the first 40 steps L2RR maintains a higher effective-prompt ratio than dynamic sampling, overtakes GRESO after about 550 steps, and after about 600 steps has lower and more stable per-step rollout time than GRESO. Over 1000 steps total rollout time is about 24 hours for L2RR, about 25 hours for GRESO, and 41 hours for dynamic sampling.

• Reported mismatch between abstract and setup. The abstract refers to "three model scales," while the experimental section describes three models spanning two parameter scales (1.5B and 7B).

Methodology in Plain English

The setup separates two parties. The trainer holds the source dataset and runs GRPO as usual. The target data owner holds a small held-out target set and never releases the prompts — it only returns aggregate feedback derived from that set.

Training alternates between two updates. First, a batch of source prompts is drawn from the current sampling distribution, rollouts are generated and scored, and the policy takes a normal GRPO step. Second, the sampling distribution is adjusted with an exponentiated-gradient step using two signals added together: a target-side signal, computed by asking how much each prompt's gradient moves the policy in a direction that improves held-out target accuracy (approximated with a one-step look-ahead), and a waste signal that penalizes prompts which have repeatedly produced identical rewards across their most recent rollouts, tracked by a counter z_{i,t} and a penalty 1 − p_e^{z_{i,t}}. Prompts are resampled until B prompts with non-degenerate rewards are collected, and all generated responses count against the rollout cost.

Rollout counts per prompt are then set so that prompts with larger estimated influence on the sampler update get more responses and persistently degenerate prompts get fewer, within a total budget scaled by the batch's average waste score.

Why This Matters

Research impact. The paper reframes selective rollout as an importance-weighting problem across a distribution shift, and shows a concrete mechanism for using a small amount of target feedback inside RLVR without exposing target data. It also connects sampler regret to gradient-estimation variance, giving a principled justification for non-uniform rollout allocation rather than a heuristic one.

Real-world applications:

• Clinical or hospital settings, where target prompts must stay on-premise and training proceeds through shared model updates only.
• Blind benchmark evaluation, where a remote evaluator returns aggregate scores without releasing test inputs.
• Domain adaptation for specialized math, science, or engineering corpora, where the public training set differs in difficulty or subdomain from deployment queries.
• Cost-constrained post-training pipelines, where rollout compute is the dominant expense and vendors need to hit an accuracy target within a fixed budget.

Industry relevance. Rollout generation is the main cost driver in RLVR training, and the reported reductions (for example 7.6M to 3.0M responses, or about 24 versus 41 hours over 1000 steps) are directly relevant to teams deciding how much inference compute a post-training run justifies. The restricted-disclosure interface also fits common enterprise arrangements where data and model ownership are split.

Future Directions

• Formal privacy guarantees. The authors state that establishing formal privacy guarantees for the directional-derivative feedback is future work; bounding what an adversary can infer from repeated target-side feedback is an open question.

• Feedback beyond gradients. The paper mentions aggregated target gradients or directional derivatives as one protocol; whether coarser feedback (for example scalar scores only) still supports useful sampler learning is untested here.

• Robustness to larger and structured mismatch. Experiments cover two source sets and six math benchmarks; whether the approach holds under severe domain shift, non-math verifiable tasks, or target sets smaller than the 4% slice tested is not reported.

• Interaction with other efficiency mechanisms. L2RR is compared against GRESO and AR3PO but not combined with response reuse; whether target-guided sampling and reuse compose, and how the rollout-allocation proxy (c_{i,t} = w_{i,t}²) should be tuned, remain open.

Target Audience

Researchers and engineers working on LLM post-training, reinforcement learning with verifiable rewards, and data-efficient fine-tuning. It is most useful to readers already comfortable with GRPO mechanics and interested in either distribution-shift correction during RL or reducing rollout cost in production-scale training pipelines. Readers seeking a beginner-level introduction to RLVR should look elsewhere first.

Authors’ abstract

Recently, reinforcement learning from verifiable rewards (RLVR) is practically very important, where the leading approach is group relative policy optimization (GRPO) that typically faces a trade-off between rollout efficiency and downstream performance. To improve efficiency while preserving performance, existing selective rollout methods mainly exploit only source-side training dynamics to optimize the discrete distribution over source prompts. However, in practice, the source training dataset is rarely identically distributed to the target test dataset. When such a mismatch exists, due to legal or privacy constraints, it is usually not acceptable to share target prompts themselves with the RLVR trainer. Nevertheless, it is still possible for the target data owner to provide limited aggregated feedback from a small held-out target set without revealing target prompts or rollout traces. In this paper, we propose learning-to-reweight rollout (L2RR) that explicitly incorporates target-side performance feedback into GRPO. Specifically, L2RR learns a discrete sampling distribution over source prompts by optimizing a bi-level objective that maximizes accuracy on a held-out target dataset, while simultaneously penalizing rollouts that provide little incremental learning information. Empirical studies on six math reasoning benchmarks and three model scales show that L2RR improves target-domain performance and increases rollout efficiency.