🤖 AI Summary
This work addresses inefficiencies in reinforcement learning with verifiable rewards (RLVR), where excessive identical-reward prompts waste generation budgets and existing pre-sampling strategies struggle to balance exploration and exploitation. The authors propose a latent-space-guided prompt sampling mechanism that dynamically partitions the candidate set and prioritizes uncertain prompts likely to yield non-zero reward variance by integrating historical returns with latent-space proximity. Rollout budgets are adaptively allocated based on the recent proportion of effective rewards. Evaluated across six mathematical reasoning and three out-of-distribution general reasoning benchmarks, the method achieves state-of-the-art average performance, improving up to 3.7% over the strongest baseline while converging faster, with only approximately two seconds of additional sampling overhead per step.
📝 Abstract
Reinforcement learning with verifiable rewards (RLVR) improves the reasoning capabilities of large language models, but prompt groups with identical rollout rewards consume generation budget without effective learning signals. Pre-rollout prompt selection can reduce this waste by screening prompts before rollout generation. However, existing pre-rollout methods struggle to balance exploitation and exploration: repeatedly exploiting historically informative prompts can narrow training coverage, whereas broader exploration can lower the fraction of informative prompts. To address these limitations, we introduce LEEPS, a Latent-Guided Explore--Exploit Prompt Sampler that adaptively balances the reuse of previously observed informative prompts with continued exploration of uncertain ones. LEEPS partitions candidates into exploit and explore portfolios and adaptively allocates rollout budget according to their recent non-trivial ratios. It further uses representation-space neighbors and historical rollout outcomes to prioritize uncertain prompts likely to yield non-zero reward variance, thereby making exploration more targeted without additional rollouts. Across six mathematical reasoning benchmarks, LEEPS achieves the highest average score at both model scales, with relative gains of 2.6\% and 3.7\% over the strongest baseline for Qwen2.5-Math-1.5B and 7B, respectively, and generally improves faster during the training process. It also achieves the highest average score across the three evaluated OOD general-reasoning benchmarks at both model scales and adds only about 2 seconds of online sampling overhead per training step. Code is available at https://github.com/ShuangLiangX/LEEPS.