Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference
Masahiro Kato, Taka Kato
econ.EM, cs.LG, math.ST, stat.ME, stat.ML
Submitted: 2026-07-20
License: http://creativecommons.org/licenses/by-nc-nd/4.0/
The gist: We propose one-step and two-step methods for policy learning with retrieval-augmented generation (RAG).
Terminology
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.
Sources
- Using Embedding Models to Improve Probabilistic Race Prediction
- GenAI-Powered Inference
- Causal Inference with Generative Artificial Intelligence: Application to Texts as Treatments
- Bridging the Gap between Empirical Welfare Maximization and Conditional Average Treatment Effect Estimation in Policy Learning
- Nearest Neighbor Matching as Least Squares Density Ratio Estimation and Riesz Regression
Related papers
- SLIM: Stochastic Learning and Inference in Nonlinear GMM Models
- High-dimensional censored MIDAS logistic regression for corporate survival forecasting
- Cross-Fitting-Free Debiased Machine Learning with Multiway Dependence
- Identifying Elasticities in Autocorrelated Time Series Using Causal Mechanism Graphs
- Policy design in experiments with unknown interference
- Can large language models assist choice modelling? Insights into prompting strategies and current models' capabilities