PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yun Qu, Qi Wang, Yixiu Mao, Heming Zou, Yuhang Jiang, Weijie Liu, Clive Bai, Kai Yang, Yangkun Chen, Saiyong Yang, Xiangyang Ji

arXiv:2602.01970v2cs.AIcs.LG

TL;DR

RLVR improves LLM reasoning but its rollout-intensive optimization is costly, and existing prompt selection either requires expensive evaluations or lacks cross-prompt generalization. GPS uses a lightweight generative PPM trained on shared optimization history, with difficulty-diversity batch acquisition and test-time computation allocation. Across reasoning benchmarks, GPS improves training and test-time efficiency, including up to 2.0× training speedup and 36.4% lower inference cost.

  • Problem

    RLVR is computationally expensive, while existing prompt selection relies on costly exact evaluations or prompt-specific predictive models that do not generalize across prompts.

  • Method

    GPS trains a lightweight generative PPM on shared optimization histories and combines intermediate-difficulty prioritization with history-anchored diversity for batch selection.

  • Results

    Across diverse reasoning benchmarks and LLM backbones, GPS improves training efficiency and performance, achieving up to 2.0× speedup and up to 36.4% lower inference cost.

  • Takeaways & Limitations

    The learned PPM supports efficient prompt selection during RL post-training and computation reallocation at test time.

  • Takeaways & Limitations

    Prompt difficulty is non-stationary during training, so predictors without mechanisms for modeling evolving dynamics can adapt with delay.

Abstract

from arXiv · show

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.

1. Introduction

RLVR improves reasoning but is computationally expensive, motivating online prompt selection. GPS addresses limits of costly evaluation and prompt-specific prediction with a generalizable PPM and difficulty-diversity batch acquisition.

  • RLVR requires substantial rollouts for policy evaluation and updates, creating high computation and memory costs.
  • Intermediate-difficulty prompts provide more informative updates, whereas overly easy or extremely hard prompts can yield vanishing gradients.
  • Online prompt selection improves training but adds substantial rollout overhead by evaluating candidate prompts before selecting a batch.
  • Existing prompt-specific predictive models fail to share information across prompts and can lag behind the evolving model state.
  • GPS uses a lightweight generative PPM trained on shared optimization histories, combining intermediate-difficulty prioritization with history-anchored diversity for batch acquisition.
  • GPS achieves up to 2.0× training speedup, up to 69% lower rollout cost than evaluation-based selection, and test-time computation reductions of up to 36.4%.It can also improve accuracy by up to 3.2% under fixed budgets.

2. Preliminaries

The preliminaries formulate RLVR over mathematical and logical prompts, where sampled responses receive verifiable rewards for policy optimization. They also introduce online selection methods that trade informative batches against rollout cost and prediction limitations.

  • 2.1. Notations: The prompt pool contains mathematical or logical problems, and each training iteration selects a batch for rollout generation and optimization.
  • 2.1. Notations: Responses are verified by a reward function, with binary correctness rewards used in the main formulation.
  • 2.1. Notations: Prompt success rate is treated as an intrinsic difficulty attribute representing the probability that the current model solves the prompt.
  • 2.1. Notations: The optimization history collects prompt-specific feedback and associated historical signals across training steps.
  • 2.2. Reinforcement Learning with Verifiable Rewards: RLVR optimizes the policy using verifiable rewards, while GRPO stabilizes post-training with group-normalized advantages and no PPO value network.
  • 2.3. Online Prompt Selection for RLVR: Evaluation-based selection uses real rollouts on an oversized candidate set, whereas prediction-based selection estimates prompt difficulty to avoid extra evaluations.

3. Method

