PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient RLVR Training via Weighted Mutual Information Data Selection

Xinyu Zhou, Boyu Zhu, Haotian Zhang, Huiming Wang, Zhijiang Guo

arXiv:2603.01907v1cs.LGcs.CL

TL;DR

RL data selection is inefficient when difficulty-based heuristics treat intermediate success rates as informative while ignoring uncertainty from limited evidence. The paper introduces INSIGHT, a Bayesian weighted-mutual-information method that selects data using stable posterior beliefs and extends to multi-rollout RLVR. Across planning, mathematics, and general reasoning benchmarks, INSIGHT consistently improves performance and training efficiency, with reported gains up to +1.41, +1.01, and ∼2.2x acceleration.

  • Problem

    Existing online RL data selection equates intermediate difficulty with informativeness and overlooks epistemic uncertainty caused by limited evidence.

  • Method

    INSIGHT models latent task success rates Bayesianly and uses weighted mutual information with mean posterior beliefs to select informative datapoints, including in multi-rollout RLVR.

  • Results

    INSIGHT consistently improves final performance and training efficiency across planning, mathematical, and general reasoning benchmarks, reaching up to ∼2.2x acceleration.

  • Takeaways & Limitations

    Information-aware selection offers a robust alternative to heuristic difficulty-based sampling, particularly in smaller-model and low-resource regimes.

  • Takeaways & Limitations

    Difficulty-only comparisons are analyzed under an evidence-dependent uncertainty framework that illustrates binary-reward RLVR with K = 1 before multi-rollout extension.

Abstract

from arXiv · show

Reinforcement learning (RL) plays a central role in improving the reasoning and alignment of large language models, yet its efficiency critically depends on how training data are selected. Existing online selection strategies predominantly rely on difficulty-based heuristics, favouring datapoints with intermediate success rates, implicitly equating difficulty with informativeness and neglecting epistemic uncertainty arising from limited evidence. We introduce InSight, an INformation-guided data SamplInG metHod for RL Training, grounded in a weighted mutual information objective. By modeling data outcomes with Bayesian latent success rates, we show that expected uncertainty reduction decomposes into complementary difficulty- and evidence-dependent components, revealing a fundamental limitation of difficulty-only selection. Leveraging this observation, InSight constructs a stable acquisition score based on the mean belief of datapoints' success rather than noisy sampled outcomes, and naturally extends to multi-rollout settings common in reinforcement learning with verifiable rewards (RLVR). Extensive experiments demonstrate that InSight consistently achieves state-of-the-art performance and improves training efficiency, including a +1.41 average gain on Planning & Mathmatics benchmarks, +1.01 improvement on general reasoning, and up to ~2.2x acceleration, with negligible additional computational overhead.

1. Introduction

RL training efficiency depends strongly on adaptive data selection, but existing methods often equate intermediate difficulty with informativeness. INSIGHT instead uses weighted mutual information to account for both difficulty and accumulated evidence.

  • Motivation: RL training is computationally expensive and uniform sampling wastes effort on mastered or currently intractable tasks.This mismatch can also reduce optimization stability by lowering the effective batch size.
  • Existing approaches: Offline curricula are fixed before training, while online methods adapt to evolving model capabilities but face signal-cost and utility-estimation trade-offs.Oversampling improves estimate reliability at substantial computational cost.
  • Existing approaches: Difficulty-based heuristics prioritize datapoints near intermediate success rates, implicitly treating difficulty as a surrogate for informativeness.Their utility may diminish as evidence accumulates and uncertainty becomes sufficiently reduced.
  • INSIGHT: INSIGHT scores datapoints with weighted mutual information to prioritize expected uncertainty reduction while jointly modeling difficulty and accumulated evidence.Observed rewards update Bayesian beliefs, enabling adaptive selection through the training process.
  • Results: INSIGHT achieves up to +1.41 gain on Planning & Mathmatics benchmarks, +1.01 improvement on general reasoning, and ∼2.2x acceleration.The reported results establish an information-aware alternative to difficulty-only selection.

