PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamics-Predictive Sampling for Active RL Finetuning of Large Reasoning Models

Yixiu Mao, Yun Qu, Qi Wang, Heming Zou, Xiangyang Ji

arXiv:2603.10887v1cs.LGcs.AI

TL;DR

Online prompt selection can improve RL finetuning but often requires costly rollouts over enlarged candidate batches. DPS predicts prompt-solving states with an HMM and online Bayesian inference, then selects likely informative prompts before rollout. Across mathematics, planning, and visual geometry, it reduces rollout use while matching or exceeding rollout-intensive Dynamic Sampling, though it relies on correctness-based rewards.

  • Problem

    Online prompt selection improves training-data adaptivity, but rollout-intensive filtering over enlarged candidate batches creates substantial computational overhead.

  • Method

    DPS models prompt-solving progress as an HMM and applies online Bayesian inference to predict state distributions for pre-rollout prompt selection.

  • Results

    Across mathematics, planning, and visual geometry, DPS achieves comparable or superior performance to Dynamic Sampling with significantly fewer rollouts.

  • Takeaways & Limitations

    DPS reduces redundant rollouts while preserving adaptive selection of informative prompts during RL finetuning.

  • Takeaways & Limitations

    DPS relies on correctness-based rewards to define solving states, although the framework may extend to dense or process-based rewards.

Abstract

from arXiv · show

Reinforcement learning (RL) finetuning has become a key technique for enhancing the reasoning abilities of large language models (LLMs). However, its effectiveness critically depends on the selection of training data. Recent advances underscore the importance of online prompt selection methods, which typically concentrate training on partially solved or moderately challenging examples under the current policy, thereby yielding more effective model updates. While significantly accelerating RL finetuning in terms of training steps, they also incur substantial computational overhead by requiring extensive LLM rollouts over large candidate batches to identify informative samples, an expense that can outweigh the finetuning process itself. To address this challenge, this work proposes Dynamics-Predictive Sampling (DPS), which online predicts and selects informative prompts by inferring their learning dynamics prior to costly rollouts. Specifically, we introduce a new perspective by modeling each prompt's solving progress during RL finetuning as a dynamical system, where the extent of solving is represented as the state and the transition is characterized by a hidden Markov model. Using historical rollout reward signals, we perform online Bayesian inference to estimate evolving state distributions, and the inference outcome provides a predictive prior for efficient prompt selection without rollout-intensive filtering. Empirical results across diverse reasoning tasks, including mathematics, planning, and visual geometry, demonstrate that DPS substantially reduces redundant rollouts, accelerates the training process, and achieves superior reasoning performance.

1 INTRODUCTION

RL finetuning benefits from adaptive prompt selection, but rollout-intensive filtering is costly. DPS predicts prompt-solving dynamics online to select informative prompts before expensive rollouts.

  • RL finetuning effectiveness depends heavily on training-data quality, motivating adaptive online prompt selection.
  • DPS models each prompt’s solving progress as a dynamical system with hidden-Markov transitions and uses lightweight inference for online selection.
  • Dynamic Sampling improves sample quality and accelerates training steps by filtering prompts using multiple rollouts.
  • Long chain-of-thought generation makes Dynamic Sampling’s enlarged-batch rollouts computationally expensive, often outweighing finetuning costs.
  • Across mathematics, planning, and visual geometry, DPS selects more informative samples and matches or exceeds Dynamic Sampling with significantly fewer rollouts.

2 PRELIMINARY

RL finetuning optimizes response returns, but prompts with uniform outcomes provide weak learning signals. Online selection methods therefore use rollouts to retain prompts with informative reward variation, at substantial computational cost.

  • RL finetuning maximizes expected response return, with correctness commonly represented by binary rewards.
  • When all sampled responses share the same reward, normalized advantages can vanish and optimization can effectively halt.
  • Dynamic Sampling rolls out enlarged candidate batches and discards prompts whose responses have identical rewards.
  • Repeated rollouts and evaluations increase generation costs proportionally with candidate-batch size, especially for long chain-of-thought tasks.

3 DYNAMICS-PREDICTIVE SAMPLING FOR ACTIVE RL FINETUNING

DPS represents prompt-solving progress with a hidden Markov model, updates state and transition beliefs online from intermittent rollout observations, and selects prompts predicted to be partially solved.

  • 3.1 GENERATIVE MODELING OF PROMPT-SOLVING DYNAMICS: Prompts are categorized as fully unsolved, partially solved, or fully solved according to the number of correct sampled responses.
  • 3.1 GENERATIVE MODELING OF PROMPT-SOLVING DYNAMICS: Partially solved prompts are prioritized because mixed correct and incorrect responses provide stronger update signals.
  • 3.1 GENERATIVE MODELING OF PROMPT-SOLVING DYNAMICS: DPS models solving states with an HMM whose categorical states evolve through a transition matrix, while observations are exact when a prompt is selected.
  • 3.2 ONLINE INFERENCE AND TRANSITION LEARNING: Online Bayesian inference updates state beliefs and transition statistics from observed outcomes, while unobserved steps preserve the prior state belief.
  • 3.2 ONLINE INFERENCE AND TRANSITION LEARNING: An exponentially decayed Dirichlet update emphasizes recent transition statistics, enabling faster adaptation when the decay factor is smaller.
  • 3.2 ONLINE INFERENCE AND TRANSITION LEARNING: The inferred posterior and transition matrix generate a next-step predictive prior through low-dimensional online computation.
  • 3.3 PROMPT SAMPLING WITH PREDICTED DYNAMICS: DPS selects the B prompts with the highest predicted probability of being partially solved before current-step rollout outcomes are observed.

