SmartRAG: Jointly Learn RAG-Related Tasks From the Environment Feedback
Authors: Jingsheng Gao
Linxu Li
Ke Ji
Weiyuan Li
Yixin Lian
yuzhuo fu
Bin Dai
Venue: ICLR 2025 Poster (2025)
Keywords: Retrieval-augmented Generation
Language Models
Reinforcement Learning
Primary area: reinforcement learning
TL;DR: We introduce SmartRAG, a joint framework to enable an LM learn when to retrieve, what to retrieve and how to answer.

Abstract: RAG systems consist of multiple modules to work together. However, these modules are usually separately trained. We argue that a system like RAG that incorporates multiple modules should be jointly optimized to achieve optimal performance. To demonstrate this, we design a specific pipeline called SmartRAG that includes a policy network and a retriever. The policy network can serve as 1) a decision maker that decides when to retrieve, 2) a query rewriter to generate a query most suited to the retriever and 3) an answer generator that produces the final response with/without the observations. We then propose to jointly optimize the whole system using a reinforcement learning algorithm, with the reward designed to encourage the system to achieve the highest performance with minimal retrieval cost. When jointly optimized, each module can be aware of how other modules are working and thus find the best way to work together as a complete system. Empirical results demonstrate that the jointly optimized system can achieve better performance than separately optimized counterparts.

