PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Vaishali Senthil, Ashutosh Hathidara, Sebastian Schreiber

arXiv:2605.29271v1cs.AIcs.IRcs.LG

TL;DR

Tool retrieval must bridge colloquial or vague user queries and technical API catalogs, but encoder fine-tuning and frozen-LLM HyDE each fail in complementary settings. CoHyDE co-trains an encoder and rewriter iteratively, improving NDCG@5 over the strongest single-component baseline by +2.5 pp on standard queries and +6.3 pp on vague queries.

  • Problem

    Tool retrieval lacks a single training approach that reliably bridges underspecified user language and technical API-catalog vocabulary.

  • Method

    CoHyDE iteratively co-trains a dense encoder and LLM rewriter, aligning rewriter outputs with encoder retrieval scores and retraining the encoder on catalog-style descriptions.

  • Results

    +2.5 pp NDCG@5 on standard queries and +6.3 pp on held-out vague queries over the strongest single-component baseline after three CoHyDE rounds.

  • Takeaways & Limitations

    Jointly evolving the encoder and rewriter outperforms either single component across standard and vague tool-retrieval queries.

  • Takeaways & Limitations

    Evaluation uses one training seed and a 10K-tool English ToolBench subset, so transfer to enterprise catalogs, non-English queries, and description-poor schemas remains untested.

Abstract

from arXiv · show

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.

1 Introduction

Tool retrieval limits language-model agents operating over catalogs with tens of thousands of APIs, because context windows cannot contain every tool’s documentation and retrieval selects the candidate set for each query. CoHyDE addresses complementary weaknesses in encoder fine-tuning and frozen-LLM query expansion by iteratively co-training a dense encoder and LLM rewriter, improving retrieval on standard and vague queries.

  • Motivation: Tool catalogs increasingly contain tens of thousands of APIs, making retrieval an upstream bottleneck that bounds agent action quality.No agent can fit every tool’s documentation into its context window, so retrieval must select a small candidate set per query.
  • Motivation: Encoder fine-tuning handles in-distribution queries but falters when surface form drifts, while HyDE-style expansion is catalog-unaware for well-formed queries.The two approaches therefore exhibit complementary failure modes: encoders rely on residual lexical cues under drift, whereas query expansion can degrade retrieval for already well-formed queries.
  • Method: CoHyDE jointly evolves the dense encoder and LLM rewriter through repeated alternating updates for tool retrieval.The LLM generates catalog-style hypothetical descriptions, the encoder receives contrastive retraining, and the LLM is preference-aligned via DPO using encoder retrieval scores.
  • Results: +2.5 pp NDCG@5 is CoHyDE’s improvement over the strongest single-component baseline on standard queries after three co-training rounds.The evaluation uses a ∼10k-tool subset of the ToolBench catalog.
  • Results: +6.3 pp NDCG@5 is the improvement on held-out vague queries, reaching as high as +8 pp on the hardest vague tier.These gains are reported after three rounds of co-training over the strongest single-component baseline.

2 Related Work

Prior tool-retrieval work either fine-tunes dense encoders or uses frozen-index LLM expansion, while related rewriter and retriever-generator methods freeze at least one component. CoHyDE instead co-trains the rewriter and encoder, targeting catalog-description style and aligning the rewriter with encoder feedback.

  • Tool retrieval: Tool-retrieval research largely splits between contrastively fine-tuned dense encoders and frozen black-box methods using LLM expansion or generative indexing.These approaches are represented by encoder training on query–tool pairs and by expansion or indexing with a frozen retrieval system.
  • Tool retrieval: CoHyDE differs from iterative rewriting by preference-aligning its rewriter via DPO, targeting catalog-description style, and retraining the encoder without real query–tool pairs.The closest prior work iteratively rewrites user instructions and retrains the encoder on rewritten-instruction–tool pairs.
  • Query expansion and trained rewriters: HyDE and Query2doc expand queries with hypothetical documents, while CSQE corrects corpus misalignment at test time by injecting retrieved sentences.CoHyDE addresses the same corpus-misalignment issue during training, and related trained rewriters use RL or DPO with frozen retrievers.
  • Dense retriever robustness and joint retriever-generator training: Synthetic-query domain adaptation commonly runs generation once with a frozen generator, whereas CoHyDE co-trains both retrieval components rather than freezing one.The related work also notes that dense retrievers are brittle off-distribution and surveys joint retriever–generator frameworks.

