Improving the Efficiency of Test-Time Search in LLMs with Backtracking

Authors: Anikait Singh, Kushal Arora, Sedrick Keh, Jean Mercat, Tatsunori Hashimoto, Chelsea Finn, Aviral Kumar
Primary area: foundation or frontier models, including LLMs
Keywords: LLMs, Reasoning, Test Time Inference, Backtracking, In-Context Verifiers

Abstract:
Solving reasoning problems is an iterative multi-step computation, where a reasoning agent progresses through a sequence of steps, with each step logically building upon the previous one to reach a desired conclusion. If the desired solution is not attained, the agent must backtrack and try reasoning chains that are quite different from previous attempts. Though prior work such as test-time search against an outcome verifier can improve performance, most search is done in parallel via Best-of-N reranking, and independently for each attempt at a problem, thus wasting a significant amount of computation in sampling multiple full solutions even beyond the point that is needed. Can we reduce the total amount of computation by sharing information and computation across multiple attempts to a given problem? In this paper, we build a novel approach combining process verifiers that predict likelihoods of success \emph{per step} with preemptive backtracking to maximize performance per generated token. To do this, the PRM can be used to identify where a problematic step in a solution trace is by using the sensitivity of the predictions of the learned verifier and allowing the model to do focused resampling of the problematic portion of a solution. This approach can significantly reduce the amount of computation by leveraging partial computation from previous revisions. To further enhance the computational efficiency of inference, we introduce in-context process supervision, where the verifier is conditioned on the history of revisions that are attempted, reducing uncertainty in the verification decisions and improving the verifier's confidence with each round of backtracking. This framework for iterative backtracking, leveraging in-context process supervision, enables an effective tradeoff between inference and model performance.

Decision: reject
Reviewer scores: 3, 5, 5, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1730644330753

thread_type:
official_review

<thread_object>
creator: reviewer_1tmn

<content>
summary:
The paper proposes to improve efficiency of test-time search in LLMs with backtracking based on PRMs.

strengths:
The paper correctly identifies a problem with independent monte carlo sampling in high dimensional spaces.

weaknesses:
This is a very old idea, well known in statistics and Markov chain optimization/inference. Sampling the whole trace is called independent sampling. Resampling from a midpoint (whether uniformly or with weighting) is called naive sitewise resampling. Sitewise resampling where there is an attempt to reuse, rather than always resample, the tail is called sitewise resampling with reweighting. This is how sitewise Metropolis Hastings (inference)/ simulated annealing (optimization) works. Reinventing the wheel is often a good idea, including in this case, provided the sources are properly cited. This paper ignores work on sitewise resampling completely. 

Performance-wise, I can assume that whoever chose to use 'linear search' for test-time computations in LLMs didn't know  better, but most likely, linear search variants were chosen because they are easily parallelizable. Sitewise resampling is not. However, there are methods that they can both benefit from PRM and are parallelizable. These methods belong to sequential Monte Carlo family. The paper omits time analysis of parallel computations and does not compare wall-clock performance of linear and backtracking search, thus ignoring this problem. If PRM were used to improve the search performance at test time while still preserving advantages of parallel computations, a variant of sequentail Monte Carlo (that is, reweighting and resampling multiple solution based of PRM) should have been worth investigation.

questions:
How are wall-clock times of linear and backtracking search solutions compared, on single queries? Can backtracking search provide interactive experience the way linear search does?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730761541485

thread_type:
official_review

<thread_object>
creator: reviewer_qzqj

<content>
summary:
This paper introduces a framework to enhance reasoning tasks using PRM-based backtracking by enabling shared reasoning information and computations across multiple process verifiers. The approach predicts the likelihood of success at each step and identifies problematic steps. By incorporating in-context process supervision, the framework reduces uncertainty in verification decisions, significantly improving inference efficiency.

strengths:
1. I like the idea that identifying problematic steps in the reasoning chain and resample only those parts, rather than starting over with each new attempt. This adaptive approach reduces computational load, allowing for more targeted corrections that effectively address irrelevant errors and enhance the overall quality of the framework.
2. Backtracking in the framework allows the model to locate the root cause of reasoning error. In-context process-supervision allows the model to adapt the prediction results during the search process and increase the confidence about the learning results. 
3. Evaluation in the MATH domain shows that this approach improves test-time computational efficiency by 15% compared to the linear search algorithm.

weaknesses:
1. The dataset is limited to MATH domain. 
2. The words in the figures are too small to read. The writing needs to be improved.

questions:
Why was a linear search algorithm chosen as the SOTA benchmark in your experiments? Which specific linear search algorithm was used, and is it the one from Snell et al. (2024)? Additionally, why were no comparisons made with the works mentioned in the related work section? Please clarify the rationale behind these choices, as the novelty and contributions of this work are unclear without this context.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730810075926

