CoHyDE: Iterative Co-Training of LLM Rewriter & Dense Encoder for Tool Retrieval

2026-05-27

AI Summary

This work addresses the challenge that large language models (LLMs) struggle to effectively retrieve structured API tools when faced with ambiguous or informal user queries. The authors propose the first unified iterative co-training framework that jointly optimizes an LLM-based rewriter and a dense retrieval encoder. The rewriter reformulates user queries into normalized forms aligned with tool description styles, while the encoder performs InfoNCE-based contrastive learning using these rewritten queries. Simultaneously, the rewriter is fine-tuned via Direct Preference Optimization (DPO), leveraging retrieval scores for preference alignment. Evaluated on a 10k-tool subset of ToolBench, the method significantly improves retrieval performance—boosting NDCG@5 by 2.5% on standard queries, 6.3% on ambiguous queries, and up to 8.0% on the most challenging ambiguous cases—demonstrating its effectiveness in handling both canonical and fuzzy query scenarios and overcoming the limitations of single-component approaches.

Abstract

Tool retrieval over large API catalogs is a core bottleneck for LLM agents: user queries arrive in colloquial, often underspecified language, while the catalog uses technical API vocabulary that no fixed encoder can bridge on its own. The two dominant training approaches, contrastive encoder fine-tuning and HyDE-style query expansion with a frozen LLM, address this problem from opposite ends and fail in complementary directions: the fine-tuned encoder excels when the query's surface form already matches the catalog but collapses when it does not, while zero-shot HyDE is more robust to underspecified queries yet generates catalog-unaware hypothetical descriptions that degrade retrieval when queries are well-formed. We introduce CoHyDE, an iterative procedure that trains the dense encoder and the LLM rewriter as a single co-evolving system: the encoder is retrained with InfoNCE on catalog-style hypothetical descriptions produced by the rewriter, and the rewriter is preference-aligned via DPO against the encoder's retrieval scores, with both sides warm-started on the tool catalog before the loop begins. On a ~10k tool subset of the ToolBench catalog, three rounds of CoHyDE improve over the strongest single-component baseline by +2.5 pp NDCG@5 on standard queries and +6.3 pp on held-out vague queries, with gains as large as +8 pp on the hardest vague tier. Ablations confirm that co-training is the key ingredient: using either component in isolation fails to match CoHyDE on both well-formed and vague queries, with losses of up to -8 pp on vague queries.

Problem

• tool retrieval
• LLM agents
• query underspecification
• API catalog
• semantic gap

Innovation

• co-training
• tool retrieval
• LLM rewriter
• dense encoder
• query expansion