3 Methodology

CoHyDE formulates tool retrieval as dense nearest-neighbor search and co-trains an encoder with an LLM rewriter through alternating updates. Both components are warm-started on catalog data before rewriter-generated descriptions drive encoder retraining and encoder scores supervise rewriter alignment.

  • Problem formulation: Tool retrieval returns a ranked k-tool subset maximizing overlap with the gold tool set, using single-vector dense encoder retrieval over rendered catalog records.The encoder maps text to unit-norm vectors, and approximate nearest-neighbor search performs retrieval.
  • Warmup: The rewriter is catalog-warm-started across five tool renderings, teaching catalog vocabulary, naming conventions, and surface-form variation.The rendering family ranges from title-only to the full tool record.
  • Warmup: The encoder warmup uses standard InfoNCE training on query-tool pairs, providing a contrastive head start before description-only retraining.This initialization is reported as the strongest encoder-only baseline.
  • Alternating co-training: Each round retrains the encoder on rewriter-generated descriptions paired with rendered tools, excluding real query-tool pairs from this stage.The generated descriptions become contrastive anchors for encoder training.
  • Alternating co-training: The encoder then scores sampled rewriter descriptions by NDCG@5, and DPO trains the rewriter from the highest- and lowest-scoring candidates.The updated rewriter regenerates training data for the next encoder round, creating a coupled trajectory.

4 Experiments & Results

CoHyDE is evaluated on standard and vague ToolBench-derived queries against sparse, dense, expansion, HyDE, and iterative-feedback baselines. Results show complementary roles for the jointly trained encoder and rewriter, with co-training improving retrieval robustness and avoiding failures from independently composing the components.

  • Main results: −39.5pp: BGE S1a’s G1 vague performance falls relative to its standard counterpart, while G3 vague reaches 14.9%.The encoder dominates standard splits but is brittle when paraphrases replace catalog-matching surface vocabulary.
  • Baseline comparisons: 6.2%: query expansion reaches only G3 vague NDCG@5, below vanilla BGE’s 8.3%, whereas HyDE’s catalog-style descriptions bridge the vocabulary gap more effectively.Both baselines use the same vanilla BGE encoder and Qwen3.5-4B generator; they differ in whether generation stays query-side or moves toward document-side descriptions.
  • Baseline comparisons: −10.8pp: combining vanilla-LLM HyDE with trained BGE S1a degrades standard NDCG because the encoder is calibrated on raw-query anchors rather than hypothetical descriptions.This motivates jointly evolving the encoder and rewriter instead of composing independently trained components.
  • Main results: CoHyDE improves over BGE S1a by average +2.5pp NDCG@5 on standard queries and +6.3pp on vague queries.At r=3, gains occur across every evaluation split.
  • Ablations: 23.6pp, 35.5pp, and 23.9pp: pairing the co-trained rewriter with vanilla BGE reduces standard G1, G2, and G3 NDCG@5 by those amounts.Vague splits also decline by 9–21pp, confirming that the rewriter requires a jointly trained encoder.
  • Ablations: 0.5pp: replacing the trained rewriter with a vanilla rewriter lowers G1 standard NDCG@5 from 86.8 to 86.3, while G2 standard rises 2.0pp from 73.6 to 75.6.The encoder carries precision on easy standard queries, while the rewriter contributes robustness on hard and vague queries.
  • Round evolution: Monotonically non-decreasing: performance improves from S1 through R3 on five of six splits, with only standard G2 retreating 0.6pp between R2 and R3.R1-to-R2 gains exceed R2-to-R3 gains across all tiers and split types.
  • Comparison with iterative baselines: 7.5pp: CoHyDE exceeds RaFe on G3 vague NDCG@5, scoring 21.1 versus 13.6 on the hardest cross-domain vague split.CoHyDE also trails RaFe by 6.9pp, 6.4pp, and 6.4pp on standard G1, G2, and G3 comparisons, respectively, according to the cited passage.

5 Conclusion

