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

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we revise the related work to discuss sitewise resampling and measure wall-clock time for linear and backtracking search. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

This paper ignores work on sitewise resampling completely.

Thank you for the references. We have added sitewise resampling to the related works section of the paper. Please let us know if there are any additional sources you would like us to add in this discussion.

Comparisons based on linear search and sequential Monte-Carlo

Thanks for the pointer to sequential Monte-Carlo methods, we have added and discussed a number of them in the related work section of the paper (in bold). Let us know if there are any sources you would recommend adding or revising in this list.

While we do agree that linear search algorithms with PRMs can be naive and improvements using SMC methods can be made in the long run, we would like to note, to best of our knowledge, this is the first work that applies ideas similar to sitewise sampling to search with PRMs for improving reasoning. Snell et al. 2024 present one of the first results that show effective results with beam search on PRMs. Our work pushes the frontier in this space and compares to most existing prior work in this space.

We agree with the reviewer that the community and the authors should look to statistics for literature on non-linear efficient search and apply them to the test-time compute framework. We would like to consider this as future work, highlighting this to the community."

We thank the reviewer for their comments and feedback, especially for agreeing that our method approaches improving the reasoning capabilities of language models in a principled manner. To address the main concerns, we have now updated the paper to improve clarity of Sections 3.2 and 3.5 and explain them briefly below. In addition, we add new experimental results showing that our approach for backtracking is effective on coding problems. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

The empirical evaluation seems to be limited to a single dataset.

We have added CodeContests, as an additional dataset to show the efficacy of our approach on competitive programming, complementing our results in MATH (Hendryks et al, 2023), which studies high school math problems. In this task, an LLM must write code for a set of contest coding problems from a variety of sources (e.g CodeForces, HackerEarth, and more). To formulate code generation as an MDP, we define a step/action as a line of code, and the state/context as the concatenation of the contest problem (consisting of the problem description, runtime requirements, and sample input/output pairs) and the code so far. Evaluation is done with a python execution server on a set of held out contest problems. We present preliminary results with an Oracle Verifier, comparing our backtracking algorithm (orange) with Best-of-N (Blue, Snell et al. 2024) and Backtracking to a random step (Green).

We find an average improvement of 10.16% from the linear search algorithm of Best-of-N, Snell et al. 2024 (backtracking from the first step) and significantly outperform backtracking to a random step. Note this is preliminary and we hope to train a PRM with Monte-Carlo Rollouts by the end of the rebuttal period or for the final version of the paper.

It is hard to understand what the plots in Figure 7 present.

In this figure, we show how backtracking can allow for optimization of the value function through subsequent revisions. The goal of backtracking is to reset to the portion of the solution right before the model went on the wrong track. With our approach, we are able to identify this reset point with the advantage,( computed as the delta in values between subsequent reasoning steps). We identify the backtracking step as one with the maximum negative delta with its subsequent step, i.e., one with minimum advantage. For example, in the 1st revision, we found that the lowest advantage is in Step 5, which corresponds to the biggest drop in the value function from step 5 to step 6. When we resample trajectories from this poor step (Step 5), we find that the new trajectory generated has higher values. Looking over multiple revisions, we see that the value function is monotonically increasing over the course of the revisions, which is a desirable property of the backtracking algorithm. Note that the value function corresponds to the probability of success of a particular state (context) and action (step) pair, i.e, a higher value at a step indicate higher chances of generating successful trajectory. Hence, increasing the value function should lead to a higher accuracy for a set of problems. We will update the paper to explain this better..

The presentation is quite dense and therefore it is not easy to follow the details of the paper.

We have updated the paper and have added clarity as you have suggested in the preliminaries and method section of the paper. We would be happy to clarify any questions you may have that leads to a better understanding of this paper.

What language models were used in the evaluation?

For our experiments, we use the Meta Llama-3.1-8B-Instruct models.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add new results showing the efficacy of backtracking in the code contests domain, clarify the linear search algorithm, and updated the paper for clarity. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

The dataset is limited to MATH domain. We have added CodeContests, as an additional dataset to show the efficacy of our approach on competitive programming, complementing our results in MATH (Hendryks et al, 2023), which studies high school math problems. In this task, an LLM must write code for a set of contest coding problems from a variety of sources (e.g CodeForces, HackerEarth, and more). To formulate code generation as an MDP, we define a step/action as a line of code, and the state/context as the concatenation of the contest problem (consisting of the problem description, runtime requirements, and sample input/output pairs) and the code so far. Evaluation is done with a python execution server on a set of held out contest problems. We present preliminary results with an Oracle Verifier, comparing our backtracking algorithm (orange) with Best-of-N (Blue, Snell et al. 2024) and Backtracking to a random step (Green).

