PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-RAG: Long Context Multi-step Retrieval via Value-based Embedder Training

Artyom Sorokin, Nazar Buzun, Alexander Anokhin, Oleg Inozemcev, Egor Vedernikov, Petr Anokhin, Mikhail Burtsev, Trushkov Alexey, Yin Wenshuai, Evgeny Burnaev

arXiv:2511.07328v2cs.LGcs.IR

TL;DR

Complex questions often require multi-step retrieval, but existing approaches can demand expensive fine-tuning of small language models. Q-RAG instead fine-tunes an embedder with reinforcement learning and achieves state-of-the-art or competitive results across long-context and open-domain QA benchmarks, including contexts up to 10M tokens.

  • Problem

    Most RAG methods use single-step retrieval, while complex questions require multi-step context interaction and existing solutions require expensive fine-tuning.

  • Method

    Q-RAG fine-tunes state and action embedders with value-based reinforcement learning for multi-step retrieval in latent text-chunk space.

  • Results

    Q-RAG achieves state-of-the-art or highly competitive results across long-context and open-domain QA benchmarks, including contexts up to 10M tokens.

  • Takeaways & Limitations

    Q-RAG provides a resource-efficient multi-step retrieval alternative that remains effective as context length grows and can pair with frozen large language models.

  • Takeaways & Limitations

    Q-RAG assumes that long contexts are pre-segmented into non-overlapping text chunks.

Abstract

from arXiv · show

Retrieval-Augmented Generation (RAG) methods enhance LLM performance by efficiently filtering relevant context for LLMs, reducing hallucinations and inference cost. However, most existing RAG methods focus on single-step retrieval, which is often insufficient for answering complex questions that require multi-step search. Recently, multi-step retrieval approaches have emerged, typically involving the fine-tuning of small LLMs to perform multi-step retrieval. This type of fine-tuning is highly resource-intensive and does not enable the use of larger LLMs. In this work, we propose Q-RAG, a novel approach that fine-tunes the Embedder model for multi-step retrieval using reinforcement learning (RL). Q-RAG offers a competitive, resource-efficient alternative to existing multi-step retrieval methods for open-domain question answering and achieves state-of-the-art results on the popular long-context benchmarks BabiLong and RULER for contexts up to 10M tokens. Code is available at https://github.com/griver/Q-RAG

1 INTRODUCTION

Q-RAG addresses the limitations of single-step RAG on complex, long-context tasks by training a compact retrieval agent with value-based reinforcement learning instead of fine-tuning an LLM. It achieves state-of-the-art performance on ultra-long-context reasoning and NIAH benchmarks while remaining faster and cheaper than existing multi-step RAG methods.

  • Motivation: LLMs face static knowledge, long-context inefficiency, attention dilution, and hallucinations, motivating Retrieval-Augmented Generation.RAG extracts relevant portions of external corpora or lengthy contexts to produce shorter, more focused inputs.
  • Motivation: Single-step retrieval works for simple tasks but is insufficient for complex questions, while multi-step methods require expensive LLM fine-tuning.This requirement makes existing approaches impractical for large models and limits accessibility.
  • Approach: Q-RAG trains a retrieval agent directly in the latent space of text-chunk embeddings using value-based reinforcement learning.The method uses temporal-difference reinforcement learning to train a compact, efficient multi-step retrieval model.
  • Results: Q-RAG achieves state-of-the-art results on long-context commonsense reasoning, multi-hop QA, and NIAH tasks with contexts up to 10 million tokens.It also performs competitively on the MuSiQue and HotPotQA open-domain QA benchmarks.
  • Approach: Q-RAG incorporates temporal information into the multi-step embedder, enabling temporal reasoning during retrieval that generalizes to long contexts at inference time.This temporal mechanism is presented as a distinct contribution of the approach.

2 RELATED WORK

