[ICML 2025/ArXiv] INSIGHT: Accelerating RLVR Training via Weighted Mutual Information

TL;DR

Reinforcement Learning (RL) for LLMs is notoriously slow due to the massive number of rollouts required. INSIGHT (INformation-guided data SamplInG metHod) cuts training time by up to 2.2x by replacing random or "difficulty-only" task selection with a mathematically grounded Weighted Mutual Information (WMI) objective. It ensures every rollout reduces the model's uncertainty, leading to superior reasoning performance (+1.41 avg gain) with zero additional compute overhead.

Background: The Cost of Being "Difficult"

In the world of Reinforcement Learning with Verifiable Rewards (RLVR), we often use "difficulty" as a proxy for "informativeness." If a model succeeds 50% of the time on a task, we assume it's the "sweet spot" for learning.

However, the authors of INSIGHT identify a critical flaw: Difficulty $eq$ Informativeness. As a model sees a 50%-success-rate task repeatedly, its epistemic uncertainty (how much it can still learn from this task) vanishes. Existing methods like MoPPS keep selecting these "dead" tasks just because they are difficult, leading to plateaued learning and wasted GPUs.

The Core Insight: Epistemic vs. Aleatoric Uncertainty

The paper's breakthrough lies in the decomposition of uncertainty reduction. By modeling the success rate ${\varphi }_{a}u$ of a task $au$ using a Beta Distribution, they show that the expected variance reduction $\Delta V(au)$ depends on two things:

1. Prior Mean (${\bar{\varphi }}_{a}u$): The task difficulty (Aleatoric).
2. Accumulated Evidence (n): How many times we've seen this task (Epistemic).

As shown in the figure below, information gain decays rapidly as evidence ($n$) grows, even if the difficulty stays at the "ideal" 0.5 mark.

Figure: Expected variance reduction decays as evidence accumulates, proving that "difficulty-only" selection is suboptimal.

Methodology: Designing the WMI Score

INSIGHT introduces a new acquisition score, Weighted Mutual Information (WMI): $\mathcal{A}(au)=w({\bar{\varphi }}_{a}u)\cdot I(R_{1:K};{\Phi }_{a}u)$

• Mutual Information ($I$): This term focuses on Epistemic Exploration. It calculates how much the reward from $K$ rollouts will reduce the uncertainty about the task's success rate.
• Weighting Function ($w$): This term handles Aleatoric Exploitation. It filters for tasks with high variance and biases them toward a target difficulty $\mu$ (usually 0.3 or 0.5), acting as a dynamic curriculum.

Figure: The INSIGHT pipeline maintains Bayesian beliefs and ranks candidates to select the top-M most informative datapoints.

Experimental Battle: SOTA Performance

The authors tested INSIGHT across models ranging from 0.6B to 7B parameters. The results across Math (AIME24, MATH500) and General Reasoning (MMLU, GPQA) were striking.

Key Findings:

1. Efficiency: INSIGHT achieved the same performance as random sampling in less than half the steps (up to 2.2x speedup on Countdown).
2. Performance Gain: On the DeepSeek-R1-Distill-Qwen-7B model, INSIGHT outperformed the random baseline by +1.08% across five math benchmarks.
3. Better than Oracle? Surprisingly, INSIGHT performed comparably to "Dynamic Sampling" (DS), which oversamples data and filters it manually—but INSIGHT did so in 12.5 hours compared to DS's 30.5 hours.

Table: INSIGHT consistently hits the highest average scores across benchmarks compared to other online selection baselines.

Critical Perspective: Why Does This Matter?

The industry is currently obsessed with "Scaling Laws" for RL (e.g., DeepSeek-R1). However, the bottleneck is often the quality and diversity of the prompts being used in rollouts.

INSIGHT proves that we don't need more data; we need smarter data selection. By using a Bayesian surrogate to track what the model "knows" and "doesn't know," we can make RL training far more accessible for researchers with limited compute.

Limitations: The method relies on a Beta-Bernoulli conjugate prior, which assumes binary rewards (correct/incorrect). Extending this to weighted or continuous rewards (e.g., partial credit) would require moving to more complex distributions like Dirichlet or Gaussian Processes.

Conclusion

INSIGHT is a masterclass in applying "old-school" Bayesian information theory to "new-school" LLM Reinforcement Learning. It effectively bridges the gap between curriculum learning and epistemic exploration, setting a new standard for training efficiency in the era of reasoning models.