Learning to Retrieve Iteratively for In-Context Learning (2024)

Abstract: We introduce iterative retrieval, a novel framework that empowers retrievers to make iterative decisions through policy optimization. Finding an optimal portfolio of retrieved items is a combinatorial optimization problem, generally considered NP-hard. This approach provides a learned approximation to such a solution, meeting specific task requirements under a given family of large language models (LLMs). We propose a training procedure based on reinforcement learning, incorporating feedback from LLMs. We instantiate an iterative retriever for composing in-context learning (ICL) exemplars and apply it to various semantic parsing tasks that demand synthesized programs as outputs. By adding only 4M additional parameters for state encoding, we convert an off-the-shelf dense retriever into a stateful iterative retriever, outperforming previous methods in selecting ICL exemplars on semantic parsing datasets such as CalFlow, TreeDST, and MTOP. Additionally, the trained iterative retriever generalizes across different inference LLMs beyond the one used during training.

## introduction a significant emergent capability of large language models (llms) is in-context learning (icl; @xcite , which facilitates few-shot learning. in icl, a set of exemplars 1 is usually provided to build the mapping relationship between inputs and outputs. these exemplars can either be hand-crafted and fixed or retrieved from a training set. however, if retrieving from the dataset, the retrievers used in such applications are typically off-the-shelf models (e.g., contriever @xcite ) that do not consider interactions among retrieved items when † johns hopkins university. performed while interning at microsoft. this approach can be formulated as markov decision processes (mdps). at each step, the action taken by the retriever is a retrieval call that fetches (potentially multiple) documents from the dataset d. 2 the policy is trained to optimally select exemplars at each step so that the overall trajectory maximizes the reward, leading to better icl performance. by leveraging the llms as environments, we create simulators that allow a policy to roll out in the environment and receive feedback on the effectiveness of the composed prompts, measured by a reward (metric). thus, exemplar selection and prompt composition can be framed as policy optimization aimed at maximizing rewards, which can be addressed through reinforcement learning. we situate our study in in-context semantic parsing due to its difficulty, popularity, and practical value. 3 we instantiate an iterative retriever and investigate the performance of policy learning under this setup. our contributions include: ## overview of an iterative retriever we consider the problem of in-context learning (icl): given a dataset @xmath0, 𝑦 𝑖 )} 𝑖 of exemplars, a retriever 𝑅 retrieves a sequence of exemplars 𝑅(𝑥) based on input query 𝑥 and generate the answer 𝑦 based on the distribution 𝑃 lm (•|𝑥; 𝑅(𝑥)). this retriever 𝑅 : @xmath1. the goal of the retriever 𝑅 is to select a sequence of exemplars ((𝑥 𝑖 , 𝑦 𝑖 )) 1≤𝑖 ≤𝐾 such that the probability of the expected output 𝑦 is maximized: however, this is a combinatorial optimization problem that is computationally infeasible to solve exactly. much of prior work resort to selecting top-𝑘 exemplars based on a scoring function 𝑆: prior work has differed on the choice of the scoring function 𝑆: bm25 @xcite , coverage @xcite , etc. however, such method did not model the interaction between the retrieved exemplars and the language model. we propose an iterative version, where we create a retrieval state 𝑠, and for each step 𝑖 one exemplar (𝑥, 𝑦) ∈ d is retrieved. this is an approximation to the optimization problem in equation 1 . after 𝐾 steps, the retrieved sequence would be 𝑅 iter (𝑥) = ((𝑥 𝑖 , 𝑦 𝑖 )) 1≤𝑖 ≤𝐾 . this formulation of an iterative retriever naturally fits in the definition of a markov decision process (mdp). here, our decision process comprises of (d * , d, 𝜏, 𝑟), where this is a maximum inner product search (mips) problem, and thus can be solved with a vector index such as faiss @xcite . this is different from a regular retriever, where a single retrieval call to the training set 𝑅(𝑥) = arg max ( 𝑥,𝑦) ∈d q • f enc (𝑥) is made. the iterative retriever navigates the encoding space of exemplars, adjusting the query vector q ′ at each step based on previously selected exemplars 𝑠, thus steering the search process to find new candidates. figure 2 demonstrates the process of such an iterative retriever. this stateful design allows for optimized retrieval results through iterative interactions, incorporating signals from both external sources (llms) and internal states (previously retrieved items tracked via state transitions). ## training environment simulator to construct feedback (or reward) from the underlying llms, we treat llms as environments where actions are performed and evaluated. we design an iterative prompting schedule within this llm environment to simulate the process of iterative retrieval and corresponding icl prompt execution. at each step 𝑖, the current sequence of chosen exemplars, 𝑠 𝑖 , is turned into an llm prompt using a predefined template, 5 then used for llm generation. this schedule effectively simulates the real-world scenario of prompting llms, allowing us to observe various execution dynamics, such as generated hypotheses and their probabilities. reward design technically, if the final task metric were available, it can be used directly as the reward to optimize for. however, such a reward is often too coarse to reflect differences in partially correct results. for example, if the metric is exact match, which is common in semantic parsing tasks, the reward would simply be the kronecker delta 𝛿(𝑦 * , 𝑦), yielding 1 only if the prediction 𝑦 exactly matches the reference 𝑦 * , and 0 otherwise. given that the llm simulator provides access to the probabilities of generated sequences, 6 we employ a more general reward design that is not taskspecific. our reward leverages the lm completion probability of the reference sequence 𝑃 lm (𝑦 * |𝑥) @xcite , which captures subtle changes in the likelihood of the lm generating the target sequence with respect to changes in the input 𝑥. in icl, more exemplars typically result in better performance before reaching saturation. inspired by @xcite , we further refine the reward to reflect the increase in the likelihood of the reference 𝑦 * given the prompt. this is a proxy value that measure how much this exemplar contribute to generating the reference parse. this design encourages the model to select exemplars that most significantly contribute to the final result 5 refer to appendix a.3 for the template used in this work. 6 this is generally accessible in many llm inference implementations such as vllm @xcite . for openai-style apis, this can be accessed using the "echo" parameter. policy optimization we employ proximal policy optimization (ppo; @xcite to train an iterative retriever for its stability and efficiency. 7 one core idea of ppo is to define a clipping term that controls the policy optimization process, so that variance is reduced. given a trajectory (𝑥 1 , • • • , 𝑥 𝑇 ), we have advantage â𝑖 at step 𝑖 describes how much better it is to take a specific action 𝑥 𝑖 at state 𝑠 𝑖 , over randomly selecting an action according to 𝜋(𝑥 𝑖 |𝑠 𝑖 ). to compute it, besides the neural model defined in §3, we follow common practice in reinforcement from human feedback (rlhf; @xcite to add a single linear layer to serve as a state-value function 𝑉 (𝑠) = v•s that maps states to values. generalized advantage estimate (gae; @xcite is then used for variance-reduced advantage estimation atop the learned state-value function: ## experimental setup datasets we validate our pilot iterative retriever for icl on a set of semantic parsing datasets, namely smcalflow @xcite , treedst @xcite , and mtop (english portion only; @xcite , following the benchclamp benchmark @xcite . samples of representations are shown in figure 3 . for statistics, see appendix a.1. baselines we compare our iterative retriever (henceforth denoted as iterr) with a range of offthe-shelf retrievers, including bm25 @xcite and a dense encoder, contriever @xcite . additionally, we benchmark against two strong baselines from prior work on improving exemplar selection: for the epr baseline, we replace the base dense retrieval encoder with contriever instead of s-bert @xcite for fair comparison. following @xcite , we use the trained epr model as initialization for ceil. similarly, the same epr checkpoint is used to initialize the text encoder in iterr. note that in iterr, we freeze the weights of the epr encoder and only train the gru-based state transition function, policy network, and value network, resulting in 4m more parameters compared to the original contriever (110m → 114m). for retrievers without iterative capabilities, we take only the top-𝑘 retrieved items and keeping their original ranks. for epr, ceil, and iterr, we selected the best performing model checkpoints on the validation set. all generation is run with 10 exemplars; i.e. 𝑘 = 10. ## results & analyses we evaluate the performance of different retrievers by comparing their downstream icl performance on semantic parsing ( the gain in em is intuitive since it aligns with the training objective, which involves the probability of generating target parses. the improvement in smatch indicates that iterr optimizes retrieval generalization across inference llms iterr benefits from interactive training with an underlying llm. while training incurs costs, these can be minimized by training only once, ideally using a smaller lm. hence in this section we investigate the generalization capabilities of iterr trained with a smaller lm 𝐴, but used for generation under a larger lm 𝐵. in the following experiments, iterr is trained with llama-2-7b as the environment, but used for (a) intra-family lms: variants within the llama-2 model family; and (b) inter-family lms: mistral @xcite ) from a different model family. we follow the setup described in §6, substituting only the llm. as shown in figure 5 , iterr significantly outperforms (> 1% gain) baselines for 75% of the settings and is comparable to a prior strong baseline (within 1% in absolute performance) for 15% of settings, demonstrating its generalization within and beyond its own model family. in intra-family generalization, performance metrics improve with larger model sizes, and iterr consistently outperforms all baselines. this improvement is most evident with larger models such as llama-2-70b and codellama-70b-instruct. for inter-family generalization, iterr maintains its advantage across datasets, though this is less pronounced than within the same model family. this is expected, as the signal from llm simulator is more representative for models sharing the same pre-training procedure. notably, with mistral, contriever performs worse than bm25 on mtop, but iterr still shows improvement. this suggests that iterr, comprising a frozen epr and additional gru layers, can learn task-specific abilities not present in the vanilla epr. ## additional related work llms as environment in @xcite used policy gradient to learn a dense retriever for icl exemplar retrieval, but the state does not contain previously selected examples, and thus is not iterative and unable to model exemplar order. @xcite used 𝑄-learning rl for icl exemplar reordering, with a similar reward design like ours. however, the proposed method does not extend to exemplar retrieval, since the policy space is too large to be handled by 𝑄-learning. few-shot semantic parsing few-shot semantic parsing using llms has shown impressive capabilities in understanding new examples with minimal training data @xcite . however, these parsers often struggle with generalization and fail to parse unobserved local structures due to their limited access to information encoded through exemplars @xcite . to this end, recent research has explored various approaches to improving exemplar selection. epr @xcite used a proxy lm to score outputs from an unsupervised retriever, enabling better training of a dense retriever. @xcite , @xcite , and @xcite emphasize learning to select exemplars based on particular criteria, such as diversity measures and coverage of local structures, to enhance compositional generalization. while these approaches have shown performance improvements in semantic parsing tasks, these are highly based on heuristics constructed from researcher's experience. our approach could be seen as an automated version (through rl) of seeking information useful for semantic parsing. ## conclusion we proposed iterative retrievers that iteratively builds a prompt to perform in-context learning. such retrievers are framed as markov decision processes and trained via policy optimization from llm feedback, where the policy directs which exemplar to append to the existing exemplar sequence. experiments on semantic parsing demonstrated performance gain of iterative retrievers over various datasets and state-of-the-art baselines, showing that they are able to construct prompts that improves incontext learning and downstream llm generation.