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

Existing tool-learning approaches are costly, restricted to predefined tools, or constrained by context length and few-shot demonstrations. ToolkenGPT learns embeddings for tools represented as tokens while keeping the LLM frozen, then prompts for arguments after a toolken is selected. Across numerical reasoning, knowledge-based question answering, and embodied plan generation, it substantially outperforms recent baselines and supports large, rapidly adaptable tool sets.

  • Problem

    Fine-tuning tool-using LLMs is costly and restricted to predefined tools, while in-context learning cannot demonstrate massive tools within limited context and may fail with many choices.

  • Method

    ToolkenGPT represents each tool as a vocabulary token with a learned embedding, selects toolkens during generation, and prompts the frozen LLM to complete tool arguments.

  • Results

    ToolkenGPT substantially outperforms recent baselines across numerical reasoning, knowledge-based question answering, and embodied plan generation, including knowledge-base settings with up to 234 tools.

  • Takeaways & Limitations

    ToolkenGPT demonstrates that frozen LLMs can use relevant tools from large tool sets while adapting new tools through expandable toolken embeddings.

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

Tool integration research balances costly, tool-specific fine-tuning against context-limited in-context learning. ToolkenGPT represents tools as embeddings, enabling large tool sets and extensive demonstrations, and improves performance across diverse tasks.

  • Motivation: Fine-tuning is costly and tied to predefined tools, while in-context learning is limited by context length and few-shot demonstrations.In-context learning can struggle when many tools must be selected or demonstrated.
  • Approach: ToolkenGPT represents each tool as a new vocabulary token with a learned embedding, avoiding LLM fine-tuning while adapting quickly to new tools.Toolkens are inserted into the language-model head and expanded as tools are added.
  • Approach: ToolkenGPT combines extensive demonstration data with support for massive tool sets by adding respective toolkens to the vocabulary.This design targets limitations of both fine-tuning and in-context learning.
  • Evaluation: The framework is evaluated on numerical reasoning, knowledge-based question answering, and embodied plan generation.The introduction reports tool use across these three problem domains.

2 Related Works

Prior tool-learning work primarily uses fine-tuning or in-context demonstrations. ToolkenGPT instead tunes lightweight toolken embeddings while keeping the language model frozen, targeting efficient adaptation to many tools.

  • Fine-tuning LLMs to use tools: Fine-tuning methods augment language models with one or a few domain-specific tools but are costly and struggle with emergent or updated tools.ToolkenGPT avoids gradient updates to the language-model parameters when adapting new tools.
  • Efficient tuning of large language models: ToolkenGPT learns lightweight embeddings for new tools without gradient calculation through the language-model parameters, with minimal GPU memory overhead.The stated training cost is similar to language-model inference.
  • In-context learning for tools: In-context learning uses tool descriptions and demonstrations in prompts and supports newly introduced tools, but context limits restrict the number of usable examples and tools.This paradigm underlies applications such as ChatGPT plugins and LangChain.
  • Efficient tuning of large language models: Existing parameter-efficient tuning methods freeze most language-model parameters, but the paper argues they are not proven suitable for efficient tool learning.ToolkenGPT focuses specifically on predicting tools as tokens for massive-tool learning.

3 ToolkenGPT for Mastering Massive Tools

ToolkenGPT augments a frozen language model by treating tools as selectable tokens, then switches to prompted argument generation when a toolken is chosen. Tool outputs re-enter text generation, while toolken embeddings can be trained from extensive demonstrations and expanded for new tools.

  • 3.1 Framework Overview: Tool execution returns results to the reasoning context, allowing generation to continue until a final answer is produced.The framework supports tools that return results or affect a real-world environment.
  • 3.1 Framework Overview: ToolkenGPT concatenates toolken embeddings with the language-model head, so the next token may be either a word token or a toolken.New tools can be added by expanding the toolken embedding matrix.
  • 3.1 Framework Overview: After a toolken is predicted, the model switches from reasoning mode to tool mode and uses demonstrations to generate the tool arguments.The tool call is then executed and its result is returned to continued reasoning.
  • 3.2 Learning Toolken Embeddings: The language-model parameters remain frozen; only the toolken embedding matrix is optimized, avoiding gradients through the model body.The authors describe this as more stable and efficient than tuning methods that update additional model components.
  • 3.2 Learning Toolken Embeddings: New toolkens can be added and refined with subsequent demonstration training, including massive demonstrations rather than only a few in-context examples.Training pairs can come from ground-truth tool calls or synthetic tool-use cases generated by language models.
  • 3.2 Learning Toolken Embeddings: ToolkenGPT trains on paired text and tool-call sequences while ignoring returned-value tokens during the objective.This keeps training aligned with inference, where the model predicts the toolken and receives the executed result afterward.

