mlsanityData to 14 Oct 2026

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

arXiv 2607.18225 ↗First posted 20 Jul 2026econ.EM cs.LG math.STcited by 0

On the map: Retrieval-augmented generation → Agentic RAG, agents and memory

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.

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.

Builds on

29 of 39 references.

Closest ancestors · 2

Same direction, newest first

  1. 2024REGENT: A Retrieval-Augmented Generalist Agent That Can Act In-Context in New Environments6 Dec
  2. 2022Retrieval-Augmented Reinforcement Learning17 Feb

Other references · 27

  1. 2020Retrieval-Augmented Generation for Knowledge-Intensive NLP Taskscited by 9,297
  2. 2020Dense Passage Retrieval for Open-Domain Question Answeringcited by 3,824
  3. 2020REALM: Retrieval-Augmented Language Model Pre-Trainingcited by 2,141
  4. 2023Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflectioncited by 1,839
  5. 2020Leveraging Passage Retrieval with Generative Models for Open Domain Question Answeringcited by 1,210
  6. 2019Generalization through Memorization: Nearest Neighbor Language Modelscited by 800
  7. 2023REPLUG: Retrieval-Augmented Black-Box Language Modelscited by 792
  8. 2022Transformers learn in-context by gradient descentcited by 732
  9. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 715
  10. 2017Nonparametric regression using deep neural networks with ReLU activation functioncited by 433
  11. 2019Are Transformers universal approximators of sequence-to-sequence functions?cited by 406
  12. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 298
15 more
  1. 2018Adaptivity of deep ReLU network for learning in Besov and mixed smooth Besov spaces: optimal rate and curse of dimensionalitycited by 162
  2. 2019Deep learning is adaptive to intrinsic dimensionality of model smoothness in anisotropic Besov spacecited by 46
  3. 2024STRAP: Robot Sub-Trajectory Retrieval for Augmented Policy Learningcited by 44
  4. 2024Transformers are Minimax Optimal Nonparametric In-Context Learnerscited by 37
  5. 2023Approximation and Estimation Ability of Transformers for Sequence-to-Sequence Functions with Infinite Dimensional Inputcited by 32
  6. 2024Pretrained transformer efficiently learns low-dimensional target functions in-contextcited by 31
  7. 2024Understanding Scaling Laws with Statistical and Approximation Theory for Transformer Neural Networks on Intrinsically Low-dimensional Datacited by 30
  8. 2022Large-Scale Retrieval for Reinforcement Learningcited by 24
  9. 2019A semiparametric instrumental variable approach to optimal treatment regimes under endogeneitycited by 22
  10. 2020Estimating Optimal Treatment Rules with an Instrumental Variable: A Partial Identification Learning Approachcited by 17
  11. 2024Causal Inference with Generative Artificial Intelligence: Application to Texts as Treatmentscited by 16
  12. 2025CausalRAG: Integrating Causal Graphs into Retrieval-Augmented Generationcited by 12
  13. 2025Nearest Neighbor Matching as Least Squares Density Ratio Estimation and Riesz Regressioncited by 7
  14. 2026Efficient and Minimax Optimal In-context Nonparametric Regression with Transformerscited by 7
  15. 2026Using Embedding Models to Improve Probabilistic Race Predictioncited by 2

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025REX-RAG: Reasoning Exploration with Policy Correction in Retrieval-Augmented Generation
  2. 2025Optimizing Retrieval for RAG via Reinforcement Learning
  3. 2026Vector Policy Optimization: Training for Diversity Improves Test-Time Search
  4. 2023Learning Action Embeddings for Off-Policy Evaluation
  5. 2026Active Causal Experimentalist (ACE): Learning Intervention Strategies via Direct Preference Optimization
  6. 2025CDF-RAG: Causal Dynamic Feedback for Adaptive Retrieval-Augmented Generation
  7. 2025Policy Learning with a Natural Language Action Space: A Causal Approach
  8. 2026The RAT: A Unified Bayesian Model for RAG Evaluation