PaperlayerView exact revision on arXiv

Source-linked AI summary

ToolkenGPT: Augmenting Frozen Language Models with Massive Tools via Tool Embeddings

Shibo Hao, Tianyang Liu, Zhen Wang, Zhiting Hu

arXiv:2305.11554v4cs.CLcs.LG

TL;DR

Tool-augmented LLMs face costly, fixed-tool fine-tuning and context limits that constrain in-context demonstrations. ToolkenGPT represents tools as learned tokens that can be added to a frozen model, and improves performance across numerical reasoning, knowledge-based question answering, and embodied plan generation.

  • Problem

    Fine-tuning tool use can be costly and tied to predefined tools, while in-context demonstrations are constrained by context length and can be challenging for learning many tools.

    Source: §1 · p. 2

  • Method

    ToolkenGPT represents each tool as a learned token embedding that can be added to a frozen LLM, then prompts the model to generate arguments when a tool token is predicted.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    ToolkenGPT significantly enhances performance across numerical reasoning, knowledge-based question answering, and embodied plan generation, and outperforms baselines across one-hop and multi-hop numerical scenarios.

    Source: §1 · p. 2, §4.1 · p. 7, §5 · p. 10

  • Takeaways & Limitations

    The findings show that frozen LLMs can rapidly adapt to and use a larger set of tools across complex scenarios.

    Source: §5 · p. 10

Abstract

from arXiv · show

Augmenting large language models (LLMs) with external tools has emerged as a promising approach to solving complex problems. However, traditional methods, which finetune LLMs with tool demonstration data, can be both costly and restricted to a predefined set of tools. Recent in-context learning paradigm alleviates these issues, but the limited context length only allows for a few shots of demonstrations, leading to suboptimal understandings of the tools. Moreover, when there are numerous tools to choose from, in-context learning could completely fail to work. In this paper, we propose an alternative approach, $\textbf{ToolkenGPT}$, which combines the benefits of both sides. Our approach represents each $\underline{tool}$ as a to$\underline{ken}$ ($\textit{toolken}$) and learns an embedding for it, enabling tool calls in the same way as generating a regular word token. Once a toolken is triggered, the LLM is prompted to complete arguments for the tool to execute. ToolkenGPT offers the flexibility to plug in an arbitrary number of tools by expanding the set of toolkens on the fly. In addition, it improves tool use by allowing extensive demonstration data for learning the toolken embeddings. In diverse domains, including numerical reasoning, knowledge-based question answering, and embodied plan generation, our approach effectively augments LLMs with tools and substantially outperforms various latest baselines. ToolkenGPT demonstrates the promising ability to use relevant tools from a large tool set in complex scenarios.

1 Introduction

ToolkenGPT enables LLMs to use large sets of tools by representing each tool as a learned token embedding, without fine-tuning the LLM. It combines lightweight adaptation with extensive demonstration learning and tool expansion.

Motivation

  • LLMs use tools such as calculators and databases to address weaknesses in world knowledge, hallucination, and symbolic operations, while rapidly emerging tools make adaptation to massive tool sets important.The introduction frames tool mastery as relevant to applications ranging from writing assistance to automated customer support.Source: §1 · p. 1
  • Fine-tuning, including Toolformer [56], is costly and tool-specific, while in-context learning [69] [49] [53] is constrained by context length and few-shot difficulty with unusual tools such as those challenging GPT-4 [6].Toolformer fine-tuned GPT-J to learn five tools; in-context learning struggles to demonstrate massive tool sets.Source: §1 · p. 2

Contribution

  • ToolkenGPT represents each tool as a vocabulary token with a learned embedding; when predicted, prompting elicits arguments, executes the tool, and reinjects outputs.The method learns lightweight toolken embeddings and can add tools by expanding the vocabulary, without fine-tuning the LLM.Source: §1 · p. 2

Results

  • Across numerical reasoning, ToolkenGPT outperforms Chain-of-Thought [65] and ReAct [69], supports over 200 relation APIs for question answering, and improves embodied-plan grounding with 58 actions and objects.The paper reports better grounding than previous in-context learning and specialized decoding methods.Source: §1 · p. 2