GPS replaces independent prompt-specific prediction with a shared-history generative PPM and combines difficulty guidance with batch diversity for efficient prompt selection. The learned PPM also supports difficulty-aware test-time computation allocation.

  • Generalizable Prompt Predictive Model: Independent prompt-specific PPMs face cold-start limitations, delayed adaptation to non-stationary difficulty, and poor generalization to unseen prompts.Their estimates are reliable mainly for frequently sampled prompts and may fail to exploit semantic structure or changing model performance.
  • Generalizable Prompt Predictive Model: GPS uses a lightweight generative PPM that models prompt difficulty from complete optimization histories shared across prompts.A global latent difficulty context transfers information across prompts, while a history-dependent prior adapts predictions to evolving optimization dynamics.
  • Extension to Test-Time Computation Allocation: The generalizable PPM predicts difficulty for unvisited prompts and supports computation reallocation at test time as the policy evolves.Difficulty-aware allocation is motivated by the marginal gain of an additional Best-of-N sample, while the evolving prior keeps predictions synchronized with policy changes.
  • Generalizable Prompt Predictive Model: Full optimization history can yield lower predictive mean squared error than estimators based only on per-prompt observations when it contains non-redundant information.The theorem characterizes the resulting estimation gap and identifies when shared history improves prediction.
  • Difficulty-Diversity Unified Prompt Batch Selection: The batch acquisition objective balances individual prompt informativeness with intra-batch dispersion and inter-step exploration.A greedy approximation makes the max-sum diversification objective scalable for large-scale LLM training.
  • Difficulty-Diversity Unified Prompt Batch Selection: GPS prioritizes intermediate-difficulty prompts because a success rate near 0.5 maximizes reward variance in binary-reward settings.The criterion can also accommodate alternative acquisition objectives.

4. Experiments

GPS is evaluated across mathematical and logical reasoning tasks, diverse LLM backbones, and multiple RLVR algorithms. It improves difficulty prediction, training efficiency, final performance, and test-time computation allocation, while its ablations support the value of diversity and latent context.

  • Experimental Setup: Experiments cover mathematical and logical reasoning benchmarks across diverse base and distilled LLM backbones, primarily using GRPO within verl.GPS is also tested with PPO and Reinforce++ on Countdown.
  • Reliable Difficulty Prediction: GPS rapidly learns prompt-difficulty rankings, with correlation improving as optimization histories accumulate and exceeding MoPPS in prediction quality.The evaluation uses Spearman rank correlation and statistical significance testing.
  • Efficient and Improved RLVR: 1.4×–2.0× training-step speedup over Uniform accompanies average gains of 1.6–1.9 points on mathematical tasks and 4.1–5.7 points on logical tasks.The lightweight PPM adds negligible overhead, while runtime increases modestly because of longer response generation.
  • Efficient and Improved RLVR: Up to 69% fewer rollouts and 28%–47% lower training time let GPS achieve comparable performance to DS at substantially reduced cost.GPS also outperforms prediction-based baselines and shows strong performance on out-of-distribution benchmarks.
  • Test-Time Generalization and Efficiency: At fixed test-time budgets, GPS improves performance by up to 3.2% relative to Default or reduces computation by up to 36.4% without performance loss.The learned PPM retains statistically significant correlations on most unseen test benchmarks without additional computational overhead.
  • Ablation Studies: Removing history-anchored diversity causes a substantial performance drop, while removing inter-step exploration or latent variables also reduces performance.These ablations indicate that intra-batch dispersion, inter-step exploration, and latent difficulty context each contribute to prompt selection.

5. Conclusion

GPS accelerates RL post-training through a small generative PPM that generalizes prompt-difficulty estimates and combines difficulty guidance with batch diversity. The learned PPM also supports test-time computation allocation, while the authors identify refined generative modeling and stronger allocation designs as future directions.

  • Conclusion: GPS accelerates RL post-training by combining a small generative PPM with difficulty-guided, batch-level diversity selection.The approach supports cross-prompt generalization and adaptation to policy evolution.
  • Conclusion: The learned PPM also aids test-time computation allocation, while future work may refine generative modeling and allocation designs.The discussion describes test-time allocation as a preliminary opportunity for further improvement.

Impact Statement

