ToolChain*: Efficient Action Space Navigation in Large Language Models with A* Search

We sincerely appreciate Reviewer 7Kjr for considering our work as novel and significant with empirical and qualitative improvements. Your suggestions and insights are very helpful for further enhancing the submission quality. We will include a comparison of concurrent work [1] and include the following changes in our updated draft. Please find the responses below:

[Question 1] What do you mean by "task-specific" heuristic functions? Are the tasks here based on "g"? If so, do you use different memory buffers for each task, and how many trajectories do you need in the memory buffer before it produces a valid score?

Similar to the definition in [2], we use “task-specific” to describe the heuristic functions for two main reasons: (1) different seed data are utilized for the long-term memory in each task (i.e., dataset); (2) the computed heuristic values are only meaningful within the context of their respective tasks, as similar questions often share underlying solution logic. All tasks are defined based on the natural language descriptions $g$, with further details in Section 2. The initial number of trajectories in the memory buffers for each dataset is as follows:

All seed data for these memory buffers are derived from ToolBench [3], which is a reasonable number of trajectories and comparable to few-shot demonstrations.

[Question 2] It seems to me that the future cost heuristic might mislead the planner since it is likely that the same action may appear at different positions (in the sequence) for different tasks. For instance: "Clean a kitchen table" and "Set a kitchen table" have their action sequences reversed.

We acknowledge that the future cost heuristic function may not always be accurate. With experimental results presented in Table 2, we have proposed two solutions to mitigate this issue: (1) we integrated the LLM imagination score with heuristic functions to mutually regularize their outputs; (2) the calculation of the heuristic future function involved averaging the relative positions of each action across different plans. Thus, if there exist two scenarios of “Clean a kitchen table → Set a kitchen table” and “Set a kitchen table → clean a kitchen table”, our method effectively incorporates information from both sequences into the heuristic future score.

[Reference]

[1] Hazra et al. “SayCanPay: Heuristic Planning with Large Language Models using Learnable Domain Knowledge.” arXiv 2023.

[2] Hao et al. “Reasoning with language model is planning with world model.” EMNLP 2023.

[3] Xu et al. “On the Tool Manipulation Capability of Open-source Large Language Models.” arXiv 2023.

We appreciate Reviewer tHXf for considering our work as a firm contribution with extensive experiments. Your suggestions and comments are very helpful for us in improving presentation quality and clarifications with wider impacts. Please find the corresponding responses below:

[Weakness 1 & Question 1] Questions about the design process of two important evaluators, g(n) and h(n) in the A$^\*$ search. Further studies exploring different forms of cost functions could provide additional insights beyond existing ablation studies.

We appreciate your insightful question regarding the design of our cost functions $g(n)$ and $h(n)$ in ToolChain*. Similar to most of the recent advances [1][2][3][4][5] in LLM agents and LLM evaluation, ToolChain* is also based on empirical findings. While lacking systematic theoretical proof, ToolChain* in designing $g(n)$ and $h(n)$ is rooted in empirical findings rather than relying solely on handcrafted functions. To validate the advantage of our proposed value function, we explored various cost function formulations, including self-evaluation prompting LLMs to generate verbalized scores [4][6][7], and vote scores in 'tree-of-thoughts' [3]. Our findings indicated that these scoring methods were not consistently accurate. We identified two key factors guiding our heuristic function design: (1) the black-box nature of state-of-the-art LLMs offering limited information for score formulation, and (2) the observation that similar tasks often share similar logical structures in their solution plans. Driven by the two motivations, we developed heuristic scores to maximize the use of additional information derived from past experiences.

The results demonstrate that a combination of LLM-generated scores and heuristic scores yields superior performance. In the initial stages, due to limited data in the long-term memory, the model predominantly relies on the LLM-generated scores for generating plans. As the long-term memory expands, the heuristic functions increase in accuracy, effectively regularizing and refining the LLM-generated scores.

