AgentIR-4B: Dual-Encoder for Retrieval
- AgentIR-4B is a dual-encoder transformer that combines agent reasoning traces with sub-queries to enhance context-aware document retrieval.
- The model leverages the Qwen3-Embedding-4B backbone with LoRA adapters and a contrastive learning objective, optimized using the DR-Synth pipeline.
- Empirical results on multi-turn benchmarks demonstrate a significant performance boost, achieving 66.27% accuracy compared to traditional retrieval methods.
AgentIR-4B is a 4-billion-parameter dual-encoder transformer designed for Reasoning-Aware Retrieval in Deep Research agent systems. Unlike prior retrieval frameworks that operate solely on user queries, AgentIR-4B is explicitly optimized to embed both an agent’s natural-language reasoning trace and its issued sub-query, enabling contextually enriched document retrieval for complex multi-turn research tasks. Developed from the Qwen3-Embedding-4B backbone and fine-tuned with a synthetic supervision pipeline (DR-Synth), AgentIR-4B delivers substantial accuracy gains on benchmarks requiring deep, multi-hop agentic search (Chen et al., 4 Mar 2026).
1. Model Architecture
AgentIR-4B employs a dual-encoder architecture, with separate but structurally similar encoders for (i) the concatenated agent reasoning trace and sub-query, and (ii) the target document. The backbone model is Qwen3-Embedding-4B, featuring approximately parameters, a hidden size of 4096, 32 transformer layers, and 16 attention heads per layer. Trainable LoRA adapters (rank 8) are injected into every attention projection (query, key, value, and output) in both encoders, contributing approximately 1–2 million additional parameters while freezing the backbone. Both encoders output L₂-normalized 1024-dimensional vectors, serving as the embedding space for similarity computation.
2. Reasoning-Aware Retrieval Paradigm
Retrieval in AgentIR-4B is formalized as nearest-neighbor search in a joint “query + reasoning” embedding space. The joint encoder function is defined as
where denotes string-level concatenation using a prompt template. The document encoder, , applies the same architecture to a document via a separate encoder head:
Documents are ranked by cosine similarity: . The reasoning trace is concatenated directly ahead of , allowing the model to integrate prior context, search intent, and intermediate deductions into the query embedding, maximizing the value of the agent’s evolving thought process.
3. Synthetic Data Generation with DR-Synth
To address the absence of labeled datasets pairing agent reasoning and sub-queries with labeled documents, AgentIR-4B introduces DR-Synth, a data synthesis pipeline. DR-Synth transforms conventional QA datasets (triplets of question, answer, positive documents) into thousands of turn-level training instances suitable for contrastive retrieval learning. For each input QA triple, an LLM-powered search agent (e.g., Tongyi-DeepResearch) executes multi-turn reasoning and querying to produce a trajectory; only successful answer-matched rollouts are kept. For each search turn, the pipeline collects (reasoning, sub-query) pairs, re-ranks candidate documents using an oracle reranker, and builds training tuples comprising the positive and several hardest negative documents. On the WebShaper QA dataset (500 questions), DR-Synth yields 5,238 turn-level examples after filtering.
4. Training Objective and Optimization
AgentIR-4B is fine-tuned exclusively via the LoRA adapters using an in-batch contrastive loss:
where 0, 1, positive 2, hard negatives 3, and 4. All vectors are L₂-normalized before similarity computation. The model employs weight decay of 5 and 0.1 dropout in the LoRA layers, with no additional regularization terms. In-batch negatives further increase contrastive pressure.
5. Empirical Performance on BrowseComp-Plus
On the multi-turn BrowseComp-Plus benchmark, AgentIR-4B, paired with the Tongyi-DeepResearch agent, achieves 66.27% end-to-end QA accuracy, a 17.6-point absolute gain over its unfine-tuned Qwen3-Embed-4B backbone (48.67%) and outperforming an 8B-parameter embedding model (50.72%) by approximately 15 points. Document recall and average search calls are 78.86 and 25.91, respectively. Comparisons include:
- BM25: 33.98% accuracy, 46.83 recall
- Qwen3-Embed-4B + LLM rerank: 55.66% accuracy
- HyDE-style Reason-Rewriter+Reason-Embed-8B: 31.08% accuracy
Ablation studies confirm the contributions of both reasoning-aware input and DR-Synth training: reasoning-only input at inference yields 55.54% accuracy, while DR-Synth finetuning (without reasoning) yields 59.40%. Combining both yields the full 66.27% result. Comparisons are robust across three agent implementations.
6. Implementation and Usage
AgentIR-4B, pretrained weights, and associated DR-Synth data are publicly available. Practical usage involves cloning the repository, installing required dependencies (Tevatron, PyTorch, HuggingFace Transformers), obtaining checkpoints, and following provided scripts for DR-Synth data generation or direct fine-tuning with WebShaper. The model can be integrated into a standard dense retrieval index (FAISS, ANNOY) and paired with agentic LLM search systems using the released prompt templates. Detailed instructions for rollout generation, model training, and evaluation are documented and provided at https://texttron.github.io/AgentIR/.
7. Significance and Implications
AgentIR-4B demonstrates that leveraging the explicit reasoning traces generated by LLM agents provides a high-value retrieval signal, closing gaps between agentic and traditional human-in-the-loop retrieval paradigms. The introduction of DR-Synth supplies critical supervision by converting abundant QA data into agent-centric retriever training material. The resulting system achieves state-of-the-art multi-hop reasoning retrieval performance among models of similar or larger size in the Deep Research agent setting, without incurring additional inference cost beyond a single embedding operation per search (Chen et al., 4 Mar 2026). A plausible implication is that further integration of agent intermediate states—beyond queries alone—could systematically improve performance in multi-turn, agent-driven research workflows.