Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
58,137 characters · 29 sections · 45 citation commands
Vector Search As Nearest Neighbor Matching: RAG-based Policy Learning in Causal Inference
\noindentKeywords: causal inference, policy learning, retrieval-augmented generation, vector search, nearest-neighbor matching, plug-in classifier, conditional average treatment effect, minimax rate, in-context learning.
Retrieval-augmented generation (RAG) has become a common architecture for action selection in systems built on language models. In a RAG algorithm, a query is embedded in a vector space, and documents or other information are retrieved according to distances in that space. Conditional on the retrieved evidence, a generative model produces an answer, which may include a recommendation or a tool choice.
Despite its widespread use, the decision-theoretic guarantees of RAG have not been fully studied. If a query is treated as a covariate and the action recommended in the generated answer is treated as the decision, RAG-based decision-making can be viewed as a policy learning problem. Action selection and policy learning have long been studied in causal inference and reinforcement learning.
In causal inference, counterfactual decision-making has been studied as a policy learning problem Athey2021policylearning. The goal is to recommend the action with the highest conditional expected outcome, so estimation of the conditional expected outcome is central. Two main approaches are the plug-in approach and the counterfactual risk minimization (CRM) approach Swaminathan2015counterfactualrisk; the latter is also called empirical welfare maximization Kitagawa2018whoshould. The plug-in approach first estimates the conditional expected outcome for each action and then chooses the action with the largest estimate. By contrast, the CRM approach estimates a value functional and learns a policy by maximizing the estimated value.
We formulate RAG-based action selection under the potential outcome framework, also known as the Neyman--Rubin causal model. The explicit two-step procedure is a plug-in policy rule because it estimates action-specific conditional expected outcomes before choosing an action. We evaluate the one-step output as a policy without imposing the same interpretation on its unobserved internal computation.
We first consider the case in which $K$ actions are given. Let $X$ denote covariates that represent the query or decision context supplied to the RAG system, let $A\in{\mathcal{K}}\coloneqq\left\{0,1,\dots,K-1\right\}$ denote the action, and let $Y(a)$ denote the potential outcome under action $a$. Our goal is to choose the action with the highest conditional expected outcome, and we write the context-specific optimal action as \[a^*(x)\in\operatorname*{arg\,max}_{a\in{\mathcal{K}}}{\mathbb{E}}\left[Y(a)\mid X=x\right].\] Because vector search retrieves documents or examples that are similar or relevant to the query, we interpret the RAG system as producing an estimate of $a^*(x)$ from the context $x$. When the retrieval database contains covariates, actions, and outcomes, vector search supplies nearest neighbors used to estimate the conditional expected outcomes that determine $a^*(x)$. Similarity is therefore used to select matched evidence for each action, rather than to choose the action that appears most often among neighboring cases. This operation corresponds to nearest-neighbor matching in causal inference.
We next extend this formulation to settings in which the available action set is not specified in advance. Although the preceding formulation assumes that $K$ actions are given, the set of available actions is often not known in advance in applications of language models. For example, in workplace applications, a user may submit a query and ask what to do next without providing a candidate action set. To reflect this setting, we propose two forms of RAG-based policy learning. The one-step method receives a query and directly returns a recommended action, while the actions considered internally remain unobserved. The two-step method proceeds as follows:
The two-step method separates action generation from action choice. This separation allows us to decompose its regret into the loss from the generated candidate set and the loss from expected-outcome estimation or ranking within that set.
This study makes three contributions:
This study is related to three streams of literature: policy learning, RAG, and nonparametric analysis of neural networks and transformers.
\paragraph{Policy learning.} We study policy learning in causal inference Swaminathan2015counterfactualrisk,Kitagawa2018whoshould,Athey2021policylearning. As in supervised learning Audibert2011fast, policy learning can be approached through plug-in estimation or counterfactual risk minimization. The latter is also referred to as empirical welfare maximization (EWM) Kitagawa2018whoshould, and a theoretical comparison of the two approaches is given by Kitagawa2018whoshould. Although that comparison gives favorable results for EWM relative to the plug-in approach under its conditions, their relative performance depends on the underlying data-generating process, as discussed by Audibert2011fast. Because the structure of RAG is more directly compatible with plug-in estimation, we formulate RAG-based policy learning using the plug-in approach.
In particular, we interpret action-specific retrieval over observed cases as a nearest-neighbor matching procedure. Nearest-neighbor match counts have density-ratio limits that are proportional to inverse propensity scores, and nearest-neighbor matching can also be related to Riesz regression Lin2023estimationbased,Kato2025nearestneighbor. The analyses of Kitagawa2018whoshould and Athey2021policylearning use inverse probability weighting (IPW) or augmented IPW estimators. The representation results in Lin2023estimationbased and Kato2025nearestneighbor therefore connect RAG-based policy learning to those analyses.
\paragraph{RAG.} Our study is also related to work on RAG. Guu2020realm jointly trains a latent document retriever and a language model. Lewis2020retrievalaugmentedgeneration combines a dense passage index with a sequence-to-sequence generator and marginalizes over retrieved passages. Dense passage retrieval maps queries and passages into a shared embedding space and ranks passages by inner-product similarity Karpukhin2020densepassage. This retrieval step is computationally related to nearest-neighbor methods. In particular, the $k$-nearest-neighbor language model constructs a nonparametric next-token distribution from nearby contextual representations and interpolates it with the distribution produced by a parametric language model Khandelwal2020generalizationthrough. Later studies develop different ways to incorporate retrieved information into generation. Fusion-in-Decoder encodes retrieved passages separately and combines their representations in the decoder Izacard2021leveragingpassage. REPLUG augments a frozen black-box language model and trains the retriever using feedback from the language model Shi2024replug. Self-RAG learns when to retrieve and uses self-reflection signals to assess whether the retrieved passages support the generated response Asai2024selfrag. These methods select external evidence according to proximity in a representation space, but they do not interpret the retrieved items as matched observations under the potential outcome framework.
Retrieval has also been incorporated into sequential decision-making. Goyal2022retrievalaugmentedreinforcement augments reinforcement learning agents with direct access to datasets of past experience. Humphreys2022largescaleretrieval uses approximate nearest-neighbor search over a large collection of expert states to support offline reinforcement learning. REGENT conditions a generalist policy on demonstrations retrieved from related tasks Sridhar2025regent, whereas STRAP retrieves relevant sub-trajectories before learning a policy for a target task Memmel2024strap. These studies place retrieval within reinforcement learning or imitation learning rather than plug-in policy learning under the potential outcome framework. In a different direction, CausalRAG incorporates causal graphs into retrieval and generation for knowledge-intensive tasks Wang2025causalrag. Its use of causal information concerns relations among concepts in the retrieved documents rather than counterfactual outcomes under alternative actions.
Unlike these approaches, we formulate RAG-based action selection under the potential outcome framework. For the two-step method, vector search supplies local evidence for outcome regression, the generator estimates action-specific conditional expected outcomes or their contrasts, and the algorithm selects the action with the largest estimate. The one-step method is evaluated at the level of its returned policy because its internal criterion is not observed. The two-step formulation also separates regret due to candidate generation from regret due to expected-outcome estimation or ranking within the generated set.
\paragraph{Nonparametric analysis of neural networks and transformers.} The statistical theory of neural networks provides a basis for viewing transformers as nonparametric estimators. Nonparametric regression with deep ReLU networks has been studied under compositional assumptions and over Besov-type function classes SchmidtHieber2020nonparametric,Suzuki2019adaptivityof,Suzuki2021deeplearning. For transformers, Yun2020aretransformers establishes universal approximation results for sequence-to-sequence functions, while Takakura2023approximationand and Havrilla2024understandingscaling examine approximation and estimation for high-dimensional sequence inputs and data with low-dimensional structure.
A related literature studies in-context learning as a statistical learning procedure. Existing work examines the function classes that transformers can learn from in-context examples and relates their predictions to regression algorithms, gradient descent, and algorithm selection Garg2022whatcan,Akyurek2023whatlearning,Oswald2023transformerslearn,Bai2023transformersas. More recent studies analyze in-context learning for nonparametric regression and adaptation to low-dimensional target functions Kim2024transformersare,Oko2024pretrainedtransformer,Ching2026efficientand. We use prediction-error guarantees from this literature as inputs to the regret analysis for transformer-based expected-outcome estimation.
Let $X\in{\mathcal{X}}\subseteq\mathbb R^d$ denote covariates, which represent a query or decision context. The covariates can combine the query with pre-action individual attributes, such as age, gender, and occupation, when those attributes are available and relevant to the decision problem. Let $A\in{\mathcal{A}}$ denote an action, where the action space ${\mathcal{A}}$ can be finite or uncountably infinite. Let $Y\in{\mathcal{Y}}\subseteq{\mathbb{R}}$ denote a scalar outcome. We denote the conditional expected outcome of $Y$ given $A=a\in{\mathcal{A}}$ and $X=x\in{\mathcal{X}}$ by $f_0(a,x)\coloneqq{\mathbb{E}}\left[Y\mid A=a,X=x\right]$.
We focus on deterministic policies. A policy is a measurable function \[ \pi\colon{\mathcal{X}}\to{\mathcal{A}}. \] Let $\Pi$ be a class of such policies. For $\pi\in\Pi$, define its value by \[ V(\pi) \coloneqq {\mathbb{E}}\left[f_0(\pi(X),X)\right]. \]
We denote a learned or otherwise data-dependent policy by $\widehat a\colon{\mathcal{X}}\to{\mathcal{A}}$, so that $\widehat a(x)$ is the action selected at covariate value $x$. Its value is \[ V(\widehat a) \coloneqq {\mathbb{E}}\left[f_0(\widehat a(X),X)\right]. \] When $\widehat a$ depends on training data, the retrieval database, or auxiliary algorithmic randomness, the expectation also averages over this randomness.
For a chosen action $\widehat a(\cdot)$, we define its value by \[ V(\widehat a)\coloneqq{\mathbb{E}}\left[f_0(\widehat a(X),X)\right]. \]
We assume that the maximum of $f_0(a,x)$ over $a\in{\mathcal{A}}$ is attained and that a measurable maximizer can be selected. For each $x\in{\mathcal{X}}$, let \[ a^*(x) \in \operatorname*{arg\,max}_{a\in{\mathcal{A}}} f_0(a,x). \] The measurable selector $a^*\colon{\mathcal{X}}\to{\mathcal{A}}$ is the optimal deterministic policy. Its value is \[ V(a^*) = {\mathbb{E}}\left[\max_{a\in{\mathcal{A}}} f_0(a,X)\right]. \]
The regret of $\widehat a$ relative to the optimal action rule $a^*$ is \[ R(\widehat a) \coloneqq V(a^*)-V(\widehat a) = {\mathbb{E}}\left[ f_0(a^*(X),X)-f_0(\widehat a(X),X) \right]. \]
In this section, we describe our proposed method, called RAG-based policy learning (RAG-PL). We consider two forms of RAG-PL. The one-step method returns an action directly from the covariates. The two-step method first generates candidate actions and then selects one action from those candidates. The two-step method has two stages. The first stage constructs the candidate set. The second stage evaluates the candidates using retrieved evidence and selects one of them.
We consider three input-output forms of a RAG system. The first returns a set of candidate actions given the covariates. he second returns an estimated conditional expected outcome for each action or ranks the candidate actions by that quantity. The third returns a recommended action directly from the covariates.
\paragraph{Action-set RAG.} Let $\mathfrak F({\mathcal{A}})\coloneqq\left\{C\subseteq{\mathcal{A}}:1\leq |C|<\infty\right\}$ denote the collection of all non-empty finite subsets of ${\mathcal{A}}$. We represent an action-set RAG as a set-valued function $g\colon{\mathcal{X}}\to\mathfrak F({\mathcal{A}})$. For each $x\in{\mathcal{X}}$, the set $g(x)\subseteq{\mathcal{A}}$ contains the candidate actions returned by the RAG system.
\paragraph{Vector search as nearest-neighbor matching.} Let $\varphi\colon{\mathcal{X}}\to\mathbb R^{d_\varphi}$ be the embedding map used for vector search, and write $H=\varphi(X)$. Let $\mathcal D=\left\{D_j\right\}_{j=1}^{N_{\mathrm{db}}}$ be a retrieval database. When ${\mathcal{A}}$ is discrete and the database contains observations of covariates, actions, and outcomes, we write $D_j=(X_j,A_j,Y_j)$ and, for an action $a$, define the action-specific index set $\mathcal I_a\coloneqq\left\{j:A_j=a\right\}$. Given $x$ and $a$, let $\mathcal N_{k,a}(x)$ be the indices of the $k$ observations indexed by $\mathcal I_a$ whose embeddings are closest to $\varphi(x)$, and let $\mathcal R_k(x,a)\coloneqq\left\{D_j:j\in\mathcal N_{k,a}(x)\right\}$ denote the retrieved evidence. We assume that $|\mathcal I_a|\geq k$ for every action evaluated through action-specific retrieval. This operation matches the current context with similar cases within each action, as in nearest-neighbor matching. For a general document corpus, $\mathcal R_k(x,a)$ denotes evidence retrieved by a query that includes the candidate action. In either case, similarity determines which evidence is used to evaluate an action. It is not itself the criterion for choosing the final action, and we do not select an action by the frequency with which it appears among neighboring cases. For continuous action spaces, exact action-specific retrieval based on $\mathcal I_a=\left\{j:A_j=a\right\}$ is generally unavailable. In that case, retrieval must use an action-aware distance or kernel on ${\mathcal{A}}\times{\mathcal{X}}$, or another procedure that borrows information across nearby actions. The nearest-neighbor analysis in Section (ref) is restricted to binary actions.
\paragraph{Expected-outcome RAG or ranking RAG.} We represent an expected-outcome RAG by a function $\widehat f\colon{\mathcal{A}}\times{\mathcal{X}}\to{\mathbb{R}}$. Given an action $a\in{\mathcal{A}}$ and covariates $x\in{\mathcal{X}}$, it returns an estimate $\widehat f(a,x)$ of the conditional expected outcome $f_0(a,x)$. When the generator and retrieved evidence are written explicitly, we use \[ \widehat f(a,x) = G_\theta\bigl(x,a,\mathcal R_k(x,a)\bigr). \] A ranking RAG instead receives $x$ and a finite candidate set $C=\left\{a_1,\dots,a_m\right\}\in\mathfrak F({\mathcal{A}})$. It retrieves evidence for the candidate actions and returns an ordered tuple $r(x,C)=(a_{(1)},\dots,a_{(m)})$, which is a permutation of the actions in $C$ from the highest to the lowest estimated conditional expected outcome.
\paragraph{Policy RAG.} A policy RAG receives covariates and directly returns a recommended action. We represent it by a function $\pi_{\mathrm{RAG}}\colon{\mathcal{X}}\to{\mathcal{A}}$.
The one-step RAG-PL method applies a policy RAG directly. Given $x\in{\mathcal{X}}$, the RAG system is instructed to recommend an action $\widehat a_{\mathrm{one}}(x)$ with the highest conditional expected outcome, and we write $\widehat a_{\mathrm{one}}(x)=\pi_{\mathrm{RAG}}(x)$. The policy implemented by the RAG system is usually not observed separately from its output. In particular, the one-step method does not expose an intermediate candidate set, expected outcomes for the candidate actions, or a separate ranking step. We therefore evaluate the returned policy without imposing a particular interpretation on its internal computation.
\paragraph{Action-set generation.} Given $x\in{\mathcal{X}}$, the action-set RAG returns a finite candidate set \[ g(x)=\left\{a_1(x),\dots,a_{M(x)}(x)\right\}\subseteq{\mathcal{A}}, \qquad 1\leq M(x)<\infty. \] Thus, even when ${\mathcal{A}}$ is uncountably infinite, the subsequent evaluation or ranking is performed only over the finite set $g(x)$. The candidate set can contain actions found in matched cases and actions generated by the language model. Its purpose is to include actions with high conditional expected outcomes, rather than to reproduce the empirical distribution of actions in the retrieved evidence.
\paragraph{Ranking or expected-outcome estimation.} For every $a\in g(x)$, the RAG system retrieves $\mathcal R_k(x,a)$. When an expected-outcome RAG is used, it evaluates $\widehat f(a,x)=G_\theta(x,a,\mathcal R_k(x,a))$. When a ranking RAG is used, it returns the ordered tuple $r(x,g(x))=(a_{(1)}(x),\dots,a_{(M(x))}(x))$, where $a_{(1)}(x)$ is the candidate with the highest estimated conditional expected outcome.
\paragraph{Action choice.} Under expected-outcome estimation, we choose \[ \widehat a_{\mathrm{two}}(x)\in\operatorname*{arg\,max}_{a\in g(x)}\widehat f(a,x), \] where a fixed rule is used to break ties. Under ranking, we choose \[ \widehat a_{\mathrm{two}}(x)\coloneqq a_{(1)}(x). \] Both methods define a policy from ${\mathcal{X}}$ to ${\mathcal{A}}$. The two-step method makes action generation, retrieval of similar cases, and action choice explicit, which allows their contributions to regret to be studied separately.
We analyze the regret of the one-step and two-step methods under the setup above. We write $R(\widehat a)$ for the regret of a selected action $\widehat a(\cdot)$. When the candidate set, the expected-outcome estimate, or the selected action is random, the expectation defining $R(\widehat a)$ and all expectations below also average over this randomness. We first separate the effect of the generated candidate set from the effect of selecting an action within that set. We then connect the second term to nonparametric prediction error. Section (ref) studies the case in which vector search is represented explicitly by nearest-neighbor matching.
For each $x\in{\mathcal{X}}$, let \[ a_g^*(x)\in\operatorname*{arg\,max}_{a\in g(x)}f_0(a,x), \] where the fixed tie-breaking rule is used when needed, and we assume that the resulting selector is measurable. The maximum is attained because $g(x)$ is finite. We define \[ R_{\mathrm{gen}}(g) \coloneqq {\mathbb{E}}\left[f_0(a^*(X),X)-f_0(a_g^*(X),X)\right] \] and \[ R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g) \coloneqq {\mathbb{E}}\left[f_0(a_g^*(X),X)-f_0(\widehat a_{\mathrm{two}}(X),X)\right]. \] The first term is the regret caused by restricting the choice to $g(X)$. The second is the regret caused by expected-outcome estimation or ranking within $g(X)$. We call $R_{\mathrm{gen}}(g)$ the candidate-set regret and $R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g)$ the within-candidate regret.
Adding and subtracting $f_0(a_g^*(X),X)$ gives
This identity does not require independence between candidate generation and action choice.
For $\varepsilon\geq0$, define the set of $\varepsilon$-optimal actions by \[ {\mathcal{A}}_\varepsilon^*(x) \coloneqq \left\{a\in{\mathcal{A}}:f_0(a^*(x),x)-f_0(a,x)\leq\varepsilon\right\}. \] Exact inclusion of $a^*(x)$ is not required. It is enough that the generated set contain an action whose conditional expected outcome is close to the optimum.
The third statement controls the value of the best action in $g(x)$ through its distance from $a^*(x)$. It does not concern the estimation error of $\widehat f$. The smoothness condition on $f_0$ converts distance in the action space into a difference in conditional expected outcomes.
Suppose that the two-step method chooses \[ \widehat a_{\mathrm{two}}(x) \in \operatorname*{arg\,max}_{a\in g(x)}\widehat f(a,x). \] For each $x\in{\mathcal{X}}$, define \[ \delta_f(x;g) \coloneqq \max_{a\in g(x)}|\widehat f(a,x)-f_0(a,x)|. \] Let $v_g^*(x)\coloneqq\max_{a\in g(x)}f_0(a,x)$ and define the smallest positive gap within the candidate set by \[ \Delta_g(x) \coloneqq \min_{a\in g(x):f_0(a,x)<v_g^*(x)} \left\{v_g^*(x)-f_0(a,x)\right\}, \] where $\Delta_g(x)=+\infty$ if every action in $g(x)$ attains $v_g^*(x)$.
The theorem is stated for the realized set $g(X)$ and therefore does not require candidate generation to be independent of $\widehat f$. A high-probability error bound can be used in the same way. If $0\leq f_0(a_g^*(X),X)-f_0(a,X)\leq B_f$ holds almost surely for every $a\in g(X)$, the margin condition holds, and \[ \Pr\left(\delta_f(X;g)>r_f\right)\leq\eta_f, \] then
For example, suppose that $|g(X)|\leq M$ and, conditional on $X$ and $g(X)$, \[ \Pr\left( |\widehat f(a,X)-f_0(a,X)|>b_n+u \mathrel{\Big|} X,g(X) \right) \leq c_1\exp(-c_2a_nu^2) \] for every $a\in g(X)$. A union bound gives \[ \Pr\left(\delta_f(X;g)>b_n+u\right) \leq c_1M\exp(-c_2a_nu^2). \] Thus, the number of candidates enters the simultaneous error through a logarithmic term when $u$ is chosen to make the right-hand side small.
For a ranking RAG, let $a_{(1)}(x)$ be the first action in $r(x,g(x))$. Then, \[ R(a_{(1)}) = R_{\mathrm{gen}}(g) + {\mathbb{E}}\left[f_0(a_g^*(X),X)-f_0(a_{(1)}(X),X)\right]. \] When the ranking orders the candidates according to $\widehat f(a,x)$, Theorem (ref) applies. If the RAG system returns only an ordering, its analysis requires a direct bound on the true conditional expected-outcome difference between the best candidate and the first-ranked candidate.
Let \[ \mathcal S_g \coloneqq \left\{s\colon{\mathcal{X}}\to{\mathcal{A}}:s\text{ is measurable and }s(x)\in g(x)\right\} \] and define the value induced by $\widehat f$ as \[ \widehat V_f(s) \coloneqq {\mathbb{E}}\left[\widehat f(s(X),X)\right]. \] If a measurable selector $\widehat s_f$ satisfies \[ \widehat s_f(x) \in \operatorname*{arg\,max}_{a\in g(x)}\widehat f(a,x), \] then, for every $s\in\mathcal S_g$, it holds that $\widehat f(\widehat s_f(x),x)\geq\widehat f(s(x),x)$. Hence, \[ \widehat s_f \in \operatorname*{arg\,max}_{s\in\mathcal S_g}\widehat V_f(s). \] The same pointwise argument applies to an empirical average over target contexts. Thus, when the policy class contains every measurable selector and its value is constructed from $\widehat f$, pointwise maximization and maximization over the policy class are the same optimization problem. They can differ when the policy class is restricted or when the value estimator uses observed outcomes through IPW or AIPW rather than being constructed from $\widehat f$ alone.
The preceding results allow a prediction rate to be inserted into the regret bound without changing the RAG-PL algorithm. Suppose that $R_{\mathrm{gen}}(g)\leq r_{g,n}$ and that the expected-outcome error on the generated set is bounded by $r_{f,n}$. Under the finite-candidate margin condition in Theorem (ref), we have \[ R(\widehat a_{\mathrm{two}}) \leq r_{g,n}+C_M(2r_{f,n})^{1+\kappa}. \] For example, if a nonparametric estimator has a uniform error bound $r_{f,n}\asymp n^{-\alpha/(2\alpha+d)}$, the second term is of order $n^{-\alpha(1+\kappa)/(2\alpha+d)}$. This substitution is valid only when the available prediction result controls the error required by the regret theorem. An integrated mean-squared error cannot be used as a uniform error without an additional argument.
The same distinction matters for transformers. The analysis of Kim2024transformersare gives a mean-squared prediction bound for a transformer composed of a learned neural representation and a linear-attention layer. In the Besov setting, a representative bound has the form \[ q_{\mathrm{ICL}} \lesssim N_{\mathrm{rep}}^{-2\alpha/d} + \frac{N_{\mathrm{rep}}\log N_{\mathrm{rep}}}{n_{\mathrm{ctx}}} + \frac{N_{\mathrm{rep}}^2\log N_{\mathrm{rep}}}{T_{\mathrm{pre}}}, \] where $N_{\mathrm{rep}}$ is the representation dimension, $n_{\mathrm{ctx}}$ is the number of in-context examples, and $T_{\mathrm{pre}}$ is the number of pretraining tasks. If this bound holds for both binary expected-outcome estimates, Theorem (ref) below gives \[ R(\widehat a) \lesssim (q_{\mathrm{ICL},0}+q_{\mathrm{ICL},1})^{(1+\kappa)/(2+\kappa)}. \] A sharper exponent $1+\kappa$ for the root prediction rate requires a uniform error bound or a pointwise exponential-deviation bound. It does not follow from an MSE result alone. The results of Oko2024pretrainedtransformer and Ching2026efficientand further describe adaptation to low-dimensional target functions and minimax MSE rates, and they enter the regret analysis according to the same distinction between error types.
The general regret decomposition itself has no sample size. For a specified retrieval database, its size and the number of matched observations have a direct statistical meaning. Transformer theory uses different quantities for the representation dimension, in-context examples, and pretraining tasks. For an already trained RAG system whose training data are not observed, the total number of pretraining tokens is not by itself a sample size for conditional expected-outcome estimation. In that case, the regret bound should be stated in terms of a prediction error that is either assumed or evaluated separately.
The regret of the one-step output is \[ R(\widehat a_{\mathrm{one}}) = {\mathbb{E}}\left[f_0(a^*(X),X)-f_0(\widehat a_{\mathrm{one}}(X),X)\right]. \] For comparison, its observed output can be represented by the singleton set $g_{\mathrm{one}}(x)=\left\{\widehat a_{\mathrm{one}}(x)\right\}$. This representation does not assert that the RAG system uses a singleton set internally.
Suppose that $\widehat a_{\mathrm{one}}(x)\in g(x)$ for $P_X$-almost every $x$. Because $a_g^*(x)$ is at least as good as every action in $g(x)$, we have \[ R_{\mathrm{gen}}(g) \leq R(\widehat a_{\mathrm{one}}). \] Combining this inequality with (ref) gives \[ R(\widehat a_{\mathrm{two}}) \leq R(\widehat a_{\mathrm{one}}) + R_{\mathrm{choice}}(\widehat a_{\mathrm{two}};g). \] If the two-step method selects the best action in $g(x)$ according to $f_0$, its regret is no larger than the regret of the one-step output. With an estimated expected outcome, Theorem (ref) instead gives \[ R(\widehat a_{\mathrm{two}}) \leq R(\widehat a_{\mathrm{one}}) + 2{\mathbb{E}}\left[\delta_f(X;g)\right]. \] Thus, a larger candidate set can contain a better action than the one-step output, but the gain can be lost when expected-outcome estimation or ranking is inaccurate.
When the action set is fixed and finite, we can set $g(x)={\mathcal{A}}={\mathcal{K}}$ for every $x$. Then, $R_{\mathrm{gen}}(g)=0$, and Theorem (ref) gives \[ R(\widehat a_{\mathrm{two}}) \leq 2{\mathbb{E}}\left[\max_{a\in{\mathcal{K}}}|\widehat f(a,X)-f_0(a,X)|\right]. \]
For the binary action set ${\mathcal{K}}=\left\{0,1\right\}$, define \[ \tau_0(x) \coloneqq f_0(1,x)-f_0(0,x), \qquad \widehat\tau(x) \coloneqq \widehat f(1,x)-\widehat f(0,x). \] Using the tie-breaking rule that selects action $1$ at equality, we have $a^*(x)=\mathbbm{1}\left[\tau_0(x)\geq0\right]$ and $\widehat a(x)=\mathbbm{1}\left[\widehat\tau(x)\geq0\right]$. For any selected action $\widehat a(x)\in\left\{0,1\right\}$,
When $f_0(a,x)={\mathbb{E}}\left[Y(a)\mid X=x\right]$, $\tau_0(x)$ is the conditional average treatment effect.
We use the margin condition
It controls the probability of contexts for which the two actions have nearly equal conditional expected outcomes Audibert2011fast.
Vector search connects RAG-PL with causal inference by selecting observations whose covariates are close to the current context. This section gives an explicit nearest-neighbor regression model in which the expected-outcome estimate in Section (ref) is formed from matched observations and its error can be substituted into the regret bounds above.
For each $a\in\left\{0,1\right\}$, let \[ \mathcal D_a = \left\{(H_{a,j},Y_{a,j})\right\}_{j=1}^{N_a} \] be a database for action $a$, where $H_{a,1},\dots,H_{a,N_a}$ are independently and identically distributed according to $Q_a$ on $\mathbb R^{d_\varphi}$ and \[ Y_{a,j}=m_a(H_{a,j})+\xi_{a,j}. \] We assume that the query $X$ is independent of the action-specific databases. For a query $x$, write $h=\varphi(x)$ and let $\mathcal N_{k,a}(h)$ be the indices of the $k$ nearest embeddings in $\mathcal D_a$. Define the matched expected-outcome estimate by \[ \widetilde f_k(a,x) \coloneqq \frac{1}{k}\sum_{j\in\mathcal N_{k,a}(h)}Y_{a,j}. \] The following result gives a standard bias and variance calculation for nearest-neighbor regression in the notation used by RAG-PL Jiang2019nonasymptoticuniform. Its use as an action-specific matching estimator follows the setup of Abadie2006largesample.
The constant $C$ depends only on the constants in the assumptions.
The term $b_{\varphi,a}$ is zero when the embedding retains all information needed for the conditional expected outcome, so $f_0(a,x)=m_a(\varphi(x))$. Otherwise, it records the difference between conditioning on $X$ and conditioning on its embedding. The local-mass condition requires enough observations under action $a$ near each query. It complements the positivity condition used for causal identification.
The displayed estimator is a benchmark for the expected-outcome RAG. Suppose that the RAG output satisfies \[ \max_{a\in\left\{0,1\right\}}\sup_{x\in{\mathcal{X}}} |\widehat f(a,x)-\widetilde f_k(a,x)| \leq \eta_n. \] Suppose that the margin condition (ref) holds. For the deviation-based bound below, also suppose that $|\tau_0(X)|\leq B_\tau$ holds almost surely. Let $N_{\min}=\min\left\{N_0,N_1\right\}$ and $b_\varphi=b_{\varphi,0}+b_{\varphi,1}$. Combining Theorem (ref) with Theorem (ref) gives
provided that the term in parentheses is sufficiently small relative to $t_0$. The MSE statement gives the alternative bound
Choosing $k\asymp N_{\min}^{2\beta/(2\beta+d_\varphi)}$ balances the nearest-neighbor bias and variance. In (ref), this gives the root prediction rate $N_{\min}^{-\beta/(2\beta+d_\varphi)}$ before the margin exponent is applied. In (ref), the term $\exp(-k/4)$ is absorbed into $1/k$.
For a general document corpus without observed outcomes, Theorem (ref) does not apply directly. In that case, the theorem describes the matching calculation that the RAG output is meant to approximate, and a regret rate requires an assumption or an evaluation of the difference represented by $\eta_n$.
Similarity is used to select comparable evidence, not to choose the action that appears most often in nearby cases. Let $\mathcal N_k^H(h)$ denote the indices of the $k$ nearest observations to $h$ in the pooled embedding database. If \[ \widehat e_k^H(a\mid h) = \frac{1}{k}\sum_{j\in\mathcal N_k^H(h)} \mathbbm{1}\left[A_j=a\right], \] then $\widehat e_k^H(a\mid h)$ estimates $\Pr(A=a\mid H=h)$ under standard nearest-neighbor conditions. At $h=\varphi(x)$, this quantity equals $\Pr(A=a\mid X=x)$ only under an additional condition such as $A\perp X\mid H$. Maximizing this quantity selects the most common past action near $x$, not the action with the largest conditional expected outcome. The joint distribution of $(X,A)$ can remain unchanged while the conditional outcomes are changed so that either action is optimal. Thus, observations of contexts and actions alone do not determine the optimal policy.
Nearest-neighbor matching also has a weighting representation. Let $X_1^{\mathrm{tar}},\dots,X_m^{\mathrm{tar}}$ be target contexts and define \[ \widetilde f_k(a,X_i^{\mathrm{tar}}) = \frac{1}{k}\sum^{N_a}_{j=1} \mathbbm{1}\left[j\in\mathcal N_{k,a}(\varphi(X_i^{\mathrm{tar}}))\right]Y_{a,j}. \] Define \[ K_{a,j} \coloneqq \sum^m_{i=1}\mathbbm{1}\left[j\in\mathcal N_{k,a}(\varphi(X_i^{\mathrm{tar}}))\right]. \] Then, changing the order of summation gives \[ \frac{1}{m}\sum^m_{i=1}\widetilde f_k(a,X_i^{\mathrm{tar}}) = \sum^{N_a}_{j=1}\frac{K_{a,j}}{mk}Y_{a,j}. \] The match count therefore acts as a weight when local predictions are averaged over target contexts. For the propensity-score interpretation, suppose that the action-specific databases are sampled from the same observational population, so $Q_a=P_{H\mid A=a}$, and that the target embeddings follow $P_H$. Bayes' rule then gives \[ \frac{p_H(h)}{p_{H\mid A=a}(h)} = \frac{\Pr(A=a)}{\Pr(A=a\mid H=h)}. \] The density ratio is therefore proportional to an inverse propensity score defined conditional on the embedding. The density ratio is proportional to an inverse propensity score. Lin2023estimationbased studies the density-ratio limit of nearest-neighbor match counts, and Kato2025nearestneighbor relates nearest-neighbor matching to least-squares density-ratio estimation and Riesz regression. These results connect matching with IPW and AIPW representations when outcomes are available. They do not make the propensity score an expected outcome.
The regret analysis is written in terms of $f_0(a,x)={\mathbb{E}}\left[Y\mid A=a,X=x\right]$. A causal interpretation requires consistency, conditional exchangeability, and positivity. Under these conditions, it holds that \[ {\mathbb{E}}\left[Y\mid A=a,X=x\right] = {\mathbb{E}}\left[Y(a)\mid X=x\right]. \] The information contained in $X$ is therefore part of the identification argument. It is natural to write $X=(Q,W)$, where $Q$ is the query and $W$ contains variables observed before the action, such as age, gender, occupation, and prior decisions. If a variable affects both past action selection and the outcome, omitting it can prevent the conditional comparison from having a causal interpretation. Including more variables does not establish conditional exchangeability by itself, and post-action variables should not be used as controls.
Vector search is usually performed on an embedding $H=\varphi(X)$. Matching on $H$ is sufficient only when the representation retains the information required for adjustment or for the conditional expected outcomes. The GenAI-Powered Inference framework combines structured covariates with low-dimensional features extracted from unstructured data Imai2025genai. This suggests keeping personal attributes explicit while representing the query text through an embedding. Embedding-powered BISG provides an example in which embeddings are used to estimate a missing personal attribute probabilistically Dasanaike2026usingembedding. Such an estimate is a probabilistic proxy rather than an observed covariate, so its estimation error and possible distribution shift remain relevant. The variables used for adjustment can also be richer than the variables allowed in the final policy. In that case, expected outcomes are first adjusted using the larger information set and then averaged over variables excluded from the policy.
Instrumental variables provide another identification strategy when conditional exchangeability fails in the data used to estimate expected outcomes. They enter through the source-data model, not by treating the RAG output as an instrument. Existing studies consider optimal treatment regimes under additional IV assumptions Cui2021asemiparametric,Qiu2021optimalindividualized and decision rules under partial identification Pu2021estimatingoptimal. These problems have different target values. Applying RAG-PL to them would require the expected outcome supplied to the RAG system to be defined for the corresponding target.
When $\widehat f(a,x)$ is available and the policy class contains all measurable selectors, maximizing $\widehat f(a,x)$ at each context is equivalent to maximizing ${\mathbb{E}}\left[\widehat f(s(X),X)\right]$ over the policy class. The same equivalence holds for an empirical average over target contexts. A restricted policy class can introduce an approximation loss, but the criterion is still constructed from the same expected-outcome estimate. The relation between EWM and least-squares CATE estimation for a reparameterized binary policy class is studied by Kato2025bridginggap.
A different analysis is needed when the policy value is constructed from observed outcomes rather than from $\widehat f$ alone. IPW uses weighted outcomes, and AIPW combines an expected-outcome estimate with an outcome-residual correction Kitagawa2018whoshould,Athey2021policylearning. In the main setting of this study, the analyst starts from a trained RAG system and may not observe unit-level outcomes and assignment probabilities. The fact that outcome-related text may have appeared during pretraining does not provide the data required to construct IPW or AIPW. If a retrieval database contains unit-level observations of $(X,A,Y)$, IPW and AIPW can instead be constructed directly from those data. For one-step RAG-PL, the internal criterion is not observed, so we evaluate the returned policy without assigning it to one of these formulations.
The margin condition is separate from identification. Identification determines whether the contrast has a causal interpretation, while the margin condition controls how often estimation error changes the preferred action. In the binary case, it bounds the probability that $|\tau_0(X)|$ is close to zero. When few contexts lie near this boundary, the regret decreases faster than the expected-outcome error.
The margin condition depends on the information used to define the conditional expected outcomes. Omitting a relevant attribute can average positive and negative effects into a contrast near zero, while adding an attribute can also reveal finer effects that lie near zero. Replacing $X$ by an embedding can change the margin for the same reason. For two-step RAG-PL, the margin within $g(x)$ is also distinct from the quality of the generated set. A large gap among poor candidates does not remove $R_{\mathrm{gen}}(g)$.
When ${\mathcal{A}}$ is finite and given, $R_{\mathrm{gen}}(g)=0$. For a large or uncountable action space, the generated set determines which actions can be compared. Exact inclusion of $a^*(x)$ is stronger than necessary; Proposition (ref) uses an $\varepsilon$-optimal action, a coverage probability, or distance in the action space. These conditions describe the candidate set through the best conditional expected outcome it contains.
Increasing the number of candidates can improve the chance of including a good action, but more expected outcomes must then be estimated or ranked. Equation (ref) keeps these effects separate. A useful candidate-set size depends on both the quality of generation and the accuracy of the second step.
We evaluate whether action-specific retrieval of cases with similar pre-action information improves policy choice. Each data-generating process first produces a structured pre-action state together with the conditional expected outcome under every available action. We convert the structured states into complete English queries and case reports before executing the notebooks. The RAG procedures observe only the resulting texts, while the numerical state is retained to calculate the optimal action and regret. Because each numerical level is mapped deterministically to a verbal category, the text retains the full discrete state information. Appendix (ref) gives the full specifications.
\paragraph{Data-generating processes.} DGP 1 and DGP 2 use two actions. Action 0 is temporary workload reduction and action 1 is weekly individual coaching. Let $X=(X_1,\ldots,X_6)$ describe recent performance; workload; task complexity; experience; schedule flexibility; and team support. Each component is independently and uniformly distributed on $\left\{-2,-1,0,1,2\right\}$. DGP 1 has a linear conditional effect and observational action assignment. DGP 2 has randomized action assignment and a smooth nonlinear conditional effect whose population mean is zero. The preferred action therefore varies with the current state even though neither constant action is favored on average. DGP 3 contains 24 interventions formed from four intervention types with three intensity levels and two durations. Its conditional expected outcome rewards agreement between the employee's stated needs and the intervention profile.
\paragraph{Methods and implementation.} The main results use GPT-5.4 mini with temperature zero. Retrieval vectors are obtained from text-embedding-3-small. The embedding input contains only a fixed-order description of the pre-action information. It excludes the historical action and the observed outcome. It also excludes the decision question. Cosine distance is used after normalization.
In DGP 1 and DGP 2, One-step RAG-PL and Two-step RAG-PL receive the same current query and the same six nearest reports under each action. The twelve reports are presented in alternating action order. The prompt also gives the exact mean of the six displayed scores under each action. Providing these means allows the model to use exact local sample averages, so the experiment does not test whether it can calculate those averages from the reports. One-step RAG-PL returns an action directly. Two-step RAG-PL returns the two action-specific conditional expected outcomes, and the experiment code selects the larger estimate. RAG without covariates receives neither the current profile nor historical pre-action profiles. It uses six randomly sampled action-and-outcome reports under each action and returns one population-level decision for the retrieval database. That decision is applied to all ten test queries in the repetition. RAG without covariates is therefore a population-level baseline that differs from RAG-PL in both the current-state information supplied and the retrieval rule. In DGP 3, both RAG-PL procedures receive the same twelve nearest reports and the complete catalog of 24 actions. One-step RAG-PL returns one action directly. Two-step RAG-PL generates five distinct candidate actions. It retrieves five reports observed under each candidate and estimates the five candidate-specific conditional expected outcomes. It then selects the candidate with the largest estimate.
\paragraph{Evaluation.} Each DGP is repeated 20 times. Each repetition contains ten test queries. DGP 1 and DGP 2 use 500 historical reports, while DGP 3 uses 600. For each repetition, we first average regret and policy value over its ten queries. Table (ref) reports the mean of these repetition-level quantities. The regret standard error is calculated across the 20 repetitions. The optimal-action probability counts any action attaining the largest true conditional expected outcome as optimal. Figures (ref)--(ref) show the repetition-level mean regret.
\paragraph{Results.} In DGP 1, One-step RAG-PL and Two-step RAG-PL both reduce mean regret from $4.5450$ under RAG without covariates to $3.3700$. The paired mean-regret reduction is $1.1750$ for each procedure, with standard errors of $0.3562$ for One-step RAG-PL and $0.3635$ for Two-step RAG-PL. The optimal-action probability rises from $0.4100$ under RAG without covariates to $0.5500$ under One-step RAG-PL and $0.5550$ under Two-step RAG-PL.
DGP 2 isolates personalization because historical actions are randomized. One-step RAG-PL has mean regret $1.5676$ and Two-step RAG-PL has mean regret $1.6071$, compared with $2.3638$ for RAG without covariates. The paired reductions relative to RAG without covariates are $0.7963$ with standard error $0.2257$ and $0.7567$ with standard error $0.2473$, respectively. Both RAG-PL procedures select an optimal action with probability $0.6400$, compared with $0.5100$ under RAG without covariates.
In DGP 3, Two-step RAG-PL lowers mean regret from $1.4985$ to $1.2154$. The paired difference, Two-step RAG-PL minus One-step RAG-PL, is $-0.2831$ with standard error $0.1431$. The regret decomposition for Two-step RAG-PL gives candidate-set regret $0.4822$ and within-candidate regret $0.7332$. Thus, both candidate generation and selection among the generated candidates contribute to the remaining regret. The generated set contains an exactly optimal action with probability $0.4450$ and an action within $0.5$ of the optimum with probability $0.6700$. Additional estimation and robustness results are reported in Appendix (ref).
We proposed one-step and two-step RAG-based policy learning methods. We interpreted the vector-search component of RAG as nearest-neighbor matching and related this matching procedure to propensity-score weighting, density-ratio estimation, and Riesz regression. This interpretation connects RAG-based decision-making to the existing literature on policy learning. We then derived regret upper bounds for the proposed methods by combining policy-learning analysis with nonparametric prediction-error bounds for nearest-neighbor estimators and transformers. In the simulation studies, both RAG-PL methods had lower mean regret than RAG without covariates in the binary settings, and the two-step method had lower mean regret than the one-step method in the 24-action setting.
\onecolumn