Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling

When solving hard reasoning problems, large language models (LLMs) often need to try multiple times to get the right answer. This "test-time scaling"—letting the model think longer by sampling many possible reasoning paths—is effective. However, it is incredibly expensive in terms of both total computation and latency (the delay before a response is received).

Current methods attempt to mitigate this by deciding when to stop sampling. They usually rely on rigid, human-designed rules or specific distributional assumptions. This paper proposes a different path: treating adaptive sampling as a reinforcement learning (RL) problem. By training a tiny, lightweight controller to manage the sampling process, the authors claim they can find a better balance between accuracy and cost.

The Problem

The status quo for test-time scaling involves methods like Self-Consistency (SC). SC simply samples a fixed number of responses and takes a majority vote. While effective at improving accuracy, it is computationally wasteful. It treats easy and hard questions identically.

Existing adaptive methods try to do better. Adaptive Self-Consistency (ASC) samples one response at a time. It stops when a certain probability threshold is met. Early-Stopping Self-Consistency (ESC) uses fixed batches to save on latency. However, as shown in, these methods are limited by their heuristics.

ASC is often too sequential. This drives up latency through excessive sampling rounds. ESC can be too aggressive. It may sacrifice accuracy to save on total samples. Most importantly, these methods struggle to navigate the tension between accuracy, latency, and compute cost.

How It Works

The authors treat adaptive sampling as a finite-horizon Markov Decision Process (MDP). Instead of a hardcoded rule, they deploy a lightweight four-layer MLP (a Multi-Layer Perceptron, a basic type of neural network) that acts as a "controller." This controller decides whether to stop or to pull more samples.

The mechanism works in three stages:

1. State Observation: At each round, the controller looks at the current pool of answers. It does not look at the prompt or the model's internal logits (the raw output probabilities). It only sees the statistics of the final answers. The state includes the counts of the top-$K$ most frequent answers, the total number of samples collected, and the entropy (a measure of uncertainty or disorder) of the answer distribution.
2. Action Selection: Based on this state, the controller selects an action. It can either stop sampling ($a=0$) or generate a specific number of additional responses in parallel ($a \in {k_1, k_2, ...}$). This allows the system to jump ahead in batches when uncertainty is high. This avoids the slow, one-by-one cadence of ASC.
3. Reward Optimization: The controller is trained via Proximal Policy Optimization (PPO) to maximize a reward function. This function balances quality and cost. The reward includes a terminal reward based on whether the majority-vote answer matches the "running majority" (the consensus if sampling continued to the maximum budget). It also includes step-wise penalties for every new sampling round or additional sample generated.

The authors prove that this RL formulation is a Lagrangian relaxation of a constrained optimization problem. This means the weights used to penalize latency and compute ($\lambda_{lat}$ and $\lambda_{comp}$) act as dual variables. These allow users to set explicit budgets for how much they are willing to spend.

Numbers

The paper reports efficiency gains across mathematical reasoning benchmarks (AIME and HMMT). The most striking result is the reduction in overhead. The authors demonstrate that RL-Guided Sampling reduces sampling rounds by 3x compared to ASC. This helps lower the time spent waiting for sequential steps. It also reduces total samples by 30% compared to ASC. Compared to ESC, it reduces sampling rounds by 10% and total samples by 35%.

These wins matter for production. Looking at, the method consistently achieves a superior Pareto frontier.

This means for any given budget of samples or rounds, it delivers higher accuracy than the baselines. The efficiency holds at the token level. shows that the method reduces both total tokens (compute) and sequential tokens (latency) across different model scales.

The hardware footprint is also small. The controller is a tiny MLP. It can be trained and deployed on a CPU. This adds negligible overhead to the GPU-heavy task of LLM inference. The authors also demonstrate strong generalization in .

A controller trained on a tiny 0.6B model can be applied directly to a massive, closed-source model like GPT-4.1-nano. It maintains competitive scaling behavior despite the change in model.

What's Missing

There are gaps that a practitioner should note. First, the state representation is quite thin. It relies exclusively on the frequency and entropy of final answers. While this makes the method "non-invasive," it ignores other potential signals. It does not use answer length or the model's own confidence scores. Adding these might improve the controller's "intuition," but it would increase integration complexity.

Second, the reward function uses "Running Majority" as a proxy for truth. An ablation study ([Table 2]) shows that using the actual ground-truth label as a reward signal actually performs worse. This happens because the controller lacks the semantic context to understand why an answer is correct. The controller is essentially a pattern matcher for convergence, not a judge of truth.

Finally, the paper does not explore how this behaves in highly heterogeneous environments. If your workload mixes very short queries with massive multi-step math proofs, a single controller might struggle. It is unclear how well one policy handles such vastly different scales of "convergence."

Should You Prototype This

Yes, if you are looking to optimize test-time scaling costs. The core value proposition is efficiency. If you are already using Self-Consistency or ASC, you might slash your inference bill by 30% without losing your reasoning edge.

The implementation seems straightforward. You do not need to modify your LLM engine. You just need a small service to sit alongside it. This service collects answer statistics and tells the engine when to stop. Code is available at https://github.com/RunpengDai/RL-Guided-Adaptive-Sampling. Since the controller is CPU-bound and tiny, the barrier to entry is low. Try it on a small open-source model first. See if the convergence patterns match your production traffic.

Figures from the paper