Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference
2026-07-20 • Machine Learning
Machine Learning
AI summaryⓘ
The authors introduce new ways to teach a system to make decisions using retrieval-augmented generation (RAG). They propose a two-step approach where the system first finds relevant information using vector search, then predicts the best action based on this info, linking it to methods from causal inference. They analyze the errors involved in this process and provide guarantees about its performance. They also test a simpler one-step method that decides directly without showing intermediate steps.
retrieval-augmented generationpolicy learningpotential outcome frameworkvector searchcausal inferencenearest-neighbor matchingregret decompositiontransformersembedding spaceconditional expected outcomes
Authors
Masahiro Kato, Taka Kato
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.