Reinforcing Compositional Retrieval: Retrieving Step-by-Step for Composing Informative Contexts (2025)

Abstract: Large Language Models (LLMs) have demonstrated remarkable capabilities across numerous tasks, yet they often rely on external context to handle complex tasks. While retrieval-augmented frameworks traditionally focus on selecting top-ranked documents in a single pass, many real-world scenarios demand compositional retrieval, where multiple sources must be combined in a coordinated manner. In this work, we propose a tri-encoder sequential retriever that models this process as a Markov Decision Process (MDP), decomposing the probability of retrieving a set of elements into a sequence of conditional probabilities and allowing each retrieval step to be conditioned on previously selected examples. We train the retriever in two stages: first, we efficiently construct supervised sequential data for initial policy training; we then refine the policy to align with the LLM’s preferences using a reward grounded in the structural correspondence of generated programs. Experimental results show that our method consistently and significantly outperforms baselines, underscoring the importance of explicitly modeling inter-example dependencies. These findings highlight the potential of compositional retrieval for tasks requiring multiple pieces of evidence or examples.

## introduction large language models (llms) have demonstrated remarkable progress in recent years, mastering text generation and diverse problem-solving. context often plays a critical role in grounding these models in task-relevant information: it helps incorporate domain-specific knowledge, clarify ambiguous queries, and provide demonstrative examples. this leads to the common adoption of retrieval-augmented frameworks @xcite figure 1 : when generating the target semantic program of a query (parsing the query into the logical form), relying on single or repetitive demonstrative context cannot cover the ground-truth program local structure and lead to poor context augmentation quality. in order to capture the whole retrieval results' collective effect to facilitate better local structure coverage and program generalization, we solve this compositional retrieval problem by modeling it as a markov decision process. 2020; @xcite , which introduce a "retrieve" step to extract relevant context segments such as documents or examples to enrich the llm's input, improving both the accuracy and reliability of its outputs. however, many complex tasks demand combining multiple pieces of evidence or examples of diverse semantics or structures. for instance, multihop queries often require several documents to accumulate all necessary insights. another scenario, as illustrated in figure 1 , involves translating a query into a logical forms, where an llm needs to assemble symbols or sub-trees from multiple demonstrative examples to construct the target program accurately. in this case, relying on a single retrieved example can only provide partial hints, while repetitive or redundant examples provide little additional information. consequently, this compositional retrieval setting necessitates coordinated selection of examples to maximize collective utility of chosen composition. to address these complexities, conventional topk retrieval is insufficient, as it operates in a singleturn manner and fails to capture how newly selected items condition subsequent choices. therefore, we propose a compositional retrieval paradigm that conducts retrieval in multiple steps, with each step adapting to the evolving context formed by previously retrieved items. this approach ensures that newly selected examples complement previously chosen ones, improving the overall coverage of retrieved content (figure 1 ). although recent works have explored sequential retrieval for selecting exemplary contexts @xcite , they typically rely on a large scoring language model for constructing training sequences, incurring significant computational overhead. to maximize the collective utility of the retrieved examples instead of selecting multiple highscoring yet potentially redundant elements from a single global ranking, we formulate compositional retrieval as a markov decision process (mdp), wherein each state reflects both the query and the previously chosen items. we adopt a tri-encoder architecture as the policy model, mapping each query and partial retrieval sequence to a probability distribution over candidate examples. to initialize this policy, we construct supervised fine-tuning data in a more efficient manner by maximizing sub-structure coverage without scoring models. in the second stage, we perform reinforcement learning (rl) and adopt a grouped policy optimization @xcite , reinforcing the policy using a task-specific reward based on local structure coverage of the generated logical form. this design captures the llm's compositional needs and progressively enhances the quality of the assembled context. we evaluate our method reinforcing compositional retrieval (rcr) on compositional generalization semantic parsing benchmarks which requires generating new combinations of familiar structures and symbols, making it an ideal initial scenario to examine how context compositions impact program generation capabilities. results show our method consistently outperforms top-k and sequential retrieval baselines. the observed improvements stem from our tri-encoder retriever's ability to model inter-example dependencies, enabling more effective context composition and enhancing program generation quality within the llm. our analysis further underscores the effectiveness of leveraging rl in refining retrieval strategies. we systematically evaluate different rl setups and advantage estimation methods, providing a comprehensive evaluation and insights into both its benefits and limitations in retrieval tasks. in summary, our key contributions are: ## problem formulation in this work, we depart from conventional topranked retrieval which treat each candidate independently. instead, we propose a compositional retrieval problem that aims to optimize the selection of the most influential context group, and each element within this group will contribute together to lead to a correct prediction. ## reinforcing compositional retrieval the training of our propose tri-encoder sequential retriever contains two stages, supervised finetuning (sft) and reinforcement learning (rl). ## related work early example selection methods retrieve semantically similar demonstrations @xcite , but semantic proximity alone does not always yield optimal exemplars. to improve selection, @xcite propose training a dense retriever with a scoring lm, ranking candidates by the loglikelihood of the ground-truth output. extensions further refine this retriever across multiple tasks @xcite . other studies highlight the importance of diversity in retrieval, such as structure or label diversity @xcite @xcite . however, many of these approaches treat each candidate independently, overlooking inter-example dependencies that influence icl effectiveness. recent empirical studies @xcite reveal that diversity (e.g., structure diversity, label diversity) also play an important role in demonstration selection. @xcite select diverse demonstrations that aims to collectively cover all of the required structures in the semantic parsing program. ye et al. (2023) learns to select diverse example set based on the conditional determinantal point process. however, previous works @xcite overlook the interplay between in-context examples as they treat each candidate example independently. this may be suboptimal as the in-context examples can influence each other, previous work indicate that icl is sensitive to the order of in-context examples @xcite . to capture the exemplar dependency, recent works @xcite and @xcite propose a data construction method by sequentially select and extend the in-context sequence from a llm, the constructed dependency-aware data will be leverage to train a retriever via contrastive loss. different from indirect method that learning dependency via constructed training data, zhang et al. (2022) directly frame icl problem as a decision making problem and use offline q-learning to actively select examples from unlabeled set to label. however, this method does not learn a retriever to acquire effective examples for each test instance, during the inference, this method will traverse the entire training set to greedly choose an action based on the q-value which is produced by a scoring lm as well. therefore, this method is costly when action space (i.e., training dataset) is large and it lacks diversity. although our method is also decision making, our method presents key differences with @xcite , as our method is retrievalbased, allowing sample and inference efficiency and to explore diversity. a similar work framing retrieval as mdp is proposed by @xcite , their reward is primarily likelihood-oriented (whitebox), whereas ours focuses on response-oriented (black-box). besides, our reward design focus heavily on structural features which is crucial for program generation. a comprehensive comparison of related methods and their properties is provided in appendix c, with a detailed summary in ## conclusion in this work, we introduce compositional retrieval, a retrieval paradigm that models the selection of multiple interdependent examples. to achieve this, we propose a tri-encoder sequential retriever, formulating retrieval as a mdp and training it in two stages sft via efficient data construction and rl for policy refinement. experiments highlight the potential of compositional retrieval for tasks requiring multiple sources of evidence and offer insights into the benefits and limitations of rl in retrieval.