MetaTool: Facilitating Large Language Models to Master Tools with Meta-task Augmentation

Recent studies have shed light on the potential of utilizing tools to augment LLMs with external factual knowledge Qin et al. (2023a); Nakano et al. (2021); Song et al. (2023); Hao et al. (2024); Shen et al. (2024); Gao et al. (2023); Wu et al. (2023); Qian et al. (2023); Zhuang et al. (2024); Schick et al. (2024) and complete tasks in complex environments Gupta & Kembhavi (2023). With the burgeoning intelligence in reasoning and perception, LLMs’ tool-use capability can be widely applied in the automation of various domains including Embodied AI Wang et al. (2024c; b), web manipulation Rawles et al. (2023); Hong et al. (2024); Yang et al. (2023); Deng et al. (2024); He et al. (2024); Zhou et al. (2023), and image/video editing Wang et al. (2024a); Argaw et al. (2022); Hang et al. (2024); Fu et al. (2023). Effectively mastering complex tools challenges the model to comprehend the precondition and potential outcome of using tools. In this paper, we aim to facilitate LLMs for tool-oriented tasks by learning robust tool understanding.

As noted by Hernik & Csibra (2009), when learning to utilize a specific tool, children perceive it as an object with particular functions, engaging in a cognitive process to understand its purpose and operation. Analogously, a comprehensive understanding of the tools’ functionalities is indispensable for enabling the controller to use tools proficiently. In real-world scenarios, tools are typically accompanied by a manual (or tutorial), which provides sufficient relevant details about their functionalities and usage. Endowed with strong few-shot learning Brown et al. (2020) and zero-shot learning Wei et al. (2021) capabilities, foundation models can be prompted to unravel tools’ functionalities and comprehend how to use them. To this end, we can construct suitable task-specific prompts either through manual design Vemprala et al. (2024) or retrieval Zhou et al. (2022). However, prompting is restricted by input context length, thus the situation may be more challenging with multiple complex tools with long descriptions. While most training-based tool learning methods rely on extensive expert-annotated solution data for goal-oriented tasks, the knowledge contained in the tool execution process itself remains unutilized. We propose a self-supervised data augmentation method to efficiently endow LLMs the comprehension of a set of tools.

A tool task can be defined as a tuple $\left\langle\mathcal{S},\mathcal{A},\mathcal{T},g\right\rangle$, where $\mathcal{S},\mathcal{A},\mathcal{T}$ is the state space, action space, and toolset, and $g$ is the goal state of the task. Toolset $\mathcal{T}=\{t\}_{N}$ consists of $N$ tools, each as a state transition function $s^{\prime}=t(s,\theta)$ that formalizes the outcome of state change when feeding the input parameters $\theta$ into the tool. An action $a=\left\langle t,\theta\right\rangle\in\mathcal{A}$ specifies the tool and its input. As an autonomous agent, an LLM should iteratively respond with actions and inputs according to the state until it reaches the goal. Broadly, when the tools can not alter any external state, tool output like retrieval results can be regarded as the state thus the goal state is the desired information.

We enhance the tool understanding of the model with self-supervised surrogate (pretext) tasks instead of in-context descriptions or demonstrations. Formally, we regard tools as external systems that implement state transition mappings. Tool understanding, therefore, involves comprehending the perception-action process of these systems (referred to as tool execution) and should be generalizable to various task objectives.

Meta-task definition. We first generate single-step tool execution data $\mathcal{D}=\{s,a,s^{\prime}\}$ by stochastically sampling initial state $s\in\mathcal{S}$ and action $a\in\mathcal{A}$, and obtaining the tool output $s^{\prime}$. Five surrogate tasks (meta-tasks) are designed based on the unsupervised dataset $\mathcal{D}$. Basically, the model is required to predict masked factors of the execution process in line with the idea of Masked Autoencoder (MAE) He et al. (2022). We define the meta-tasks as below:

• •

  Effect: The model predicts the outcome state $P(s^{\prime}|a,s)$ given the initial state and the action.

• •

  Decision-making: The model decides a feasible action $P(a|s,s^{\prime})$ given the initial and outcome state.

• •

  Reversion: The model deduces the initial state $P(s|a,s^{\prime})$ given the action and the outcome state.

• •

  Input Boundary: The model determines whether an action can be successfully executed given the current state: $P(\mathbbm{1}_{s^{\prime}\neq s}|a,s)$.