The paper frames GPS as a small-steer-large approach for reducing RL post-training computation through shared prompt-difficulty modeling and history-aware selection. Its scope depends on coarse semantic structure, history-dependent context, and regularization choices, while the framework does not introduce new model capabilities or outputs.

  • Impact: GPS targets RL post-training efficiency by avoiding redundant rollouts through online predictive prompt selection, without directly altering model outputs.The approach operates at the data-selection and optimization-strategy levels.
  • Conceptual Contribution: A shared lightweight PPM transfers difficulty signals across prompts, contrasting with prior prompt-specific modeling and forming the paper’s small-steer-large paradigm.The shared model is presented as an early attempt to guide a larger LLM using a smaller predictive model.
  • Scope and Assumptions: The conditional generative model is an instantiation of shared, history-conditioned modeling rather than the primary technical contribution.Alternative generative architectures or diversity regularizers could be used within the broader framework.
  • History-Anchored Diversity: History-anchored diversity counters over-exploitation of a narrow prompt subset by encouraging coverage across current and historical selections.Without diversity regularization, concentrated sampling can increase overfitting risk and weaken generalization to underexplored prompts.
  • Difficulty Modeling: Prompt difficulty is modeled from fixed semantic embeddings and a time-evolving latent context summarizing optimization history.Semantic representations provide a stable but coarse cross-prompt anchor, while latent context captures time-varying optimization effects.
  • Scope and Assumptions: Semantic representations require only a weak structural assumption: under a fixed policy, they organize prompts coarsely rather than determining difficulty on their own.An embedding-similarity baseline reaches approximately 0.2 correlation, motivating history-dependent context.

C. Theoretical Proof

The theoretical section formalizes why conditioning prompt-difficulty prediction on shared optimization history can reduce estimation risk. It motivates a latent-context surrogate and documents the experimental settings used to evaluate the framework.

  • Theoretical Proof: Theorem 3.1 states that shared-history prediction improves risk over prompt-specific prediction by an amount represented by C(τ) ≥ 0.The improvement is strict when shared history contains non-redundant predictive information.
  • Theoretical Proof: The risk difference is positive exactly when full optimization history provides predictive information unavailable from prompt-specific feedback alone.The proof uses conditional-expectation projections and an orthogonal L2 decomposition.
  • Practical Surrogate: Because conditioning on the full history is generally intractable, GPS approximates it with a time-evolving latent prior and a shared decoder conditioned on prompt semantics.The surrogate seeks a stable relative difficulty ordering without assuming the latent variable fully characterizes optimization history.
  • Experimental Settings: The experiments train on DeepScaler mathematics and Countdown tasks using several LLM families, GRPO as the default algorithm, and specified reward and optimization settings.Evaluation includes in-distribution and held-out mathematical or logical benchmarks, with Euclidean embedding distance used for diversity.

D.4. Implementation of the Generalizable Prompt Predictive Model

The implementation uses a lightweight variational predictive model that shares history across prompts to estimate difficulty, then reallocates test-time computation by predicted-difficulty quantiles. The design prioritizes efficiency and feasibility, while leaving architectural optimization and allocation strategy search for future work.

  • D.4. Implementation of the Generalizable Prompt Predictive Model: The predictive model jointly learns an encoder, decoder, and history-conditioned prior through a variational ELBO objective.
  • D.4. Implementation of the Generalizable Prompt Predictive Model: Prompt embeddings and success-rate feedback are processed by lightweight Transformer encoders over short sequences of recent optimization history.
  • D.4. Implementation of the Generalizable Prompt Predictive Model: A shared MLP decoder combines latent context with prompt embeddings to predict difficulty, enabling experience transfer across prompts.
  • D.4. Implementation of the Generalizable Prompt Predictive Model: The predictive model has 20M parameters, less than 1% of most LLMs, and introduces negligible computational overhead across experimental settings.
  • D.4. Implementation of the Generalizable Prompt Predictive Model: The model architecture is intentionally simple and stable rather than optimized through extensive architecture search or scaling studies.
  • D.5. Extension to Test-Time Computation Allocation: At test time, fixed per-prompt sampling is replaced by difficulty-guided allocation because easy prompts saturate quickly while extremely hard prompts offer negligible gains.
  • D.5. Extension to Test-Time Computation Allocation: The learned PPM is reused to estimate difficulty without costly test-time estimation, including on unseen benchmarks where prompt-specific models do not generalize.
  • D.5. Extension to Test-Time Computation Allocation: Under a global budget, each prompt receives a minimum allocation and remaining computation is distributed using predicted-difficulty quantiles and a budget-determined coefficient.

E.1. Computational Complexity Analysis

