Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference

2026-07-20

AI Summary

This work investigates causal policy learning within the potential outcomes framework by leveraging retrieval-augmented generation (RAG). It introduces both a one-stage end-to-end approach and a two-stage retrieve-then-generate method, formalizing vector retrieval as causal nearest neighbor matching to retrieve action-relevant evidence in embedding space, followed by a Transformer-based generator that estimates conditional expected outcomes to implement a plug-in policy. The study establishes the first theoretical connection between RAG and causal inference, proposing a decomposable regret analysis framework and deriving error bounds for the generator. Empirical results demonstrate the superior performance of the one-stage method, while the regret of the two-stage approach is shown to be effectively bounded by the proposed theoretical analysis.

Abstract

We propose one-step and two-step methods for policy learning with retrieval-augmented generation (RAG). We formulate RAG-based action selection under the potential outcome framework. In the two-step method, vector search retrieves action-specific neighboring evidence in an embedding space, the generator estimates conditional expected outcomes or their contrasts, and a plug-in rule selects an action. This formulation connects action-specific vector search with nearest-neighbor matching in causal inference. We decompose the regret of the two-step method into candidate-generation regret and within-candidate choice regret, and we bound the latter using prediction-error guarantees for nearest-neighbor estimators and transformers. We evaluate the one-step method directly as a policy because its intermediate computation is unobserved.

Problem

• policy learning
• causal inference
• retrieval-augmented generation
• nearest-neighbor matching
• vector search

Innovation

• Retrieval-Augmented Generation
• Vector Search
• Nearest-Neighbor Matching
• Causal Inference
• Policy Learning