# Maximally-Informative Retrieval for State Space Model Generation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 39 of its 47 references that are in the map.
- First posted on arXiv 13 Jun 2025.

## Builds on (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Toolformer: Language Models Can Teach Themselves to Use Tools](https://mlsanity.com/paper/2302.04761/) (2023)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [MS MARCO: A Human Generated MAchine Reading COmprehension Dataset](https://mlsanity.com/paper/1611.09268/) (2016)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Constructing A Multi-hop QA Dataset for Comprehensive Evaluation of Reasoning Steps](https://mlsanity.com/paper/2011.01060/) (2020)
- [Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://mlsanity.com/paper/2212.03533/) (2022)
- [M3-Embedding: Multi-Linguality, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation](https://mlsanity.com/paper/2402.03216/) (2024)
- [Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions](https://mlsanity.com/paper/2212.10509/) (2022)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [Learning To Retrieve Prompts for In-Context Learning](https://mlsanity.com/paper/2112.08633/) (2021)
- [MTEB: Massive Text Embedding Benchmark](https://mlsanity.com/paper/2210.07316/) (2022)
- [Precise Zero-Shot Dense Retrieval without Relevance Labels](https://mlsanity.com/paper/2212.10496/) (2022)
- [Active Retrieval Augmented Generation](https://mlsanity.com/paper/2305.06983/) (2023)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models](https://mlsanity.com/paper/2405.17428/) (2024)
- [A Survey on RAG Meeting LLMs: Towards Retrieval-Augmented Large Language Models](https://mlsanity.com/paper/2405.06211/) (2024)
- [LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders](https://mlsanity.com/paper/2404.05961/) (2024)
- [Improving Text Embeddings with Large Language Models](https://mlsanity.com/paper/2401.00368/) (2023)
- [RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation](https://mlsanity.com/paper/2310.04408/) (2023)
- [Generative Representational Instruction Tuning](https://mlsanity.com/paper/2402.09906/) (2024)
- [Task2Vec: Task Embedding for Meta-Learning](https://mlsanity.com/paper/1902.03545/) (2019)
- [SGPT: GPT Sentence Embeddings for Semantic Search](https://mlsanity.com/paper/2202.08904/) (2022)
- [Fine-Tuning or Retrieval? Comparing Knowledge Injection in LLMs](https://mlsanity.com/paper/2312.05934/) (2023)
- [Long-Context LLMs Meet RAG: Overcoming Challenges for Long Inputs in RAG](https://mlsanity.com/paper/2410.05983/) (2024)
- [Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach](https://mlsanity.com/paper/2407.16833/) (2024)

## Abstract (arXiv)

Given a query and dataset, the optimal way of answering the query is to make use all the information available. Modern LLMs exhibit impressive ability to memorize training data, but data not deemed important during training is forgotten, and information outside that training set cannot be made use of. Processing an entire dataset at inference time is infeasible due to the bounded nature of model resources (e.g. context size in transformers or states in state space models), meaning we must resort to external memory. This constraint naturally leads to the following problem: How can we decide based on the present query and model, what among a virtually unbounded set of known data matters for inference? To minimize model uncertainty for a particular query at test-time, we introduce Retrieval In-Context Optimization (RICO), a retrieval method that uses gradients from the LLM itself to learn the optimal mixture of documents for answer generation. Unlike traditional retrieval-augmented generation (RAG), which relies on external heuristics for document retrieval, our approach leverages direct feedback from the model. Theoretically, we show that standard top-k retrieval with model gradients can approximate our optimization procedure, and provide connections to the leave-one-out loss. We demonstrate empirically that by minimizing an unsupervised loss objective in the form of question perplexity, we can achieve comparable retriever metric performance to BM25 with no finetuning. Furthermore, when evaluated on quality of the final prediction, our method often outperforms fine-tuned dense retrievers such as E5.