2. Related Work

Prior work treats data selection as an efficiency and generalization problem across learning settings, with RL posing additional costs because it requires repeated rollouts, evaluation, and credit assignment.

  • RL for LLMs: RL research also includes methods for mitigating training challenges and infrastructure efforts for scalable RL training.These complementary directions address different aspects of the RL training pipeline.
  • Data Selection for RL: Data selection has long targeted informative or high-quality samples in pretraining and supervised fine-tuning.This literature motivates selecting training data rather than sampling uniformly.
  • Data Selection for RL: RL makes uniform sampling especially inefficient because training requires repeated rollouts, policy evaluation, and credit assignment.These costs distinguish RL data selection from lower-cost supervised settings.

3. Preliminary

The preliminary setup models each reasoning task as a bandit arm with an unknown Bayesian success rate and binary reward observations. The framework selects tasks for information gain rather than cumulative reward.

  • Setup & Notations: Each datapoint is a mathematical or logical reasoning problem drawn from the full dataset pool.The selected batch contains M datapoints from the pool.
  • Data Difficulty Modeling: At each training step, the policy generates K independent response rollouts for selected datapoints and receives binary rewards.Rewards indicate whether each response succeeds or fails against the task criterion.
  • Data Difficulty Modeling: Each datapoint has an unknown latent success rate ϕτ ∈ [0, 1] that represents its difficulty under the current policy.The success rate is inferred from observed binary feedback.
  • Bayesian Online Data Selection: The Bayesian surrogate treats datapoints as stochastic bandit arms but selects them for informative learning signals rather than maximum cumulative reward.This objective targets model updating through reward observations.
  • Bayesian Online Data Selection: A Beta prior provides closed-form inference and conjugate posterior updates for the Bernoulli reward model.A uniform prior is typically initialized with parameters (1, 1), while temporal discounting is applied for stability.
  • Difficulty-Based Selection: Earlier heuristics rank datapoints whose sampled success rates lie closest to a target difficulty ϕ∗ ≈ 0.5.This corresponds to selecting problems considered mid-level in difficulty.

4. Limitation of Difficulty-Only Data Selection

Difficulty-only selection favors outcome variability but ignores epistemic uncertainty and accumulated evidence. The analysis shows that expected uncertainty reduction depends on both prior mean and evidence, then decays as evidence grows.

  • Limitation of Difficulty-Only Data Selection: Difficulty-only heuristics target aleatoric uncertainty while ignoring epistemic uncertainty from limited rollouts.A datapoint can remain difficult even as uncertainty about its latent success rate decreases.
  • Motivation: Prioritizing expected uncertainty reduction can produce more reliable and discriminative reward signals for policy optimization.These signals can yield gradients that better reflect true task difficulty.
  • Expected Variance Reduction: Expected variance reduction depends on the prior mean ¯ϕτ and accumulated evidence n, not solely on whether ¯ϕτ ≈ 0.5.The figure examines expected variance reduction as a function of prior mean.
  • Analysis: The analysis illustrates the common RLVR setting with binary rewards and K = 1 to characterize evidence-dependent uncertainty reduction.The full derivation is provided in the appendix.
  • Expected Variance Reduction: When accumulated evidence is large, uncertainty reduction becomes negligible even for datapoints with ¯ϕτ ≈ 0.5.This evidence-dependent decay exposes a fundamental limitation of difficulty-only selection.
  • Analysis: The difficulty heuristic also ignores sampling variability and may repeatedly select well-estimated, mid-level-difficulty data with limited information gain.The sampled success rate concentrates around the posterior mean in expectation.

5. Method