4 EXPERIMENTS

DPS is evaluated across reasoning tasks, model settings, sampling baselines, predictive accuracy, training performance, efficiency, and robustness factors. It selects informative prompts effectively while substantially reducing rollout and runtime costs relative to rollout-intensive sampling.

  • Experimental setup: Experiments cover competition-level mathematics, numerical planning, visual geometry, and multiple language or multimodal model capacities.Models are finetuned with GRPO and evaluated using average Pass@1 accuracy.
  • Prediction accuracy: DPS predicts prompt-solving states online and concentrates sampled batches on partially solved prompts more effectively than US and HR.The predictor is evaluated with overall accuracy plus Class 2 precision, recall, and F1, while effective sample ratio measures partially solved prompts.
  • Training performance: DPS improves policy faster than US and HR, reaches higher final performance, and achieves performance comparable to DS across tasks.DPS slightly surpasses DS on MATH, while US and HR degrade on MATH; DS is the rollout-intensive oracle baseline.
  • Generalization: DPS matches or exceeds DS generalization performance across tasks while consistently outperforming US and HR.Evaluations include mathematics, Countdown, general reasoning, held-out data, and harder variants.
  • Efficiency: DPS typically uses less than 30% of DS’s rollout budget and about half its runtime on standard MATH while maintaining comparable or better results.Its sampling and prediction operations are negligible; the modest runtime difference versus US and HR comes from longer response generation.
  • Robustness and analysis: Non-stationary decay improves performance and prediction accuracy, whereas equal weighting of all historical observations degrades both.The experiments also vary solving-state partitions, transition priors, and exploration behavior.

5 CONCLUSION AND LIMITATIONS

DPS models prompt-solving progress as a hidden Markov dynamical system and uses lightweight online inference to select informative prompts without rollout-intensive filtering. The method reduces redundant rollouts and accelerates training, but its current state definition relies on correctness-based rewards.

  • Conclusion: DPS represents solving extent as a state and models its transitions with a hidden Markov model for online prompt prediction and selection.The framework uses lightweight inference to avoid rollout-intensive filtering.
  • Conclusion: Across diverse reasoning tasks, DPS reduces redundant rollouts, accelerates training, and achieves superior reasoning performance.
  • Limitations: DPS relies on correctness-based rewards to define solving states, although the framework can extend to dense or process-based rewards.The authors also identify top-k selection as a potentially non-optimal strategy for future work.

ETHICS STATEMENT

The paper frames efficient RL finetuning around adaptive prompt selection, computational cost, and curriculum effects, while reporting public-data and reproducibility practices.

  • Ethics and reproducibility: The paper states that experiments use public datasets, follow dataset licenses, and introduce no additional ethical risks beyond standard LLM finetuning.
  • Ethics and reproducibility: The authors provide theoretical derivations and experimental details and commit to releasing complete code for reproduction.
  • Curriculum: The selection strategy induces a self-paced curriculum that shifts from easier partially solved prompts toward harder prompts as model capacity grows.

C PROOF AND DERIVATION

The supplied derivation passages specify Bayesian inference for prompt-solving states in a hidden Markov model and describe the experimental domains and benchmarks used to evaluate DPS.

  • Bayesian derivation: The transition-posterior update is defined as ξ_t(i, j) = P(z_{t−1} = j, z_t = i | y_1:t) for observed outcomes.
  • Bayesian derivation: The posterior joint state distribution factors into the previous-state belief, transition probability, and current observation likelihood.
  • Bayesian derivation: The deterministic emission model sets p(y_t | z_t = i) = δ(y_t, i), so only the state matching the observed outcome has nonzero emission probability.
  • Bayesian derivation: When the current outcome is unobserved, setting ξ_t = 0 makes the Bayesian update default to the prior without new evidence.
  • Experimental scope: Experiments cover competition mathematics, numerical planning, and visual geometry across language and multimodal models with varying capacities and architectures.
  • Experimental scope: Evaluation includes mathematics, planning, and visual-geometry benchmarks, with separate training datasets and held-out test sets for the task domains.

D.2 IMPLEMENTATION DETAILS

