Small Generalizable Prompt Predictive Models Can Steer Efficient RL Post-Training of Large Reasoning Models

TL;DR

Reinforcement learning with verifiable rewards can enhance large language model reasoning but is computationally expensive. The authors introduce GPS, a small generative prompt predictive model that shares experience across prompts via a global latent context and couples this with a history-aware batch acquisition strategy to select informative prompt batches. GPS achieves faster RL post-training, better final performance, and lower test-time costs across math and logic benchmarks, while maintaining compatibility with multiple RLVR algorithms. The approach demonstrates strong cross-prompt generalization and enables test-time computation allocation without extra training cost, offering a scalable path for efficient reasoning with large models.

Abstract

Reinforcement learning enhances the reasoning capabilities of large language models but often involves high computational costs due to rollout-intensive optimization. Online prompt selection presents a plausible solution by prioritizing informative prompts to improve training efficiency. However, current methods either depend on costly, exact evaluations or construct prompt-specific predictive models lacking generalization across prompts. This study introduces Generalizable Predictive Prompt Selection (GPS), which performs Bayesian inference towards prompt difficulty using a lightweight generative model trained on the shared optimization history. Intermediate-difficulty prioritization and history-anchored diversity are incorporated into the batch acquisition principle to select informative prompt batches. The small predictive model also generalizes at test-time for efficient computational allocation. Experiments across varied reasoning benchmarks indicate GPS's substantial improvements in training efficiency, final performance, and test-time efficiency over superior baseline methods.

Figures (13)

• Figure 1: Framework overview. Unlike prompt-specific modeling like MoPPS qu2025prompt, GPS proposes a generalizable prompt predictive model (PPM) to estimate difficulty and track LLM evolution. The comprehensive prompt batch acquisition accounts for intermediate difficulty prioritization and batch-level diversity, improving RL post-training efficiency while mitigating PPM overfitting.
• Figure 2: Spearman’s rank correlation and $p$-value during training between predicted prompt difficulty and empirical success rate.
• Figure 3: Training curves of GPS and baselines across different scenarios and backbone models versus training steps. DS serves an oracle baseline with respect to training steps, but incurs substantially higher rollout costs. Training curves plotted against the number of rollouts are provided in Fig. \ref{['fig:performance_rollout']}.
• Figure 4: Test-time computation allocation results across benchmarks. $\mathrm{pass@k}$ versus the number of generated samples $k$ is shown, with insets reporting Spearman’s rank correlation $\rho$ (* $p<0.05$, ** $p<0.01$, *** $p<0.001$) and other statistical metrics.
• Figure 5: (a) Training curves on Countdown with PPO and Reinforce++. GPS is compatible with both algorithms and consistently outperforms Uniform. (b) Ablation study on Countdown, including removing history-anchored diversity (w/o hisdiv), removing inter-step exploration only (w/o his), and replacing the PPM with a deterministic PPM without latent variables (w/o $z$).

Theorems & Definitions (4)

• Theorem 3.1: Better Prediction with Shared History
• Theorem C.1: Better Prediction with Shared History
• proof
• Remark C.1: Realization via Global Latent Context