Alphazero-like Tree-Search can guide large language model decoding and training

We thank you for your valuable feedback and comments. We have modified several parts of our paper according to your suggestions.

Question 1: The performance improvement is unclear and somewhat unconvincing.

Thanks for your comment. We totally agree with your observation. There are three things we want to clarify:

1. We want to clarify that TS-LLM’s DFS/BFS variants are different from that in Tree of Thought (ToT). They are TS-LLM’s variants instead of ToT baselines. The main difference is that we use a learned value function instead of prompting advanced LLMs, which is one of our important contributions. We have shown in Table 4 the result of BFS and MCTS-alpha on Game24 by prompting LLaMA2-7B model as value function and ORM. It can exactly be regarded as an instance of ToT over small LLMs. Our conclusion is that prompting a small LLM can not offer accurate value estimation and the performance of tree-search is even worse than simple CoT baseline.
2. For path@1 results in Table 2, MCTS-alpha and MCTS-rollout do consume more tokens and it is reasonable to question whether it is mainly because of more token consumptions. Our Q4’s experiment can help to clarify such an issue. We additionally conduct experiments on search aggregation shown in Figure 3, where you can scale up your token consumption by running multiple times of search in MCTS/MCTS-alpha/MCTS-rollout/DFS or setting larger beam size in BFS. Specifically, the second row of Figure 3 can clearly demonstrate the performance of each tree-search algorithm when given the same token consumption. Basically, MCTS-alpha and MCTS-rollout are still dominant in most cases.
3. At last, we never claim in our paper that MCTS-alpha/rollout is the new SOTA over all evaluation tasks. Our experiments are not designed to show new SOTAs or the most efficient tree search algorithm. We treat different search algorithms as TS-LLM’s variants and equally present their pros and cons. We believe every algorithm has its limitations and no one can claim superiority in all scenarios. For example, as we illustrate in our analysis, MCTS, BFS, and DFS generally perform well in shallow search problems while MCTS-α and MCTS-Rollout are dominant in deep search problems.

We thank you for your comment and understand that readers may have similar misunderstandings. To make this more clear, we’ve also modified this and incorporated all three points in our rebuttal revision. Specifically, we modify the Benchmark algorithms subsection in Section 4.1 to include more clarifications among benchmark algorithms and our experiment objectives. We’ve also modified the results analysis in Q1, to incorporate the Path@N result shown in Figure 3 for a better presentation.

Question 2: Misapplication of RLHF Task: This paper's approach to the RLHF task seems misaligned with its intended purpose. Instead of tackling the fundamental challenges of RLHF, the study opts for a smaller agent (125M) and compensates with a higher token count to achieve elevated rewards through an open-source Reward Model. While this method may yield higher scores during training and potentially quicker convergence due to reduced sampling, the extended time required for "search" processes negates these benefits. Ultimately, this approach fails to address the crucial issues inherent to RLHF.

Thank you for your comment, we fully agree with you that (1) we do not offer a solution that can solve the RLHF core problem like reward function learning (since we choose an open-source reward model) (2) The tree-search operations bring in additional computational burdens.

Here what we are trying to provide is an alternative algorithm, in contrast to PPO or rejection sampling-based training, when given a fixed reward function. Since we are not a paper studying RLHF, this is just a showcase to present the applicability of the TS-LLM framework over tasks beyond reasoning. This is also a good example to show MCTS-alpha can do deep tree search in LLM generation (since we set a really wide and deep tree in the RLHF experiment.) In addition, there is one additional advantage of the tree-search algorithms. Section 3.2 has presented that we do not need to finetune the LLM generator using Rejection sampling finetuning or PPO, and can let the model generate reliable answers by pure inference.

Question 3: In Figure 1, the presentation should be further clarified. The authors did not evaluate Game24 on a token-level, however the figure may mislead readers.

Thanks for your comment, we have added the clarification in the footnote of Figure 1.

Question 4: Is the term k-node used in the context of BFS equivalent in meaning to k-node in the context of MCTS?

Exactly. The k is mentioned in the last paragraph of Section 3.1, as the number of nodes we subsample during the leaf node expansion. In other words, it is how many child nodes we have given a parent node.

Question 5: In Table 4, why do DFS/BFS/MCTS yield identical results for the RLHF task? This paper should clarify this.

Thanks for your comments and suggestions. We do clarify this in our paper in the last two sentences of the first paragraph in Q1’s answer.