2 Related Works

Related work on tool use has relied on fine-tuning or in-context demonstrations, while parameter-efficient tuning adapts only part of a frozen model. ToolkenGPT instead learns embeddings for tools as tokens, aiming to capture tool knowledge efficiently and support many tools.

Fine-tuning LLMs to use tools

  • Prior work fine-tuned LLMs for individual-domain tools or general tool collections, but such tuning is costly and struggles with emergent or updated tools.Examples include retriever-based REALM, RAG [35], and RETRO [3], WebGPT [46], and general-tool systems TALM [50] and Toolformer [56]; ToolkenGPT learns new toolken embeddings without gradient calculations for LLM parameters, with training cost similar to inference and minimal GPU memory overhead.Source: §2 · pp. 2–3

In-context learning for tools

  • In-context tool use supplies descriptions and demonstrations, but performs poorly in complex scenarios with unfamiliar or numerous tools; ToolkenGPT embeddings can encode semantics from demonstrations beyond surface text.Unlike Li et al. [36], which retrieves tools using text embeddings of their documents, ToolkenGPT uses toolken embeddings; it is also compatible with prompting such as Chain-of-Thought (CoT) [65].Source: §2 · p. 3

Efficient tuning of large language models

  • Parameter-efficient fine-tuning freezes most LLM parameters and tunes a subset, but existing methods have not proven suitable for efficient tool learning; ToolkenGPT presents token prediction as an efficient approach for massive tools.Examples include Adapters [22], BitFit, prompt tuning [34] [64], and LoRA [23]; the paper states it is the first to explore efficient tuning methods for predicting tools as tokens.Source: §2 · p. 3

3 ToolkenGPT for Mastering Massive Tools

ToolkenGPT represents tools as learned tokens that the language model can select alongside ordinary words, then uses demonstrations to generate tool arguments. It keeps the language model frozen and trains only toolken embeddings.

3.1 Framework Overview

  • ToolkenGPT concatenates learned toolken embeddings with the language-model head so the next token can be a word or toolken, and new tools can be added by expanding the embedding matrix.This lets the model consider tools and words uniformly during generation.Source: §3.1 · p. 4
  • After predicting a toolken, the model switches to tool mode, uses demonstrations to generate arguments, executes the call, and resumes reasoning with the returned result.The demonstrations concern the selected tool and use a special syntax for tool calls.Source: §3.1 · p. 4, Figure 1 · p. 4

3.2 Learning Toolken Embeddings

  • Only toolken embeddings are optimized while the original language-model parameters remain frozen, keeping GPU memory use near inference levels and allowing training from massive demonstrations.New tools can be added by expanding the embedding matrix and refining the new embedding with tool demonstrations.Source: §3.2 · pp. 4–5
  • Training pairs mark tool calls with a toolken and replace returned-result tokens with [N/A], which the loss ignores so the model learns to predict the call rather than its result.At inference, the tool's returned value is filled back into the text.Source: §3.2 · p. 5
  • Training pairs can come from supervised data with ground-truth tool calls or from LLM-synthesized demonstrations generated using tool documentation and example calls.The paper frames synthetic demonstrations as distilling knowledge from the LLM into toolken embeddings.Source: §3.2 · p. 5

4 Experiments

Across arithmetic, knowledge-base querying, and embodied planning, the experiments report that ToolkenGPT learns extensive tool sets and improves task performance over advanced prompting methods.

4.1 Numerical Reasoning

  • ToolkenGPT outperforms all baselines on FuncQA one-hop and multi-hop, while matching or surpassing in-context learning with four operators on GSM8K-XL.Its embeddings, trained only on one-hop synthetic data and without Chain-of-Thought examples, also improve multi-hop performance and can be combined with Chain-of-Thought prompting. GSM8K-XL tests 568 cases with magnified numbers; FuncQA requires at least 13 operators.Source: §4.1 · pp. 5–7, Table 2 · p. 6

