Grounding by Trying: LLMs with Reinforcement Learning-Enhanced Retrieval
Authors: Sheryl Hsu
Omar Khattab
Chelsea Finn
Archit Sharma
Venue: ICLR 2025 Poster (2025)
Keywords: LLMs
Reinforcement Learning
Information Retrieval
Primary area: foundation or frontier models, including LLMs

Abstract: The hallucinations of large language models (LLMs) are increasingly mitigated by allowing LLMs to search for information and to ground their answers in real sources. Unfortunately, LLMs often struggle with posing the right search queries, especially when dealing with complex or otherwise indirect topics. Observing that LLMs can learn to search for relevant facts by $\textit{trying}$ different queries and learning to up-weight queries that successfully produce relevant results, we introduce $\underline{Le}$arning to $\underline{Re}$trieve by $\underline{T}$rying (LeReT), a reinforcement learning framework that explores search queries and uses preference-based optimization to improve their quality. LeReT can improve the absolute retrieval accuracy by up to 29\% and the downstream generator evaluations by 17\%. The simplicity and flexibility of LeReT allows it to be applied to arbitrary off-the-shelf retrievers and makes it a promising technique for improving general LLM pipelines.

Published as a conference paper at ICLR 2025
GROUNDING BY TRYING: LLMs WITH REINFORCE- MENT LEARNING-ENHANCEED RETRIEVAL
Sheryl Hsu', Omar Khattab'2, Chelsea Finn 1.3 & Archit Sharma 1.4
'ttanford University""attabrikk Databricks, PPhysical Intelligence, DeepMind fsherylh, wwwhiitsthsssseeeeessteteee..tee
ABSTRACT
The hallucinations of large language models (LLMs) are increasingly mitigated by allowing LLMs to search for information and to ground their answers in real sources. Unfortunately, LLMs often struggle with posing the right search queries, especially when dealing with complex or otherwise indirect topics. Observing that LLMs can learn to search for relevant facts by trying different queries and learn- ing up- up-weight queries that successfully produce relevant results, we introduce Learning to Retrieve by Trying (LcReT)), a reinforcement learning framework that explores search queries and uses preference-basee optimization to improve their quality. LeReT can improve the absolute retrieval accuracy by up to 29% and the downstream generator evaluations by 17%. The simplicity and flexibility of LeReT allows it to be applied to arbitrary off-the-shelf retrievers and makes it a promising technique for improving general LLM pipelines. Project website: http: Ihsu. com/LeReT/
INTRODUCTION
Despite tremendous progress, large language models (LLMs) still often hallucinate, motivating sig- nificant interest in grounding LLM answers in verified sources (Guu et al.) 2020 Komeili et al.) 2022 PerplexityA]] 2024 Google [024]] OpenAl| 2024). Unfortunately, simply retrieving seman- tically similar documents to the user question, as is prevalent in (RAG: [ewis et al. 2020) pipelines, tends to fail for complex information needs not answered di- rectly by any individual document. To tackle this, multi-hop retrieval pipelines gather information incrementally over multiple steps of search. For example, if a user asks What is good dinner place driving from the Bay Area to Lake Tahoe on Friday night to avoid traffic?, a grounded system might need to learn about towns en route Lake Tahoe from the Bay Area, followed by traffic forecast on I-80 and finally, restaurants in Auburn (and other towns).
generation
However, doing this successfully is hard as off-the-shelf LLM performance is often unsatisfactory, and producing supervision for the best search queries to generate in sequence of "hops" is non- trivial and expensive. Recent work tackles this via prompt optimization and rejection fine-tuning given a downstream signal. For example, Khattab al. (2023) "bootstrap" trajectories of reasoning and search queries and collect trajectories that lead to a high downstream answer accuracy, using them either to search for effective few-shot prompting examples or to finetune the LLM responsible for query generation. We observe that the problem of teaching LLM to generate effective search queries inherently a reinforcement learning (RL) problem and ask can RL improve the grounding of annwers generated by L.Ms when wielding open-ended tools like search engines?
If we can observe the retrieved documents for different search queries and compute rewards for find- ing relevant documents, we can train the LLM to produce queries that lead to better outcomes. Such learning from trial-and-error naturally lends itself to RL formalism, going beyond imitation-based methods in prior works. Indeed, we find that naïve sampling from LLMs with high temperature and using the observed data for RL is not effective. Instead, our proposed framework, Learning to Retrieve by Trying, LeReT, induces diverse search queries for each question by diversifying the few-shot examples in the prompts of the system. After this diversified sampling of search queries and the resulting retrieval, LeRe'T applies context distillation (Snell et al., 2022) followed by opti-

---

Published as a conference paper at ICLR 2025
Performance of Different Algorithms for Grounding LLMs
Retrieval
Grounded Generator
80
Few-Shot
70
LeNuT-CD (Durs)
LoReT (Ours)
65-
50
Accuracy (%)
Recall (%)
55
20
Figure LeReT significantly improves retrieval and generation. LeReT provides a reinforce- ment learning based framework for improving grounding and performance of LLM generated an- swers by improving the retrieval of relevant factual data.
mizing the guery-generating LLM using preference-based RL. We use identity policy optimization (IPO: Azar al. 2023 Rafailov al. 2024b), though other RL algorithms can be substituted.
Our main contribution is LeReT, a framework for improving grounding of LLM answers by lever- aging retrieval annotations to improve multi-hop retrieval accuracy. On question-answering datasets, LeReTT considerably outperforms prior methods like few-shot prompting and supervised fine-tuning in both retrieval quality and downstream generation quality, with stronger generators like GPT-4 benefiting more from the improvements in retrieval. We experiment with iterative version of LeReT and find that its performance improves over iterations. Our analysis reveals that prompt-driven diverse sampling is critical for LeReT to be effective, and we also analyze different ways to generate rewards for retrievals. Finally, our experiments find that LcReT can be used across retrievers, and thus, provides simple and general framework for improving retrieval. While focus on retrieval for grounding LLM answers in this work, the core method behind LeReT can be straightforwardly extended to other agentie pipelines in which LLMs control blackbox tool, so long as a reward can be formulated on its outputs.
2 RELATED WORK
Retrieval-Augmented Generation (RAG). Over the past few years, interest has been growing in
conditioning L.M outputs retrieved information (Chen et al. 2017 Lee et al. 2019) Guu al., 2020 Lewis et al. 2020 Lazaridou et al. 2022: Asai et al. (2024). This strategy seeks to make LLM systems more efficient, updatable, and transparent by decoupling the system's knowledge from the model parameters. This makes it easy to update the knowledge corpus and also makes it possible to inspect the sources relied upon when LLMS produce factual statements.
Previously, Nogueira & Cho 2017) trained query reformulator for retrieval queries using rein- forcement learning. While this the closest analogue our work, this paper does not use genera- tive language model but instead presents a very specific query selection architecture (understandably. since it is 2017) and corresponding training recipe. straightforward application of RI. similar to this applied modern LLMs would result in performance similar to few-shot prompting.
Multi-Hop Retrieval. The standard RAG formulation is best suited for "simple" questions, where a direct search can find all the information required for producing responses. Beyond these, bench- marks such as HotPotQA (Yang et al. 2018), MuSiQue (Trivedi et al. 2022), and HoVer (Jiang gathering and synthesizing information from several independent
et al., 2020) assess systems documents within a massive corpus like Wikipedia. To tackle retrieval in this setting, early systems like MDR (Xiong et al.) and Baleen (Khattab et al. 2021) introduced bespoke strategies for fine-
2

---