Furthermore, we have conducted explorations to utilize additional information from open-sourced LLMs in order to formulate a more precise cost function. These explorations include: (1) action likelihood, calculated as the product of token log probabilities within an action; (2) answer probability, determined by the token probability of 'yes' in response to the question 'Is this action step correct?'; and (3) a reward-style probability, which applies weight decay based on the length of tokens in the actions and the total number of actions in a plan. Given the constraints of computational resources and limited time available during the rebuttal session, we conducted additional experimental results using LLaMA-2 (7B) for additional insights:

The results indicate that our proposed cost functions outperform other value functions from open-source LLMs. We greatly appreciate your constructive feedback and will include further exploration of incorporating open-sourced LLM information into cost function computation in our future work.

[Reference]

[1] Yao et al. “ReAct: Synergizing Reasoning and Acting in Language Models.” ICLR 2023.

[2] Sun et al. “AdaPlanner: Adaptive Planning from Feedback with Language Models.” NeurIPS 2023.

[3] Yao et al. “Tree of thoughts: Deliberate problem solving with large language models.” NeurIPS 2023.

[4] Hao et al. “Reasoning with language model is planning with world model.” EMNLP 2023.

[5] Lu et al. “Chameleon: Plug-and-play compositional reasoning with large language models.” NeurIPS 2023.

[6] Xie et al. “Self-Evaluation Guided Beam Search for Reasoning.” NeurIPS 2023.

[7] Xiong et al. “Can LLMs Express Their Uncertainty? An Empirical Evaluation of Confidence Elicitation in LLMs.” arXiv 2023.

Thank you so much for your thoughtful and constructive feedback, as well as your plan to raise the score. We greatly appreciate your confirmation on the findings and contributions of our work, and the understanding that our results are interesting enough.

[Question 2] Is the LCS calculation in $g_{t,1}$ over the sequence of API actions ($a_i\in\mathcal{A}$) alone or does it include additional information about the task or problem in the text string? Do the API actions include the parameters or just the method names?

We apologize for the potential confusion. We only leverage the sequence of API actions in our calculations of the LCS calculation $g_{t,1}$. The API functions used in the LCS calculation include both the API function name and the parameters. We will add further clarifications in the updated manuscript.

[Question 3] How much does overall performance vary with the prompt and the other hyper-parameters of the algorithm (dataset, heuristic score implementation, etc.)?

For additional experimental results and clarifications, we kindly invite the reviewer to refer to our response in Weakness 2.2 for more details.

[Question 4] How much performance comes from the use of the LLM in different parts of the algorithm vs classical (non-LLM) best-first search? Specifically, how is performance impacted by non-LLM expansion with LLMs only used for node evaluation? Is it possible to construct a search-based baseline which doesn't use any LLMs but otherwise incorporates the same intuition (long-term memory, heuristic scoring, etc.)? If yes, how might it perform?

LLM agents have demonstrated strong planning and reasoning capabilities in complicated problem-solving [1][2][3][4][6]. Following the most recent research in this line, our primary scope is to develop a planning method for LLM agents to enhance solution quality and search efficiency. Considering the very recent advances [1][2][3][4][6], approaches involving non-LLM best-first search may not lead to an optimized design, and thus not considered in our current research scope: To incorporate non-LLM best-first search (non-LLM agent) into ToolChain*, we need to consider all potential actions as candidates at each step. Given the extensive action space formed by API functions and their parameters, this approach would necessitate evaluating thousands of nodes at every step, leading to impracticality and inefficiency. In contrast, LLM agents utilize their inherent reasoning and planning capabilities to narrow down to the most promising potential next steps, allowing us to apply our proposed cost function to a more manageable set of options. This significantly reduces the need to exhaustively traverse the entire action space at each step. To assess the feasibility of a non-LLM approach, we experimented with randomly expanding 50 potential actions at each step (where >50 actions are less impractical due to the aforementioned efficiency issue), while keeping the cost function unchanged from the LLM agent settings:

