EconBase
← All papers

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

Masahiro Kato, Taka Kato

arXiv 20 Jul 2026 · Econometrics

arXiv:2607.18225 · PDF · Extracted main text

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.

Citation extraction

41
references
68
in-text mentions
41
distinct cited
2
self-citations
9,076
main-text words

appendix boundary found by appendix_command · 60% of the source is main text. Read the extracted text to check this.

Most heavily cited references

The works this paper leans on most, across its whole bibliography — not restricted to papers in our corpus. Ranked by composite intensity, which combines how often a work is mentioned, how many sections mention it, and how much of that falls in the main text rather than the appendix.

ReferenceIntensityMentionsSectionsMain text
1Toru Kitagawa and Aleksey Tetenov (2018) Who should be treated? empirical welfare maximization methods for treatment choice0.87462100%
2Masahiro Kato (2025) Nearest neighbor matching as least squares density ratio estimation and riesz regression, 2025b self0.8434375%
3Zhexiao Lin, Peng Ding, and Fang Han (2023) Estimation based on nearest neighbor matching: from density ratio to average treatment effect0.8434375%
4Susan Athey and Stefan Wager (2021) Policy learning with observational data0.81142100%
5Michelle Ching, Ioana Popescu, Nico Smith, Tianyi Ma, William G. Und… (2026) Efficient and minimax optimal in-context nonparametric regression with transformers0.7373367%
6Juno Kim, Tai Nakamaki, and Taiji Suzuki (2024) Transformers are minimax optimal nonparametric in-context learners0.7373367%
7Kazusato Oko, Yujin Song, Taiji Suzuki, and Denny Wu (2024) Pretrained transformer efficiently learns low-dimensional target functions in-context0.7373367%
8Jean-Yves Audibert and Alexandre B. Tsybakov (2007) Fast learning rates for plug-in classifiers0.73732100%
9Noah Dasanaike and Kosuke Imai (2026) Using embedding models to improve probabilistic race prediction, 20260.5112250%
10Alexander Havrilla and Wenjing Liao (2024) Understanding scaling laws with statistical and approximation theory for transformer neural networks on intrinsically low-dimens…0.5112250%

Showing the top 10 of 41 scored citations.