Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling
Summary
This paper formulates adaptive sampling for large language models as a Markov decision process and trains a lightweight RL controller to balance correctness, latency, and computational cost, achieving improved trade-offs.
View Cached Full Text
Cached at: 06/03/26, 07:36 AM
Paper page - Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling
Source: https://huggingface.co/papers/2606.03102
Abstract
Adaptive sampling for large language models is formulated as a Markov decision process and optimized using reinforcement learning to balance correctness, latency, and computational cost.
Test-time scaling improves the reasoning performance oflarge language modelsbut incurs substantial cost in both total computation and latency. Existingadaptive samplingmethods 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 formulateadaptive samplingas aMarkov decision process(MDP). We train a lightweight sampling controller withreinforcement 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 theLagrangian relaxationof aconstrained optimizationproblem 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.
View arXiv pageView PDFAdd to collection
Get this paper in your agent:
hf papers read 2606\.03102
Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash
Models citing this paper0
No model linking this paper
Cite arxiv.org/abs/2606.03102 in a model README.md to link it from this page.
Datasets citing this paper0
No dataset linking this paper
Cite arxiv.org/abs/2606.03102 in a dataset README.md to link it from this page.
Spaces citing this paper0
No Space linking this paper
Cite arxiv.org/abs/2606.03102 in a Space README.md to link it from this page.
Collections including this paper0
No Collection including this paper
Add this paper to acollectionto link it from this page.
Similar Articles
LEEPS: Latent-Guided Explore-Exploit Prompt Sampling for Efficient RLVR in Large Language Models
This paper introduces LEEPS, a latent-guided explore-exploit prompt sampler for efficient reinforcement learning with verifiable rewards (RLVR) in LLMs. It adaptively balances reuse of informative prompts and exploration of uncertain ones, improving reasoning benchmark scores by 2.6-3.7% over baselines while adding only ~2 seconds of overhead per training step.
Towards Robust Reinforcement Learning for Small-Scale Language Model Agents
This paper systematically investigates instability in reinforcement learning for small language model agents (70-500M parameters), identifying three failure modes and proposing robust techniques including a merge-and-reinitialize adapter approach and safety mechanisms; it achieves stable convergence and improved win rates.
Towards Robust Reinforcement Learning for Small-Scale Language Model Agents
This paper systematically investigates failure modes in reinforcement learning for small language models (70-500M parameters) using PPO, identifies silent LoRA freezing, numerical overflow, and catastrophic policy collapse, and proposes a robust system with merge-and-reinitialize adapters, float32 precision, and a safety mechanism. The approach converges stably and outperforms baselines with less data.
Scaling with Confidence: Calibrating Confidence of LLMs for Adaptive Test Time Scaling
The paper proposes C3RL, a reinforcement learning algorithm that calibrates LLM confidence while maintaining accuracy, and CAS, a confidence-based adaptive test-time scaling strategy that reduces inference costs by up to 12.33 times.
Towards Scalable Multi-Task Reinforcement Learning with Large Decision Models
This paper introduces LDM-v0, a large decision model trained offline on trajectories from thousands of diverse reinforcement learning environments, demonstrating that a single transformer policy can match the performance of task-specific policies across robotics, autonomous driving, inventory management, cybersecurity, trading, and video games.