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

Masahiro Kato and Taka Kato. 2026. Preprint (arXiv:2607.18225). 日本語版

This study formulates a mechanism that chooses actions with RAG inside the potential outcome framework of causal inference. Gathering similar past cases by vector search, separately for each action, has the same structure as nearest neighbor matching in causal inference. We use that correspondence to give theoretical guarantees on the regret of the resulting choices.

The problem

LLM systems routinely decide what to answer or what to do by retrieving past cases. Almost no theoretical basis existed for judging how good those choices are. Which action yields the better outcome is exactly the treatment selection problem that causal inference has worked on. Rewriting RAG-based action selection in the language of causal inference connects the tools of the two fields.

The proposed method

The two-stage method first collects nearby cases for each action by vector search in an embedding space. A generative model then estimates the conditional expected outcome, or the difference between such outcomes, from those cases, and a rule based on the estimates picks the action. Reading the per-action vector search as nearest neighbor matching lets us split the regret into a part that comes from generating the candidates and a part that comes from choosing among them. The second part is bounded above using nearest neighbor estimators together with prediction error guarantees for transformers. A one-stage method, whose intermediate computation cannot be observed, is evaluated directly as a policy.

Main results and conditions

The main results are the regret decomposition for the two-stage method and an upper bound on each of its terms. The guarantees assume that retrieval acts as nearest neighbor matching in the embedding space, and that prediction error guarantees are available for estimating the outcome. When retrieval is poor and no suitable candidate is collected, the regret from candidate generation remains.

Relation to earlier work

RAG, which builds retrieval into generation, was proposed by Lewis and coauthors (2020, Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks). Treatment effect estimation by nearest neighbor matching, on the other hand, is a classical method of causal inference whose theory was worked out by Abadie and Imbens (2006, Large Sample Properties of Matching Estimators for Average Treatment Effects). For learning a policy from data there is the statistical theory of Athey and Wager (2021, Policy Learning with Observational Data). The new contribution here is to reread the retrieval step of RAG as matching, which carries this causal inference theory over to action selection in LLM systems.

Where it applies

Think of automating work in which the next move is decided by searching past cases, where someone also wants a statistical assessment of whether the choice was a good one. The more actions there are to choose among, and the more cases have accumulated for each action, the more this formulation gives you.

Paper and materials

BibTeX

@misc{rag-policy-learning-vector-search,
  author       = {Masahiro Kato and Taka Kato},
  title        = {Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference},
  year         = {2026},
  eprint       = {2607.18225},
  archivePrefix = {arXiv},
  url          = {https://arxiv.org/abs/2607.18225},
}