Prior work addresses complex long-context retrieval through fine-tuning-free agentic systems, LRM-based multi-step retrieval, retriever fine-tuning, and architectures designed for extremely long sequences. Q-RAG is positioned closest to feedback-based retriever training while distinguishing itself through reinforcement-learning-based embedder training for multi-step retrieval.

  • Retrieval for long-context tasks: Fine-tuning-free LLM agents combine off-the-shelf retrievers with LLMs, while knowledge-graph methods improve retrieval quality but can be extremely slow at inference.Examples include Search-o1, GraphReader, HippoRAG, and AriGraph.
  • Retrieval for long-context tasks: LRM-based approaches fine-tune models to generate intermediate search queries for multi-step retrieval over long contexts.IM-RAG pioneered this direction by fine-tuning an LLM with a frozen embedder using PPO; later work includes R1-Searcher, SearchR1, RAG-RL, and ReSearcher.
  • Retrieval for long-context tasks: Feedback-based retriever fine-tuning, exemplified by RePlug, is closest to Q-RAG, but RePlug does not address multi-step reasoning or reinforcement learning in this setting.BeamRetriever instead trains a reranker for BeamSearch-style planning and targets short-context QA.
  • Long-context architectures: Recurrent and state-space architectures support extremely long-sequence processing through recurrence, with Mamba offering linear-time scalability but weaker in-context learning and less expressive token-to-token interaction than Transformers.The passage introduces recurrent-memory alternatives including Mamba and the Recurrent Memory Transformer.
  • Long-context architectures: Long-context recurrent models expand memory or modify attention: ATLAS outperforms RMT and Titans, while ARMT achieves the best long-context scores among recurrent models.MemUP uses uncertainty to identify events requiring long-term memory.
  • Long-context architectures: LongRoPE2 extends pretrained LLMs’ effective context window to 128K tokens through RoPE rescaling and mixed-window training while retaining short-context performance.It addresses the positional encoding bottleneck.

3 METHODS

Q-RAG formulates multi-step retrieval as a finite-horizon MDP and trains a value-based agent whose state and action embedders score candidate chunks. Relative positional encoding incorporates relationships among retrieved evidence and candidate chunks, supporting both fact-finding and narrative retrieval.

  • MDP formulation: Q-RAG models retrieval over pre-segmented long contexts as a finite-horizon MDP, where the agent selects chunks needed to answer a query or verify a claim.The dataset consists of triples (C, q, y), with q representing either a user question or a generated factuality-consistency claim.
  • MDP formulation: States contain the query and previously selected chunks ordered by their original document positions, with deterministic transitions and termination by a step budget or STOP action.The initial state is s0 = [q], and each transition appends the selected chunk before reordering the state.
  • Value-based agent: The value-based agent uses maximum-entropy reinforcement learning, approximating Q-values with state and action embedders whose dot product scores candidate chunks.The action embedder uses rotary position embeddings to encode candidate chunk content and position, while the temperature α controls exploration.
  • Training: PQN replaces DQN’s replay buffer, avoiding repeated chunk re-embedding and reducing the training slowdown and memory requirements associated with large contexts.Q-RAG additionally uses λ-returns rather than one-step returns to improve training stability and learning speed.
  • Relative positional encoding: Relative positional encoding lets the Q-function exploit candidate-to-evidence spatial relationships while retaining local order within intervals, enabling narrative as well as disjoint-document retrieval.The reported experimental example uses δ=10 and ℓ=9 for inter-interval steps and within-interval resolution.

4 EXPERIMENTS

Q-RAG is evaluated across reasoning, needle-retrieval, and open-domain multi-hop QA tasks over contexts from 4K to 10M tokens. It achieves strong long-context retrieval and multi-hop QA performance, including robust scaling to million-token contexts and best overall results across HotPotQA and MuSiQue.

  • Evaluation Setup: Q-RAG is evaluated on BabiLong, RULER, HotPotQA, and MuSiQue across context lengths from 4K to 10M tokens.The tasks cover commonsense and temporal reasoning, Needle-in-a-Haystack retrieval, and open-domain multi-hop question answering.
  • BabiLong: Q-RAG achieves the highest average BabiLong performance in ultra-long contexts ranging from 1 to 10 million tokens.This demonstrates stronger long-context generalization than other specialized long-context methods.
  • BabiLong: On BabiLong’s hardest QA3 task, Q-RAG shows virtually no degradation as context length increases, unlike alternative long-context approaches.QA3 requires multi-step search over at least three facts and temporal reasoning.
  • RULER: Q-RAG achieves near-perfect performance on all RULER Needle-in-a-Haystack subtasks and generalizes from 4K-length training documents to 1M-token contexts without accuracy loss.For RULER multi-hop QA, Q-RAG is evaluated out-of-distribution after fine-tuning on HotPotQA.
  • RULER: On RULER multi-hop QA, Q-RAG significantly outperforms all baselines at every evaluated context length, with degradation beginning only at 1M tokens.The comparison includes Beam Retriever and other methods fine-tuned on RULER or related synthetic Needle-in-a-Haystack datasets.
  • Open-Domain Multi-Hop QA: Q-RAG matches Beam Retriever’s fact retrieval accuracy, surpasses other HotPotQA baselines, matches Search-R1 on HotPotQA, and outperforms all alternatives on MuSiQue.These results yield the best overall performance across the HotPotQA and MuSiQue benchmarks.

