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.
72,439 characters · 16 sections · 62 citation commands
Distributionally Robust Treatment Effect
In empirical economics, causal analysis serves two distinct but related objectives. Retrospective studies primarily address internal validity, focusing on the identification of causal effects within a given sample. In contrast, prospective policy analysis concerns external validity, requiring extrapolation of treatment effects to new populations, locations, or time periods. The latter problem is inherently more challenging, as the counterfactual distribution of outcomes under alternative environments is unobserved.
Most empirical work evaluates retrospective causal effects and implicitly treats these estimates as informative for future policy decisions. Such extrapolation, however, relies on strong assumptions about the stability or exchangeability of the underlying population distribution. In many settings, these assumptions are difficult to justify: populations may differ systematically across locations or evolve over time, rendering the original sample unrepresentative of the policy-relevant target population. As a result, external validity remains a central and unresolved challenge.
This paper studies out-of-population prediction of treatment effects that is robust to distributional shift between the source population and an unobserved target population. Our approach contributes to the growing literature on transfer learning, but departs from existing methods by requiring minimal information about the target population. Essentially, our estimator can be used to formulate a prediction of the individual treatment effect for the same treatment/policy implemented in another location or time period with only retrospective data from a single sample. We consider this setting as a common scenario in practice.
In this paper, we consider the case where there are only outcome variables and treatment status but no covariates.\footnote{Extensions incorporating covariates are left for future work.} In other words, we study how to generalize the findings from randomized experiments. In some cases, there is no universally agreed-upon target population, or the baseline covariates that can be collected from the target population are very limited. As a result, we find ourselves in a situation with limited information, where little is observed from the target population. This scenario arises when we consider expanding the same program to a different location without collecting the necessary data or conducting a cost-benefit analysis to determine whether the policy should be continued in the near future. Instead, we construct an ambiguity/uncertainty set centered around the source/reference distribution—the population distribution from which our sample is drawn. By choosing the radius of the ambiguity set wisely, our hope is that the target population is included in the class of distributions within the neighborhood of the source distribution.
Formally, we consider a distributionally robust optimization (DRO) problem that minimizes the worst-case mean squared error (MSE) of treatment effect prediction over all distributions within a Wasserstein ball centered on the source distribution. Our objective function is conservative here, as we have limited information on the target distribution. The distribution within the ambiguity set that leads to the largest MSE is considered the least favorable distribution. Within the DRO literature, there are many ways to measure the distance between distributions in the ambiguity set. Our use of the Wasserstein distance is motivated by its flexibility: unlike $\phi$-divergences with the Kullback–Leibler divergence serving as a leading example, it does not require the source and target distributions to share common support and admits a natural metric interpretation. See Section S.3.1 in gu2024wasserstein for further comparisons between Wasserstein distance and $\phi$-divergences. To maintain tractability while defining the Wasserstein neighborhood, we focus on continuous outcome variables in this paper.
Our choice of the quadratic loss function involves the joint distribution of two potential outcomes. Nevertheless, due to the fundamental missing data problem in the potential outcomes framework, the joint distribution of the potential outcomes remains unidentified but lies within the Fr\'echet class of distributions consistent with the observed marginals. By Sklar’s theorem, this class can be equivalently represented by the set of all copulas linking the marginal distributions of the two potential outcomes. This leads to partial identification. As a result, we select the worst and best distributions within the copula set as our optimistic and pessimistic optimization objects. We pick the current pair of loss function and ambiguity set to balance the goals of generality, tractability, and non-trivial solutions.
We need to solve a modified three-layer minimax optimization problem coupled with partial identification. The inner maximization primal problem with respect to the ambiguity set can be transformed into a dual problem by minimizing a penalized MSE. Due to the partial identification of the joint distribution of potential outcomes, we derive sharp upper and lower bounds for the minimax optimizer using the Fr\'echet-Hoeffding inequality.
Our analysis yields several economically interpretable results. The robust predictor preserves the sign of the conventional estimator -- the average treatment effect (ATE) under the source distribution -- but shrinks its magnitude toward zero. This shrinkage reflects a precautionary adjustment for distributional uncertainty. Importantly, the extent of shrinkage depends on the degree of treatment effect heterogeneity. When treatment effects are homogeneous, we see a delayed shrinkage. Namely, the robust predictor coincides with the naive ATE within a certain neighborhood of the source distribution, implying no adjustment for small distributional shifts. Only when the target distribution is sufficiently different from the source distribution does the best-predicted treatment effect begin to shrink toward zero. By contrast, under heterogeneous treatment effects, even small deviations from the source distribution induce immediate shrinkage. These patterns align with the intuition that heterogeneity amplifies sensitivity to distributional changes.
In Section (ref), we propose bound estimators based on an M-estimation formulation of the dual problem and establish its consistency and asymptotic distribution. To conduct inference on the partially identified minimax parameter, we develop a two-step procedure that combines ideas from the partial identification literature (see imbens2004confidence and stoye2009more) with Bonferroni-type corrections. In Section (ref), we discuss practical considerations for selecting the radius of the Wasserstein ball, which governs the degree of robustness. Monte Carlo simulations and synthetic data illustrate the finite-sample performance of our method in Section (ref). Section (ref) concludes.
Our paper is related to two strands of literature. The first is on transfer learning and external validity. A seminal contribution is hotz2005predicting, which provides a framework for extrapolating causal effects across populations. In addition to standard identification assumptions for internal validity, they impose a form of locational unconfoundedness, under which differences across locations arise solely from shifts in covariates, while the conditional distribution of the potential outcomes remains invariant. Under this assumption, treatment effects in the target population can be recovered by reweighting conditional average treatment effects using the target covariate distribution. This approach -- often referred to as covariate shifts -- has been extended in subsequent work, including spini2021robustness, huang2023leveraging, jin2024tailored, huang2024sensitivity, and menzel2023transfer. A growing body of evidence, however, suggests that such assumptions are restrictive in practice. As emphasized by allcott2015site and jin2025beyond, unobserved differences across locations may invalidate conditional exchangeability. In contrast to this literature, we allow for distributional shifts in potential outcomes, rather than restricting attention to covariate shifts alone.
Recent work has begun to relax the location unconfoundedness assumption by allowing for richer forms of distributional change. For example, guo2024statistical and zhang2024minimax consider settings in which the conditional distribution of outcomes in the target population lies within a convex combination of conditional distributions observed across multiple sites. Similarly, jeong2024out model site-level heterogeneity as random perturbations around a common target distribution. These approaches leverage multisite data or partial information about the target population to reduce uncertainty about the target distribution. By contrast, our setting is intentionally more limited: we consider a single observed sample and allow for minimal or no information about the target population, a scenario that is common when extrapolating to new environments or future periods. This distinction leads us to adopt a different approach to modeling uncertainty. We use a Wasserstein neighborhood instead of a linear combination of source sites.
The second strand of literature is distributionally robust optimization (DRO). DRO has been widely studied across operations research, statistics, and machine learning; see, for example, blanchet2019quantifying, duchi2021learning, gao2023distributionally, and fan2025quantifying. More recently, DRO devices have been incorporated into econometric applications. bertsimas2022distributionally use DRO to conduct sensitivity analysis with respect to unobserved confounding. qu2024distributionally develop distributionally robust instrumental variables estimator that is resilient to weak or invalid instruments in finite samples. chen2024robust apply DRO to relax rational expectations in moment restrictions. In addition, DRO has been used to study robust policy learning; see, for example, mo2021learning, adjaho2022externally, kido2022distributionally, and lei2023policy. Other work, such as christensen2023counterfactual and gu2024wasserstein, employs DRO to assess the sensitivity of counterfactuals to parametric assumptions about the latent variable distribution in a class of structural models.
Our contribution differs from existing DRO-based approaches in several dimensions. First, unlike the literature on individualized policy learning, which typically evaluates the robustness of treatment rules, we focus on predicting treatment effects under distributional shift. This distinction leads to a different objective function: we minimize the worst-case mean squared error of treatment effect prediction, rather than optimizing policy performance. Second, our analysis explicitly addresses partial identification arising from the unobserved joint distribution of potential outcomes, a feature that is largely absent in existing DRO applications. Third, we provide a formal inference procedure for our proposed estimator, while inference is overlooked in most DRO literature.
Suppose we have access to a sample randomly drawn from a single cross section of the source/reference population distribution. In the sample, we observe a binary treatment variable $T_i$ and a realized outcome variable $Y_i=T_iY_i(1)+(1-T_i)Y_i(0)$, where $Y_i(1)$ and $Y_i(0)$ denote a pair of potential outcomes. Let us denote the joint distribution of $(Y(1), Y(0))$ from the source population by $P$. We are nevertheless interested in making inference about the treatment effect in a target distribution $Q$ of potential outcomes $(\tilde{Y}(1), \tilde{Y}(0))$.\footnote{We use $(\tilde{Y}(1),\tilde{Y}(0))$ to denote the potential outcomes under distribution $Q$ to differentiate it from $P$.} The distribution $Q$ can be different from the source distribution $P$ because it is from a different location or a future time period. We do not observe the outcome variables from $Q$ and probably not even the covariates. Hence, the treatment effect under $Q$ is not identified.
For instance, we collect data from a job training program, given that participation is randomly assigned, we can identify $\tau^*=\mathbb{E}_P[Y(1)-Y(0)]$ for the source distribution. However, in addition to the evaluation of the program in a specific location in the past, we are interested in examining whether the job training program can be expanded to another location or should be implemented on a long-term basis. Without actually implementing the job training program at a target site and, in particular, implausible to do so for the future period, our goal is to predict the worst-case treatment effect under $Q$ distribution using a sample from the source population. Therefore, we propose an additional step to the usual program evaluations. Using the same dataset in an empirical research study, following a typical causal analysis, we provide a formal procedure for generalizing the causal estimates under distributional shift.
In the hypothetical scenario where we could observe a sample from distribution $Q$, the solution to the minimization of the MSE, $\mathbb{E}_Q\big[(\tilde{Y}(1)-\tilde{Y}(0)-\tau)^2\big]$, turns out to be $\tau^Q=\mathbb{E}_Q\big[\tilde{Y}(1)-\tilde{Y}(0)\big]$. The ATE $\tau^Q$ can be seen as the best prediction of the individual treatment effect under the target distribution $Q$. If we had access to a sample from $Q$, the prediction of the individual treatment effect and the identification of the ATE would coincide. However, this coincidence breaks down when $Q$ is unknown. Without input from $Q$, we construct a class of distributions $\mathcal{Q}=\{Q: D(P,Q)\leq \delta^2\}$ centered around the source distribution within distance $\delta^2$. The target distribution $Q$ is considered to be contained in the ambiguity set $\mathcal{Q}$ when the neighborhood radius $\delta$ is carefully chosen.
Based on the observation in Remark (ref), we use a slightly different combination of loss function and ambiguity set to make the task of treatment effect prediction even harder. Let us consider the nonparametric quadratic loss $(\tilde{Y}(1)-\tilde{Y}(0)-w\tau)^2$ with additional weighting $w$. We modify the loss function in this way to maintain the regression specification, ensuring the problem remains tractable. We also augment the source distribution $P$ with a constant weighting of 1 and the target distribution $Q$ with an adversarial weighting $w$. The augmented distributions are denoted by $\bar{P}$ and $\bar{Q}$. We use
to measure the distance between the augmented source distribution and target distribution, where the cost function is the quadratic of the $L_p$ norm $\| (\tilde{Y}(1),\tilde{Y}(0),w)-(Y(1),Y(0),1) \|_p^2$ for $p\in (1,\infty]$. In the definition of Wasserstein distance, $\Pi(\bar{P},\bar{Q})$ is the set of couplings of $\bar{P}$ and $\bar{Q}$. Since the adversarial weighting $w$ is unitless, we normalize the potential outcomes by the standard deviation of the realized outcome under the source distribution, following the penalized regression literature, to make it scale-invariant. The optimizer will eventually be scaled back using the same standard deviation when reported. Consequently, the neighborhood radius $\delta$ can be assessed in the magnitude of the standard deviation of the realized outcome $Y$ when we use $L_2$ norm. We apply this normalization in both the simulation and the empirical illustration based on synthetic data below.
Since we only impose that $\bar{Q}$ belongs to $\mathcal{Q}=\{\bar{Q}: D(\bar{P},\bar{Q})\leq \delta^2\}$, to be conservative we pick the distribution within the ambiguity set $\mathcal{Q}$ that leads to the largest MSE for prediction. The corresponding distribution is considered the least favorable distribution. Such a procedure is robust in the sense that, for any $\tau$, the MSE of predicting the individual treatment effect for all distributions in $\mathcal{Q}$ will be bounded by the worst-case MSE. The adversarial weighting $w$ makes the MSE minimization problem more conservative because the adversary could rescale $\tau$ via the weighting $w$ to make the prediction error of the individual treatment effect larger, which inflates the adversarial loss. In this sense, the minimax solution $\tau^{DR}$ is considered the optimal worst-case out-of-population prediction of individual treatment effect. The prediction $\tau^{DR}$ results from a bias-variance tradeoff by minimizing the MSE and hence is different from $\tau^Q$, which is not identifiable under our framework. Relatively stable $\tau^{DR}$ along the increase of the level of robustness, $\delta$, is an indicator of generalizability of the treatment effect to new populations, even with possible distributional shift.
If we had a hypothetical sample from $Q$, we only need marginal distributions of potential outcomes to identify the ATE under $Q$. However, for the prediction problem we set up, the joint distribution $P$ is involved in the definition of the Wasserstein ambiguity set. This is induced by our choice of the nonparametric quadratic loss, which involves the second moment of the individual treatment effect. Even though joint distribution of $(Y(1), Y(0))$ exists in the sense that $P(y_1,y_0)=C^*(P_1(y_1), P_0(y_0))$ for some copula $C^*: [0,1]^2\mapsto [0,1]$, where $P_1$ and $P_0$ are the marginal distributions of $Y(1)$ and $Y(0)$, it can never be identified using our sample. As a result, our context presents an additional layer of complexity compared to the general transfer estimates literature.
With that said, we know the joint distribution $P(y_1,y_0)$ must belong to the Fr\'echet class of joint distributions with marginals $P_1$, $P_0$, parameterized by the set of copulas $\mathcal{C}(P_1,P_0)$. Therefore, we can pick one distribution in $\mathcal{C}(P_1,P_0)$ that gives the smallest worst-case MSE and another one that leads to the largest worst-case MSE. These two cases are considered optimistic and pessimistic cases, respectively. This step is related to the partial identification of the joint distribution of potential outcomes. Visually, we can imagine there is a set $\mathcal{C}(P_1,P_0)$. Each point within the set $\mathcal{C}(P_1,P_0)$ is a joint distribution of potential outcomes. Centered around each point, there is a Wasserstein neighborhood with radius $\delta$. For each point, the Wasserstein neighborhood is defined as ((ref)) with $P$ replaced by the copula.
To assess the robustness of the treatment effect, we start with the following two objective functions:
and
The solution to ((ref)) and ((ref)) is denoted by $\tau_p$ and $\tau_o$ respectively, which predict the treatment effect by minimizing the worst-case MSE. Essentially, we need to find a solution to a minimax optimization problem. The middle layer of $\inf$ and $\sup$ arises due to the partial identification issue.
The primal problem in ((ref)) and ((ref)) appears to be initially difficult to solve, as it involves optimization with respect to an infinite number of distributions. Inspired by the results in, for example, blanchet2019robust and gao2023distributionally, we can derive a closed form of the dual problem of the inner maximization problem.
where $1/p+1/q=1$ so that $q\in[1,\infty)$.
We can see that the right-hand side of ((ref)) is the quadratic of the square root of the MSE under the copula $C$ plus a penalty term with the penalization parameter being the radius of the Wasserstein neighborhood. The penalization term involves the sum of a constant two and the parameter $\tau$, which is different from the usual penalization of targeting parameters alone. The MSE under $C$ in the right-hand side of ((ref)) can be further decomposed into two terms.
Because of the observation in ((ref)), we only need to find the copula that leads to the largest variance of individual treatment effect for the pessimistic case and the copula that leads to the smallest variance for the optimistic case.
For the outer minimization problem, minimizing the quadratic is equivalent to minimizing the terms within the curly bracket. Compared to the bridge estimator for a linear regression in ((ref)) below (see, for instance, knight2000asymptotics), we take the square root of the loss function, and the penalization is not purely applied to the targeting parameter. Thus, solutions to the minimization of the right-hand side of ((ref)) and ((ref)) can be considered as a square-root bridge-type estimator.
Define the solution to the outer minimization problem as \[ f(V,\delta)=\text{argmin}_\tau \sqrt{V+(\tau^*-\tau)^2}+\delta(2+|\tau|^q)^{1/q}. \]
Based on a simple observation, when there is no distribution shift, $\delta=0$ and $f(V,\delta)=\tau^*$. Namely, the prediction of the treatment effect is the ATE under the source distribution when the target population coincides with the source population. On the other hand, when $\delta>0$, we allow for a shift in the distribution. We can easily see that $f(V,\delta)$ shares the same sign as $\tau^*$ but shrinks toward zero, as in any regularized estimation. By being conservative, our prediction of treatment effect under $Q$ is always no larger than $\tau^*$ in magnitude.
Since the joint distribution $P(y_1,y_0)=C^*(P_1(y_1), P_0(y_0))$ also belongs to the copula set $\mathcal{C}(P_1,P_0)$, $V_o\leq V^P\leq V_p$. Therefore, as long as we can find a pair of ($V_p$, $V_o$), we have found a bound for the minimax optimizer with respect to the unknown joint distribution $P$.
When the treatment effect is homogeneous, we immediately know the joint distribution of the potential outcomes from the marginal distributions. This is the case where we do not need to worry about finding a copula, and the pessimistic and the optimistic cases coincide. On the other hand, if we would like to avoid the complication of finding the smallest variance of the individual treatment effect, a naive lower bound for $V_o$ is simply zero.
Given $Var_P\big(Y(1)-Y(0)\big)=0$, the dual objective function reduces to
Proposition (ref) implies delayed shrinkage of the minimax optimizer, which is in contrast to the typical pattern of regularized estimation but intuitive in the context of causal analysis. For any Wasserstein neighborhood radius less than or equal to $\bar{\delta}$, our prediction of the individual treatment effect under $Q$ is always $\tau^*$. Figure (ref) is a graphical illustration of delayed shrinkage, where we set $q=2$ and consider two values of $\tau^*$: $\tau^*=2$ and $\tau^*=1$. The corresponding cutoffs are $\bar{\delta}_2=1.22$ and $\bar{\delta}_1=1.73$.
We can imagine there is more generalizability of our causal estimates if the treatment effect is homogeneous. Only when the target distribution $Q$ is sufficiently different from the reference distribution $P$ by setting a relatively large $\delta$, our best prediction of the treatment effect starts to shrink toward zero. Furthermore, the boundary radius $\bar{\delta}$ decreases in $\tau^*$. This is also intuitive as it would be harder to maintain a larger treatment effect given distributional shift.
When we switch to heterogeneous treatment effect, $Var_{C}\big(Y(1)-Y(0)\big)>0$ and $|f(V,\delta)|<|\tau^*|$ for $\delta>0$. This implies that the solution $f(V,\delta)$ shrinks toward zero immediately whenever there is a distribution shift, even for a tiny shift. Figure (ref) illustrates the immediate shrinkage when the population variance $V$ is 5 with $q=2$ and $q=3$ respectively. We can also see that shrinkage occurs to a greater extent for smaller values of $q$.
Given the pessimistic and optimistic cases, we need to find bounds of $Var_{C}\big(Y(1)-Y(0)\big)$. Using Fr{\'e}chet-Hoeffding inequality, we can find sharp bounds of the variance; see, for instance, fan2010sharp. This technique has also been used by aronow2014sharp and imbens2021causal to derive sharp bounds for the design-based variance-covariance matrix. Let $C^L(u,v)=\max(u+v-1,0)$ and $C^U(u,v)=\min(u,v)$. Fr{\'e}chet-Hoeffding inequality implies that \[ Cov_{C^L}(Y(1),Y(0))\leq Cov_P(Y(1),Y(0))\leq Cov_{C^U}(Y(1),Y(0)). \] As a result, the sharp upper bound of the variance of individual treatment effect is $V_p=V_U(P_1,P_0)=V\big(C^L(P_1,P_0)\big)$, where potential outcomes $Y(1)$ and $Y(0)$ are perfectly negatively dependent. Similarly, the sharp lower bound of the variance is $V_o=V_L(P_1,P_0)=V\big(C^U(P_1,P_0)\big)$, where the two potential outcomes are perfectly positively dependent.
Ultimately, we are interested in predicting the treatment effect under $Q$ using information from the reference distribution $P$. If we can identify the joint distribution $P$, we can form our objective function as
Define the solution to ((ref)) to be \[ \tau^{\texttt{DR}}\equiv f(V^P,\delta)=\text{argmin}_\tau \sqrt{Var_P\big(Y(1)-Y(0)\big)+(\tau^*-\tau)^2}+\delta(2+|\tau|^q)^{1/q}. \] According to Proposition (ref), we have found sharp bounds for $|f(V^P,\delta)|$, which is $\left(|f\left(V\big(C^L(P_1,P_0)\big),\delta\right)|, |f\left(V\big(C^U(P_1,P_0),\delta\right)|\right)$.\footnote{The sign of $f(V^P,\delta)$ and the lower and upper bounds depend on the sign of $\tau^*$.} Therefore, $f(V^P,\delta)$ can be partially identified.
splawa1990application proposes another set of variance bounds, which uses Cauchy–Schwarz inequality to derive the bounds for the covariance between two potential outcomes. This pair of bounds is easier to compute and is given below. \[ V_p^N=Var_{P_1}(Y(1))+Var_{P_0}(Y(0))+2\sqrt{Var_{P_1}(Y(1))Var_{P_0}(Y(0))} \] \[ V_o^N=Var_{P_1}(Y(1))+Var_{P_0}(Y(0))-2\sqrt{Var_{P_1}(Y(1))Var_{P_0}(Y(0))} \]
To proceed with estimation, we need to find the sample counterpart of the objective functions ((ref)) and ((ref)). We first need a set of internally valid estimators.
Assumptions (ref)-(ref) are standard assumptions in the causal inference literature. In particular, Assumption (ref) is the usual random assignment and overlap condition for experimental data.
We can estimate $\tau^*$ using various estimators $\hat{\tau}^*$. For instance, we can use the difference-in-means estimator $\bar{Y}_1-\bar{Y}_0$, where $\bar{Y}_1$ and $\bar{Y}_0$ are sample averages of the treated and untreated outcomes. Or, we can use the inverse probability weighting estimator, $\frac{1}{n}\sum^n_{i=1}\frac{T_iY_i}{e}-\frac{1}{n}\sum^n_{i=1}\frac{(1-T_i)Y_i}{1-e}$.
As for $\hat{V}$, it is easy to compute $\widehat{Var}(Y(1))$ and $\widehat{Var}(Y(0))$ using the random sample of treated and control units. To estimate the sharp bounds of the covariance term, $\widehat{Cov}(Y(1),Y(0))$, we use \[ \int_0^1\hat{P}^{-1}_1(u)\hat{P}^{-1}_0(u)du-\bar{Y}_1\bar{Y}_0 \] and \[ \int_0^1\hat{P}^{-1}_1(u)\hat{P}^{-1}_0(1-u)du-\bar{Y}_1\bar{Y}_0, \] where $\hat{P}_1(y)=\frac{1}{n_1}\sum^n_{i=1}T_i\mathbbm{1}\{Y_i\leq y\}$ and $\hat{P}_0(y)=\frac{1}{n_0}\sum^n_{i=1}(1-T_i)\mathbbm{1}\{Y_i\leq y\}$ are empirical CDFs for the treated and untreated units and $\hat{P}_1^{-1}(u)=\inf\{y:\hat{P}_1(y)\geq u\}$ and $\hat{P}_0^{-1}(u)=\inf\{y:\hat{P}_0(y)\geq u\}$ are their inverse functions. As a result, \[ \hat{V}_o=\widehat{Var}(Y(1))+\widehat{Var}(Y(0))-2\left(\int_0^1\hat{P}^{-1}_1(u)\hat{P}^{-1}_0(u)du-\bar{Y}_1\bar{Y}_0 \right) \] and \[ \hat{V}_p=\widehat{Var}(Y(1))+\widehat{Var}(Y(0))-2\left(\int_0^1\hat{P}^{-1}_1(u)\hat{P}^{-1}_0(1-u)du-\bar{Y}_1\bar{Y}_0\right). \]
Estimators of the Neyman bounds are straightforward to construct based on the variance estimators of $Var_{P_1}(Y(1))$ and $Var_{P_0}(Y(0))$. They are denoted by $\hat{V}_o^N$ and $\hat{V}_p^N$: \[ \hat{V}_p^N=\widehat{Var}(Y(1))+\widehat{Var}(Y(0))+2\sqrt{\widehat{Var}(Y(1))\widehat{Var}(Y(0))}, \] \[ \hat{V}_o^N=\widehat{Var}(Y(1))+\widehat{Var}(Y(0))-2\sqrt{\widehat{Var}(Y(1))\widehat{Var}(Y(0))}. \]
Based on the dual problem, the outer minimization of ((ref)) and ((ref)) becomes an M-estimation problem. Therefore, we can apply the empirical process theory in EmpiricalProcess to derive the asymptotic properties. Let $\tau_p=f(V_p, \delta)$ and $\tau_o=f(V_o, \delta)$. Moreover, let $\hat{\tau}_p$ and $\hat{\tau}_o$ be the solution to the sample minimization problem. Before showing the asymptotic properties of $(\hat{\tau}_p,\hat{\tau}_o)$, let us prove an intermediate result.
Lemma (ref) can be summarized as:
To describe the asymptotic distributions, we need more notation. From now on, let $V_p$ and $V_o$ be the variance bounds (either Neyman or sharp). For $b\in\{p,o\}$, let $\tau_b$ be the minimizer of $M(\tau) \coloneqq A(\tau) + \delta B(\tau)$, where
Let $\partial_{b} A(\tau)$ and $\partial_{\tau^*} A(\tau)$ denote the gradient of $A(\tau)$ with respect to $V_b$ and $\tau^*$ at $\tau$. In the following, both $A$ and its derivatives are evaluated at $\tau_b$:
For $b\in\{p,o\}$, define $ \mathbb{D}_b(\tau_b) = [ M''(\tau_b) ]^{-1} D_b$, where $M''$ denotes the second-order derivative of $M(\tau)$.
The bounds estimators $\hat{\tau}_p$ and $\hat{\tau}_o$ are asymptotically jointly normal. Both $\hat{\tau}_p$ and $\hat{\tau}_o$ are subject to the same source of estimation error, but with different “loading” terms $\mathbb{D}_b(\tau_b)$. Hence, their asymptotic covariance can be easily calculated based on Theorem (ref).
When $\tau_b \neq 0$, the estimation errors of $V_p$, $V_o$, and $\tau^*$ will collectively influence $\hat{\tau}_b$. However, when $\tau_b = 0$, only the estimation error of $\tau^*$ contributes to the variability of $\hat{\tau}_b$, as shown in the following theorem.
Theorem (ref) is a summary of Theorem (ref) in Appendix (ref), which also covers the cases of $q\in(1,2)$ and $q=1$. The asymptotic behavior of the distributionally robust estimator becomes more intricate when the corresponding $\tau_b$ is zero, even though the primary source of stochasticity comes solely from the estimation of $\tau^*$. This complexity arises due to the interplay between the penalty term and the stochastic component of the objective function, which varies depending on the value of the penalty order $q$.
When $q \geq 2$, the objective function is locally quadratic around zero, allowing standard techniques for M-estimation to apply. Under these conditions, the estimator retains asymptotic normality, a hallmark of well-behaved quadratic penalties. Notably, when $q = 2$, the stochastic term $A(\tau)$ exhibits a functional form nearly identical to the penalty term $B(\tau)$ with respect to $\tau$. This structural similarity enables the penalty to exert a slight shrinkage effect on the asymptotic variance, thereby improving estimation efficiency compared to the unpenalized estimator $\hat{\tau}^\ast$.
In this section, we rule out the case $p=\infty$ to construct a unified inference procedure. In the dual problem, $p=\infty$ implies $q=1$, which imposes a significant penalty on the minimax prediction in a similar spirit to a LASSO estimator as shown in Case 2-3 of Theorem (ref) in Appendix (ref).
Theorem (ref) implies that the bound estimators are not necessarily asymptotically normal when $\tau_b=0$, which is uninteresting and also causes trouble for inference. As a result, we would like to rule out this extreme case. We propose the following two-step inference procedure. In the first step, we test the null $H^P_0: \tau^*=0$. Since the solution $\tau_b=0$ holds if and only if $\tau^*=0$ is true when $p\in (1,\infty)$, testing $H^Q_0: \tau^{\texttt{DR}}=0$ is equivalent to testing $H^P_0: \tau^*=0$, which is pretty straightforward. Intuitively, if we do not find any statistically significant internally valid treatment effect based on the source data, we typically would not bother assessing the external validity of our findings. On the other hand, if the null $H^P_0$ is rejected at a small significance level, we would like to see how robust our findings are under distributional shift. This is our goal in the second step.
Since $\tau^{\texttt{DR}}$ is never zero when $\tau^*\neq 0$, zero would not be a meaningful hypothesized value in the second step. Instead, policymakers might have a breakdown point in mind, for instance, the cost to implement the policy. A reasonable hypothesis would be whether the confidence interval of our minimax predictor contains the breakdown point. This is a nonstandard inference problem since $\tau^{\texttt{DR}}$ is only partially identified. Fortunately, we have shown in Section (ref) that the upper and lower bound estimators $\hat{\tau}_o$ and $\hat{\tau}_p$ are asymptotically jointly normal. As a result, we can apply the approach in imbens2004confidence (IM hereafter) and stoye2009more to construct a confidence interval for $\tau^{\texttt{DR}}$.
Nevertheless, such a two-step procedure comes with a caveat. We only proceed with the second step if $H^P_0$ is rejected in the first step, which introduces pre-testing bias. To solve this problem, we modify a Bonferroni-type correction approach in the literature for nonstandard inference; see, for instance, staiger1997instrumental, romano2014practical, mccloskey2017bonferroni, and guo2024statistical.
Let the size of the test be $\alpha$. In the first step, we construct a $1-\beta$ confidence interval for $\tau^*$, $I^*_n(1-\beta)$, where $\beta\in [0,\alpha]$ is some small value. Based on the value of $\tau^*=t$, $\forall\ t\in I^*_n(1-\beta)$, we construct the second-step confidence interval based on equation (6) in imbens2004confidence with confidence level $1-\alpha+\beta$, $I_n^{[t]}(1-\alpha+\beta)$. In practice, one can create a fine grid of $I^*_n(1-\beta)$ in the first step and compute the corresponding confidence intervals in the second step. Lastly, we take a union of the second-step confidence intervals, $\mathcal{I}_n=\cup_{t\in I^*_n(1-\beta)}I_n^{[t]}(1-\alpha+\beta)$. Each interval in the second step accounts for the uncertainty of variance bound estimation only since we fix $\tau^*=t$, and the union step accounts for the uncertainty of $\hat{\tau}^*$. The inclusion of $\beta$ in the confidence level $1-\alpha+\beta$ accounts for the possibility that $\tau^*$ may not lie in $I_n^*(1-\beta)$.
To show the coverage of the two-step inference procedure, we first strengthen the pointwise convergence result in Lemma (ref) to uniform convergence.
An open question that we have not yet addressed is how to determine the radius of the Wassertein neighborhood. In the DRO literature, a data-driven approach has been proposed. The radius $\delta$ is chosen to be decreasing in sample size; see blanchet2022confidence and lin2022distributionally. This approach assumes an i.i.d. sample from the unknown target distribution, which does not comply with our setting. Also, in the limit $\delta$ approaches zero, which corresponds to the case without distributional shifts. We are instead interested in the generalizability of our causal estimates, given distributional shift, even when the sample size is large.
With the connection to the regularized estimation literature via the dual problem, $\delta$ can be considered as the penalization parameter. Following this literature, one might be tempted to find the optimal $\delta$ through cross-validation. This approach typically picks the penalization parameter as the one that minimizes the prediction error in subsamples. Since the target distribution $Q$ is unknown, we do not have a good criterion to assess the performance of different $\delta$.
The nature of $\delta$ resembles that of the sensitivity parameter in the literature on sensitivity analysis that deals with the potential failure of the unconfoundedness assumption. Instead of the concern about internal validity in the sensitivity analysis literature, we assume internal validity but examine the robustness of the internal causal estimates under distributional shift. There is no single best choice for the sensitivity parameter. The general idea is to find some benchmark. In the sensitivity analysis literature, if the sensitivity parameter that nullifies the results is larger than a reasonable benchmark, then the causal findings are considered to be insensitive to the unobserved confounders. We follow the same spirit in finding benchmarks to help with the economic interpretation of $\delta$.
There are many ways one can form benchmarks. Below, we present a few possibilities. Even though in the target distribution $Q$ we do not observe $Y(1)$ since no treatment has been implemented yet, we might still be able to observe $Y(0)$. Hence, one can compute the Wasserstein distance for $Y(0)$ between $P$ and $Q$ distributions. In robust prediction, $\delta$ can be set to multiples of the Wasserstein distance for $Y(0)$ to approximate the true distance between $P$ and $Q$. Without access to the data from a target distribution, we can use the heterogeneity of $P$ as a benchmark for the distributional shift from $P$ to $Q$. If we observe covariates, we can split the sample based on these covariates and compute the Wasserstein distance of the potential outcomes across the resulting subsamples.
For example, in the analysis of job training programs, it has been demonstrated that the pre-intervention employment record is one of the most important predictors of heterogeneous treatment effects; see hotz2005predicting and gupta2023s. As a result, we split the data into two subsamples, one previously employed and another previously unemployed. For the $L_2$ norm, the square root of the sum of the squared 2-Wasserstein distances of the marginal distributions serves as a lower bound for the Wasserstein distance of the joint distributions of potential outcomes. As a preview, for the job training program data used in Section (ref) below, such a lower bound based on the $L_2$ norm cost function is \$1,154 between the two subsamples with or without a previous employment record, which is about 0.23 standard deviation of the post-treatment earnings.
In practice, we recommend using a spectrum of $\delta$ as a stress test. Even though our robust prediction of the treatment effect can never be exactly zero if $q>1$, we can set the minimum level of the treatment effect that can offset the cost as the threshold. The radius $\delta$ leading to a prediction equal to the threshold would be an interesting cutoff, which is considered a breakdown point. We can use one or multiple benchmarking approaches proposed above to assess whether this $\delta$ is considered too small. If so, then our internal estimates might not be robust to distributional shift. In other words, there is not much external validity to our causal findings.
For our general theory, we allow for any $q\in[1,\infty)$ in the definition of the Wasserstein neighborhood. All of our theoretical results hold, regardless of the value of $q$. Therefore, another loose end is how to choose $q$. The behavior of the minimax optimizer follows the same pattern as long as $q> 1$. The only difference is that $f(V,\delta)$ shrinks toward zero more slowly with the increase of $\delta$ as $q$ becomes larger. Moreover, $f(V,\delta)$ never reaches zero when $q>1$. On the other hand, $f(V,\delta)$ can be exactly zero when $q=1$ if $\delta$ is sufficiently large, behaving like a Lasso estimator.
Figure (ref) plots the predicted treatment effect when $\tau^*=2$ for different values of $q$. Figure (ref) shows the behavior of the minimax optimizer under heterogeneous treatment effect with $V=5$, and Figure (ref) shows the case for homogeneous treatment effect with $V=0$.
Parallel with the regularized estimation literature, multiple values of $q$ have been proposed, such as $q=1$ for Lasso and $q=2$ for Ridge. There is no single answer for the optimal $q$. In practice, we recommend using $q=2$ because of its tractability and clear interpretation. With the $L_2$ norm, which leads to $q=2$, $\|(Y(1),Y(0),1)-(\tilde{Y}(1),\tilde{Y}(0),w)\|_2^2=\|(Y(1),Y(0))-(\tilde{Y}(1),\tilde{Y}(0))\|_2^2+(1-w)^2$. Based on the worst-case distribution $\tilde{Q}$, we can quantify the distribution shift resulting from $(Y(1),Y(0))$, which turns out to be $\frac{2}{2+\tau^2}\delta^2$. Thus, when $\tau$ is small, the distribution shift is primarily driven by the change in potential outcomes.
We study the finite sample performance of our two-step confidence intervals in simulation exercises. Potential outcomes $(Y(1),Y(0))$ are drawn from a bivariate normal distribution but truncated to $[-6, 6]^2$, \[ N\Bigg(
,
\Bigg). \] We set $\rho=0.7$, $\sigma_1=2$, $\sigma_0=1$, $\mu_1=\sigma_1$, and $\mu_0=0.2\sigma_0$, unless otherwise noted. Treatments are randomly assigned with probability 0.3.
We consider six cases: (\romannumeral 1) $p=2$, $\delta=0.1$; (\romannumeral 2) $p=2$, $\delta=1$; (\romannumeral 3) $p=2$, $\delta=1$, $\sigma_1=2$, $\sigma_0=0.01$; (\romannumeral 4) $p=2$, $\delta=0.1$, $\mu_1=0.2\sigma_1$, $\mu_0=0.1\sigma_0$; (\romannumeral 5) $p=1.5$, $\delta=0.1$; (\romannumeral 6) $p=3$, $\delta=0.1$. The first two cases serve as baselines with a small radius and a relatively large radius. The third case resembles near-point identification, where the sharp bounds of the variance of individual treatment effect after normalization are [1.92, 1.96]. Point identification can pose a threat to valid uniform inference under partial identification; see, for instance, imbens2004confidence. Case (\romannumeral 4) examines the scenario where the variance of the outcome is significantly larger than the average treatment effect, making it challenging to estimate the ATE under the source distribution $P$ precisely. The first four cases use the $L_2$ norm for the Wasserstein distance. In contrast, the last two cases change the $L_p$ norm in the cost function of the Wasserstein neighborhood but otherwise remain the same as the baseline cases. We use the plug-in variance estimator $\hat{\bm{\Sigma}}$ based on the influence functions in Appendix (ref).
Table (ref) reports the coverage rate of the IM confidence intervals (CIs) with or without Bonferroni correction across 2,000 replications, the average CIs, and the average length ratio of the two-step CIs over the non-corrected CIs. To proceed with our proposed two-step CIs, replications with first-step CIs containing zero are dropped. Whenever this occurs, the reported results are averages across the remaining replications. We report the results with both sharp variance bounds and Neyman variance bounds. We set $\alpha=0.05$ and $\beta=0.045$. In our simulations, the two-step CIs are less conservative when $\beta$ is closer to $\alpha$, but there is not much improvement when $\beta$ is larger than 0.045.
In a finite sample with 500 observations, the non-corrected IM confidence intervals (CIs) exhibit slight under-coverage in cases (\romannumeral 3) and (\romannumeral 4). Nonetheless, the two-step CIs consistently achieve the nominal coverage rate, as expected. When the sample size increases to 1,000, the coverage rate of the non-corrected IM CIs exceeds 0.95 in case (\romannumeral 4). However, in case (\romannumeral 3), the coverage rate of the non-corrected CIs remains below 0.92 even with 2,000 observations per sample. These results remain qualitatively unchanged even with 5,000 replications. When both the non-corrected and two-step CIs attain the nominal coverage rate, the two-step CIs are 16–25% wider than the non-corrected CIs, and their coverage rate can approach one as the neighborhood radius increases. The performance of the CIs is stable across different choices of $L_p$ norms and is consistent between sharp and Neyman variance bounds.
We illustrate our prediction method in the context of a job training program. During the mid-1970s, the National Supported Work Demonstration program randomly assigned qualified applicants to training positions. Pioneered by lalonde1986evaluating and followed by dehejia1999causal and many others, this dataset has been extensively studied to evaluate the performance of different causal estimators. The outcome variable is earnings in 1978 for men, and the treatment variable is participation in the job training program.
Instead of the original experimental data, we use the artificial data of size 1,000,000 generated by (Wasserstein) Generative Adversarial Networks in athey2024using as our population. There are 445 observations in the experimental sample used in dehejia1999causal, comprising 185 men who were treated and 260 men who were untreated. We generate 2,000 samples randomly drawn from our artificial population, while maintaining the fixed ratio of treatment and control units. The treatment subsample size is $n_1=n*185/(185+260)$ and the control subsample size is $n_0=n*260/(185+260)$. We try different sample sizes $n$.
The population ATE is 1,333 dollars. Because the artificial population comes with counterfactuals, we can compute the population variance of individual treatment effect, $Var_P\left(Y(1)-Y(0)\right)$. Table (ref) reports the population variance, the sharp bounds, and the Neyman bounds of the variance. The Neyman bounds are wider than the sharp bounds, as expected. Both lower bounds are pretty close to zero.
We first compute the population prediction with respect to the joint distribution $P$, the perfectly positively dependent copula, and the perfectly negatively dependent copula, respectively. Figure (ref) depicts the predictions for a range of $\delta$ with $q=2$. As expected, $\tau_p\leq \tau^{\texttt{DR}}\leq \tau_o$.
Next, we compare the average prediction across 2,000 samples with the population prediction. The bound parameters are obtained based on sharp variance bounds using population data. We examine the upper and lower bounds of the prediction in the sample, respectively, using either Neyman variance bounds or sharp variance bounds. For the left panel of Figures (ref) and (ref), the sample size is 445. The sample predictions are close to the population prediction for the lower bound. However, there are noticeable gaps for the upper bound estimator based on the sharp variance bound. Increasing the sample size by tenfold in the right panels leads to sample predictions aligning much more closely with the population prediction.
For the sample size $n=4,450$, we also examine the IM confidence intervals based on sharp variance bounds and Neyman variance bounds, with and without the Bonferroni correction. The results are plotted in Figure (ref). The solid black line collects the value of $\tau^{\texttt{DR}}$ in the population corresponding to different radii of the Wasserstein neighborhood. Each dot on the upper and lower curves is the average of the confidence interval endpoints across 2,000 replications. The Neyman bound CIs are pretty similar to the sharp bound CIs. Not surprisingly, our two-step CIs are wider than the non-corrected CIs, but they are not a lot wider.
The length ratios of the CIs with or without Bonferroni correction are reported in Table (ref) below. The ratios are quite reasonable, ranging from 1.16 to 1.27. The length ratios based on the Neyman variance bounds are overall smaller, implying that two-step CIs based on the Neyman bounds are less conservative compared with the two-step CIs based on the sharp bounds.
In the literature, it has been recorded that the average cost of providing job training services ranges from \$953, \$919, \$430 to \$118 per trainee across various locations in the US in the early 1980s; see hotz2005predicting. For the job training program on which our artificial population data is based, lalonde1986evaluating records that the program cost is at least \$2,700 per trainee. The average cost per trainee across the four locations is \$605. Based on the first-order condition of the outer minimization problem, the lower bound $\tau_p=605$ when $\delta=0.92$ using the sharp bound of the variance. Using the benchmark \$1,154 calculated in Section (ref) based on the heterogeneity between two subsamples, which is equivalent to 0.23 standard deviation of the realized outcome, this $\delta$ is sizable, indicating some robustness of the treatment effect against distributional shift.
We propose a method for out-of-population prediction of treatment effect with only retrospective data. Although our robust prediction is partially identified, we provide a confidence set for the prediction through a two-step procedure.
In the current paper, we consider only distributional shifts in potential outcomes and do not include covariates. In practice, however, covariates play an important role in observational data. Extending our framework to incorporate covariates is an important direction for future research. Ideally, we would allow for distributional shifts both in covariates and in the conditional distribution of potential outcomes.
We focus on a single cross section in this paper. However, panel data have been used extensively in empirical works. A popular method for identifying causal effects is the difference-in-differences approach, which utilizes panel or pooled cross-sectional data. There are typically multiple periods post treatment. It would be interesting to generalize our method to short panel data.