Let us clarify this more. When evaluating path@1 performance, we only generate one path given a search algorithm. It means MCTS cannot take the backward step (because it only conducts backward operation after it reaches the terminal node), BFS beam size is 1, and DFS cannot prune node and only get 1 path. In this case, MCTS, BFS, and DFS degenerate to greedy search over value. Results are different on the first 3 tasks because they are using sentence-level node expansion, so the action nodes are based on sampling. For the RLHF task, we derive the node by conducting a forward process and extract the top 50 tokens, it is deterministic as long as the input is the same.

We fully agree that we need more explanations and have modified the analysis results in the first paragraph of our Q1’s answer. We’ve also clarified the Path@1 setting and additionally added a new result analysis by referring to the Path@N result shown in Figure 3. In this setting, MCTS, BFS and DFS will be different so we can conduct a better analysis. We’ve also added more explanations in the caption of Table 2 to better illustrate the reason why DFT/BFS/MCTS yield identical results.

Question 6: Do BFS and DFS utilize the same Large Language Model (LLM) as MCTS-alpha and MCTS-Rollout, or do they employ GPT-4 as their original Tree of Thought (ToT) configuration?

All search algorithms utilize the same LLM, which is LLaMA2-7B in the first three tasks and GPT-2 in the RLHF alignment task. We want to show that our framework is applicable to all different LLMs of any size and directly applying ToT pipeline to them may not work.

Question 1: The performance improvement is unclear and somewhat unconvincing.

In your response to the question regarding performance improvement, you emphasize that your paper does not claim that MCTS-alpha/rollout establishes a new state-of-the-art (SOTA) across all evaluation tasks, and that your experiments were not designed to demonstrate new SOTAs or the most efficient tree search algorithm. While this clarification is appreciated, it might be beneficial for the paper to more explicitly state the specific contexts or scenarios where MCTS-alpha and MCTS-Rollout demonstrate their strengths, beside from "going deeper".

considering the use of value networks in MCTS is not a novel approach, a deeper exploration of how your implementation improves upon existing models would further strengthen your argument.

Question 2: Misapplication of RLHF Task

If the primary goal is not to address the core problems of RLHF but to present an alternative algorithm, it might be more beneficial for readers if you explicitly compare your method with existing approaches like PPO and rejection sampling-based training. This comparison would help in understanding the unique advantages or limitations of your method in the RLHF domain. Moreover, while using an open-source reward model simplifies the process, it might also be worthwhile to discuss how this choice influences the generalizability and applicability of your findings. If your method shows promise in specific scenarios or configurations, clarifying these conditions would add valuable context for your readers, helping them understand where your approach fits in the broader landscape of RLHF solutions.

Clarity on BFS and MCTS Variants

In your text, you replaced "BFS" with "BFS-V," suggesting a variant. However, this change might not be clear to readers that it is equivalent to Beam Search. For greater precision and clarity, changing the term to "Beam Search" would be more appropriate and straightforward. This adjustment ensures that the concept is immediately recognizable and accurately represented.

Similarly, the use of MCTS needs clarification. If the paper employs a variant of MCTS from Sampled MuZero, mention this earlier. This distinction is crucial to avoid confusion, especially since it differs not only from traditional MCTS but also from the MCTS used in "AlphaZero."

Fairness in Baseline Comparison

Your paper uses CoT-greedy as a baseline, which may not provide a fair comparison with your methods. If your research builds upon ToT and RAP, it's essential to include and compare their results as well. While it's noted that these studies use GPT-4 and your method employs a smaller model, this contrast is precisely what readers need to understand. It's crucial to demonstrate why your method is preferable, outlining any trade-offs or improvements they might expect.

We thank you for your valuable feedback and comments.

Question 1: While this clarification is appreciated, it might be beneficial for the paper to more explicitly state the specific contexts or scenarios where MCTS-alpha and MCTS-Rollout demonstrate their strengths, beside from "going deeper". considering the use of value networks in MCTS is not a novel approach, a deeper exploration of how your implementation improves upon existing models would further strengthen your argument.

We do explicitly state the strengths of MCTS-alpha and MCTS-Rollout in our analysis of Q1 and mention that these two algorithm variants dominate when dealing with deep-search problems while the rest three perform quite well in shallow one.

Could you elaborate more on ' implementation improves upon existing models '? Or do you mean existing algorithms such as TOT and RAP? If so, please check the following answer to Question 5.

Question 2: Misapplication of RLHF Task

We do explicitly compare our methods to existing approaches like PPO and rejection sampling-based training. That is mainly in Q5 and Table 6 when we discuss how such a process can guide training.

Thank you for your suggestion, we will clarify more about the choice of the reward model and further clarify the specific configuration.

Question 3: In your text, you replaced "BFS" with "BFS-V," suggesting a variant. However, this change might not be clear to readers that it is equivalent to Beam Search. For greater precision and clarity, changing the term to "Beam Search" would be more appropriate and straightforward.