The implementation uses shared GRPO-based finetuning and compares sampling procedures under task-specific configurations, with DPS requiring only a small number of tuned assumptions.

  • Shared setup: All methods use GRPO within verl, and evaluation averages pass@1 accuracy over 16 independent completions per prompt.
  • Task configuration: Batch sizes are 256 for MATH and Countdown and 512 for Geometry3k, with task-specific mini-batches and output-length limits.
  • Task configuration: MATH uses binary rewards, whereas Countdown and Geometry3k add a 0.1 format bonus for incorrectly answered but correctly formatted responses.
  • Task configuration: Experiments run on 8 NVIDIA A100 GPUs with 80GB memory.
  • Sampling implementations: DS filters prompts with zero reward variance, HR excludes prompts yielding all correct responses in an epoch, and DPS predicts states from historical evidence.
  • DPS configuration: DPS initializes uniform state beliefs and Dirichlet transition priors without prior knowledge, tuning only the non-stationary decay ratio λ.

E.1 ADDITIONAL PREDICTION RESULTS

Additional analyses examine DPS’s state prediction and its effect on sampled batches. The results indicate improved prediction discrimination and fewer ineffective prompts than uniform sampling.

  • Prediction analysis: DPS predicts each prompt’s solving state online, enabling adaptive prioritization of partially solved examples during training.
  • Batch composition: DPS yields fewer fully solved and fully unsolved prompts than uniform sampling across all evaluated tasks.
  • Batch composition: HR produces fewer fully solved prompts but more fully unsolved prompts than DPS because it treats the fully solved state as absorbing.
  • Prediction analysis: Confusion-matrix diagonals strengthen and off-diagonal errors diminish as training progresses, indicating improved state-label discrimination.
  • Prediction analysis: The center confusion-matrix cell becomes more prominent in predictions and ground truth, indicating greater emphasis on the target region.
  • Overall finding: The analyses conclude that lightweight inference tracks solving progress and concentrates training on desired prompts.

E.2 ADDITIONAL EVALUATION RESULTS

Across additional evaluations, DPS maintains strong performance and generalization while using fewer rollouts than Dynamic Sampling. Its behavior remains robust across models, response lengths, transition priors, and prompt or response-length distributions.

  • Additional task and model evaluations: DPS matches or outperforms baseline methods on Geometry while requiring significantly fewer rollouts than DS.
  • Additional task and model evaluations: DPS shows consistent improvements over baseline methods on ARC-c and MMLU-Pro after MATH training.
  • Additional task and model evaluations: With Llama-3.2-3B-Instruct, DPS is comparable to DS and surpasses HR and US in test accuracy and effective sample ratio.
  • Additional task and model evaluations: Under a 32k response budget, DPS surpasses US and HR and slightly outperforms DS, indicating generalization to longer responses.
  • Transition-prior analysis: Structured transition priors can slightly improve early prediction and performance, but their advantages diminish as more training data becomes available.
  • Response and prompt length: During MATH training, DPS response lengths rise toward DS-like trajectories, while its sampled prompts become longer as model competence increases.

E.6 EMPIRICAL ANALYSIS ON COMPUTATIONAL SCALING BEHAVIOR

DPS’s runtime and memory costs scale approximately linearly with dataset size, while LLM training and generation remain substantially more expensive in the evaluated settings. At larger dataset scales, however, DPS overhead can become non-negligible.

  • Dataset-size scaling: DPS sampling and HMM updates scale approximately linearly in runtime and memory with dataset size.
  • Dataset-size scaling: LLM training and generation require about 1100s of runtime and 600 GiB of GPU memory in the 7B evaluation.
  • Dataset-size scaling: DPS overhead could become non-negligible when dataset size exceeds 10^8, although such scales are beyond typical practical settings.
  • LLM-size scaling: As model size increases from 1.5B to 7B, total LLM training and generation runtime rises from roughly 370s to 1100s, while DPS has no additional rollout-cost dependence.

E.7 SENSITIVITY ANALYSIS ON THE RESPONSE GROUP SIZE

DPS remains effective across response group sizes and outperforms simpler sampling baselines, especially when fewer responses are available per prompt. Additional entropy regularization provides limited benefit, while process-reward extension remains preliminary.

  • Response group size: For k = 4, DPS test accuracy exceeds US by more than twofold, while DPS decreases only about 4% relative to its larger-k performance.
  • Response group size: Smaller response groups reduce US’s likelihood of sampling effective prompts, whereas DPS maintains high prediction accuracy at k = 4.
  • Entropy-regularized selection: DPS+Entropy performs best with coefficient 0.1 but does not noticeably improve test accuracy or effective sample ratio over DPS.
  • Additional baselines: DPS outperforms the Var+EMA predictive baseline and diversity-based sampling in test accuracy and effective sample ratio.
  • Continuous process rewards: Applying active sampling to process rewards requires establishing how process rewards relate to sample informativeness; the reported extension is preliminary.

F DATA EXAMPLES

The paper illustrates its experimental tasks with representative prompts from mathematics, numerical planning, and visual geometry. The examples show the task-specific formats used for model responses and evaluation.

  • Task examples: MATH examples use prompts about mathematical definitions and problem solving, with reasoning and final answers requested in specified formats.
  • Task examples: Countdown examples ask models to construct an equation from given numbers using basic arithmetic operations.
  • Task examples: Geometry3k examples present geometric questions requiring internal reasoning and a boxed final answer.
Loading 2603.10887v1…