INSIGHT selects RLVR data using a weighted mutual-information score that combines epistemic uncertainty with expectation-based difficulty. The method extends mutual-information analysis to multiple rollouts and selects the highest-scoring datapoints for training.

  • Difficulty-only selection conflates outcome variability with informativeness, whereas INSIGHT separates difficulty-dependent and evidence-dependent uncertainty reduction.
  • Mutual information measures expected uncertainty reduction about a datapoint’s latent success rate from observed rewards.The formulation uses entropy and beta-distribution quantities to characterize this reduction.
  • INSIGHT extends the uncertainty-reduction analysis to K independent rollouts, where the observed success count follows a binomial distribution.
  • As evidence accumulates, information gained from an additional observation decreases, reflecting diminishing epistemic uncertainty for well-understood datapoints.
  • The weighted score uses the prior mean success rate, filters for high outcome variability, and smoothly biases selection toward difficulty level µ with sharpness controlled by η.
  • INSIGHT computes WMI scores, ranks a larger candidate set, and uses the top-M datapoints for rollouts and policy training.The procedure is presented as an online selection algorithm designed to integrate with RLVR algorithms.

6. Experiments

INSIGHT is evaluated across planning, mathematics, and general-reasoning tasks using multiple LLM scales and RL baselines. It consistently improves accuracy and reaches competitive performance with fewer training steps.

  • Experimental Setup: INSIGHT is evaluated on planning, mathematics, and general-reasoning tasks using GRPO across LLMs ranging from 0.6B to 7B parameters.Planning uses Countdown; mathematics uses AIME24, AMC23, MATH500, Minerva Math, and OlympiadBench; general reasoning uses MMLU and GPQA-Main.
  • Planning & Mathematics: INSIGHT consistently outperforms RANDOM and MOPPS across model scales, achieving the highest average accuracy in every planning and mathematics setting.Gains reach +5.13 on CountDown and +2.30 on AIME24 over RANDOM.
  • General Reasoning: +1.14 on MMLU-STEM and +3.16 on GPQA are notable INSIGHT gains across general-reasoning models.INSIGHT achieves the strongest or near-strongest performance across model sizes.
  • Discussion: Improvements decrease with model scale, while EXPECTED-DIFFICULTY generally outperforms MOPPS and INVERSE-EVIDENCE can match or underperform RANDOM.The results indicate that mean difficulty is more stable than sampled difficulty, whereas epistemic uncertainty alone is insufficient without task difficulty.
  • Training Efficiency: Up to ∼2.2x, ∼1.5x, and ∼1.6x speedups are achieved on Qwen3-0.6B, Qwen3-4B, and R1-Distill-Qwen-7B, respectively.On Countdown, INSIGHT reaches competitive or superior performance with fewer training steps and shows more consistent early-stage gains than MOPPS.

7. Ablation Analysis

Ablations show that INSIGHT benefits from combining difficulty weighting with mutual information and from using mean posterior beliefs rather than sampled success rates. A moderate difficulty target performs best overall.

  • WMI Components: Combining mutual information and difficulty weighting yields the most consistent performance across benchmarks.Mutual information alone lacks difficulty awareness, while difficulty weighting alone lacks explicit uncertainty awareness.
  • Mean Versus Sampled Beliefs: Mean success-rate estimation consistently outperforms sampled success rates across models and benchmarks.The analysis attributes this to less noisy data ranking and more reliable acquisition signals.
  • Difficulty Bias: The difficulty bias µ = 0.3 achieves the strongest or near-strongest performance in the sensitivity study.Extreme settings, µ = 0.1 or 0.7, tend to underperform, especially for smaller models.

8. Conclusion

The paper introduces INSIGHT as an information-guided framework for efficient LLM reinforcement learning. Its experiments show consistent gains in efficiency and final performance, especially for smaller and low-resource regimes.

  • Conclusion: INSIGHT combines task difficulty and epistemic uncertainty through weighted mutual information and uses stable mean posterior beliefs for data selection.The framework also extends to multi-rollout RLVR settings.
  • Conclusion: INSIGHT consistently improves training efficiency and final performance over strong online baselines across planning, mathematics, and general reasoning.The strongest benefits occur particularly in smaller-model and low-resource regimes.