• •

  Output Boundary: The model determines whether a state can be reached with any action given the current state: $P(\mathbbm{1}_{\exists(t,\theta),s^{\prime}=t(s,\theta)}|s,s^{\prime})$.

• •

  Counterfact: The model predicts the new outcome state $P(s^{\prime\prime}|a,s^{\prime},a^{\prime})$ if a new action $a^{\prime}$ were executed given that the current action $a$ results in the current outcome $s^{\prime}$.

Effect, decision-making, reversion meta-tasks emphasize the causality of a tool, regarding the action as the intervention to the state Pearl (2009); Pearl & Mackenzie (2018) and the outcome as the causal effect is determined by the tool mechanism. On top of that, counterfact task is the composition of reversion and effect, further imagining the outcome altered from the fact in effect task and raising higher requirements for causal reasoning Bareinboim et al. (2015); Zhang & Bareinboim (2016). When implemented as APIs, tools may receive non-executable inputs and result in ineffective outcomes. Thus the input and output domains are also unique features of a tool as a function. We consider Input boundary meta-task that emphasizes the tool affordance that refers to what actions can be executed considering the situation and the precondition. Output boundary meta-task emphasizes the functionality of tools, that is, what goals can and cannot be achieved given the current state.

Metaset construction. Based on the single-step data $\mathcal{D}$, datasets of meta-tasks (referred to as metasets) are constructed by automatically generating question-answering pairs, showcased in Figure 2. For each sample and each meta-task, we insert the variables of states and actions into 5 templates (diversified with GPT-4) to obtain diverse QA data.

With the data of meta-tasks, we explore several manners to augment the tool manipulation ability to achieve task goals: 1) In-context learning: To enhance the tool understanding of LLMs in a training-free way, we incorporate several demonstrations of each meta-tasks in the system prompt. Existing works may include demonstrations in the tool cookbook, yet not in a systematic way. 2) Self-supervised Learning: Since we aim to build the model’s tool understanding as the foundation of tool-oriented learning, an intuitive manner is to train the LLM first on the metasets as the surrogate tasks and then on the solution data of tool-oriented tasks. In order to maintain the general ability of the model in the first stage, only the parameters of the query and value projection layers of the Transformer are updated instead of full-parameter training. We also propose to train the model on metasets separately to build the tool-use ability step by step. 3) Data augmentation: We also utilize the metasets as the augmented data of conventional instruction tuning methods that the metasets are mixed with solution data and the model is trained uniformly. The model trained on the mixed data is referred to as MetaTool.

To evaluate LLMs’ ability on tool-oriented tasks, we develop 3 tasks emphasizing closed toolset manipulation rather than long-term planning and reflection which are also crucial abilities for AI agents. The key challenge of these tasks is effectively utilizing tools to achieve the goal, which requires the model to understand the rules (preconditions) and the mechanisms of the toolset. The task definition and dataset construction are elaborated below.

SpellAnyWord (SAW). In this task, the agent needs to sequentially construct a string that contains the target string as a continuous substring. The initial state of the task is a void string. Two non-degradable tools (functions) are avaliable: hinzufügen: to add two adjacent letters in the alphabet to the end of the current string. The tool input $\theta$ should be the preceding letter (e.g. passing ’a’ to hinzufügen on current string ” will result in ’ab’). Swap: to swap the position of two adjacent letters in the current string. The input should be the preceding letter (e.g. passing ’a’ to Swap on ’ab’ will result in ’ba’). An example task: The target string is ’any’. A successful action sequence can be [hinzufügen(’a’), hinzufügen(’n’), hinzufügen(’y’), Swap(’a’), hinzufügen(’o’)], which will result in a state sequence [’ab’, ’abno’, ’abnoyz’, ’banoyz’, ’banyoz’] and the final string ’banyoz’ has ’any’ as a substring.

BlocksWolrd (BW). In this scenario, the agent needs to stack several blocks on the table into a target state with one hand. Only one block can be moved at a time. Two tools (functions) are avaliable: Pick: to pick a block in the hand. The tool input should be the target block indicated by its color (e.g. Pick(’yellow’)). Blocks cannot be picked if there are blocks on top of them or there’s already a block in the hand. Stack: to stack the block in the hand onto the target block or table. The input should be the color of the target block or ’table’ (e.g. Stack(’white’), Stack(’table’)). Blocks cannot be stacked on a block with another block already on top of it or there’s no block in the hand.

