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

Abstract: LLMs have demonstrated powerful decision-making and planning capabilities in solving complicated real-world problems. LLM-based autonomous agents can interact with diverse tools (e.g., functional APIs) and generate solution plans that execute a series of API function calls in a step-by-step manner. The multitude of candidate API function calls significantly expands the action space, amplifying the critical need for efficient action space navigation. However, existing methods either struggle with unidirectional exploration in expansive action spaces, trapped into a locally optimal solution, or suffer from exhaustively traversing all potential actions, causing inefficient navigation. To address these issues, we propose ToolChain*, an efficient tree search-based planning algorithm for LLM-based agents. It formulates the entire action space as a decision tree, where each node represents a possible API function call involved in a solution plan. By incorporating the A* search algorithm with task-specific cost function design, it efficiently prunes high-cost branches that may involve incorrect actions, identifying the most low-cost valid path as the solution. Extensive experiments on multiple tool-use and reasoning tasks demonstrate that ToolChain* efficiently balances exploration and exploitation within an expansive action space. It outperforms state-of-the-art baselines on planning and reasoning tasks by 3.1% and 3.5% on average while requiring 7.35x and 2.31x less time, respectively.

• The paper presents ToolChain* which integrates A* search to balance exploration and exploitation in large decision spaces.
• The method utilizes a task-specific cost function combining cumulative and future cost heuristics to efficiently prune high-cost paths.
• Experiments on ToolBench and GSM8K demonstrate that ToolChain* outperforms baselines in both efficiency and accuracy.

ToolChain*: Efficient Action Space Navigation in LLMs with A* Search

The paper presents ToolChain$^*$, a novel approach for navigating expansive action spaces in LLMs using an A* search-based algorithm. It addresses the inefficiencies and limitations of existing methods by effectively balancing exploration and exploitation. The focus is on improving decision-making in LLMs through a tree search-based framework.

Introduction to ToolChain$^*$

ToolChain$^*$ proposes an A* search integrated with a task-specific cost function to navigate the decision spaces of LLMs formulated as decision trees. Each node in the tree signifies a potential API function call, allowing the method to efficiently prune high-cost paths and pursue the most promising routes. The integration of A* search with LLMs enables the efficient handling of large action spaces by focusing on cost-effective paths.

Figure 1: A comparison of existing methods that leverage LLMs for decision-making from a searching space perspective.

Methodology

ToolChain$^*$ operates through an iterative process comprising three phases: selection, expansion, and update (Figure 2).

• Selection: Identify and expand the node with the lowest estimated cost from the frontier nodes.
• Expansion: Generate potential actions for the next step using LLMs, expanding the decision tree.
• Update: Re-evaluate cost functions with task-specific heuristics to prioritize nodes for future exploration.

  Figure 2: ToolChain$^*$ framework of three phases: (a) selection, (b) expansion, and (c) update.

Cost Function Design

The cost function $f(n) = g(n) + h(n)$ combines:

• Cumulative Cost $g(n)$: Evaluates the cost of actions in the current plan using a geometric mean of task-specific heuristics and self-consistency frequency.
• Future Cost $h(n)$: Estimates the remaining path cost using task-specific heuristics and the LLM's imaginative capabilities.

Experiments

ToolChain$^*$ was evaluated on tool-use tasks from ToolBench and reasoning tasks from GSM8K, demonstrating superior performance and efficiency.

• On ToolBench datasets, ToolChain$^*$ outperformed baselines such as ReAct and MCTS, achieving higher success rates.
• In GSM8K, ToolChain$^*$ showed improved accuracy over baseline methods including Chain-of-Thought and Self-Consistency.

Figure 3: Time efficiency evaluation on ToolBench and GSM8K shows ToolChain$^*$ efficiency over other methods.

Case Studies

ToolChain$^*$'s capacity to explore expansive decision spaces while mitigating error propagation is illustrated through qualitative assessments. Real-world case studies highlight its adaptability in revising plans and avoiding error paths, distinguishing it from unidirectional methods like Chain-of-Thought.

Figure 4: ToolChain$^*$ vs. Chain-of-Thought on GSM8K dataset showcasing exploration space.

Conclusion

ToolChain$^*$ effectively addresses the challenges of navigating extensive action spaces in LLM-based decision-making. It mitigates the limitations of prior methods by integrating an efficient search framework with robust cost function design, enabling practical applications in real-world scenarios. Its ability to handle complex tasks with improved efficiency makes it a promising approach for future development in AI planning and reasoning capabilities.