Thank you for your feedback and suggestions. For the name of BFS-V, it doesn’t mean it’s just a variant. As we mentioned in Section 3.2.2, page 5, it means “with value function based tree-pruning”, suggesting the pruning is done by a learned value function which is different from Tree of Thought (ToT). We do agree with you that it can be called “Beam Search”, however, we maintained the name “BFS” to keep consistency with Tree of Thought paper. And then we added a suffix “-V” to show the main difference between the BFS(Beam Search) algorithm in TSLLM and Tree of Thought. To make it more clear, we also say in the main paper that “BFS-V can be regarded as a beam-search with cumulative reward as the objective” in Section 3.2.2, page 5.

We thank you for your valuable feedback and comments. We have modified several parts in our paper according to your suggestions.

Question 1: Based on the result in Table 2, it is unclear if the proposed method is better than the existing method such as BFS (ToT). For instance, in GSM8K, the BFS is better than the proposed methods such as MCTS-\alpha and MCTS-Rollout… As a result, I am not convinced that this approach will work better or not.

Thank you for your comment. There are three things we want to clarify:

1. We want to clarify that TS-LLM’s DFS/BFS variants are different from that in Tree of Thought (ToT). They are TS-LLM’s variants instead of ToT baselines. The main difference is that we use a learned value function instead of prompting advanced LLMs, which is one of our important contributions. We have shown in Table 4 the result of BFS and MCTS-alpha on Game24 by prompting LLaMA2-7B model as value function and ORM. It can exactly be regarded as an instance of ToT over small LLMs. Our conclusion is that prompting a small LLM can not offer accurate value estimation and the performance of tree-search is even worse than simple CoT baseline.
2. Table 2 represents path@1 results. A much more comprehensive result can be seen in our Q4’s experiment. We additionally conduct experiments on search aggregation shown in Figure 3, where you can get Path@N results and scale up your token consumption by running multiple times of search in MCTS/MCTS-alpha/MCTS-rollout/DFS or set larger beam size in BFS. Specifically, the bottom line of Figure 3 can clearly demonstrate the performance of each tree-search algorithm when given the same token consumption. Basically, MCTS-alpha and MCTS-rollout are still dominant in most cases.
3. At last, we never claim in our paper that MCTS-alpha/rollout is the new SOTA over all evaluation tasks. Our experiments are not designed to show new SOTAs or the most efficient tree search algorithm. We treat different search algorithms as TS-LLM’s variants and equally present their pros and cons. We believe every algorithm has its limitations and no one can claim superiority in all scenarios.

To make this more clear, we’ve also modified this and incorporated all three points in our rebuttal revision. Specifically, we modify the Benchmark algorithms subsection in Section 4.1 to include more clarifications among benchmark algorithms and our experiment objectives. We’ve also modified the results analysis in Q1, to incorporate the Path@N result shown in Figure 3 for a better presentation.

Question 2: Second, it seems that most of the experiments are about improving the value function and reward function, not about fine-turning the LLM. For example, if we do fine-tuning Table 2, are we able to improve the performance? Essentially, I want to see more results based on question 5 written in the paper: Can TS-LLM further train LLM?

We do present experimental results by finetuning Table 2 and Figure 3. The specific results are presented in Table 5 and Table 6. In Table 5, we present how we can leverage tree-search examples in further finetuning the LLM-based value function. In Table 6, we show more comprehensive results over GSM8K and RLHF alignment and present how tree-search examples can increase LLM generation by finetuning LLM policy $\pi_{\theta_0}$ to $\pi_{\theta_1}$, LLM value $V_{\theta_0}$ to $V_{\theta_1}$ and both. Does this address your concerns?

Question 12: In general, how hard should it be to adapt the proposed methods to another domain?

One of the contribution of TS-LLM is that it is generally applicable. The minimal requirement of our proposed framework is a set of problems/tasks and an oracle reward function(maybe rule-based) on them. Therefore, this requirement can be satisfied for a wide range of NLP tasks and decision-making tasks.

Question 13: What are the computational requirements for only training the value/reward models without fine-tuning?

First, we need to do inference on the training problem set with the LLM policy to sample enough data for value/reward training. The value/reward model are then trained on such data. Since the network architecture of the value/reward model we used was almost the same as the LLM policy, finetuning such a model has a similar computational requirement as finetuning an LLM. Although we did update all weights of the value/reward model which have the same size as the LLM policy which indeed requires advanced GPUs like A800, for users with limited computational resources, using a smaller decoder-only transformer or techniques like LoRA would be an efficient and effective solution.

