PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding the Behaviors of Environment-aware Information Retrieval

Ruifeng Yuan, Chaohao Yuan, David Dai, Yu Rong, Hong Cheng, Hou Pong Chan, Chenghao Xiao

arXiv:2606.16817v1cs.CLcs.IR

TL;DR

Different retrievers may require different query formulation strategies, a gap this paper studies by training LLMs with reinforcement learning to adapt their queries. It finds retriever-dependent optimal styles, with strategies failing to transfer effectively across retrievers, and introduces branching rollouts for stable multi-step training.

  • Problem

    Current agentic RAG research inadequately investigates how LLMs should adapt query formulation to diverse retrieval backends.

  • Method

    The paper systematically analyzes reinforcement-learning-based query adaptation across retrievers and introduces branching rollouts to stabilize multi-step training.

  • Results

    Retrievers exhibit distinct optimal query styles, such as descriptive queries for Contriever and question-like queries for Qwen3-Embedding, with strategies transferring poorly across retrievers.

  • Takeaways & Limitations

    Retriever-aware query formulation provides actionable guidance for RAG practitioners and supports more robust multi-step retrieval optimization.

  • Takeaways & Limitations

    The study evaluates only text-only retrievers and one- or two-step retrieval, leaving multimodal retrieval and longer trajectories for future work.

Abstract

from arXiv · show

Recent retrieval-augmented generation (RAG) approaches have demonstrated strong capability in handling complex queries, yet current research overlooks a critical challenge: different retrievers require fundamentally different query formulation strategies for optimal performance. In this work, we present the first systematic analysis of how LLMs can learn to adapt their query formulation strategies for different retrievers via reinforcement learning (RL). Our empirical study reveals that RL effectively teaches an LLM to tailor its queries to specific retriever characteristics. We discover that different retrievers exhibit surprisingly distinct optimal query styles (e.g., descriptive vs. question-like), suggesting strategies learned for one retriever ineffective for another. We further show that performance can be enhanced by incorporating retriever-specific human guidance and by scaling model size. To facilitate learning over multi-retrieval-step trajectories, we introduce a branching-based rollout technique that improves training stability. Our work provides the first empirical evidence and actionable insights for building truly retriever-aware RAG systems. Code and resources are available at https://github.com/LCO-Embedding/Envs-aware-Information-Retrieval.

1 Introduction

Agentic RAG must adapt query formulation to diverse retrieval backends rather than treating retrieval as a uniform tool call. This work systematically studies retriever-aware query formulation through reinforcement learning and introduces branching-based rollouts for multi-step retrieval.

  • Motivation: Agentic RAG decomposes complex queries and iteratively retrieves information, but real-world systems must support diverse backends, including local repositories.Examples include private financial reports and patient case databases.
  • Research scope: The work presents the first systematic analysis of how LLMs adapt query formulation strategies to different retrievers via reinforcement learning.The study examines retriever types, human guidance comprehensiveness, model size, and retrieval-step scaling.
  • Retriever-specific strategies: Different retrievers require distinct optimal query styles: Contriever favors long-form descriptive queries, whereas Qwen3-Embedding performs better with question-like queries.Strategies optimized for one retriever cannot be effectively transferred to different retrievers.
  • Human guidance: More comprehensive retriever-specific human guidance improves performance by guiding reinforcement-learning exploration in query formulation.The guidance ranges from general query rewriting strategies to retriever-specific formulations provided by human experts.
  • Scaling and optimization: Larger models achieve better performance, can discover policies experts do not anticipate, and benefit from branching rollouts that isolate learning signals to the final turn.Increasing retrieval steps enhances query formulation capabilities and supports multi-step retrieval trajectories.

2 Methodology

The methodology frames retriever-aware query adaptation as reinforcement learning, using retriever feedback to optimize query rewrites and branching rollouts to stabilize multi-turn credit assignment.

  • Retriever-aware query adaptation: An LLM query-rewriter policy learns retriever-specific search queries through GRPO, which stabilizes optimization without a value-function critic.The policy generates queries for a black-box retriever environment conditioned on a user question.
  • Retriever-aware query adaptation: The reward is nDCG@10 against the ground-truth passage, prioritizing retrieval of the correct document and its high ranking over simple recall.This reward mirrors the ranking requirements of real-world RAG systems.
  • Multi-turn query refinement: Multi-turn trajectories interleave queries and masked search results, with weighted aggregation assigning greater importance to later rounds.The trajectory contains queries and reasoning at each turn, while intermediate search results are masked during training.
  • Branching Rollout: Branching Rollout addresses multi-turn reward variance by stabilizing credit assignment in the two-turn reinforcement-learning setting.The method replaces independent linear trajectories such as 16× 1 with branching structures such as 4 × 4.
  • Branching Rollout: Branch-specific estimation averages N second-turn branches for first-turn values and uses history-specific baselines to isolate second-turn rewrite quality.These mechanisms provide Monte Carlo estimation for turn 1 and variance reduction for turn 2.

3 Experiments