We find an average improvement of 10.16% from the linear search algorithm of Best-of-N, Snell et al. 2024 (backtracking from the first step) and significantly outperform backtracking to a random step. Note this is preliminary and we hope to train a PRM with Monte-Carlo Rollouts by the end of the rebuttal period or for the final version of the paper.

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)?

In Figure 5, backtracking to the first step corresponds to Best-of-N Sampling with a PRM (Snell et al, 2024, Charniak & Johnson, 2005). We want to clarify that our comparisons consist of the most prevalent search algorithms typically used in conjunction with PRMs to the best of our knowledge. In particular, the linear search algorithm consists of Best-of-N and it is the one from Snell et al. 2024. Note we make use of parallel sampling in our backtracking algorithm after we localize the incorrect step with the PRM (Section 3.2 - Suffix Generation and Stopping Criteria). Thus, for comparison, the choice of the parallel sampling algorithm is consistent for baselines and the backtracking approach. For works that alter the proposal distribution (RISE, SCoRe, and Self-Refine), these works are complementary to our work as our proposal distribution can be substituted with the modified learned distributions presented in these works. We hope to add additional linear search approaches such as beam search for the final iteration of the paper.

The writing needs to be improved.

Thank you for the feedback, we have uploaded an updated version of the paper for your review. Please let us know if you have any additional suggestions or comments about this?

Experimental setup could be extended to evaluate performance on other reasoning tasks

We have added CodeContests, as an additional dataset to show the efficacy of our approach on competitive programming, complementing our results in MATH (Hendryks et al, 2023), which studies high school math problems. In this task, an LLM must write code for a set of contest coding problems from a variety of sources (e.g CodeForces, HackerEarth, and more). To formulate code generation as an MDP, we define a step/action as a line of code, and the state/context as the concatenation of the contest problem (consisting of the problem description, runtime requirements, and sample input/output pairs) and the code so far. Evaluation is done with a python execution server on a set of held out contest problems. We present preliminary results with an Oracle Verifier, comparing our backtracking algorithm (orange) with Best-of-N (Blue, Snell et al. 2024) and Backtracking to a random step (Green).

We find an average improvement of 10.16% from the linear search algorithm of Best-of-N, Snell et al. 2024 (backtracking from the first step) and significantly outperform backtracking to a random step. Note this is preliminary and we hope to train a PRM with Monte-Carlo Rollouts by the end of the rebuttal period or for the final version of the paper.

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.

In Figure 5, backtracking to the first step corresponds to Best-of-N Sampling with a PRM (Snell et al, 2024, Charniak & Johnson, 2005). We want to clarify that our comparisons consist of the most prevalent search algorithms typically used in conjunction with PRMs to the best of our knowledge. In particular, the linear search algorithm consists of Best-of-N and it is the one from Snell et al. 2024. Note we make use of parallel sampling in our backtracking algorithm after we localize the incorrect step with the PRM (Section 3.2 - Suffix Generation and Stopping Criteria). Thus, for comparison, the choice of the parallel sampling algorithm is consistent for baselines and the backtracking approach. For works that alter the proposal distribution (RISE (Qu et al 2024), SCoRe (Kumar et al 2024), and Self-Refine (Madaan et al, 2023)), these works are complementary to our work as our proposal distribution can be substituted with the modified learned distributions presented in these works as stated in the related works section. We hope to add additional linear search approaches such as beam search for the final iteration of the paper.

Also, it was not clear to me how the 15% improvement is computed.

We use the accuracy vs tokens plot in Figure 5 for the learnt PRM to compute the average percent improvement over tokens of Backtracking with the Multi-Turn Value Function over the best performing baseline, Baseline 1: Backtracking from the First Step (Best-of-N). Therefore, the Multi-Turn Value Function has approximately a 15% improvement from Baseline 1. We will modify the text to avoid this confusion.

[...] additional LMs (currently using Llama 8B, would be interesting to see performance on additional, larger models)

