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

Research has not adequately examined how LLMs can adapt query formulation to different retrievers, despite agentic RAG systems using diverse retrieval backends. This paper studies reinforcement-learning-based adaptation and finds that retrievers favor distinct query styles that do not transfer effectively between them.

  • Problem

    Research had not adequately investigated how LLMs can adapt query formulation across diverse retrieval backends, despite agentic RAG's reliance on retrieval beyond commercial search engines.

    Source: §1 · p. 2

  • Method

    The paper systematically studies reinforcement-learning-based LLM adaptation to different retrievers and introduces branching rollouts for stable multi-step query-strategy optimization.

    Source: §1 · p. 2

  • Results

    Different retrievers have distinct optimal query styles—descriptive for Contriever and question-like for Qwen3-Embedding—so strategies optimized for one do not transfer effectively to another.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings provide RAG practitioners actionable guidance to account for retriever-specific query strategies.

    Source: §1 · p. 2

    The study covers only text retrievers and text-based document collections, not multimodal retrieval.

    Source: Limitations · p. 9

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 work across diverse retrieval backends, but different retrievers may need different query strategies. This introduction presents a systematic study of using reinforcement learning (RL) to make LLMs adapt their queries to the retriever they are using.

Study scope

  • The study systematically examines how retriever type, human guidance, model size, and retrieval-step count affect RL-based query formulation.It is presented as the first systematic analysis of how LLMs learn retriever-specific query strategies through RL.Source: §1 · p. 2

Findings

  • Contriever favors long-form descriptive queries, whereas Qwen3-Embedding favors question-like queries, and strategies optimized for one do not transfer effectively to the other (Izacard et al., 2022; Zhang et al., 2025).The study finds that RL can adapt an LLM’s query formulation strategy to local retrievers.Source: §1 · p. 2

Human guidance

  • More comprehensive, retriever-specific human guidance in prompts leads to better performance by guiding RL exploration in query formulation.The study compares three guidance levels, from general rewriting strategies to expert-provided retriever-specific formulations.Source: §1 · p. 2

Scaling and multi-step retrieval

  • Larger models perform better and can discover query policies that human experts do not anticipate.The study also finds that additional retrieval steps improve query formulation because models can use earlier retrieved documents to refine later queries.Source: §1 · p. 2
  • A branching rollout technique isolates the learning signal to the final turn, supporting stable RL training over multi-step retrieval trajectories.The technique is introduced to support multi-step query-strategy optimization.Source: §1 · p. 2

2 Methodology

The methodology trains an LLM query-rewriting policy to adapt to a black-box retriever, using relevance-ranked retrieval rewards and single- or multi-turn reinforcement learning. Branching rollouts address unstable credit assignment in multi-turn training.

2 Methodology

  • The query-rewriting policy is optimized with GRPO (Shao et al., 2024), which supports stable learning without a value-function critic.The policy generates search queries conditioned on the user’s question and the black-box retriever environment.Source: §2 · pp. 2–3

2.1 Reward Function Formulation

  • Reward is nDCG@10 against the ground-truth passage, encouraging the model to retrieve the correct document and rank it highly.This ranking-oriented reward is selected over simple recall to reflect real-world RAG requirements.Source: §2.1 · p. 3

2.2 Scenario 1: Single-Turn Retrieval

  • In single-turn retrieval, the policy samples a group of query rewrites per input and normalizes each reward against group statistics to compute its advantage.The GRPO objective uses an importance-sampling ratio between the current and previous policies.Source: §2.2 · p. 3

2.3 Scenario 2: Multi-Turn Iterative Retrieval

  • Multi-turn retrieval forms trajectories of queries and search results, then aggregates rewards with greater weight on later rounds.Search results are masked during training, and the weighting values the final answer while encouraging useful intermediate steps.Source: §2.3 · p. 3

2.4 Branching Rollout Strategy

  • Branching rollouts stabilize multi-turn credit assignment by averaging first-turn outcomes across branches and grouping second-turn rewards by shared history.For the two-turn setting, the method branches M first-turn actions into N second-turn completions, reducing intermediate-reward variance and isolating second-turn rewrite quality.Source: §2.4 · p. 3

3 Experiments

The experiments evaluate query rewriting on RAGBench and test transfer to BEIR and a finance-specific retrieval environment. They also describe retriever-specific query behavior and the multi-turn RL training setup.

3.1 Settings

  • RAGBench evaluation aggregates documents into a global retrieval corpus, uses questions as queries and associated contexts as positives, and deduplicates documents with Qwen-Embedding-0.6B at similarity above 0.99.RAGBench contains more than 100,000 samples across general and specialized domains (Friel et al., 2024).Source: §3.1 · p. 3
  • For zero-shot out-of-distribution evaluation, the study tests query rewriters trained on general-domain RAG data on BEIR (Thakur et al., 2021).Source: §3.1 · p. 3
  • The finance transfer case study uses FinAgentBench (Choi et al., 2025) to test general-domain-trained rewriters under terminology shifts and documents combining text with tables.The passage motivates finance as a stress test through distribution shift and its practical retrieval relevance.Source: §3.1 · p. 4