Logistics (LOG). The agent needs to solve a logistics problem by arranging trucks and airplanes to transport the package to the target location. Locations are grouped by cities. Trucks can be used to move packages between locations in the same city and planes can be used to move packages between cities. Two tools (functions) are available: Truck: to transport the truck and the package (if there is any) from one location to another. Plane: to transport the airplane and the package (if there is any) from one location to another. The tool input should be the starting and ending location indicated by numbers. (e.g. Truck(1,2), Plane(2,4)). An action is invalid when there is no truck or airplane at the starting location.

Datasets collection. For the SAW task, we randomly sample 2k target strings (from 2 letters to 10 letters) as task goals. We modify the BW and LOG tasks from the prior LLM benchmark Valmeekam et al. (2024) into the tool-use version, thus 2k goals for each task are adopted following the original configuration. Optimal action sequences are obtained with heuristic strategy as the solution data. For each annotated action, we generate a thought with one of 5 templates (diversified with GPT-4). The thoughts analyze the situation and what to do following ReACT Yao et al. (2022) to leverage the model’s reasoning ability. Thus the solution contains a sequence of thought-tool-input tuples. Besides the 3 tasks we define, MetaTool can be easily generalized to other tool-oriented tasks by writing the system prompt and developing the tool functions or APIs as the external environment. The meta-task data can then be generated automatically.

Our model is fine-tuned based on LLaMA3-8b-instruct AI@Meta (2024) with parameter-efficient fine-tuning method Qlora Dettmers et al. (2024) on 8 A100 GPUs. We utilize the instruction tuning version of LLaMA3 since comprehending tool-oriented tasks with specific objectives is the basis of tool understanding and manipulation. For tool-oriented task training, we construct instruction-solution pairs from the task goals and annotated solutions. For meta-task training, we formulate the task objectives and the tool execution outcome as question-answering pairs. For each task, we train the model on 10k meta-task data and 10k solution data for 3 epochs with AdamW optimizer and the learning rate of 2e-4. The models are tested in a simulated environment that receives the action of using a tool and returns the outcome and current state. We evaluated the model performance on 100 unseen cases of each three tasks.

Overall comparison. We evaluate the success rate (SR%) of completing each task and show the performances of several models in Table 1. Overall, SOTA closed-source LLMs show impressive zero-shot performance on tool-oriented tasks compared with open-source LLMs including LLaMA3 and Vicuna. By training on both meta-tasks and solution data, our model MetaTool gains significant improvement (+22.7%SR on average) compared with LLaMA3 (baseline) and surpasses GPT-4/GPT-3.5-turbo in the SAW/BW tasks (+3.5%/11.0%SR). Both GPT and LLaMA3 show weaker performances when provided with meta-task demonstrations (IC) since demonstrating limited cases can be redundant or misleading without proper design. LLaMA3-SS that trained on meta-tasks first gains limited improvement compared with the baseline. We conjecture that learning meta-tasks without practicing tool manipulation (training on action sequences) cannot effectively facilitate tool-use ability with tool understanding. Also fine-tuning with specific QA data may affect the basic linguistic ability of the model.

Ablation study. We study the ablation of different data components and report the performances in Table 2. Merely training on solution data improves the model performance to an extent compared with the baseline LLaMA3 (+8.5% on average). It’s worth noticing that only training on meta-tasks can improve the model’s zero-shot performance on tool-oriented tasks (line 2), contrary to providing demonstrations of meta-tasks in the system prompt (LLaMA3-IC in Table 1). When removing QA data from each meta-task, the model performance shows varying degrees of degradation, which verifies the profits of meta-tasks. The meta-tasks of effect and decision-making have a relatively greater influence on the model’s tool understanding capability. Theoretically, these meta-tasks represent the causal mechanism of tools, which is the basis of high-level understanding or skills.

Case study. As showcased in Figure 3, the agent is required to construct stacks containing a blue block on top of a yellow block from a pile of 4 blocks. With mere descriptions of tools in the prompts, LLaMA3 fails to understand the precondition of using tools resulting in invalid actions. Training on tool-oriented solution data, LLaMA3-solution attempts to lift the yellow block successfully but fails to sequentially achieve the task goal and falls into repetitive loops. The proposed MetaTool model achieves the target state with an effective action sequence (although still not the optimal efficiency) and corresponding reasoning. These 3 models correspond to the 3 paradigms illustrated in Figure 1. The results show that LLMs can learn tool manipulation better on the basis of robust tool understanding.