4.2 Knowledge-based Question Answering

  • On KAMEL, ToolkenGPT (sup) achieves the highest results by a large margin, and ToolkenGPT (syn) beats all baselines across relation subsets without in-domain training data.In-context learning drops sharply beyond 30 tools due to context limits, while description-only in-context learning generally fails to map descriptions to relation identifiers.Source: Figure 2 · p. 7, §4.2 · pp. 7–8

4.3 Embodied Plan Generation

  • ToolkenGPT achieves the highest VirtualHome success rate, learning to predict valid actions and objects and to choose SIT <chair> rather than the invalid SIT <desk>.The baselines’ desk prediction reflects a misunderstanding of VirtualHome’s SIT action, which means “sit on”; a desk is not sittable.Source: Table 3 · p. 9, §4.3 · p. 9, Figure 3 · p. 9

4.4 Analysis

  • On FuncQA, LoRA performs slightly better than ToolkenGPT, but fine-tuning takes significantly longer; ToolkenGPT also enables plug-and-play massive tools through decoupled parameters.Both methods are based on LLaMA-7B.Source: §4.4 · pp. 9–10, Table 4 · p. 10
  • Adding tool mode improves ReAct’s argument-completion accuracy, but ToolkenGPT still outperforms the improved baseline by a large margin, indicating better tool selection.In tool mode, the LLM receives demonstrations using only the selected tool.Source: Table 5 · p. 10, §4.4 · p. 10
  • On KAMEL, supervised training data performs better than synthetic data at the same size, while increasing training data benefits both approaches.The authors suggest a distribution gap between synthetic training examples and the test set may limit synthetic-data performance.Source: §4.4 · p. 10

5 Conclusion · A Details of Numerical Reasoning · A.1 GSM8K-XL

ToolkenGPT augments frozen LLMs with tool embeddings, enabling token-like tool calls and use of a larger tool set with extensive demonstration data. The paper reports improved performance across several task domains and highlights rapid adaptation to new tools.

5 Conclusion

  • ToolkenGPT significantly enhances LLM performance across numerical reasoning, knowledge-based question answering, and embodied plan generation.The conclusion attributes these gains to toolken embeddings.Source: §5 · p. 10
  • Each tool receives a toolken embedding, letting frozen LLMs call and use tools as easily as generating word tokens.The approach avoids expensive fine-tuning and is designed to accommodate a much larger tool set while using extensive demonstration data.Source: §5 · p. 10
  • ToolkenGPT rapidly adapts to and leverages new tools, addressing the evolving landscape of available tools.Source: §5 · p. 10
  • Future work could learn robust toolken embeddings from tool descriptions and input-output records as well as demonstration data.The conclusion cites [24] as an example of other rich forms of experience.Source: §5 · p. 10
  • The authors also propose integrating toolken embeddings with advanced planning techniques to develop an autonomous agent for complex real-world problems.The conclusion cites [19] in connection with advanced planning techniques.Source: §5 · p. 10

A Details of Numerical Reasoning

  • For numerical-reasoning comparisons, CoT and ToolkenGPT reasoning use the same prompts; ReAct3 prompts label calculations, and tool-mode prompts use four sampled tool examples.The sampled examples are transformed into ReAct-style prompts; prompts are placed in the supplementary file because many tools are used.Source: §5 · p. 10, A Details of Numerical Reasoning · p. 16

A.1.1 Data Synthesis

GSM8K-XL is synthesized by converting GSM8K examples into variable-based templates, validating the rewrites, and solving them after magnifying eligible numbers. Percentage values remain unchanged when cubic scaling would distort their meaning.

A.1.1 Data Synthesis

  • ChatGPT (gpt-3.5-turbo), prompted with two examples, rewrites GSM8K question-answer pairs into variable-based templates.The prompt asks for a more general form using variables instead of specific numbers.Source: §A.1.1 · p. 16
  • A verification function evaluates formulas in the solutions, substitutes original numbers into the rewrites, and excludes pairs whose results differ from the original answers.This check is used to validate rewrites before magnification.Source: §A.1.1 · p. 16
  • After validation, original numbers are replaced with their cubic values, except percentage values, which remain unchanged when cubic scaling would be inappropriate.For example, directly cubing 10% would yield 1000%, so the original percentage is preserved.Source: §A.1.1 · p. 16
  • The magnified values are inserted into the templates, and a solving function processes the resulting formulas to produce GSM8K-XL.The solving function returns solutions for the magnified question-answer pairs.Source: §A.1.1 · p. 16