GPS reduces rollout and selection overhead while improving difficulty estimation and prompt-batch informativeness. Its framework also extends beyond binary rewards, although continuous-reward acquisition requires an additional reward-variance criterion.

  • GPS performs prompt selection with lightweight computation, contrasting with evaluation-based methods that incur rollout overhead over enlarged candidate sets.
  • GPS sampling and PPM updates take less than 16s on a synthetic 1M-prompt pool, with sublinear overhead relative to LLM training.
  • GPS achieves higher difficulty-prediction correlation than MoPPS under uniform sampling, whose prompt-specific observations become sparse at large scale.
  • GPS maintains a consistently higher effective sample ratio than Uniform and MoPPS, indicating more informative selected prompt batches.
  • GPS achieves comparable training performance to DS using only 31%–34% of DS’s generated rollouts.
  • Under continuous process rewards, GPS replaces success-rate difficulty with predicted reward variance and consistently outperforms uniform prompt selection.
  • History-anchored diversity and latent difficulty context both improve performance, with their ablations degrading results across Countdown and DeepScaler.

E.7. Evaluation with Different LLM families

GPS maintains reliable difficulty prediction and better performance on Llama-3.2-3B-Instruct, extending beyond the Qwen and DeepSeek model families. The ablation evidence also supports the contribution of its diversity and predictive-model components.

  • On Countdown with Llama-3.2-3B-Instruct, GPS continues to provide reliable difficulty prediction and consistently better performance.
  • The ablation compares GPS components including the generative predictive model and history-anchored diversity on Countdown and DeepScaler.
  • The evaluation is framed as evidence that GPS is not specialized to one LLM series and may generalize across heterogeneous model families.

E.8. Training Dynamics

GPS changes training dynamics by selecting intermediate-difficulty prompts, producing longer reasoning responses and stronger exploration than uniform sampling. Performance remains robust across diversity weights, candidate sizes, domains, and embedding choices.

  • Training Dynamics: The training-dynamics evaluation includes response length, entropy, and training reward under different prompt-selection strategies.
  • Training Dynamics: GPS and DS produce longer responses than uniform sampling, consistent with intermediate-difficulty prompts eliciting longer reasoning chains.
  • Training Dynamics: In DeepScaler, GPS and DS increase entropy, with GPS inducing the stronger increase associated with more diverse exploration.
  • Training Dynamics: GPS keeps average training reward near 0.5 throughout training, reflecting sustained focus on informative prompts compared with uniform sampling.
  • Hyperparameter Sensitivity: Very small diversity weights produce redundant selection, whereas excessively large weights prioritize diversity over utility and degrade training performance.
  • Candidate Batch Size: Increasing candidate batch size improves performance until gains saturate beyond a moderate threshold, indicating diminishing returns.
  • Mixed-Domain Evaluation: GPS consistently outperforms Uniform in mixed-domain mathematics and code-generation training.
  • Embedding Sensitivity: GPS remains robust to embedding choice, significantly outperforming uniform sampling across evaluated embedding models despite slight absolute-performance variation.

E.13. Applicability beyond Reasoning: Evaluation on RLHF

GPS extends beyond binary-reward reasoning tasks to RLHF with continuous rewards, where it maintains reward-prediction reliability and improves training efficiency and final performance over uniform sampling. The study identifies theoretically optimal continuous-reward sampling as future work.

  • GPS is evaluated on HH-RLHF with Qwen3-1.7B-Base in an RLHF setting with continuous rewards.
  • GPS maintains a strong Spearman correlation with empirical rewards while using average continuous reward prediction and quantile-threshold sampling.
  • 1.24× training speedup and improved final performance are achieved over uniform sampling on RLHF.
  • A theoretically optimal sampling criterion for continuous rewards remains an open direction for future work.

F. Data Examples

The paper uses benchmark-specific prompt formats that combine reasoning instructions with output constraints. Examples cover mathematics and general reasoning tasks, including explicit ground-truth and answer-format requirements.

  • DeepScaler & Mathematics Benchmarks: DeepScaler and mathematics prompts append a chain-of-thought instruction and require the final answer within \boxed{}.
  • General Reasoning Benchmarks: General reasoning benchmarks follow PRIME’s prompt template and request reasoning in <think> tags with the final answer in <answer> tags.
  • DeepScaler & Mathematics Benchmarks: A mathematics example defines an operation, asks for a composed-expression difference, and gives the ground-truth answer as −2.
Loading 2602.01970v2…