Impact Statements

The paper aims to improve reinforcement-learning efficiency through principled data selection rather than introduce new model capabilities or application domains. Its technical appendix derives variance-reduction and Beta-distribution entropy expressions used in the analysis.

  • Impact Statement: The stated impact goal is improving reinforcement-learning efficiency through principled data selection.The paper focuses on algorithmic efficiency and does not introduce new model capabilities or application domains.
  • Technical Derivation: The variance-reduction derivation models latent success rates with Beta distributions and updates their posterior using observed binary rewards.The appendix applies Bayes’ rule to obtain the posterior distribution.
  • Technical Derivation: The Beta-distribution entropy is expressed in closed form using the beta function and digamma function.The derivation substitutes closed-form expectations involving ψ(·).
  • Technical Derivation: For large accumulated evidence nτ, the Beta posterior is approximated by a Gaussian to approximate mutual information through expected log variance reduction.The approximation relies on posterior variance concentration and a small-variance Taylor expansion.

D.1. Datasets Details

The experiments cover Countdown, DeepScaler mathematics, and WebInstruct-verified reasoning tasks, using GRPO as the default RL algorithm. INSIGHT precomputes selection scores and achieves comparable or better performance than dynamic sampling at substantially lower runtime.

  • Datasets and setup: Countdown uses 2,000 training samples from CountDown-34 and a 512-problem held-out test set, with rewards based on correctness and formatting.Problems contain three or four source numbers and require arithmetic operations to reach a target.
  • Datasets and setup: DeepScaler trains on approximately 40,000 mathematics problem-answer pairs and evaluates average pass@1 across five mathematics benchmarks.The benchmarks are AIME24, AMC23, MATH500, Minerva Math, and OlympiadBench, using binary correctness rewards.
  • Datasets and setup: WebInstruct-verified provides diverse reasoning problems spanning domains including physics and chemistry, evaluated with LM-Evaluation-Harness and rule-based rewards.The dataset is described as high-quality and intended to enhance robust reasoning capabilities across broad domains.
  • Training configuration: GRPO implemented in VeRL is the default algorithm, while INSIGHT is designed to integrate with other RLVR algorithms.Training uses 100 steps and eight responses per prompt for advantage estimation; evaluation uses pass@1 from 16 generations per prompt.
  • Runtime comparison: INSIGHT precomputes WMI scores, avoiding dynamic sampling's oversampling, full generation, and reward evaluation during data filtering.This design is described as significantly faster than dynamic sampling under otherwise matched training settings.
  • Runtime comparison: More than 2x higher total training hours are required by dynamic sampling, while INSIGHT consistently improves over RANDOM and matches or exceeds dynamic sampling on most metrics.Dynamic sampling's extra cost comes from oversampling followed by complete generation and reward evaluation; INSIGHT adds no extra training-time overhead.

F. Ablation Analysis of the Effect of Larger Candidate Batch Size ˆ

The optimal larger candidate batch size depends on model capacity. Smaller models favor moderate candidate pools, whereas larger models generally benefit from the largest tested pool, 16x.

  • Candidate batch-size ablation: The ablation compares larger candidate batch sizes of 8x, 12x, and 16x across models on mathematics benchmarks.Table 7 reports these comparisons, with bold values indicating the best results.
  • Larger models: For larger models such as QWEN3-4B and R1-DISTILL-7B, performance improves more consistently with larger candidate pools, making 16x strongest across most benchmarks.The paper attributes this pattern to larger models' ability to exploit candidate diversity and estimate epistemic uncertainty more accurately.
  • Overall finding: Overall, optimal candidate batch size varies with model capacity: moderate pools suit smaller models, while larger models can leverage broader pools to improve selection quality.This capacity-dependent behavior is reported across the tested mathematical reasoning settings.
Loading 2603.01907v1…