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

arXiv: https://arxiv.org/abs/2607.18225 · HTML: https://mlsanity.com/paper/2607.18225/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Agentic RAG, agents and memory, one of 15 directions in Retrieval-augmented generation.
- Builds on 29 of its 39 references that are in the map.
- First posted on arXiv 20 Jul 2026.

## Placed in

- Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/)

## Builds on (most cited first, 29 of 29)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [REPLUG: Retrieval-Augmented Black-Box Language Models](https://mlsanity.com/paper/2301.12652/) (2023)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Nonparametric regression using deep neural networks with ReLU activation function](https://mlsanity.com/paper/1708.06633/) (2017)
- [Are Transformers universal approximators of sequence-to-sequence functions?](https://mlsanity.com/paper/1912.10077/) (2019)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Adaptivity of deep ReLU network for learning in Besov and mixed smooth Besov spaces: optimal rate and curse of dimensionality](https://mlsanity.com/paper/1810.08033/) (2018)
- [Retrieval-Augmented Reinforcement Learning](https://mlsanity.com/paper/2202.08417/) (2022)
- [Deep learning is adaptive to intrinsic dimensionality of model smoothness in anisotropic Besov space](https://mlsanity.com/paper/1910.12799/) (2019)
- [STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learning](https://mlsanity.com/paper/2412.15182/) (2024)
- [Transformers are Minimax Optimal Nonparametric In-Context Learners](https://mlsanity.com/paper/2408.12186/) (2024)
- [Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Input](https://mlsanity.com/paper/2305.18699/) (2023)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Data](https://mlsanity.com/paper/2411.06646/) (2024)
- [Large-Scale Retrieval for Reinforcement Learning](https://mlsanity.com/paper/2206.05314/) (2022)
- [A semiparametric instrumental variable approach to optimal treatment regimes under endogeneity](https://mlsanity.com/paper/1911.09260/) (2019)
- [Estimating Optimal Treatment Rules with an Instrumental Variable: A Partial Identification Learning Approach](https://mlsanity.com/paper/2002.02579/) (2020)
- [Causal Inference with Generative Artificial Intelligence: Application to Texts as Treatments](https://mlsanity.com/paper/2410.00903/) (2024)
- [REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments](https://mlsanity.com/paper/2412.04759/) (2024)
- [CausalRAG: Integrating Causal Graphs into Retrieval-Augmented Generation](https://mlsanity.com/paper/2503.19878/) (2025)
- [Nearest Neighbor Matching as Least Squares Density Ratio Estimation and Riesz Regression](https://mlsanity.com/paper/2510.24433/) (2025)
- [Efficient and Minimax Optimal In-context Nonparametric Regression with Transformers](https://mlsanity.com/paper/2601.15014/) (2026)
- [Using Embedding Models to Improve Probabilistic Race Prediction](https://mlsanity.com/paper/2604.22555/) (2026)

## Abstract (arXiv)

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.