Due to compute constraints, it would be difficult to evaluate with additional larger models in time for the rebuttal. We plan to run evaluation with more models for the final iteration of the paper. Note in prior work (e.g Snell et al 2024, Setlur et al 2024) a single PaLM/Gemma model was utilized for the entirety of the experiments they present in their work. One additional evaluation we hope to run is an oracle verifier with different models for the code contests domain by the end of the rebuttal period or for the final version of the paper.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add new results showing that our approach is effective on coding problems from the CodeContests dataset. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

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.

We have updated Algorithm 1 in the new revision of the paper (for your reference).The advantage threshold, referred now by $\omega$, is used and we keep the initial solution trajectory as $\tau$. We have removed the revision margin $t$ for clarity and added the stopping criterion.

To walk through the psuedocode, we start with an input trajectory $\tau = \{s_1, s_2, \dots, s_N\}$. Next, we compute an advantage function $A(s_i, a_i)$ for each step within the trajectory. In each iteration, we identify the step $s_{i_{\text{revise}}}$ with the lowest advantage. To allow for estimation errors, we smooth the advantage function and choose a more conservative step, $i_{\text{smooth}}$, to backtrack to, increasing the likelihood that the step we bactrack to and its previous steps are correct. Then, the trajectory is then resampled starting from the step $T=i_{\text{smooth}} - 1$ by using a linear search algorithm (e.g., Best-of-N) as a subroutine. If the new trajectory $\tau{\prime}$ improves the solution, we update the best trajectory seen so far $\tau_{\text{best}}$. This process is repeated until either a maximum number of iterations M is reached or the stopping criterion is met (using the PRM’s value at the final step of the trajectory to evaluate correctness). The algorithm finally outputs the revised trajectory $\tau_{\text{best}}$, representing the most improved version of the input trajectory.

Please let us know if there is anything more we can do to improve the clarity of the algorithm.

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.)

We would like to clarify that the main contribution of our work is in showing that sequential improvements on a response when backtracking with a PRM and running sequential sampling. Concretely, we summarize our contribution as follows:

1. Sequential improvement by backtracking: We propose a novel approach to sequentially improve model reasoning by using the advantage function to identify the incorrect step and revising our solution from there. This helps, both, avoid excessive compute used by best-of-N and beam search in re-generating the correct part of the response and continuing after a fatal error. This improves, both computational efficiency and accuracy of the search.

2. Improving the value function calibration: We propose multiple approaches for improving the value function calibration for backtracking. These include 1.) Using an in-context verifier to reduce uncertainty of the value function on future revisions 2.) fine-tuning with a balanced dataset to deal with the dataset imbalances arising from backtracking and 3.) smoothing of the advantage function and conservatively choosing a step to revise to, ensuring that the step is in a correct portion of the solution.

Our approach provides substantial improvements in terms of performance as a function of inference-compute budget (Figure 5 + Figure 1 Rebuttal, CodeContests), which we believe should be of value and interest to the community interested in LLMs and scaling test-time compute. We additionally show that given a small inference token budget, we can obtain the best response.

We do note that Section 3.5 does include a set of practical considerations of learning a PRM a portion of which has been adopted from previous literature primarily because other details are needed to train an effective verifier. However, other practical conditions we address (such as issues with dataset balance for backtracking as well as smoothing of the advantage function) are additional novel contributions. We do not claim that learning the verifier itself is novel but rather using it in this way for search is. For completeness, we still wanted to identify and communicate details that helped the training of value verifiers in the context of the backtracking algorithm.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add new results showing that our approach is effective on coding problems from the CodeContests dataset. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

