Q-RAG: Long Context Multi‑Step Retrieval via Value‑Based Embedder Training
Authors: Artyom Sorokin
Nazar Buzun
Aleksandr Anokhin
Egor KONSTANTINOVICH VEDERNIKOV
Petr Anokhin
Mikhail Burtsev
Evgeny Burnaev
Venue: ICLR 2026 Oral (2026)
Keywords: Reinforcement Learning
RL
QA
Long-context
RAG
NLP
Primary area: reinforcement learning

Abstract: 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.

Published as a conference paper at ICLR 2026
Q-RAG: LONG CONTEXT MULTI-STEP RETRIEVAL VIA VALUE-BASER EMBEDDER TRAINING
Artyom Sorokin 1.2.1 Nazar Buzun' 1.3, Alexander Anokhin', Egor Vedernikov
Petr Anokhin' MMikhail Burtsev", Evgeny Burnaey 1.2
AXXX, Moscow, Russia
2 Applied AT Institute, Moscow, Russia
3 Research Center of the Artificial Intelligence Institute, Innopolis University, Innopolis, Russia
'London Institute for Mathematical Sciences, London, UK
Correspondence to: griver29@gmail com', n. w..buun@seeeia..c
ABSTRACT
Retrieval-Augmenttd Generation (RAG) methods enhance LLM performance by efficiently filtering relevant context for LLMs, reducing hallucinations and infer- ence 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 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 Em- bedder model for multi-step retrieval using reinforcement learning (RL). Q-RAG offers a competitive, resource- resource-efficient alternative to existing multi-step retrieval methods for open-domain question answering and achieves tatte-of-the-arr results on the popular long-context benchmarks Babil.ong and RULER for contexts up to 10M tokens. Code is available at: https: /tttp:///ithubs com grrggriiiireerrr RAG
INTRODUCTION
Large language models (LLMs) have achieved impressive results across wide range of tasks NNovikov et al., 2025 Guo et al.) 2025 Yang et al. (025). However, they still face some sev- eral fundamental limitations such as static knowledge, computational inefficiency on long contexts, degraded performance caused by attention dilution, and hallucinations (Hsieh et al. 2024 Kuratov cet al. 2024 Liu et al., 2025). Aegmeeval-uugmented Generation (RAG) is one of the most widely used techniques to address these issues Yu et al. 2024).
RAG works by extracting only the most relevant parts from large external corpus or context, such as newly added knowledge lengthy texts. This allows LLMs operate shorter and more focused inputs, improving efficiency and output quality. Most current RAG methods rely on single-step re- trieval. This setup performs well in relatively simple tasks like Reedle-in-a-Haystaae (Hsieh al.) 2024). Still, more complex problems require multi-step interaction with the context. Multi-step retrieval can be viewed form of search-based reasoning. There are several existing approaches to multi-step retrieval reasoning. One direction involves constructing a knowledge graph from the retrieved information (Ma et al.] 2025 Li al.] 2024). These methods are often slow at infer- ence time, since the LLM must process the entire context to build the graph for each new input. Another line of work uses L.M agents, which interleave RAG queries with --n-erreeeeeed instruc- tions (Singh et al.) 2025 Anokhin et al., 2024). These systems are sensitive to noisy or inaccurate retrieved passages, which may disrupt the generation of future queries. This shows the need for joint optimization of the retrieval and generation components. Recently, methods have emerged that fine-tune LLMs to interact more effectively with retrieval tools (Song et al.) 2025 Jin et al. 2025] Chen al. 2025). These methods tend to perform better, but they require expensive fine-tuning of the LLM itself. This makes them impractical for large models and limits accessibility for most researchers and practitioners.

---

Published as a conference paper at ICLR 2026
In this work, we focus developing resource-efffcient multi-step RAG approach using reinforce- ment learning. Instead of fine-tuning an LLM, we train an agent that performs retrieval directly in the latent space of text chunk embeddings. This allows to learn compact and efficient model using value-based RI. methods.
Our approach achieves ofate-of-the-ar results long-context commonsence reasoning, multi-hop QA, and NIAH tasks with contexts up to 10 million tokens. It also performs competitively on open-domain QA benchmarks such as Musique and HotporQA (Yang et al. 2018 R). while being significantly faster and cheaper to train and run compared to existing multi-step RAG methods. Our contributions are the following:
We propose new method for training a multi-step retrieval agent using temporal difference reinforcement learning.
We achieve state-of-the-art results on benchmarks that require commonsense reasoning and NIAH tasks over ultra long contexts (up 10M tokens).
We introduce enabling temporal reasoning during retrieval. Our temporal reasoning mechanism general- izes well long contexts at inference time.
way incorporate temporal information into the multi-step embedder,
RELATED WORKS
There are several main directions for tackling complex retrieval scenarios long context tasks.
A highly popular approach involves building tine-tuning free LLM Agents that combine off-the-shelf retrievers with LLMs, such as Search-ol (Li et al. 2025). Many of these works further enhance retrieval quality by constructing large knowledge graphs over the context, which, while requiring little additional training, are extremely slow at inference due to the need for LLMs process the entire context, e.g. GraphReader (Li et al. 2024), HippoRAG (Jimenez Gutierrez et al. 2024), AriGraph (Anokhin et al. 2024).
Another line of work fine-tunes LRMs to perform multi-step retrieval, allowing the model to gener- ate intermediate search queries inside the reasoning for long contexts. The first work apply this idea was IM-RAG (Yang et al. 2024). which finettund the LLM with frozen embedder using PPO (Schulman et al. [2017]. More recent papers, such as R1-Searcher (Song al. (025), Search- R1 (Jin et al., 2025). RAG-RI (Huang et al., 2025), and ReSearcher (Chen et al. 2025), extended this direction by employing GRPO (Shao et al. (224)) for the task. Unlike these methods, which freeze the embedder and fine-tune the LLM, our approach fine-tunes only the embedder, allowing it to pair with LLMS of any size, including proprictary ones, while keeping fine-tuning efficient and inexpensive.
different approach is to fine-tune the retriever itself using feedback from the LLM, as in Re- Plug (Shi al.) 2024). This direction is most similar ours, but RePlug did not address multi-step reasoning or use reinforcement learning in this setting. BeamRetriever (Zhang et al.] 2024) achieves state-of-the-art results short-context QA by training reranker for BeamSearch-styyee planning. In contrast, Q-RAG trains the embedder with reinforcement learning, enabling faster inference and better scalability long contexts through efficient vector similarity instead of transformer-based trajectory scoring.
Extremely long-sequence processing is demonstrated by models that combine recurrence with Transformer architecture. The Mamba family of state space models (Gu & Dao 2024) replaces at- tention with structured recurrent dynamics, offering linear-time scalability and strong performance on long sequences, though often at the cost of weaker in-context learning and less expressive token- to-token interaction compared to Transformer-base
architectures. The Recurrent Memory Trans- former (RMT) (Bulatov et al. (022) introduces segment-level recurrence by passing memory tokens between fixed-size segments, enabling Q&A on sequences up to 10M tokens. Titans (Behrouz al.) 2024) frames recurrent memory training as a meta-learning problem and uses surprise prioritize information that should be retained over very long contexts, showing scaling beyond 2M tokens. Relatedly, MemUP (Sorokin al.) 2022) used uncertainty to identify events that require long- mer- memory in recurrent models. Similar to Titans, ATLAS (Behrouz et al. 2025) increases memory capacity, achieving better long-context performance than both RMT and Titans. The Associative
2

---

