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

TL;DR

This work tackles the challenge of compositional retrieval for tasks requiring multiple pieces of evidence by reframing retrieval as a Markov Decision Process. It introduces Reinforcing Compositional Retrieval (RCR), a tri-encoder sequential retriever trained in two stages—supervised fine-tuning (SFT) to maximize local-structure coverage and reinforcement learning (RL) with Group Relative Policy Optimization (GRPO) to align with downstream LLM preferences using a structure-based reward. Empirically, RCR outperforms top-$k$ and sequential baselines on compositional generalization benchmarks GeoQuery and COVR-10, with ablations highlighting the value of explicit inter-example dependencies and RL refinement. The method demonstrates the potential of explicit contextual composition for semantic parsing and related tasks, offering a scalable approach to assembling informative contexts for LLMs.

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.

Figures (8)

• 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.
• Figure 2: Tri-encoder model to retrieve step-by-step.
• Figure 3: Local structures when size $l=4$. Following levy-etal-2023-diverse, we exclude more distant familial ties.
• Figure 4: Group Relative Policy Optimization (GRPO) process for optimizing the sequential retriever.
• Figure 5: Performance over various numbers of training samples in GRPO. Results are evaluated in GeoQuery-Template2.