A.1.2 Training Details · A.1.3 Prompt for GSM8K-XL Dataset · A.2 FuncQA

The appendix reports the toolken-embedding training setup and gives GSM8K-XL prompt formats for direct prompting and reasoning with arithmetic tool operators.

A.1.2 Training Details

  • Toolken embeddings use 5,063 training examples and 1,000 validation examples, with a 5e-4 learning rate and early stopping within 10 epochs.Early stopping is based on the development set.Source: §A.1.2 · p. 17

A.1.3 Prompt for GSM8K-XL Dataset

  • Direct prompting asks ChatGPT to solve each math problem step by step and end with “So, the answer is xxx.”The prompt uses a [QUESTION] placeholder.Source: §A.1.3 · p. 17
  • The CoT and ToolkenGPT reasoning prompt instructs the model to answer questions step by step.The examples show worked solutions followed by a final numeric answer.Source: §A.1.3 · p. 17
  • Direct-prompt demonstrations show worked arithmetic for fish counts, contest points, store revenue, and protein totals.Each example concludes with a numeric answer.Source: §A.1.3 · pp. 17–18
  • The operator prompt directs the model to answer questions using <add>, <subtract>, <multiply>, and <divide>.The operator-form prompt also provides a [QUESTION] placeholder.Source: §A.1.3 · pp. 18–19
  • Operator-annotated demonstrations insert tool calls into worked solutions for fish counts, contest points, store revenue, and protein totals.The examples use <multiply>, <add>, and <divide> expressions and retain the final numeric answers.Source: §A.1.3 · pp. 18–19

A.2.1 Training Details

Toolken embeddings are trained on a subset of examples, with validation-based early stopping.

A.2.1 Training Details

  • Training uses 611 examples, reserves 39 for validation, sets the learning rate to 1e-4, and applies development-set early stopping for at most 20 epochs.Source: §A.2.1 · p. 19

A.2.2 Prompt for Synthetic Training Data … B Details of Knowledge-based QA

The appendix specifies how the authors synthesize tool-training examples and format one-hop and multi-hop FuncQA prompts, then introduces the knowledge-based QA data and training details. It also describes filtering generated examples and removing duplicates.

A.2.2 Prompt for Synthetic Training Data

  • The synthesis pipeline starts from two manually crafted examples, then filters nonconforming generations and removes duplicates.The prompt does not guarantee that generated examples follow the required format.Source: §A.2.2 · p. 19
  • The synthetic-data prompt requests varied real-world, single-step questions and answers using a described function, with numbers replaced by [ARG] and [ANSWER].It presents two examples and uses [FUNC] and [DESCRIPTION] placeholders.Source: §A.2.2 · p. 19

A.2.3 Prompt for FuncQA One-Hop

  • The zero-shot ChatGPT prompt asks for a math solution and a final answer formatted as “So, the answer is xxx.”The question is inserted at [QUESTION].Source: §A.2.3 · pp. 19–20
  • The CoT and ToolkenGPT one-hop prompts demonstrate worked solutions, with ToolkenGPT encoding operations such as <add>, <subtract>, and <divide> as tool calls.The prompt lists available operation tokens and ends with a [QUESTION] placeholder.Source: §A.2.3 · p. 20

A.2.4 Prompt for FuncQA Multi-Hop

  • The multi-hop zero-shot ChatGPT prompt requests step-by-step solutions and a final answer in the “So, the answer is xxx.” format.The examples cover problems requiring multiple arithmetic steps.Source: §A.2.4 · p. 21
  • The multi-hop ToolkenGPT prompt lists operation tokens and demonstrates their use in worked solutions involving probability, geometry, logarithms, and time conversion.The examples encode operations such as <power>, <choose>, <divide>, and <multiply>.Source: §A.2.4 · pp. 21–22