These results offer insights into the comparative performance of LLM-based and non-LLM agents in the context of ToolChain*, highlighting the significant contributions of LLM agents to the efficiency and effectiveness of our method. We observe significantly lower performance (i.e., 0) by the non-LLM agent, because the non-LLM agent struggles to perform effectively when choosing the best one from the numerous available actions, with the evaluations being enabled by our LLM-based score functions. Interestingly, the non-LLM best-first search method demonstrates marginally superior performance on Google Sheets across other datasets, primarily due to the inclusion of single-step solutions in its test set. Other traditional non-LLM A* approaches, as referenced in [7], typically involve training for one-step expansion using additional annotated data, which is not always available in the evaluation sets considering the huge amount of actions.

[Weakness 2.4] The paper mentions the importance of the cost function ("this efficiency gain may stem from the proposed superior cost function") but does not investigate deeper. As a result, it becomes difficult to clearly identify the overall impact and broader utility of ToolChain$^\*$, especially compared to similar search-based ideas proposed in Tree-of-Thought.

To elucidate its importance, we refer to our experimental results in Table 2 and additional experiments addressing Weaknesses 2.1. These experiments demonstrate that alternative cost function designs (by comparing row 2 with rows 3,4,5 in additional experiments in Weaknesses 2.1) or the removal of any component from the cost function (by comparing all results in Table 2) results in a notable average performance decrease of 7.9% and 5.0%, respectively. Moreover, these adjustments to the cost functions in ToolChain* lead to an approximate increase of 1.52x in computation time on average. These empirical findings highlight the significance of the cost function in ToolChain* for plan quality and efficiency.

In addition to the empirical results above, the overall impact and broader utility of our paper can be summarized as follows: (1) We propose ToolChain*, a novel A*-like tree search algorithm, to develop autonomous LLM-based agents for complex planning and reasoning tasks; (2) ToolChain* formulates the action space as a decision tree, effectively mitigating error propagation and expanding search space; (3) We leverage a combination of LLM-based functions and heuristic functions to efficiently navigate the agents in the action space. (4) Our extensive experiments demonstrate the effectiveness and efficiency of ToolChain* in tool usage across five distinct datasets including both planning and reasoning tasks (e.g., home searching, trip booking, mathematics, etc.).

Using exactly the same score function, we compare ToolChain* with the tree-of-thoughts (ToT) approach [3], which relies on depth-first and breadth-first search methods. We summarize the advantages of ToolChain* as follows: (1) ToolChain* is based on A* search, utilizing future cost functions for goal-directed node prioritization. This estimation results in enhanced performance and efficiency compared to the ToT, as demonstrated in Table 1 and Figure 5; (2) ToolChain* significantly reduces the number of nodes that need exploration. While ToT-BFS explores all nodes at a given depth before moving to the next level and ToT-DFS potentially follows irrelevant paths, ToolChain* focuses only on the most promising paths through A* search; (3) ToolChain* is designed for more complex, real-world applications, demonstrating its broader utility in LLM agent tool use and reasoning scenarios.

[Question 1] Please explain the details of the Imagination Score in more detail. For example, what does "imagine more concrete future steps" mean? Where is the "imagined plan" generated in Algorithm 1? (I think in the cost calculation f(n) but not sure). How many extra calls to the LLM?

The introduction of the imagination score is available in Section 3.3. We would like to provide further clarifications on the future steps and imagined plans with a detailed example. Consider a current plan $p=(a_0,a_1,\cdots,a_t)$. The future score $h(\cdot)$ aims to estimate the potential future cost required to complete the plan and reach the target. The imagination score is a component $h_{t,2}(\cdot)$ of the future score $h(\cdot)$. To calculate $h_{t,2}(\cdot)$, we initially prompt LLMs to imagine subsequent steps to the current plan, resulting in an imagined plan $p_{im}=(a_{t+1},a_{t+2},\cdots,a_T)$. We then calculate the proportion of current steps present in the imagined plan as the imagination score: $h_{t,2}(n)=t/T$. The imagined plan is generated in the cost function calculation $f(n)$ in Algorithm 1.