Question 14: Are there any meaningful combinations that were not tested due to complexity?

There are a lot of potentials to try and we have listed some in Appendix B discussing limitations and future work. For example, more complex tree construction strategies such as mixed sentence/token-level node expansion, adaptive node-max-width w.r.t. depth, and rejecting semantically similar action nodes to increase search efficiency and diversity.

Question 15: Were the samples in the dataset always selected from the train set?

Yes. For the finetuning of both the LLM policy and the value/reward model, we construct the dataset from the samples in the training set.

Question 16: Is Path@N counting the number of paths up to a leave node or the number of separated paths? For instance, for beam-search with size K, the number of leaves is at most K, but there were other paths partially explored that were discarded along the way.

For BFS(beam-search) and DFS, the answer is the number of separated paths. For MCTS variants, this means running the search algorithm N times, they are not ensured to be distinct.

Question 17: How is the same level of token computation maintained? For instance, it could be the maximum average of the tokens used for all the variations of MCTS. Or perhaps, it was run with multiple values of N, and then return bigger N under the token of the other method. But I’d expect that to be done per sample.

We just choose N which is close to the most tokens used in all tree search algorithms. We want to clarify here such a process is reasonable. We’ve counted the number of tokens for the CoT outputs of each task, we found that the number of tokens among every problem in all tasks has small standard deviations compared to the average values (e.g. for GSM8k it’s $101 \pm 36$, for Game24 it’s $79 \pm 3$, for ProntoQA it’s $92 \pm 17$).

Question 7: Why defining $g$ in page 4? Section 3.1 uses \phi_theta

We guess you mean $\pi_\theta$ in Section 3.1. They have different meanings here, $\pi_\theta$ is the policy that takes an input $x$ and outputs actions. While $g$ is defined in the dynamics transition function of the problem, for the tasks we are testing now the transition function is deterministic, i.e. appending output tokens to the prefix/prompt string. The transition function $g$ can be considered being influenced by the LLM policy $pi_\theta$ as $g(s, a) = g(s, a \sim \pi_\theta(\cdot | s))$

Question 8: Is there any restriction about the context size of the LLM and the value/reward models?

The context size of the LLM and value/reward models are restricted by the pretrained length of the LLM and the decoder-only transformer in the value/reward model.

Question 9: Was any fine-tuning or training using multiple-GPUS or the A800 were used to run independent experiments?

As we mentioned in Appendix D.2, the finetuning and training were conducted on 8 NVIDIA A800 GPUs.

Question 10: Is this a single decoder-only model that produces the intermediary values and the final rewards, or are they different models?

Yes, we use one decoder-only transformer model with a value head to output scalars at each token.

Question 11: How were the hyper-parameters setup. How much budget was used to set these numbers? How sensitive are the results to the particular parameters?