Published as a conference paper at ICLR 2025
Multi-Hop Retrieval
Hop
Hop
Hop
LLM (Generator)
User
Question
Answer
documents
documents
documents
An overview of the standard multi-hop hop retrieval pipeline we study in this work. A user asks question to the system. In each hop, the LLM generates search queries for the retriever and receives collection of documents. The overall set of retrieved documents and the user question are then given to a downstream LLM for grounded answer generation.
Figure 2:
tuning the retrieval models that produce representations of queries and documents, adapting them directly for compositional search queries. Unfortumately. fine-tuning the retriever representations is data hungry approach that is also challenging to scale massive corpora like the Web, re-training the retriever often requires re-indexing the corpus. Increasingly, research in this space (Trivedi et al,,
2023 Yao al. 20233 Khattab al. 2022 Press al.) 2023) relies off-the-shelf retrievers such as the Wikipedia API and focuses on improving LLMs ability to to generate effective queries.
Optimizing LLM Programs & Agents. Recent work tackles this by employing prompt optimiza- tion and rejection fine-tuming fijettion,inee-tuiing using downstream signal. For example, in Khattab et al. (2023), the authors "bootstrap" trajectories of reasoning and search queries and use the trajectories that lead to high downstream answer accuracy as candidate examples. The trajectories are then used either to search for effective few-shot prompting examples or finetune the LLM responsible for query gen- eration. This approach was extended in OOsahl-Ong et al. (2024) and [Soylu et al. (2024) in which the authors also explore using these trajectories to search for free-form instructions for prompting or to nest forms of rejection fine-tuning and prompt optimization, respectively.
Additionally, prompt-based techniques have been developed to further improve retrieval and down- stream generations. For example, Query2doc (Wang al.) (023) prompts the LLM for hypothetical documents to concatenate with the query. Rethinking with Retrieval (He et al. 2022) uses de- composed reasoning steps queries. On the downstream generation side, there have been methods (Yu et al.) 2024 Lan et al.) 2023] Kim et al.) 2023) that enable the generator to better ref- erence the retrieved information and produce better answers. As general RI. framework, work complementary to these prior works.
Beyond retrieval LLM programs, similar techniques can be used for optimizing agent behavior. For example, in [Song et al. (2024), the authors train agents to navigate the web, simulate science experiments, perform household tasks by collecting failure and success trajectories and training the LLM using preference optimization. Our work finds that the quality of exploration data sampled from LLMs is critical the success of RL in agentic pipelines, and introduces a prompt- diver- sification strategy that samples diverse and high high-quality exploration data (discussed in Section 4.] analysis in Section 5.4).
3 PRELIMINARIES
Multi-Hop Retrieval Setup. An overview retrieval shown in Figure 2 We assume access to retriever that maps search query g to set of N most similar documents D [d. JN J where d, denotes an individual document. A user asks a question u, a LLM म, generates a query qı for the retriever, which results in ordered set of document Dj. In the next hop, the LLM म, takes and D, as input and outputs query q2. This repeats for H hops. The final ordered set of retrieved documents DR D, U D2 J... DH given as input the LLM generator ny along with the
3

---

chRandomSeea rch (BFRS), to self-generate a number of in- dependently optimized few-shot, chain-of-thoughh prompts P {PI.....PP } for LLM π,. The independently optimized prompts would naturally lead to diverse samples from T,, and DSPy's op- reuse Published as a conference paper at ICLR 2025
question u, which generates the final answer for the query. In this work, restrict ourselves to fine-tuning the LLM ny and treat the retriever and LLM generator Tg as blackbox models.
Language Models and Reinforcement Learning. Reinforcement learning has become the de facto tool for aligning large language models and has inspired considerable work at the intersection of language models and RL (Stiennon et al., 2020 Ouyang al. [022 Zhho et al.
2023 Rafailov et al. 2024b). We briefly review direct alignment methods (Rafailov et al. 2024b). Given a dataset
of preferences D,- - fx', vlw V/ where x denotes the dialogue history, V., denotes the preferred response, and of denotes the dispreferred response. Bradley- -erry (Bradley & Terry 1952) offers model that connects choice to an implicit goodness score, useful for learning a reward model ria
LBT (. loggPo(sa
wi) ((, [log (ro(2, yw) ro m))], (1)
where a denotes the sigmoid function p y) denotes the probability of yu being preferred over Typically a L.LM æ is trained to maximize this learned reward model using RI. as described by max, (x-+(()))) (ro (z. y) 3KL( metr)] where Tee denotes a fixed reference policy. How- ever, Rafailov et al. (2024b) shows that parameterizing r v) = / log (TT( <y) x²) and
optimizing Eq [ implicity optimizes the RLHF objective exactly, removing the need for a sepa-
rately parameterized reward model an explicit RL training loop. However, |Azar al. (2023) and
Rafailoy et al. (2024a) have found that optimizing a DPO parameterized reward uing Eqq can lead to overoptimization, and suggest optimizing the following objective:
LPo (z. Mi >r
(F vw) - (z(2, 91 0.57 )"
(2)
where 7 is a hyperparameter controlling the target margin between the implicit rewards for You and yr, and (. y) log (Ta(y x)/ Ter (y ()). Minimizing the objective in Eq 2 leads to identity
policy optimization (IPO: Azar et al. 2023).
4 LERET: LEARNING TO RETRIEVE BY TRYING
We introduce Learning to Retrieve by Trying, or LeReT, a novel framework for improving the grounding of LLM generations by training the search query LLM म, using preference optimization. In hop i, we sample a query qi from the LLM based on the user question and documents seen in hops < i, and observe reward signal for the retrieval quality. Both sampling from the LLM and retrieval make it hard to backpropagate directly from the reward signal, making RL a more suitable optimization framework for this setting. We first discuss how to generate a dataset of queries and retrieved documents that is suitable for RL optimization in Section 4.1 We then discuss how to convert the reward-annotated dataset into dataset of preferred and dispreferred queries and use that dataset to optimize ,, with IPO in Section 4.2) We also briefly discuss an iterative version of LeReTT that alternates between sampling and optimization. Finally, we combine the elements and give a
practical overview in Section 4.3
4.1 PROMPT DRIVEN DIVERSE QUERY GENERATION
Given a dataset of questions, we want to "try" a set of search queries and observe the retrieved documents. What queries would be good to observe the retrieved documents for? This roughly cor- responds to the exploration problem in RL. As our experiments later in Section 5.4 also indicate, a good distribution of queries would result in diverse outcomes (for better exploration), but is impor- tant that some queries produce high quality retrievals. To sample such diverse and effective queries, LeReT moves beyond high- temperatur sampling and uses a diverse set of examples to few-shot prompt the LLM म,. We use DSPy (Khattab et al., 2022: 2023 's prompt optimizers, specifically a simple Boot rentssstaaeeSShhhtee
timization ensures that the prompts are still resulting in relevant retrievals. Note that the same set of prompts across all questions throughout the dataset.
For hop h, LeReT does the following: For every question u E U, LeReT samples search queries conditioned on each of the prompts p E P. resulting in a set of search queries Qn = ( Pi, u, Ch-1)) p. E P). where Ch-1 denotes the context from previous hops. For
4

---

Published as a conference paper at ICLR 2025
Prompt Sampling
For given quedioo context:
Question
LUM
Guary
Contect
Veward
Context
Sample next hop context
Quenn
Qpe
proportionally to reward
Figure 3: Overview of prompt driven diverse sampling and data generation. LeReT induces diverse but effective search queries by bootstrapping several few-shot prompts for query generation and uses the retrieval reward to collect preferred and dispreferred queries for each question's hop.
every query qi E Qh. we retrieve set of documents denoted Cιi and compute the reward corresponding each query by evaluating the retrieved documents, that Trex(2, Ch-1. q:) (RCh))) (more on retrieval reward R in Section (.3). The final dataset for this hop consists of Dr = {(u, C,-1, r(u, Ca-1, u U. E EQ), where every entry consists of the user ques- tion, the context from the previous hop, the sampled query, and the reward computed by running the retriever. Though the prompts we generate are used to sample diverse high quality search queries at training time, we leverage context distillation to remove the need for optimized prompting at test-time.
When training in multi-hop setting like this, must select which contexts to use for the next hop- Naïvely generating queries for every possible context leads to an exponentially growing dataset with respect to the number of hops, which becomes computationally infeasible quickly. At the end of hop h, we have P different contexts Chi for a given question. LeReT randomly selects one of the P contexts to for the next hop, creating Ch. To do this, first filter out contexts that have achieved the optimal reward (assumed to be 1), as no more relevant documents can be retrieved. For the remaining contexts, weigh each context by the reward and sample one of them. This biases the data towards trajectories that achieve higher reward, while still containing trajectories where the model can recover from poor retrieval earlier hops. The final training dataset is simply the union of the dataset from each hop, that is Drr = D1 U.. The pseudocode for dataset sampling is given Algorithm
4.2 MODEL OPTIMIZATION
Given the training dataset Drr, we want to update the LLM ,.. First, we make a simplification by optimizing every hop greedily, that is, based on the reward obtained in that hop alone. Ideally, updating тr on data hop accounts for rewards obtained all future hops. Retrieving relevant documents in earlier hops is likely always better, and intuitively, retrieving irrelevant documents in earlier hops will rarely lead to better overall retrieval. We verify this assumption empirically: For two sets of retrieved documents, low reward and high reward set, low reward retrieval leads to a higher total reward in only 0.026% of the cases (Appendix B.I.. Thus, the greedy approach considerably simplifies the optimization, without any evident theoretical empirical sacrifice.
Given the dataset Dr. we can optimize म,. using any RL algorithm. In this work, we consider preference-based RL approaches, specifically IPO described in Section 3 because of its simplicity and effectiveness. To optimize using IPO, we need to transform Dr into preference dataset. This can be done straightforwardly by comparing two search queries qi and gj for the same user question and context and choosing the query with higher reward as the preferred response and the other query as the dispreferred response. However, before we can optimize म,, using IPO, we must account for the fact that the search queries were sampled using few-shot prompting, and at test-time we will not using the prompt. To do so, we leverage context distillation (Snell et al., (022) by fine-tuning
the search queries without the context. We observe in experiments that this roughly matches the performance of few-shot prompting. After context distillation and converting the training dataset into preference dataset, optimize n, using IPO.
5

---

Published as a conference paper at ICLR 2025
Iterative-LeReT. Thus far, we have assumed that sampling search queries and model training done as two separate steps. However, we can alternate between the two steps, leveraging the im- provements in previous iterations to sample better data for the next iterations. Iterative-LRReT closely follows iterative-DPO (Xu et al. 2024). We partition the dataset of user questions 2I and run LeRe'T each partition, sampling from the model fine-tuned ically, we have data partitions uI,..., .,,, We start with the LLM To and apply LeReTT on UI to obtain the fine-tuned LLM Ti so we have LeReT (To, + ITI We use LeReT U2, using क to sample the training data and continue fine-tuning from, that is, LeRe'T' (TI, this for all partitions until we get the final model æy. We find that iterative sampling and training can be effective as models may not achieve accurate and relevant retrievals in the initial iterations, and later models may be able to generate better exploration data.
the previous partitions. Specif-
We can repeat
4.3 REWARD LABELING FOR RETRIEVED DOCUMENTS
In order to construct the required preference datasets, we need reward signal R to score the doc- uments retrieved by search query. How do we compute this reward signal? There are broadly two ways to get such supervision: direct supervision, where human provides oracle documents to ground the answers in the training dataset or explicitly reviews the relevance of documents retrieved by the search query and indirect supervision, where the supervision comes from evaluations of the downstream generator such as preference feedback on the final answers or some answer verification. do not obtain any explicit supervision for retrieval, but only receive
The latter is indirect because information about how the generator performed after being conditioned on the retrieved documents. We run short study in Section B.5 comparing the two forms of supervision, and find that direct supervision results in better performing models. However, a full study comparing direct and indirect forms of supervision and their trade-off's is beyond the scope of this paper, and potentially requires novel algorithmic considerations. For the majority of the paper, we assume some form of direct supervision, as allowed by commonly used datasets and benchmarks.
Algorithm Prompt Driven Diverse Sampling Training
I: Input: Number hops H, Number of few-shot prompts P, LLM TT, Retriever X., Datasct M 2: Initialize: C1 = 2; [pi,...,PP .P) as few-shot prompts; Dpet = Ø
3: for h in range(1, H) do
4: for in do
for i range(0, P) do
6:
Sample qi ~ σ-(( Retrieve Cl. + Ch-1 U (((g))
u, Ch-1)
7:
8:
Compute reward end for
R(Ch))
10:
for range(0, P) do
11:
for j in range(i + 1, P) do
If * ry: Add (u, Cn-1, 2)) to preference dataset pret end for for 15: Sample context for next hop On ~ Sample x vi)
12:
13:
16: end for
17:
end for
18:
WLaREFED SFT(Tr, Dpeet)
19: MLRET = IPO (FLLRRED. Dreat)
5 EXPERIMENTAL EVALUATION
We now evaluate how LeReTT impacts the quality of retrieval and of downstream generation. We first test LeReT on two multi-hop question answering datasets, finding that LeReT significantly outperforms baselines such as few-shot prompting and supervised fine-tuning. We also find that applying LeReT iteratively leads to further improvement over iterations. We analyze prompt driven diverse sampling in contrast with sampling using high temperature and also discuss different reward functions for the retrieval step. Finally, we evaluate LeReT's adaptability for various pipelines by testing it against retrievers.

---

Published as a conference paper at ICLR 2025
Model Dataset
1 Hop
Hops
3 Hops
Hops
Generator
Method
RE
RE
AP
RE
FM
Base
42.3
38.8
54.7
41.9
41.0
Few-shot
49.9
45.6
64.8
53.9
47.1
Llama 8b HotporQA
Few-shot all
50.2
46.4
63.5
50.5
45.2
Query2Doe
47.7
43.7
61.4
44.8
LeReT-CD
51.4
47.3
69.8
58.0
49.3
LeReT
56.7
52.5
77.1
66.3
52.5
Base
52.2
48.4
70.9
57.7
51.0
Few-shot
54.4
50.4
66.7
57.8
48.5
Gemma 9b
Few-shot all
54.6
50.5
69.6
58.8
50.0
HoipotQA
Query2Doc
43.5
40.0
51.2
43.9
35.9
LeReT-CD
53.5
49.6
71.9
59.2
51.4
LeReT
67.0
54.3
Base
37.9
34.8
45.6
37.9
48.5
38.3
50.0
39.3
61.5
Few-shot
45.6
42.2
53.4
45.9
56.0
46.0
57.3
46.1
64.6
Llama 8b
Few-shot all
38.8
35.8
51.9
44.4
57.5
46.3
59.7
45.9
64.7
Ho Ver
Query2Doe
39.7
36.4
57.1
45.7
64.0
LeReT-CD
42.9
39.8
56.6
48.4
63.2
52.2
66.9
54.3
67.5
LcReT
45.8
42.5
65.4
56.1
72.8
61.4
76.9
64.3
69.8
Base
40.8
37.7
45.5
38.1
48.8
39.6
50.1
40.4
61.7
Few-shot
46.3
42.9
57.9
59.3
48.5
64.3
Gemma 9b HoVer
Few-shot all
46.3
42.8
55.8
48.7
64.1
52.6
68.2
54.1
67.5
Qucry2Doc
40.5
37.6
43.3
39.1
45.5
40.2
46.7
40.7
60.3
LeReT-CD
45.2
41.7
59.5
50.7
65.3
54.1
69.0
56.2
67.2
LcReT
47.0
43.7
67.5
57.6
75.2
63.1
79.4
66.1
71.5
Table 1: LeReT improves the performance of Llama 3 8b and Gemma 9b on HotpotQA and HoVer. We compare models trained with LeReT versus the base model, the base model with few- shot prompting, and the base model with Query2Doc. We measure the recall (RE) and average precision (AP) of the retrieved documents (higher is better) along with the exact match of generations produced using the retrieved documents.
Datasets. We test LeReT on HotpotQA Yang et al. 2018) and HoVer (Jiang et al. 2020). Both datasets are based on a Wikipedia knowledge base and are multi-hop, meaning that models must reason across multiple articles arrive at the correct answer. The datasets provide both the correct answer supporting articles. HotpotQA question-ans answering dataset that requires up to 2 hops, while Ho Ver is fact verification dataset that requires up to hops. Both datasets are relatively large, allowing us to train on over 10,000 questions.
Evaluation metrics. We measure retrieval performance using recall and average precision. Recall the number of correctly retrieved documents over the total number of correct documents. Average precision (Eq. in the appendix) takes into account the ordering of the documents, i.e. if the correct 3 documents are ranked as the last 3 out of the score will be lower. For generation, we measure both exact match the entire answer and F1 at the word level.
Baselines. For our baselines, we compare against using the base (general-purposs instruction- tuned) model to generate queries and also prompting the base model using bootstrapped few shot prompts optimized by DSPy. For the main few shot prompting baseline, we use the few shot prompts used during prompt driven diverse sampling. These prompts are created by optimizing the first hop with DSPy and then using that prompt for all hops. To demonstrate that the gains of LeReT are additive on top of these, we report the maximum achieved few-shot prompt performance achieved by any bootstrapped fewshot prompt pI,..., Pp.. For some experiments, we also report few-shot all, which where we use DSPy to optimize over the entire pipeline and bootstrap different examples for each hop. We also Que (Wang et al.) 2023), prompting technique that asks the LLM to generate hypothetical document in addition
the query, as a baseline. We also report LeReT-CD as baseline for some experiments. This is the performance of the model after the SFT step (but before IPO) and as explained in Section 4.2 is the same as context distillation.
Experiment setup. Unless otherwise specified, use Llama 3 8b Instruct Gemma 2 9b it the base model for query generation. We use ColBERTv2 (Santhanam et al. 2022) as the retriever. For the reward function, use the average precision of the retrievals, so R AP(Cn:).

---

Published as a conference paper at ICLR 2025
1 Hop
2 Hops
3 Hops
4 Hops
Dataset
Model
Method
RF
AP
RE
AP
Standard
56.1
52.2
79.9
67.0
HotpotQA
Gemma 9b
Iteration
55.7
51.7
78.2
63.6
Iteration
57.6
53.5
82.3
70.5
Standard
45.8
42.5
65.4
56.1
72.8
61.4
76.9
64.3
HoVer
Llama 8b
Iteration
45.1
41.7
62.7
54.1
69.6
59.1
73.5
62.1
Iteration 2
44.9
41.6
65.3
55.5
73.4
61.2
78.2
64.4
Table 2: Iteratively applying LeReT leads to performance gains compared to standard LeReTT
Gemma 9b and Ilama 8b are tested with two iterations and recall and average precision are measured (higher is better).
5.1 RESULTS ON HOTPOTQA & HOVER
In terms retrieval recall, LeReT improves recall by 9-22% on HotPotQA and 27-29% on HoVer relative to the Llama and Gemma unadapted instruct models ("base"). This substantially exceeds the gains achieved via few-shot prompting alone, showing that sampling from multiple few shot prompt ensembles and training the model with RL is crucial. The gains also compounds because lower quality search queries at a given hop distract future steps. We feed the improved retrievals into Llama 3.1 70b, asking it to generate response the question using the provided context. We find that improving retrieval produces corresponding improvement in generations, with the generator exact match increasing at approximately half the rate of recall.
hops, possibly
5.2 EHERATIVE-EERE
We evaluate the performance of applying LeReT for two iterations. Training with only half the data (iteration 1) results in slightly worse performance compared to standard non-iterative LeReT, but after the second iteration the model performs better than the LeReTT model. That is, sampling data that is both more on-policy and higher scoring in the second iteration leads to improvement.
5.3 FACTUALITY WITH DIFFERENT GENERATORS
Base EM
(RE 54.15)
Few-shot (RE 63.60)
LeReT (RE 80.40)
Generator Model
FI
EM
FI
EM
FI
Gemma
13.8
21.6
16.6
24.5
18.9
(+5.1)
28.5
Llama 3 8b
35.1
44.6
40.4
50.2
46.9
(+11.8)
58.6
Llama 3.1 70b
38.1
47.7
45.6
56.3
53.5
(+15.4)
64.9
GPT4
33.4
43.3
41.6
52.7
50.7
(+17.3)
62.9
Table 3: The stronger the generator model, the more it benefits from improved retrieval. We test different generator models using retrievals sampled from HotpotQA on the base model, few- shot prompted model, and LeReT-trained model. We report the recall of the retrievals (RE). We measure the exact match and Fl scores of the generated answers (higher better).
Improving retrieval seeks improve LLM grounding. Intuitively, stronger models with better rea- soning capabilities should benefit more from having the correct documents to reason with than weaker models that may not be able to generate the correct answer even with the right documents. To assess this, we take the retrieval output by various Llama 3 8b models for HotpotQA and con- dition various generator models with them. As seen Table 3 stronger generators deliver higher quality and larger gains when supplied with LeReT-traines retrieval contexts. We note that although GPT4 has the largest improvement, it does not have the highest score. Examining the generations, see that GPT4 very closely followed the instructions to "base your answers only the provided context" and would output statements such as "answer cannot be found the context" instead of trying to answer it anyway the way weaker models did.
8

---

Published as a conference paper at ICLR 2025
Data size
Gold (%)
Unique AP
AP Std Dev
Model
Hotpot
HoVer
Hotpot
HoVer
Hotpot
HoVer
Hotpot
HoVer
@ temp 0.3
44,705
28.575
37.5
12.1
1.66
1.25
0.09
0.029
@ temp 0.5
55,564
32,091
39.4
12.6
2.10
1.27
0.10
0.033
@ temp 0.7
69,361
37,354
43.9
12.6
2.10
1.28
0.14
0.033
@ temp 1.35
92,264
48,516
46.3
12.0
2.35
1.32
0.16
0.039
@ temp 2.0
81,017
58,281
41.0
10.3
2.36
1.34
0.16
0.045
Fixed few-shot @ temp 2.0
93,613
92,542
47.7
13.5
2,41
1.32
0.16
0.038
Diverse few-shot @ temp 2.0
95.373
71,433
47.5
12.3
2.47
1.34
0.17
0.044
Diverse few-shot temp 0.7
105,506
49,304
54.2
13.8
2.35
1.28
0.15
0.031
Table 4: Sampling with higher temperature results in greater diversity of responses (higher unique ap, ap std dev) while few shot prompting results in better data (higher gold rate). Specifically, gold rate is defined as the percentage of questions for which we have gold star sponse (a query that results in the maximal score). Unique AP is the number of unique average precision scores there are for question, and AP std dev the standard deviation of the average precision scores for question. Data size is measured in terms of the number of preference pairs.
5.4 DIVERSE FEW-SHOT PROMPTING VERSUS HIGH- TEMPERATURE SAMPLING
While prior work typically uses high temperature sampling to generate diversity, LeReTT leverages few-shot prompting to generate diverse exploration data. To evaluate the effectiveness of few-shot prompt diversification, we consider alternate sampling strategies, particularly sampling from n,, at different temperatures with no few-shot prompting. a fixed few-shot prompt (fixed), or multiple few shot prompts (diverse).
First, we compute some statistics about the rewards generated by the sampled search queries; we compute the average number of unique rewards per question and the standard deviation of the re- wards as a proxy for diversity, and we measure the percentage of questions where at least one query (gold star answer) achieves maximal reward Table 4 that while sampling with higher temperature improves diversity in search queries, few-shot prompting leads to significantly higher quality data and using multiple few-shot prompts provides comparable diversity.
proxy for quality of sampled data. We find in
We train on four of the sampled datasets: (1) queries sampled with diverse few-shot prompting at standard temperature (0.7), (2) queries sampled at high temperature (2.0), (3) queries sampled with diverse few-shot prompting high temperature (2.0), and (4) queries sampled with single fixed few-shot prompt at high temperature (2.0). We find that training a model using data sampled with few-shot prompting at temperature 0.7 results in the best performance, which also the sampling strategy that results in the largest percentage of questions with at least one gold star answer. This suggests that exploration critical to the success of RL training and justifies the extra effort of bootstrapping few shot prompts.
Hop
2 Hops
Model
RE
AP
RE
AP
Base model
42.3
38.8
54.7
41.9
Few-shot
49.87
45.58
64.77
53.86
LeReTT diverse few-shot @ temp 0.7
55.74
51.76
78.14
67.67
LeReT @ temp 2.0
49.12
44.81
69.86
58.56
LeReTT fixed few-shot @ temp 2.0
50.89
46.69
70.79
59.93
LeReT diverse few-shot @ temp 2.0
51.73
47.79
73.67
62.60
Table 5: Diversifying few-shot prompts when sampling search queries result in more effective training datasets for RL. The standard LeReT with few-shot prompting at temperature 0.7 results in better performance than training on dataset sampled temperature 2.0 without few-shot prompting or a dataset sampled at at temperature of 2.0 with fixed few-shot prompting or a dataset sampled at temperature of 2.0 with diverse sho- prompting.

---

Published as a conference paper at ICLR 2025
Hop
Hops
Hops
Hops
Dataset
Method
RE
AP
RE
AP
RE
AP
RE
AP
Llama 8b
10.6
9.1
18.5
15.0
Hotpot
Few-shot
39.6
34.7
50.7
40.8
LeReTT
43.8
38.9
60.0
48.6
Llama
15.6
14.0
22.8
19.2
27.8
22.0
31.2
23.5
Ho
Few-shot
37.1
33.2
45.9
38.3
49.8
40.1
52.1
39.6
LeReT
39.1
35.2
51.9
44.5
58.6
48.8
62.6
51.3
Table 6: LeReT greatly improves the performance Llama 8b on Hotpot and HoVer with
Azure Al Search used as the retriever. We perform the same sampling and training pipeline as all other experiments use Azure AI Search instead of ColBERT.
5.5 DIFFERENT RETRIEVERS
Finally, we test whether LeReT applicable to general RAG systems by swapping our retriever from ColBERT over Wikipedia Azure AI Search, applied with the default configuration for full text search. We observe that the base Llama model performs very poorly compared to its retrievals with ColBERT. This is likely because Azure is not specialized to the Wikipedia index which is helpful for our multi-hop tasks. However, the query generating LLM can adapt to compensate for this weaker retriever,
significant improvement with few-shot prompting and LeReT. This demonstrates the power of LeReTT to adapt to general blackbox tools in the pipeline. Given the poor performance of the base model, few-shot prompting based exploration (Section 5.4 found to be necessary, versus simply sampling with high temperature.
DISCUSSION
In this work, we introduced LeReT, framework for improving the quality of multi-hop retrieval via reinforcement learning and thus enabling better grounding for LLM systems. Beyond retrieval specifically, this can be extended to learning for agentic systems or LLM programs that use other tools. LeReT conducts diversified exploration varied optimized few-shot prompts. It then uses this to construct preference dataset for every hop consisting of queries that lead to a diverse set of retrieved outcomes. conducts context distillation followed by an iterative application of the IPO objective. Experimental evaluation on the HotpotQA and Ho benchmarks with two different retrieval models reveals that LeReT can improve the quality of Llama 8b- and Gemma 9b-based systems by up to 29% in recall.
search queries in each hop by sampling using
train the model, first
Limitations & Future Work. While in this work we have used direct supervision for retrieval, a fruitful effort would be to enable learning from indirect supervision such as the correctness of the final generative response. Another promising direction is learning by updating the tools themselves like training the retriever model used encode the search queries and documents. Doing this would require changes to the sampling algorithm and addressing the signal-to-noie ratio but would likely lead to significant gains.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-RAG: Learning to retrieve, generate, and critique through self-reflection. The Twelfth International Confer- ence Learning Representations, 2024. httph:ts:://// http:::://penreeeeeeeeteeeeeeeee hSyW5go0v8
Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. A general theoretical paradigm to understand learning from human preferences, 2023. URL httph:ttps:aaaat org/
wttp///ggab.aaeee abaaabssb/// 12036
Ralph Allan Bradley and Milton E Terry. Rank analysis of incomplete block designs: I. the method of paired comparisons. Biometrika, 39(3/4):324-444 1952.
Danqi Chen, Adam Fisch, Jason Weston, and Antoine Bordes. Reading wikipedia to answer open- domain questions. In Proceedings of the 55th Anmual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers), pp. 1870-1879, 2017.
Google. Generative ai is coming to google search. http:://bloo..oogllpprrttttttt gearch/generativv aiiee-ai--- ai- ai-google seoogle-seacch search mmay-20241 2024.
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.
Hangfeng He, Hongming Zhang, and Dan Roth. Rethinking with retrieval: Faithful large language model inference, 2022. URL /eatts://arxtt wwtp///////.ooo/. httpss//aaxxxx..rrr//////03000 2w0.2331100030333 000003
Yichen Jiang, Shikha Bordia, Zheng Zhong, Charles Dognin, Maneesh Singh, and Mohit Bansal. HoVer: A dataset for many-hop fact extraction and claim verification. In Trevor Cohn, Yulan He, and Yang Liu (eds.). Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 3441-3460, Online, November 2020. Association for Computational Linguistics. doi: 10. 10.18833/11/2000 http:///020.nnnnnn-eee..p..... E findings-ent 1. 309.
URL. httts:://acl/ aaaclannhelogy
Omar Khattab, Christopher Potts, and Matei Zaharia. Baleen: Robust multi-hop reasoning at scale via condensed retrieval. Advances in Neural Information Processing Systems, 34:27670-27682,
2021.
Omar Khattab, Keshav Santhanam, Xiang Lisa Li, David Hall, Percy Liang, Christopher Potts, and Matei Zaharia. Demonsstaate-eaarh--rredd Composing retrieval and language models for knowledge-intensive NLP. arXiv preprint urXiv:2212. 14024, 2022.
Omar Khattab, Arnav Singhvi, Paridhi Maheshwari, Zhiyuan Zhang, Keshav Santhanam, Sri Vard- hamanan, Saiful Haq, Ashutosh Sharma, Thomas T. Joshi, Hanna Moazam, Heather Miller, Matei Zaharia, and Christopher Potts. Dspy: Compiling declarative language model calls into self- improving pipelines. arXiv preprint arXiv:2310.03714, 2023.
Gangwoo Kim, Sungdong Kim, Byeongguk Jeon, Joonsuk Park, and Jaewoo Kang. Tree of clarifi cations: Answering ambiguous questions with retrieval-augmentee large language models, 2023.
URL htttss:/ org/ http://abb/22.....0
Mojtaba Komeili, Kurt Shuster, and Jason Weston. Internet-augmented dialogue generation. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (eds.), Proceedings of the 60th An- nual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). pp. 1460-84788 Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. URL https://aaclaatttooooo.oo///////
1865
acl long.
579
Tian Lan, Deng Cai, Yan Wang, Heyan Huang, and Xian-Ling Mao. Copy is all you need, 2023.
//tpxiv////vvvv.// org 2307.069622.00 06962
11

---

Published as a conference paper at ICLR 2025
Angeliki Lazaridou, Elena Gribovskaya, Wojciech Stokowiec, and Nikolai Grigorev. Internet- augmented language models through few-shot prompting for open-domain question answering. arXiv preprint arXhv:2233.05115, 2022.
Kenton Lee, Ming- Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 6086-6096, 2019.
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 mnowledge-intensive nip tasks. Advances in Neural Information Processing Systems, 33: 9474, 2020.
Rodrigo Nogueira and Kyunghyun Cho. Task-oriented query reformulation with reinforcement learning. In Martha Palmer, Rebecca Hwa, and Sebastian Riedel (eds.), Proceedings of the 2017 Conference Empirical Methods in Natural Language Processing, pp. 574-583, Copen- hagen, Denmark, September 2017. Association for Computational Linguistics. doi: 18.18653/v// D17-1061. URL wtt.ttss//aallaattooooooooo///o/
OpenAl. Searchgpt is a prototype of new ai search features. searchept-protttyyyes 2024.
Krista Opsahl-Ong, Michael Ryan, Josh Purtell, David Broman, Christopher Potts, Matei Zaharia, and Omar Khattab. Optimizing instructions and demonstrations for multi-stage language model programs, 2024.
11695
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730-27744, 2022.
PerplexityAL, Perplexity mai. httpss//www.prrppleetttt
2024.
Ofir Press, Muru Zhang, Sewon Min, Ludwig Schmidt, Noah A. Smith, and Mike Lewis. Measuring and narrowing the compositionality gap in language models, 2023. URL https://arxiv.
org brg/aba/2210 03350
Rafael Rafailov, Yaswanth Chittepu, Ryan Park, Harshit Sikchi, Joey Hejna, Bradley Knox, Chelsea Finn, and Scott Niekum. Scaling laws for reward model overoptimization in direct alignment algorithms. arXiv preprint urXiv:2406.02900 2024a.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Advances in Neural Information Processing Systems, 36, 2024b.
Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Col- bertv2: Effective and efficient retrieval via lightweight late interaction, 2022. URL https:
01488
Charlie Snell, Dan Klein, and Ruiqi Zhong. Learning by distilling context, 2022. URL https: warxiv.org/abs/2222
15189
Yifan Song, Da Yin, Xiang Yue, Jie Huang, Sujian Li, and Bill Yuchen Lin. Trial and emror Exploration-based trajectory optimization for Ilm agents, 2024. WRRlhhtts:///axiv.... abs/ abs/2403 02502
Dilara Soylu, Christopher Potts, and Omar Khattab. Fine-tuning and optmiizzaiooonnn Two great wrw. http:/bb//2400 10930
steps that work better together, 2024. URI.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in Neural Information Processing Systems, 33:3008-3021, 2020.
12

---

Published as a conference paper at ICLR 2025
Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. Musique Multi- hop questions via single-hop question composition, 2022. URL htttps://axxii.. org/ abs/
Harsh Trivedi, Niranjan Balasubramanian, Tushar Khot, and Ashish Sabharwal. Interleaving re- trieval with chain-of-thought reasoning for mhowledge-intensive multi-step questions, 2023, URL
https://arxiv. /ttp:///axxvvvvv wwt.//rggaaaa...oo 2212.10509
Liang Wang. Nan Yang, and Furu Wei. Query2doc: Query expansion with large language mod- els. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 9414-9423, Singapore, December hete:3.emmllmmmmm main.585. URL
2023. Association for Computational Linguisties. doi: 10. 10.18653/112023.commm attpss//accaantioooyy.ooo www.:///oooo.....2.....o2 wmnlp- mammalpmmain.... main.
585
Wenhan Xiong, Xiang Li, Srini Iyer, Jingfei Du, Patrick Lewis, William Yang Wang, Yashar Mehdad, Scott Yih, Sebastian Riedel, Douwe Kiela, et al. Answering complex open-domain questions with multi-hop dense retrieval. In International Conference on Learning Representa- tions.
Jing Xu, Andrew Lee, Sainbayar Sukhbaatar, and Jason Weston. Some things are more cringe than others: Iterative preference optimization with the pairwise cringe loss, 2024. URL https:
/
Zhilin Yang. Peng Qi. Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. HotpotQA: A dataset for diverse, explainable multi-hop question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and Jun ichi Tsujii (eds.), Proceed- ings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp- 2369- 2380, Brussels, Belgium. doi: 10.18653/v1/D18.
2018. Association for Computational Linguistics. URL /RL/pttpss//aaaattoooooo....
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models, 2023. URL httptt/ss//:xxii
org/ 2015/abs/200 2abb/2210.00 23622
Wenhao Yu, Hongming Zhang, Xiaoman Pan, Kaixin Ma, Hongwei Wang, and Dong Yu. Chain- of-note Enhancing robustness in retrieval-augmeeted language models, 2024. URL https:
wwx.iv.rrg/abs//00
Yao Zhao, Rishabh Joshi, Tianqi Liu, Misha Khalman, Mohammad Saleh, and Peter J Liu. Slic-hf: Sequence likelihood calibration with human feedback. arXiv preprint arXiv:2305. 10425, 2023.
13

---

Published as a conference paper at ICLR 2025
A SAMPLING AND TRAINING DETAILS
We sample with P= 4 for HotpotQA and P = 3 for HoVer. We implement our sampling pipeline on top of DSPy (Khattab et al., 2023), specifically defining a single hop as a program and sampling data using the evaluate functions. We also use chain-of-thought prompting when generating queries.
We use learning rate of 1 for SFT/context distillation all our experiments, and use T 0.05 and learning rate of 1.-7. We train SFT for I epoch, and only distill the best performing prompt. We train IPO for epochs.
A.1 DATA SCALING ANALYSIS
We conduct data scaling experiments for L.ReT. We evaluated a training run of Llama 3 8b on the full HotpotQA training set (90,447 questions), which resulted 494,208 preference pairs after prompt driven diverse sampling. We find that the majority of the improvement occurs relatively quickly. Based on this, we only use a quarter of the HotpotQA training set for subsequent experi- ments. However, data scaling likely depends on host of factors, including the task complexity and the base model, and we conducted the data scaling experiment to reduce the computational cost of experiments.
Performance Over Training Examples
2 Hop RE
2 Hopp AP
0.75
Hop RE
Hop
0.70
Percentage
0.65
0.60
0.55
0.50
20052
s9p0e
09930
119900
209000
239616
260588
200220
229472
359426
199300
119920
Figure 4: The model performance saturates quickly. Mcasuring the test performance of Llama
3 8b as training progresses on the preference dataset collected using LeReTT on the full HotpotQA train set (90,447 HotpotQA questions, 494,208 preference pairs).
For the retriever, we set up local instance of ColBERTv2 with Wiki 2017 abstracts index. For Azure AI Search, ereated a custom index by uploading all the abstracts from the Wiki 2017 abstracts index. We use the default settings, only using standard text search with no semantic or vector search.
B ADDITIONAL EXPERIMENTAL DETAILS
B.I JUSTIPYING GREEDY OPTIMIZATION
We run LeReT for two hops on all 90,447 questions in HotpotQA. After sampling the first hop, chose two sets of documents that have different rewards. We then the second hop with no few-shot prompting and evaluate the reward after the second hop. We find that the lower reward set of documents resulted in higher reward after the second hop in only 0.026% of cases.
14

---

Published as a conference paper at ICLR 2025
B.2 AVERAGE PRECISION
Average precision defined according to Eq. where is the total number relevant documents, P(k) is the precision of the first & documents, and rel(k) is 1 if the leth document is relevant and otherwise:
N P(k) rel(k)
1 AP
(3)
R k=1
B.3 MULTI-HOP SAMPLING
Can we get away without training on data from all hops? We run ablation to determine the neces- sity sampling across multiple hops. Sampling across multiple hops requires computationally and less parallelizable. Specifically, we train Llama 8b on Hotpot and HoVer, sampling only the data from the first half the hops. For Hotpot, this amounts to sampling from just the first hop and for Ho Ver, this amounts sampling data from the first two hops.
Hop
Hops
Hops
Hops
Dataset
Method
RE
AP
RE
AP
RE
AP
RE
AP
1 hop
58.8
54.7
70.8
62.5
Hotpot
All hops (2)
56.7
52.5
77.1
66.3
2 hops
45.5
42.2
62.8
54.3
69.7
59.3
73.6
62.2
HoVer
All hops (4)
45.8
42.5
65.4
56.1
72.8
61.4
76.9
64.3
Table 7: Sampling from all hops instead of only the first few significantly improves perfor- mance. We train Llama 8b on preferences from only the first half of the hops in a pipeline, and compare to training on the full dataset containing preferences over queries from the hops.
We find training on data from all hops leads to substantial improvement in performance across both the datasets, compared to training on only the first half leads to performance gains. The gain is larger for Hotpot where the task only has two hops, and thus never sees data where the context has additional documents from previous hops.
B.4 LONG FORM GENERATIONS
We present a preliminary attempt at long form generation. The majority of LLM use cases are for long form generation, and such want to test LeReT's ability to improve long form generations. In addition, it is more to difficult to evaluate the factuality of long form answers, meaning that eval- uating the relevance of the documents conditioned its answer on as LeReT does may be simpler. Current retrieval datasets focus on short question answering, leading us to generate our own long form dataset.
To create dataset with open ended questions that still had correct retrievals, we prompted GPT for 20 broad topics. From each of those 20 broad topics, we prompted GPT for 500 topics, giving us a total of 10,000 topics such as "Injuries in American football" or "Effects of mobile radiation on human health". We then fed those topics into Colbert and retrieved the top wikipedia abstracts. We then prompt GPT with the 10 wikiepdia abstracts and asked it come up with a question that required students to use exactly of the articles. This gave it the freedom to choose articles that were closely related and led more natural questions than forcing it to given 3 articles.
We then train Llama 3 8b on this dataset using LeReT. We find an approximately 8.47% improve- ment in document retrieval. We create long form generations by feeding the retrieved documents into Llama 3.1 70B. We find that the LcReT-generatin are superior to few shot prompting with a 55.56% win rate.
15

---

Published as a conference paper at ICLR 2025
B.5
POSSIBLE REWARD FUNCTIONS
Disagree (%)
Data Size
Dataset
Hard
Soft
Generator
Retrieval
Hotpot
25.09
38.46
96,766
163,644
Hover
31.48
33.03
10,488
88,448
Table 8: Sampling data using the generator F1 score leads to poor quality data with many wrong preference pairs (high disagree values) and less data (lower data size). We sample pref- erence datasets using the Fl score of the generated answer. Hard disagree is preference pairs where the ranking by the retrieval reward is swapped compared to the ranking by the generator reward. Soft disagree preference pairs where the retrieval reward is equal but the generator reward has a ranking. Data size is the size of the preference dataset generated using each method.
Do we need direct supervision in LeReT for computing the reward function that outputs reward for a given question and set of retrieved documents or can we get away with indirect methods for super- vising retrieval quality? We currently more direct supervision for retrieval but requires knowing correct set of documents advance, as is available in Hotpot/HoVer But, there settings where the optimal retrieved documents may be hard to specify in advance. In such cases, indirect supervision may be easier to provide, where the final generated answer conditioned on the retrieved documents is reviewed, and reward is gener- ated based on the verification of the final answer. Such supervision can be quite weak and have a high amount of noise, as the generator may answer correctly even when conditioned on incorrect documents (because of internal knowledge) or provide incorrect answers even when conditioned on the right documents.
average precision of retrieved documents, which provides
To explore these different settings, we apply LeRe'T using the F1 score of the final generated answer as the reward. We condition the LLM on retrieved documents along with the actual question to generate an answer, and compare that to the correct answer. Formally, instead of using AP(Cha) as the reward, we use = (1( (d, Cha)).
We find that the F1 score of the generator does not provide a very strong signal. For the preference dataset that we constructed using the generator, over 50% of the preference pairs are wrong. We split these incorrect pairs into two categories: hard and soft disagree. Hard disagree means that according to the F1 generation reward, d > qj but according to the AP retrieval reward, l < Qj- Soft disagree means that according to the F1 generation reward, q. 9j but according to the AP retrieval reward, Oj- To reduce the number of soft disagrees, we experiment with adding a threshold for the difference in Fl score to form preference pair, but find that over half the questions HotpotQA and all the questions in Ho have one word answers so this is not effective.
1 Hop
2 Hops
Model
RE
AP
RE
AP
Base model
44.2
40.6
55.5
43.0
LeReT-RRcrieeer
56.7
52.5
77.1
66.3
49.6
45.5
64.5
53.8
Table 9: With the weaker signal of generation, LeReT is able to improve upon the base model but does not match the performance of using the retriever reward. We take the dataset sampled on Hotpot and train Llama 3 8b.
We also find that the datasets generated using the generator reward are significantly smaller than those generated using the AP retriever reward, for example about 8.4 times smaller in case of Ho Since Ver has one word answers, the generator F1 score is less fine grained than the retrieval accuracy over documents. This leads to smaller preference dataset as LeReT provides contexts to the next hop only if they have not achieved the maximum score. In the one-word answer case, there could be queries that retrieve only some (or none) of the correct documents but give the generator enough context to guess use prior knowledge of coarse reward, it is much more likely that a question will be excluded from subsequent hop even though the model has not output optimal query to the retriever.
output the correct answer. Thus, with this sort
16

---

Published as a conference paper at ICLR 2025
When train model this data, find that it significantly improves the base model, but substantially under performs the case where the reward signal is derived from average precision of gold documents. We find that the first hop data is far better than the second hop data, (11.43% hard disagree compared to 37.70, 37 .40 soft disagree compared to 39.78) which likely contributes to the decreased performance the second hop.
B.6 GENERALIZATION
We test how well LeReTT-trained models are able to generalize. Specifically, we take Gemma 9b trained on Hover and evaluate performance on Hotpotqa. Similarly, we take Ilama 8b trained on Hotpolqa and test in on Hover. We find that while these models do not perform well those specifically trained on the dataset, they outperform few-shot prompting and as such with LeReT, models are learning to search in ways that are broadly applicable to different datasets.
Hop
Hops
3 Hops
Hops
Test Dataset
Model
Method
RE
AP
RE
AP
RE
AP
RE
AP
Few-shot
49.9
45.6
64.8
$3.9
HotpotQA
Gemma 9b
LeReT-Hotpot
56.7
52.5
77.1
66.3
LeReT-Hover
51.3
47.2
73.8
60.7
Few-shot
45.6
42.2
53.4
45.9
56.0
46.0
57.3
46.1
HoVer
Llama 8b
LeReT-Hover
45.8
42.5
65.4
56.1
72.8
61.4
76.9
64.3
LeReT-Hotpot
42.6
39.3
60.1
51.6
67.5
56.6
71.2
59.2
Table 10: Models trained with LeReT on given dataset lead to improved performance the other dataset. Gemma 9b trained on Hover tested on Hotpot and Llama 8b trained on Hotpot tested Hover. Both models outperform few shot prompting.
17


Reviews:
[{"review_id": "wFTVJhp0wH", "reviewer": "Reviewer_NW9U", "summary": "In multi-hop question answering, a model needs to perform multiple retrieval steps before arriving at an answer.  Since the reward (answer) is not known until the last step, this problem lends itself well to reinforcement learning (RL).  The paper proposes to optimize one component (the question generator for retrieval) in the multi-hop QA pipeline using RL.  The method uses direct supervision (gold retrieved docs) to first train a reward function.  Then a diverse set of queries are sampled from the model using varied prompts.  The rewards for these generations are fed into an RL algorithm (IPO) to improve the query generator.  The method improves substantially on pure SFT methods on HotpotQA and HoVer.", "questions": "* Has similar methods been used for single-hop QA or RAG in general?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* Well written paper, clear presentation\n* Results are strong and convincingly support the claim that using RL to improve query generator works better than pure SFT.", "weaknesses": "* A bit difficult to judge the novelty of the contribution.  Has similar methods been used for single-hop QA or RAG in general?  If so, the novelty here might be marginal, especially since the paper relies on direct supervision for the reward model.\n* The two multi-hop datasets used are not natural and somewhat out-of-date.  It would be great to see the methods usefulness on more relevant tasks and benchmarks.  The long-form generation attempt in the appendix is interesting, and could perhaps be developed more and moved into the main text?", "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": "5BYIVRgRmF", "reviewer": "Reviewer_v1Yv", "summary": "This paper presents a framework called \"Learning to Retrieve by Trying\" (LeReT), which aims to improve the grounding of large language models (LLMs) through reinforcement learning (RL)-based retrieval. LeReT enables LLMs to generate queries, learn from trial-and-error, and enhance retrieval by using diverse few-shot prompts combined with preference-based reinforcement learning. LeReT's flexibility makes it adaptable across different retrieval systems, with potential applicability in broader retrieval-augmented generation (RAG) contexts.", "questions": "1. While iterative training improves retrieval accuracy, does it potentially overfit the model to specific multi-hop tasks, thereby impacting generalizability in other retrieval-augmented scenarios?\n2. What is involved in adapting LeReT to new domains or types of queries that it was not originally trained on?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Introduces a unique reinforcement learning framework to improve retrieval accuracy in LLMs, especially for complex multi-hop queries.\n2. Demonstrates the effectiveness of iterative training in enhancing the retrieval and grounding abilities of LLMs.\n3. Compatible with various retrieval systems, including ColBERTv2 and Azure AI Search, indicating its broad applicability.", "weaknesses": "1. Primarily relies on direct supervision for labeling relevant documents, which may limit its scalability in cases where explicit relevance labels are unavailable.\n2. Requires extensive computation due to multi-hop retrieval and diverse query sampling, making it resource-intensive.\n3. The need for sampling across multiple hops is computationally intensive and less parallelizable, reducing scalability.", "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": "JLah0zFXwh", "reviewer": "Reviewer_cwbp", "summary": "The paper presents LeReT, a framework for improving LLM retrieval and answer grounding using reinforcement learning. It uses prompt-driven diverse query generation, model optimization with preference-based RL, and reward labeling for retrieved documents. Experiments on HotpotQA and HoVer datasets show significant improvements in retrieval and downstream generation compared to baselines. Analysis reveals the importance of diverse few-shot prompting and LeReT's applicability across different retrievers. Limitations include the use of direct supervision, and future work may focus on indirect supervision and tool updating.", "questions": "- Line 76 \"If LLMs can observe the retrieved documents for different search queries, they can learn which queries lead to better outcomes.\" What supports this claim? And if this is true, then why do you need the direct or indirect supervision to teach LLM \"how to query\"?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- LeReT is applicable to general retrieval-augmented generation (RAG) systems and can adapt to different retrievers.\n- LeReT significantly improves retrieval accuracy. Compared to the unadopted Llama and Gemma instruction models, the recall rate increases by 9-22% on HotPotQA and 27-29% on HoVer.\n- It can be used iteratively: applying LeReT for two iterations shows that the model performance after the second iteration is better than that of the standard non-iterative LeReT.", "weaknesses": "- The novelty assertion of the proposed method lacks clarity. Regarding the related work spanning from line 139 to 148, it remains ambiguous as to how the proposed method differentiates itself from those other methods. I comprehend that the proposed approach employs diverse query generation and IPO for preference learning. However, these seem to be more of incremental enhancements within an existing framework rather than representing a distinct novelty.\n- Also, the experiments lack comparisons with the most relevant recent works. Only the basic few-shot prompt baseline is compared.", "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": "6ASUmT4XM6", "reviewer": "Reviewer_9WtD", "summary": "The paper addresses the issue of hallucination in language models by enhancing their capability to generate effective queries for retrieving relevant facts, thereby improving the accuracy of model responses. The proposed approach, LeReT, generates diverse search queries by incorporating few-shot examples. It constructs comparative queries guided by the supervision of retrieval quality and optimizes the model using the IPO algorithm. The study empirically validates the LeReT framework, exploring various methods of collecting retrieval rewards and assessing performance across different retrievers.", "questions": "1. On lines 218–219, are there any concrete examples illustrating the differences between BFRS-generated queries and high-temperature sampling queries?\n\n2. Regarding line 260, what are the examples used in few-shot prompting? Why is few-shot prompting not feasible during testing? Is it due to considerations of inference efficiency?\n\n3. As for section 4.3, are there any details about how rewards are computed?\n\n4. For Iterative-LeReT and the discussion in Section 4.3, this multi-hop search process seems compatible with online RL training, such as using PPO, by incorporating process rewards at each hop and a final reward for the conclusion. Have the authors explored online RL training?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The motivation that enhancing RAG performance through learning to generate better search queries is clear and reasonable.\n2. The experiments present the improvement in both retrieval quality and finally performance, which validate the soundness of proposed algorithm.", "weaknesses": "1. The writing and presentation of results in this work could be enhanced. On line 254, I think it should be, \"In this work, we consider XXX.\"\n\n2. Some baselines (https://aclanthology.org/2023.emnlp-main.585.pdf, https://openreview.net/pdf?id=vDvFT7IX4O, and https://aclanthology.org/2023.acl-long.99.pdf) could be added to both the experiments and the related works sections. Although these works may not directly focus on multi-hop retrieval, they fit broadly within the same topic as this work, i.e., \"query expansion.\" These works focus on generating better search queries to enhance RAG performance.\n\n3. Relying on annotated golden documents, referred to as \"direct supervision\" in this work, limits the vision of the study. In more general scenarios, collecting \"indirect supervision\" is more feasible. Moreover, signals from \"indirect supervision\" are the ultimate indicators for downstream tasks.", "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)