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

Technical Summary: RL-Guided Adaptive Sampling for Test-Time Scaling

Problem Statement
Test-time scaling has emerged as a critical technique for enhancing the reasoning capabilities of Large Language Models (LLMs) without additional training. Methods such as Self-Consistency (SC), Tree-of-Thoughts, and Best-of-N sampling improve answer quality by allocating more inference-time computation. However, these gains incur substantial costs in both total computation and latency. Existing adaptive sampling methods, such as Adaptive Self-Consistency (ASC) and Early-Stopping Self-Consistency (ESC), attempt to mitigate these costs by dynamically deciding when to stop sampling. Yet, these approaches suffer from significant limitations: they often rely on human-designed heuristics or distributional assumptions, require auxiliary signals (e.g., internal model confidence, hidden states, or question difficulty) that may be unavailable, and can be invasive to standard inference pipelines. There is a need for a lightweight, principled method that balances answer correctness, latency, and computational cost without requiring auxiliary features or model intervention.

Methodology
The authors propose RL-Guided Sampling, a framework that formulates adaptive sampling as a finite-horizon Markov Decision Process (MDP).

• MDP Formulation:

  • State ($s_t$): The state is constructed solely from the statistics of the current pool of sampled answers ($D_t$). It includes the counts of the top-$K$ most frequent answer classes, the total number of samples generated, and the entropy of these counts. Crucially, the state does not rely on query-specific information or internal model signals.
  • Action ($a_t$): At each round, the controller selects an action from a discrete set: either stop sampling (action 0) or generate a specific number of additional candidate answers in parallel (actions $k_1, k_2, \dots$).
  • Reward Function: The reward balances final answer quality against inference costs.
    • Step-wise Penalty: A penalty is incurred for every additional sampling round ($\lambda_{lat}$) and for every new sample generated ($\lambda_{comp}$).
    • Terminal Reward: Upon termination (either by the controller stopping or hitting the maximum budget), the model receives a reward based on the correctness of the majority-vote answer. The target correctness ($y^\star$) is defined as the majority vote that would be obtained if sampling continued to the maximum budget $N$, decoupling the reward from ground-truth labels.
  • Optimization: The controller is trained using Proximal Policy Optimization (PPO) to maximize the expected cumulative reward. The authors note that this objective can be interpreted as the Lagrangian relaxation of a constrained optimization problem where the goal is to maximize accuracy subject to latency and computation budgets.

• Architecture and Deployment:
  The controller is a lightweight four-layer Multi-Layer Perceptron (MLP). It requires only the final generated answers to function, making it non-invasive and deployable on CPUs. It does not require fine-tuning the underlying LLM.

Key Contributions

1. MDP Formulation for Adaptive Sampling: The paper introduces a novel formulation of adaptive sampling as an MDP, moving away from heuristic rules toward a learned policy that explicitly navigates the performance-cost trade-off.
2. Lightweight and Non-Invasive Design: The method relies purely on the statistics of the sampled answer pool, requiring no auxiliary signals (like confidence scores) and no modification to the LLM's reasoning process. The controller is small enough to be trained and deployed on CPUs.
3. Theoretical Interpretation: The authors provide a theoretical connection between their RL objective and constrained optimization, showing that the weighted reward function acts as a Lagrangian relaxation for maximizing accuracy under strict resource constraints.
4. Generalization: The framework demonstrates strong transferability, where policies trained on one dataset or model (e.g., Qwen-3-0.6B) can be applied to different benchmarks and even different, more capable models (e.g., GPT-4.1-nano) with minimal performance degradation.

Experimental Results
The method was evaluated on three mathematical reasoning benchmarks (AIME24, AIME25, HMMT 2025) using various LLM samplers (Qwen-3 family and GPT-4.1-nano).

• Performance vs. Baselines: Compared to strong baselines like ASC and ESC, RL-Guided Sampling achieves superior trade-offs:
  • vs. ASC: Reduces sampling rounds by approximately 3–4$\times$ and total samples by ~30% while maintaining comparable accuracy. This addresses the high latency bottleneck of ASC's sequential sampling.
  • vs. ESC: Reduces total sample requirements by ~33% and sampling rounds by ~10%, achieving higher accuracy by avoiding overly aggressive early stopping.
  • vs. SC: While all adaptive methods reduce total sample counts compared to fixed-budget SC, RL-Guided Sampling achieves this with significantly fewer rounds than ASC and better efficiency than ESC.
• Scaling Behavior: Scaling curves indicate that RL-Guided Sampling consistently outperforms baselines across various parameter settings, offering better accuracy for the same or fewer samples and rounds.
• Resource Allocation: Analysis shows the learned policy intuitively allocates more samples to queries with higher answer entropy (diversity), aligning with the need for more samples to converge on a reliable majority.
• Ablation Studies: The "Running Majority" target (majority vote at the max budget $N$) was found to be superior to using ground-truth labels (Real Label), which degraded performance by introducing noise and preventing the learning of universal stopping criteria.

Significance and Claims
The paper claims that RL-Guided Sampling offers a principled, lightweight alternative to existing adaptive sampling methods. By framing the problem as an MDP and training a small controller via RL, the method effectively balances answer correctness, latency, and computation cost without requiring invasive modifications to the LLM or access to internal model states. The authors emphasize that the approach is particularly valuable for practical deployment, as it can be trained on inexpensive open-source models and transferred to reduce the inference costs of stronger, proprietary, or closed-source systems. The work represents an initial step toward RL-based adaptive sampling, with future directions suggested for incorporating richer state signals and real-world cost metrics.