Table of Contents
Fetching ...

Adaptive Rectification Sampling for Test-Time Compute Scaling

Zhendong Tan, Xingjun Zhang, Chaoyi Hu, Yancheng Pan, Shaoxun Wang

TL;DR

This work tackles test-time compute scaling by enabling fine-grained, step-level self-correction in large language models. It introduces Adaptive Rectification Sampling (AR-Sampling), which leverages a Process-Supervised Reward Model (PRM) as a verifier and uses trigger sentences to prompt rethink only at steps deemed likely to be erroneous, controlled by a threshold $p$ and a maximum rethink count $m$. By integrating with existing best-of-N strategies and enabling per-step verification, AR-Sampling improves solution accuracy on math reasoning benchmarks like GSM8K and MATH500 while preserving token efficiency. The approach demonstrates that targeted, adaptive rethinking can mitigate overthinking and enhance reasoning reliability without excessive token expansion, offering a practical path for scalable test-time reasoning improvements.

Abstract

The newly released OpenAI-o1 and DeepSeek-R1 have demonstrated that test-time scaling can significantly improve model performance, especially in complex tasks such as logical reasoning. Common test-time scaling methods involve generating more chains of thought (CoTs) or longer CoTs with self-correction. However, while self-correction can improve performance, it may lead to significant token waste and reduce readability of the CoT if the reasoning steps are already correct. To demonstrate that large language models (LLMs) can rectify errors at a more fine-grained level, we propose Adaptive Rectification Sampling (AR-Sampling), which can guide the LLMs to self-correction at the appropriate step. AR-Sampling leverages a process-supervised reward model (PRM) as a verifier and constructed trigger sentences to guide the model in adaptive step-level rethinking. Through the experiments on GSM8K and MATH500, it indicates that our approach enables the models to rethink in more fine-grained level, improving the accuracy of solutions, while generating a reasonable number of additional tokens.

Adaptive Rectification Sampling for Test-Time Compute Scaling

TL;DR

This work tackles test-time compute scaling by enabling fine-grained, step-level self-correction in large language models. It introduces Adaptive Rectification Sampling (AR-Sampling), which leverages a Process-Supervised Reward Model (PRM) as a verifier and uses trigger sentences to prompt rethink only at steps deemed likely to be erroneous, controlled by a threshold and a maximum rethink count . By integrating with existing best-of-N strategies and enabling per-step verification, AR-Sampling improves solution accuracy on math reasoning benchmarks like GSM8K and MATH500 while preserving token efficiency. The approach demonstrates that targeted, adaptive rethinking can mitigate overthinking and enhance reasoning reliability without excessive token expansion, offering a practical path for scalable test-time reasoning improvements.

Abstract

The newly released OpenAI-o1 and DeepSeek-R1 have demonstrated that test-time scaling can significantly improve model performance, especially in complex tasks such as logical reasoning. Common test-time scaling methods involve generating more chains of thought (CoTs) or longer CoTs with self-correction. However, while self-correction can improve performance, it may lead to significant token waste and reduce readability of the CoT if the reasoning steps are already correct. To demonstrate that large language models (LLMs) can rectify errors at a more fine-grained level, we propose Adaptive Rectification Sampling (AR-Sampling), which can guide the LLMs to self-correction at the appropriate step. AR-Sampling leverages a process-supervised reward model (PRM) as a verifier and constructed trigger sentences to guide the model in adaptive step-level rethinking. Through the experiments on GSM8K and MATH500, it indicates that our approach enables the models to rethink in more fine-grained level, improving the accuracy of solutions, while generating a reasonable number of additional tokens.

Paper Structure

This paper contains 10 sections, 3 equations, 2 figures.

Figures (2)

  • Figure 1: The framework of AR-Sampling. AR-Sampling uses PRM as a verifier to check each step. If the score is lower than the threshold, we consider this step unfavorable for the reasoning and use a trigger to force the model to rethink from the current step. By adjusting the threshold score and the maximum number of rethinks, we can dynamically control the generation budget.
  • Figure 2: Samples of the PRM score. The x axis is the normalized step index. We classify them into correct (a) and wrong (b) according to the final answers. For the wrong cases, there is greater fluctuation, and the scores of some steps are very low. We believe that these steps are more likely to lead to wrong answers.