Maximally-Informative Retrieval for State Space Model Generation

TL;DR

This paper introduces RICO, a gradient-based retrieval method for state space models that optimizes document selection at inference time, improving answer accuracy without additional fine-tuning.

Contribution

We propose Retrieval In-Context Optimization (RICO), a novel gradient-based retrieval approach that directly leverages model feedback to select relevant data for inference.

Findings

RICO achieves comparable performance to BM25 without fine-tuning.

Our method often outperforms fine-tuned dense retrievers like E5.

Theoretically, top-k retrieval with gradients approximates our optimization.

Abstract

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 \emph{no finetuning}. Furthermore, when evaluated on quality of the final prediction, our method often outperforms fine-tuned dense retrievers such as E5.