ToolWeaver: Weaving Collaborative Semantics for Scalable Tool Use in Large Language Models

AI summary

Overview

• Research area: Tool-augmented large language models, specifically how an LLM represents, selects, and executes external APIs at large scale.
• Technical level: Advanced. The paper assumes familiarity with vector quantization (RQ-VAE), residual codebooks, optimal transport, autoregressive fine-tuning, and retrieval metrics such as NDCG.
• Scope: The paper proposes a generative tool-use framework that replaces one-token-per-tool identifiers with hierarchical, collaborative-aware code sequences, and evaluates it on the ToolBench benchmark of 46,985 APIs.

What This Paper Is About

Existing generative tool-use systems map each tool to a single new token in the LLM vocabulary. On a benchmark like ToolBench, this means adding nearly 47,000 tokens to a model such as Llama-3-8B, whose vocabulary is 128,256 — a linear expansion that inflates memory, risks damaging pretrained language ability, and gives the model no shared structure for learning which tools work together.

ToolWeaver instead represents each tool as a short sequence of hierarchical codes drawn from multiple codebooks. The goal is to make vocabulary growth logarithmic in the number of tools, and to let functionally related tools share codes so the model can learn collaborative patterns from dense code co-occurrence rather than rare co-occurrence of isolated tool IDs.

Key Contributions

1. Compositional tool codes. A framework that represents each tool as a hierarchical sequence of discrete codes. With L codebooks of K code vectors each, the representation capacity is up to K^L tools while only requiring L × K new vocabulary tokens.
2. Collaborative-aware tokenization. An unsupervised structured tokenization process built on RQ-VAE that weaves a tool's intrinsic semantics (from its name and description embedding) together with its extrinsic co-usage patterns (from a tool-tool co-occurrence similarity matrix), plus a uniform-mapping conflict mitigation step so distinct tools do not collide on the same code sequence.
3. Multi-stage generative alignment. A two-stage fine-tuning procedure — tool retrieval alignment on query-tool pairs, then tool usage trajectory alignment on full interaction trajectories — that teaches the LLM to natively generate the hierarchical code sequences.
4. Large-scale evaluation. Experiments on ToolBench (over 16,000 tool collections comprising 46,985 unique APIs) showing gains over retrieval-based and generative baselines, while better preserving the base model's general language capabilities.

Main Findings

• Tool retrieval (Table 1): ToolWeaver achieves the highest NDCG across all three query complexities. On the hardest I3 scenario it reaches NDCG@1 of 88.00, NDCG@3 of 85.80, and NDCG@5 of 90.12, versus ToolGen's 81.00, 80.83, and 85.83. On I1 it scores 91.16 / 91.14 / 93.48 and on I2 89.76 / 89.70 / 91.80.
• End-to-end completion (Table 2): ToolWeaver records the largest advantage in the multi-tool I3 setting, with a Solvable Pass Rate of 52.19 compared to ToolGen's 36.34. Its I3 Solvable Win Rate against the GPT-4o-mini reference is 59.02 versus ToolGen's 49.18. On I1 it reports SoPR 53.17 and SoWR 40.49; on I2, SoPR 44.03 and SoWR 48.11. Generalization splits are also competitive: I1-Tool SoPR 54.85, I1-Cat SoPR 57.41, I2-Cat SoPR 46.24.
• Baseline comparison: GPT-4o-mini records SoPR of 52.25 (I1), 40.41 (I2), and 24.86 (I3). ToolLlama-2 records SoPR of 28.94, 24.69, and 10.93 across I1–I3, and SoWR of 25.15, 30.19, and 24.59.
• Collaborative weight sensitivity (Figure 2): Average NDCG@k across I1–I3 improves as the collaborative regularization weight λ rises from 0.01 to 1, peaks at λ = 1, and declines when λ is raised further to 10 — indicating that an overly strong collaborative prior distorts tool representations.
• Component ablation (Figure 3): Adding semantic initialization alone produces a leap of over 20 NDCG points over the baseline lacking it. Adding collaborative guidance on top gives a further gain whose magnitude grows with task complexity — modest for I1 and most pronounced for I3.
• Tokenization strategy comparison (Figure 4): Compared against Atomic, Numerical, Hierarchical, and Semantic alternatives, ToolWeaver wins on both retrieval and end-to-end evaluation. The Hierarchical and Semantic strategies struggle to beat the Atomic baseline, showing that naive structure alone is insufficient.
• Preservation of language ability (Table 3): The base Llama-3-8B scores 6.34 perplexity on WikiText-2, with BERTScore F1 of 85.35 (CNN/DM) and 85.05 (XSum). ToolGen degrades to 104.54 perplexity (roughly 16 times the base model), 82.93, and 82.53, with an average drop of 2.47. ToolWeaver stays at 25.36 perplexity, 85.07, and 84.18, with an average drop of 0.57.

Methodology in Plain English

The pipeline has three parts.

Step 1 — Turn tool documentation into embeddings. Each tool's name and description are passed through a pretrained text encoder to produce a dense semantic vector.