3.2 Implementation Details

  • Multi-turn RL weights first-round retrieval at 0.5 and second-round retrieval at 1, and computes loss only on model-generated reasoning and rewritten queries.The implementation also compares n^2 rollouts over full two-turn trajectories with n rollouts over second-turn rewriting conditioned on first-turn rewrites and search results.Source: §3.2 · p. 4
  • Four selected retrievers—Qwen3-Embedding, all-MiniLM-L6-v2, Contriever, and BM25—span supervised dense models, unsupervised representation retrieval, and sparse keyword retrieval.The dense models respectively represent decoder/high-capability/large and encoder/low-capability/small systems; BM25 is sensitive to phrasing and spelling variants on non-massive corpora.Source: §3.2 · p. 4
  • Under Contriever, a 14B agent discovered question-to-statement rewriting, while the description contrasts intuitive document-style queries with better performance from question-style queries.Contriever is described as trained on randomly cropped Wikipedia spans; the reported policy finding is framed as a retriever-specific strategy.Source: §3.2 · p. 4

4 Results

Across retrievers and domains, learned rewrites improve retrieval, while optimal query styles vary by backend; prompt guidance, model scale, and branching also shape learning outcomes.

4.1 Main Results

  • On BEIR, rewriting raises Contriever’s average from 28.84 to 34.98 zero-shot; in-distribution adaptations outperform raw retrievers, and FinAgentBench improves across all tested systems.Gemini-2.5-Flash rewrites under aligned prompts instead reduce performance, while Appendix D reports that retrieval gains also improve end-to-end RAG.Source: §4.1 · p. 5

4.2 Measuring the Discrepancy of Optimal Queries

  • Semantic RE-MMD stays typically below 2.0, but structural RE-MMD often exceeds 20.0, indicating transfer failures are mainly stylistic rather than differences in information need.The contrast is especially pronounced between sparse BM25 and dense Contriever; policies optimized for one retriever generally do not transfer reliably to another.Source: Figure 3 · p. 6, §4.2 · p. 6

4.3 Human Guidance in Prompt

  • For Contriever, specific document-style guidance wins finally after exploratory prompts collapse early, whereas for BM25 exploration beats specific and general prompts.Explicit expert guidance helps when retriever behavior is not sufficiently known to the model; exploratory prompts can suffice when that behavior is represented in pretraining.Source: §4.3 · p. 7

4.4 Scaling of Model Size and Retrieval Steps

  • The 14B rewriter escapes the 4B and 8B plateaus and discovers statement-style queries: “Being white and Latino” scores nDCG@10 0.965 versus 0.542 for question-style.The 14B model’s statement-style strategy emerges under exploratory prompts and differs from both question-form and document-style rewrites.Source: §4.4 · pp. 7–8
  • The 4 × 4 branching strategy significantly outperforms the 16 × 1 independent strategy in training stability.The comparison addresses credit assignment under the shared reward function.Source: Figure 5 · p. 8, §4.4 · p. 8
  • Larger models explore more broadly and can begin with lower average rewards, but this exploration can uncover more effective policies and escape local optima.The 14B model’s greater capacity and world knowledge help mitigate premature policy collapse.Source: §4.4 · pp. 7–8
  • Branching isolates second-turn learning by canceling the fixed first-turn reward, and averages four continuations to reduce first-turn reward variance and prevent training collapse.The four continuations also estimate the value of a first-round action more stably than a single second-round sample.Source: §4.4 · p. 8

5 Related Work

Prior RAG work addresses query mismatches and multi-turn retrieval, but often assumes uniform retrieval environments or encounters training instability. This work focuses on retriever-specific query adaptation and introduces branching rollouts to stabilize multi-step optimization.

5 Related Work

  • RAG augments LLMs on tasks exceeding their knowledge boundaries; to bridge behavioral gaps between LLMs and retrievers, prior approaches mitigate semantic mismatches via query rewriting or zero-shot expansion (Lewis et al., 2020; Xiao et al., 2025a; Chan et al., 2024; Gao et al., 2023).Source: §5 · p. 8
  • Prior RL rewrite methods typically assume a uniform retrieval environment, despite evidence that retrievers have distinct length preferences; this work analyzes retriever-specific adaptation to address structural drift.Structural drift describes policies tailored for one backend failing to transfer to others.Source: §5 · pp. 8–9
  • Agentic frameworks such as Self-RAG (Asai et al., 2024) and Adaptive-RAG (Mallen et al., 2023) use multi-turn retrieval, but long-horizon training can be unstable; branching rollouts stabilize multi-step query-strategy optimization.The approach complements work on the retriever-generator preference gap (Ke et al., 2024).Source: §5 · p. 9

6 Conclusion

