Re-Initialization Token Learning for Tool-Augmented Large Language Models

TL;DR

This paper introduces a novel re-initialization token learning method that aligns tool tokens with word embeddings, significantly improving tool integration and problem-solving in large language models across various tasks.

Contribution

The proposed method enhances tool token learning by aligning tool embeddings with word space, leading to better tool call accuracy and improved performance on complex reasoning tasks.

Findings

Improved tool call accuracy across tasks

Enhanced performance on numerical reasoning and QA

Outperforms recent baselines like CoT and REACT

Abstract

Large language models have demonstrated exceptional performance, yet struggle with complex tasks such as numerical reasoning, plan generation. Integrating external tools, such as calculators and databases, into large language models (LLMs) is crucial for enhancing problem-solving capabilities. Current methods assign a unique token to each tool, enabling LLMs to call tools through token prediction-similar to word generation. However, this approach fails to account for the relationship between tool and word tokens, limiting adaptability within pre-trained LLMs. To address this issue, we propose a novel token learning method that aligns tool tokens with the existing word embedding space from the perspective of initialization, thereby enhancing model performance. We begin by constructing prior token embeddings for each tool based on the tool's name or description, which are used to initialize and regularize the learnable tool token embeddings. This ensures the learned embeddings are well-aligned with the word token space, improving tool call accuracy. We evaluate the method on tasks such as numerical reasoning, knowledge-based question answering, and embodied plan generation using GSM8K-XL, FuncQA, KAMEL, and VirtualHome datasets. The results demonstrate clear improvements over recent baselines, including CoT, REACT, ICL, and ToolkenGPT, indicating that our approach effectively augments LLMs with tools through relevant tokens across diverse domains.