Step 2 — Quantize those vectors into hierarchical codes. A Residual-Quantized VAE with L codebooks of K centroid vectors each compresses a tool's embedding into a sequence of L code indices. At each level the closest centroid is selected and subtracted, leaving a residual for the next level. Two extra ingredients shape the codes: a graph Laplacian regularization term derived from a tool-tool co-occurrence matrix (build via cosine similarity on co-occurrence counts) pulls frequently co-used tools toward nearby representations, and a final-layer uniform assignment constraint — solved as an optimal transport problem with the Sinkhorn-Knopp algorithm — spreads tools evenly across the last codebook so no two tools collide on the same identifier. The overall tokenization loss combines reconstruction, quantization, and collaborative terms, with a commitment weight β typically set to 0.25.

Step 3 — Teach the LLM to speak the codes. Each code sequence becomes a short string of new tokens (for example <T1_1><T2_1>) added to the vocabulary with randomly initialized embeddings. The model is fine-tuned first on query-tool pairs to generate the right code for a query, then on full interaction trajectories covering reasoning, tool calls with parameters, and final answers, with loss computed only over the assistant's tokens. At inference, a constrained beam search guided by a precomputed prefix tree (trie) of valid code sequences masks invalid next tokens during tool selection only.

Why This Matters

• Impact on research: The paper reframes tool selection as a tokenization design problem rather than purely a retrieval or fine-tuning problem. It argues that prior collaborative-semantics work (for example LC-Rec, SeLLa-Rec, and CFRAG) relies on post-hoc alignment, introducing semantically isolated identifiers and then teaching the model their meaning separately; ToolWeaver instead bakes collaborative structure into the tokenization step itself.
• Real-world applications:
  • Travel and trip-planning agents that must combine weather, air quality, flight, hotel, and map APIs within one orchestration.
  • Customer-service or enterprise assistants that need to route queries across thousands of internal and third-party REST APIs without context-window limits.
  • Developer tooling that recommends and chains APIs from large public hubs such as RapidAPI.
  • General-purpose agent platforms where new APIs are added continuously and cannot be retrained into the vocabulary one token at a time.
• Industry relevance: The linear-versus-logarithmic vocabulary argument is directly about deployment cost. Adding tens of thousands of tokens raises memory and latency, and Table 3 quantifies the collateral damage to a model's general language behavior — a practical concern for any team that fine-tunes a general-purpose model for tool use and still needs it to write and reason normally.

Future Directions

• Reinforcement learning for collaborative discovery. The conclusion names this explicitly as a future direction: letting the model autonomously discover collaborative patterns rather than deriving them from static co-occurrence statistics.
• Open question — where the λ trade-off settles at other scales. The sensitivity analysis peaks at λ = 1 for this benchmark; whether that optimum transfers to other tool corpora or model sizes is not established in the main text.
• Open question — generality across base models. The main results use Llama-3-8B, with supplementary results on the Qwen model series placed in Appendix B.3 (not included in the provided text), so the breadth of architecture transfer is only partially reported here.
• Open question — deployment safety. The ethics statement notes that more capable agents could be directed toward malicious APIs and that ToolBench was not audited for biases or privacy risks; it frames any real-world deployment as requiring robust safety protocols and human oversight.

Target Audience

Researchers and engineers working on tool-augmented LLM agents, generative retrieval, and recommendation-style semantic ID design will get the most from this paper. It is also relevant to practitioners who fine-tune LLMs against large API catalogs and need to weigh vocabulary expansion costs against task performance, and to readers interested in vector-quantization tokenization methods that bridge collaborative signals and language models.

Authors’ abstract

Prevalent retrieval-based tool-use pipelines struggle with a dual semantic challenge: their retrievers often employ encoders that fail to capture complex semantics, while the Large Language Model (LLM) itself lacks intrinsic tool knowledge from its natural language pretraining. Generative methods offer a powerful alternative by unifying selection and execution, tasking the LLM to directly learn and generate tool identifiers. However, the common practice of mapping each tool to a unique new token introduces substantial limitations: it creates a scalability and generalization crisis, as the vocabulary size explodes and each tool is assigned a semantically isolated token. This approach also creates a semantic bottleneck that hinders the learning of collaborative tool relationships, as the model must infer them from sparse co-occurrences of monolithic tool IDs within a vast library. To address these limitations, we propose ToolWeaver, a novel generative tool learning framework that encodes tools into hierarchical sequences. This approach makes vocabulary expansion logarithmic to the number of tools. Crucially, it enables the model to learn collaborative patterns from the dense co-occurrence of shared codes, rather than the sparse co-occurrence of monolithic tool IDs. We generate these structured codes through a novel tokenization process designed to weave together a tool's intrinsic semantics with its extrinsic co-usage patterns. These structured codes are then integrated into the LLM through a generative alignment stage, where the model is fine-tuned to produce the hierarchical code sequences. Evaluation results with nearly 47,000 tools show that ToolWeaver significantly outperforms state-of-the-art methods, establishing a more scalable, generalizable, and semantically-aware foundation for advanced tool-augmented agents.