The paper presents the first systematic study of adapting LLM query formulation to specific retrievers through reinforcement learning. It finds retriever-dependent strategies and introduces scaling and branching-rollout approaches that support retriever-aware RAG.

  • Optimal query strategies are highly retriever-dependent, and their stylistic transfer failure is quantified as “Structural Drift.”Source: §6 · p. 9
  • Scaling model size enables discovery of non-intuitive query strategies that outperform expert heuristics.Source: §6 · p. 9
  • The proposed branching rollout technique stabilizes multi-turn training, paving the way for robust, retriever-aware RAG systems.Source: §6 · p. 9

Limitations · A RAGBench Details · B FinAgentBench and its Adaptation

The study is limited to text-only retrieval and one- or two-step retrieval trajectories. The appendices describe RAGBench and explain how FinAgentBench was adapted into a conventional retrieval dataset.

Limitations

  • The study covers text-only retrievers and text-based collections, leaving multimodal retrieval of heterogeneous content such as images or audio for future work.The authors identify extending the analysis to multimodal retrieval as an important direction.Source: Limitations · p. 9
  • Evaluation of retrieval-step scaling is limited to one- and two-step retrieval; longer trajectories remain for future study.The authors say these settings demonstrate multi-step retrieval benefits, but call for investigation of learning dynamics and performance trends with longer trajectories.Source: Limitations · p. 9

A RAGBench Details

B FinAgentBench and its Adaptation

  • FinAgentBench (Choi et al., 2025) contains over 18,000 expert-curated samples from SEC filings dated 2023–2024.Each original sample combines ranking instructions, a financial question, candidate document chunks, and graded relevance labels in one prompt.Source: B FinAgentBench and its Adaptation · p. 11
  • To adapt FinAgentBench, the authors extracted questions and chunks, pooled them globally, retained exact-match unique instances, and mapped inherited relevance labels to question–document pairs.This produces a conventional retrieval dataset of questions, documents, and graded relevance judgments; no normalization or model-based processing was applied.Source: B FinAgentBench and its Adaptation · p. 11

C RE-MMD Empirical Estimation

RE-MMD is estimated with the kernel trick because directly computing high-dimensional query embeddings is often intractable. The method uses successful query trajectories and separates semantic intent drift from lexical and stylistic drift.

C RE-MMD Empirical Estimation

  • The kernel trick estimates RE-MMD without directly computing often-intractable high-dimensional query embeddings.It uses k(q, q′) = ⟨ϕ(q), ϕ(q′)⟩H to estimate the distance between mean embeddings.Source: C RE-MMD Empirical Estimation · p. 11
  • The empirical estimator uses successful query trajectories from source and target environments, Qsrc and Qtgt.The passage specifies an unbiased estimator for trajectory sets Qsrc = {x1, ..., xn} and Qtgt = {y1, ..., ym}.Source: C RE-MMD Empirical Estimation · p. 12
  • A Gaussian RBF kernel over dense embeddings such as MiniLM measures intent drift through nonlinear semantic similarity.This is the semantic kernel.Source: C RE-MMD Empirical Estimation · p. 12
  • A linear kernel over sparse TF-IDF features measures lexical and stylistic drift as centroid distance, capturing vocabulary and query-style divergence.This is the structural kernel.Source: C RE-MMD Empirical Estimation · p. 12

D Full RAG Performance

In HotpotQA full-RAG experiments, improvements in nDCG@10 consistently translate into full-RAG performance gains that remain stable across different numbers of retrieved documents.

D Full RAG Performance

  • The HotpotQA evaluation compares a contriever/Qwen2.5-7B-Instruct rewriter agent with vanilla RAG and Gemini-2.5-flash general rewrites.It uses the BEIR HotpotQA corpus and original questions, reporting Accuracy, F1, and Exact Match (EM).Source: D Full RAG Performance · p. 12
  • Full-RAG performance gains from nDCG@10 improvements remain stable across different numbers of retrieved documents.The authors describe this stability as evidence of the robustness of retriever-aware query rewriting.Source: D Full RAG Performance · p. 12

E Zero-shot Multi-turn Retrieval Case Study

In a four-turn case study, an agent trained with two-turn branching adapts its rewrites using retrieval feedback, reaches nDCG@10 93.74, and generalizes beyond its training horizon.

E Zero-shot Multi-turn Retrieval Case Study

  • Despite training with only a two-turn branching strategy, the agent demonstrated zero-shot retrieval across four turns, extending beyond its training horizon.The example suggests learned retrieval behaviors can extend to more complex multi-turn interactions.Source: E Zero-shot Multi-turn Retrieval Case… · p. 12
  • The first broad hypothetical-document rewrite scored nDCG@10 0.0; after seeing returned documents, the agent pivoted to specific entities, enhancing the next-round score.Its second-turn query focused on de-dollarization, the Chinese yuan, and global financial systems; the passage does not give the resulting score.Source: E Zero-shot Multi-turn Retrieval Case… · p. 13
  • nDCG@10 reached 93.74 after the agent aligned its rewrite with the financial-reporting style and domain of the corpus.The agent recognized that available documents were corporate 10-K filings and financial disclosures, then used more specific financial terminology.Source: E Zero-shot Multi-turn Retrieval Case… · p. 13

F Training Prompts

The section documents the training prompts used in different experimental settings.

Related papers

Loading 2606.16817v1…