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

AI summary

Overview

Research area: Reinforcement learning (RL) finetuning of large language models (LLMs) for reasoning, specifically online prompt/data selection during training.

Technical level: Advanced. The paper assumes familiarity with policy gradient RL (GRPO), hidden Markov models, and Bayesian inference over Dirichlet-Categorical models.

Scope: The paper proposes Dynamics-Predictive Sampling (DPS), a method that predicts which training prompts are most informative — before spending compute on rollouts — by modeling each prompt's solving progress as a hidden Markov dynamical system.

What This Paper Is About

RL finetuning improves the reasoning abilities of large language models, but its efficiency depends heavily on which training prompts are selected. State-of-the-art online selection methods such as Dynamic Sampling (DS) work well in terms of training steps, but they require generating many extra responses over enlarged candidate batches to find useful prompts, and this rollout cost can exceed the cost of finetuning itself. The paper's goal is to keep the adaptivity of online prompt selection while eliminating the rollout-heavy filtering step, by predicting prompt informativeness instead of measuring it.

Key Contributions

1. A dynamical-systems formulation of prompt solving. The authors define an implicit three-state label for each prompt at each training step — State 1 (fully unsolved, all responses incorrect), State 2 (partially solved, some correct and some incorrect), and State 3 (fully solved, all correct) — and model its evolution with a hidden Markov model whose transition matrix is a random variable.

2. An online Bayesian inference scheme for intermittent observations. Because a prompt is only observed when it is actually rolled out, the method uses a degenerate emission model (exact state revealed when observed, no constraint when unobserved) and performs a three-stage update per step: Bayes update of the state belief, Dirichlet posterior update of the transition model, and forward propagation to produce a next-step predictive prior.

3. A non-stationary extension. A decayed Dirichlet update, α_t(i,j) = λ·α_{t−1}(i,j) + (1−λ)·α_0(i,j) + ξ_t(i,j) with λ ∈ (0,1), lets the transition model forget outdated statistics and adapt as the model's competence changes, and also acts as an implicit exploration mechanism.

4. An efficient selection pipeline. At each step, the B prompts with the highest predicted probability of being in State 2 are selected for rollout and training, using only low-dimensional matrix operations rather than large-scale generation for filtering.

Main Findings

• Prediction accuracy holds up during training. Evaluated as an online classification task over the three states, the predictor maintains high overall accuracy and strong precision and recall for Class 2 (partially solved), the state the paper identifies as most critical for training efficiency.

• Higher concentration of informative prompts. DPS consistently produces a higher proportion of partially solved prompts in sampled batches than Uniform Sampling (US) and History Resampling (HR), reaching approximately 90% in many tasks.

• Better mathematics performance at far lower rollout cost. On MATH-trained models (Table 1), DPS reaches an average of 52.13 with R1-Distill-1.5B and 63.13 with R1-Distill-7B, compared with 52.00 and 62.42 for DS (Oracle), 48.57 and 59.31 for US, and 49.23 and 60.38 for HR. DPS uses 737k rollouts for the 1.5B model and 287k for the 7B model, while DS (Oracle) uses 2933k and 1147k respectively.

• Runtime roughly halved versus DS. On the standard MATH dataset, DPS generally uses about half of DS's runtime (for example, 32h versus 89h for R1-Distill-1.5B and 39h versus 73h for R1-Distill-7B). DPS runtime is slightly longer than US and HR (27h and 28h for the 1.5B model), which the authors attribute to longer generated responses at higher performance rather than to the prediction and selection operations.

• Countdown results track the same pattern. On Countdown with Qwen2.5-3B, DPS scores 74.27 on CD-34 and 47.78 on CD-4 using 246k rollouts, versus DS (Oracle) at 74.95 and 47.67 using 1141k rollouts. With Qwen2.5-7B, DPS scores 81.15 and 59.61 at 246k rollouts versus DS (Oracle) at 81.26 and 60.77 at 1006k rollouts.

• Rollout efficiency is the headline gain. Across tasks, DPS typically matches or exceeds DS's results using less than 30% of DS's rollout budget.

• Non-stationary decay matters. In the λ ablation, both removing decay entirely (λ = 1) and relying only on the most recent feedback (λ = 0) reduce both performance and prediction accuracy; moderate values perform best.

• Three states are the right granularity. Performance and effective sample ratios decline with either a coarser two-state partition or finer partitions with more than three states — coarse partitions mix distinct dynamics, and fine partitions spread limited observations too thinly.

• Smaller response groups favor DPS. Testing k ∈ {4, 8, 16} on Countdown, DPS consistently beats US in performance and effective ratio, with the largest gap at k = 4, since a smaller k makes it less likely that uniform sampling encounters a mix of correct and incorrect responses.