B Details of Knowledge-based QA

  • The knowledge-based QA appendix describes converting Wikidata relation identifiers such as P1346 into natural-language descriptions, then outlines data synthesis and training settings.Source: B Details of Knowledge-based QA · p. 22

B.1 Getting Text Description

KAMEL generates text descriptions for each relation by sampling three dataset facts, turning them into question–answer examples, and prompting ChatGPT to define an API from the relation template and examples.

B.1 Getting Text Description

  • For each relation, the method randomly samples three dataset facts and instantiates them as question–answer pairs for ChatGPT.These pairs are used to generate a description for the relation.Source: §B.1 · p. 22
  • The prompt asks ChatGPT to define an API that answers a question given its template and example answers.It supplies three example question–answer pairs before requesting the API.Source: §B.1 · p. 22
  • Examples pair artwork-language questions with original_language(title) and person-language questions with spoken_languages(name).The prompt describes these APIs as retrieving an artwork's original language and a person's spoken languages, respectively.Source: §B.1 · p. 22

B.2 Synthetic Data

The section describes two prompts for synthesizing diverse training data, with samples from both aggregated. The examples show tool calls embedded in continuing text and distinguish ToolkenGPT’s needs from those of in-context learning baselines.

B.2 Synthetic Data

  • Two prompts synthesize diverse training examples, and samples from each prompt are aggregated.Source: §B.2 · p. 22
  • Natural-language tool descriptions are unnecessary for ToolkenGPT’s demonstration-based training but crucial to in-context baselines, especially ICL (desc).Source: §B.2 · p. 22
  • One prompt embeds varied function calls and returned values in explanatory sentences, then requests 10 more examples for a named function.Source: §B.2 · p. 23
  • The other prompt illustrates weather and calculator calls embedded in text, then requests 10 more examples for a named function.Source: §B.2 · p. 23

B.3 Training Details … C.2 Prompts

The appendix reports toolken embedding training settings and describes VirtualHome preprocessing and the prompts used for embodied plan generation.

B.3 Training Details

  • Toolken embeddings use a learning rate of 1e-4, early stopping on the development set, and a maximum of 5 training epochs.Source: §B.3 · p. 23

C Details of Embodied Plan Generation

C.1 Preprocessing

  • The VirtualHome scripts are filtered for executability and state changes, deduplicated by goal and instruction, screened for ambiguous object names, and limited to frequent verbs and objects.The retained verbs and objects occur more than 10 times in the data.Source: §C.1 · p. 24
  • Unlike Huang et al. [25], the preprocessing treats scripts with the same goal but different instructions as distinct tasks and evaluates them by exact resulting-state matching rather than human judgment.Different instructions may imply different action sequences and final state graphs; evaluation uses the Evolving Graph 5 to compare generated and ground-truth states.Source: §C.1 · p. 24

C.2 Prompts

  • All methods use the same prompts for plan generation in this experiment.Source: §C.2 · p. 24
  • The prompt specifies the robot's available actions and presents a task with its location, manipulable objects, goal, hint, and example plan.One example asks the robot to write an email and shows actions such as [WALK], [FIND], [TURNTO], and [SWITCHON].Source: §C.2 · p. 24

C.3 Training Details · D Computational Resources · E Safeguard Statement

Toolken embeddings use early-stopped development-set training, and the paper reports its GPU setup and cautions that future tool-learning applications may pose misuse risks.

C.3 Training Details

  • Toolken embeddings are trained at a 1e-4 learning rate, with development-set early stopping and a maximum of 10 epochs.Source: §C.3 · p. 25

D Computational Resources

E Safeguard Statement

  • The paper focuses on mathematical, knowledge-based, and embodied planning applications, while warning that future border applications of tool learning may be misused.It recommends careful consideration of safety before real-world application.Source: E Safeguard Statement · p. 25

Related papers

Loading 2305.11554v4…