5 ABLATION STUDY

The ablation study evaluates Q-RAG on the challenging, arbitrary-length BabiLong-QA3 benchmark against non-fine-tuned, supervised, and architectural variants. It also examines retrieval-budget sensitivity on HotPotQA, finding that increasing retrieval steps from 2 to 3 improves retrieval and answer quality across three LLMs.

  • Ablation setup: Q-RAG ablations were conducted on BabiLong-QA3, selected for its difficulty and support for evaluation at arbitrary context lengths.All baselines used three random seeds, with results reported across multiple context lengths.
  • Hyperparameter sensitivity: Q-RAG sensitivity was evaluated for the λ-return parameter and entropy coefficient α, with Figure 3 also comparing inference runtime across context lengths.α controls the strength of entropy regularization in the soft Q function.
  • Ablation setup: The study compares Q-RAG with the same retrieval pipeline using pretrained embedders, supervised trajectory fine-tuning, and variants without target networks or soft-Q objectives.These baselines isolate the effects of reinforcement-learning embedder training, target networks, and entropy-regularized value functions.
  • Retrieval-budget sensitivity: Increasing retrieval steps from 2 to 3 improves both correct-fact retrieval and answer quality across all three LLMs on 1000 HotPotQA samples.The system used an Alibaba-NLP/gte-multilingual-base embedder trained on HotPotQA+MuSiQue, with EM measuring correct retrieved chunks and F1 accounting for noise.
  • Retrieval-budget sensitivity: Within a reasonable retrieval-count range, final answer accuracy is primarily dependent on the number of correct supporting facts retrieved.The study also considers dynamic stopping based on a Q-value threshold, analyzed in Appendix B.

6 CONCLUSION

Q-RAG provides resource-efficient multi-step retrieval through reinforcement-learning-based embedder training, achieving strong long-context and open-domain QA results while remaining compatible with powerful frozen LLMs. The paper also identifies future extensions and provides code, data, instructions, and pretrained checkpoints for reproduction.

  • Results: Q-RAG achieves state-of-the-art or highly competitive results across BabiLong, RULER, MuSiQue, and HotPotQA, with advantages widening as context length grows.Performance shows minimal degradation at ultra-long scales.
  • Efficiency: All training used a single A100 GPU with 80 GB memory, compared with about eight-A100 clusters typically reported for Search-R1/R1-Searcher.Q-RAG fine-tunes only the embedder while keeping the LLM frozen.
  • Deployment: Q-RAG’s frozen-LLM design supports pairing with powerful pre-trained or proprietary LLMs, enabling efficient training and flexible deployment over very long contexts.The method preserves strong retrieval while reducing training requirements.
  • Future work: Future work includes structured LLM feedback rewards, stronger compositional and temporal reasoning in embedding space, and tighter integration with generation.These directions aim to preserve Q-RAG’s efficiency and scalability.
  • Reproducibility: The repository provides reproducibility materials for BABILong, HotPotQA, MuSiQue, and RULER, while pretrained Q-RAG checkpoints are available through Hugging Face.The released implementation fine-tunes publicly available embedders, including multilingual-e5-large, Alibaba-NLP/gte-multilingual-base, and facebook/contriever.

A INNER PRODUCT APPROXIMATION FOR Q-FUNCTION

