Source-linked AI summary
Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling
Runpeng Dai, Tong Zheng, Rui Liu, Chengsong Huang, Hongtu Zhu
TL;DR
Test-time scaling improves LLM reasoning but increases computation and latency, while existing adaptive sampling often depends on heuristics or distributional assumptions. The paper formulates adaptive sampling as an MDP and trains a lightweight RL controller using final-answer statistics to balance correctness, latency, and computation. Across benchmarks and samplers, it reports improved trade-offs over ASC and ESC, including fewer sampling rounds and total samples.
Problem
Test-time scaling improves reasoning performance but increases inference cost, and existing adaptive methods often rely on heuristics or distributional assumptions.
Method
The paper trains an RL controller for an adaptive-sampling MDP that uses answer statistics to decide whether to stop or acquire samples while jointly penalizing latency and computation.
Results
RL-Guided Sampling consistently improves the accuracy–efficiency trade-off over ASC and ESC, reducing sampling rounds and total samples across benchmarks and samplers.
Takeaways & Limitations
The CPU-friendly, non-invasive controller provides a lightweight approach to balancing answer correctness, sampling latency, and computation cost.
Takeaways & Limitations
The current formulation uses simple answer statistics and fixed penalties rather than richer signals or direct real-world time and money costs.
Abstract
from arXiv · showhide
Test-time scaling improves the reasoning performance of large language models but incurs substantial cost in both total computation and latency. Existing adaptive sampling methods partially mitigate this issue by dynamically deciding when to stop sampling, yet they typically rely on heuristic rules or rely on distribution assumptions. In this work, we formulate adaptive sampling as a Markov decision process (MDP). We train a lightweight sampling controller with reinforcement learning (RL) to jointly balance answer correctness, latency, and computation cost. At each round, the controller decides to stop sampling or to acquire additional samples. Our method is lightweight which only relies on statistics of final answers, and can be trained and deployed on CPU. We further show that the resulting framework admits an interpretation as the Lagrangian relaxation of a constrained optimization problem with explicit budget constraints. Experiments against strong baselines such as ASC and ESC show that our method achieves improved trade-offs among answer correctness, sampling rounds, and total samples required.
1 Introduction
Test-time scaling improves LLM reasoning but increases inference cost, motivating adaptive sampling methods that balance answer quality with computation and latency. RL-Guided Sampling formulates this problem as an MDP and learns a lightweight, non-invasive controller from answer statistics.
- Test-time scaling improves LLM reasoning performance without additional training but increases inference cost.
- Existing adaptive sampling methods often rely on human-designed heuristics, distributional assumptions, or auxiliary signals unavailable in some settings.
- RL-Guided Sampling formulates adaptive sampling as an MDP and trains a four-layer MLP controller with reinforcement learning to balance correctness, latency, and computation cost.
- The controller uses only statistics from sampled final answers, avoiding model confidence features and intervention in the LLM reasoning process.
- The weighted RL objective corresponds to a Lagrangian relaxation of maximizing accuracy under latency and computation budgets.
- Compared with ASC, RL-Guided Sampling reduces sampling rounds by 3 times and total samples by 30%; compared with ESC, it reduces them by 10% and 35%, respectively.These improvements are reported across three benchmarks and multiple language-model samplers, with consistency across trade-off levels.
2 Method
The method models adaptive sampling as a finite-horizon MDP in which answer-pool statistics determine whether to stop or acquire more samples. Its reward combines answer quality with latency and computation penalties, while its objective admits a constrained-optimization interpretation.
- Problem Setup and MDP Formulation: At each round, the controller observes answer-pool statistics and decides whether to stop or acquire additional samples.
- Reward Design: The reward combines a step-wise penalty for additional rounds and samples with a terminal reward based on the final prediction’s correctness.
- Problem Setup and MDP Formulation: The state contains top-K answer-class counts, the total number of sampled candidates, and the entropy of those counts.
- Problem Setup and MDP Formulation: Action 0 stops sampling and returns the current majority-vote answer, while positive actions generate additional candidate answers in parallel.
- Reward Design: The terminal target is the majority-vote answer obtained at the maximum budget, encouraging stopping when the current answer distribution converges without using ground-truth labels.
- Optimization: PPO is used to optimize expected cumulative reward, while the same MDP can support value-based or other policy-gradient algorithms.
- Lagrangian View: The objective is equivalent to a Lagrangian relaxation in which non-negative penalties act as dual variables for expected sample and round constraints.
3 Experimental Settings
The evaluation uses three challenging mathematical-reasoning benchmarks, diverse LLM samplers, and adaptive-sampling baselines. Performance is assessed across accuracy, computational cost, and latency metrics, with a lightweight controller configuration.
- Benchmarks: Evaluation covers AIME24, AIME25, and HMMT 2025, while controller training uses a random 200-question subset of DAPO.
- Language-Model Samplers: The sampler set includes Qwen-3 variants and GPT-4.1-nano across 0.6B, 1.7B, and 4B scales, reasoning and instruct types, and open-source and proprietary settings.
- Baselines: Baselines include parallel Self-Consistency, sequential ASC with a p-value threshold, and batch-based ESC with intra-batch consistency stopping.
- Metrics: The five metrics cover accuracy, computational cost through total samples and total tokens, and latency through sampling rounds and sequential tokens.
- Environment and Training Setup: The environment uses K = 5, a 7-dimensional state space, actions {0, 1, 2, 4}, and a four-layer MLP policy network.
4 Results and Analysis
RL-Guided Sampling improves the accuracy–efficiency trade-off across sampling budgets, benchmarks, and model settings. It reduces sampling requirements, allocates resources according to answer statistics, and transfers across models while reward design affects performance.
- Main Results: 3–4ˆ fewer average sampling rounds and approximately 30% fewer total samples than ASC are achieved while preserving comparable accuracy.The reported comparison identifies reduced latency and computation relative to ASC.
- Main Results: Approximately 10% fewer sampling rounds, roughly 33% fewer total samples, and higher accuracy are achieved relative to ESC.RL-Guided Sampling avoids ESC’s overly aggressive early stopping while improving the overall trade-off.
- Main Results: All adaptive methods reduce total sample count versus SC, but their sampling rounds exceed SC because adaptive early stopping requires sequential evaluations.ASC’s sequential sampling creates latency, while ESC reduces rounds through coarse-grid evaluation at a performance or sample-cost trade-off.
- Scaling Curves: RL-Guided Sampling consistently achieves a favorable accuracy–efficiency trade-off across parameter configurations, outperforming strong baselines from both sample-budget perspectives.Figure 2 compares accuracy against sampling rounds and total samples on AIME24 and AIME25.
- Explanatory Analysis: The policy allocates more samples to queries with higher Answer Entropy, but allocation does not strictly follow entropy, indicating sensitivity to sampling trajectories.The study correlates average samples per query with Answer Entropy and Answer Accuracy using Qwen3-0.6B responses from the DAPO-subset.
- Ablation Study: Running Majority achieves the best overall reward-signal performance, whereas Real Label substantially degrades accuracy and increases sampling rounds and total generated samples.Running Majority also slightly outperforms Full Majority in sampling efficiency.
- Generalization Analysis: Controllers trained using different models maintain highly competitive scaling behavior on GPT-4.1-nano and Qwen3-4B-Instruct despite distribution shift.The cross-model evaluation supports robust transferability across the evaluated settings.
5 Related works
Related work addresses efficient test-time computation through adaptive sampling, confidence-weighted reasoning, and multi-objective or constrained reinforcement learning. RL-Guided Sampling adopts linear scalarization within this broader optimization perspective.
- Adaptive sampling: Adaptive sampling methods reduce fixed-budget test-time scaling costs by terminating once a consensus criterion is reached or adapting budgets to query difficulty.Other methods weight reasoning paths by confidence to recover high-quality answers from fewer samples.
- Test-time reasoning: Tree search, query diversification, and iterative refinement allocate additional test-time computation to improve complex reasoning.These approaches aggregate reasoning paths, diversify query formulations, or bootstrap self-correction.
- Multi-objective reinforcement learning: Multi-Objective Reinforcement Learning optimizes conflicting goals by discovering Pareto-optimal policies across multiple reward signals.Linear scalarization aggregates multiple objectives into a single optimization target.
- Multi-objective reinforcement learning: The proposed method relies on the linear scalarization paradigm because of its algorithmic simplicity and robust empirical performance.
- Constrained reinforcement learning: Constrained Reinforcement Learning maximizes a primary objective while satisfying limits such as safety boundaries or resource budgets.Lagrangian relaxation connects constrained reinforcement learning to multi-objective reinforcement learning.
6 Conclusion
The paper introduces RL-Guided Sampling, a lightweight reinforcement-learning framework that formulates adaptive sampling as a Markov decision process. It optimizes correctness, computation cost, and latency while outperforming strong baselines and transferring across datasets and samplers.
- RL-Guided Sampling trains an RL controller to optimize answer correctness, computational cost, and latency in adaptive test-time sampling.The framework uses final-answer statistics and does not require auxiliary signals or intervention in the LLM reasoning process.
- The CPU-friendly controller consistently outperforms ASC and ESC while reducing total samples and sampling rounds.The learned policy also exhibits strong transferability across different datasets and samplers.
Limitations
The current formulation improves the accuracy–latency–compute trade-off but remains deliberately limited in its state representation and reward design. Future refinements could add richer answer signals and real-world generation costs.
- The current formulation is an initial step toward reinforcement-learning-based adaptive LLM sampling and can be further refined.The lightweight controller already shows consistent improvements in the accuracy–latency–compute trade-off.
- The state representation uses simple sampled-answer statistics, leaving room for answer confidence and average answer length as richer signals.
- Reward design could directly incorporate real-world generation costs, including time and money, as penalties.These extensions are described as complementary to the framework and may improve alignment with deployment costs.
A Implementing Details
The implementation generates offline candidate pools, trains the RL adapter with PPO, and evaluates methods on shared sampled response sets across repeated random seeds. Results include accuracy, sampling rounds, total samples, and token-level cost metrics.
- Sampling setup: Each question has 128 candidate responses, and trajectories sample N = 32 responses from that offline pool during training and evaluation.
- Training configuration: Table 3 lists the sampling configurations used to generate candidate responses for each model.
- Training: The RL adapter is a four-layer neural network trained with PPO using Stable-Baselines3 in an OpenAI Gym-based MDP environment.
- Evaluation: The evaluation uses the same sampled response sets for all methods, repeats the process over 100 random seeds, and reports mean accuracy and computation cost.
- Training configuration: Table 4 reports the training hyperparameters of RL-Guided Sampling.
- Results reporting: Table 1 reports mean sampling rounds, total samples, and accuracy across five controllers trained with different random seeds for each model.The default training penalties are λ_lat = 0 and λ_comp = 0.0075.
- Results reporting: Token-level evaluation measures computation cost by Total Tokens and latency by Sequential Tokens, confirming the reported accuracy–efficiency trade-off.
B.2 Additional Details and Results for Scaling Analysis
The additional scaling analysis compares RL-Guided Sampling with SC, ASC, and ESC across token and sampling budgets, using parameter sweeps for each adaptive method.
- Parameter sweeps: ASC is evaluated over nine stopping thresholds, while ESC is evaluated over chunk sizes K ∈ {2, 3, 5, 7}.These sweeps provide comparison points across adaptive-sampling settings.
- Scaling comparisons: RL-Guided Sampling is evaluated against SC, ASC, and ESC across accuracy–token scaling curves on Qwen3-4B-Instruct with AIME24 and AIME25.Figure 5 reports the token-level comparison.
B.3 Additional Results for Explanatory Analysis
The explanatory analysis extends scaling and query-level comparisons to additional models, datasets, and token metrics, including AIME24 results and three-benchmark reporting.
- Explanatory analysis: Figure 6 examines correlations between total samples per query, Answer Entropy, and Answer Accuracy for AIME24 queries sampled with Qwen3-0.6B.The figure places total samples against Answer Entropy and Answer Accuracy.
- Benchmark metrics: Table 5 compares test-time scaling approaches across three benchmarks using accuracy, Sequential Tokens, and Total Tokens.Acc. denotes accuracy; Seq. Tokens and Tot. Tokens denote latency- and computation-related token metrics, respectively.
- Scaling analysis: Accuracy–token and accuracy–sampling scaling curves compare SC, ASC, ESC, and RL-Guided Sampling on Qwen3-0.6B across AIME24 and AIME25.Figure 7 provides both token- and sampling-level views.
C Proof of Proposition 1
The proof shows that the RL objective is algebraically equivalent to the Lagrangian relaxation of adaptive sampling under latency and computation budgets, while accuracy is represented by agreement with a maximum-budget consensus.
- Accuracy proxy: The MDP’s expected accuracy proxy measures agreement with the maximum-budget consensus because the environment is decoupled from ground-truth labels.With terminal rewards in {1, −1}, the expectation is an affine transformation of the matching probability.
- RL objective: The controller maximizes expected episodic return, formed by terminal answer-quality reward and intermediate sampling-cost penalties.The return sums rewards through the stopping step.
- Constrained formulation: The constrained problem maximizes final expected accuracy proxy subject to predefined latency and computation budgets.These budgets correspond to sampling rounds and total computation costs.
- Lagrangian relaxation: Non-negative dual variables λlat and λcomp weight the latency and computation constraints in the Lagrangian relaxation.The resulting expression combines expected accuracy, latency, computation, and constant budget terms.
- Equivalence: Because the scaling mappings and offset terms are policy-independent affine transformations, maximizing the Lagrangian is equivalent to maximizing the RL objective.This establishes that the MDP’s step-wise rewards solve the constrained problem’s Lagrangian relaxation.