For example, assume that we have a current plan “Agent.WalkTo(dining_room) $\to$ Agent.WalkTo(food) $\to$ Agent.Find(food) $\to$ Agent.Grab(food)”. We first prompt LLMs to use their imagination to complete the plan, resulting in “Agent.WalkTo(dining_room) $\to$ Agent.WalkTo(food) $\to$ Agent.Find(food) $\to$ Agent.Grab(food) $\to$ Agent.Find(table) $\to$ Agent.Put(food, table)”. We then use the proportion of current steps in the imagined plan as the imagination score, where $h_{t,2}(n)=4/6=0.67$, leading to an imagination-based future cost of $1-h_{t,2}(n)=0.33$.

Concerning the additional API calls, each action requires two separate calls to calculate both the self-consistency sampling score and the LLM imagination score. Consequently, by leveraging the total API calls noted in our response to Question 5, we estimate the number of additional API calls as approximately 53.7 per question, which is similar to the number of extra API calls in other tree search-based methods [3][4].

We appreciate Reviewer 52p8 for considering our work as interesting and important with a potentially broad impact. Your suggestions for clarification and improvement are constructive. Please find the responses below:

[Weakness 1] Some of the implementation details of the four components of the heuristic function are not clearly explained. These are central to evaluating the algorithmic contributions of the paper and reproducibility so a clearer description would be very useful. I found it particularly difficult to understand the implementation details of the "imagination score". See the questions for details.

Details of the imagination score are explained in Section 3.3 of our original paper. For further clarifications and an illustrative example, we kindly invite you to refer to our response to Question 1.

[Weaknesses 2.1] Although the experiments have a good amount of detail including ablation experiments, the paper does not clearly identify the main sources of performance improvement. For example, how much does a better scoring function matter compared with a better LLM or search procedure?

We would like to emphasize that the improvement of our approach over the state-of-the-art approaches (i.e., baselines) [1][2][3][4] is not from the capability of a better LLM. This has been validated by the comparison between our approach and baselines using exactly the same LLM in Table 1. Instead, the improvement comes from (a) the scoring function design and (b) the search procedure. Although using a more advanced LLM may lead to improved performance, in our comparisons, we always keep using the same LLMs for our approach and baselines in Table 1.

To further understand the improvement by (a) the scoring function design and (b) the search procedure, as well as how significant the improvement is compared to the improvement by a better LLM, we have conducted a more thorough ablation study on ToolChain*. This extensive ablation study includes different base LLM models (GPT-3.5, GPT-4), scoring functions (ToT-Value, ToT-Vote, LLM scores, our proposed cost function $f(n)$), and search procedures (ToolChain*, ToT):

In this table above, $f(n)$ indicates our cost function design in Section 3.1, while ToT-Value and ToT-Vote are two value functions utilized in ToT [3].

Our main contributions lie in the development of a more effective and efficient planning algorithm for LLM agents to address complicated real-world problems. We summarized the following observations from the table above: (1) It is important to note that the design of the scoring function is an integral component of a search procedure. While it is challenging to isolate and compare their individual contributions, our experiments indicate that the scoring function design has slightly more impact than the search algorithm itself. In addition, when compared with baselines under the same base LLMs (GPT-3.5 or GPT-4), both our proposed searching procedures and scoring functions have contributed to model performance (see details in Table 1). (2) ToolChain* shows better performance with GPT-4 than GPT-3.5, suggesting that employing a more advanced base LLM could lead to performance improvement, while this improvement is considered moderate compared to the improvement by our searching procedures and scoring functions.

Dear Reviewer 52p8,

We sincerely appreciate your valuable and constructive feedback. We would like to kindly remind you that the author/reviewer discussion phase concludes on November 23rd. We hope our responses and additional experiments have addressed your concerns and improved the paper’s quality. If you have any further suggestions or comments, please feel free to share them. We are looking forward to a constructive discussion during the rebuttal phase.

