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.
56,826 characters · 11 sections · 66 citation commands
Unveiling the Potential of Robustness in Selecting Conditional Average Treatment Effect Estimators
The escalating demand for decision-making has sparked an increasing interest in Causal Inference across various research domains, such as economics farrell2015robust, chernozhukov2018double, kitagawa2018should, abadie2023should, statistics wager2018estimation, li2022random, foster2023orthogonal, kennedy2023towards, healthcare zhang2019medical, foster2011subgroup, qian2021synctwin, bica2021real, kinyanjui2022adcb, and financial application bottou2013counterfactual, chu2021graph, huang2021causal, donnelly2021correction, huang2023towards, fernandez2023comparison, li2024causal. The primary goal in personalized decision-making is to quantify the causal effect of a specific treatment (or policy/intervention) on the target outcome, and understanding such causal effects is closely connected with identifying the Conditional Average Treatment Effect (CATE). In observational studies, identifying the CATE faces a significant and fundamental challenge: the absence of counterfactual knowledge. According to Rubin Causal Model rubin2005causal, the CATE is determined by comparing potential outcomes under different treatment assignments (i.e., treat and control) for a specific covariate. Nonetheless, in real-world applications, we can only observe the potential outcome under the actual treatment (i.e., factual outcome), while the potential outcome under the alternative treatment (i.e., counterfactual outcome) remains unobserved. The unavailability of the counterfactual outcome is widely recognized as the fundamental problem in causal inference holland1986statistics, making it difficult to accurately determine the true value of the CATE.
The advancement of machine learning (ML) has opened up a promising opportunity to improve the CATE estimation from observational data. Several innovative CATE estimation approaches, such as meta-learners and causal ML models, have been proposed to tackle the fundamental challenge in causal inference and enhance the predictive accuracy of CATE estimates (as discussed in Section (ref)). Nevertheless, the emergence of various CATE estimation methods has brought forth a new question: Given multifarious options for CATE estimators, which should be chosen? In observational data, treatment is often non-random and propensity scores remain unknown. Conventional model validation procedures, unfortunately, are not suitable for CATE estimator selection in this case due to the absence of ground truth CATE labels. Therefore, exploring proper metrics for CATE estimator selection remains an essential yet challenging research topic in causal inference.
Recent research has emphasized the significance of model selection for CATE estimators, as highlighted in schuler2018comparison, pmlr-v202-curth23b, mahajan2022empirical. These works have proposed and summarized two types of criteria for CATE estimator selection, the plug-in metric $\mathcal{R}_{\tilde{\tau}}^{plug}(\hat{\tau})$ and the pseudo-outcome metric $\mathcal{R}_{\tilde{Y}}^{pseudo}(\hat{\tau})$:
One can establish a plug-in estimator $\tilde{\tau}$ or construct a pseudo-outcome estimator $\tilde{Y}$ using the validation data to select CATE estimator $\hat{\tau}$. The previous studies schuler2018comparison, pmlr-v202-curth23b, mahajan2022empirical have shown that these metrics offer some assistance in identifying well-performing CATE estimators. However, two additional challenges are still encountered in these two metrics.
\paragraph{Challenge 1: How to determine the metric form and underlying ML models for nuisance parameters?} As previously discussed, plug-in and pseudo-outcome metrics have various forms, and both of them rely on estimating nuisance parameters $\tilde{\eta}$ using ML algorithms such as linear models, tree-based models, etc. Plug-in metrics even need to fit an additional ML model for the plug-in learner $\tilde{\tau}$. However, selecting the suitable metric form and ML algorithms can be very difficult without the knowledge of true data generating process. Consequently, we might go round in circles as this challenge leads us back to the original estimator selection problem pmlr-v202-curth23b.
\paragraph{Challenge 2: These metrics are not well-targeted for selecting robust a CATE estimator.} In potential outcome framework rubin2005causal, the factual distribution $P^{F}$ and the counterfactual distribution $P^{CF}$ for $t \in \{0,1\}$ can be defined as follows:
The above (ref) reveals that the covariate shift $P(X|T=t) \neq P(X|T=1-t)$ leads to a distribution shift between $P^{F}$ and $P^{CF}$ - and such distribution shift can be further exacerbated once the unconfoundedness assumption $P(Y^t|X,T=t)=P(Y^t|X,T=1-t)$ is violated. It is widely recognized that ML models often struggle when the training and test data do not adhere to the same distribution. Therefore, it becomes essential to select a CATE estimator learned on $P^{F}$ that demonstrates robust performance to the counterfactual distribution $P^{CF}$. This need for robustness holds even greater significance than the pursuit of an ideal “stellar” estimator because striving for the perfect estimator can be futile in the absence of ground truth counterfactual labels.
\paragraph{Contributions.} In this paper, we propose a Distributionally Robust Metric (DRM) for CATE estimator selection. The main contributions are summarized as follows: (1) The proposed DRM method is nuisance-free, eliminating the need to fit models for nuisance parameters (outcome function, propensity function, and plug-in learner). (2) The DRM method is designed to prioritize selecting a distributionally robust CATE estimator. (3) We provide a finite sample analysis of the proposed distributionally robust value $\hat{\mathcal{V}}^t(\hat{\tau})$ for $t \in \{0, 1\}$, showing it decays to $\mathcal{V}^t(\hat{\tau})$ at a rate of $n^{-1/2}$. (4) Experimental results validate the effectiveness of the DRM method in selecting a CATE estimator that is robust to the distribution shift incurred by covariate shift and hidden confounders.
Suppose the observational data contain $n$ i.i.d. samples $\{(x_i, t_i, y_i)\}_{i=1}^{n}$, with the associated random variables being $\{(X_i, T_i, Y_i)\}_{i=1}^{n}$. For each unit $i$, $X_i \in \mathcal{X} \subset \mathbb{R}^d$ is $d$-dimensional covariates and $T_i \in \{0,1\}$ is the binary treatment. Potential outcomes for treat ($T=1$) and control ($T=0$) are denoted by $Y^1, Y^0 \in \mathcal{Y} \subset \mathbb{R}$. The observed (factual) outcome is $Y=TY^1 + (1-T)Y^0$. The propensity score rosenbaum1983central is defined as $\pi(x):=P(T=1 \mid X=x)$. The conditional mean potential outcome surface is defined as $\mu_t(x) :=\mathbb{E}\left[Y^t \mid X=x\right]$ for $t \in \{0,1\}$. The true CATE is defined as
Following the standard and necessary assumptions in potential outcome framework rubin2005causal, we impose Assumption (ref) that ensure treatment effects are identifiable.
The goal of CATE estimator selection is to select the best CATE estimator, denoted by $\hat{\tau}_{best}$, from a set of $J$ candidate estimators $\{\hat{\tau}_1, \dots, \hat{\tau}_J\}$:
Here, $\mathcal{R}^{oracle}(\hat{\tau})$ is associated with $\mathbb{E}[(\hat{\tau}(X)-\tau_{true}(X))^2]$, known as the Precision of Estimating Heterogeneous Effects (PEHE) w.r.t. $\hat{\tau}$ hill2011bayesian, shalit2017estimating. Note that $\mathcal{R}^{oracle}(\hat{\tau})$ cannot be employed to evaluate CATE estimators' performances in real applications as we do not have access to $\tau_{true}$. Previous studies have introduced plug-in and pseudo-outcome metrics to aid in CATE estimator selection, as shown in equation (ref). Then, the CATE estimator $\hat{\tau}_{select}$ is selected on validation data by
Notably, both the plug-in and pseudo-outcome metrics necessitate the fitting of nuisance parameters $\tilde{\eta}$ (e.g., $\tilde{\eta}=(\tilde{\mu}_1, \tilde{\mu}_0, \tilde{\pi})$) using off-the-shelf ML models. While some papers like cui2024selective address the selection of nuisance parameters for Averate Treatment Effect (ATE) estimators, e.g., the doubly robust estimator farrell2015robust, chernozhukov2018double, huang2022robust, our paper focuses on the selection of CATE estimators rather than nuisance parameters. For the plug-in metric, $\tilde{\tau}$ can be constructed using any CATE estimator discussed in Appendix (ref), yielding metrics such as plug-T, plug-DR, etc. For the pseudo-outcome metric, $\tilde{Y}$ can be constructed using a specific formula discussed in Appendix (ref), yielding metrics such as pseudo-DR, pseudo-R, etc. The metrics based on the influence function alaa2019validating and the R-learner objective nie2021quasi are categorized into the pseudo-outcome metric. The categorization of plug-in and pseudo-outcome metrics maintains consistency with pmlr-v202-curth23b, mahajan2022empirical.
\paragraph{CATE estimation.} Recent advancements in ML have emerged as powerful tools for estimating CATE from observational data, and researchers pay particular attention to meta-learners and causal ML models. Existing meta-learners mainly include traditional learners such as S-learner, T-learner, PS-learner, and IPW-learner, as well as new learners such as X-learner kunzel2019metalearners, U-learner fisherinverse, nie2021quasi, DR-learner kennedy2023towards, foster2023orthogonal, R-learner nie2021quasi, and RA-learner curth2021nonparametric. The specific details of these meta-learners are stated in Appendix (ref). Additionally, some studies also focus on developing innovative causal ML models for CATE estimation, such as Causal BART hahn2020bayesian, Causal Forest wager2018estimation, athey2019generalized, oprescu2019orthogonal, generative models like CEVAE louizos2017causal and GANITE yoon2018ganite, representation learning nets including SITE yao2018representation, TARNet shalit2017estimating, Dragonnet shi2019adapting, FlexTENet curth2021inductive, and HTCE bica2022transfer, disentangled learning nets like D$^2$VD kuang2017treatment, kuang2020data, DeR-CFR wu2022learning, and DR-CFR hassanpour2019learning, and representation balancing nets such as BNN johansson2016learning, CFRNet shalit2017estimating, DKLITE zhang2020learning, IGNITE guo2021ignite, BWCFR assaad2021counterfactual, DRRB huang2023towards, and DIGNet huang2024dignet. Recent surveys guo2020survey, yao2021survey, nogueira2022methods have also conducted a systematic review of various causal inference methods.
\paragraph{CATE estimator selection.} Compared to the diverse range of CATE estimation methods, selecting CATE estimators has received limited attention in existing causal inference research. Current methods for selecting CATE estimators can be broadly classified into two main categories. The first category, which is also considered in this paper, involves using plug-in and pseudo-outcome methods to evaluate CATE estimators. These methods share two common characteristics: 1) Both methods require fitting ML models for nuisances (e.g., outcome function, propensity function, CATE function) on a validation set and then implementing the learned ML models in either the plug-in surrogate or the pseudo-outcome surrogate; 2) Both methods serve as surrogates for the expected error between the CATE estimator and the true CATE, i.e., $\mathcal{R}^{oracle}(\hat{\tau})$ in equation (ref). The difference between the two methods is that the plug-in method directly approximates the true CATE function, where only covariate variables are involved, while the pseudo-outcome method typically constructs a specific formula incorporating covariates, treatment, and outcome variables. For example, the pseudo-DR proposed in saito2020counterfactual is constructed by the outcome predictors learned with representation balancing objective shalit2017estimating, johansson2022generalization. Recent research schuler2018comparison, pmlr-v202-curth23b, mahajan2022empirical has conducted thorough empirical investigations into exploring these two methods for selecting CATE estimators. Their findings suggest that no single selection criterion can universally outperform others in all scenarios in the task of selecting CATE estimators. More details of the two selection methods are stated in Appendix (ref). The second category considers leveraging the data generating process (DGP) to generate synthetic data with the known true CATE function, allowing the validation of CATE estimators' performance on this synthetic data. For example, authors in advani2019mostly find that placebo and structured empirical Monte Carlo methods are helpful for estimator selection under some restrictive conditions. In addition, researchers in schuler2017synth, athey2021using, parikh2022validating focus on training generative models to enforce the generated data to approximate the distribution of the observed data. However, the DGP-based method still faces some limitations in CATE estimator selection: 1) it only guarantees the resemblance of the generated data to the factual distribution, without considering the counterfactual distribution; and ii) there is a potential risk of the method favoring estimators that closely resemble the generative models curth2021really.
In this section, we introduce the Distributionally Robust Metric (DRM) for CATE estimator selection. First, we capture the uncertainty in PEHE in a distributionally robust manner (Section (ref)). We then establish the DRM based on the distributionally robust value of PEHE (Section (ref)).
Proposition (ref) indicates that the PEHE is equal to four terms, where $\mathbb{E}[\hat{\tau}(X)^2]$, $\mathbb{E}[\hat{\tau}(X)Y^0]$, and $\mathbb{E}[-\hat{\tau}(X)Y^1]$ depend on $\hat{\tau}$, while $\zeta$ is a constant that is independent of $\hat{\tau}$. The term $\mathbb{E}[\hat{\tau}(X)Y^t]$ for $t \in \{0,1\}$ can be further decomposed as follows:
Equation ((ref)a) can be computed empirically since the potential outcome $Y^t$ is observable in the group of $T=t$. However, equation ((ref)b) is empirically uncomputable due to the unavailability of $Y^t$ in the group of $T=1-t$. The unknown term $\mathbb{E}[\hat{\tau}(X)Y^t|T=1-t]$ therefore determines the uncertainty in PEHE. To capture such an uncertainty, we therefore establish distributionally robust values for $\mathbb{E}[\hat{\tau}(X)Y^0|T=1]$ and $\mathbb{E}[-\hat{\tau}(X)Y^1|T=0]$ based on a Kullback-Leibler (KL) ambiguity set.
By setting an adequately large ambiguity radius in Definition (ref), the following inequalities hold for $\mathbb{E}[\hat{\tau}(X)Y^0|T=1]=\mathbb{E}^{P_T}[\hat{\tau}(X)Y^0]$ and $\mathbb{E}[-\hat{\tau}(X)Y^1|T=0]=\mathbb{E}^{P_C}[-\hat{\tau}(X)Y^1]$:
To provide a clearer understanding, let us consider the example of $\mathbb{E}^{P_T}[\hat{\tau}(X)Y^0]$. Since the term $\mathbb{E}[\hat{\tau}(X)Y^0]$ is computable on its factual distribution $P_C$ but uncomputable on its counterfactual distribution $P_T$, we can construct an ambiguity set centered around the distribution $P_C$ such that it is large enough to contain the distribution $P_{T}$. By doing so, we can capture the uncertainty of $\mathbb{E}^{P_T}[\hat{\tau}(X)Y^0]$ w.r.t. $\hat{\tau}$. In other words, the value of the uncomputable quantity $\mathbb{E}^{P_T}[\hat{\tau}(X)Y^0]$ will be at most $\mathcal{V}^0(\hat{\tau})$. Similarly, the value of the uncomputable quantity $\mathbb{E}^{P_C}[-\hat{\tau}(X)Y^1]$ will be at most $\mathcal{V}^1(\hat{\tau})$. Obviously, the uncertainty in PEHE will be larger if the distribution shift between factual and counterfactual distribution is severer. Consequently, we can obtain the distributionally robust value of PEHE in Corollary (ref), which measures the uncertainty in PEHE.
As Corollary (ref) provides the distributionally robust (worst-case) value of PEHE, it can naturally measure the robustness of the CATE estimator $\hat{\tau}$ against distribution shift between counterfactual distribution and factual distribution. In this section, we will provide two steps involved in using Corollary (ref) to construct the DRM method for CATE estimator selection.
\paragraph{Step 1: Establishing computational tractability of $\mathcal{V}^t(\hat{\tau})$.} The distributionally robust values $\mathcal{V}^0(\hat{\tau})$ and $\mathcal{V}^1(\hat{\tau})$ in equation (ref) are initially defined as supremum problems over infinite support, presenting a substantial computational challenge. Theorem (ref) reformulates the infeasible supremum problems into tractable minimum problems.
In the finite-sample scenario, $\mathcal{V}^0(\hat{\tau})$ and $\mathcal{V}^1(\hat{\tau})$ can be empirically approximated as follows:
Note that in equation (ref), the potential outcomes $Y^0$ and $Y^1$ are replaced by the observed outcome $Y$ due to the fact that $(1-T)Y^0=(1-T)Y$ and $TY^1=TY$, which aligns with the Consistency assumption in Assumption (ref). We then provide a finite-sample analysis of the gap between $\hat{\mathcal{V}}^t(\hat{\tau})$ and $\mathcal{V}^t(\hat{\tau})$ in the following Theorem (ref), which suggests the gap decays at a rate of $n^{-1/2}$.
\paragraph{Step 2: Finalizing Distributionally Robust Metric for CATE estimator selection.} We first define two functions that are useful in obtaining $\mathcal{V}^0(\hat{\tau})$ and $\mathcal{V}^1(\hat{\tau})$:
Here, $Z$ denotes $\hat{\tau}(X)Y$ for notational simplicity. We then use the Newton-Raphson method to find the empirical solution for $\mathcal{\hat{V}}^t(\hat{\tau})$, exploiting the convexity of $\hat{F}_t(\lambda_t, \epsilon_t;\hat{\tau})$ w.r.t. $\lambda_t$. Based on the distributionally robust value of PEHE, i.e., $\mathcal{V}_{PEHE}(\hat{\tau})$ in equation (ref), we finally obtain the selected estimator $\hat{\tau}_{select}= \mathop{\arg \min}_{\hat{\tau} \in \{\hat{\tau}_1, \dots, \hat{\tau}_J\}}\mathcal{R}^{DRM}(\hat{\tau})$ such that
Algorithm (ref) provides complete procedure of using the DRM method for CATE estimator selection.
\paragraph{Discussion on the ambiguity radius $\epsilon$.} The ambiguity radius $\epsilon$ plays a critical role in real-world applications mohajerin2018data, ma2020understanding, pflug2023multistage. However, determining an appropriate value for $\epsilon$ can be challenging as it requires striking a balance between ensuring the bound in equation (ref) holds and maintaining its tightness. Specifically, if $\epsilon$ is set too small, it fails to guarantee that the counterfactual distribution is contained within the ambiguity set centered at factual distribution (the bound in Corollary (ref) can hold). On the other hand, if $\epsilon$ is set too large, even though the ambiguity set can encompass more distributions to ensure the counterfactual distribution is contained, the bound in Corollary (ref) can be less tight. In general, selecting a proper ambiguity radius is an open problem in distributioanlly robust optimization (DRO) literature hu2013kullback, mohajerin2018data, kuhn2019wasserstein, levy2020large, wang2021sinkhorn.
In this paper, we provide a guidance for determining the ambiguity radius for our DRM method. Based on the above discussion, an ideal radius should be $\epsilon_1^*=D_{KL}(P_C||P_T)$ and $\epsilon_0^*=D_{KL}(P_T||P_C)$, which ensures that the bound in Corollary (ref) holds and is tight. However, as defined in equation (ref), both $P_C$ and $P_T$ involve counterfactual information, making it unattainable to directly compute $D_{KL}(P_C||P_T)$ and $D_{KL}(P_T||P_C)$. To overcome this challenge, we demonstrate that Proposition (ref) provides an intriguing alternative approach to acquire $D_{KL}(P_C||P_T)$ and $D_{KL}(P_T||P_C)$ when unconfoundedness in Assumption (ref) is satisfied.
Proposition (ref) provides an important insight that the uncomputable term $D_{KL}(P_C||P_T)$ (or $D_{KL}(P_T||P_C)$) can be replaced by a computable quantity $D_{KL}(P^C_X||P^T_X)$ (or $D_{KL}(P^T_X||P^C_X)$), where $P^C_X$ and $P^T_X$ are empirically observable. Consequently, the ideal ambiguity radius can be set as $\epsilon_1^*=D_{KL}(P^C_X||P^T_X)$ and $\epsilon_0^*=D_{KL}(P^T_X||P^C_X)$. While the KL divergence can be approximated using empirical algorithm (e.g, Nearest-Neighbor wang2006nearest, noh2014bias), we recommend setting the ambiguity radius larger than the empirically approximated KL divergence (see specific explanations in Appendix (ref)). This is necessary because it ensures that the ambiguity set is large enough to contain the target distribution. It is also important to note that though the Algorithm (ref) involves approximating $\epsilon_1^*=D_{KL}(P^C_X||P^T_X)$ and $\epsilon_0^*=D_{KL}(P^T_X||P^C_X)$, the DRM itself remains free of nuisances, as this approach only determines the ambiguity radius but does not involve learning any nuisance function such as the outcome function, propensity function, and plug-in learner.
\paragraph{Estimators & Selectors.} We consider a total of 36 CATE estimators, comprising the combination of 4 base ML models and 9 meta-learners. Specifically, the base ML models are Linear Regression (LR), Support Vector Machine (SVM), Random Forests (RF), and Neural Net (Net). We consider these ML models for CATE estimators because they are representative of both rigid and flexible models, with each encoded distinct inductive biases, as highlighted by curth2021inductive, pmlr-v202-curth23b. Note that for the LR method, we employ Ridge regression for regression tasks and Logistic regression for classification tasks. As for the remaining methods, we utilize their corresponding regressors and classifiers for regression and classification tasks, respectively. Regarding the meta-learners, we select a set of both traditional basic learners (S-, T-, PS-, and IPW-learners) and recently developed learners (X-, DR-, U-, R-, and RA-learners), as detailed in Appendix (ref). We consider 14 CATE selectors, consisting of 9 plug-in methods that rely on the above 9 learners, 3 pseudo-outcome methods (pseudo-DR, -R, and -IF), the random selection, the factual selection (from the 6-learner pool with S-, T-), the Nearest-Neighbor Matching rolling2014model, and our proposed DRM. The specific details of baseline selectors are stated in Appendix (ref). We employ the eXtreme Gradient Boosting (XGB) chen2016xgboost as the underlying ML model for both plug-in and pseudo-outcome methods. We choose XGB because: i) it demonstrates superior performance in various scenarios, ensuring a good performance of baseline selectors; ii) the need to avoid potential congeniality bias that may arise from using the similar ML models employed in CATE estimators pmlr-v202-curth23b; iii) aligning with alaa2019validating where XGB is used for their proposed pseudo-IF metric. The details of hyperparameters for nuisance models are stated in Section (ref) of Appendix.
\paragraph{Dataset.} Since the ground truth of CATE is unavailable in real-world data, previous studies commonly utilize semi-synthetic datasets to compare model performance. In line with curth2021inductive, pmlr-v202-curth23b, we collect the covariates with $n=4802$ data points from ACIC2016 dataset dorie2019automated. Then, we generate treatment with $T_i|X_i \sim Bern(1/(1+\exp(-\xi (\beta_T' X_i+3))))$, where $Bern$ indicates the Bernoulli distribution. The potential outcomes are generated by a linear function with interaction terms:
The coefficient values are set as follows: $\beta_T, \beta_{j}, \beta_{j,k}, \beta_{j,k,l} \sim Bern(0.2)$, $\gamma_j \sim Bern(\rho)$, and $\epsilon_i \sim \mathcal{N}(0, 0.1)$. The parameter $\xi$ in treatment assignment represents the level of selection bias, and the parameter $\rho$ in $\gamma_j$ represents the complexity of the CATE function. We adopt the above data generating process to randomly generate 100 datasets, each with a training/validation/testing ratio of 49%/21%/30%.
\paragraph{Settings.} In this section, we mainly investigate whether the estimator selected by DRM can demonstrate robustness to the selection bias and unobserved confounders. In addition, as demonstrated in curth2021inductive, pmlr-v202-curth23b, the complexity of CATE function also affects relative performance of estimators and selectors. Given these considerations, we design the following three settings to compare the CATE selectors. Setting A: With the unconfoundedness assumption, let $\rho$ vary in $\{0, 0.1, 0.3\}$ with fixing $\xi=1$. Setting B: With the unconfoundedness assumption, let $\xi$ vary in $\{0, 1, 2\}$ with fixing $\rho=0.1$. Setting C: Without unconfoundedness assumption, fix $\rho=0.1$ and $\xi=1$. Then randomly remove $\lfloor m \cdot d \rfloor$ covariates such that the dimension of observed covariates is $d - \lfloor m \cdot d \rfloor$, where $m$ denotes the ratio of missing covariates varying in $\{0.1, 0.5, 0.9\}$. All the experiments are run on Dell 3640 with Intel Xeon W-1290P 3.60GHz CPU.
\paragraph{Comparison criteria.} The CATE estimator $\hat{\tau}$ is believed better if it achieves a smaller difference between $\mathcal{R}^{oracle}(\hat{\tau})$ and $\mathcal{R}^{oracle}(\hat{\tau}_{best})$, where $\hat{\tau}_{best}$ is the actual best estimator in equation (ref). We therefore use the following Regret criteria to compare estimators chosen by different selectors:
To further assess the ranking ability of each selector, we calculate the Spearman rank correlation between the rank order determined by the oracle metric $\mathcal{R}^{oracle}(\hat{\tau})$ and the rank order determined by each selector. All the reported values (Mean $\pm$ Standard deviation) are computed over 100 runs.
\paragraph{Regret comparison.} The results presented in Table (ref) demonstrate consistently good performance from the DRM selector across various settings. In setting A, the DRM selector outperforms other selectors as the CATE complexity ($\rho$) varies. Additionally, Plug-R, Plug-S, and Plug-PS also perform well in terms of the Regret criterion, which aligns with prior findings in schuler2018comparison that the R-objective is excellent in many cases. Note that the strong performance of Plug-S and Plug-PS may be due to less pronounced heterogeneity in the CATE function compared to the outcome function in the data generating process. We also compare the PEHE performance (i.e., $\mathcal{R}^{oracle}(\hat{\tau}_{select})$) of different selectors in Table (ref) of Section (ref). The results indicate that Plug-R, Plug-S, and Plug-PS tend to exhibit better PEHE as the CATE complexity decreases, aligning with the findings in pmlr-v202-curth23b. In setting B, the DRM selector demonstrates robustness against selection bias (controlled by $\xi$) compared to many baselines. However, for the case $\xi=2$, DRM selects a poor estimator 1 or 2 times out of 100 experiments, as shown in Figure (ref). Although this weakens its overall performance, DRM still outperforms many baselines in this scenario. In the scenario $\xi=0$ where no selection bias is present, the factual selection criterion performs better in this specific setting. In this case, DRM does not demonstrate a significant advantage, as there is no distribution shift caused by selection bias. In setting C where the unconfoundedness assumption is violated, most selectors exhibit inferior performance. In contrast, DRM demonstrates consistent outperformance across all three cases, and its superiority becomes particularly significant as $m$ increases to 0.9, showcasing its robustness against the distribution shift arising from unobserved confounders.
\paragraph{Ranking ability.} In Table (ref), the DRM method demonstrates favorable performance in ranking estimators, surpassing certain Plug- (e.g., U, T, IPW, DR, RA) and Pseudo- (e.g., DR, IF) selectors. In comparison to other nuisance-free baselines (Random, Fact, and Matching), DRM achieves significantly superior ranking ability. However, compared to Plug-S, -PS, -X, and -R, it does not exhibit remarkable performance in ranking CATE estimators, possibly due to the fact that DRM selects estimators based on their distributionally robust (worst-case) performance. Indeed, the definition of ranking inherently involves the concept of expected (average) performance, which is not determined solely by either the best or worst performance. While distributionally robust performance serves as a suitable criterion for selecting players to participate in the Olympics, it may not be a reasonable standard for ranking players' average performance. Therefore, it would be intriguing to explore some ways in future research that can enhance the ranking ability of our DRM selector.
\paragraph{Variance analysis.} Table (ref) indicates that baseline selectors tend to exhibit higher variances in Regret performance. This is primarily due to the wide range of PEHE performances across the 36 CATE estimators. If a selector consistently selects either good or bad estimators, the variance would not be very large. To investigate this further, we sorted all 36 estimators in ascending order based on their $\mathcal{R}^{oracle}(\hat{\tau})$ values, resulting in the sorted list: $[\mathcal{R}^{oracle}(\hat{\tau}_1), \dots, \mathcal{R}^{oracle}(\hat{\tau}_J)]$. We then determine the actual rank of the selected estimator within this list and visualize the distribution of these 100 ranks using a stacked bar chart. Figure (ref) shows that many baseline methods tend to select CATE estimators from various percentile ranges, leading to high variance across the 100 selections. Notably, the DRM selector consistently chooses higher-ranked (i.e., better performing in PEHE) estimators, demonstrating its robustness in CATE estimator selection.
\paragraph{Potential improvements.} There are several potential improvements based on the current experimental settings. First, the existing results suggest that Plug-S performs better than Plug-T, indicating that the complexity of CATE function is relatively simple. It would help to provide more comprehensive analysis if investigating how DRM compares to baselines when the CATE function is more complex. Second, since the impact of selection bias can vary with sample size alaa2018limits, it is important to compare different selectors when the sample size is sufficiently large. Third, considering baselines that are specifically designed for addressing hidden confounders could provide valuable insights for testing different selectors under such conditions. We encourage deeper investigation of causal model selection without assuming unconfoundedness. Finally, it would be good if future studies will apply DRM and other selectors in Healthcare, Economics, and Business applications with real-world data, as CATE estimator selection plays an important role in personalized decision makings.
This paper sheds lights on the potential of robustness in CATE estimator selection. We propose a distributionally robust metric (DRM). The proposed metric is nuisance-free, eliminating the need to fit models for nuisance parameters (outcome function, propensity function, and plug-in learner). Additionally, it is well-targeted for selecting a robust CATE estimator. We provide a finite sample analysis that demonstrates the gap between $\hat{\mathcal{V}}^t(\hat{\tau})$ and $\mathcal{V}^t(\hat{\tau})$ reduces at a rate of $n^{-1/2}$ for $t \in \{0, 1\}$. The experimental results showcase that the CATE estimator selected by DRM demonstrate robustness to the distribution shift incurred by covariate shift and hidden confounders.
\paragraph{Limitations and future work.} This paper explores the potential of robustness in CATE estimator selection. However, we must acknowledge that our DRM method is not a silver bullet, as consistent estimation on the CATE are never attainable chernozhukov2018generic. Here, we outline some challenges and suggest future research directions. First, while Proposition (ref) provides useful guidance for setting ambiguity radius in the DRM algorithm, we cannot guarantee that the empirically-computed radius is optimal due to potential bias in the algorithm's approximation of KL-divergence. Second, as discussed in Section (ref), enhancing the ranking capability of DRM is a promising area for further research. Moreover, our findings are based on KL-divergence. However, using other divergences, such as the Wasserstein distance, to construct the ambiguity set could incorporate more diverse distributions, despite the challenges in solving the dual formulation of the Wasserstein distributionally robust value. Simultaneously, exploring whether alternative divergences can yield a tighter bound for the PEHE error is also interesting pmlr-v80-alaa18a. Finally, inspired by cui2024selective, understanding how nuisance parameters influence metrics like plug-DR and pseudo-DR might be helpful in CATE estimator selection. We hope our methods and findings will spur interest in model selection for causal inference, as well as in related fields like domain adaptation and out-of-distribution generalization.
Qi WU acknowledges the support from The CityU-JD Digits Joint Laboratory in Financial Technology and Engineering, The Hong Kong Research Grants Council [General Research Fund 11219420/9043008], and The CityU APRC Grant 9610643. The work described in this paper was partially supported by the InnoHK initiative, the Government of the HKSAR, and the Laboratory for AI-Powered Financial Technologies. We finally thank all the anonymous reviewers for their constructive suggestions.