The experiments evaluate retriever-aware query rewriting in-distribution on RAGBench, zero-shot on BEIR, and under financial-domain transfer. They use four representative retrievers and examine how their characteristics favor different query styles.

  • In-distribution Evaluation: RAGBench provides the in-distribution evaluation, with more than 100,000 samples spanning general and specialized domains, document structures, and information needs.The benchmark is adapted for retrieval tasks by aggregating its documents.
  • OOD Evaluation: BEIR is used for zero-shot evaluation of whether query rewriters trained on general-domain RAG datasets transfer to commonly used retrieval tasks.The transfer evaluation targets models trained on datasets such as RAGBench.
  • Domain-specific Transfer: Financial-domain transfer tests generalization under terminology and distribution shifts that make finance a challenging cross-domain environment.The study presents finance as an in-depth case study and stress test for transfer from general-domain datasets.
  • Retriever-specific Query Strategies: Under Contriever, a 14B agent discovers question-to-statement rewriting, producing a document-style query while achieving superior performance with a question-style query.The result illustrates that query formulation can vary substantially with retriever environment.

4 Results

RL-trained query rewriters adapt their formulations to retriever-specific behaviors, improving retrieval across in-distribution, out-of-distribution, and domain-shifted settings. Transfer failures primarily reflect structural query-style mismatch, while human guidance, model scaling, and branching rollouts further improve learning.

  • RQ1.1: Retriever-specific query rewriting consistently outperforms using retrievers alone, demonstrating that LLMs learn distinct query “dialects” for different backends.The gains occur across diverse retrievers on the in-distribution RAGBench dataset.
  • Generalization: 6.14 points: the query rewriter raises Contriever’s average BEIR performance from 28.84 to 34.98 in zero-shot evaluation.Gemini-2.5-Flash instead produces ineffective rewrites and lowers performance.
  • Generalization: Query rewriting improves every tested FinAgentBench system, including Contriever, BM25, and Qwen3-Embedding, under substantial domain shift.Retrieval improvements also increase the likelihood that the gold document is highly ranked and improve end-to-end RAG performance.
  • RQ1.2: Structural RE-MMD is often > 20.0 while semantic RE-MMD is typically < 2.0, showing that transfer failures mainly reflect query-style mismatch rather than differing information needs.The contrast is especially pronounced between sparse BM25 and dense Contriever retrievers.
  • RQ2: Human guidance improves RL when the model lacks retriever-specific prior knowledge, whereas exploratory prompts suffice or outperform specific prompts when that behavior is well represented in pretraining.For Contriever, exploratory prompts ultimately achieve higher performance; for bm25, exploratory prompts outperform specific prompts, with both exceeding general prompts.
  • RQ3.2: 4 × 4 Strategy (branching) significantly outperforms 16 × 1 Strategy (independent) in training stability by reducing reward variance and preventing collapse through improved credit assignment.Branching averages outcomes from four second-round attempts conditioned on each first-round attempt under Rtotal = 0.5Rturn1 + Rturn2.

5 Related Work

Prior work addresses LLM–retriever mismatches through query rewriting and reinforcement learning, but generally assumes a uniform retrieval environment. This work instead studies retriever-specific adaptation and stabilizes multi-step query formulation with branching rollouts.

  • Prior approaches mitigate LLM–retriever semantic mismatches through query rewriting or reinforcement learning aligned with generation goals.
  • Different retrievers exhibit distinct length preferences, and policies tailored to one backend can fail to transfer to others.This transfer failure is described as structural drift.
  • Agentic retrieval frameworks often face training instability in long-horizon trajectories, motivating branching rollouts to stabilize multi-step query formulation optimization.

6 Conclusion

The paper presents the first systematic study of adapting LLM query formulation to specific retrievers through reinforcement learning. It shows that strategies are highly retriever-dependent, while model scaling and branching rollouts support stronger and more stable retriever-aware RAG.

  • 6 Conclusion: The work studies adapting LLM query formulation to specific retrievers via reinforcement learning.It is presented as the first systematic study of this adaptation problem.
  • 6 Conclusion: Optimal query strategies are highly retriever-dependent, with stylistic transfer failure quantified as “Structural Drift.”The conclusion identifies retriever dependence and names the associated phenomenon Structural Drift.
  • 6 Conclusion: Scaling model size enables non-intuitive strategies that outperform expert heuristics.The conclusion attributes discovery of these strategies to increased model scale.
  • 6 Conclusion: Branching rollout stabilizes multi-turn training, supporting robust, truly retriever-aware RAG systems.The proposed technique is presented as enabling progress toward robust retriever-aware systems.

Limitations · A RAGBench Details · B FinAgentBench and its Adaptation

