Using retrieval-augmented generation as nearest‑neighbor matching to learn policies from example data
Researchers propose a way to use retrieval-augmented generation (RAG) systems to choose actions in situations where we care about causal outcomes. The aim is to pick the action that gives the highest expected result for a given situation, based on past observed cases. The paper shows how a two-step RAG procedure — first finding candidate actions with vector search, then estimating their expected outcomes and choosing the best — can be viewed as a form of nearest‑neighbor matching from causal inference.
The authors place the idea inside the potential‑outcome (Neyman–Rubin) framework. They treat a decision context as covariates X, actions as A, and the outcomes that would occur under each action as Y(a). In the two-step procedure the system first generates a finite set of candidate actions. For each candidate, a vector search in an embedding space retrieves nearby past cases that match that action. A generative model then estimates the conditional expected outcome for each candidate, and a plug‑in rule picks the action with the largest estimated outcome. The paper also describes a one‑step RAG method that returns an action directly; that method is evaluated only at the level of the final policy because its internal retrieval and scoring steps are not observable.
A key contribution is interpreting action‑specific vector search as nearest‑neighbor matching. Rather than choosing the action that appears most often near the query, the method uses the retrieved nearest neighbors to estimate expected outcomes for each action. This view lets the authors separate two kinds of error when the chosen policy is worse than the best possible one. They call these candidate‑generation regret (loss from a weak candidate set) and within‑candidate choice regret (loss from estimating or ranking candidates poorly). The paper gives bounds on the within‑candidate regret using prediction‑error guarantees for nearest‑neighbor estimators and for transformer‑based generators.