MassTool: A Multi-Task Search-Based Tool Retrieval Framework for Large Language Models

2025-07-01

AI Summary

Existing tool retrieval methods prioritize tool representation learning while neglecting precise query understanding, resulting in poor robustness to out-of-distribution queries. To address this, we propose MassTool—a multi-task, search-enhanced tool retrieval framework tailored for large language models (LLMs). Its core innovation lies in a dual-encoder architecture that jointly optimizes query understanding and tool matching, incorporating a Query-Centric Graph Convolutional Network (QC-GCN) to model tool semantic structure, Search-Enhanced Intent Modeling (SUIM), and Adaptive Knowledge Transfer (AdaKT) to realize a two-stage decision process. MassTool integrates contrastive regularization, listwise ranking loss, and search-log-driven user intent modeling. Extensive experiments demonstrate that MassTool significantly improves retrieval accuracy across multiple benchmarks, substantially enhancing the reliability and generalization capability of LLM-based tool invocation.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Querying, indexing, and retrieval in Web-related graphs

Abstract

Tool retrieval is a critical component in enabling large language models (LLMs) to interact effectively with external tools. It aims to precisely filter the massive tools into a small set of candidates for the downstream tool-augmented LLMs. However, most existing approaches primarily focus on optimizing tool representations, often neglecting the importance of precise query comprehension. To address this gap, we introduce MassTool, a multi-task search-based framework designed to enhance both query representation and tool retrieval accuracy. MassTool employs a two-tower architecture: a tool usage detection tower that predicts the need for function calls, and a tool retrieval tower that leverages a query-centric graph convolution network (QC-GCN) for effective query-tool matching. It also incorporates search-based user intent modeling (SUIM) to handle diverse and out-of-distribution queries, alongside an adaptive knowledge transfer (AdaKT) module for efficient multi-task learning. By jointly optimizing tool usage detection loss, list-wise retrieval loss, and contrastive regularization loss, MassTool establishes a robust dual-step sequential decision-making pipeline for precise query understanding. Extensive experiments demonstrate its effectiveness in improving retrieval accuracy. Our code is available at https://github.com/wxydada/MassTool.

Problem

• Enhances query representation for precise tool retrieval
• Improves tool retrieval accuracy via multi-task framework
• Handles diverse queries with search-based intent modeling

Innovation

• Two-tower architecture for query and tool matching
• Search-based user intent modeling for diverse queries
• Adaptive knowledge transfer for multi-task learning