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

Masahiro Kato and Taka Kato。2026年。プレプリント（arXiv:2607.18225）。 English

RAG を使って行動を選ぶ仕組みを、因果推論の潜在結果の枠組み（potential outcome framework）で定式化した研究です。行動ごとに類似事例をベクトル検索で集める操作が、因果推論の近傍マッチングと同じ構造を持つことを示し、この対応を使って行動選択の後悔（regret）に理論的な保証を与えます。

研究が扱う問題

LLM システムでは、過去の事例を検索して回答や行動を決める構成が広く使われています。しかし、その行動選択がどれだけ良いのかを評価する理論的な土台は、これまでほとんどありませんでした。どの行動を取れば成果が高いかという問いは、因果推論が扱ってきた処置選択の問題そのものです。この研究は、RAG による行動選択を因果推論の言葉で書き直すことで、両分野の道具を接続します。

提案した方法

この研究では、検索と推定を分ける二段階の方法と、中間計算が観測できない一段階の方法を扱います。二段階の方法では、まず行動ごとに、埋め込み空間上のベクトル検索で近い事例を集めます。次に、生成モデルがその事例から条件付き期待成果またはその差を推定し、推定値に基づく規則で行動を選びます。行動別のベクトル検索を近傍マッチングとみなすこの対応により、後悔を、候補生成に由来する部分と候補内の選択に由来する部分に分解できます。後者は、近傍推定量と transformer の予測誤差保証を使って上から抑えます。一段階の方法は、方策そのものとして直接評価します。

主な結果と成立条件

主な結果は、二段階の方法に対する後悔の分解と、その各項の上界です。保証は、検索が埋め込み空間上の近傍マッチングとして働くこと、および成果の推定に予測誤差の保証が使えることを前提としています。検索の質が悪く適切な候補が集まらない場合、候補生成に由来する後悔は残ります。

先行研究との関係

検索を生成に組み込む RAG は Lewis らが提案しました（2020年、Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks）。一方、近傍マッチングによる処置効果の推定は Abadie と Imbens が理論を整えた、因果推論の古典的な方法です（2006年、Large Sample Properties of Matching Estimators for Average Treatment Effects）。データからの方策学習には Athey と Wager の統計理論があります（2021年、Policy Learning with Observational Data）。この研究は、RAG の検索操作をマッチングとして読み替えることで、これらの因果推論の理論を LLM システムの行動選択に持ち込んだ点が新しい貢献です。

利用できる状況

過去の対応事例を検索して次の打ち手を決める、という業務の自動化で、選択の善し悪しに統計的な評価を付けたい場面を想定しています。行動の候補が多く、行動ごとの事例が蓄積されているほど、この定式化が生きます。

論文と資料

BibTeX

@misc{rag-policy-learning-vector-search,
  author       = {Masahiro Kato and Taka Kato},
  title        = {Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference},
  year         = {2026},
  eprint       = {2607.18225},
  archivePrefix = {arXiv},
  url          = {https://arxiv.org/abs/2607.18225},
}