4 Experiments

Experiments evaluate ToolkenGPT for numerical reasoning, knowledge-based question answering, and embodied plan generation. Across these settings, it improves tool use with many tools, while toolken embeddings offer efficiency and support for tool selection and argument completion.

  • Experimental Settings: ToolkenGPT is evaluated on arithmetic tools, database APIs, and robot actions across numerical reasoning, knowledge-based question answering, and embodied plan generation.The experiments assess both tool-calling accuracy and task-solving success.
  • 4.1 Numerical Reasoning: ToolkenGPT and ReAct handle large-number GSM8K-XL calculations better than ChatGPT and few-shot Chain-of-Thought, while both generally select correctly among four basic operators.GSM8K-XL magnifies calculation numbers to make the task more challenging.
  • 4.1 Numerical Reasoning: As the number of tools increases, ReAct misses, misselects, or mis-arguments tools, whereas ToolkenGPT outperforms all baselines on one-hop and multi-hop FuncQA.ToolkenGPT embeddings trained on one-hop synthetic data also improve multi-hop contexts and can be combined with Chain-of-Thought prompting.
  • 4.2 Knowledge-based Question Answering: On KAMEL, ToolkenGPT (sup) achieves the highest results, while ToolkenGPT (syn) outperforms all other baselines across subsets without in-domain training data.In-context learning suffers sharp drops beyond 30 tools because of context-length limits and difficulty selecting relations.
  • 4.3 Embodied Plan Generation: In VirtualHome, ToolkenGPT achieves the highest success rate and naturally predicts valid actions and objects, unlike baselines that often produce ungrounded or physically invalid plans.Translation and Grounded Decoding improve grounding but do not substantially improve action and object understanding.
  • 4.4 Analysis: Fine-tuning slightly outperforms ToolkenGPT on FuncQA, but takes significantly longer than training toolken embeddings; ToolkenGPT also supports plug-and-play massive tools.The framework attributes this flexibility to decoupled parameters for different tools.
  • 4.4 Analysis: Adding tool-mode prompting improves ReAct argument completion, yet ToolkenGPT remains substantially better, indicating toolken embeddings aid decisions about when and which tool to call.The ablation isolates contributions from tool selection and argument completion.
  • 4.4 Analysis: Under the same data budget, supervised training performs better than synthetic training, while increasing the training set benefits both data sources.The authors attribute weaker synthetic-data performance to a distribution gap between synthetic examples and the test set.

5 Conclusion

The paper presents ToolkenGPT as a way to augment frozen LLMs with massive external tools without expensive fine-tuning. Across numerical reasoning, knowledge-based question answering, and embodied plan generation, toolken embeddings improve performance and support rapid adaptation to new tools.

  • 5 Conclusion: ToolkenGPT augments frozen LLMs with massive external tools without expensive fine-tuning.It represents each tool with a toolken embedding that allows tool use like generating a word token.
  • 5 Conclusion: ToolkenGPT enables LLMs to call different tools as easily as generating word tokens.The approach uses extensive demonstration data to learn toolken embeddings while accommodating a larger tool set.
  • 5 Conclusion: ToolkenGPT significantly enhances LLM performance across numerical reasoning, knowledge-based question answering, and embodied plan generation.The reported tasks span arithmetic tools, database-related question answering, and embodied planning.
  • 5 Conclusion: ToolkenGPT rapidly adapts to and leverages new tools.The paper frames this capability as useful for keeping pace with the evolving landscape of massive tools.
  • 5 Conclusion: The experimental workflow includes testing-dataset construction, synthesized training data, method prompts, and training settings.The appendix also states that CoT and ToolkenGPT reasoning mode use the same prompts for fair comparison.