Published as a conference paper at ICLR 2025
SMARTRAG: JOINTLY LEARN FROM THE ENVIRONMENT FEEDBACK
TASKS
Jingsheng " Linxu Li Ke Ji Weiyuan Li Yixin Lian", Yuzhuo Fu Bin Dai Shanghai Jiao Tong University 8Xiaobing.Al
& The Chinese University of Hong Kong, Shenzhen
jiggshng,,,,, a edu. on {keji]elink, cuhk. edu. [lilinxu, liweiyuan, lianyixin, daibin]exiaobing..
ABSTRACT
RAG systems consist of multiple modules to work together. However, these mod- ules are usually separately trained. We argue that system like RAG that incor- porates multiple modules should be jointly optimized to achieve optimal perfor- mance. To demonstrate this, we design a specific pipeline called SmartRAG that includes a policy network and retriever. The policy network can serve as 1) a de- cision maker that decides when to retrieve, 2) query rewriter generate a query most suited to the retriever, and 3) an answer generator that produces the final response with/without the observations. We then propose to jointly optimize the whole system using a reinforcement learning algorithm, with the reward designed to encourage the system to achieve the best performance with minimal retrieval cost. When jointly optimized, all the modules can be aware of how other modules are working and thus find the best way work together complete system. Empirical results demonstrate that the jointly optimized SmartRAG can achieve better performance than separately optimized counterparts.
INTRODUCTION
Although large language models([.LMss) (Chowdhery et al.) 2023 Touvron et al. 2023 Chung al.)
2024) have demonstrated exceptional capabilities across various domains, addressing knowledge- related issues beyond model parameters remains a challenging task (Mallen et al. 2203b Min al.,
2023 Retrieval-auge auggeettttiititioi generation(RAGr effectively enhances model performance in these
scenarios by retrieving additional information from external tools (Ram al. 2023).
RAG systems usually consist of multiple modules including at least retriever and a generator. Some systems may have other modules like a reranker (Glass et al. 2022), a decision maker deciding when to retrieve (Jeong et al.] [2024] Wang et al. 2023). query rewriter (Ma et al. [2033 Tan et al. (024) or a verifier (Lewis et al. 2020; Izacard et al. 2003).. These modules are often hand-designed and separately optimized. One of the issues is that the golden answer of the intermediate modules are usually not accessible. What is worse, sometimes the golden answer is model-dependent or sirieveer-deendeeet For example, Asai al. (2024) uses the result of GPT4 (Achiam et al. 2023) as the ground truth for the decision maker, which can be suboptimal. A question that GPT4 can answer without retrieval may need retrieval for other base models, meaning that the golden answer for the decision maker is model-dependent. In this sense, optimizing each module separately will obviously be suboptimal.
We argue that a system like RAG that incorporates multiple modules should be jointly optimized in an end-to-end manner. For this purpose, design a RAG system, namely SmartRAG, and jointly optimize it using reinforcement learning, with environment feedback as supervision.
Specifically, SmartRAG includes two core building blocks, a policy network, and a retriever, as shown in Figure The policy network can serve as three different roles. Firstly, it should decide
Work done during an internship at Xiaobing.Al
Corresponding Author

---

Published as a conference paper at ICLR 2025
[Answer] John Adams
[Question] Who was the first president who lived the white hause? [Observation] None [Searched Text]
Reward
Action
Policy T
[Search] First president lived in the white house
State
Reward
update
Q
[Searched] John Adams became the first president to move into the White House
retrieve
Observation
Database
Figure 1: Overview SmartRAG. The policy network takes the original question and optional observations as input and generates action that can be interpreted as a decision-maker and the corresponding parameters. If the policy decides to retrieve, the retriever will be called to obtain the observation that updates the state. Otherwise, the policy network will output the final answer. SmartRAG is optimized using a reinforcement learning algorithm with a properly designed reward.
whether to retrieve not based on the input state, which includes the input question and optional ob- servations. The second role is query rewriter. If the policy network decides to retrieve the database, it will also output the search query that better suits the retriever than the original question. The last role is answer generator. If the policy network thinks is already able to answer the question based on the current information, it should generate appropriate final response.
To optimize such system, the only supervision should come from the environment, rather than a proxy "golden" answer generated by another model. For SmartRAG, the primary goal is to correctly answer the question while the secondary goal is to minimize the retrieval cost. With this guideline, can design reward function that encourages the whole system correctly answer the question with minimal retrieval trials.
We validate the effectiveness of SmanRAG on various datasets, demonstrating that our SmartRAG outperforms its counterparts with separately optimized modules. We also systematically analyze the properties of SmartRAG, showcasing that learns when retrieve, what retrieve and how to answer based on the observations. The integration of these competencies necessitates a high degree of inter-module awareness, which in turn enhances the system's 's overall performance.
Our contribution lies in three aspects:
We design SmartRAG, a system that includes policy network and retriever. The policy network can serve three different roles related to RAG. With this design, our SmartRAG system can be optimized using an RL framework.
2. We propose using PPO jointly optimize our SmartRAG achieve better performance than training each related module separately.
3. Empirical results demonstrate the effectiveness our method on various tasks. Systematic analysis also applied to help us better understand how SmartRAG works and why outperforms previous algorithms.
RELATED WORK
2.1
RETRIEVALLAUGGEENTER GENERATION
Retrieval augmentation techniques have been utilized to acquire external knowledge, which aids language models in achieving superior performance across a broad spectrum tasks (Guu et al.) 2020 Gao al. 2023 Hu & Lui 2024). Previous research has primarily concentrated on "what
to retrivee (Khandelwa et al. 2019 Ram et al. 2023) and "how to utilize the retrieved informa-
tion" (Khattab et al. 2022).
Recently, some studies have begun to explore when to use retrieval to meet the varying requirements of diverse tasks (Tiang al. 2023). For instance, Mallen al. (2023b) assesses the popularity
2

---

Published as a conference paper at ICLR 2025
of query based the frequency of its entities and recommends activating retrieval modules only when the entity frequency drops below a predetermined threshold. Wang et al. (2023) enhances the model's retrieval performance by using self- lnffleeewedge to decide what they know and do not know. Self-Rag (Asai et al. 2024) leverages special tokens to adaptively retrieve external knowl- edge and confirm the output's relevance efficacy. Additionally, some researchers incorporated supplementary modules or classifiers to determine whether a query necessitates additional knowl- edge for resolution (Liu et al. 2024 Jeong al. 2024). Wang et al. (2024) create compositional unknown dataset (CuQA), and utilize the confidence words or scores to decide whether to retrieve.
Unlike previous methods, we have implemented an end-to-end approach to jointly optimize all the related tasks such that the whole system can achieve better performance.
2.2 LEARNING FROM FEEDBACK
Training LLMs from feedback (uyang et al. 2022 Ji et al. 2023 has proven effective in enabling LLMs to understand the impact of their actions and adapt their behavior accordingly, which effec- tively aligns the model's behavior with human intentions. Based on the different forms of explicit guidance, it can be categorized into four types: label-based (Guerdan et al. 2023). reward-based (Wu jet al. 2024), demonstration- -based (Dasari al. 2023). and comparison- based approaches (Ouyang iet al. 2022). One of the most effective training paradigms among them is reinforcement learn- ing (Ge al. 2024). The significant reason that it is relatively easier to collect and evaluate the quality of responses compared to expert answers (Rafailvv et al. 2024). improving proficiency in translation (Kreutzer al. 2018). summarization (Liu al.] 2020). and instruction-folloiinn Ra- mamurthy et al. (2023). In retrieval augmentation, Ma et al. 2023 employs Proximal Policy Opti- mization (PPO) (Schulman et al. (017b) to refine the search query, using the match score between the predicted answer and the gold standard answer as the reward. Furthermore, this use of gold answers as feedback enables the model to Iearn when to utilize tools to expand its capabilities and interact with the real world (Qiao et al. 2024).
Different from most of the previous methods, which are solving a context bandit problem, our pro- posed SmartRAG contains multiple steps of interactions between the policy and the environment and learns from the environment feedback rather than human feedback.
3 METHODOLOGY
Let x be the input question and y be the golden answer. For RAG systems, there is usually companion retriever RR that takes query q as input and returns an observation iee o = R(q).
The core part of SmartRAG a policy network no. which is an LLM with the parameters being 0. The policy network takes the state s as input, which is combination of the original question and the optional observations, i.e.
[x, os)
where os is the concatenation of all the historical observations returned from the retriever while the action sampled from the output of the policy network. In SmartRAG, there are two different types of actions: answer and retrieve. Let an be the first token of a. Then ag is a special token that can be either [ANSWER] or [RETRIEVE] If ag = [RETRIEVE], means that the policy thinks it necessary to retrieve with query a1:, which is the rest part of a. In this case, the retriever will be called obtain the new observation R R(ar:) that will be appended to os. Then the policy network will be called again to produce the next action by ① On the other hand, if ag = [ANSWER]. it indicates that the policy network thinks there no need for an extra retrieve and is able to produce the right answer with a1.- Thus we obtain the final answer ý a1:-
The Sma pipeline is shown in Algorithm I We can set the quota for retrieval as N. When the number of retrieves has already met the quota, can force the policy network to generate the answer by forcing ag. which is the first output token, to be [ANSWER] Thus we can avoid the policy network to keep retrieving if no satisfying observation is obtained from the retriever. Of course the retrieval quota is not a mandatory setting. We can remove it if unnecessary in certain applications.
3

---

Published as a conference paper at ICLR 2025
Algorithm SmartRAG Pipeline
Require: Policy Network Te. Retriever R
1: Input: input question x, retrieve quota N, observation os <- []. retrieve count n
2: while π S N do
3: if 22 = N then,
4:
To ([x, os)
s.t. a0 [ANSWER]
5:
else
6:
To ([x, os)
7:
n +-++1 n +1
8:
if ag [ETRIEVE] then,
9:
+ R(a)). os + [os, 0)
10:
else
11:
return a1:
3.1
WARM-UP USING SUPERVISED FINETUNING
Since our policy network is initialized using base LLM, the output of which does not necessarily follow our design that the first output token is either [RETRIEVE] or [ANSWER] To achieve the output format requirements, we first warm-up the base LLM with our constructed dataset.
For each question x, we construct three different types of input-outpus pairs for the warm-up step. The first pair is a direct answer. In this case, the input is the original question while the output the golden answer prefixed with the special token [ANSWER]. Thus we have the first data point, i.e.
x, os M] a [[ANSWER]. y].
The second data point encourages the policy network to apply retrieval action. To achieve a diverse initial state for the future RL algorithm, ask GPT-3355 [ to rewrite the original question to produce a diverse query for the retriever. The second data point then becomes
$ = [x, os = =], a [RETRIEVE], GPT-Rewriter(x))
(3)
The third data point enables the initial policy network to know how to answer with the observations, it is constructed as
= x, os ((GGTTRRwwiter(x)) a [[ANSWER], y)
After finetuning the base L.LM with the above-constructed dataset, we can obtain an initial policy network To, that can output with our desired format and generate reasonable results. It should be noticed that the warm-up step only provides reasonable initial policy. It does necessarily mean that the constructed dataset from ( to (4) is the "golden" answer since the behavior of the policy network will be further optimized in the following reinforcement learning step.
We can further enhance the initial policy by slightly modifying the warm-up dataset. For the data points that the base LLM can already correctly answer, we remove the constructed SFT data points defined in (3) and 4.. Otherwise, we remove the constructed SFT data point defined in (2). By doing so, the warm-up dataset already contains the information about whether the base LLM has the corresponding knowledge for the question, leading to better initial policy, denoted Toz-
3.2 JOINTLY OPTIMIZATION USING REINFORCEMENT LEARNING
After obtaining the initial policy mez in the SFT warm-up stage, we then start jointly optimizing the whole framework using reinforcement learning. Following Algorithm I we can generate trajectory (s(0) (0) (0) s()) (T) r(T)), where s(t) (!) and (!) represents the state, action and reward the t-th step respectively. T' stands for the trajectory length and we then define the details of the reward ((t).
gpt-3.--turbo-1100
4

---

Published as a conference paper at ICLR 2025
Our reward function is only related to the action ( If the policy produces related answer, the system will get a positive reward. If the policy uses a retrieval opportunity, it should be penalized for the retrieval cost. Thus the reward function can be written
I
(A) a1: v) + Fi
(1) al: 8)
if a (1) [ANSWER],
(())
(5)
a
if = [RETRIEVE].
Here T is the identity function. It equals to 1 if the arguments identical and equals otherwise. Besides the exact match reward term, we also add a soft match term in ag
(t) = [ANSWER]
a case, which is the Fi score between (4) and y. In addition, is a positive constant penalizing the retrieval cost. In some applications, where retrieval cost is not an issue, we can simply set a to be 0
remove the penalty. The overall objective then becomes
E
(6)
where y is the discount factor.
By such design, encourage the whole SmartRAG system to produce the correct answer with a minimal number of retrieval trials. As a result, the system will only retrieve when the model does not inherently know the answer to the question while the retriever happens to have the related knowledge. The system will not call the retriever when the base LLM already knows the answer or neither the base LLM nor the retriever does not know the correct answer. All the modules in the AmartRAG pipeline work as whole system rather than as separate independent parts.
We optimize the objective 6 using proximal policy optimization (PPO) (Schulman et al. 2017
4 EXPERIMENTS
Dataset. Following Ma et al. (0023). we use three open-domain QA datasets for evaluation, namely PopQA (Mallen et al.) 2023a). AmbigNQ (Min et al.) (2200 and HotpotQA (Yang et al. (018)) Besides, we also use other datasets like TrivialQA (Joshi et al.] 2017). OpenBookQA (Mihaylov et al.) 2018), MedQA-cn (Jin et al.] (021)) and ARC-c (Clark et al.) (218) for more analysis and comparisons. All the details of the related datasets are shown in Appendix A
Retriever. We utilize Bing search engine as our retriever R, which offers a broad knowledge scope and ensures up-to-date factual information. The search engine will return multiple results given query. We select the top K results and concatenate the snippets from these results as the observation. We leave how to utilize these snippets the answer generator. In our experiments, we set K without other clarifications.
Training Details. In the SFT warmup stage, we finetune the network for 1 epoch with the learning rate being 3e The batch size is be In the reinforcement learning stage, we use an on- policy sampling strategy. For each iteration, of all these trajectories 5120. The policy network is then trained for epoch on these samples with the learning rate being 2e 6 and batch size being 32. The retrieval penalty in (5) is set to 0.2 and the retrieval quota N is specified to be to fairly compare with other baselines.
sample many trajectories such that the total length
More implementation details are introduced in Appendix A lowing baselines.
We compare method with the fol-
Vanilla RAG retrieves the search engine using the input question. Then asks the LLM to generate the final answer with the observations returned from the search engine in few-shot manner.
SFT RAG uses the same pipeline as our SmartRAG. Similar to Vanilla RAG, it uses the input question to retrieve the search engine. We then finetune the answer generator given the input ques- tion along with the observations. SFT RAG can be regarded as only training the answer generator and leave the other two modules as they are.
GPT4 Rewriter SFT RAG adds query rewrite module than SFT RAG. We ask GPT4, which one of the most powerful LLMs, to rewrite the search query before feeding it to the search engine. In this case, the answer generator and the query rewriter separately optimized.
5

---

Published as a conference paper at ICLR 2025
PopQA
AmbigNQ
HotpotQA
Average
Model
EM
F1
EM
F1
EM
F1
EM
F1
Flan-T5 large
ICL
7.03
10.16
2.47
7.30
12.46
20.42
7.32
12.63
SFT
12.76
17.70
4.56
10.92
12.72
19.72
10.01
16.11
Vanilla RAG
34.36
38.80
30.72
40.16
22.06
32.57
29.05
37.18
SFT RAG
38.15
42.05
31.63
41.05
24.01
34.03
30.93
38.68
GPT4 Rewriter SFT RAG
39.27
42.68
33.43
42.40
24.12
34.33
32.27
39.80
SmartRAGGours)
42.50
45.95
33.71
42.75
25.54
35.23
33.29
41.25
LlaMa-2 7B
ICL
21.79
24.61
18.63
27.12
17.62
25.60
19.35
29.11 11
SFT
27.91
31.54
22.47
31.28
20.74
29.52
23.71
30.78
ReAct Yao et al. 2023
27.21
31.49
22.45
31.24
21.91
29.83
23.86
30.85
IRCaT (Tnvedi et 2023
9.22
17.93
13.52
25.94
17.42
27.13
13.39
23.67
Self-RAG (Asi al.
2024
1.07
23.49
6.12
26.48
6.80
17.53
4.66
22.50
SKR Wang et al. 2023
36.04
39.59
35.41
44.75
23.63
33.44
31.69
39.26
GPT4 Rewniter + SKR
39.25
42.94
35.03
44.68
23.39
33.61
32.73
40.41
Vanilla RAG
32.26
40.06
33.12
46.20
24.18
34.94
29.82
40.40
SFT RAG
38.71
42.33
36.58
46.09
25.59
35.15
33.64
41.19
GPT4 Rewriter SFT RAG
42.22
46.04
36.13
46.02
25.29
35.38
34.54
42.49
SmartRAG(ours)
44.32
47.60
37.89
47.76
26.60
36.42
36.27
43.93
Table 1: Performance open-domain QA datasets. Bold numbers indicate the best result while underlined number indicates the second- result.
Self-RAG (Asai ct al. 2024) leverages GPT-4 to construct large-scale dataset encompassing tasks such determining whether retrieval necessary, assessing the relevance retrieved content, and formulating appropriate responses. Subsequently, SFT is employed to enable the LLM to adaptively perform retrieval and generation. Retrieval is triggered when the model predicts special retrieval token exceeding a predefined threshold, followed by answer generation.
SKR (Wang al. 2023) is also adaptive retrieval algorithm that utilizes the model's correctness in responding to the training set as a supervisory signal: correct responses indicate that retrieval is unnecessary, while incorrect responses suggest the need for retrieval. Based this, classifier is then trained to predict whether sample requires retrieval.
Yao et al. 2023) uses the inherent ability of the base network to decide how to retrieve and answer using sfain-of-thougee (CoT) (Wei al. 2022). It can also self-reflect the previous queries and answers to arrive at the final answer.
ReAct
IRCoT (Trivedi et al. 2023) enhances each step of the CoT generation process by incorporating knowledge retrieval steps during the generation process. Note that both ReAct and IRCoT use multi- step retrieval such that the retrieval cost is higher than other baselines and our SmartRAG.
All the above baselines use at least one time of retrieve. To further show the baseline without any retrieval, also compare with two extra baselines, namely ICL (In-Context Learning) which is simply a few-shot inference, and SFT which is directly finetuned on the golden answer.
Following Jeong al. (2024) ; Asai et al. (2024)
two different types of base LLMs, i.e. Flan-T5 large (Chung et al. 2024h and LlaMa2-7B aB (Touvron et al. (023). Two different metrics are adopted in the comparison: the exact match (EM) and the Fl score.
4.1 ADVANTAGE OF JOINT OPTIMIZATION
We compare our proposed SmartRA with all the baselines in Table [ In this table, ICL and SFT show the results without retrieval. The performance is very poor, indicating that retrieving

---

Published as a conference paper at ICLR 2025
AmbigNO
SPT
SPT
SCS
Bowntttere
301-
GPT: Rourtteree
SCE
De
Se11AC
20
Figure 2: Fl Score of different retrieval percentage across three datasets on Flan-T--
is necessary for these datasets. All the algorithms that have retrieve-related modules significantly increase the performance.
By jointly optimizing all the related modules, our SmartRAG is able to outperform all the counter- parts that use separately optimized modules. This validates our argument that joint optimization necessary to make all the modules properly work together to achieve optimal performance.
Specifically, SFT RAG only trains the answer generator while GPT4 Rewriter + SFT RAG sepa- rately learns both the answer generator and the query rewriter. Both methods performs worse than SmartRAG on all the datasets. In addition, we find that GPT4 Rewriter + SFT RAG sometimes performs worse than SFT RAG, meaning that the rewritten query provided by GPT4 may be worse than the original question. In this sense, the query rewritten by GPT4 is just reasonable but not golden.
Compared to SFT RAG, SKR further learns the decision maker to decide whether to retrieve, thereby reducing retrieval costs while incurring degradation in performance. GPT4 Rewriter + SKR separately learns all three modules and thus performs better than vanilla SKR. Yet both EM and Fl score are worse than SmartRAG.
We only provide the result of ReAct and IRCoT for the LlaMa2 base model because these two methods produce extremely poor results on Flan-T5. Without any adaptation the real applica- tion, which includes tasks that are either model-dependent or retriever-dependent, simply using the inherent reasoning ability of the base models has no means of achieving optimal performance.
4.2 SMARTRAG LEARNS WHEN TO RETRIEVE
The first token of the action a indicates whether to retrieve or not. In practice, the policy network out- puts Beroulli distribution over [[ANSWER], [RETRIEVE]], we can manually set up threshold on the logit of [ANSWER] to control the retrieval ratio. Figure shoows the Fl score of SmartRAG with different thresholds. We compare AmG with four baselines. The first one is SFT. Since does not have any decision maker for whether to retrieve, it only has two states: full retrieve and no retrieve. The other three baselines are SKR (Wang et al. 2023). GPT4 Rewriter + SKR and Self-RAG (Asai et al. 2024). We can see that the results of Am is to the upper left side of the other methods, meaning that SmartRAG can produce similar Fl score with less retrieval cost. On both PopQA and HotpotQA datasets, SmartRAG can even produce a higher F1 score with less retrieval percentage than SKR. To achieve this, SmartRAG should put more retrieval efforts on the scenarios when retrieval does bring benefit. These are instances where the base L.MMs lack inherent knowledge of the correct answer, yet the necessary information is encompassed within the database.
A deduction about the above arguments is that SmartRAG should not retrieve anything if the database hardly contains any useful knowledge. This is indeed the case for the Bing search engine on OpenBookQA, MedQA-cn and ARC-c datasets. We list the retrieval percentage and accuracy on these datasets in Table 2 SFT and SET RAG produce extremely close accuracy since the database contains very little knowledge of these datasets. The SKR method still produces high retrieval percentage, which wastes the retrieval resources. Our SmartRAG is aware of the property of the database and learns not retrieve it in these cases.

---

Published as a conference paper at ICLR 2025
Retrieval Percentage MedQA-cn
Accuracy MedQA-cn
Method
OpenBookQA
ARC-c
OpenBookQA
ARC-c
SFT
0.00
0.00
0.00
67.20
31.66
54.95
SFT RAG
100
100
100
68.20
31.11
56.99
SKR (Wang et al. 2023
5.60
84.29
9.98
67.20
31.89
54.86
SmartRAG
0.00
0.00
0.00
66.60
33.31
55.72
Results of different one-choice questions on Flan-TS-larre. All the methods produce similar accuracy since the database does not contain the useful answer. Our SmartRAG learns not to retrieve for it does not bring any benefit in these cases.
Table 2:
Question Who was the producer of 9?
Golden Answer: Shane Acker Observation with original question: 5 Things to Know About THE9, China 2019s New Pop Idol Group: Predicting the shelf-life of pop act is never straightforward, but it will be interesting to see just how long THE9 stick around given the rapid demise of boyband Nine Percent, who were spawned by 2018-2019s Idol Producer.
Output: [ANSWER] Tim Burton
Rewritten Query:
Who produced the movie 9?
Observation with rewritten query: (2009) IMDb: Shane Acker first made 9 (2009) as a ten minute short film while he was still at UCLA. It was nominated for Best Animated Short at the Oscars, and although didn't win. Acker was offered the chance expand it into feature film. Output: [A NSWER Shane Acker
Question 2: When was the contruction of the makkah royal clock tower hotel started? Golden Answer: 2004
Observation with original question: Ryugyong Hotel: Ryugyong Hotel was designed by North Korean firm Baikdoosan Architects and Engineers, which is also the contractor of the project before 1992, the firm had built the building to
height before the construction suspended 1992. ... However, considering Ryugyong Hotel 112 meters shorter than Willis Tower, 360,000 square meters is still ... Output: [ANSWER] 1992
Rewritten Query: Start date construction for the Makkah Royal Clock Tower Hotel Observation with rewritten query: Makkah Royal Clock Tower, Saudi Arabia World Construction Network: The Makkah Royal Clock Tower complex, also known the Abraj Al-Bait Towers, located near Masjid Haram Mecca, Saudi Arabia, is mixed residential and hotel complex. Construction on complex was started in 2004 and finished in 2012.
Output: [ANSWER] 2004
Table 3: Comparison between observations with orignal questions and rewritten queries.
4.3 SMARTRAG LEARNS WHAT TO RETRIEVE
Once SmartRAG decides retrieve the database, the next task is what retrieve. Directly searching the original question produces suboptimal performance, as Table [ already demonstrates. In this section, showcase how SmartR AG rewrites the retrieving query in Table 3 with two examples. In the first question Who was the producer of 9, directly searching the question cannot provide any useful information because the term 9 is quite confusing. As result, the retriever does not return any useful information and thus the final answer is wrong. The rewritten query Who produced the movie 9 aware that 9 is movie name and explicitly uses the term movie retrieve. With this modification, the returned observation is more related and contains the correct answer, which is then output by the answer generator in the next step. In the second example, the subject of the original question is makkah royal clock tower hotel. Directly searching the question returns information about the hotel with no messages related to the start date of construction for the hotel. The subject modified to start date of the construction by the policy network, making the search more accurate.
To further validate that the query rewrite module does play its role in enhancing the overall perfor- mance, we replace all the retrieving queries in SmartRAG with the original question. The evaluation result is shown in Table 4 (comparing lines SmartRAG and Replace Query). Replacing the rewritten query in SmartRAG with the original question does degrade the performance on all the datasets over both exact match and Fl score metrics. We also introduce a new metric hit. It tracks the ratio of observations that contain the golden answer.
8

---

the NBA's single-game scoring record with 100 points in Published as a conference paper at ICLR 2025
PopQA
AmbigNQ
HotpotQA
Method
EM
F1
hit
EM
F1
hit
EM
F1
hit
Replace Query
39.13
42.98
50.49
33.41
42.19
49.31
24.31
34.04
33.83
Replace Generator
40.95
43387
52.59
33.23
42.02
49.78
25.63
35.04
35,89
SmartRAG 42.50
45.95
52.59
33.71
42.75
49.78
25.54
35.23
35,89
Table 4: Comparison to replace the query with the original questions and replace the generator with the pre-RI. warm-up model.
Question Who was the screenwriter for On the Beach? Golden Answer: Nevil Shute
Observation: On the Beach, Summary and Facts, Britannica Encyclopedia Britannica: On the Beach, American dramatic film, released in 1959, that was in the aftermath of imagined World War III. It apocalyptic novel of the same name by Nevil Shute. The fatal fallout of nuclear war in year 1964 serves as the fictional backdrop for romance between a navy On the Beach (1959)- FilmAftinity: On the Beach is a film directed by Stanlcy Kramer with Gregory Peck, Ava Gardner, Fred Astaire, Anthony Perkins
was based
Before PPO Output: [ANSWER] Stanley Kramer
After PPO Output: [ANSWER] Nevil Shute
Question 2: What's the most points scored in nba game by single team?
Golden Answer: 186
Observation: NBA Single Game Leaders and Records for Points: Checkout the complete list of NBASin- gle Game Leaders and Records for and more on Baakettall-eeferenee.comm Year-by- Year Top 10: Checkout the complete list of NBA and ABAYear-by- bearrr Top 10 Leaders and Records for and more on Basketball- Reference.com ... shows the teams that scored the most points in one game in NBA history, including regular season and Playoffs. Most points scored by a team in an NBA game: Pistons set the single-gand game scoring record for an NBA team with 186 points in an all-time shootout against the Nuggets. Detroit defeated Denver 186- 184 in... Who has scored the most points in a single NBA game? ESPN: On March 2, 1962, Wilt Chamberlain
Dec. 13, 1983, Detroit
the Philadelphia Warriors' 169- 147 win over the New York Knicks. The NBA didn't 't add the...
Before PPO Output: [ANSWER] 370
After PPO Output: [ANSWER] 186
Table 5: Improvement of answer from observation between before and after PPO.
4.4 SMARTRAG LEARNS How TO ANSWER
Given the correct time to retrieve and the proper query for retrieval, the system should also learn how to well utilize the observations from the retriever produce the final answer. To validate that SmartRAG learns how to answer with the observations, we replace the final answer generation step with the version before PPO. Note that this generator before PPO is already finetuned with data points that contain observations from the retriever. So it also learns how to answer based on the observations. The only difference that it is not jointly optimized with other modules.
Table 4 (comparing lines SmartRAG and Replace Generator) shows the replacement also degrades the performance on most of the metries. The hit metric remains the same because they use the same queries for retrieving and thus have exactly the same observations. Two examples are presented in Table In both examples, the correct answer is included in the observations (highlighted in green). SmartRAG can correctly find the answer from the lengthy information. However, the answer gen- erator before PPO either finds the wrong information (the first example) or hallucinates the wrong answer (the second example).
4.5 DATASET TRANSFER
When jointly optimized, SmartRAG's policy network gains awareness of both the companion knowledge base and its own capabilities. Specifically, the policy network knows what can or cannot be retrieved from the knowledge base. In addition, it is also aware of what itself knows or does not know. When transferring the optimized SmartRAG system to a previously unseen dataset, such kind of awareness can help the system achieve better performance. In this section, we test the model

---

Published as a conference paper at ICLR 2025
TriviaQA
2WikiMuliHHopOA
MuSiQue
Method
EM
F1
EM
EM
SFT RAG
47.72
56.99
23.08
28.21
3.27
11.15
SKR (Wang et al. 2023
37.51
45.35
22.61
27.65
3.02
SmarlkA
48.87
59.17
23.55
29.26
4.84
13.33
Table 6: Dataset Transfer Flan-T5-large.
PopOA
PopOA
intall polic1
Intal policyl
install polry?
Install po
policy
policy
25
final policy
final pelicy
Exact Match
15-
Retrieval Percentage
Retrieval Percentage
Figure 3: Influence of different initial policies for SmartRAG.
on three previously unseen datasets, namely TrivialQA (Joshi et al. 2017). 2 WikiiuultiHopO/ (Ho al. 2020) and MuSiQue (Trivedi al. 2022). We compare SmartRAG with two other baselines including SFT RAG and SKR(Wang et al. (202)). The result is shown Table SmartRAG achieves better transfer performance than both basclines, validating our argument.
trained in Section 4.
To further demonstrate the awareness of both the knowledge base and the policy model itself, categorize the TrivialQA test set into three types: 1) questions answerable without retrieval, 2) those requiring retrieval for a correct answer, and 3) those unanswerable even after retrieval. We fix the retrieving threshold to be 0.5 and evaluate the retrieval ratio for each of the categories. Theoretically speaking, retrieval is only needed for the second type data points. The model should not retrieve for the first type because it knows it already knows the answer. It will also not retrieve for the third type because it knows the knowledge base does not have such information. The retrieval ratio on all three types of data points for SmartRAG is 10.59%, 42.59% and 29.04% respectively, with the highest probability of retrieval for the second category, thereby affirming our assertion. Using different thresholds and alternative datasets may yield different results, but we always find that the second category has the highest retrieving ratio.
4.6 INFLUENCE OF THE INITIAL POLICY
We have proposed two initial policies in Section 3.1 To study how the initial policies can influence the final performance, we start from both initial policies and go through the same PPO procedure. The exact match and score of the initial policy and final policy for both initializations are shown in Figure 3 Too (initial policy 1) clearly performs worse than тæ; (initial policy 2) because the corresponding warm up dataset contains less information. Applying PPO enhances both policies. The final policy with better initialization still outperforms its counterpart with worse initialization, indicating that dell-ne initial policy benefits the PPO procedure.
CONCLUSION
We introduce SmartRAG, a novel framework that can decide when to retrieve, what to retrieve, and how to answer with/without the retrieved observations. All these functions are accomplished by policy network and a companion retriever. SmartRAG jointly optimizes the whole system using reinforcement learning. Empirical results show that SmartRAG outperforms previous methods whose modules are separately optimized. For systems comprising multiple modules, it is necessary optimize them from holistic perspective- challenge reserve for future research endeavors.
10

---

Published as a conference paper at ICLR 2025
ETHICAL CONCERNS
This paper primarily focuses on an end-to-end RAG framework that addresses the questions of when to retrieve, what to retrieve, and how to answer within single LLM. Our method utilizes open- source datasets and Bing retrieval for training and testing, with all data being publicly accessible. While our approach demonstrates the ability to effectively enhance the model's decision-making capabílities and improve the accuracy of retrieved information, as well the final answers, we hope that this design-more aligned with real-world applications- can enhance the efficiency of RAG pipelines. Furthermore, we aspire for this promising framework to applied across various L.MM decision-making scenarios.
REPRODUCIBILITY STATEMENT
The details of experiment settings are provided in Section and more detailed description and implementation setting can be found in Appendix [ All datasets and models used in our work are publicly available. Additionally, we have uploaded our code in the supplementary materials and will open-source it upon acceptance of the paper.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-RAG: Learning to retrieve, generate, and critique through self-reflection. In The Twelfth International Confer- ence on Learning Representations, 2024. URL aattps://openzevate net/ for.mmidrmm hSyW5go0v8
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Yunxuan Li, Xuezhi Wang. Mostafa Dehghani, Siddhartha Brahma, al. Scaling instruction-finetus lan- guage models. Journal of Machine Learning Research, 25(70):1-53, 2024.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 10X3::8800.05455 2018.
Sudeep Dasari, Abhinav Gupta, and Vikash Kumar. Learning dexterous manipulation from exem- plar object trajectories and pre-grasps. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 3896. IEEE, 2023.
Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, and Haofen Wang. Retrieval-and -uggeettd generation for large language models: A survey. arXiv preprint aXXIv22312.10997, 2023.
Yingqiang Ge. Wenyue Hua, Kai Mei, Juntao Tan, Shuyuan Xu, Zelong Li, Yongfeng Zhang, et al. Openagi: When Ilm meets domain experts. Advances in Neural Information Processing Systems, 36, 2024.
Michael Glass, Gaetano Rossiello, Md Faisal Mahbub Chowdhury, Ankita Naik, Pengshan Cai, and Alfio Gliozzo. Re2g: Retrieve, rerank, generate. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. pp. 2701-2715, 2022.
Luke Guerdan, Amanda Coston, Zhiwei Steven Wu, and Kenneth Holstein. Ground (less) truth: A causal framework for proxy labels human-alee algorrith decision-making In Proceedings of the 2023 ACM Conference Fairness, Accountability, and Transparency, pp. 688-704, 2023.

---

Published as a conference paper at ICLR 2025
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938. PMLR. 2020.
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi- hop QA dataset for comprehensive evaluation of reasoning steps. In Donia Scott, Nuria Bel, and Chengqing Zong (eds.), Proceedings of the 28th International Conference on Computational Linguistics, pp. 6609-6625, Barcelona. Spain (Online), December 2020. International Com- mittee on Computational Linguistics. doi: 10. 10.111666///22000 100016866/1/0000...ooggmmmm coling-main.580. URL https:
//aclancholooy.orr//200000 www.g//gm....iiii............... 580
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.0966 2021.
Yucheng Hu and Yuxing Lu. Rag and rau: A survey on rerrieallaugmmented language model in natural language processing. arXiv preprint arXiv:2404. 19543, 2024.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini. Fabio Petroni, Timo Schick, Jane Dwivedi Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models. Journal of Machine Learning Research, 24(251): 1-43, 2023.
Soycong Jeong. Jinheon Back, Sukmin Cho, Sung Ju Hwang, and Jong C Park. Adaptive-rag: Learning to adapt retrieval- -ugieete large language models through question complexity. 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. 7029-7043, 2024.
Jiaming Ji, Tianyi Qiu, Boyuan Chen, Borong Zhang, Hantao Lou, Kaile Wang, Yawen Duan, Zhonghao He, Jiayi Zhou, Zhaowei Zhang, al. Ai aligmment: comprehensive survey. arXiv preprint arXiv:2310. 230.:22910.98222 2023.
Zhengbao Jiang, Frank Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi-Yu, Yiming Yang. Jamie Callan, and Graham Neubig. Active retrieval augmented generation. arXiv preprint arXiv:2305.06988 069883...68..... 2023.
Di Jin, Eileen Pan, Nassim Oufattole, Wei-Hung Weng, Hanyi Fang, and Peter Szolovits. What dis- ease does this patient have? a large-scale open domain question answering dataset from medical exams. Applied Sciences, 11(4))6422, 2021.
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. arXiv preprint arXiv: 1705.03551, 2017.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generaliza- tion through memorization: Nearest neighbor language models. In International Conference on Learning Representations, 2019.
Omar Khattab, Keshav Santhanam, Xiang Lisa Li, David Hall, Percy Liang, Christopher Potts, and Matei Zaharia. Demonstrate-searchh preddit:: Composing retrieval and language models for knowledge-intensive nip. arXiv preprint arXiv:2212. 14024, 2022.
Julia Kreutzer, Joshua Uyheng, and Stefan Riezler. Reliability and learnability of human bandit feedback for sequence-tor seeeuueee reinforcement learning. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1777_ 1788, 2018
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466 2019.
12

---

Published as a conference paper at ICLR 2025
Patrick Lewis, Ethan Perez. Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih. Tim Rocktäschel, et al. eetrieval-augment genera- tion for www.ledge-intensiy nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
Fei Liu et al. Learning to summarize from human feedback. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, 2020.
Yanming Liu, Xinyue Peng. Xuhong Zhang. Weihao Liu, Jianwei Yin, Jiannan Cao, and Tianyu Du RA-ISF: Learning to answer and understand from retrieval augmentation via iterative self- feedback. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Findings of the Association for Computational Linguistics ACL 2024, pp. 4730-4749, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics. doi: 10.18653/v122024 findings acl. 281. URL htwp::tpp////m Teietps://accaanthooooyoo org 2 htt.:o///22.....nnssaaaaaaaaeeaa act.nngsssa.....0. 281
Xinbei Ma, Yeyun Gong, Pengcheng He, Hai Zhao, and Nan Duan. Query rewriting in retrieval- augmented large language models. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 5303-5315, 2023.
Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das. Daniel Khashabi, and Hannaneh Hajishirzi. When not to trust language models: Investigating effectiveness of parametric and non non-parametric memories. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 9802-9822, 2023a.
Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. In Anna Rogers, Jordan Boyd-Graben and Naoaki Okazaki (eds.), Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Pa- pers), pp. 9802-9822, Toronto, Canada, July 2023b. Association for Computational Linguis- tics. doi: 10.18653/v 1/t012.aclllooog..e..
URL https: http://///aaaaathhhoooooo org.2023.
ac ong. 6
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02788 2018.
Sewon Min, Julian Michael, Hannaneh Hajishirzi, and Luke Zettlemoyer. Ambigqa: Answering ambiguous open-domain questions. arXiv preprint arXiv:2004. 10645, 2020.
Sewon Min, Kalpesh Krishna, Xinxi Lyu, Mike Lewis, Wen-tau Yih, Pang Koh, Mohit Tyyer, Luke Zettlemoyer, and Hannanch Hajishirzi. Factscore: Fine- ceee-graind atomic evaluation of factual pre- cision in long form text generation. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 12076-12100, 2023.
Long Ouyang. Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Gray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models follow instructions with human feedback. In Ad- vances in Neural Information Processing Systems, 2022.
Torummidd-T5xAAXXB
Shuofei Qiao, Honghao Gui, Chengfei Lv, Qianghuai Jia, Huajun Chen, and Ningyu Zhang Making language models better tool learners with execution feedback. In Proceedings of the 2024 Con- ference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 3550-3563 2024.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton- Leyton-Brown, and Yoav Shoham. In-context retrieval- auggeetedd language models. Transactions of the Association for Computational Linguistics, 1 11:1316-1331, 2023.
13

---

Published as a conference paper at ICLR 2025
Rajkumar Ramamurthy, Prithviraj Ammanabrolu, Kianté Brantley, Jack Hessel, Rafet Sifa, Chris- tian Bauckhage, Hannanch Hajishirzi, and Yejin Choi. Is reinforcement learning (not) for natural language processing: Benchmarks, baselines, and building blocks for natural language policy op- timization. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL pdf ? Pilfiid-8aa BadzzasuuduuU
net./
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford. and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.06347, 2017a.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017b. URL https://arxiv:
org/ http//bb/1177...o 06347
Jiejun Tan, Zhicheng Dou, Yutao Zhu, Peidong Guo, Kun Fang, and Ji-Rong Wen. Small models, big insights: Leveraging slim proxy models decide when and what to retrieve for llms, arXiv preprint arXiv:2402. 2024.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. musique: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 10:539-554, 2022.
Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. Interleaving re- trieval with -fannoof-though reasoning for snowledge-intensing multi-step questions. In Pro- ceedings of the 6/st Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers), pp. 10014-10037, 2023.
Hongru Wang, Boyang Xue, Baohang Zhou, Tianhua Zhang, Cunxiang Wang, Guanhua Chen, Huimin Wang, and Kam-fai Wong. Self-dc: When to retrieve and when to generate? self divide- and-conquer for compositional unknown questions. arXiv preprint arXhv:2402.13514, 2024.
Yile Wang, Peng Li, Maosong Sun, and Yang Liu. Self-knowledge guided retrieval augmentation for large language models. In Findings of the Association for Computational Linguistics; EMNLP 2023, pp. 10303-10315, 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe Le, Denny Zhou, et al. Chain-of- prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
Zeqiu Wu, Yushi Hu, Weijia Shi, Nouha Dziri, Alane Suhr, Prithviraj Ammanabrolu, Noah A Smith, Mari Ostendorf, and Hannaneh Hajishirzi. Fine-grained human feedback gives better rewards for language model training. Advances in Neural Information Processing Systems, 36, 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov, and Christopher D Manning. Hotpotqa: dataset for diverse, explainable multi-hop question answering. arXiv preprint / 18990096000 2018.
Shunyu Yao. Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. In International Conference on Learning Representations (ICLR), 2023.
14

---

Published as a conference paper at ICLR 2025
APPENDIX
A
SmartRAG Details
16
A.I Details of Raw Datasets
16
A.2 Details of Warm Dataset Collection
16
A.3 Training and inference Details
17
A.4
Further Explanation about Reward Design
19
A.5
Inference Latency and Training Time
19
B
Analysis
19
B.
Training Process
19
B.2
Retrieval Number
20
B.
Query Generator
B 4
Adaptive Retrieval
21
B.5
Qualitative Cases
22
15

---

Published as a conference paper at ICLR 2025
A SMARTRAG DETAILS
A.1 DETAILS OF RAW DATASETS
concise summary of the utilized datasets presented below:
PopQA (Mallen et al., 2023a) consists of 14k questions designed to encompass factual information in the long tail, potentially overlooked by widely used QA datasets. It is an cetity-centric, open-domain QA dataset featuring entities with varying levels popularity.
AmbigNQ (Min et al.] 2020) provides disambiguated version of Natural Ques- tions(NQ) [Kwiatkosski et 2019). For ambiguous questions in NQ, minimal constraints are introduced to refine them into multiple similar yet more specific questions
HotporQA (Yang al. 2018
multi-hop dataset derived from Wikipedia. The ques- tions are diverse and not limited by any predefined knowledge bases or schemas. Hot- potQA is an English Wikipedia question-answering dataset consisting of approximately 113K crowd-sourced questions, designed to require the introductory paragraphs of two Wikipedia articles for answering. Each question is accompanied by two gold paragraphs and set of sentences within these paragraphs, identified by crowd workers supporting facts essential for answering the question.
TriviaQA (Joshi et al. 2017) is a large-scale, text-based question-answeriigg dataset con- sisting of 950K question-answer pairs derived from 662K documents, sourced from both Wikipedia and the web.
MuSiQue (Trivedi et al. 2022 is constructed by combining multiple single-hop queries, where each reasoning step fundamentally depends on information from another. In our ex- periments, we simplified the task by treating it as single-hop problem for testing purposes.
2WikiMultiHopQA (Ho et al., 2020 multihop QA dataset derived from Wikipedia. requires two hops to answer a question. Similarly, in our paper, we treat it as a single-step test.
OpenBookQA (Mihaylov al. 2018) consists of 5,957 multiple-choic questions, each with four possible answers. The dataset is supplemented with external fundamental scien- tific facts, requiring a deep understanding of these facts to answer the questions accurately.
MedQA-cn (Tin al. 2021) compiles questions from the National Medical Board Exam- inations of Mainland China. It serves as challenging benchmark, encompassing wide range of medical knowledge, including patient profiles, disease symptoms, and drug dosage requirements. This diversity necessitates a contextual understanding to accurately respond the questions presented.
ARC-c (Clark al., 20188) contains complex questions that demand deep reasoning, inte- gration of multiple information sources, and advanced inferential capabilities. This subset pushes Al models beyond basic text matching, requiring them address ambiguity and engage in sophisticated bobliim-ollvng, testing the depth of their reasoning and compre- hension.
The scale of the data used in the paper is presented in Table 7 To reduce the number of training iterations, we combined the three datasets for joint training. It worth noting that this cross-dataset training presents a more challenging task, and we employed multi-dataset training during both the warm-up and PPO training phases. To mitigate the long-tail issues arising from distribution differ- ences among the datasets, we utilized only 20k samples from the HotpotQA dataset to align with the scales of the PopQA and AmbigNQ datasets. Additionally, for the multiple-choice questions, we limited our training to 2k samples from both OpenBookQA and MedQA-en.
A.2 DETAILS OF WARM DATASET COLLECTION
In the earlier mentioned process of handling the warm-up dataset, we utilized datasets of the same scale as those used in training for multi-task warm-up. Each question presents three options: direct response, generating a query, and responding based on the retrieved content. The answer content comes directly from the dataset, while the search query is generated by GPT-3.5. The prompt
16

---

Published as a conference paper at ICLR 2025
Datasets
Category
Training Sct
Test Set
PopQA
Open Domain
13.0k
0.7k
AmbigNQ
Open Domain
99.4k
5.8k
HotpotQA
Open Domain
90.4k 4k
7.4k
TriviaQA
Open Domain
78. 8k
8.8k
OpenBookQA
Scientific Knowledge
5.0k
0.5k
MedQA-en
Medical Domain
10.2k
1.3k
ARC-c
Commonsense reasoning
1.1k
1.2k
Table 7: Details of raw datasets.
provided to GPT-3.5 is: 'Assuming you are a professional retrieval model, rewrite the following question into a query that is more suitable for retrieval: {Question)'. Based on this, the warm-up model responds to each question with 50% probability of either providing direct answer or an answer after retrieval, making it difficult determine the importance of retrieval for each question. However, this setup offers reinforcement learning and balanced initial state to explore different decision-making strategies.
As mentioned in Section 4.6 better initial state can lead to improved results; therefore, we refined the SFT data for initial policy 2. We divided the data into two categories based on the model's ability to provide direct answers in the training set (FI score 0.2): one where the model can answer correctly to some extent, and another where the answers are entirely incorrect. During warm-up, for the first category, the model only provides direct responses, while for the second, it generates query and answers based on the retrieved content. Under this warm-up training strategy, the model acquires certain level of retrieval decision-making capability. This reasoning stems from the idea that if the model can provide correct answers on its own, there is no need to rely on external retrieval. Through joint training our framework, the model further becomes aware of the effectiveness of the database, optimizes its query retrieval, and refines the answer generation, allowing it to explore a better state beyond the warm-up phase.
A.3 TRAINING AND INFERENCE DETAILS
Retriever We utilize the Bing search engine our retriever, which does not require candidate index construction like dense retrievers, nor predefined candidates as traditional resources. Instead, offers broad knowledge scope and ensures up-to-date factual information. concateaate snippets from all retrieved web pages, selecting relevant sentences identified by Bing-much like entering query into browser, then gathering the text displayed the search results page. The results obtained through this concatenation method resemble summaries of each webpage, encapsulating the core content of the retrieved information within a constrained length. The search engine will return multiple results given a query. We select the top K results and concatenate the snippets from these results as the observation. In our experiments, we K 4 without other clarifications.
Model Since our framework employs a single model handle all tasks within the RAG system, employed models of varying sizes for training. To further explore the effects of different ar- chitectures and model sizes, we conducted experiments on the Flan-T5 series models. In addition, trained a Llama2 7B model to validate the effectiveness of our framework is noteworthy that, since we use PPO for training, it requires optimizing the policy model, value model, and reference model simultaneously. To reduce memory consumption and improve training speed, we applied LoRA (Hu et al. 2021) for tuning LlaMa-7B.
larger model.
Settings We use 4 Nvidia A100 with 80GB memory to train our models. For the SFT warmup stage, we use the following training parameters: a learning rate of 3e-4, a batch size of 8, and to prevent overfitting, the number of epochs is set to 1. For reinforcement learning, we set the sampling steps to 5120, 10 threads, 512 steps for each. After sampling, the policy network is trained for epoch, with a learning rate of 2e-6 and batch size of 32.
Retrieval reward For the penalized retrieval reward a in equation 5 assign value of -0.2. This decision is motivated by the substantial influence this parameter exerts on the training dynamics. A high setting would drive the model to predominantly opt for retrieval, whereas low setting would
17

---

Published as a conference paper at ICLR 2025
discourage retrieval altogether. Given the explicit nature of this retrieval reward, any significant deviation from the retrieval-gente generated reward during the model's sampling process may prompt the model to adopt overly simplistic strategy. Consequently, this could hinder the reinforcement learning process, limiting the model to either a retrieval-only or non-retrieval optimization approach, thus compromising its ability sample diverse examples effectively.
Direct Answer(Query) Instructions
You will be presented with a question. If you know the answer, please respond directly. If you don't know the answer, use the Bing search engine to find the necessary information and then answer the question based your observation.
Question: input
Please format your output as follows:
1. If you choose answer the question directly, please use: "[Answer] YOUR-ANSWER
2. If you choose to use the Bing search engine, please use: '[Search]
QUERY"
Please output:
Retrieval Answer Instructions
You will be presented with a question. If you know the answer, please respond directly. you don't know the answer, the Bing search engine to find the necessary information and then answer the question based on your observation.
Question: input
Observation: search
Please format your output as follows:
1. If you choose to answer the question directly, please use: '[Answer] YOUR ANSWER If you choose use the Bing search engine, please use: '[Search] YOUR_SE WWR.REEEARHHQUUEEEE
Please output:
Table Instructions for two-step RAG in SmartRetriverX.
Generation Details During the reinforcement learning training phase, ensure sampling diversity and obtain samples reflecting different decision- makion strategies, we employ Top-K K decoding with K set to 50. This means that each time step, the top most probable candidate words from the probability distribution are selected, while the probabilities of the other words are set to zero. The probabilities of these K candidate words are then renormalized probability distribution. A word is then randomly sampled from this new distribution as the output for the current time step. This process is repeated until an end token generated or the maximum length is reached.
sum to forming new
During inference, we use Beam Search for decoding. For the first token that determines when to retrieve, we apply a threshold-based approach (Asai et al. [2024]: if the confidence score of the generated query exceeds the threshold, retrieval is performed; otherwise, the model directly provides an answer.
List of Instructions The entire inference process consists of either one two steps. The first step involves either directly answering or generating a query. a query is generated, second step follows, which entails answering based on the content retrieved according to the query. Table
18

---

Published as a conference paper at ICLR 2025
presents the instructions utilized in both steps. To ensure consistency in instruction during the fine-tuning process, we designed the two instructions to maintain as similar a format as possible; the instruction for the second step incorporates only one additional component: the retrieved con- tent observation. We employed the same data format across all datasets, and for multiple-choiee questions, we incorporated the candidate answers into the instruction of the question.
A.4 FURTHER EXPLANATION ABOUT REWARD DESIGN
Next, we will illustrate through various examples how our reward design and training strategy enable SmartRAG to achieve optimal performance with minimal retrieval attempts.
Suppose a question that the LLM cannot corretly answer without retrieval, which will result in two distinct trajectories.
Trajectory-1: [Answer[[ncorreet, reward=0)
Trajectory-2: [eerrieee] (reward=-0.2) [AAnser) (correct, reward=2)
In the first trajectory, the value of the state [Answer] since gives the wrong answer. In the second trajectory, the value of the state [Retrieve] is -.. + 2 where is the discount factor. In this case, the model will learn to use trajectory 2 to correctly answer the question since the value of the second trajectory is while that of the first trajectory is only.
When the LLM can directly answer the question without relying on retrieval, two different trajecto- ries also emerge:
Trajectory-1: [Answer] (correct, reward=2)
Trajectory-2: [Retrieve] (reward=-0.2) [Answer] (correct, reward=2)
In the first trajectory, the value of the state [Answer] is 2. In the second trajectory, the value of the state [Retrieve] is -0.2 + 2 * 7(< 2). In this case, the model will learn to use trajectory to answer the question. That being said, algorithm not only learns how correctly answer the question, but also learns how to use the minimal cost.
Therefore, the retrieval reward does not entirely discourage the model from performing retrieval. Instead, it facilitates a tradeoff between the benefits of retrieval and its associated costs, enabling the model achieve higher answer accuracy with fewer retrieval attempts. Simultaneously optimizing the three tasks not only enhances retrieval efficiency but also improves the model's 's response quality, ultimately resulting an LLM capable of addressing end-to-ent RAG challenges effectively.
A.S INFERENCE LATENCY AND TRAINING TIME
Our inference pipeline is the same many of the baselines such GPT4 rewriter + SFT RAG and SKR. CoT-related algorithms like ReAct and IRCoT are much less efficient because they need extra computational cost for the intermediate thought and reflection. Considering that SmartRAG learns when to retrieve, it can automatically skip the retrieval step when unnecessary. As figure 3 shows, SmartR AG can achieve similar (or even better) performance than the baselines using less computa- tional cost. In addition, table 2 demonstrates that SmartRAC refuses to retrieve when the database contains too few information such that it process achieves lower model costs and reduced retrieval costs compared to all other baselines.
computational cost. As result, inference
In terms of training, our approach includes additional RL training phase compared standard SFT. While this introduces an extra training step, the overall improvement in inference performance and cost efficiency justifies this additional effort.
B ANALYSIS
B. TRAINING PROCESS
Our training process includes both warm-up training and reinforcement learning training. The figure illustrates the changes in iteration reward, KI. reward, policy loss, and value loss during the rein-
19

---

Published as a conference paper at ICLR 2025
Cobode Time
Sobode Ranurrd Over
0.00
0.33
0.025
MMAM A
0.30
Spagte howast
11 |
M
0.40
0.23
0.50
Itaratoon
cp Oest
Lasa Tirma
0.275
-0.032
-0.03
8.200
-..00
0.100
-0.00
Karaton
Figure 4: Iteration reward, KL reward, policy loss and value loss during training.
PopOA
PoPOA
K-
Rec Recleca Percentage
Samrbval Prrcentage
kumiiaal
Figure 5: Exact Match, F1 Score, and Retrieval Threshold of different retrieval number K on Flan- Ta-large.
forcement learning phase. The iteration reward, corresponding to the reward in the equation, shows an upward trend throughout the training process, indicating that the model is exploring in the right direction. The increase in KL divergence reward indicates that the model's strategy is being updated, with growing divergence from the previous strategy. At the same time, the loss trends of the policy model and value model are almost the opposite of the reward trend, which aligns with the typical training process of oscillating convergence followed by stable oscillations.
B.2 RETRIEVAL NUMBER
In the retrieval-augmeented framework, the number of retrieved texts also influences the model's performance. Generally, when the number of retrieved documents not large, increase in the number of retrieved documents leads to greater gains for the model. This because when the retrieved information falls within the model's capacity to process, larger set of documents is more likely to contain useful information. Therefore, we investigated the model's performance when the retrieval size is K=1 and K=4 based on initial policy 1. When K=4, it represents retrieval of more documents and stronger retrieval capacity, allowing us to explore the impact of varying- varying-value value database on our method's training results. The results, as shown in the figure, indicate that retrieving more documents leads to better performance. The figure also illustrates the relationship between the retrieval ratio and the threshold under different retrieval intensities. It is evident that when the retrieval database is more likely contain relevant information, our method awareness of the database, which in turn drives the model to rely more on retrieval. This aligns with real-world scenarios, where higher retrieval efficiency should increase the likelihood of the model
demonstrate better
20

---

Published as a conference paper at ICLR 2025
PopQA
AmbigNQ
HotpotQA
Method
EM
F1
EM
F1
EM
F1
LlaMa2-7B
42.22
45.97
34.33
43.88
25.33 25.54
35.32 35.23
GPT35
42.50
45.95
33.71
42.75
Table 9: Performance between different query generator on Flan- Saam-T5-large
SPT SCA Smartse
STT Smarchg
STT SCA
PrsiA
AmbighiQ
HotzergA
SFT
STT
STTT
Tel Natioval Parcoriage
Tackoal Parcentage
netroval Perceriage
Figure 6: Fl Score of different retrieval percentage across three datasets on Flan-T5-las of initial policy
engaging in retrieval, and vice versa. In contrast, the previous adaptive retrieval method was unable to recognize the impact of different databases.
B.3 QUERY GENERATOR
To further evaluate the impact of different query generators on our results, we replaced GPT-3.5 with LLaMA-7B generate the query dataset while keeping all other settings unchanged. We then trained our SmartRAG model on Flan-T5-large, and the results are presented in Table 1. It can be observed that the overall performance does not change significantly. This is because using different L.LM to generate the SFT query dataset is analogous to adopting different initial policy for the RI step. Such variations the initial state exert minimal influence on the final performance, provided the LLM is not exceptionally suboptimal. Moreover, the results from SFT GPT-4 rewriter demonstrate that queries directly rewritten by GPT-4 are not optimal. This limitation arises because GPT-4 lacks prior exposure the knowledge base. In contrast, the SmartRAG framework, which integrates a knowledge base into its training paradigm, enables query optimization that is better aligned with the target knowledge base.
B.4 ADAPTIVE RETRIEVAL
Section 4.2 presents the Fl score results under the initial policy2 state. Here, provide more comprehensive results with a different initial state, where the results for both initiate policy2 and initiate policyl are shown in Tables 6 and 7 respectively. Nevertheless, approach enables the model to learn the concept of when to retrieve, resulting in improved final answers, which stem from the optimization of what retrieve and how answer.
observe that
21

---

Published as a conference paper at ICLR 2025
Ambigho
2CR Smartss
SCt
8
PopQA
27T
2TT
STT
SCa
SmartMG
Pet Percertage
Serieaal Pretertage
Figure 7: Exact Match and Fl Score of different retrieval percentage across three datasets on Flan- Ta-large of initial policy 1.
B.5 QUALITATIVE CASES
Sections and have already demonstrated that SmartRAG significantly improves both the 'what to retrivee' and 'how answer' tasks through quantitative and qualitative analyses. Table TO and I1 presents additional cases from three datasets. Similarly, it can be observed that the reformu- lated queries are more effective in retrieving relevant results, and under the same retrieval conditions, the model trained with PPO also demonstrates superior performance
22

---

Published as a conference paper at ICLR 2025
Question 1: In what city was Rafael Reyes born?
Golden Answer: Cotija
Observation with original question: Rafael Reyes - Political leader, Conservative party Britan- nica: Rafael Reyes (bomn 1850, Santa Rosa, New Granada201 4died February 19, 1921, Bogot00e1, Colombia) was an explorer and statesman who was president and dictator of Colombia from 1904 to 1909. He attempted to give his nation strong me--man rule that would attract foreign investment and foster domestic industrialization.
Output: [ANSWER] Santa Rosa
Rewritten Query:
Birth city of Rafael Reyes
Observation with rewritten query: Rafael Reyes (Lcafar Seyer) Age, Wife, Children, Family, Biog- raphy .... Quick Info2192. Hometown: Cotja. Michoacan, Mexico. Wife: Kat Von D. Age: 47 Years. Some Lesser Known Facts About Rafacl Reyes aka Leafar Seyer.
Output: [ANSWER] Cotija
What other series is the actress who plays Malory Archer well known for? Golden Answer: Arrested Development
Question 2:
Observation with original question: How 'Archer' said goodbye to Jessica Walter on its season finale: Walter, who voiced the character of spy agency chief Malory Archer the animated comedy series, died in March at the age of 80. Deadline subsequently reported that Walter had been working...
Output: [ANSWER] The Simpsons
Rewritten Query: Which other series is the actress known for playing Malory Archer known for?
Observation with rewritten query: Jessica Walter Archer Wiki - Fandom: She performed the voice of Malory Archer, who was strongly influenced by Walter's aforementioned role as Lucille Bluth Arrested Development, in Archer for the first 12 seasons. Walter passed away on March 24th 2021. Season 12 served as her final television role and was released following her death. Its final episode was dedicated to her... ...
Output: [ANSWER] Arrested Development
Question 3: Wheelock Whitney is just one member of the Whitney Family. Where did this American family originate from? Golden Answer:
London, England
Observation with original question: Minnesota businessman Wheelock Whitney Jr. dies - AP News: Wheelock Whitney, whose fraternity brothers at Yale University included former President George H.W W. Bush, was born to prominent family in St. Cloud. He was CEO from 1963-1972 of the investment banking firm Dain and Co., which he helped turn into regional financial powerhouse and eventually became RBC Wealth Management.
Output: [ANSWER] St. Cloud
Rewritten Query:
What is the origin of the Whitney Family, including Wheelock Whitney?
Observation with rewritten query: Whitney History, Family Crest and Coats of Arms: Origins Avail- able: England, Ireland, Whitney is one of the many names that the Normans brought with them when they conquered England in 1066. The Whitney family lived in Herefordshire, at the village of Whitney. Early Origins of the Whitney family.
Output: [ANSWER] England
Table 10: Comparison between observations with orignal questions and rewritten queries.
23

---

Published as a conference paper at ICLR 2025
Question 1: In what country is Dell? Golden Answer: America
Observation: Where Are Dell Computers Made?: Dell 2018s Current Manufacturing Footprint and Supplier Ecosystem. Flash forward to 2023, and Dell now operates major PC manufacturing plants in the Chinese cities Xiamen, Chengdu, Wuhan, and Shenzhen. Dell also owns enormous PC produc- tion facilities just over border from mainland China in Taipei
Taoyuan, Taiwan... The other Dell assembly plants are located in China, Malaysia, Brazil, and the United States. Here's 's Where Dell Computers And Laptops Are Made Today: Aside from the United States, Dell's ODM and final as- sembly facilities in China, Vietnam, Taiwan, Mexico, Brazil, Ireland. Malaysia, India, Poland, and Thailand. Where Dell Computers Assembled? Device Mexico: Dell outsources the as- sembly of its computers to Foxconn in Mexico. Malaysia: Dell2019s assembly factory is located in Penang, Malaysia. Other places where Dell computers are assembled include Ireland, the United States of America, Singapore, Vietnam, Japan, etc.
Before PPO Output:
[ANSWER] Canada
After PPO Output:
[ANSWER] America
Question 2: What is the religion of Paul?
Golden Answer: Christianity
Observation: Paul the Apostle World History Encyclopedia: A Founder of Christianity. In the last century, scholars have come appreciate Paul the actual founder of the religious movement that would become Christianity. Paul was Diaspora Jew, a member experienced a revelation of the resurrected Jesus. Paul and Jesus, The Origin of Paul's Religion, John Gresham Machen
party Pharisees,
Before PPO Output: [ANSWER] Anglicanism
After PPO Output:
[ANSWER] Christianity
Question 3:
Where the think tank founded by neoconservative political analyst born in 1952 established?
Gold Answer:
Washington, D.C.
Search Query: Establishment location of the think tank founded by 1952-born ncoconservative political analyst
Observation: Think tanks and the knowledge201 3pollcy nexus in China: For decades, institutional- ized linkages with the party-state have shaped intellectual activity and subsequently, think tank agency; think tanks were anointed as either establishment oriented, being officially yoked to the party apparatus, or semi-estee aomi-establismment operating with a modicum of independence but with the consent of governmen- tal patrons. Think tanks and the mhowledge2013pole linkages with the strty-state have shaped intellectual activity and subsequently, think tank agency; think tanks were anointed as either establishment oriented, being officially yoked to the party apparatus, or semi-establishment, operating with modicum of independence but with the consent of governmental patrons.
nexus China: For decades, institutionalized
Before PPO Output: [ANSWER] New York City
After PPO Output: [ANSWER] Washington, D.C.
Table 11: Improvement of answer from observation between before and after PPO.
24


Reviews:
[{"review_id": "OFBRtNE9PT", "reviewer": "Reviewer_mFWs", "summary": "The paper introduces SmartRAG, a retrieval-augmented generation (RAG) framework designed to enhance performance through joint optimization of its modules. Unlike traditional RAG approaches, SmartRAG uses a policy network that dynamically decides when to retrieve information, what queries to use, and how to synthesize answers based on retrieved observations. Reinforcement learning with environment feedback is employed to jointly train the system, improving coordination between modules and reducing retrieval costs. Experimental results on multiple QA datasets show that SmartRAG outperforms separately optimized systems.", "questions": "a) There exists many works focused on the question of \"When to retrieval\" in RAG, but in Section 4.1, only one baseline was selected for comparison. Incorporating a set of baseline examples into these studies would enhance comparative and analytical insights. \n\nb) The current baselines primarily apply process-wide optimizations, such as fine-tuning the generator, without optimizing individual modules (e.g., rewriter, generator, and decision-maker) separately. This setup limits the ability to demonstrate how joint optimization in SmartRAG resolves potential conflicts and enhances coordination across modules.\n \nDetailed comments:\na) A dedicated section on the limitations and assumptions underlying SmartRAG, particularly regarding the potential latency issues in real-world applications, would be beneficial. \n\nb) While effective, the joint optimization approach may introduce complexity that could limit the model’s scalability, especially when applied to very large or real-time systems. Additional discussion on computational overhead and training time would strengthen the paper. \n\nc) Considering the potential for optimizing PPO's efficiency or exploring alternative RL algorithms to reduce training costs, it could be beneficial to further discuss these aspects, which may enhance the practicality of the framework.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "a) The paper introduces reinforcement learning to jointly optimize RAG components, integrating retrieval decisions, query rewriting, and answer generation as policy-driven actions, representing progress in RAG architectures.\n\nb) The introduction of reinforcement learning further demonstrates the flexibility of RAG and explores the potential of integrating RAG with other AI technologies.\n\nc) Evaluation across several datasets demonstrates the effectiveness of the approach.", "weaknesses": "a) There exists many works focused on the question of \"When to retrieval\" in RAG, but in Section 4.1, only one baseline was selected for comparison. Incorporating a set of baseline examples into these studies would enhance comparative and analytical insights. \n\nb) The current baselines primarily apply process-wide optimizations, such as fine-tuning the generator, without optimizing individual modules (e.g., rewriter, generator, and decision-maker) separately. This setup limits the ability to demonstrate how joint optimization in SmartRAG resolves potential conflicts and enhances coordination across modules.\n \nDetailed comments:\na) A dedicated section on the limitations and assumptions underlying SmartRAG, particularly regarding the potential latency issues in real-world applications, would be beneficial. \n\nb) While effective, the joint optimization approach may introduce complexity that could limit the model’s scalability, especially when applied to very large or real-time systems. Additional discussion on computational overhead and training time would strengthen the paper. \n\nc) Considering the potential for optimizing PPO's efficiency or exploring alternative RL algorithms to reduce training costs, it could be beneficial to further discuss these aspects, which may enhance the practicality of the framework.", "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": "x2FFPnyQ1K", "reviewer": "Reviewer_47ez", "summary": "This paper proposes a new Smart RAG framework that consists of two parts. (1)Policy Network and (2) Retriever. The policy network decides whether to retrieve or not based on the input state (question + initial observation), and generates an appropriate query to obtain the required information form the retriever. The system is trained in an end-end method using supervised fine tuning and PPO.", "questions": "1. Clarification on the training dataset. \n    1. Was all of the training data used for SFT and PPO?\n    2. On which dataset was the EM and F1 evaluations done? Was this used as a part of the SFT or PPO process? \n2. Is SFT RAG the performance of $\\pi_{\\theta^*_0}$ in the setting described in Figure 1? (Was SFT RAG evaluated in the same setting as SmartRAG? where the policy network was allowed one retrieval? )\n3. Assuming the availability of a powerful model (GPT-3.5) for generating the dataset for SFT is not always true. Have the authors tried using the base LLM to generate the SFT dataset?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The main contribution of this work is the end to end trining of the RAG framework. The authors do a good job in analyzing the proposed framework (Section 4.1-4.4) \n\n2. The paper is generally well written.", "weaknesses": "The proposed method is not generalizable. \n\n1. The proposed method is very domain specific and requires the training of the entire module for each dataset.\n2. The generalizability of the approach is further curtailed by the availability of datasets with true final answers. (Which is needed for SFT and  PPO) \n3. The end to end framework requires complete retraining if any part of the framework is changed.", "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": "ZIwvQE4D19", "reviewer": "Reviewer_a9rm", "summary": "This paper proposes an integrated framework that includes multiple components of RAG. The intuition is that jointly training these components can yield optimal performance. The experiments are conducted on four benchmarks with several baselines and two backbone models. The author also conducts ablation studies to investigate SmartRAG's behavior in terms of when to retrieve, what to retrieve, and how to answer.", "questions": "1. Why does the policy network serve three roles in the model? What's the motivation behind it? Won't these three roles conflict with each other?\n2. The optimization of the whole system is dependent on the feedback from the RL framework. How to distinguish whether a low reward is caused by a wrong policy or inaccurate retrieval?\n3. Why is the retrieval percentage of SKR in Figure 2 fixed at 0.8?\n4. In Section 4.1, the author states that *OpenBookQA, MedQA-cn, and ARC-c datasets* have nothing to retrieve from the Bing search engine. If that's the case, why does SFT RAG yield the best performance in 2 out of 3 cases with the highest retrieval ratio?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The proposed SmartRAG is clear and easy to understand. Some of the designs are intuitive.\n2. The paper writing is relatively smooth and easy to follow.", "weaknesses": "1. Part of the setups of SmartRAG are not well-justified. Please take a look at the questions.\n2. Certain experimental observations are not elaborated clearly. Please take a look at the questions.\n3. The RAG baselines seem not extensive, compared with references on this topic.", "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": "GawhwFHCvy", "reviewer": "Reviewer_FNgY", "summary": "The paper proposes a jointly optimization method called SmartRAG which combines the retriever, the query rewriter and the answer generator. The authors mainly use a policy network which is powered by a LLM to decide whether to retrieve or answer directly. They set up a quota limitation, the policy network can only answer the query when the quota is reached. A supervised warm-up with constructed dataset which contains direct answer, retrieved answer and observed answer is firstly applied to the policy network. This step ensures the policy network with the ability to output a [retrieve] or [answer] action. Then they design a reward at time step t of either answer or retrieve and use PPO to optimize the policy network. SmartRAG can decide the action and generate answers with/without retrieved observations.", "questions": "1. In equation (1) where os is the concatenation of all the historical observations returned from the retriever, how do you solve the problem of LLM token limit?\n2. Why [retrieve] action suffers from a 100% penalty in equation (5) of the reward design? If retrieval cost contributes to a negative reward, how about the correct answer generated from retrieved observations? Will that lead to an increase in the reward?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. I like the idea of this paper, it consider a real problem LLM applications where the pre-trained language models and downstream RAG applications are built and optimized separately.\n2. The paper is well-written, most ideas are clearly illustrated.\n3. SmartRAG uses reinforcement learning methods to optimize its policy network, this is reasonable since states, actions and rewards observed in previous time steps can be used as environmental feedback instead of human feedback.", "weaknesses": "1. In the methodology part, the authors should address why there is a pre-specified quota for retrieval. RAG technicals are designed to address the lack of specific domain knowledge in the training data of a pre-trained LLM. It doesn’t make sense for the LLM-driven policy network being forced to generate the [answer] action instead of [retrieve] action when the retrieve times met the quota. I’ve never seen such a setting in industrial RAG applications. Retrieval seems necessary and provides solid answers, thus shouldn’t cause penalty even the retrieval cost of similarity search in a knowledge base do exist.\n2. The optimization, however, seems not fully end-to-end. The knowledge base, which serves as a fundamental role in a RAG system, is ignored in the the design (Algorithm1 SmartRAG pipeline) and rarely mentioned in the paper. If the knowledge base is involved in the process, in methodology part, the optimization should not be tailored for specific query datasets, but remains effective even for unseen queries.", "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 (Poster)