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

This paper investigates the decision-theoretic foundations of Retrieval-Augmented Generation (RAG) when used for action selection. By framing RAG as a policy learning problem within the potential outcome framework (the Neyman-Rubin causal model), the authors seek to understand the guarantees behind RAG-based systems that recommend actions based on retrieved evidence.

The authors propose two methods for RAG-based policy learning: a one-step method, where the model directly outputs an action, and a two-step method, which separates the process into candidate generation and action selection. In the two-step approach, the system first generates a set of candidate actions, then retrieves action-specific evidence using vector search, and finally selects the action with the highest estimated conditional expected outcome. The authors interpret this vector search as a form of nearest-neighbor matching, allowing them to apply statistical tools from causal inference to analyze the system's performance.

The study provides a formal decomposition of the regret associated with the two-step method into two components: candidate-generation regret (the loss from failing to include the optimal action in the candidate set) and within-candidate choice regret (the loss from incorrectly ranking the candidates). By leveraging prediction-error guarantees from nonparametric statistics and transformer theory, the authors derive bounds for the within-candidate regret. This framework clarifies how RAG systems function as plug-in estimators for optimal policy selection, distinguishing between the retrieval of evidence and the final decision-making step.

As RAG systems become increasingly prevalent for automated decision-making and tool use, understanding their causal properties is essential for reliability. This work bridges the gap between modern generative AI architectures and classical causal inference, providing a rigorous mathematical language to evaluate whether RAG-based recommendations are truly optimal or merely artifacts of the retrieval process.