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

arXiv: https://arxiv.org/abs/2504.11420 · HTML: https://mlsanity.com/paper/2504.11420/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 32 references that are in the map.
- First posted on arXiv 15 Apr 2025.

## Builds on (most cited first, 26 of 26)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Unsupervised Dense Information Retrieval with Contrastive Learning](https://mlsanity.com/paper/2112.09118/) (2021)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [Atlas: Few-shot Learning with Retrieval Augmented Language Models](https://mlsanity.com/paper/2208.03299/) (2022)
- [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://mlsanity.com/paper/2204.06745/) (2022)
- [Learning To Retrieve Prompts for In-Context Learning](https://mlsanity.com/paper/2112.08633/) (2021)
- [Measuring Compositional Generalization: A Comprehensive Method on Realistic Data](https://mlsanity.com/paper/1912.09713/) (2019)
- [Active Example Selection for In-Context Learning](https://mlsanity.com/paper/2211.04486/) (2022)
- [Compositional Generalization and Natural Language Variation: Can a Semantic Parsing Approach Handle Both?](https://mlsanity.com/paper/2010.12725/) (2020)
- [Compositional Exemplars for In-context Learning](https://mlsanity.com/paper/2302.05698/) (2023)
- [Unified Demonstration Retriever for In-Context Learning](https://mlsanity.com/paper/2305.04320/) (2023)
- [Diverse Demonstrations Improve In-context Compositional Generalization](https://mlsanity.com/paper/2212.06800/) (2022)
- [Improving Text-to-SQL Evaluation Methodology](https://mlsanity.com/paper/1806.09029/) (2018)
- [How Do In-Context Examples Affect Compositional Generalization?](https://mlsanity.com/paper/2305.04835/) (2023)
- [Improving Compositional Generalization with Latent Structure and Data Augmentation](https://mlsanity.com/paper/2112.07610/) (2021)
- [Unobserved Local Structures Make Compositional Generalization Hard](https://mlsanity.com/paper/2201.05899/) (2022)
- [DemoRank: Selecting Effective Demonstrations for Large Language Models in Ranking Task](https://mlsanity.com/paper/2406.16332/) (2024)
- [Learning to Retrieve Iteratively for In-Context Learning](https://mlsanity.com/paper/2406.14739/) (2024)
- [Does In-Context Learning Really Learn? Rethinking How Large Language Models Respond and Solve Tasks via In-Context Learning](https://mlsanity.com/paper/2404.07546/) (2024)

## Built on by (most cited first, 1 of 1)

- [The Imperfect Learner: Incorporating Developmental Trajectories in Memory-based Student Simulation](https://mlsanity.com/paper/2511.05903/) (2025)

## Abstract (arXiv)

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.