thread_type:
official_review

<thread_object>
creator: reviewer_pd42

<content>
summary:
The problem focuses on solving reasoning problems using LLMs which typically involves complex, multi-step generation where small mistakes in early steps can ultimately lead to incorrect outputs. While previous approaches have focused on linear search in parallel using a fixed budget, the proposed approach is non-linear and based on backtracking using learned process-based verifiers which allows localizing errors and proposing target revisions for these errors. Experiments show improvement of 15% over linear search algorithms.

strengths:
Strength:
- Interesting and novel approach utilizing backtracking to improve test-time generation efficiency
- Paper is mostly clear and written well
- Experiments show promising results for the proposed approach

weaknesses:
Weaknesses:
- I could not find clear comparison to linear search approaches like beam search, Best of N, etc., e.g., in terms of performance vs. budget. For example, we can compare a large number of candidates generated in parallel vs. additional turns; the paper explains why this may be inefficient in Section 3, but I did not see experiments validating this claim. Also, it was not clear to me how the 15% improvement is computed.

- I found parts of Algorithm 1 to be a bit confusing. It takes as an input "Advantage threshold" and "Revision margin t", however: (1) advantage threshold uses the same symbol as the initial trajectory; (2) both do not seem to be used in the algorithm; (3) I imagine one of them corresponds to the "Advantage smoothing for effective backtrack step selection" but it is not clear which one of them.

- Limited technical novelty: while the concept of backtracking to the correct step is novel, most elements including the PRM that provides per step estimates and its training have been adopted from previous literature (Snell et al., Wang et al.)

- Experimental setup could be extended to evaluate performance on other reasoning tasks and additional LMs (currently using Llama 8B, would be interesting to see performance on additional, larger models)

Minor:
Acronyms: "PRM" appears in abstract before definition, "SFT" in appendix without definition

questions:
See point under "weaknesses" above. In particular, I would appreciate clarification regarding the comparison with linear search algorithm.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730818874673

thread_type:
official_review

<thread_object>
creator: reviewer_ykyw

<content>
summary:
This paper aims to improve the efficiency of reasoning during inference. Specifically, this paper first uses process verifiers to predict likelihoods of success per step (for example using the process reward model to estimate the value function), which can be used to identify a problematic step. This significantly reduces the amount of computation. To further enhance the computational efficiency of inference, the authors introduce in-context process supervision.

strengths:
(1) This proposed method improve the computational efficiency if one can use process reward model to get a good value function estimator and identify the problematic step correctly. And the empirical result in mathematical reasoning validate this.

(2) This method works well for multi-step tasks.

weaknesses:
(1) The performance of the backtracking search method is largely dependent on the reward shaping methods, i.e., process reward model (PRM).  If the estimated rewards (and value functions) is not accurate enough, it will result in great challenges in backtracking search process. As the true reward is sparse (only 0,1 in the final), it seems not easy to estimate the process reward well.  

(2) The authors evaluate how backtracking would perform on reasoning problems (Hendrycks Math dataset). Would this generalize well to other types of data or reasoning tasks?

Minors:

In the abstract and introduction, the abbreviation "PRM" appears before its full term is introduced, it is a little confusing.

questions:
See above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730819399899

thread_type:
official_review

<thread_object>
creator: reviewer_rbmd

<content>
summary:
This work focuses on the inference acceleration of LLMs by using a backtracking method when the LLMs outputs the inference intermediate steps. The authors propose to detect errors in intermediate reasoning steps using PRMs and resamples the incorrect part of a solution, saving computational resources. The authors also introduce the so-called in-context process supervision, where the model adjusts verification based on past attempts. This further reduces the computation of LLMs.

strengths:
1. This paper proposes new methods to localize the incorrect part of the responses generated by the LLMs. The proposed in-context process supervision adapts verification based on historical data, which improves the performance of LLMs inference.

2. The efficacy of the proposed methods are verified on the mathematical problems.

weaknesses:
1. The experiments of the work mainly focus on mathematical problems. However, the applicability of the proposed algorithm to other domains, e.g., code and legal reasoning, remains untested.

2. The efficacy of the proposed algorithm heavily depends on the accuracy of PRM. The incorrect error localization could lead to inefficient backtracking and degenerate performance. The influence of the PRM accuracy is not well-explored in the current work.

questions:
Same as the Weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730821832599

thread_type:
official_review

<thread_object>
creator: reviewer_vxtv

<content>
summary:
The paper presents a new method for improving the reasoning capabilities of large language models. The reasoning process can be viewed as a search process through a sequence of steps with each step logically building upon the previous ones until the desired conclusion is reached. This process can be formally modelled as a Markov Decision Process and solved via reinforcement learning. A key component of the solving approach is a process verifier that can predict the likelihood of success per step and thus facilitate backtracking to maximise performance per generated token. These process verifiers can be learned from relatively small datasets. The empirical evaluation seems to indicate that the proposed approach improves the reasoning process by employing preemptive backtracks.

