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

2026-07-20 · 17:57 UTC ·Masahiro Kato, Taka Kato ·1 min read

Authors

Masahiro Kato, Taka Kato

Categories

econ.EM, cs.LG, math.ST, stat.ME, stat.ML

arXiv

https://arxiv.org/abs/2607.18225v1

PDF

https://arxiv.org/pdf/2607.18225v1

Brief

Kato & Kato formulate retrieval-augmented generation (RAG) action selection within the potential-outcome causal framework and propose one-step and two-step policy-learning procedures. The two-step method retrieves action-specific neighbors in embedding space, estimates conditional outcomes or contrasts, and applies a plug-in decision rule; they relate this to nearest-neighbor matching, decompose regret, and bound the within-candidate term via prediction-error guarantees for nearest neighbors and transformers.

Why it matters

Masahiro Kato and Taka Kato (arXiv:2607.18225v1, 2026-07-20) propose one-step and two-step RAG-based policy-learning methods framed under the potential-outcome model; the two-step pipeline uses action-specific vector search to retrieve embedding-space neighbors, a generator to estimate conditional expected outcomes or contrasts, and a plug-in rule to select actions.

Key details

The authors connect action-specific vector search to nearest-neighbor matching in causal inference, decompose two-step regret into candidate-generation regret and within-candidate choice regret, and prove bounds on the within-candidate regret using prediction-error guarantees for nearest-neighbor estimators and transformers.

They evaluate the one-step method directly as a learned policy because its intermediate retrieval/estimation computations are unobserved in deployment (abstract and PDF available on arXiv).

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.