A.1.1 Data Synthesis

The data-synthesis workflow rewrites numerical question-answer pairs with variables, verifies the rewritten formulas, magnifies suitable numbers, and reconstructs GSM8K-XL examples.

  • A.1.1 Data Synthesis: The workflow rewrites numerical questions and answers using variables instead of specific numbers.ChatGPT is prompted with examples to replace concrete numbers with appropriate placeholders.
  • A.1.1 Data Synthesis: A verification function recalculates embedded formulas and excludes rewrites whose substituted results disagree with the original answers.This validation step checks the rewritten question-answer pairs before magnification.
  • A.1.1 Data Synthesis: The workflow replaces suitable original numbers with cubic values while preserving percentage values when cubic magnification would be inappropriate.The stated exception avoids changing values such as 10% into 1000%.
  • A.1.1 Data Synthesis: The magnified placeholders are substituted back into the rewritten pairs, and a solving function produces the GSM8K-XL dataset.The process obtains solutions for the magnified formulas after reconstructing the examples.
  • A.1.1 Data Synthesis: The examples include arithmetic transformations for egg sales and truck-load capacity problems.The rewritten forms use variables such as v_1, v_2, and v_3 while retaining explicit solution formulas.

A.1.2 Training Details

The appendix presents prompts for direct prompting, CoT, and ToolkenGPT reasoning, using arithmetic operator calls or step-by-step solutions on example problems.

  • A.1.2 Training Details: Direct prompting asks the model to answer a question without additional examples or reasoning instructions.The prompt format is simply “Question: [QUESTION] Answer:”.
  • A.1.2 Training Details: CoT and ToolkenGPT reasoning mode use prompts that request step-by-step answers.The appendix labels this prompt family for Chain of Thought and ToolkenGPT reasoning mode.
  • A.1.2 Training Details: The reasoning examples solve fish, contest, clothing-store, and protein problems through explicit intermediate calculations.These examples show arithmetic reasoning before the final answer, including multiplication, addition, and division.
  • A.1.2 Training Details: The direct-prompting arithmetic examples represent operations with <add>, <subtract>, <multiply>, and <divide> operators.Examples cover fish multiplication, contest-point arithmetic, clothing-store earnings, and protein calculation.
  • A.1.2 Training Details: The appendix includes a concise final-answer format for direct question answering.The specified response format is “So, the answer is xxx.”

A.2.1 Training Details

Toolken embeddings are trained on a small training subset with separate validation data, a specified learning rate, early stopping, and a maximum epoch limit.

  • A.2.1 Training Details: Toolken embeddings use 611 training examples and 39 validation examples.Training uses a learning rate of 1e-4, early stopping on the development set, and at most 20 epochs.

A.2.2 Prompt for Synthetic Training Data

The appendix specifies prompts for synthesizing one-step and multi-hop mathematical training examples, including filtering generated data and formatting tool calls. It also introduces the knowledge-based QA preprocessing context by converting Wikidata relation identifiers into natural-language descriptions.

  • A.2.2 Prompt for Synthetic Training Data: Training-data synthesis begins with two manually crafted examples, uses a ChatGPT prompt to generate more, and filters non-conforming outputs.The generation prompt does not guarantee strict adherence to the required format.
  • A.2.2 Prompt for Synthetic Training Data: The synthetic prompt asks teachers to generate one-step real-world questions and answers using a specified function, with numbers left for teachers to fill.Function descriptions and argument and answer placeholders guide generation.
  • A.2.3 Prompt for FuncQA One-Hop: Zero-shot ChatGPT examples present arithmetic reasoning in natural language and end with the answer in a fixed sentence format.Examples cover addition, subtraction, and division.
  • A.2.3 Prompt for FuncQA One-Hop: CoT and ToolkenGPT one-hop prompts provide worked examples and request answers using explicit mathematical operation tokens such as <power> and <add>.The examples include annual growth, addition, subtraction, and division.
  • A.2.4 Prompt for FuncQA Multi-Hop: Multi-hop prompts ask models to solve problems step by step and represent intermediate operations with tools such as <power>, <choose>, <divide>, and <multiply>.Examples cover coin-toss probability, cube-painting cost, logarithms, and clock-hand motion.
  • A.2.4 Prompt for FuncQA Multi-Hop: The multi-hop ToolkenGPT format uses a question placeholder, stepwise instructions, and final answers separated by ####.Parallel examples show the same calculations without tool-call markup.
  • B Details of Knowledge-based QA: The appendix also frames knowledge-based QA preprocessing as converting Wikidata relation identifiers such as P1346 into natural-language descriptions.This subsection introduces subsequent data synthesis and training settings.