The paper’s scope is limited to text-only retrieval and one- or two-step evaluation, while its benchmarks cover broad-domain RAG data and an expert-curated financial retrieval dataset adapted into a conventional format.

  • Limitations: The study examines only text-based retrievers and document collections, excluding multimodal retrieval of heterogeneous content such as images or audio.The authors identify multimodal retrieval as an important direction for future work.
  • Limitations: Evaluation of retrieval-step scaling is restricted to one-step and two-step retrieval.
  • A RAGBench Details: RAGBench contains more than 100,000 samples spanning domains including finance, law, and healthcare.The benchmark aggregates extensive question-answering datasets, including CovidQA, PubMedQA, HotpotQA, MS Marco, CUAD, EManual, TechQA, FinQA, and TAT-QA.
  • A RAGBench Details: RAGBench aggregates QA datasets from diverse domains and established benchmarks.The supplied description names CovidQA, PubMedQA, HotpotQA, MS Marco, CUAD, EManual, TechQA, FinQA, and TAT-QA.
  • B FinAgentBench and its Adaptation: FinAgentBench contains over 18,000 expert-curated financial retrieval samples drawn from SEC filings between 2023 and 2024.Its samples receive graded relevance labels for candidate document chunks.
  • B FinAgentBench and its Adaptation: The original FinAgentBench format combines ranking instructions, a financial question, and candidate document chunks into one large prompt.Candidate chunks are annotated with graded relevance labels.
  • B FinAgentBench and its Adaptation: The adaptation maps each question-chunk pair to its inherited relevance label, producing a conventional retrieval dataset.The resulting dataset consists of questions, documents, and graded relevance judgments.

C RE-MMD Empirical Estimation

The section estimates RE-MMD empirically because direct computation of high-dimensional mean embeddings is often intractable. It uses kernel-based estimators with semantic and structural kernels to distinguish intent drift from lexical or stylistic drift.

  • Kernel-based estimation: Directly computing the embedding ϕ(q) is often intractable, so RE-MMD is estimated using the kernel trick k(q, q′) = ⟨ϕ(q), ϕ(q′)⟩H.The kernel trick converts inner products in the high-dimensional feature space into computable kernel evaluations.
  • Kernel-based estimation: The empirical estimator uses successful query trajectories Qsrc = {x1, ..., xn} from Esrc and Qtgt = {y1, ..., ym} from Etgt.The estimator is explicitly constructed from source and target trajectory sets.
  • Drift diagnosis: A semantic kernel measures intent drift through a Gaussian RBF kernel over dense embeddings such as MiniLM.This kernel captures non-linear semantic similarities between queries.
  • Drift diagnosis: A structural kernel measures lexical and stylistic drift through a Linear kernel over sparse TF-IDF features.It is equivalent to Euclidean centroid distance in lexical space and captures vocabulary divergence and query formulation style.

D Full RAG Performance

Full-RAG experiments on HotpotQA test whether retriever-specific query rewriting improvements transfer to end-to-end generation. The results show stable gains across retrieval depths, supporting robust retriever-aware rewriting.

  • Experimental setup: Full-RAG experiments use HotpotQA with Contriever retrieval and Qwen2.5-7B-Instruct generation to evaluate retriever-specific rewriting.The rewriter agent is compared with vanilla RAG and general, nonretriever-specific rewrites from Gemini-2.5-flash.
  • Results: Improvements in nDCG@10 consistently translate into gains in full-RAG performance.This finding connects retriever-specific retrieval improvements with end-to-end RAG outcomes.
  • Results: The full-RAG improvements remain stable across different numbers of retrieved documents, demonstrating robust retriever-aware query rewriting.Robustness is observed across varying retrieval depths.

E Zero-shot Multi-turn Retrieval Case Study

A rewrite agent trained with only a two-turn branching strategy generalized zero-shot to a four-turn retrieval interaction. Across turns, it adapted from broad hypothetical generation to entity-specific and domain-aligned financial language, ultimately reaching nDCG@10 of 93.74.

  • Zero-shot multi-turn retrieval: The agent generalized zero-shot from two-turn training to a four-turn multi-turn retrieval interaction.This indicates the learned retrieval behaviors extend beyond the training horizon.
  • Retrieval adaptation: Turn 1’s overly general hypothetical document generation produced nDCG@10 of 0.0.The agent then used the returned documents to learn how corpus documents are written.
  • Retrieval adaptation: In turn 2, feedback exposed specific entities, prompting a pivot toward de-dollarization, the Chinese yuan, and global financial systems.This increased nDCG@10, though the supplied passage does not state the resulting value.
  • Retrieval adaptation: 93.74 nDCG@10 was achieved after the agent aligned its rewrite with corporate 10-K filings and financial disclosures using specific financial-reporting terminology.The passage characterizes this as a near-perfect score.

F Training Prompts

The paper uses setting-specific training prompts to teach query rewriting for retrieval. The general prompt allows rewritten queries to be context-rich passages rather than only simple questions and guides method selection.

  • Prompt Design: Figure 6 presents the training prompts used across different experimental settings.
  • General Prompt: The general prompt asks the model to rewrite and extend a question into one new query for retrieving supporting documents.
  • General Prompt: Rewritten queries may be multi-sentence passages providing context, and the model selects an appropriate query-rewriting method.
Loading 2606.16817v1…