During our experiments, we found that the performance and token consumption are sensitive to several hyperparameters: first is the general hyperparameters of building the search trees, e.g. tree-max-depth and tree-max-width, for the former, we choose the max-depth according to the statistics of the distribution of the number of steps from the dataset sampled by the LLM policy on the training set; while for the latter we’ve shown in Table 3 that there is a trade-off between token consumption and the size of the search space. MCTS variants (MCTS-alpha, MCTS-rollout and MCTS) are also sensitive to hyperparameters, we keep most hyperparameters as the default values used in AlphaZero and MuZero except for $c_{init}$ which we select a value from two possible ones {0.3, 3.0} when computing $c_{puct}$ to control exploration and exploitation in MCTS variants. For the Dirichlet noise, this might be important since there is some guess that DeepMind sets this hyperparameter according to the statistics of the dataset collected for each task. However, due to the limit of computation resources, we just adopted the default value in AlphaGo as 0.3, which is specified for chess (refer to Page 14 in the AlphaZero paper (https://arxiv.org/pdf/1712.01815.pdf) and Appendix B in MuZero(https://www.nature.com/articles/s41586-020-03051-4)).

We really thank you for your valuable feedback and helpful suggestions. Here’s our response to your feedbacks and questions.

Question 0: Writing issues and more clarifications for TS-LLM.

According to your suggestion, we’ve made the following modifications and clarifications to our paper to address the writing issue you mentioned:

1. We’ve corrected the wrong version of the citation of CoRe and cited the self-consistency paper in Section 4.1.
2. We’ve modified the description of MCTS variants in the paper in Section 3.2.2 and added a comparison of the tree-search algorithms in TS-LLM in Appendix C.2.
3. We’ve changed the “juxtapose” on page 6 Section 3.4 into a more precise expression.
4. We’ve defined the state $s$ in Section 3.2.1 outside the transition function $g$.
5. We’ve defined SFT as “supervised finetuning” in Section 4.1.
6. We’ve modified the analysis to make it more clear and precise in Q3’s analysis in Section 4.2. We’ve added the necessary clarification about whether intra- or inter-tree search is used in Table 2 and Table 3 and clarify the detailed setups in Appendix D.5.

Question 1: The notion of aggregation is misleading (Sect 3.2.3). Aggregation refers to using a set of things to get something new, but in this case, we are just selecting an output, not combining them.

We totally understand your concern that the notion of ‘aggregation’ might be misleading. While we adopted it from the previous papers CoT-SC(https://arxiv.org/abs/2203.11171) and RAP(https://arxiv.org/abs/2305.14992) for consistent terminology.

Question 2: No details on the value/reward model besides the mention that it is a decoder-only model.

As we mentioned in section 3.2.1, the value/reward model architecture is simple, we append a value head (which is an MLP that outputs a scalar, e.g. a linear layer) to a decoder-only transformer. For GSM8k, Game24, and ProntoQA, the decoder-only transformer is the pretrained LLaMA2-7b without the lm_head module. For the alignment task, the decoder-only transformer is the pretrained GPT2-small without the same output lm_head. And we used the pretrained weights as the initialization.

Question 3: No specific statistics about the behaviour of the algorithms beyond the aggregated number of tokens. For instance, knowing the number of roll-outs would be useful.

Thank you for your comment. We want to confirm the meaning of ‘roll-outs’ here. If it means the number of sequences a certain search algorithm outputs, we’ve already shown it in the first row of Figure 3 and Appendix D.6. We hope for your further explanation about the definition of ‘roll-outs’ and we will try to show you the results.

Question 4: No discussion of the wall-time for decoding

Thank you for your comment. We’ve added different algorithms’ decoding wall-time of GSM8k, Game24 and ProntoQA in Table 13~15 in Appendix D.9. The results demonstrated the gap of computational efficiency between CoT/CoT-SC decoding and tree-search algorithms in TS-LLM, due to the complicated search procedures and extra computation introduced by calling value functions in the intermediate states.

We totally agree that computational efficiency is an important issue and address it in Appendix B. But note that our current implementation is just an algorithm prototype without any engineering acceleration. We are continuously working on accelerating our tree search framework codebase including caching past key/values on the tree to eliminate repeated computations and offloading tensors on RAM, etc.

Question 5: Given that MCTS combines both the policy and the value, perhaps a comparison with BFS or DFS is not the more meaningful. Algorithms based on A* rely on the accumulated cost g —in this case inverse to the likelihood of the prefix— and h, an heuristic estimating the value of the rest of path. While the value function in the paper does not take into account the length, it’s still possible to consider their combination.

Yes, we agree with your ideas. This can actually be viewed as a dense reward setting in which, at each step, there is a nontrivial reward. It is natural and heartening to extend our method to this setting like extending from AlphaZero(https://arxiv.org/pdf/1712.01815.pdf) to MuZero(https://www.nature.com/articles/s41586-020-03051-4).

Question 6: Is BFS taking into account the likelihood of the previous path? The text only says to keep the ones with maximal value. If it completely ignores the likelihood, then why does the decoding do something useful?

No, BFS can be viewed as beam search with our learned value and ORM. However, the likelihood is not totally ignored. It will implicitly influence the search process since the actions proposed by LLMs are sampled w.r.t. the likelihood and then treated as i.i.d. sentences/actions in BFS. The decoding is useful since BFS will leverage the value function to choose the top K sentences/actions at each step.

Thanks for your feedback and helpful suggestions. Here’s our response to your questions.

Question 1: What is the k used for BFS? Is it the Tree Max Width in table 1? Did you test different amount of pruning for BFS? How sensitive are the results to the selection of k for BFS and related parameters in the other algorithms?

Let us explain this more clearly. Tree-max-width and k are two different hyperparameters. Tree Max Width determines the number of child nodes every time the tree expands while k represents the number of child nodes we maintain (based on top k value) from the whole candidates. For example, assuming we maintain k nodes at depth d, when expanding nodes for depth d+1, we will expand these k nodes to have (k * tree-max-width) child nodes, from which we select the top k nodes. After this, we have k nodes at depth d+1 and the search continues. The k can also be treated as beam size, which is a hyperparameter to control how many distinct paths are to be sampled on the tree in the BFS(Beam-Search) algorithm.

And Yes, we do test it in Q1/Q4 and Figure 3. The beam size k controls the max number of returned sequences and it corresponds to $N$ in Path@N experiments, as is shown in the first row of Figure 3. To make the setting of Path@N results clearer, we’ve also added explanations about the detailed settings used in the aggregation experiments in Appendix D.9.

About the sensitivity of k, our experiment in Figure 3 has shown that, in GSM8k/ProntoQA/Game24, the performance of BFS is proportional to the size of k (It is reasonable because we have an ORM to help rerank all final k solutions to determine a final solution, and larger k represent larger candidate space). RLHF experiment is a special case where we find BFS does not work in deep search problems. So if you ignore the extra computational burden brought by increasing k/N, in most cases, you can obtain better performance by using larger K. In this case, you won’t need to determine the specific k. However, with k growing, you also have to consider the extra computational burden because you need to expand/evaluate more nodes in BFS. So this is also a trade off between larger search candidates (larger k) and heavier computation demands.

Question 2: The sensitivity of hyperparameter such as how the max width was set.

We fully understand your concern about how sensitive are the results to the hyperparameters. This is exactly what we want to present in our paper. We want to present a comprehensive and systematic analysis about the key components of tree-search guided inference and training in LLMs including alternatives of tree-search algorithms, training of value function, aggregation (or maybe called reranking) of multiple searches, potentials to continuously improve the total framework, and the possible trade-off between different tree construction(expansion) settings (i.e. choices of sentence-/token-level action nodes, the tree-max-width and tree-max-depth).

We agree with you that tree construction hyperparameters might greatly influence the performance and computation cost of tree-search algorithms. Therefore, we conducted ablations on Game24, which is, among the four tasks tested in the experiments, a typical task to show how the selection of tree-max-width could affect the search space upper bound and the final performance. That’s because if one intermediate reasoning step (computation of two remaining numbers) is wrong (impossible to reach the final result of 24), the total subtree from this node is meaningless since all of them should be wrong answers. This phenomenon was also pointed out in the Tree of Thought paper that most failures often happened in the very early steps of Game24.

Let us illustrate more on how we choose the hyperparameter of tree-max-width. For RLHF, we choose it as the default hyperparameter in LLM decoding( top_k=50). For GSM8K/Game 24/ProntoQA, we first experiment with a small tree-max-width and increase it to see the performance gain by increasing the search space upper bound. This is similar with k mentioned in Q1. The larger the search space, the better the performance. Finally, we determine the tree-max-width by fairly considering the tradeoff between the search space limit and computational burden (This is why in Table 3, width=50 is better than width=20 but we still choose width=20 in Table 1). The procedure we follow is a common way for most search algorithms when given a new environment and it is not hard to tune at all. You can always customize your tree search space based on your requirements for the performance and computation budget.

Good point. If somehow BFS was considering all the tokens, then it would not be using the LLM likelihoods. That BFS would be out of the scope of the paper. However, BFS must consider all the tokens.

So, the algorithm called BFS in the submission is not BFS because it’s using the LLM likelihood to prune the tree. So it’s actually a two steps process: that might be mimicking A*: the token continuation discarded can be seen has having infinite cost g. The ones preserved have a finite cost depending on the depth. The models proposed in the submission are used to decide among the new nodes, such that older nodes are always expanded first. This can be exactly A* if we consider the cost g increasing by a number bigger that the maximum value that the models can provide. I think it’s, so just increasing g by 2 with each token or sentence (don’t remember now) should be enough to see the BFS of the paper as an instance of A*.

So, I agree now that BFS is in a the scope of the paper. (Here is for me a demonstration of the value of an interactive discussion with the authors).

However, the name BFS not right, even if the previous work using BFS for LLM use it. As for the aggregation, the fact a few papers of a few years ago used a terminology doesn’t mean that new papers should continue to pollute with confusion the archive of accepted papers. It’s totally ok to use a different term, as far as the manuscript mentions that previous recent work use another term.

I suggest not using aggregation and I suggest that just calling this algorithm BFS leads to confusion. Perhaps it’s a matter of emphasizing that this is doing BFS in the space after pruning with the LLM.

That view suggests that the results might be very different depending on the amount of pruning. No pruning doesn’t use the LLM. Prune greedily, keeping one solution is greedy search. In the middle, more or less pruning might lead to very different behaviour. Less pruning is like higher recall while BFS chooses among them.

New question: what’s the fundamental difference between beam search and BFS? Perhaps beam search is just fixing how many solutions are preserved while BFS might take a different number of nodes at different points of the search. Perhaps flexible beam search would be a more precise name.

new question: did you test different amount of pruning for BFS? Without such variation, this quote in the last comment doesn’t hold:

DFS/BFS can help us verify the quality of value function.

Because there selection of the amount of pruning might depend on the domain, and finding a the magic amount of pruning be non trivial, especially if going off distribution during testing.

Let’s call Oracle-BFS the one using the optimal pruning by LLM for producing the result. The performance of Oracle-BFS could dominate in some domains variations ot MCTS that doesn’t have such oracle and doesn’t commit to a specific amount of pruning. On the contrary, BFS as in the paper might not say much because perhaps the pruning is close to optimal or suboptimal. The demonstration is interesting but I don’t think it’s more than a demonstration.

I remain open to the possibility that further explanation on the pruning might convince me otherwise.

Thanks

Thank you for your comments.

On BFS: since this is not using the LLM likelihood at all, the algorithm seems to be out of the scope of the paper. Right? If you want to keep it, this needs further justification.

Instead of BFS, the natural algorithm to compare is A*. For DFS, the alternative would be IDA*, which might be tricky to implement. I think that A* is the baseline for comparison, no BFS.

A* uses the accumulated cost (cost is inverse to the accumulated likelihood) and an estimation of the future.

Regarding the roll-out, we generally want to understand the search effort besides the implementation detail. A typical metric in search is the number of expansions: how many LLMs calls. Another one is the number of calls to the other models. As the different algorithms use some models and not others, this would reinforce the explanation of the algorithm. The roll-outs may only be for the cases where the leaves are reached.

These metrics are crucial to understand the trade-off of the algorithms.

Really thank you for your feedback!

For the BFS/DFS algorithm, we do not think it is out of the scope of our paper. Firstly, we want to clarify that, in our paper, DFS/BFS are TS-LLM’s variants instead of ToT baselines. We never claim in our paper that MCTS-alpha/rollout is the new SOTA over all evaluation tasks. We believe every algorithm has its limitations and no one can claim superiority in all scenarios. For example, as we illustrate in our analysis, BFS/DFS generally perform well in shallow search problems. Secondly, for MCTS/MCTS-alpha and MCTS-rollout, they need the LLM likelihood and LLM value function, while for DFS/BFS, they still need the LLM value function to help prune the tree (and implicitly rely on the log probability to expand nodes by sampling). Learning a value function is one of our important contributions and DFS/BFS can help us verify the quality of value function. Could you explain more on why it is necessary to explicitly include LLM likelihood in the search process?

For the A* algorithm, we fully agree this is another search algorithm we should include in TS-LLM's variants. We are working on adding this algorithm, but it takes some time to rerun all our tasks and tune the search hyperparameters.

For the rollout, the token number we monitor in our experiments (token number in Table 2 and the x-axis in the second row of Figure 3), can exactly represent the node expansion times. This is another contribution we have - we want to conduct a fair comparison (Section 3.4) which prior works neglect. We sum up all the new tokens generated by LLM every time the node expands. In addition, we think the statistic of the number of generated tokens is better than mere node expansion times. Note that it also takes the expense of each expansion into consideration since it is possible that different node expansions may generate different lengths of sentences (number of tokens), resulting in varying computation expenses.

About another metric like the number of calls to other models, as we described in TS-LLM, we do have calls of the LLM value/reward model. But again, the computation expense of calling value/reward model after each expansion can be estimated accurately by the number of newly generated tokens.

Question 3: For instance, Table 3 suggests that if pruning of aggressive (only 6), then perhaps BFS would be simpler, cheaper and better.

Firstly, when we set the tree-max-width as 6, we are not conducting pruning, we are limiting the search space upper bound of all algorithms by limiting the number of child nodes an intermediate node can expand. And it is totally fine that BFS in some cases is better than MCTS-alpha/MCTS-rollout. We have addressed this problem in our global response and our previous response to you. BFS is one of TS-LLM’s variants instead of merely a baseline. Specifically, we said:

We never claim in our paper that MCTS-alpha/rollout is the new SOTA over all evaluation tasks. Our experiments are not designed to show new SOTAs or the most efficient tree search algorithm. We treat different search algorithms as TS-LLM’s variants and equally present their pros and cons. We believe every algorithm has its limitations and no one can claim superiority in all scenarios.

And the final experimental results validate our belief. BFS/DFS/MCTS algorithms perform well in shallow search problem but struggle in deep search problem (RLHF).

Question 4: What about calling it “reward beam search” instead of BFS?

Thank you for your comment. We think reward beam search is also an alternative. But to align with previous papers such as Tree of Thought, we also think it’s reasonable to mention the name BFS. So to consider both, what do you think about terms like BFS-V/DFS-V (abbreviation for BFS/DFS with value pruning)?

Question 5: The term “node construction” is also confusing.

We understand your concerns. We have replaced the term with “node expansion” in Q2.

About the Value function:

Question 8: Clarification that the proposed models are implemented as an additional head on top of the same representation. When processing a sample, are the models added to the same architecture, so the new model produces both the new tokens and then estimations? Or are those independent calls with shared weights?

Thank you for your suggestions. We choose to use the term “LLM-based” value in our paper to help clarify that we use an LLM based value function. We also have made more clarifications in Section 3.2.1 and the model and training detail subsection in Section 4.1. Specifically, we describe that “typically, LLM value’s decoder is adapted from original LLM policy $\pi_\theta$’s decoder, or alternatively, the LLM value $v_\theta$ and policy $\pi_\theta$ can have a shared decoder”.

Let me clarify our current implementation. In our current implementation, we use a separate LLM policy and LLM value function and their decoders are not shared. This is because our critic learning needs to have training samples from the SFT policy (So our critic can be seen as an on-policy value function for the SFT sampling policy). Our experimental procedure is: We first finetune the base LLM using CoT examples in the training set, to form SFT policy. Then we sample from the SFT policy to construct the critic learning training examples. And finally we get our LLM-based value function by finetuning the base LLM (with additional value heads) using these critic training examples. Though the decoders are not shared among policy and value, they are adapted from the base model. So they may still be quite similar.

And we agree with you that the shared structure might be better for computational efficiency since you do not need to call the model twice and maybe a shared decoder can get rid of all burdens brought by value function evaluation (you can reuse most computation from the policy generation). We also mention it in Appendix D.11 when discussing engineering potentials. There are several ways to achieve it. We can:

1. Instead of using SFT-tuned model, we can use few-shot prompted based model to generate the samples for value function training. And then we train the SFT tuned policy and value function with shared decoders.
2. We have some offline data that help me train the value function in advance.
3. We are doing the training in an online manner so the critic function can be iteratively trained, similar to the setting in AlphaZero.

Currently, we are working on adding a new experiment about the shared decoder by using the same data we train the separate value and policy LLM. We want to show how such model sharing might influence the performance and also efficiency.

Question 9: There is no need to test in other LLMs. My point was that the paper should be written as it happened: it was always the same representation, and that’s ok. Cite state of shared weights in AlphaZero and other work.

Thank you for your suggestions. Please refer to our response in Q8 and we have cited the state of shared weights in AlphaZero (Figure 1.b in AlphaZero’s paper).

Other topics:

Question 10: Make sure that for all the hyper-parameters, the protocol for selection is discussed, as well as the sensibility to the parameter.

We add a new subsection in Appendix D.10, to fully discuss our protocol for choosing all tree-search-related hyperparameters in our paper.

Question 11: Discussion about engineering challenge.

Thank you for your suggestion. We added 5 points in Appendix D.11 (previous D.9) about the engineering challenges and how we can further increase the efficiency with more engineering efforts.

We thank you for your valuable feedback and comments.

Question 1: This paper needs to have more experiments to show the improvement. For example, more models and more tasks. For more tasks, maybe you can find some here “https://github.com/Ber666/llm-reasoners”.

Thank you for your suggestions. In fact, we do refer to “https://github.com/Ber666/llm-reasoners” for some of our experiments including GMS8K, ProntoQA and Game24. We also add a brand new RLHF task to demonstrate the general applicability of our TS-LLM framework. You can have a more detailed look at their repo and it turns out that most environments they provide do not contain any benchmark results (the table in their Experiment Results section in README).

Question 2: The paper should discuss which kinds of tasks are better for using MCTS rollout and which are better for MCTS alpha.

Thank you for your comment. The only difference between MCTs-Rollout and MCT-alpha is that MCTS-rollout conducts the search from the beginning while MCTS-alpha conducts the search from the last search node. We think MCTS-rollout is just an offline version of MCTS-alpha since it reconducts the full search from the beginning every time. Our conclusion is, basically MCTS-rollout and MCTS-alpha will be good at the same kind of task, while MCTS-rollout can scale up token consumption to search for better generation. According to your suggestion, we have added a discussion to the MCTS-Rollout description in section 3.2.2.

Question 3: Have you tried using both rollout and value together? If the speed of rollout and the speed of the value network are very different, you can reference the paper “Multiple Policy Value Monte Carlo Tree Search.”

Thank you for your comment and we will refer to this paper. Currently, we use two separate networks to conduct rollout and value generation respectively and the speed of rollout and the speed of value estimation are similar. But note that we can also utilize a shared LLM decoder for both policy and value network so the rollout and value estimation can be conducted simultaneously and be more efficient.