Towards Completeness-Oriented Tool Retrieval for Large Language Models

TL;DR

COLT tackles the completeness challenge in tool retrieval for large language models by introducing a two-stage, model-agnostic framework that first aligns query-tool semantics and then captures high-order tool collaborations via a dual-view, graph-based learning process. It constructs three bipartite graphs among queries, scenes, and tools and employs LightGCN-based propagation with cross-view contrastive losses, culminating in a final score that combines scene-centric and tool-centric similarities. The authors also introduce ToolLens, a dataset designed for multi-tool scenarios, and COMP@K, a metric for completeness in top-K tool retrieval, both of which enable rigorous evaluation of multi-tool usage in LLMs. Empirical results show COLT consistently outperforms state-of-the-art dense retrievers on ToolLens and ToolBench and improves downstream tool learning, even when paired with compact PLMs like BERT-mini. The work provides practical significance by offering a robust retrieval framework and new resources to advance tool-augmented LLM research.

Abstract

Recently, integrating external tools with Large Language Models (LLMs) has gained significant attention as an effective strategy to mitigate the limitations inherent in their pre-training data. However, real-world systems often incorporate a wide array of tools, making it impractical to input all tools into LLMs due to length limitations and latency constraints. Therefore, to fully exploit the potential of tool-augmented LLMs, it is crucial to develop an effective tool retrieval system. Existing tool retrieval methods primarily focus on semantic matching between user queries and tool descriptions, frequently leading to the retrieval of redundant, similar tools. Consequently, these methods fail to provide a complete set of diverse tools necessary for addressing the multifaceted problems encountered by LLMs. In this paper, we propose a novel modelagnostic COllaborative Learning-based Tool Retrieval approach, COLT, which captures not only the semantic similarities between user queries and tool descriptions but also takes into account the collaborative information of tools. Specifically, we first fine-tune the PLM-based retrieval models to capture the semantic relationships between queries and tools in the semantic learning stage. Subsequently, we construct three bipartite graphs among queries, scenes, and tools and introduce a dual-view graph collaborative learning framework to capture the intricate collaborative relationships among tools during the collaborative learning stage. Extensive experiments on both the open benchmark and the newly introduced ToolLens dataset show that COLT achieves superior performance. Notably, the performance of BERT-mini (11M) with our proposed model framework outperforms BERT-large (340M), which has 30 times more parameters. Furthermore, we will release ToolLens publicly to facilitate future research on tool retrieval.

Figures (6)

• Figure 1: An illustration of tool retrieval for LLMs with tool learning and varied responses using different tools.
• Figure 2: The architecture of the proposed two-stage learning framework COLT for tool retrieval.
• Figure 3: An overview of the dataset construction pipeline of ToolLens. Human verification is included at each step.
• Figure 4: Comparison of different model sizes of PLM.
• Figure 5: Performance comparison regarding different sizes of ground-truth tool sets.