strengths:
- Improving the reasoning capabilities of existing language models is definitely a topic of considerable interest in the AI community. The proposed approach seems to address this issue in a principled manner.

weaknesses:
- The presentation is quite dense and therefore it is not easy to follow the details of the paper. The example from Figure 2 is nice and helps to understand the basic idea. However, Sections 3.2 to 3.5 which contain the main contribution are not easy to understand. Perhaps including more illustrative examples here would improve the quality of the presentation.

- The empirical evaluation seems to be limited to a single dataset. In addition, it is hard to understand what the plots in Figure 7 present.

questions:
- Is it possible to do some kind of transfer learning, namely learning a process verifier in one domain and either use it directly or fine-tuning for another domain?

- What language models were used in the evaluation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1733209812608

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overall response to All Reviewers

comment:
We sincerely thank the reviewers for their constructive feedback and insightful suggestions, which have significantly improved our paper. Below, we summarize the major changes and address specific comments in individual responses.

**Key Updates:**

- *Expanded Empirical Evaluation on CodeContests*
(vxtv, rbmd, ykyw, pd42, QZQJ)
We added results using the oracle verifier on CodeContests, demonstrating the efficacy of backtracking. Additionally, we included [results](https://imgur.com/a/CLTTIAD) with a stronger model (Yi-Coder-9B-chat), showing that the findings extend beyond the Llama family of models. We will include the full set of results with a learnt PRM for the final revision of the paper.
- *Addressing Perceived Limitations Due to PRM Accuracy*
(vxtv, rbmd, ykyw)
We conducted an [additional experiment](https://imgur.com/gallery/QXt1hN8) highlighting how even simpler search algorithms, such as Best-of-N (leveraged in prior work such as Snell et al 2023), are affected by inaccuracies in the PRM. We also discussed practical considerations for learning an effective PRM in section 4.2 for error localization and search and proposed in-context value verification in section 4.3 to reduce verifier uncertainty during subsequent revisions.
- *Latency*
(1tMN)
We measure the generation latency by comparing sequential revisions to parallel sampling. Despite seeing an increase in latency due to the sequential generation, our approach leads to a favorable tradeoff between performance (accuracy) and FLOPs, an established metric for computation cost as used in prior work such as Sardana et al. (2023). Additionally, both prior academic work such as RISE (Yu et al, 2024) and Snell et al 2024 as well as SOTA reasoning systems such as O1-preview, Deepseek r1-lite, and Qwen QWQ, similarly incur latency trade-offs due to sequential sampling aimed at improving the performance of the overall reasoning system. Therefore, this is a metric we shouldn’t be penalized for.
- *Improved Clarity Through Rewriting*
(general)
We significantly rewrote Sections 3 and 4 to enhance clarity and better articulate the contributions of our method. We are open to making further revisions to refine the paper’s writing and presentation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes to improve efficiency of test-time search in LLMs with backtracking based on PRMs. This reasoning process of LLMs can be formally modelled as a Markov Decision Process and solved via reinforcement learning. A key component of the solving approach is a process verifier that can predict the likelihood of success per step and thus facilitate backtracking to maximise performance per generated token. 

In general, reviewers felt that the work has some merit, as it localizes the incorrect parts of LLMs' outputs. Experiments on Math datasets are also illuminating. 

However, reviewers were unanimous in their assessment that this paper doesn't cross the bar for ICLR. Let me summarize some of their reservations. 

1) The experiments of the work mainly focus on mathematical problems. This is a major limitation in view of the broad applicability of LLMs. 
2) The exposition of the paper and presentation of the ideas in the paper were not ideal, making it hard to comprehend. I acknowledge that this improved during the course of the discussion, but feel that the discussion phase is not the right point to refine the paper. It requires another round of thorough review.
3) Limited technical novelty: while the concept of backtracking to the correct step is novel, most elements including the PRM that provides per step estimates and its training have been adopted from previous literature (Snell et al., Wang et al.).
4) Comparison to linear search is inadequate. This is was mentioned by multiple reviewers. 
5) Language models used are not sufficiently large.

Personally, I am not an expert at the topic. Hence, my assessment is based on the reviewers' feedback, which is unanimous. I acknowledge the authors' message to me regarding their objections to Reviewer 1TMN viewpoint regarding the applicability of linear search in this setting. However, this point is mentioned several times by various reviewers. Thus, I cannot ignore it and encourage the authors to take this comment seriously into consideration in their comparisons to the baselines in future work. At present, the work is premature and does not pass the bar for ICLR.