(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.

Note, that the underlying assumption for using verification is that there exists an insignificant gap between the generator and the verifier. This clearly exists in domains such as Math, Programming, Puzzles, and more. Thus, even if noisy, learning a PRM allows for some benefit when inference-time computation is utilized in tandem with it.

The accuracy of the verifier is indeed important for the algorithm's success. However, all search-based approaches that leverage a verifier (e.g Snell et. al) have this inherent requirement. Below, we consider the simpler search algorithm of Best-of-N. Here, we artificially add label noise to the Oracle verifier and see what the performance of the algorithm looks like for a subset of the math dataset. As seen in the following plot, as the verifier accuracy decreases, there is a large drop in the performance of a simpler search algorithm such as Best of N. Thus, given that this is a limitation for prior approaches, we should not be penalized for needing a verifier, with sufficient accuracy.

For this reason, we discuss in detail what it takes to practically learn such a verifier. For example, we outline in Section 3.5 how to practically learn a PRM to attain a higher verification accuracy. Additionally, we propose in-context value verification in Section 3.4 to further reduce the uncertainty of the verifier when performing revisions. Finally, we smooth the verifier during backtracking as seen in section 4.2 to be conservative with the step we backtrack to, ensuring that the step chosen to restart generation is indeed in a correct portion of the previous solution. This reduces the inefficiencies that come from backtracking to an incorrect step, and the smoothness can be relaxed given a stronger verifier.

(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?

We have added CodeContests, as an additional dataset to show the efficacy of our approach on competitive programming, complementing our results in MATH (Hendryks et al, 2023), which studies high school math problems. In this task, an LLM must write code for a set of contest coding problems from a variety of sources (e.g CodeForces, HackerEarth, and more). To formulate code generation as an MDP, we define a step/action as a line of code, and the state/context as the concatenation of the contest problem (consisting of the problem description, runtime requirements, and sample input/output pairs) and the code so far. Evaluation is done with a python execution server on a set of held out contest problems. We present preliminary results with an Oracle Verifier, comparing our backtracking algorithm (orange) with Best-of-N (Blue, Snell et al. 2024) and Backtracking to a random step (Green).

We find an average improvement of 10.16% from the linear search algorithm of Best-of-N, Snell et al. 2024 (backtracking from the first step) and significantly outperform backtracking to a random step. Note this is preliminary and we hope to train a PRM with Monte-Carlo Rollouts by the end of the rebuttal period or for the final version of the paper.

We thank the reviewer for their comments and for engaging with the paper. To address the concerns, we add new results showing that our approach is effective on coding problems from the CodeContests dataset. We are happy to clarify any questions you may have. Please let us know if your concerns are addressed and if so we would be grateful if you would be willing to raise your score. We would be happy to discuss if you have any concerns.

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.

We have added CodeContests, as an additional dataset to show the efficacy of our approach on competitive programming, complementing our results in MATH (Hendryks et al, 2023), which studies high school math problems. In this task, an LLM must write code for a set of contest coding problems from a variety of sources (e.g CodeForces, HackerEarth, and more). To formulate code generation as an MDP, we define a step/action as a line of code, and the state/context as the concatenation of the contest problem (consisting of the problem description, runtime requirements, and sample input/output pairs) and the code so far. Evaluation is done with a python execution server on a set of held out contest problems. We present preliminary results with an Oracle Verifier, comparing our backtracking algorithm (orange) with Best-of-N (Blue, Snell et al. 2024) and Backtracking to a random step (Green).

We find an average improvement of 10.16% from the linear search algorithm of Best-of-N, Snell et al. 2024 (backtracking from the first step) and significantly outperform backtracking to a random step. Note this is preliminary and we hope to train a PRM with Monte-Carlo Rollouts by the end of the rebuttal period or for the final version of the paper.

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.

Note, that the underlying assumption for using verification is that there exists an insignificant gap between the generator and the verifier. This clearly exists in domains such as Math, Programming, Puzzles, and more. Thus, even if noisy, learning a PRM allows for some benefit when inference-time computation is utilized in tandem with it.

The accuracy of the verifier is indeed important for the algorithm's success. However, all search-based approaches that leverage a verifier (e.g Snell et. al) have this inherent requirement. Below, we consider the simpler search algorithm of Best-of-N. Here, we artificially add label noise to the Oracle verifier and see what the performance of the algorithm looks like for a subset of the math dataset. As seen in the following plot, as the verifier accuracy decreases, there is a large drop in the performance of a simpler search algorithm such as Best of N. Thus, given that this is a limitation for prior approaches, we should not be penalized for needing a verifier, with sufficient accuracy.

For this reason, we discuss in detail what it takes to practically learn such a verifier. For example, we outline in Section 3.5 how to practically learn a PRM to attain a higher verification accuracy. Additionally, we propose in-context value verification in Section 3.4 to further reduce the uncertainty of the verifier when performing revisions. Finally, we smooth the verifier during backtracking as seen in section 4.2 to be conservative with the step we backtrack to, ensuring that the step chosen to restart generation is indeed in a correct portion of the previous solution. This reduces the inefficiencies that come from backtracking to an incorrect step, and the smoothness can be relaxed given a stronger verifier.