• Exploration is implicitly maintained. Smaller decay parameters produce more uniform sample counts with lower variance and a higher minimum across the dataset. In a probe on the Countdown 3B task using three sets of 256 prompts (fewest past samples, highest predicted unsolved probability, and random), difficult prompts were predicted even more accurately than random ones, and under-sampled prompts showed only a small accuracy gap.

• Baseline behavior. US and HR show degradation on MATH, which the authors attribute to entropy collapse from too few effective samples per batch; HR's limitations are attributed to its rigid epoch-level absorbing transition assumption and its filtering of only fully solved prompts, which are rare early and mid-training.

Methodology in Plain English

The intuition is that a prompt is most useful for training when the model sometimes gets it right and sometimes gets it wrong. The paper calls this the "partially solved" state and wants to train on those prompts.

The problem is that you only know a prompt's state after you pay for rollouts. Dynamic Sampling solves this by generating many responses, throwing away prompts with uniform rewards, and training on what remains — accurate but expensive.

DPS instead treats each prompt's state as something that changes over time in a predictable way. At any training step, a prompt can move between fully unsolved, partially solved, and fully solved. The authors model these moves as a hidden Markov model: the state is hidden, and you occasionally get a peek at it when the prompt is selected for training.

Each step, for every prompt, the method does three cheap things. If the prompt was rolled out, it updates its belief about that prompt's current state using Bayes' rule and nudges the transition model with the observed transition. If the prompt was not rolled out, the belief simply carries forward unchanged. Then it pushes the updated belief through the transition model to get a prediction for the next step. The transition model itself is learned online using Dirichlet priors — essentially counting observed transitions — with a decay factor so old counts fade and the model adapts to a changing policy.

Selection is then simple: pick the B prompts with the highest predicted probability of being partially solved. All of this involves only tiny matrix operations, so it costs almost nothing compared to generating responses.

Why This Matters

Impact on research. The paper reframes prompt selection in RL finetuning as state estimation in a dynamical system rather than as a filtering problem requiring extra generation. This shifts the cost of "active" data selection from the expensive LLM rollout to inexpensive low-dimensional inference, and it connects online data curation to the existing machinery of Bayesian HMMs and Dirichlet-Categorical conjugacy. It also provides an alternative framing for why partially solved prompts help — as a state that can be predicted and tracked — and shows that the predicted dynamics are measurably non-stationary, since removing decay hurts.

Real-world applications (drawn from the tasks evaluated):

• Training reasoning models for competition-level mathematics, evaluated on AIME24, AMC23, MATH500, MinervaMath, and OlympiadBench.
• Training models for numerical planning, evaluated on Countdown (CD-34 and the harder CD-4 variant).
• Training multimodal models for visual geometric reasoning, evaluated on Geometry3k.
• General reasoning evaluation, reported on ARC-c and MMLU-Pro in the paper's Table 4.

Industry relevance. Rollout generation dominates the cost of RL finetuning for reasoning models because chain-of-thought generations are long. Using under 30% of DS's rollout budget and roughly half its runtime for comparable or better accuracy directly reduces the GPU-hours needed per training run, which matters for any organization finetuning large reasoning models at scale. The method is algorithm-agnostic in the sense that it selects batches rather than modifying the RL objective, and it was implemented with GRPO inside the verl framework.

Future Directions

1. Extending beyond correctness-based rewards. The paper states as a limitation that the method relies on correctness-based rewards to define solving states, and notes that the framework should extend to dense or process-based rewards by partitioning cumulative return intervals. Testing that extension is an open task.

2. Replacing top-k selection. The authors state that the straightforward top-k strategy may not be optimal and plan to explore more sophisticated criteria such as entropy-based prioritization of uncertain samples.

3. Understanding and tuning the exploration-exploitation balance. The decay parameter implicitly drives revisiting of under-sampled prompts. Whether a more explicit exploration mechanism would improve the sample-count distribution across the dataset remains open.

4. Scaling analysis. The paper examines how LLM training, LLM generation, and DPS sampling and updates scale with dataset size and LLM size, but the truncated content does not report those results; the underlying question of how far the approach extends to much larger models and datasets remains to be tested.

Target Audience

Researchers and engineers working on reinforcement learning for LLM reasoning, especially those concerned with training efficiency and data curation; practitioners finetuning reasoning models who are paying for large rollout budgets; and readers interested in probabilistic modeling — hidden Markov models and online Bayesian inference — applied to training-dynamics problems. Readers without background in policy-gradient RL or Bayesian state estimation will need to consult the preliminary and appendix sections first.

Authors’ abstract

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.