CoHyDE resolves the complementary failures of contrastive encoder fine-tuning and HyDE-style generation by iteratively co-training the encoder and rewriter together. The encoder’s retrieval scores supervise rewriter alignment, while catalog-aligned rewriter outputs anchor encoder training.

  • CoHyDE uses iterative co-training to evolve the encoder and rewriter together, addressing their complementary failure modes.Naively composing the methods worsens performance because their representation spaces are calibrated to different input distributions.
  • The encoder’s NDCG@5 scores supervise rewriter training through DPO.
  • The rewriter’s catalog-aligned outputs provide training anchors for the encoder.

Limitations

The evaluation is limited by single-seed reporting and a narrow dataset: experiments use a 10K-tool English ToolBench subset skewed toward consumer-facing RapidAPI REST endpoints, so transfer to enterprise catalogs remains uncertain.

  • Evaluation limitations: All reported numbers use a single training seed, while bootstrap confidence intervals capture evaluation-set rather than training-side variance.Multi-seed retrains were not run because each co-training round is computationally costly.
  • Evaluation limitations: Experiments use a 10K-tool English subset of ToolBench skewed toward consumer-facing RapidAPI REST endpoints.
  • Evaluation limitations: Transfer of the co-training gains to enterprise catalogs remains untested.

Ethical Considerations … E Per-Stage Hyperparameter Summary

The paper reports compliance with ethics and research-integrity guidance, validates its vague-query split through LLM checks and human spot-checking, and specifies preprocessing, prompting, and hyperparameter procedures across the pipeline.

  • Ethical Considerations: The experiments followed the ACL Ethics Policy and relevant research-integrity guidelines, with no remaining unaddressed ethical risks reported.
  • A Vague-Query Construction and Validation: The vague-query split is a held-out paraphrase of the official 1,092-query evaluation set, preserving gold tools while testing query-side distribution shift.Construction follows Chen et al. (2026), substituting claude-4.5-opus for GPT-4o.
  • A Vague-Query Construction and Validation: 50 paraphrases passed human verification on all three criteria, implying a 6% rule-of-three upper bound on the true failure rate at 95% confidence.The annotator was not blinded to paraphraser identity and was not separately compensated beyond regular wages.
  • B Cleaning Operator: The cleaning operator strips reasoning traces, rejects unclosed traces by reverting to the original query, removes conversational preambles, and normalizes whitespace.It is applied identically during SFT-target construction, DPO-candidate scoring, and inference.
  • C HyDE-Style Rewriter Prompt: The HyDE-style prompt describes ordered API-tool pipelines, focusing on tool functions, inputs, outputs, and how each step feeds the next.It is used in optional SFT, S2 generation, S4 DPO-candidate sampling, and inference-time HyDE evaluations.
  • D Query-Rewriting Prompt: The query-rewriting prompt appears only in the Appendix L prompt-style ablation and is not used in CoHyDE.Its instructions request more specific natural-language queries containing relevant tool names, API capabilities, and technical terms.
  • E Per-Stage Hyperparameter Summary: Table 3 consolidates load-bearing hyperparameters for every main training and inference stage, while appendices provide objectives, optimizers, schedules, ablation context, optional SFT details, and software versions.The encoder details are in Appendix F, rewriter details in Appendix G, optional HyDE-style SFT in Appendix H, and software versions in Appendix O.

F Encoder Training Hyperparameters … K Bootstrap CI Protocol

The appendices specify the encoder and rewriter training procedures, explain why the main pipeline omits an optional HyDE-style SFT pass, define evaluation reporting, and formalize bootstrap confidence intervals. Encoder retraining uses rewriter-generated descriptions, while round-3 evaluation and uncertainty estimation follow explicit protocols.

  • F Encoder Training Hyperparameters: The encoder trains with InfoNCE during S1a warmup and once per round at S3r, differing in anchor source and checkpoint initialization.S1a uses real query–tool pairs, whereas S3r uses rewriter-generated descriptions and continues from the previous round.
  • F Encoder Training Hyperparameters: S3r trains purely on rewritten-description–tool pairs, retains the best validation-NDCG@5 checkpoint, and tests the encoder on real queries at inference.No real query–tool pair is used during S3r; the retrain uses the same core optimization settings as S1a.
  • G Rewriter Training and Inference Hyperparameters: DPO selects preferred and rejected descriptions using freshly retrained encoder NDCG@5 scores, drops ties, and merges each trained LoRA adapter before the next round.The previous round’s merged rewriter serves as the DPO reference policy and next bootstrap generator.
  • H Optional SFT Stage (HyDE-Style Bridging): The optional HyDE-style SFT pass is excluded from the main pipeline because the five-format encoder and rewriter warmups provide catalog-style alignment directly.The reported variant with optional SFT does not improve over the main pipeline.
  • I Evaluation Metrics: Evaluation metrics are averaged over queries in each relevant tier, with definitions matching standard ir_measures implementations.The evaluation uses gold tool sets and retrieved rankings to define the reported metrics.
  • J Round-3 k-Sweep: Round-3 evaluation reports hit@k, recall@k, and NDCG@k for k ∈ {1, 5, 10, 20} across standard and vague splits stratified by tier.NDCG@1 equals hit@1 by construction, while Recall@1 is capped by 1/|T*_q| for multi-tool queries.