Published as a conference paper at ICLR 2026
Recurrent Memory Transformer (ARMT) (Rodkin et al. 2024) employs quasi-lineat associative at- tention in each layer and attains the best long-context scores among recurrent models. Our approach outperforms all of these models on contexts beyond 1M tokens while belonging to different class of methods.
LongRoPE2 (Shang al. 2025) tackles the positional encoding bottleneck, extending the effective context window of pre-trained LLMs to 128K tokens while retaining short-context performance through RoPE rescaling and mixed-window training.
3 METHODS
Q-RAG Agent
values
State Embedder
$
q
Qa(s,,c0
ag c
e,
-T((1(s)
eco
C c
Ecg
Cos
2a(s..c)
cil.
Embedder
La alterdee
He traly the
helhought slarsayy.
But spar sale.
ideve valued
sayed neg repeter's plaen.
Long-Contest Documant
tmestop
reward function entic
unal Life kaee Leavet calt. parmiay
S1+1
F141 S1t1
Environment
Figure 1: Q-RAG agent interacts with multi-step retrieval environment. The starting state sy contains the initial query At the start of the episode, the agent embeds all chunks of the long context C. At each step i, the agent computes a vector embedding of the current state &t. which includes and all previously selected chunks. For every chunk ef E A,, the utility of retrieving it is evaluated by the Q-function Qa(st, a c²). The policy ma selects the next chunk from As with probability proportional to its Qa(st, c²) value.
3.1 PRELIMINARIES
Let D be dataset of triples (C, y), where C is a long context, q is an initial query, and y is the gold answer. The query can be either a user question about C or generated claim whose factuality or consistency with earlier parts of must verified. We assume C is pre-segmented in document order. The agent's goal is identify the information in C that is missing from g but necessary to produce the correct answer y. We model multi-step retrieval is the action space, $ is the state space, is the reward function, is the (deterministic) transition function, and 7 [0. 1] the discount factor. an action at E A; selects one chunk. At later steps, previously selected chunks are removed so Superscripts indicate document positions and subscripts indicate episode timesteps. The notation (equivalently cl denotes the chunk/action at position in the document; selecting the chunk with index i step is written af. Symbols and used interchangeably, depending on context.
into non-overlapping text chunks C { p
finite-horizon Markov Decision Process, or MDP (S, A. p. r,7), where &
step 0, the action set is A C, where
= (ao....,.:......................
States are ordered lists that always begin with the query, s; ord ag. 1t-1)), where ord(-) sorts by the original document order to avoid permutation ambiguity; the initial state contains only the query, 8o [q)- Transitions are deterministie, p( 0t) ord([g, ag...., at-1, a.]). An episode terminates either when a step budget T' is reached or when a special STOP action is taken.
Chhnk overlapping may complicate the explanation but does affect our proposed solution.
3

---

Published as a conference paper at ICLR 2026
When supervision provides a set of support facts F C C, reward at all intermediate steps, and at the end of the episode it is 1 if all support facts are included in the final state (otherwise 0). When only answer supervision is available, one could instead use an LLM to generate ú from the final state and define a terminal reward via an answer- quality metric (e.g., exact match FI). In this work do not pursue L.M--based rewards; all reported experiments rely on the support-face signal, and exploring LLM-based reward design is left for future work.
sparse terminal reward: the
3.2 VALUE-BASED RL FOR EMBEDDER FINE-TUNING
Action selection in entropy reinforcement learning (Ziebart 2010 Haarnoja al.) 2018) is adopted together with the corresponding definitions the soft Q" and value functions for a policy π.
multi-step retrieval is performed by a value-based agent. Specifically, maximum-
(= a) = r(s,(s a) + (5 = p(s, a))
(1)
= (=(s)=Eou<(1) (s a) a log T(a
Here, 0 is a temperature that controls the strength of exploration. This choice primarily mo- tivated by the need for effective exploration in the long-context multi-step retrieval environment. In Q-RAG, Q function is approximated using two embedders for states and actions. The state embed- der E, (s;; 01) 1d produces vector embedding for the current state St, while the action embedder Eg (a', i; (2) E IRd employ rotary position embeddings to encode both the candidate chunk content and its secument-positiinn index i. Q values are then estimated by inner product between two em- beddings: Qo(s, a') = (s; f2), (a , i; 02))- This factorization is theoretically grounded; we derive its convergence guarantees with explicit rates in Appendix A Given Qe, the chunk selection probability is computed using a Boltzmann policy:
exp (Qo(se, de) -q)
T(arls:) E.EA, exp _ (Qo(se, a)
with q = A, (s( a) and temperature annealed from an initial value to zero during train- ing (proportionally to the learning rate).
As the backbone Temporal Difference learning algorithm, we adopt the recent PQN method by Gallici et al.) Compared to DQN (Mnih et al. 2015), PQN removes the need for replay buffer. In our setting with a large number chunks, a replay buffer would require re-embedding all document chunks for each sample drawn from the replay buffer to estimate V/Q values for subsequent states
81+1- Which significantly slows the training process and increases memory space requirements.
Using PQN enables on-policy value-based training that avoids these costs. The key departures in Q-RAG, relative to the original PQN backbone, are the use of soft value functions and target networks. Ablation results demonstrating the benefit of these choices are reported in Section
As the training target, rather than the one-step return (see r.h.s. in Eq. D. a X-return used to improve stability and learning speed:
T-1-1
G. = (1- A) Σ
x n-1 Gettn + AT-
n=1
where Gettn Σ k== 1 re+h Vor (8t4z). The approximation of the state value function can be computed from Q values in the case of discrete actions:
(5t, a)
(s,) log Σ exp
(4)
a
Here 0' denotes slowly updated target network parameters. The model parameters 0 are finetuned to minimize the mean squared error to the A-returns:
Le [[(QQ(ss, 0z) -
(5)
The Q-RAG pseudocode is presented in Algorithm
4

---

Published as a conference paper at ICLR 2026
Algorithm Q-RAG
1: Hyperparameters:
2: Environments count K, retrieval steps T, temperature a, TD parameter X, EMA 7.
3: Initialize:
4: State embedder [s;
5: Action embedder Fig (a', i; (2) with position
6: Critic Qe(s, a²) = (s; 01) Ea
7: Critic target Qe,
8: procedure COMPUTETARGETS/
oe, re It= 17+1)
9:
Initialize X-returns Gr PT + yU++1 for / = T downto do
11:
G₁ - + ((1 A)e + AG(+1] end for
13: return [G+]
end procedure
15: Training (one update step)
16: for env k E l...., ; K parallel do
17:
$1,A1 ResetQuery RuseeQueeyAAnCConeee AndCotext()
Compute Egg Eg (A; 0) and Ey = (A( for step E .,...,
19:
do
20:
at ~ softmax 1 (3:0)²
21:
U: a log MaeA exp Es; EE,
22:
T = ComputeReward((s,
23:
81+1 concatenate (st, (4)
24:
A++1 = Ae \ {at)
end for
26:
B = {st, T2, ut JeT
27:
{G] )
(onputeTargeets(B)
28:
end for
29:
VLQ= TK
E Update 9 using VLQ
ET= ve(Qo(sf att Gb)²
31:
Update target parameters: 6 +- r0 + +1 -)
3.3 TEMPORAL REASONING FOR LONG-CONTEXT SEARCH
When dealing with narrative text, the information contained text chunk e may be insufficient determine whether e helps answer the question q. For example, we may need to know what happened before some specific event. A standard retriever can find several relevant text chunks that specify the character's location, but choosing the correct account temporal information. To address this, we propose a relative postional encoding of chunks that explicitly encodes their position with respect the facts already extracted into the state. At step t, let S, {ur < ... it} be the (sorted) document indices of selected chunks and A, the set of available actions. The indices in S, partition the document into k+1 disjoint intervals: "before the earliest selected fact", "between consecutive selected facts", and "after the latest selected fact." The relative positional mapping pt N + R assigns to every original chunk index real-valued index that (i) identifies the interval it belongs to and (ii) preserves the relative order between chunks. This mapping makes explicit between which extracted facts a chunk lies, while remaining invariant to global shifts of absolute positions.
can be impossible without taking into
Formally, the interval boundaries are defined as by=1, by=i, for j=1:k, and bx41=m+1 for (c(0) ym To compute relative index (i) for a chunk c², find the unique j such that by - i < hi+1 and set
b Pr (+()) jō bit1 by
where 8 > 0 is the inter-interval step and E (0, 8) controls the within-interval resolution (c.g., 1=10, 8= in our experiments). In the action embedder, the absolute position is replaced by the
5

---

Published as a conference paper at ICLR 2026
relative
Ea (a', 02 = Ea(a', 82).
(7)
which allows the Q-function to exploit the spatial relation candidates to already retrieved evidence while retaining local order within each interval. This design allows the retrieval agent to perform strongly not only on fact-finding over disjoint document collections, but also on long-form narrative
tasks, enabling Q-RAG compete with recurrent transformers (Bulatov et al. 2022 Rodkin al. 2024
Behrouz et al. 2025
2024) and other long context approaches.
4 EXPERIMENTS
4.1
EXPERIMENTAL SETUP
We evaluate our approach, Q-RAG, on tasks that cover commonsence reasoning, temporal reason- ing, a bunch of needle in a haystack tasks and open-domain multi-hop question answering tasks on context lengths that range from 4k tokens 10M tokens per sample. For commonsence and tempo- ral reasoning we use BabiLong benchmark (Kuratov et al. (002)), for Needle-in saaee-inn---aystta we use RULER benchmark Hsieh et al. (2024). For open-domain multi-hop QA we use HotpotQA Yang et al. (2018), Musique Trivedi et al. (2022) and RULER benchmarks. BabiLong and RULER re- quire long contexts. Musique and HotpotQA use short contexts.
Baselines differ by task. Computing uniform set of baselines across all datasets is difficult and time-consuming. Many methods do not release code. Some methods were evaluated only on some of these datasets. Even when the tasks match, the experimental settings often differ for the same benchmarks. Some baselines provide code but require heavy resources (e.g., at least AxA100 GPUs Jin et al. (2055; ; Song et al. 2025)) ; Huang et (2025) to fine-tune, which unavailable for Therefore, we report three types of baselines, and we mark each baseline in tables accordingly:
Ablation: baselines that test the effectiveness of our proposed modifications.
Reproduced: baselines that we finetuned and/or evaluated on our datasets using released code or publicly available checkpoints.
Reported: baselines whose scores we take directly from the original papers.
4.2 COMMONSENSE REASONING ON ULTRA-LONG CONTEXTS
On the BabiLong Kuratov et al. 20244) benchmark, we compared our method with the state-of-the- art long-context processing approaches, including Titans Behrouz et al. (2024), Atlas Behrouz et al. (2025), ARMT Rodkin et al. (2024), RMT Bulatov et al. (202)), as well as proprictary LLMs and LLM-based agents. The results for most of these baselines were taken directly from the respective original papers. As shown in Figure 2b our approach achieves the highest average performance
16-
BabiLong QA1-0A5
suaporting
08-
0-4
LOKK
22M
Context Length
Context Length WAYT
1A0 Dwars 48
-DYEV-22
:PP44
(a)
(b)
Figure 2: Comparison of answer accuracy on the long-context benchmark Babil Babil.ong. Solid lines de- note methods fine-tuned on the BabiLong, while dashed lines denote zero-shot methods. Average performance across tasks QI-QA5. b) Performance on the hardest task, QA3, which requires the longest reasoning chain and temporal awareness.

---

Published as a conference paper at ICLR 2026
on BabiLong in ultra-long contexts ranging from 10 million tokens, demonstrating superior generalization to long contexts compared to other specialized long-context methods.
present separate results for the QA3 subtask, which is the hardest subtask in the BabiLong benchmark, which specifically requires the multistep search of at least different facts and temporal reasoning. Experimental results show that the majority of models perform worst the QA3 subtask. As the results indicate, alternative long-context approaches show even greater perfor- mance degradation on this task with increasing context length. In contrast, Q-RAG shows virtually no degradation, with the largest performance gap over all baselines observed on this most chal- lenging subtask. We additionally fine-tuned the Beam-Reeriee subtasks, given its strong performance on open- ppen-domain QA datasets. However, this method failed to solve the task. Note that some methods, such as Titans Behrouz et al. (2024) and Atlas [Behrouz et al.. 2025). are absent from the Figure as they did not report detailed breakdowns by subtask.
In Figure 2al
baseline specifically on the QA3
4.3 NEEDLE IN A HAYSTACK AND LONG CONTEXT QA
While reasoning tasks are crucial for evaluating advanced retrieval systems, substantial portion of real-world applications reduces to ieedle-in-a--ee -atta--aaytaakk (NIAH) problems, making it equally important that models deliver consistently strong performance on these tasks. RULER is dataset that includes many long-context tasks. Most of these tasks follow the NIAH formulation. The NIAH setup evaluates the ability to retrieve specific "needle" from a long distracting "haystack". For RULER benchmark
use Titans Behrouz et al. (2024), Atlas Behrouz et al. (2025). Mamba2 Waleffe et al. (2024), and LongRope2 Shang et al. (2025) as baselines. Titans, Atlas are recurrent transformers. Mamba2 is state space model (SSM) that combines transformer components with SSM. LongRope2 is a method for extending the effective context window of LLMss. All methods
Table 1: Results on the RULER benchmark, evaluating long-context retrieval performance across various context lengths. S (Single-needle): Find one value for one key, MK (Multi-keys): Find one value for one key among many. MV (Multi-values): Find all values for one key, MQ (Multi-query): Answer multiple questions over the context. MH QA: open domain multi-hop question answering.
SH QA: single-hop question answering.
S
MK
NIAH
QA
Len
Methods
MV
MQ
Avg.
1-st
2-nd
3-rd
1-st
2-nd
3-rd
SH
MH
Tiitans
98.4
99.8
89.4
n/a
n/a
n/a
n/a
n/a
n/a
n/a
n/a
Atlas
99.2
100
90.6
n/a
n/a
n/a
n/a
n/a
n/a
n/a
n/a
"Mamba2-Hybrid
100
100
95.7
89.5
95.5
96
97.9
97.6
96.5
56.5
48.8
4K
LLogRoPP2-88
100
100
99
100
100
100
99
99.7
99.7
79
60
Beam-Retriever
100
100
98
98
98
97
98
99
98.5
29.0
39.0
Q-RAG
100
100
100
100
100
100
100
100
100
62
67
TTitans
96.2
80.2
n/a
n/a
n/a
n/a
n/a
n/a
n/a
n/a
n/a
*Atlas
97
84
n/a
n/a
n/a
n/a
n/a
n/a
n/a
n/a
n/a
"Mamba2-Hybrid
100
100
81.5
92
92.2
83
89.8
90.2
91.1
48.8
44
16K
"LongRoPo2-8B
100
100
100
99
100
98
95
98.2
98.8
69
58
Beam-Retriewe
100
100
97
96.5
96
95
80
98
95.3
24.0
35.0
Q-RAG
100
100
100
100
100
100
100
100
100
59
64
"Manba2-Hybrid
100
100
96.7
84
76.5
81.5
84.3
80.9
88.0
41.8
38.5
32K
"LongRoPe2-8B
100
100
100
99
98
100
98
96.2
98.9
72
55
Q-RAG
100
100
100
100
100
100
100
100
100
59
65
PLongRoPe2-8B
100
100
96
94
96.5
97
96.7
56
50
128K
Q-RAG
100
100
100
100
100
100
100
100
100
55
65
IM
Q-RAG
100
100
100
100
98.5
99.0
100
100
99.7
52
61

---

Published as a conference paper at ICLR 2026
were fine-tuned either directly RULER (Titans, Atlas, Mamba2) related synthetic NIAH- style datasets (LongRope2). Q-RAG was also fine-tuned on the NIAH subtasks. For the Multi-hop QA RULER subtask, Q-RAG fine-tuned on HotpotQA and evaluated the Multi-hop QA subtask out-of-distribution.
The results shown in Table Q-RAG achieves near-perfect performance all NIAH subtasks. Q-RAG embedder was trained on 4K-length documents and generalizes to context lengths up to 1M tokens without loss of accuracy. On the Multi-hop QA subtask, Q-RAG shows significantly better results than all our baselines at all context lengths we consider. Some degradation with increasing context length starts only from 128K.
4.4 DHEN-DOMAIN QUESTION ANSWERING
For our experiments on the HotPoiQA and Musique datasets, we compared our method against several strong baselines. The first baseline is Beam Retriever, which enables multi-step retrieval by training model score sequences of retrieved chunks. During evaluation, Beam-Retriee is given the oracle number of supporting facts (i.c., the gold hop count) and always retrieves exactly that many facts. Although this approach is slower than traditional retrieval methods and does not scale well to longer contexts, it achieves state-of-the-are results HotPotQA. Another baseline we considered is SearchR1, a recent method from a family of approaches that train the LLM itself to compose text queries for multi-step retrieval. Additionally, we evaluated the performance of LLM-agent-basse methods, including GraphReader. Q-RAG and on HotPotQA and evaluated Musique for cuttrf-distribution testing. Baseline numbers were taken directly from the corresponding papers. Missing entries indicate metrics not reported by the original authors.
were fine-tuned
The comparison results are presented in Table Our method achieves fact retrieval accuracy on par with Beam Retriever, surpasses all other baselines HotPotQA, and matches the performance of full-LLM-tuning Search-RI while outperforming all alternatives on the out-of-distribbtiion Musique dataset, resulting in the best overall performance across benchmarks. Results also include another Q-RAG version Plan Q-RAG that combines Q-RAG value function and beam search based planning (see Appendix Cp. Plan Q-RAG showed similar performance to vanilla Q-RAG, For both methods involving retrieval mechanism fine-tuning (Q-RAG and Beam Retriever), we used the QwQ-32B model produce the final answer.
Table 2: Comparison of methods on HotPotQA and Musique benchmarks. Bold text and underline denote the best and second best scores respectively.
HotPotQA
Musique (OOD)
Avg
Methods
Fact F1 Fact EM Ans F1 Ans EM Fact FI Fact EM Ans Fl Ans EM Ans FI
Ans EM
Finetuned on HotPorQA
Plan Q-RAG
0.95
0.91
0.76
0.60
0.69
0.53
0.51
0.36
0.64
0.48
Q-RAG
0.93
0.89
0.76
0.59
0.71
0.55
0.52
0.37
0.64
0.48
Ream-RRtriever
0.97
0.94
0.77
0.61
0.61
0.36
0.40
0.27
0.59
0.44
Search-rl
0.81
0.66
0.65
0.52
0.71
0.55
0.51
0.39
0.58
0.46
"RAG-RL
0.82
0.69
0.47
0.58
"Multistepp RAG w.o. FT
0.73
0.54
0.65
0.50
0.51
0.30
0.40
0.27
0.53
0.39
Zero Shot methods
GraphReader
0.46
0.24
0.40
0.20
0.43
0.22
Single step RAG
0.53
0.39
0.28
0.17
0.41
0.28
ABLATION STUDY
To assess the impact of the architectural choices in Q-RAG, an ablation study was conducted on the BabiLong- -A3 QA3 task. This benchmark was selected because it is among the most challenging long-context tasks used in the experiments and it supports evaluation at arbitrary context lengths. The following baselines were compared against Q-RAG:
8

---

Published as a conference paper at ICLR 2026
Inference Time (QWO-322)
Beam Butrover
2.97
207
O-AG
TLime (tec)
Ono 328
0.50
202
2.03
0.68
0.88
042
233
Q63
0.02
0.61
0:51
0.5
Tempanitue parameter
TD paramittr
20k
120k
140k
(a)
(b)
(c)
Figure 3: Ablation for (a) policy entropy coefficient (a) in soft Q function and (b) for X-return parameter. Inference runtime comparison (c), context length, tokens on x-axes.
Multi-step RAG w.o. FT. This baseline reproduces the full Q-RAG retrieval pipeline and uses the same state and action embedders, but relies their original pretrained weights without any reinforcement learning fine-tuning. This setting tests whether RL. fine-tuning of the embedders is beneficial for multi-step retrieval quality.
Multi-step RAG W. SFT. This baseline applies supervised fine-tuning using ground-truth support facts as supervision. The loss follows the objective used Retrivere for trajectory supervision, adapted to the multi-step retrieval setting. This setting isolates the effect of RL by comparing it to supervised learning on the same supervision signal.
Q-RAG w.o. target. This variant removes target networks from the PQN-based value learning, following the original PQN recipe without target parameters. It measures the contribution of target networks stability and performance in the Q-RAG training loop.
Q-RAG w.o. Soft-Q. This variant replaces the maximum-entropy (soft) value functions with stan- dard (non-entropy-ee regularzee)) Q-learning objectives. evaluates the effect entropy regulariza- tion and the soft value formulation retrieval performance.
All baselines were evaluated with three random seeds. Table B reports results at a 32ketoken con- text length on QA3. Figure 3 shows the sensitivity of Q-RAG the X-return parameter and the temperature a (the strength of entropy regularization) on QA2 and QA3.
Table 3: Ablation results on BabiLong QA3, Table shows Fl score for support facts retrieval. All values are averaged over runs with different seeds.
Method
1K
4K
32K
128K
1M
Q-RAG
97.8 0.17
9714
97.1 0.08
96.8
1
0.08
96.5
0.16
"QRAAG w.o. Soft-Q
95.9 - 0.70
95.5
94.5
0.50
94.0
0.30
93.3 1 0.45
Q-RAG w.o. Target
79.2 = 26.0
78.1
77.4 77.6
75.9 28.2
Multi-Step RAG w. SFT
20.33 = 0.32
20.87
0.200
18.30
0.36
Multi-Step RAG w.o. FT
15.52 1 0.11 16.38 1 0.10
15.51
0.16
15.34 1.3410.15 0.12
5.1 SENSITIVITY TO RETRIEVAL BUDGET
We investigate the dependence of final model performance on the number of Q-RAG retrievals (i.e., the retrieval budget). For this analysis, we used Q-RAG system with an Aligaha-NL.P/tt- multilingual-ban embedder, trained embedder supports contexts of up to 8192 tokens, enabling the use of a larger retrieval budget. We evaluated the system 1000 samples from the HotpotQA dataset. The final generation of the answers was performed by three LLMss: Qwen3-4B, Qwen3-14B 14BB, and Qwen3-32B.
combination of the HotpotQA and Musique datasets. This
The results presented in Table4] Here, EM (Exact Match) indicates the number of correct (ground-trutte truth supporting) chunks retrieved, while Fl accounts for the inclusion noise (non- supporting) chunks. The table shows that increasing the number of retrieves from 2 to 3 improves

---

Published as a conference paper at ICLR 2026
Table 4: Sensitivity to the number of retrieves. Dataset: HotpoiQA (1000 samples). Embedder
Aibbab--NLPggt-muullllng mareeei-ullllilllbbaaa was trained on Hotpotqa+Musique.
Facts
Qwen3-4B
Qwen3-14B
Qwen3-32B
Retrievals
EM
F1
EM
F1
EM
F1
EM
F1
2
0,832
0.903
0.439
0.620
0.556
0.708
0.504
0.675
3
0,935
0.771
0.481
0.657
0.570
0.730
0.510
0.692
4
0.962
0.652
0.493
0.664
0.577
0.734
0.513
0.695
5
0.978
0.565
0.481
0.656
0.584
0.744
0.512
0.692
both the number of correct facts retrieved and the answer quality across all three LLMs. These experiments suggest that, within reasonable range of retrieval counts, final answer accuracy is pri- marily dependent the retrieval of correct chunks and is not degraded by the presence of noise chunks.
In addition the fixed-budget setting, we also studied alternative stopping criterion in which the agent stops dynamically according to Q-value threshold. A detailed analysis of this Q-value-based early stopping is presented in Appendix B
CONCLUSION
This work introduced Q-RAG, a resource-efficient method for multi-step retrieval trained with re- inforcement learning directly in the latent space of text-chunk embeddings. Across long-context benchmarks (c.g.. BabiLong, RULER) and open-domain QA datasets (e.g., Musique, HotporQA), Q-RAG attains state-of-the-ar or highly competitive results. Its advantage over baselines widens as context length grows, and performance shows minimal degradation even at ultra-long scales.
A key practical benefit is compute efficiency: all training was performed on a single A100 GPU with 80 GB memory, whereas recent RL-based multi-step step retrievers such Search-RI -1/---RRRRRI--h-- typically report training on clusters of about eight A100 GPUs. By fine-tuning only the embedder while keeping the L.MM frozen, Q-RAG remains easy to pair with powerful pre-trained or proprietary LLMs, enabling efficient training, flexible deployment, and strong retrieval over very long contexts.
Looking ahead, promising directions include using structured LLM feedback reward signal, strengthening compositional and temporal reasoning directly in the embedding space, and exploring tighter infegration with generation while preserving the method's efficiency and scalability.
REPRODUCIBILITY STATEMENT.
We make all results reproducible by providing a code package with exact configs and run scripts; Code is available at: htttps://github.. wttp://wt///goo//oeeoo.eee
The package in- cludes utilities to download and minimally preprocess the public HotPotQA and MuSiQue datasets and to re-run every experiment and table with fixed random seeds (Fact FI/EM, An- swer F1/EM). We fine-tune only publicly available embedders - mwltilingaalee--eaaa. Alibaba NNLP /tte ceeeemuliill wwtlttiiinual-bbba lowing the hyperparameters and schedules described in Appendix F All reported runs are repro- ducible single GPU; our main experiments executed one A1 11000-8 80CB device.
and wwweebook/cottrieee
strictly fol-
ACKNOWLEDGMENTS
We thank Oleg Inozemcev for running additional experiments and assisting with hyperparameter tuning during the rebuttal stage, which helped improve several results. We also thank Alexey Trushkov and Wenshuai Yin for their valuable feedback and constructive criticism, which helped improve the paper. The work partially supported by the grant for research centers in the field of Al provided by the Ministry Economic Development of the Russian Federation in accordance with the agreement 000000C31 30SC333099554F000 and the agreement N' 19-1002025-033
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Petr Anokhin, Nikita Semenov, Artyom Sorokin, Dmitry Evseev, Andrey Kravchenko, Mikhail Burt- sev, and Evgeny Burnaev. Arigraph: Learning knowledge graph world models with episodic memory for Ilm agents. arXiv preprint arXiv:2407.0453 2024.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time, arXiv preprint arXiv:2501.0066 2024.
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Raza- viyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time. arXiv preprint arXiv:2505.23735. 2025.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:1 1079-11091, 2022.
Mingyang Chen, Tianpeng Li, Haoze Sun, Yijie Zhou, Chenzheng Zhu, Haofen Wang, Jeff Z Pan, Wen Zhang, Huajun Chen, Fan Yang, et al. Learning to reason with search for Ilms via reinforce- ment learning. arXiv preprint arXiv:2503. 19470, 2025.
Matteo Gallici, Mattie Fellows, Benjamin Ellis, Bartomeu Pou, Ivan Masmitja, Jakob Nicolaus Foerster, and Mario Martin. Simplifying deep temporal difference learning. In The Thirteenth International Conference on Learning Representations.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024.
URL http:ttss////xxii....o wttp:/wwwwwwogg/o/////eb// 00752
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.1294 2025.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning. pp. 1861-1870. Pmir, 2018
Cheng-Ping Hsieh, Simeng Sun. Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's 'sss the real context size of your long-context language models? arXiv preprint arXiv:2404.0666 2024.
Jerry Huang, Siddarth Madala, Risham Sidhu, Cheng Niu, Hao Peng. Julia Hockenmaier, and Tong retrieval-augmentte augmented generation via rl and curriculum learning. arXiv
Zhang. Rag-rl: Advancing retrieval preprint arXiv:2503. 12759 2025.
Bernal Jimenez Gutierrez, Yiheng Shu, Yu Gu, Michihiro Yasunaga, and Yu Su. Hipporag: Neurobi- ologically inspired long-term memory for large language models. Advances in Neural Information Processing Systems, 37:59532-59566 2024.
Bowen Jin, Hansi Zeng, Zhenrui Yue, Jinsung Yoon, Sercan Arik, Dong Wang, Hamed Zamani, and Jiawei Han. Search-rl: Training 1lms to reason and leverage search engines with reinforcement learning. arXiv preprint arXiv:2503.09516, 2025.
Yury Kuratov, Aydar Bulatov, Petr Anokhin, Ivan Rodkin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. Babilong: Testing the limits of Ilms with long context reasoning- ioninn--i-aa in-a-haystacc. Advances in Neural Information Processing Systems, 11111065991065555 2024.
Shilong Li, Yancheng He, Hangyu Guo, Xingyuan Bu, Ge Bai, Jie Liu, Jiaheng Liu, Xingwei Qu, Yangguang Li, Wanli Ouyang, et al. Graphreader: Building graph-based agent to enhance long- context abilities of large language models. In Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 12758-12786. 2024.
Xiaoxi Li, Guanting Dong, Jiajie Jin, Yuyao Zhang, Yujia Zhou, Yutao Zhu, Peitian Zhang, and Zhicheng Dou. Search-ol: Agentic search-enhanced large reasoning models. arXiv preprint NXiv:2501.05366 2025.
11

---

Published as a conference paper at ICLR 2026
Jiaheng Liu, Dawei Zhu, Zhiqi Bai, Yancheng He, Huanxuan Liao, Haoran Que, Zekun Wang, Chenchen Zhang, Ge Zhang. Jiebin Zhang, et al. A comprehensive survey on long context lan- guage modeling. arXiv preprint 25X3::5503.7707. 2025.
Chuangtao Ma, Yongrui Chen, Tianxing Wu, Arijit Khan, and Haofen Wang. Large language mod- els meet knowledge graphs for question answering: Synthesis and opportunities, 2025. URL
https::/arxiv.com wwwp:iii.rrgaab//////e/o 2505. 20099
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 8(7540)::29-533 2015.
Alexander Novikov, Ngån Vü, Marvin Eisenberger, Emilien Dupont, Po-Sen Huang, Adam Zsolt Wagner, Sergey Shirobokov. Borislav Kozlovskii, Francisco JR Ruiz, Abbas Mehrabian, et al. Alphaevolve: A coding agent for scientific and algorithmic discovery. arXiv preprint arXiv:2506. 13131, 2025.
Ivan Rodkin, Yuri Kuratov, Aydar Bulatov, and Mikhail Burtsev. Associative recurrent memory transformer. arXiv preprint wxXiv:2407.04844, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Ning Shang, Li Lyna Zhang, Siyuan Wang, Gaokai Zhang, Gilsinia Lopez, Fan Yang. Weizhu Chen, and Mao Yang. Longrope2: Near-lossless Ilm context window scaling. arXiv preprint 2000, 2025.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song. Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.033300 2024.
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Richard James, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. Replug: Retrieval-augmenned black-box language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies (Volume Long Papers), pp. 8364-8377, 2024.
Aditi Singh, Abul Ehtesham, Saket Kumar, and Tala Talaci Khoei. Agentic retrievallauugmente generation: A survey on agentic rag, 2025. URL
09136
Huatong Song, Jinhao Jiang, Yingqian Min, Jie Chen, Zhipeng Chen, Wayne Xin Zhao, Lei Fang, and Ji-Rong Wen. RI-searcher: Incentivizing the search capability in 1lms via reinforcement learning. arXiv preprint arXiv:2503.0553 2025.
Artyom Sorokin, Nazar Buzun, Leonid Pugachev, and Mikhail Burtsev. Explain my surprise: Learn- ing efficient long-term memory by predicting uncertain outcomes. Advances in Neural Informa- fion Processing Systems, 35:36855-36888 2022.
Harsh Trivedi, Niranjan Balasubramaniaan Tushar Khot, and Ashish Sabharwal. Musique Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539-55: 2022.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba- based language models. arXiv preprint arXiv:2400.078887 2024.
Chenghan Yang, Ruiyu Zhao, Yang Liu, and Ling Jiang. Survey of specialized large language model. arXiv preprint arXiv:2508. 19667, 2025.
Diji Yang, Jinmeng Rao, Kezhen Chen, Xiaoyuan Guo, Yawen Zhang, Jie Yang, and Yi Zhang. Im-rag: Multi-round attrieval-augment generation through learning inner monologues. In Pro- ceedings of the 47th International ACM SIGIR Conference on Research and Development in In- formation Retrieval, pp. 730-740, 2024.
12

---

Published as a conference paper at ICLR 2026
Zhilin Yang. Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing, pp. 2369-2380, 2018.
Tan Yu, Anbang Xu, and Rama Akkiraju. In defense of rag in the era of long-context language models. arXiv preprint arXiv:2409. 2411:2409.01666, 2024.
Jiahao Zhang, Haiyang Zhang. Dongmei Zhang. Liu Yong. and Shen Huang. End-to-end beam retrieval for multi-hop question answering. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies (Volume 1: Long Papers). pp. 1718-1731, 2024.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mellon University, 2010.
13

---

Published as a conference paper at ICLR 2026
A INNER PRODUCT APPROXIMATION FOR Q-FUNCTION
The classical Universal Approximation Theorem (UAT) asserts that sufficiently expressive neural networks approximate any continuous function compact domain arbitrarily well. In this section, we prove variant of the UAT for functions decomposed as an inner product between state embeddings and action embeddings modulated by positional block-diagonal matrix.
Let X C Rd,, Y C Rdy and 7 C R be compact sets and define K = X x Y x T. We will approxi- mate any E C(K, RR) in the uniform norm. One may identify X with the environment state space S. Y with the set of available actions A (see Section B). and interpret E T as relative positional encoding for action Under this correspondence, a function f(z, y,t) represents ground-truth Q-function. If Q- 2-fuuctiin does not depend on position t, one can simply take f to be constant in t.
Our starting point is the real-valued score function
F (x, y, 1) = (1x((x), R. gr((y))
(8)
where m E N is arbitrary, /ıg X + R and gR Y + 222m are continuous encoders (e.g., neural networks), and R2 E 2m x2m is dosition-dependen block rotation acting independently on each coordinate pair. The standard Rotary Position Embedding (RoPE) is precisely family t + Rs of this type.
A useful reformulation is that every real-valued score of the form equation can be written as the real part of a complex inner product after identifying 2m C Under this identification, the RoPE block rotation Re corresponds to multiplication by diagonal complex matrix A(t), which both clarifies the structure and motivates the more general complex-diagone score class considered below:
F(x, V,, = Re((M()), (())())))
A(I) = ( @m(()))
where h E C(X, ("), E C(Y, C"), and each diagonal entry Φy is drawn from a function algebra C(T,C).
We first show that every real block-rotation score equation 8 (in particular, RoPE) can be written in complex-diagonal form, with (t)) = diag(e"f1t, ..... cm"mtt for suitable fixed frequencies 0z.
We then prove that, whenever 4 is sufficiently rich (self-adjoint) subalgebra of C(T, C), the in- duced
class is dense in C(K, R) in the uniform norm. This yields the desired universal approximation property; standard RoPE is recovered as special case by choosing & to contain the relevant exponentials.
Real-valued block rotations as complex-diagona operator. Fix N. Define the real-to- complex identification p 2m + C. by
p((an 02 ....... a2m- -1 02m) T) = (a1 + iaz, a3 + ias, 02m + id2m
(10)
m We use the standard Hermitian inner product on C"". (u, v)em u²v k= Wk Uk- Then for any a, b E 2m
(a, b) = Re((p(a),p(())
Let .1,..., . On IR. be fixed frequencies and let Rd 2m2m22m be the block-diagonal rotation
(cos sin
R = dia)( R(0ml)),
R(a)
sin cos a
Define the complex diagonal matrix
A(t) = diag(e 1014
e Cmxm
A direct check on each 2 x 2 block shows that pl R22) = A(t) for all z E R Consequently, for any real-valued encoders hg
2m and g3 . 2m
(/R (I), RegR (y)) = Re((o(ha(T)), A(t) (1))
(11)
In Section 3.3 the positional encoding is denoted by p(i). Here avoid this notation to prevent confusion with the map p IR
T used below identify R22h ex C.m
14

---

Published as a conference paper at ICLR 2026
In particular, whenever algebra & C C(T, C) contains the exponentials t + eliat, the real-valued block-rotation score equation 8 (and hence standard RoPE) is a special case the complex-diagond score class As defined next.
Theorem (Complex Inner-product approximation with diagonal positional matrix). Let X C IRd,, Y C. Rd and T C. IR be compact sets, and K = X x Y x T. Let d C C(T, C) be a self-adjoint subalgebra (t.e,, & E Φ + Φ E )) which contains constants and separates points of
I D
A4 ta 3 ((h1) + ((t2)
Define the function class
Ap = U ((2x0t)==RR(((e))A(((00)0)))))) (( g C(Y,C)) A E A
MEN
(12)
where
Ag.m = { A Cmxm A(t) dt)
OR o).
and (u, t)em == u u. Then As is dense in C(K, RR. in the uniform norm. Equivalently, for any f ((K,K) R) and any > there exist m, h, and A as above such that
sup (x..))K
f(z,8,t) R(((bba)) <E.
(13)
Proof. Consider the set
JEN, ua E C(X, C), E C(Y,)) de E c0) cC(K,C)
B- un(x)) vx(v) dx(l)
Density of C(K, C). First we show that B is a self-adjoint subalgebra of C(K, C) that con- tains constants and separate points. Closure under addition and scalar multiplication is immediate. To check closure under products, take
y' b(c,,,,1) = Σ ūz(x) ( 5c().
b(r v. () = Σ (()
k=1
with uk, ue E C(X, C), Uki 8e C(Y,C), or, δe Φ. Then
(b-b))=,
(ux (uxūe( (x) (Ux (v) (t)
Since C(X, C) and C(Y C) are algebras under pointwise multiplication, u E C(X,C) and Un C(Y, C). Because d is subalgebra, Φz E 4. Hence b D' E B, so B is an algebra.
Also 1 E B by taking u 1, v = 1, ø Self-adjointness follows from pointwise conjugation:
u u<))
((z) (())
and the assumptions ò & and i C(X C), T C).
Additionally B separates points K. Let (z1, V/, (1) * (I2, 92, t2). If 21 * x2, choose coordi- nate projection u(x)= xj with æ1,5 * ø2.j and set If t * f2, use the assumption that @ separates points of T: pick ø E with ((t1) # (t2)) 1.
ò = 1; then u(x1) * u(x2). Similarly if
v1 + y2- and set
By the complex points), B C(K,C) C).
theorem (self-adjoint subalgebra, contains constants, separates
15

---

Published as a conference paper at ICLR 2026
Real parts of B lie in Ag. Take any b B:
U1, Σ (1))
Let m J, h(x) = (,())))))
s(v) = w,(())
A(t) = (t).
Using (a, 6) = aab,
A(t))) Σ k=
hx( (x()) Σ 1om
= b(2, v.),,
(14)
hence Re(b) As.
Approximation of real-valued functions. Let f C(K, R) and O. View element of C), Pick b E B such that |lf bl < E. Define = Re(b). Then, using [Rez| s |2],
|| - ol = ||f Re(b)]l = |Re({ - b)ll> E |I - bll00 <.
(15)
As shown above F E Ag. which proves density in C(K,RR)
The qualitative approximation property established in Theorem I does not reveal how the required feature dimension scales with the desired accuracy. first step toward quantitative understanding Lemmal which gives rate dds/(da+d,) for approximating stationary kemels in H Y). By incorporating this lemma into a Fourier analysis of the temporal variable, we obtain the following theorem that handles the full (z, y, t) dependent case and achieves the rate d==/d=/=sdd =dd_,
Lemma (L² low-rank approximation of Sobolev kernels). Lei X Rdd and Y C BRdy be bounded Lipschitz domains, and set D dz + dy- Let > O and let H* x Y) be real(or complex)- valued. Then for every integer d 2 1 there exist measurable feature maps
h (XXC()
9 gE LY
(16)
such that
- - (h(-)) g.)e (ee SCdD /a [|]|==((xx))-
(17)
Here C > 0 depends only on s, dus dy and the geometry of X,Y.
Proof. Consider the bounded Lipschitz domain X Y RD. By Sobolev extension
theorem for bounded Lipschitz domains, there exists bounded linear operator
:: (Ω) (R.)
(18)
such that Eulo= for all (Q). Apply it to let == Ea H*RR"), then
| ([()) S Cexa|lallH--00)
(19)
Choose a cube Q C IRd with side length L such that H7 C Q. Take a cut-off function x E C (Q)
satisfying X on Ω and < x 1, with supp X E Q (in particular, dist(supp X. aQ) 0). Define u := X a. Then u E HB)),,, supp C Q, and u/e a. By boundedness of multiplication by a smooth compactly supported function,
| S CxI S Cezalalall
(20)
Now periodise u. Since supp is compactly contained Q. the function vanishes in a neighbour- hood of aQ Identify Q with fundamental domain of the D-dimensional torus TQ. == R (..)). Define the T.-periodic extension
Uper (2) = Σ
u(z Lk), 2 Z D
(21)
kezD
16

---

Published as a conference paper at ICLR 2026
The sum is locally finite and defines L-periodic function on R Because u vanishes near aQ, there is no jump across the faces of Q: hence uper defines a function on TÉ. Moreover, the standard estimate for periodisation yields
U| (T)) S (p where Oper depends only on L (equivalently, on the choice of Q).
(22)
On TL the - can be expressed via Fourier coefficients:
| (TE)
XL
E (1+ 1p))"
(23)
KEED
where 02 are the Fourier coefficients the expansion
Uper (z)
Σ
2rikss/
== (x,,) E TD
(24)
kEBP
For given integer N 2 consider the truncated sum
Px(2): = Σ
Tg 2
(25)
||A|| SN
a trigonometric polynomial of degree N. By Parseval's identity,
Uperr PN Pwllia(TT)* = Σ
|
(26)
|| 2N >N
For ||k|loo N we have [x]] N. hence (1 + | S (1 + - so
22 2 S (1+
[K|3)* üg 2
Inserting this into equation 26 and using equation 23 gives
|p PN [(2)( I C(1 + N2) Uper 12 (TP)*
(27)
Taking square roots and using (1+ N2) )
8/2 < C'N yields
||upr Px (x S CNN U (T
(28)
Combining equation
19 equation
20 equation 22 and equation 28 we obtain
Upee
PNll S
(29)
Since Uper
Ω. restriction does not increase the -
||a PN)l $ |peer - Pv PWllea(T)) I (all"(a
(30)
Now count the number of retained modes:
d= # [k E zd ||A|| SN = (2N + + Hence N d and equation impliess
(31)
||a PNIl I Cyd /aae=/Dllall (allaaaa
(32)
It remains to write PN as an inner product of feature maps. Write h = (kiz, ky) with he E Zd. ky E gd,, that Lz,w/L 22ikk L/L v/ Enumerate all admissible indices & with || N as ((1)) k(d) and set set
(x(2) = حلح e 2xrkk()
(y( = 20ris))--//1
(33)
Define the veluer-valued maps h(x) = (hn (x),..... (aab...,,(())))) ha(x)) and g(v) (gn (v): (v...... :aa......al)))) ga(y)). Then h E (,2 Cd) and E T2; and using the standard Hermitian inner product (a, ()
Ej oe Bi we obtain
(Mx),
Σ Tgone 2mi(k() x410 PN
(34)
Together with equation 32 this yields the claim. The constant C3 depends only des dy and the geometry of X, (through the norms of the extension, cut-off, and periodisation operators and the choice of Q).
n
17

---

Published as a conference paper at ICLR 2026
Theorem 2 (Approximation by RoPE-type feature maps in L2), Let X C Rd and Y R be bounded Lipschitz domains, set D dz + dy and let T = R [2m2] with endpoints identified.
Let > O be an integer. consider real-valued function f and assume
f (T, (T:
(35)
Define
M = = + 1 fll.lllll. (T,
(36)
Then there exists constant C depending only on 8, day dy and the geometry of X, Y, such that for every integer d 2 1 one can find feature maps h E L2 (1x and E LY with d' S d, and family of unitary matrices (N(t))ET C Cd xd' of the form
A(t) = deag(e hwrt elatt)
W/ W
(37)
such that
[ Re
SCMd
8 D+1
(38)
Proof. We begin by expanding f a Fourier series with respect to the periodic variable t E T. Define the Fourier coefficients
ag (x y) = 2x /000 f(z,y,t)e ikt dt,
1 2
ke E,
(39)
which belong to (2 Y) because f (T: (2 xX)).. Then f can be written as
f(x,u,t) = Σ
(2 y)e
(40)
kez
and Parseval's 's identity gives the equalities
= 2Tr
| (Mll12((xxY)
(41)
k
| |8
2n
Ax| ||x
(a(xx))
(42)
k
These relations are the starting point for both the temporal truncation and the low-rank spatial ap- proximation.
Temporal truncation. For a cut-off frequency N E N we consider the truncated Fourier sum
(x,
ax tke
(43)
I
Using equation 41 and equation 42 we estimate the error made by this truncation:
||s - 11
= 2T Σ ||a.1?..
< 2m No.N-2
1
1ax11
(44)
|R
< 2 11
Taking square obtain
|| | (X(XXXXX)
(45)
18

---

Published as a conference paper at ICLR 2026
Low-rank approximation of the spatial coefficients. Because f (T: Y)), each coefficient aş belongs to (X Y) and the H norms satisfy the Parseval-type relation
2x
Σ kEZ
lax EllllBiexxYs II
(46)
We now fix an integer r 2 1 (the same for all frequencies [|]] S N) and apply Lemma to every
with rank parameter r. The lemma supplies functions h4 X - C' and Ih C" such that
| (he: go) (L S Cor +/D ||ax (X-x)))
(47)
Using these approximations we define a function that mimics f(A);
(o,
Σ
eth (ha(t)))
(48)
JAIEN
Because the Fourier modes eakt are orthogonal in (T), the error between f(M) and F(N) can be expressed without cross-terms:
-
f(N) E [XXXX 2π
Σ ||ox (hessx))
(49)
|A|SN
Inserting the estimate equation 47 and using equation 46
we obtain
| - | S rir (I
(50)
where C1 is a constant that absorbs Co and the factor coming from the sum of the squared H norms.
Assembling a global RoPE-type representation. single vectors. Set
Let us now collect all the component maps into
d' = (2N + 1)r,
(51)
and note that d' by the choice r [d/ 1)J. Define
h(x) = (x) hn (x))
Cd g(v)= (g-N(U)..... ,9(())) E ead
(52)
For the rotation matrices we take the block-diagonal operator
A(4) = diag(e -iN diag(e-iN& Ir, -iNN-1))
e I,),
(53)
which is clearly of the form diag(e etwart) with integer frequencies w/ (each frequency repeated times). direct computation shows that
(h ((t)))) =
Σ ikt ININ
(54)
Choice of the cut-off and final error estimate. Combining the estimates equation tion 50 and the identity equation 54 we obtain from the triangle incquality
45 equa-
Ilf (h, (1))))l < -
f( ((x) + +1 f(.)
-
(55)
S ~ +ir
C/ (*(T
(56)
Recall that r is related to the total dimension d by r [d/(2N -)]; for large d we have x d/N. Substituting this asymptotic relation into equation 55 yields
|f - (h, ll < C2 N 1 (d/N) -/D 11
We now choose N so that the two terms balance. Setting N = [dd/(0+1) gives
N x d #/(++1)
(d/N) =/D x d s/0+11)
Consequently,
| - (4.. & c(10f1l + I --//(411)
(57)
19

---

Published as a conference paper at ICLR 2026
Finally, note that the quantity [|0 + Il is bounded by a constant multiple of the norm M appearing in the statement of the theorem (because the I norms are controlled by the corresponding C(T; H*) norms). Thus we arrive at the desired estimate
|| - Re( ~ s II (h, A00))ll S C
(58)
Remark. Theorem 2 provides an ( x Y x 7) error bound, which is natural for mean-square losses and leads to clean rate because orthogonality in / (Parseval) can be fully exploited. A uniform-in- y. bound.
f I (1..A0))ll
(59)
is strictly stronger and requires additional ingredients beyond the L2 proof. There are two standard routes.
1. Uniform bounds from Sobolev embedding (typically slower rates under H only). Assume only the spatial Sobolev regularity used in Theorem namely f(..,,) E (xXXY) with > D/2. To pass from an Hd esstimateee _ one uses Sobolev embedding H Y) 4 L L Y) with any 8o > D/2. A typical spectral-truncation argument then gives, for any fixed 80 E (D/2,8) 8), a low-rank bound of the form
||a rank-d]l==((xx) I C Cdd (-- Cd-
ao)/ [| [[x])=((xx))
(60)
Choosing Su = D/2 + E yields the more explicit (but slightly weaker) rate
||a - rank-d]l I C. d (e-D/2-e) (xxx)-)
(61)
Plugging this - low-rank estimate into the RoPE construction (and using vector-valued Jackson bound followed by Sobolev embedding in (x, y)) leads to a uniform approximation rate
s(s D/2 -)
|l Ah, (o C. Md
Be
(62)
Ds + D/2
In particular, one generally pays an e-loss and a smaller exponent than in the T,2 theorem when only H H reguularyy is assumed in space.
2. Recovering clean _ rates by strengthening spatial smoothness. If one strengthens the spatial regularity of f (uniformly in f) so that -tabbee approximation theory applies, then the uniform RoPE approximation can match the clean algebraic rates. One convenient sufficient condition is additional Sobolev smoothness: assume that for some
0.
f C(T; H (X x Y)),
0 05653,
(63)
and retain 5 D/2 to guaranty H x Y) + (X x Y). Then one can approximate each Fourier coefficient ag in H at rate o 1 ||x|| and after embedding obtain
| (h, (0)
SCM, (N (N** +)=/)r--/))
(64)
with M,, maxogess 10ff/lo and d/ + 1) as in the proof. Optimizing N then
yields an exponent
== Ds + 0
(65)
and in the balanced case a = s this recovers 3 = s/(D in L (at the cost of requiring 22 smoothness (z, y)).
Theorem only requires Cd) and E Cd( Uniform-in-(zr. v) bounds typically require additional regularity (and sometimes stronger boundary regularity of X, Y) if one wants h, g to be continuous on X, Y.
The 2 theorem is sharp and technically robust under minimal assumptions. Uniform - control stronger, but either yields slower rate under pure H² assumptions or requires stronger spatial smoothness assumptions retain the clean exponent.
20

---

Published as a conference paper at ICLR 2026
B EARLY STOPPING EXPERIMENTS
In this section, we study a simple early stopping rule for the retrieval agent. Let
(a1, 02 ar)
be the full sequence of chunks the agent would select if stopping threshold were applied, and let G be a set of ground-truth chunks for the current question.
For each step t, the agent outputs Q-value Qa for taking the next retrieval action. Given fixed Q-value threshold Qirreshold we simulate an early-stopping policy that keeps taking actions while Qi ihrreshol and terminates soon as Q mhreshold We denote by /stop the number of actions actually taken under this policy, i.e. the number of selected chunks:
stop number of steps until the first t with Qi < thrreshold
Independently of the stopping rule, define tearlies the earliest step at which all ground-truth chunks have already heen collected:
farlies = min ft fai..... ar) 2 G).
If the agent never collects all ground-truth chunks, i.e. such t does not exist, discard this episode from the analysis below.
For comparison, we also consider oracle stopping policy that is allowed to look at the ground truth: it knows fearliest for each episode and simply stops at this step. By construction, this oracle policy never stops too early or too late.
Depending on the relation between (stop and learliest we distinguish three outcomes.
Early stop ("early"". If tsop fearlissts the stopping rule terminates before all ground ground-truth chunks have been selected. In this case the error is due to stopping too early and missing potentially useful chunks.
Perfect stop (""erfect")). If stop fealiesss the stopping rule terminates exactly at the first step when the set of selected chunks already contains all ground-truth chunks. this case, the stopping behavior is optimal with respect our definition.
Late stop ("late"). If Istop > learliest, then at some earlier step the agent had already collected all ground-truth chunks but continued retrieve additional chunks. This corresponds to stopping too late and taking unnecessary steps.
Figure (top row, panel (a)) shows how the proportions of early and late errors change function of the Q-value threshold Qihreshola on HotPotQA. For small thresholds, the agent almost never stops early but may continue retrieve redundant chunks, which leads to late errors. As the threshold increases, late errors decrease, but the probability of stopping too early grows.
Panel (b) of Figure 4 reports the proportion of "perfect" stopping events, peaking around thresholds Qieeshold 2 0.1-0.3. Panel (c) shows the average number of selected chunks (episode length) under the same policy. Larger thresholds lead shorter episodes, but once the threshold becomes too high, the early-stop error rate rapidly increases and performance degrades.
Table5 summarises these trade-offs quantitatively on HotPotQA for the GTE embedder with penall 1 ze extra steps=True and never termi mate=True We report the fraction of early, late and perfect stops, the average episode length, and the final Fact EM and Fact F1 scores, well as the corresponding true positive rate (TPR) and false positive rate (FPR) for the stopping rule viewed as a binary classifier. The best Fact Fl is achieved at Qthreshold = 0.2, confirming that moderate thresholds provide good balance between taking enough retrieval steps and avoiding unnecessary ones.
Using the TPR and FPR columns of Tables and and 6 we can plot the receiver operating characteristic (ROC) curves of the carly-stopping rule, shown in Figure Panel (a) corresponds to HotPotQA and panel (b) to BabiLong QA2. Each point on the curves corresponds particular Q-value threshold
21

---

Published as a conference paper at ICLR 2026
criterion
Perfect stooping
Sarly supping Late expaing
Pear Los
0.3
2.6
0.2
0.0
0.0
0
0_
O.thrahold
(a)
(b)
(c)
Stopplng criterion emors
Perfect stopping rate
Hear eptsode length
3.0
1.0
0.0
IN roke
Ere
0.4
0.4
0.2
Porfect stppiig
0.5
1.8
(d)
(e)
(f)
Figure 4: Early stopping analysis on HotPoIQA (top row) and BabiLong QA2 (bottom row). Panels (a.d) show the proportions of early and late errors function of the Q-value threshold 1hreshhld- Panels (b,e) show the proportion of perfect stops. Panels (c,f) show the average number of selected chunks (episode length).
Table 5: HotPotQA early stopping experiments
Q-value threshold
stopped early
stopped later
perfect stop
TPR
IPR
Episode len
Fact EM
Fact
Ans EM
Ans F1
-0.1
0.979
0.021
0.983
0.380
4.99
0.968
0.563
0.588
0.759
0.0
0.015
0.395
0.590
0.976
0.1
2.82
0.954
0.843
0.592
0.761
0.1
0,061
0.060
0.879
0.952
0.041
2.23
0.910
0.915
0.593
0.756
0,2
0.088
0.020
0.892
0.937
0.032
2.13
0.883
0.917
0.587
0.752
0.3
0,104
0.006
0.890
0.927
0.029
2.08
0.868
0.915
0.585
0,747
0.002
0.919
2.05
0.575
0.737
0.5
0,132
0
0.867
0.910
0.025
2.02
0.840
0.907
0.571
0.734
0.730
0.7
0.157
0
0.843
0.891
0.023
1.96
0.817
0.895
0.564
0.724
0.8
0,202
0.798
0,840
0.017
1.82
0,773
0.847
0.546
0,702
0.9
0.417
0.583
0.611
0.006
1.27
0.565
0.620
0.444
0.588
1.0
0,910
0.090
0.105
0.000
0.21
0.088
0.111
0.266
0.385
1.1 1.000
0
0
Qthreshold- The red star in each panel marks the oracle stopping policy introduced above, which knows fearlies and stops exactly at that step; this point serves as an upper bound on the achievable trade-off between TPR and FPR. On HotPotQA the area under the curve (AUC) 0.96, and BabiLong QA2 -0.97.
Figure (bottom row) and Table 6 report the same analysis on BabiLong QA2. Qualitatively, the behaviour of the stopping rule is similar to HotPotQA: higher thresholds lead to shorter episodes and more early stops, while lower thresholds reduce early-stop errors at the cost of more late stops and longer episodes.
However, the transition between these regimes much sharper on BabiLong QA2. For thresholds in the range Qthreshold E [0.2, 0.6] the fraction of perfect stops remains very high (sx 0.95-0.99), while the average episode length is reduced from about 6 to roughly 2.2 retrieval steps. In this region Fact EM and Fact Fl stay close to their maximum values (Fact F1 around 0.95), and answer accuracy (Ans EM/F1) is also near- optimall. Only when the threshold approaches 1.0, performance collapses, as the agent stops almost immediately and misses relevant chunks.
22

---

Published as a conference paper at ICLR 2026
ROC Curve
ROC Curve
1.0
1.0
True Positive Rate
True Positive Rate
0.6
0.6
0.4
0.4
0.2
ROC (area
0.970)
ROC curve (area 0.961)
Oracle Point
Oracle Point
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.6
0.8
1.0
False Positive Rate
False Pasitive Rate
(a) HotPotQA
(b) BabiLong QA2
Figure 5: ROC curves for the early-stopping rule. Panel shows HotPorQA: panel (b) shows Ba- biLong QA2. The dashed line indicates random performance. Each point corresponds to a different Q-value threshold Qhrreshodd- The red star denotes the oracle stopping policy that always stops fearliest i.c. exactly when the last ground chunk has been retrieved.
Table 6: BabiLong QA2 early stopping experiments.
Q-value threshold
stopped early
stopped later
perfect stop
Episode len
Fact
Ans
Ans
-0.10
0.000
0.994
0.006
6.00
0.996
0.499
0.884
0.884
0.00
0.000
0.994
0.006
0.996
0.499
0.884
0.884
0.10
0.000
0.498
0.502
2.86
0.996
0.845
0.944
0.944
0,20
0.000
0.036
0.964
2.29
0.996
0.949
0.976
0.976
0.30
0.006
0.010
0.984
2.25
0,990
0.952
0.970
0.970
0.40
0.008
2.24
0.988
0.953
0,970
0.50
0.008
0.000
0.992
2.23
0.988
0.954
0.972
0.972
0.60
0.016
0.000
0,984
2.21
0,980
0.948
0.968
0.968
0.70
0.042
0.000
0.958
2.16
0.954
0.934
0.948
0.948
0.80
0.112
0.000
0.888
2.06
0.884
0.905
0.884
0.884
0.90
0.177
0.823
0.861
0.830
1.00
0.930
0.000
0.070
0.22
0.070
0.107
0.230
0.230
1.10
1.000
0.000
0.000
0.00
0,000
0,000
0.000
0.000
C PLANNING FOR MULTI-STEP RETRIEVAL
We can apply planning at the multi-step retrieval stage, formulating source selection as search over the space of action trajectories; $ 4.4 for application. In the spirit of Retrieee, run beam search where candidates are ranked by the learned action-value Qa(s, a). However, our planning is computationally cheaper because Qo is computed embeddings, Q0(s, a) (s), (a)), so no new transformer forward passes are required for each candidate chunk, whereas Beam-Retriever relies a transformer reranker over trajectories, incurring fresh forward passes at every expansion. Details of the embedding- -based scoring are provided in $ 3.2. At inference, we perform beam search over Q and deterministically expand the top-k actions by Qo.
dol product of state and action
D METHOD COMPLEXITY AND EFFICIENCY
Q-RAG produces final answer using two main components. The first is multi-step retrieval agent that performs iterative search over the full document to collect all context-relevant evidence (see 3.2 The second is LLM Answerer that conditions the retrieved chunks and generates the final response. Importantly, only the retrieval agent interacts with the original long context; the effective context length seen by the I.LM Answerer depends solely the retrieval hyperparameters
23

---

Published as a conference paper at ICLR 2026
(e.g., number of retrieval steps T. maximum chunk length). Consequently, the time and memory complexity the LLM Answerer with respect to the original context length N are both O(1). Retrieval agent consists of two embedders: state embedder Es and action embedder E, (see sec. 3.2
Chunks embedding. The action embedder computes embeddings for chunks of the original doc- ument. the document has length and the chunk size is ne. embedding the entire document takes 0 fact). where fact is the embedding time per chunk (treated as a constant). The action embedder performs single pass over all chunks per retrieval episode; thus its complexity is linear in NN, i... O(N).
State Embedding. The state embedder processes the state K times per episode (once per search step). From the construction of the state (see fig. [T. the total cost over an episode is O(K tstate), where state embedding time fstate depends ne and K, but not N. Hence, the state embedder (1) with respect to document length N.
Search Policy. To select the next chunk at each step, compute the inner product between the cur- rent state embedding and all action embeddings. With naive implementation, selecting all K actions over the episode requires O K demb tee ONN), where demb is size of embedding vectors. This can be reduced using approximate ANN methods that achieve sub-linear query time in practice (??).
Overall time complexity. Summing the terms above yields
N =
tact + K tstate + K demb
-ON),
Re.
since K, fact. tstate, and demb do not depend on N.
Space complexity. The main part that direcly depends on document length the nubmer of chunk embeddings we need to store: O demb N O(N). In practice, embeddings are lightweight; The
GPU memory is mainly consumed by the LLM weights and the action embedder forward passes. By capping the action embedder's ' batch size (parameter chunk batch), the growth of peak memory with becomes negligible.
Training Time Efficiency. A critical practical advantage of the Q-RAG framework is its efficient and rapid training convergence, as demonstrated in Figure The learning curves depict the model's performance evolution
distinct and challenging benchmarks: BabiLong QA3 and HotPorQA. The curves show a sharp initial rise in evaluation metric scores, followed by a stable plateau, indi- generation task. Notably, this
cating that the model quickly learns the core convergence is achieved within approximately 12 hours of training time on GPU setup.
Babilong QA2
HotPotQA2
1.0
1.0
0.9
0.9
0.8
Eval return
Eval return
0.7
0.7
0.6
0.5
0.4
0.3
0.3
Time (hours)
Time (hours)
Figure 6: Learning curves for HotPotQA and BabiLong QA3 runs. Both graphs the average episodic return with respect to training time.
24

---

Published as a conference paper at ICLR 2026
E EXTRA QA RESULTS
Table 7 compares multi-step retrieval methods on HolPoiQAAdisttrctorrs,, Musique (in-distribution), and Musique coutoof-distributionn
It reports both fact-retrieval (Fact FI, Fact EM) and answer-generation (Ans F1, Ans EM) scores. Q-RAG and its planned variant (Plan Q-RAG) achieve strong overall results, especially out-of-distribution data, while Retriverr leads Hot- PorQA but generalizes less robustly. Methods with missing entries did not report results for the corresponding dataset metric.
Table 7: Comparison of methods on CloPorQA-distraacee Musique (in-distribution). and Musique (OOD). Bold text and underline denote the best and second best scores respectively.
HotPotQA
Musique
Musique (OOD)
Average
Methods
Fact Fl
Fact EM
Ans FI
Ans EM Fact FI
Fact EM Ans F1
Ans EM
Fact Fl
Fact EM Ans Fl
Ans EM Ans FI Ans EM
Plan QQRAG + QwQQ322
0.95
0.91
0.76
0.60
0.84
0.76
0.60
0.44
0.69
0.53
0.51
0.36
0.62
01.46
Q-RAG+QuQ-12B
0.93
0.89
0.76
0.59
0.81
0.73
0.59
0.43
0.71
0.55
0.52
0.37
0.62
0.46
Beam- Retnever+QuQ--aa
0.97
0.94
0.77
0.61
0.86
0.89
0.39
0.43
0.61
0.36
0.40
0.27
0.59
0.44
Searchorl
Search-ol
GraphReader
HippoRAG
F TRAINING DETAILS
We trained the model with AdamW (learning rate 1.5 x 10 10055,,1000 =09,, 32-0. 1-10- weight decay 5 x 10- The learning rate followed a linear schedule: we used warm-up of 1,000 steps, then linearly decayed the rate 10% of its initial value over the remaining training steps. We applied gradient clipping with maximum & norm of 2.0 and used gradient accumulation for 8 steps. The base mini-batch size was 12; with accumulation this yields an effective batch size of 12 per update (scaled by the number of devices if using distributed training).
96
In the objective and algorithmic components we set 7=0.99, a=0.05, A=0.5, and 7-0.22. Action representations were capped at maximum length of 220 tokens.
The end-to-end training of a single model did not exceed 12 hours on a single A100- GPU.
Models per benchmark. For open-domain QA benchmarks (HotPorQA, Musique), we trained multilingual-ee-lare and Alibaba-NUP For Ruler and BabiLong, we trained httpbbb-wwu//ggeemmlltllla -base encoders.
G EVALUATION DETAILS
LLM Models for generation. To compute answer-level metrics (Ans EM and Ans F1), we con- dition the QwQ-32B model on the question and the retrieved text chunks. All answer-generation results reported for Q-RAG and Plan Q-RAG on the HotPotQA and Musique benchmarks were ob- tained under consistent generation settings: decoding with temperature 0.0 and maximum output length of okens = 8000. For the BabiLong and RULER experiments, we instead used Qwen-4B with max tokens 512 and reasoning disabled (enable thinking False).
Retrieval configuration. For Q-RAG we limit the number of retrieval steps to T = 2 on Hot- PotQA, RULER and Babilong we use T 4. The same step limits are used when evaluating Search-R1 and Beam Retriever.
We split documents into fixed-length, non-overlapping chunks, aiming not break sentences across chunk boundaries. The chunk length is primarily determined by the context window of the embed- ders used in our main experiments (512 tokens) and the number of retrieval steps. For Needle-in-a- Haystack and BabiLong we use chunk length of 64 tokens. For open-domain QA tasks we set the chunk length as a function of the number of retrieval steps i.e. for HotPotQA we segment the corpus action chunks of most 110 tokens
into chunks of at most 220 tokens (T 2); for Musique (T = 4). In additional experiments with a 'Alibaba- Allibab--NLP/ttee. muutliiillll base (8k context length) use chunk length of 256 tokens.
25

---

Published as a conference paper at ICLR 2026
Dataset
Settigg
Chunk size
T
Backbone retriever
Answering LLM
HotPoRQA
Q-RAG/ Plan Q-RAG
220
2
ls-lagg
QwQ-32B
HotPaiQA
(early stopping)
mi iinguud multi wwu.tiilinmaalee--aar
QwQ:32B
Musique
Q-RAG/ Plan Q-RAG
110
4
QwQ:32B
Babilong
Q-RAG
Facobaok/contrieea. facebsok/contrieve:
Qwen3 4B
RULER
Q-RAG
Qwen3-4B
Table 8: Retrieval and generation configuration for each dataset. Chunk size is in tokens; T the maximum number of retrieval steps.
Fact-level metrics. Let Sgi be the set of ground-truth supporting facts and Spred be the set of predicted supporting facts returned by the retriever. Our Fact EM metric is defined
1, if Sgi S Speed
Fact-EM
0, otherwise.
Equivalently, in code: em 1.0 if gt.sf. issubset (pred.sf) else 0.0 O. Thus Fact EM gives full credit whenever the prediction covers all ground-truth facts, even if it also con- tains additional, irrelevant chunks; it does not require the predicted and ground-truth sets to be exactly equal.
26


Reviews:
[{"review_id": "YhsVuKIUjB", "reviewer": "Reviewer_rFyU", "summary": "Q-RAG formulates multi-step retrieval as a MDP process where the state is the embedding representation of the query and past retrieved chunks, the action is the embedding of the candidate next chunk to retrieve, and Q is defined as the dot product of the state and the action embedding. Reward is defined as whether all ground truth chunks are retrieved and are propagated back via TD to train the state and action embedders. The main contribution of Q-RAG is to train a multi-hop retriever using temporal-difference RL. The authors also claim Q-RAG to be more efficient in training and inference, as only the retriever is trained and iteratively used during inference", "questions": "see the weakness point 3", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method is intuitively sound and novel, with theoretical justification of using dot product as an approximation for the Q-function in the appendix. \n2. The experimental setup is comprehensive with good coverage on multiple tasks requiring multi-step retrieval, spanning commonsense reasoning, NIAH, and multi-hop retrieval. \n3. Numerous baselines are included for different datasets, and the baselines are very recent. \n4. Good ablation study to demonstrate the necessity of various key design choices in Q-RAG.", "weaknesses": "1. In all experiments, the choice of the total number of retrieval steps T is very important because you need to know when to stop retrieving and have the reader LLM answer the question. However, the choice of T is not discussed at all. \n2. The authors claimed efficiency in training and inference. Though it is intuitively true, the discussion regarding training efficiency was insufficient and was not studied. Only the inference runtime was partially studied in Figure 3c. \n3. The experimental setups are pretty vague. What off-the-shelf models were used as the state embedder and action embedder? What reader LLM are used for different baselines for different datasets? Why the different choices? The paper can also benefit from brief introductions to each dataset and elaborating on what chunking strategy was used for each dataset for readers not familiar with them. \n4. Presentation: Figure 3c’s bottom has been cut off and typo “Babylon” on line 42", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "veb4tVSJyp", "reviewer": "Reviewer_sUPb", "summary": "This paper introduces Q-RAG, a novel and resource-efficient framework for multi-step retrieval in Retrieval-Augmented Generation (RAG) systems. The authors' primary contribution is a paradigm shift away from fine-tuning the Large Language Model (LLM) itself, focusing instead on training a smaller, more efficient embedder model using value-based reinforcement learning (RL). The paper claims three main contributions: (1) a new method for training a multi-step retrieval agent using temporal difference RL, which is exceptionally resource-efficient (trainable on a single GPU); (2) state-of-the-art performance on ultra-long context benchmarks (up to 10M tokens) like Babilong and RULER; and (3) a novel relative positional encoding mechanism that enables temporal reasoning during the retrieval process. This approach is presented as a modular, flexible, and scalable solution to the challenges of multi-hop reasoning and \"needle-in-a-haystack\" problems in extremely long documents.", "questions": "- Given that creating datasets with ground-truth support facts is infeasible for most real-world applications, what is your proposed concrete path toward adapting Q-RAG to new domains where such supervision is unavailable?\n\n- Your paper suggests using LLM-based rewards as a future direction. How do you plan to address the now well-established problem of \"reward hacking,\" where RL agents learn to generate superficial \"master keys\" to fool LLM judges? Without a robust defense against this, wouldn't the training process collapse?\n\n- Could you elaborate on the model's limitations in tasks requiring explicit compositional reasoning (e.g., \"Which of these two products has a higher rating?\" which requires retrieval of two facts and a comparison)? How does Q-RAG's performance depend on the reasoning capability of the final, frozen LLM used for generation?\n\n- Have you considered the performance implications of integrating the Q-RAG retriever as the first stage in a more complex pipeline that includes, for example, a cross-encoder reranker? Would the greedy, step-by-step retrieval policy still be optimal in such a setting?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- **Novelty and Technical Soundness:** The core idea of fine-tuning the embedder instead of the LLM for multi-step retrieval is a significant and elegant contribution. The formulation of the problem as a Markov Decision Process (MDP) is well-defined, and the factorization of the Q-function into an inner product of state and action embeddings is a clever design choice, theoretically justified in the appendix. The introduction of dynamic relative positional encoding to handle temporal reasoning is particularly innovative and is empirically shown to be highly effective on the Babilong QA3 task.\n\n- **Exceptional Resource Efficiency and Accessibility:** The ability to train the entire system on a single A100 GPU is a major practical advantage. This dramatically lowers the barrier to entry for research in multi-step RAG, making it accessible to a much wider community compared to LLM-tuning methods that require large GPU clusters.\n\n- **Modularity and Flexibility:** By decoupling the retriever from the generator, the trained Q-RAG embedder can be paired with any LLM, including powerful proprietary models like GPT-4. This \"plug-and-play\" nature is a significant strength in the current AI ecosystem.\n\n- **Strong Empirical Results and Scalability:** The paper demonstrates state-of-the-art or highly competitive performance across a range of challenging benchmarks, including ultra-long context (Babilong, RULER) and standard multi-hop QA (HotpotQA, Musique). The model's ability to maintain high performance as context length scales to 10 million tokens is particularly impressive and clearly demonstrates the scalability of the proposed method.\n\n- **Reproducibility:** The authors provide detailed experimental setups and ablation studies justifying their architectural choices (e.g., using soft Q-learning and target networks). They also state their intention to release code, which strengthens the paper's contribution.", "weaknesses": "- **Critical Dependence on Supervised Reward Signal:** The paper's most significant weakness is its reliance on a sparse terminal reward based on the retrieval of a pre-defined set of \"ground-truth support facts ($F^*$).\" While this provides a stable training signal, it severely limits the method's real-world applicability. Such meticulously labeled datasets are extremely rare and prohibitively expensive to create for new domains (e.g., legal, medical, enterprise knowledge). This dependency relegates the current method to a powerful tool for academic benchmarks rather than a general-purpose, deployable solution.\n\n- **Underestimation of the LLM-as-a-Judge Problem:** The paper acknowledges the limitation of supervised rewards and briefly suggests using an LLM to generate a reward signal as \"future work.\" This suggestion dangerously overlooks the well-documented and critical failure modes of the \"LLM-as-a-Judge\" paradigm. Recent research has extensively shown that LLM judges are systematically vulnerable to \"reward hacking\" and \"master key\" attacks, where superficial, semantically meaningless inputs (e.g., punctuation, phrases like \"Thought process:\") can elicit false positive rewards. An RL agent like Q-RAG would rapidly learn to exploit these vulnerabilities, leading to a collapsed policy that generates \"master keys\" instead of performing meaningful retrieval. The paper's failure to address this critical challenge makes its proposed path to unsupervised application seem naive and unviable in its current form.\n\n- **Limited Compositional Reasoning in the Retriever:** While the relative positional encoding enables temporal reasoning, the retriever itself does not perform more complex compositional reasoning, such as comparison or calculation. The task of synthesizing information and drawing a final conclusion is entirely offloaded to the final, frozen LLM. The system's performance is therefore contingent on the capabilities of this final LLM, and it may fail on questions that require intricate reasoning over the retrieved facts.\n\n- **Isolation from the Broader Advanced RAG Ecosystem:** The proposed method operates in isolation from other well-established advanced RAG techniques like hybrid search (combining dense and sparse retrieval), parent-child retrieval, or sophisticated reranking modules. Integrating Q-RAG's RL-based retriever into a more comprehensive, production-style pipeline could be beneficial, but this is not explored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "weg8QUacbL", "reviewer": "Reviewer_nFRL", "summary": "The paper presents Q-RAG, a framework that formulates multi-step retrieval in retrieval-augmented generation (RAG) as a reinforcement learning problem. Two transformer encoders are trained: a state encoder that represents the current retrieval context (the initial query and previously retrieved chunks), and an action encoder that represents each candidate document chunk. The similarity between the two embeddings defines a Q-value, which indicates the estimated utility of retrieving that chunk in the current context. Training uses Soft Q-Learning with λ-return, where the model receives a binary outcome-based reward depending on whether all supporting facts are retrieved. During inference, the model iteratively embeds the current state, scores candidate chunks by their Q-values, selects the most relevant one, and updates the state until retrieval is complete. Experiments are conducted on HotpotQA, MuSiQue, RULER (1M tokens), and BabiLong (10M tokens) to evaluate retrieval accuracy and computational efficiency.", "questions": "- Do the authors observe any evidence of policy collapse or repeated retrieval of the same chunks during training?\n- How well does the learned policy generalize to unseen domains or corpora without retraining?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper presents a clear formulation of multi-step retrieval as a reinforcement learning process, replacing traditional query rewriting with a learned retrieval policy. This modeling perspective provides a coherent computational framework for retrieval-augmented generation.\n- The proposed model adopts a dual-encoder design, where a state encoder and an action encoder jointly estimate Q-values to guide multi-step retrieval decisions. Both encoders are compact transformer models, showing that a small embedding-based retriever can learn retrieval behavior previously handled by large language models.\n- The training framework integrates Soft Q-Learning with λ-return, which helps stabilize optimization under sparse outcome rewards. The method shows that an embedding model can reach retrieval accuracy comparable to reinforcement-learning–finetuned large language model retrievers while operating at roughly 10–20× lower computational cost.\n- The experimental evaluation covers multiple datasets, including HotpotQA, MuSiQue, and synthetic long-context settings (RULER, BabiLong), demonstrating consistent retrieval performance and scalability to contexts up to 1M–10M tokens.", "weaknesses": "- The model’s architecture and training setup restrict its ability to directly handle long or unstructured contexts. Both encoders have short input windows, and the approach depends on datasets with annotated supporting facts, which limits generalization to open-domain or unlabeled corpora.\n- The paper lacks theoretical analysis of the proposed Q-learning framework. No formal discussion of convergence, optimality, or error bounds is provided, and the training stability is supported only by empirical observations.\n- The reward design relies on a single outcome-based binary signal. While λ-return mitigates the sparsity issue, the paper does not discuss other possible reward shaping or exploration strategies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BSvvBLVXzu", "reviewer": "Reviewer_LyWb", "summary": "This paper proposes an approach that fine-tunes an embedder model for multi-step retrieval through reinforcement learning. The proposed method, Q-RAG, achieves state-of-the-art results on long-context benchmarks including Babilong and Needle-in-a-haystack tasks.", "questions": "Can you explain all the points I raised in the Weaknesses section?", "limitations": "", "rating": 2, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- By decoupling the retriever from the LLM and training only a lightweight embedder, the proposed approach is practical, compared to other approaches that require fine-tuning the entire LLMs.\n- This paper introduces a relative positional encoding scheme for context embeddings. This allows the agent to reason about the temporal relationships between chunks.\n- The experiments demonstrate scalability and robustness, achieving state-of-the-art results on benchmarks with contexts up to 10 M tokens.", "weaknesses": "- The paper’s analysis of the multi-step retrieval process is under-explored. While the method relies on a maximum step budget, T, this critical hyperparameter is not specified for the experiments, hindering reproducibility. Furthermore, the paper lacks an investigation into the agent’s learned stopping behavior, i.e., it is unclear if the agent effectively utilizes the STOP action to terminate retrieval early. A sensitivity analysis on how performance varies with different values of T is also absent. This omission is significant, as the optimal number of retrieval steps is often unknown in practice, leaving the method’s robustness and practical applicability in question.\n- The paper’s presentation could be significantly improved. Several key results and concepts are not properly introduced or referenced, which disrupts the reader’s flow. For instance, the “Plan-Q-RAG” varian appears abruptly in Table 2 without prior introduction or motivation in the main body of the paper. Later, I found its explanation in the appendix. Additionally, Figure 2 is not explicitly mentioned anywhere in the text.\n- The empirical validation for the proposed temporal reasoning mechanism appears narrow. The authors claim this mechanism is a major technical contribution, yet its effectiveness is demonstrated exclusively on the Babilong benchmark. To substantiate the generalizability and impact of this technique, the evaluation should have included a broader range of tasks where temporal or sequential understanding is paramount. For example, testing on datasets that require reasoning over chronologically ordered documents or procedural texts (e.g., [1]) would provide more robust evidence and strengthen the authors’ claims.\n\n[1] Karpinska, Marzena, et al. \"One Thousand and One Pairs: A “novel” challenge for long-context language models.\" Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing. 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)