The section establishes that positional block-rotation scores, including standard RoPE, admit a complex-diagonal inner-product representation whose function class is universally dense under suitable algebraic conditions. It also derives approximation-rate results, with a sharp L2 rate and a slower or more demanding uniform L∞ guarantee.

  • Complex representation: Real block-rotation scores, including RoPE, can be represented as complex inner products with diagonal positional matrices.The representation uses Λ(t) = diag(eiθ1t, . . . , eiθmt) for suitable fixed frequencies θk.
  • Universal approximation: Standard RoPE is a special case whenever Φ contains the exponentials t 7→eiθkt associated with the block rotations.Thus the RoPE score belongs to the complex-diagonal score class AΦ under the stated algebraic inclusion condition.
  • Universal approximation: If Φ is a self-adjoint subalgebra containing constants and separating points of T, the induced complex-diagonal class is dense in C(K, R) in the uniform norm.This is the universal approximation theorem stated for compact state, action, and positional domains.
  • Quantitative rates: The quantitative L2 result provides feature dimension d′ ≤ d and error bound C M d−β with β = s/(D+1), where D = dx + dy.The construction uses L2 feature maps on X and Y and unitary positional matrices Λ(t).
  • Quantitative rates: The uniform L∞ bound achieves Cε M d−βε with βε = s(s −D/2 −ε)/(D s + s −D/2 −ε), while pure Hs assumptions yield a slower rate or require stronger spatial smoothness.The section characterizes the L2 theorem as sharp and technically robust under minimal assumptions, whereas uniform control is stronger.

B EARLY STOPPING EXPERIMENTS · C PLANNING FOR MULTI-STEP RETRIEVAL

The early-stopping analysis defines threshold-based retrieval termination and shows a trade-off between shorter episodes and premature stopping, with a broad high-performing regime on BabiLong QA2. The paper also proposes computationally cheap beam-search planning for multi-step retrieval by ranking trajectories with embedding-based action values.

  • B EARLY STOPPING EXPERIMENTS: The agent continues retrieving while Qt ≥ Qthreshold and stops at the first step with Qt < Qthreshold, defining tstop as the number of selected chunks.This stopping rule is evaluated against tearliest, the earliest step at which all ground-truth chunks have been collected.
  • B EARLY STOPPING EXPERIMENTS: Stopping is classified as early when tstop < tearliest, perfect when equal, and late when greater.The oracle policy stops at tearliest and therefore never stops too early or too late.
  • B EARLY STOPPING EXPERIMENTS: As Qthreshold increases on HotPotQA, late errors decrease while early-stop errors increase, and episodes become shorter.Perfect stops peak around Qthreshold ≈0.1–0.3; overly high thresholds rapidly increase early-stop errors and degrade performance.
  • B EARLY STOPPING EXPERIMENTS: The best Fact F1 on HotPotQA is achieved at Qthreshold = 0.2 for the GTE embedder under penalize extra steps=True and never terminate=True.Table 5 reports stopping fractions, episode length, Fact EM/F1, and classifier TPR/FPR for this setting.
  • B EARLY STOPPING EXPERIMENTS: AUC is 0.96 on HotPotQA and 0.97 on BabiLong QA2, with the oracle marking an upper bound on the TPR–FPR trade-off.The ROC curves assign each point to a Q-value threshold and include the oracle stopping policy.
  • B EARLY STOPPING EXPERIMENTS: On BabiLong QA2, Qthreshold ∈[0.2, 0.6] keeps perfect stops at ≈0.95–0.99 while reducing episode length from about 6 to roughly 2.2 retrieval steps.Fact F1 remains around 0.95 and answer accuracy stays near-optimal in this interval; performance collapses only as the threshold approaches 1.0.
  • C PLANNING FOR MULTI-STEP RETRIEVAL: Planning formulates multi-step source selection as a search over action trajectories and can use beam search ranked by Qθ(s, a).Because Qθ(s, a) = ⟨Es(s), Ea(a)⟩ is computed as a state–action embedding dot product, the planning requires no new transformer forward passes.

D METHOD COMPLEXITY AND EFFICIENCY