L Design-Choice Ablations: Details

The ablations show that CoHyDE’s gains depend on format diversity, description-only encoder retraining, catalog-style rewriting, and diverse preference candidates. Several seemingly helpful alternatives instead reduce robustness or closed-loop improvement, especially on vague queries.

  • Format diversity: Single-format encoder training matched its corresponding rendering but underperformed on other renderings, while the 5-format encoder remained more robust to varying rewriter-output lengths.Even training only on ϕ5, the rendering closest in length to rewriter outputs, was less robust than the 5-format mixture.
  • Encoder retraining objective: Combined q+d retraining produced no improvement over description-only training and slightly degraded vague-query performance by pulling the encoder toward the earlier q-anchored fixed point.The description-only objective intentionally induces distribution shift away from the on-distribution fixed point established at S1.
  • Rewrite prompt: Replacing catalog-style ρHyDE with user-style ρrewrite at any loop stage lost on every standard metric, with the largest gap on cross-domain G3.Even with relevant tool names as anchors, ρrewrite generated plausible user queries that mismatched the contrastive-pair style used during S3 retraining.
  • SFT schedule: 1,000 / 5,000 SFT steps, versus 100, closed the SFT train loss but reduced candidate diversity, DPO margin, and the final closed-loop gain.Over-fitted rewriters concentrate temperature-0.7 sampling mass, weakening the preference signal.
  • HyDE-concat: HyDE-concat helped slightly on G1 standard but hurt vague queries by retaining the lexical surface that the original query is intended to escape.The method concatenates q with the hypothetical description before encoding, following the Query2doc approach.

M Xu et al. 2024 Re-implementation

The reimplementation follows Shao et al.’s three-stage iterative retrieval pipeline for a head-to-head comparison, while using stronger and different components than the original setup. It should therefore be interpreted as faithful to the pipeline structure, not an exact prompt reproduction, and as charitable to Xu’s method.

  • Encoder: The encoder is trained once contrastively and instantiated with the stronger S1a InfoNCE BGE-large-en-v1.5 checkpoint rather than Sentence-BERT base.It is trained on real (q, ϕ5(t)) pairs and never updated again at inference time.
  • LLM refiner: The reimplementation uses Qwen3.5-4B via vLLM with greedy decoding, a 400-token stage limit, and the same temperature across all three prompted stages.This deliberately matches the main paper’s rewriter setup, whereas Shao et al. use GPT-3.5.
  • Iteration schedule: T = 3 refinement rounds retrieve K = 10 tools per round, with final ranking from the last round and Top-50 saved for evaluation at k ∈{1, 5, 10, 20}.Each round uses Comprehension, Assessment, and Refinement prompts; refinement or N/A becomes the next retrieval input.
  • Three-stage prompts: Comprehension summarizes goals and retrieved tools, Assessment labels goals SOLVED or UNSOLVED and checks ranking importance, and Refinement outputs N/A or an enriched instruction.The prompts approximate the three-stage description because the original paper does not provide verbatim prompt text.
  • Caveat: The head-to-head is charitable to Xu’s method because it substitutes both a stronger encoder and a different LLM, and is not an exact reproduction of Xu’s prompts.The stated substitutions are S1a InfoNCE BGE-large versus Sentence-BERT base and Qwen3.5-4B versus GPT-3.5.

N Compute Budget and Infrastructure … Q.1 Upstream Artifacts and Licenses