B.1 Getting Text Description

KAMEL creates question-answer pairs from relation-specific templates and sampled facts, then prompts ChatGPT to define an API that answers the resulting questions.

  • B.1 Getting Text Description: KAMEL provides a question template for each relation and randomly samples three dataset facts to instantiate question-answer pairs.The instantiated examples are used as input for generating descriptions.
  • B.1 Getting Text Description: The prompt asks ChatGPT to define an API that can answer questions, illustrated with examples mapping questions and answers to functions such as original_language(title).The examples connect natural-language questions with an API name and argument signature.

B.2 Synthetic Data

Synthetic knowledge-based QA data is aggregated from two prompts, while function-generation examples demonstrate how tool outputs are inserted into generated text.

  • B.2 Synthetic Data: Two prompts are used to synthesize diverse training data, and their samples are aggregated.The passage describes diversity through combining outputs from both prompts.
  • B.2 Synthetic Data: Natural-language tool descriptions are unnecessary for ToolkenGPT but crucial for in-context-learning baselines, especially ICL (desc).ICL (desc) can understand tools only through language descriptions.
  • B.2 Synthetic Data: Function-generation examples instruct the model to continue a sentence with the returned value after a function call, using tools such as star_rating(product).The example embeds the call and returned value in a sentence about an iPhone product.
  • B.2 Synthetic Data: Additional examples use current_weather(city) to insert weather results into natural-language responses.The examples show sunny and cloudy outputs for Beijing and San Diego.

B.3 Training Details

Training and VirtualHome preparation combine bounded toolken optimization with executable, deduplicated task scripts and shared action-generation prompts.

  • B.3 Training Details: Toolken embeddings use a 1e-4 learning rate, development-set early stopping, and a maximum of 5 epochs.These are the stated training controls for embedding optimization.
  • B.3 Training Details: VirtualHome preprocessing removes non-executable or state-neutral scripts, deduplicates same-goal scripts, filters object-name conflicts, and retains frequent verbs and objects.The frequency threshold is more than 10 occurrences.
  • B.3 Training Details: Unlike Huang et al., the preprocessing treats scripts with the same goal but different instructions as distinct tasks.Different instructions can produce different action sequences and final state graphs.
  • B.3 Training Details: All methods in the VirtualHome experiment use the same prompts to generate plans.The prompt section introduces the shared experimental setup.
  • B.3 Training Details: The household robot prompt defines an action vocabulary including navigation, perception, object manipulation, and device-control actions.Listed actions include [WALK], [FIND], [GRAB], [SWITCHON], and [CLOSE].
  • B.3 Training Details: VirtualHome examples pair an initial room, manipulable objects, a goal, and a hint with an action plan, as in writing an email from a bathroom context.The plan uses actions such as [WALK], [FIND], and [TURNTO].
  • B.3 Training Details: Other examples specify computer-use tasks such as turning on a computer and performing work with a mouse and keyboard.The task is set in a home office and includes a goal of Work.
  • B.3 Training Details: A further task example asks the agent to pick up a phone after hearing it ring and running to the living room.The example illustrates a goal and hint paired with an embodied action sequence.

C.3 Training Details

ToolkenGPT training uses early stopping with a ten-epoch limit and was run on LLaMA models with different GPU configurations. The paper focuses on mathematical, knowledge-based, and embodied planning applications and recommends safety consideration for future border applications.

  • Toolken embeddings use a learning rate of 1e-4, early stopping on the development set, and a maximum of 10 epochs.
  • Training and testing use LLaMA-13B on 2 Nvidia RTX 3090 GPUs and LLaMA-33B on 4 Nvidia RTX 3090 GPUs.
  • The paper focuses on mathematical, knowledge-based, and embodied planning applications, while recommending careful safety consideration before future border applications of tool learning.
Loading 2305.11554v4…