Q-RAG confines interaction with the original long context to its retrieval agent, while the LLM Answerer’s complexity is independent of document length. Retrieval is linear in document length under naive search, while capped batching makes peak-memory growth negligible and training converges rapidly.

  • Architecture: Q-RAG uses a multi-step retrieval agent to collect relevant evidence, followed by an LLM Answerer conditioned only on retrieved chunks.Only the retrieval agent accesses the original long context; the Answerer’s effective context depends on retrieval hyperparameters.
  • Answerer Complexity: O(1) time and memory complexity with respect to original context length N apply to the LLM Answerer.Its cost depends on retrieval settings such as the number of steps T and maximum chunk length, rather than N.
  • Retrieval Complexity: O(N) complexity arises from embedding all document chunks once per retrieval episode and naively scoring candidate actions.Approximate kNN methods can reduce action-selection query time to sub-linear in practice.
  • Space Complexity: O(N) storage is required for chunk embeddings, but capping the action embedder batch size makes peak-memory growth with N negligible in practice.GPU memory is mainly consumed by LLM weights and action-embedder forward passes.
  • Training Time Efficiency: Rapid convergence on BabiLong QA2 and HotPotQA is shown by sharp initial evaluation-score gains followed by a stable plateau.Figure 6 reports average episodic return with respect to training time.

E EXTRA QA RESULTS

Table 7 compares multi-step retrieval methods across HotPotQA-distractors, MuSiQue in-distribution, and MuSiQue out-of-distribution using fact-retrieval and answer-generation metrics. Q-RAG and Plan Q-RAG perform strongly overall, while Beam-Retriever leads on HotPotQA but generalizes less robustly.

  • Datasets: Table 7 evaluates HotPotQA-distractors, MuSiQue in-distribution, and MuSiQue out-of-distribution.The comparison covers multi-step retrieval methods across three datasets.
  • Metrics: The evaluation reports Fact F1, Fact EM, Ans F1, and Ans EM.These metrics measure fact retrieval and answer generation.
  • Results: Q-RAG and Plan Q-RAG achieve strong overall results, especially on out-of-distribution data.Plan Q-RAG is identified as a planned variant of Q-RAG.
  • Results: Beam-Retriever leads on HotPotQA but generalizes less robustly.The passage contrasts its HotPotQA performance with broader generalization.

F TRAINING DETAILS

Q-RAG training used AdamW with scheduled learning-rate decay, gradient accumulation, and clipping. The objective included fixed algorithmic hyperparameters, benchmark-specific encoders, and training completed within 12 hours on one A100-80GB GPU.

  • Optimization: AdamW training used a 1.5 × 10−5 learning rate, 1,000-step warm-up, decay to 10% of the initial rate, and gradient clipping at ℓ2 norm 2.0.Gradient accumulation ran for 8 steps, producing an effective batch size of 96 from a base mini-batch size of 12.
  • Objective and representations: The objective and algorithmic components used γ=0.99, α=0.05, λ=0.5, and τ=0.02, with action representations capped at 220 tokens.
  • Training cost: 12 hours was the maximum end-to-end training time for a single model on one A100-80GB GPU.
  • Benchmark-specific encoders: Multilingual-e5-large and Alibaba-NLP/gte-multilingual-base were trained for HotPotQA and MuSiQue, while facebook/contriever was trained for RULER and BabiLong.

G EVALUATION DETAILS

The evaluation uses fixed generation settings across QA benchmarks, dataset-specific retrieval-step limits, and chunking tailored to embedder context and retrieval depth. It also evaluates retrieval quality with a coverage-based Fact EM metric that credits predictions containing all ground-truth facts.

  • Generation configuration: QwQ-32B generates HotPotQA and MuSiQue answers with temperature 0.0 and max tokens = 8000, while BabiLong and RULER use Qwen-4B with max tokens = 512 and reasoning disabled.These settings condition generation on the question and retrieved text chunks.
  • Retrieval configuration: T = 2 retrieval steps are used for HotPotQA, versus T = 4 for RULER and BabiLong, matching the limits used for Search-R1 and Beam Retriever.The step limit is the maximum number of retrieval steps.
  • Fact-level metrics: Fact EM equals 1.0 when predicted supporting facts contain all ground-truth supporting facts, granting full credit even when extra irrelevant chunks are included.It therefore tests subset coverage rather than exact equality between predicted and ground-truth fact sets.
Loading 2511.07328v2…