The implementation runs on a single 8-H200 node, with each major training and inference stage fitting on one GPU, while three-round reproduction requires roughly 50 GPU-hours and the full study 400–500 GPU-hours. The paper specifies its software stack, reports single-seed results with bootstrap uncertainty estimates, and identifies the upstream datasets, catalog, encoder, and rewriter artifacts and licenses.

  • N Compute Budget and Infrastructure: All reported experiments ran on a single node with 8 H200 GPUs, and encoder training, rewriter SFT/DPO, and HyDE inference each fit on one GPU.Multi-GPU parallelism was opportunistic rather than required.
  • N Compute Budget and Infrastructure: ∼3 hours covers S1a encoder InfoNCE warmup, while each round costs ∼2 hours for S2 regeneration, ∼1.5 hours for S3 retraining, ∼6 hours for S4 data generation, and ∼4 hours for S4 DPO training.The S4 data-generation stage uses N = 4 candidates per query, and S4 training uses approximately 4,371 steps with LoRA r = 64.
  • N Compute Budget and Infrastructure: 400–500 GPU-hours cover three co-training rounds plus baselines, ablations, and rejected variants, whereas reproducing the main result requires roughly 50 GPU-hours.The hardware is H200-class.
  • O Software Versions: The software stack uses PyTorch 2.4, Transformers 4.46, TRL 0.11, PEFT 0.13, vLLM 0.6, native PyTorch bf16, and in-house retrieval scoring.Metric definitions match standard ir_measures implementations and are provided in closed form in Appendix I.
  • P Single-Seed Caveat: All reported numbers use a single training seed, without multi-seed variance estimates because of the per-round compute cost.The per-round trajectory is described as a partial stability proxy, with smooth tier-averaged behavior and mildly non-monotonic per-cell behavior.
  • P Single-Seed Caveat: ±2pp, ±3pp, and ±5–6pp are the approximate paired-bootstrap 95% CI halfwidths for NDCG@5 on G1, G2, and G3, respectively.These intervals capture evaluation-set sampling uncertainty rather than training-seed variance; cell differences should be compared with bootstrap CIs of differences.
  • Q.1 Upstream Artifacts and Licenses: ToolBench supplies the API pool and official G1/G2/G3 queries under Apache 2.0, while ToolGen supplies the 46,980-tool catalog and Dtrain pairs under Apache 2.0.The study derives its 10K subset from the ToolGen catalog.
  • Q.1 Upstream Artifacts and Licenses: BGE-large-en-v1.5 initializes the encoder under MIT, and Qwen3.5-4B initializes the rewriter under Apache 2.0.The corresponding initializations are θ0 and ψ0, respectively.

Q.2 Data Coverage and Privacy · Q.3 Risks

The benchmark is limited to an English, consumer-oriented slice of RapidAPI’s public catalog, without end-user PII but without a dedicated PII scan. It retains a small minority of adult-content-tagged APIs, and retrieval gains may amplify downstream agent behavior and inherited catalog-selection biases.

  • Q.2 Data Coverage and Privacy: The ToolBench/ToolGen catalog is entirely English-language and sourced from RapidAPI’s public catalog, with a skew toward consumer-facing REST APIs.
  • Q.2 Data Coverage and Privacy: No non-English text appears in queries, tool descriptions, or rewriter outputs.
  • Q.2 Data Coverage and Privacy: Tool records contain public API metadata rather than end-user PII, but the catalog did not undergo a dedicated PII scan.
  • Q.2 Data Coverage and Privacy: A manual review of 100 vague-paraphrase outputs found no inadvertent generation of personal information.
  • Q.2 Data Coverage and Privacy: The catalog retains a small minority of adult-content-tagged APIs to preserve benchmark composition and comparability with prior work.
  • Q.2 Data Coverage and Privacy: The benchmark contains 10,000 tools, 104,224 training examples, and 1,092 evaluation queries, with equally sized vague-paraphrase queries.
  • Q.3 Risks: Retrieval improvements may amplify desirable or undesirable downstream agent behavior, depending on catalog tools and the agent’s policy.
  • Q.3 Risks: The ToolGen catalog inherits selection biases toward consumer-facing REST APIs, English-language descriptions, and un-audited underlying API content.
Loading 2605.29271v1…