CoHyDE: Co-treinamento Iterativo de Reescritor de LLM e Codificador Denso para Recuperação de Ferramentas

Resumo

A recuperação de ferramentas em grandes catálogos de API é um gargalo central para agentes de LLM: as consultas dos usuários chegam em linguagem coloquial, frequentemente subespecificada, enquanto o catálogo utiliza vocabulário técnico de API que nenhum codificador fixo consegue preencher por conta própria. As duas abordagens de treinamento dominantes, ajuste fino contrastivo do codificador e expansão de consulta no estilo HyDE com um LLM congelado, abordam esse problema de extremos opostos e falham em direções complementares: o codificador ajustado por fine-tuning se destaca quando a forma superficial da consulta já corresponde ao catálogo, mas colapsa quando não corresponde, enquanto o HyDE zero-shot é mais robusto a consultas subespecificadas, mas gera descrições hipotéticas alheias ao catálogo que degradam a recuperação quando as consultas são bem formadas. Apresentamos o CoHyDE, um procedimento iterativo que treina o codificador denso e o reescritor de LLM como um único sistema co-evolutivo: o codificador é retreinado com InfoNCE em descrições hipotéticas no estilo do catálogo produzidas pelo reescritor, e o reescritor é alinhado por preferência via DPO com base nas pontuações de recuperação do codificador, com ambos os lados pré-inicializados no catálogo de ferramentas antes do início do loop. Em um subconjunto de ~10 mil ferramentas do catálogo ToolBench, três rodadas do CoHyDE superam a linha de base de componente único mais forte em +2,5 p.p. de NDCG@5 em consultas padrão e +6,3 p.p. em consultas vagas retidas, com ganhos de até +8 p.p. no nível de vagueza mais difícil. Ablações confirmam que o co-treinamento é o ingrediente chave: usar qualquer um dos componentes isoladamente não consegue igualar o CoHyDE tanto em consultas bem formadas quanto em consultas vagas, com perdas de até -8 p.p. em consultas vagas.

English

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.