Best Regards,
The Authors

We sincerely appreciate Reviewer gdT8 for considering our work as novel and promising with solid empirical results. Your suggestions are decent insights for us to revise the submission draft. Please find the responses below:

[Weakness 1] Cost functions rely on heuristic components like long-term memory, self-consistency sampling, and LLM imagination which may not always be accurate.

We recognize that individual heuristic components of our cost functions, such as (1) long-term memory, (2) self-consistency sampling, and (3) LLM imagination, might not always yield perfectly accurate results. Similar to recent works [1][2][3] in the field of LLM agents, our work leverages heuristic functions to evaluate actions during the search for solutions. In this paper, we innovatively identify three intrinsically complementary components and further discover that a combination of these components enhances model performance compared to using any single component in isolation (see Table 2 for empirical details). This improvement stems from each component providing a unique perspective in evaluating the nodes within the search tree. For example, in scenarios with limited plans in long-term memory, the other two components (self-consistency sampling and LLM imagination) become more influential in shaping the cumulative function $f(n)$ and the future score $g(n)$. Conversely, as more plans accumulate in the long-term memory, the heuristic scores gain reliability and effectively regularize the scores generated by LLMs.

[Weakness 2] Still trails open-loop systems in efficiency, albeit outperforming other tree search methods. There is a tradeoff between search depth and solution quality.

We agree with the reviewer's point about the efficiency trade-offs of our method in comparison to open-loop systems. As acknowledged in Appendix A, our method, while effective, does not yet match the efficiency of open-loop systems. Our ToolChain* aims to strike an optimal balance between efficiency and solution quality. While open-loop systems are more efficient by directly generating the entire plan, they often compromise on plan correctness due to limited exploration in the potential action space. To fundamentally address the limited exploration issue of open-loop systems, similar to recent Tree-of-Thought (ToT) [1] and Monte Carlo Tree Search (MCTS) [2], our approach is based on a tree search system with inevitable additional costs, resulting in lower efficiency. When compared with recent state-of-the-art tree-search systems [1][2], Toolchain* stands out by achieving the best efficiency (as discussed in Section 4.2). Enhancing model efficiency without significantly compromising performance remains a key direction for our future work.

[Weakness 3] Limited analysis on how the approach generalizes to even more expansive action spaces and longer planning horizons. However, I think that the above weaknesses are some hard open questions in NLP with LLM.

We acknowledge the importance of further analysis in more expansive action spaces and longer planning horizons. Besides Table 5 in Appendix C.2 (page 15), we provide more details of our tested environments:

The # TooIs represents the number of different tools in the environment. The Avg. # APIs/Solution represents the average number of different API functions called in each solution. For our evaluated environments, while the number of tools may appear limited, the combination of these tools with flexible parameters significantly expands the action space, which is comparable to ToolLLM [4], the largest known dataset (concurrent to our work) in the tool-use domain. Considering the multi-tool scenarios only compose tools within the same category, the ToolLLM dataset is of a similar scale to ours, averaging around 70 different tools in each category. In terms of solution length, the ToolLLM dataset features solutions with 6 to 15 actions, a range comparable to our datasets. Notably, some solutions in our virtual home dataset extend to approximately 36 actions, indicating a more complex environment.

Thank you for your reply! The authors have addressed my primary concerns effectively.

We express our gratitude to all reviewers for their valuable and constructive feedback. In response to the concerns raised, we have summarized our main modifications as follows:

• We performed additional experiments to benchmark our approach against various alternative cost function designs;
• We included hyperparameter studies of $\alpha$ and $\beta$ for model robustness;
• We conducted comparisons with additional variants of non-LLM cost functions and non-LLM agents.

More details are explained in the individual response for each reviewer. We thank the reviewers again and look forward to further suggestions or discussions.