Title: Unveiling the Potential of Robustness in Selecting Conditional Average Treatment Effect Estimators
Venue: NeurIPS 2024
Accepted: yes
Keywords: Causal inference

Treatment effect

CATE estimator selection

Model Validation

Robustness
Abstract: The growing demand for personalized decision-making has led to a surge of interest in estimating the Conditional Average Treatment Effect (CATE). Various types of CATE estimators have been developed with advancements in machine learning and causal inference. However, selecting the desirable CATE estimator through a conventional model validation procedure remains impractical due to the absence of counterfactual outcomes in observational data. Existing approaches for CATE estimator selection, such as plug-in and pseudo-outcome metrics, face two challenges. First, they must determine the metric form and the underlying machine learning models for fitting nuisance parameters (e.g., outcome function, propensity function, and plug-in learner). Second, they lack a specific focus on selecting a robust CATE estimator. To address these challenges, this paper introduces a Distributionally Robust Metric (DRM) for CATE estimator selection. The proposed DRM is nuisance-free, eliminating the need to fit models for nuisance parameters, and it effectively prioritizes the selection of a distributionally robust CATE estimator. The experimental results validate the effectiveness of the DRM method in selecting CATE estimators that are robust to the distribution shift incurred by covariate shift and hidden confounders.

Unveiling the Potential of Robustness in Selecting Conditional Average Treatment Effect Estimators

Abstract

The growing demand for personalized decision-making has led to a surge of interest in estimating the Conditional Average Treatment Effect (CATE). Various types of CATE estimators have been developed with advancements in machine learning and causal inference. However, selecting the desirable CATE estimator through a conventional model validation procedure remains impractical due to the absence of counterfactual outcomes in observational data. Existing approaches for CATE estimator selection, such as plug-in and pseudo-outcome metrics, face two challenges. First, they must determine the metric form and the underlying machine learning models for fitting nuisance parameters (e.g., outcome function, propensity function, and plug-in learner). Second, they lack a specific focus on selecting a robust CATE estimator. To address these challenges, this paper introduces a Distributionally Robust Metric (DRM) for CATE estimator selection. The proposed DRM is nuisance-free, eliminating the need to fit models for nuisance parameters, and it effectively prioritizes the selection of a distributionally robust CATE estimator. The experimental results validate the effectiveness of the DRM method in selecting CATE estimators that are robust to the distribution shift incurred by covariate shift and hidden confounders.

1 Introduction

The escalating demand for decision-making has sparked an increasing interest in Causal Inference across various research domains, such as economics [23, 13, 43, 1], statistics [70, 49, 26, 41], healthcare [78, 27, 61, 9, 42], and financial application [11, 15, 37, 21, 35, 24, 50]. 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 [64], 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 [33], 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 3). 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 [66, 20, 53]. These works have proposed and summarized two types of criteria for CATE estimator selection, the plug-in metric R plug ˜ τ (ˆ τ ) and the pseudo-outcome metric R pseudo ˜ Y (ˆ τ ) :

One can establish a plug-in estimator ˜ τ or construct a pseudo-outcome estimator ˜ Y using the validation data to select CATE estimator ˆ τ . The previous studies [66, 20, 53] 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.

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 ˜ η 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 ˜ τ . 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 [20].

Challenge 2: These metrics are not well-targeted for selecting a robust CATE estimator. In potential outcome framework [64], the factual distribution P F and the counterfactual distribution P CF for t ∈ { 0 , 1 } can be defined as follows:

̸

The above (2) reveals that the covariate shift P ( X | T = t ) = 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.

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 ˆ V t (ˆ τ ) for t ∈ { 0 , 1 } , showing it decays to V t (ˆ τ ) 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.

2 Background of CATE Estimator Selection

Suppose the observational data contain n i.i.d. samples { ( x i , t i , y i ) } n i =1 , with the associated random variables being { ( X i , T i , Y i ) } n i =1 . For each unit i , X i ∈ X ⊂ R d is d -dimensional covariates and T i ∈ { 0 , 1 } is the binary treatment. Potential outcomes for treat ( T = 1 ) and control ( T = 0 ) are denoted by Y 1 , Y 0 ∈ Y ⊂ R . The observed (factual) outcome is Y = TY 1 + (1 -T ) Y 0 . The propensity score [63] is defined as π ( x ) := P ( T = 1 | X = x ) . The conditional mean potential outcome surface is defined as µ t ( x ) := E [ Y t | X = x ] for t ∈ { 0 , 1 } . The true CATE is defined as

Following the standard and necessary assumptions in potential outcome framework [64], we impose Assumption 2.1 that ensure treatment effects are identifiable.

Assumption 2.1 (Consistency, Overlap, and Unconfoundedness) . Consistency: If the treatment is t , then the observed outcome Y equals Y t . Overlap: The propensity score is bounded away from 0 to 1 , i.e., 0 < π ( x ) < 1 , ∀ x ∈ X . Unconfoundedness 3 : Y t ⊥ ⊥ T | X, ∀ t ∈ { 0 , 1 } .

The goal of CATE estimator selection is to select the best CATE estimator, denoted by ˆ τ best , from a set of J candidate estimators { ˆ τ 1 , . . . , ˆ τ J } :

Here, R oracle (ˆ τ ) is associated with E [(ˆ τ ( X ) -τ true ( X )) 2 ] , known as the Precision of Estimating Heterogeneous Effects (PEHE) w.r.t. ˆ τ [32, 68]. Note that R oracle (ˆ τ ) cannot be employed to evaluate CATE estimators' performances in real applications as we do not have access to τ true . Previous studies have introduced plug-in and pseudo-outcome metrics to aid in CATE estimator selection, as shown in equation (1). Then, the CATE estimator ˆ τ select is selected on validation data by

Notably, both the plug-in and pseudo-outcome metrics necessitate the fitting of nuisance parameters ˜ η (e.g., ˜ η = (˜ µ 1 , ˜ µ 0 , ˜ π ) ) using off-the-shelf ML models. While some papers like [16] address the selection of nuisance parameters for Averate Treatment Effect (ATE) estimators, e.g., the doubly robust estimator [23, 13, 36], our paper focuses on the selection of CATE estimators rather than nuisance parameters. For the plug-in metric, ˜ τ can be constructed using any CATE estimator discussed in Appendix A.1, yielding metrics such as plug-T, plug-DR, etc. For the pseudo-outcome metric, ˜ Y can be constructed using a specific formula discussed in Appendix A.2, yielding metrics such as pseudo-DR, pseudo-R, etc. The metrics based on the influence function [5] and the R-learner objective [55] are categorized into the pseudo-outcome metric. The categorization of plug-in and pseudo-outcome metrics maintains consistency with [20, 53].

3 Related Work

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 [47], U-learner [25, 55], DR-learner [41, 26], R-learner [55], and RA-learner [18]. The specific details of these meta-learners are stated in Appendix A.1. Additionally, some studies also focus on developing innovative causal ML

models for CATE estimation, such as Causal BART [30], Causal Forest [70, 8, 58], generative models like CEVAE [51] and GANITE [77], representation learning nets including SITE [76], TARNet [68], Dragonnet [69], FlexTENet [19], and HTCE [10], disentangled learning nets like D 2 VD [44, 45], DeR-CFR [74], and DR-CFR [31], and representation balancing nets such as BNN [39], CFRNet [68], DKLITE [79], IGNITE [29], BWCFR [6], DRRB [35], and DIGNet [38]. Recent surveys [28, 75, 56] have also conducted a systematic review of various causal inference methods.

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., R oracle (ˆ τ ) in equation (3). 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 [65] is constructed by the outcome predictors learned with representation balancing objective [68, 40]. Recent research [66, 20, 53] 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 A.2. 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 [2] find that placebo and structured empirical Monte Carlo methods are helpful for estimator selection under some restrictive conditions. In addition, researchers in [67, 7, 59] 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 [17].

4 The Distributionally Robust Metric

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 4.1). We then establish the DRM based on the distributionally robust value of PEHE (Section 4.2).

4.1 Capturing the Uncertainty in PEHE

Proposition 4.1. The PEHE w.r.t. the CATE estimator ˆ τ can be decomposed as follows:

Proposition 4.1 indicates that the PEHE is equal to four terms, where E [ˆ τ ( X ) 2 ] , E [ˆ τ ( X ) Y 0 ] , and E [ -ˆ τ ( X ) Y 1 ] depend on ˆ τ , while ζ is a constant that is independent of ˆ τ . The term E [ˆ τ ( X ) Y t ] for t ∈ { 0 , 1 } can be further decomposed as follows:

Equation (6a) can be computed empirically since the potential outcome Y t is observable in the group of T = t . However, equation (6b) is empirically uncomputable due to the unavailability of Y t in the group of T = 1 -t . The unknown term E [ˆ τ ( X ) Y t | T = 1 -t ] therefore determines the uncertainty in PEHE. To capture such an uncertainty, we therefore establish distributionally robust values for E [ˆ τ ( X ) Y 0 | T = 1] and E [ -ˆ τ ( X ) Y 1 | T = 0] based on a Kullback-Leibler (KL) ambiguity set.

Definition 4.2 (KL ambiguity set) . Given two distributions Q and P and the ambiguity radius ϵ > 0 . The KL ambiguity (uncertainty) set B ϵ ( P ) is defined as

Here, D KL ( Q || P ) denotes the KL divergence of some arbitrary distribution Q from the reference distribution P . Now we define the distribution of ( X,Y 0 , Y 1 ) in the treated and controlled groups as

By setting an adequately large ambiguity radius in Definition 4.2, the following inequalities hold for E [ˆ τ ( X ) Y 0 | T = 1] = E P T [ˆ τ ( X ) Y 0 ] and E [ -ˆ τ ( X ) Y 1 | T = 0] = E P C [ -ˆ τ ( X ) Y 1 ] :

To provide a clearer understanding, let us consider the example of E P T [ˆ τ ( X ) Y 0 ] . Since the term E [ˆ τ ( 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 E P T [ˆ τ ( X ) Y 0 ] w.r.t. ˆ τ . In other words, the value of the uncomputable quantity E P T [ˆ τ ( X ) Y 0 ] will be at most V 0 (ˆ τ ) . Similarly, the value of the uncomputable quantity E P C [ -ˆ τ ( X ) Y 1 ] will be at most V 1 (ˆ τ ) . 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 4.3, which measures the uncertainty in PEHE.

Corollary 4.3. Let V 0 (ˆ τ ) and V 1 (ˆ τ ) be the quantities defined in equation (9) , ζ be the constant given in Proposition 4.1, u 1 := P ( T = 1) , and u 0 = 1 -u 1 = P ( T = 0) . The distributionally robust value of PEHE w.r.t. ˆ τ is defined as V PEHE (ˆ τ ) such that

4.2 Establishing Distributionally Robust Metric

As Corollary 4.3 provides the distributionally robust (worst-case) value of PEHE, it can naturally measure the robustness of the CATE estimator ˆ τ against distribution shift between counterfactual distribution and factual distribution. In this section, we will provide two steps involved in using Corollary 4.3 to construct the DRM method for CATE estimator selection.

Step 1: Establishing computational tractability of V t (ˆ τ ) . The distributionally robust values V 0 (ˆ τ ) and V 1 (ˆ τ ) in equation (10) are initially defined as supremum problems over infinite support, presenting a substantial computational challenge. Theorem 4.4 reformulates the infeasible supremum problems into tractable minimum problems.

Theorem 4.4. The distributionally robust values V 0 (ˆ τ ) and V 1 (ˆ τ ) in equation (9) are equivalent to

In the finite-sample scenario, V 0 (ˆ τ ) and V 1 (ˆ τ ) can be empirically approximated as follows:

Algorithm 1 Using DRM for CATE Estimator Selection

Input: The candidate CATE estimators { ˆ τ 1 , . . . , ˆ τ J } . The validation dataset with n i.i.d. observational samples { ( X i , T i , Y i ) } n i =1 . The number of iterations K . The initialization λ (0) 0 and λ (0) 1 . The ambiguity radius ϵ 0 and ϵ 1 .

1: for j = 1 to J do

2: for k = 0 to K -1 do

3: Compute ˆ F t ( λ ( k ) t , ϵ t ; ˆ τ j ) for t ∈ { 0 , 1 } by equation (14a).

4: Compute ∂ ˆ F t ( λ ( k ) t , ϵ t ; ˆ τ j ) /∂λ ( k ) t for t ∈ { 0 , 1 } by equation (14b).

5: λ ( k +1) t ← max { λ ( k ) t -ˆ F t ( λ ( k ) t , ϵ t ; ˆ τ j ) / ( ∂ ˆ F t ( λ ( k ) t , ϵ t ; ˆ τ j ) /∂λ ( k ) t ) , 0 } for t ∈ { 0 , 1 } .

6: Save ˆ V t (ˆ τ j )[ k ] = ˆ F t ( λ ( k +1) t , ϵ t ; ˆ τ j ) for t ∈ { 0 , 1 } .

7: Return ˆ V t (ˆ τ j ) = arg min k ∈{ 0 ,...,K -1 } ˆ V t (ˆ τ j )[ k ] for t ∈ { 0 , 1 } .

8: Use ˆ V 0 (ˆ τ j ) and ˆ V 1 (ˆ τ j ) to compute R DRM (ˆ τ j )

ˆ τ ∈{ ˆ τ 1 ,..., ˆ τ J }

by equation (15). Output: ˆ τ select = arg min R DRM (ˆ τ ) .

Note that in equation (12), 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 2.1. We then provide a finite-sample analysis of the gap between ˆ V t (ˆ τ ) and V t (ˆ τ ) in the following Theorem 4.5, which suggests the gap decays at a rate of n -1 / 2 .

Theorem 4.5. Let u t := P ( T = t ) for t ∈ { 0 , 1 } . Assume 0 < ¯ λ ≤ λ 0 , λ 1 ≤ ¯ λ and ˆ τ ( X ) Y is bounded within the range of ¯ M to ¯ M . Define C exp = 1 { ¯ M ≤ ¯ M ≤ 0 } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) + 1 { ¯ M ≤ 0 , ¯ M ≥ 0 } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) + 1 { 0 ≤ ¯ M ≤ ¯ M } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) . For n ≥ 2 /u 2 log(2 /δ ) and t ∈ { 0 , 1 } , with probability 1 -δ , we have

Step 2: Finalizing Distributionally Robust Metric for CATE estimator selection. We first define two functions that are useful in obtaining V 0 (ˆ τ ) and V 1 (ˆ τ ) :

Here, Z denotes ˆ τ ( X ) Y for notational simplicity. We then use the Newton-Raphson method to find the empirical solution for ˆ V t (ˆ τ ) , exploiting the convexity of ˆ F t ( λ t , ϵ t ; ˆ τ ) w.r.t. λ t . Based on the distributionally robust value of PEHE, i.e., V PEHE (ˆ τ ) in equation (10), we finally obtain the selected estimator ˆ τ select = arg min ˆ τ ∈{ ˆ τ 1 ,..., ˆ τ J } R DRM (ˆ τ ) such that

Algorithm 1 provides complete procedure of using the DRM method for CATE estimator selection.

Discussion on the ambiguity radius ϵ . The ambiguity radius ϵ plays a critical role in real-world applications [54, 52, 60]. However, determining an appropriate value for ϵ can be challenging as it requires striking a balance between ensuring the bound in equation (9) holds and maintaining its tightness. Specifically, if ϵ 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 4.3 can hold). On the other hand, if ϵ is set too large, even though the ambiguity set can encompass more

distributions to ensure the counterfactual distribution is contained, the bound in Corollary 4.3 can be less tight. In general, selecting a proper ambiguity radius is an open problem in distributioanlly robust optimization (DRO) literature [34, 54, 46, 48, 72].

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 ϵ ∗ 1 = D KL ( P C || P T ) and ϵ ∗ 0 = D KL ( P T || P C ) , which ensures that the bound in Corollary 4.3 holds and is tight. However, as defined in equation (8), 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 4.6 provides an intriguing alternative approach to acquire D KL ( P C || P T ) and D KL ( P T || P C ) when unconfoundedness in Assumption 2.1 is satisfied.

Proposition 4.6. Let P T X := P ( X | T = 1) and P C X := P ( X | T = 0) denote the covariates distribution in the treat and control group, respectively. Assuming that random variables ( X,T,Y 1 , Y 0 ) satisfy the unconfoundedness in Assumption 2.1, we have

The proof is deferred to Appendix B.2.

Proposition 4.6 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 ϵ ∗ 1 = D KL ( P C X || P T X ) and ϵ ∗ 0 = D KL ( P T X || P C X ) . While the KL divergence can be approximated using empirical algorithm (e.g, Nearest-Neighbor [73, 57]), we recommend setting the ambiguity radius larger than the empirically approximated KL divergence (see specific explanations in Appendix C.1). 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 1 involves approximating ϵ ∗ 1 = D KL ( P C X || P T X ) and ϵ ∗ 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.

5 Experiments

5.1 Experimental Setup.

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 [19, 20]. 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 A.1. 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 NearestNeighbor Matching [62], and our proposed DRM. The specific details of baseline selectors are stated in Appendix A.2. We employ the eXtreme Gradient Boosting (XGB) [12] 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 [20]; iii) aligning with [5] where XGB is used for their proposed pseudo-IF metric. The details of hyperparameters for nuisance models are stated in Section C.2 of Appendix.

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 [19, 20], we collect the covariates with n = 4802 data points from ACIC2016 dataset [22]. Then, we generate treatment with T i | X i ∼ Bern (1 / (1 + exp( -ξ ( β ′ 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: β T , β j , β j,k , β j,k,l ∼ Bern (0 . 2) , γ j ∼ Bern ( ρ ) , and ϵ i ∼ N (0 , 0 . 1) . The parameter ξ in treatment assignment represents the level of selection bias, and the parameter ρ in γ 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%.

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 [19, 20], 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 ρ vary in { 0 , 0 . 1 , 0 . 3 } with fixing ξ = 1 . Setting B: With the unconfoundedness assumption, let ξ vary in { 0 , 1 , 2 } with fixing ρ = 0 . 1 . Setting C : Without unconfoundedness assumption, fix ρ = 0 . 1 and ξ = 1 . Then randomly remove ⌊ m · d ⌋ covariates such that the dimension of observed covariates is d -⌊ m · d ⌋ , 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.

Comparison criteria. The CATE estimator ˆ τ is believed better if it achieves a smaller difference between R oracle (ˆ τ ) and R oracle (ˆ τ best ) , where ˆ τ best is the actual best estimator in equation (3). 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 R oracle (ˆ τ ) and the rank order determined by each selector. All the reported values (Mean ± Standard deviation) are computed over 100 runs.

5.2 Experimental Results

Regret comparison. The results presented in Table 1 demonstrate consistently good performance from the DRM selector across various settings. In setting A, the DRM selector outperforms other selectors as the CATE complexity ( ρ ) varies. Additionally, Plug-R, Plug-S, and Plug-PS also perform well in terms of the Regret criterion, which aligns with prior findings in [66] 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., R oracle (ˆ τ select ) ) of different selectors in Table 3 of Section C.3. The results indicate that Plug-R, Plug-S, and Plug-PS tend to

SettingA(p=0)

Setting A (p=0.1)

Setting A(p=0.3)

SettingB(E=0)

Matching

DRM

Matching

DRM

Matching

DRM

Matching

DRM

Random

Fact

Random

Fact

Random

Fact

Random

Fact

Pseudo-IF

Pseudo-IF

Pseudo-IF

Pseudo-IF

Pseudo-R

Pseudo-R

Pseudo-DR

Plug-RA

Plug-R

Pseudo-DR

PlUg-RA

Pseudo-DR

Plug-RA

Plug-R

Pseudo-DR

Plug-RA

Pldg-R

Plug-IPW

Plug-DR

Plug-IPW

Plug-IPW

Plug-DR

Plug-IPW

Plug-DR

Plug-X

Plug-X

Plug-X

Plug-X

P

L-onld

T-bnld

P

T-bnld

PU

Plug-T

n-onld

n-bnld

n-bnId

n-bnId

0.1.2.3.4.5.6.7.8.91

0.1.2.3.4.5.6.7.8.91

0.1.2.3.4.5.6.7.8.91

0.1.2.3.4.5.6.7.8.91

Ratio

Ratio

Ratio

Ratio

Setting B(=2)

Setting C (m =0.1)

Setting C(m=0.5)

Setting C(m=0.9)

Matching

Random

DRM

Fact

Matching

Random

DRM

Fact

Matching

Random

DRM

Fact

Matching

Random

DRM

Pseudo-IF

Pseudo-IF

Pseudo-IF

Pseudo-IF

Pseudo-DR

Pseudo-R

Plug-RA

Pldg-R

Pseudo-DR

Pseudo-R

Plug-RA

Plug-R

Pseudo-DR

Pseudo-R

Pseudo-DR

Pseudo-R

PlUg-RA

Pldg-R

Plug-DR

Plug-DR

Plug-DR

Plug-DR

Plug-IPW

Plug-X

Plug-IPW

Plug-X

Plug-IPW

Plug-X

Plug-IPW

Plug-X

P

L-onld

P

F-onid

Plug-T

n-onld

O-bnld

n-bnid

0-6nid

0.1.2.3.4.5.6.7.8.91

0.1.2.3.4.5.6.7.8.91

0.1.2.3.4.5.6.7.8.9 1

0.1.2.3.4.5.6.7.8.91

Ratio

Ratio

Ratio

Ratio

Is the Rank 1-3

Is the Rank 4-11

Is the Rank 12-19

Is the Rank 20-27

Is theRank28-36

exhibit better PEHE as the CATE complexity decreases, aligning with the findings in [20]. In setting B, the DRM selector demonstrates robustness against selection bias (controlled by ξ ) compared to many baselines. However, for the case ξ = 2 , DRM selects a poor estimator 1 or 2 times out of 100 experiments, as shown in Figure 1. Although this weakens its overall performance, DRM still outperforms many baselines in this scenario. In the scenario ξ = 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.

Ranking ability. In Table 2, 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.

Variance analysis. Table 1 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 R oracle (ˆ τ ) values, resulting in the sorted list: [ R oracle (ˆ τ 1 ) , . . . , R oracle (ˆ τ 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 1 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.

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 [3], 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.

6 Conclusion

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 ˆ V t (ˆ τ ) and V t (ˆ τ ) reduces at a rate of n -1 / 2 for t ∈ { 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.

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 [14]. Here, we outline some challenges and suggest future research directions. First, while Proposition 4.6 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 5.2, 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 [4]. Finally, inspired by [16], 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.

References

Alberto Abadie, Susan Athey, Guido W Imbens, and Jeffrey M Wooldridge. When should you adjust standard errors for clustering? The Quarterly Journal of Economics , 138(1):1-35, 2023.

Arun Advani, Toru Kitagawa, and Tymon Słoczy´ nski. Mostly harmless simulations? using monte carlo studies for estimator selection. Journal of Applied Econometrics , 34(6):893-910, 2019.

Ahmed Alaa and Mihaela Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design. In International Conference on Machine Learning , pages 129-138. PMLR, 2018.

Ahmed Alaa and Mihaela van der Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning , volume 80 of Proceedings of Machine Learning Research , pages 129-138. PMLR, 10-15 Jul 2018.

Ahmed Alaa and Mihaela Van Der Schaar. Validating causal inference models via influence functions. In International Conference on Machine Learning , pages 191-201. PMLR, 2019.

Serge Assaad, Shuxi Zeng, Chenyang Tao, Shounak Datta, Nikhil Mehta, Ricardo Henao, Fan Li, and Lawrence Carin. Counterfactual representation learning with balancing weights. In International Conference on Artificial Intelligence and Statistics , pages 1972-1980. PMLR, 2021.

Susan Athey, Guido W Imbens, Jonas Metzger, and Evan Munro. Using wasserstein generative adversarial networks for the design of monte carlo simulations. Journal of Econometrics , 2021.

SUSAN ATHEY, JULIE TIBSHIRANI, and STEFAN WAGER. Generalized random forests. The Annals of Statistics , 47(2):1148-1178, 2019.

Ioana Bica, Ahmed M Alaa, Craig Lambert, and Mihaela Van Der Schaar. From real-world patient data to individualized treatment effects using machine learning: current and future methods to address underlying challenges. Clinical Pharmacology & Therapeutics , 109(1):87100, 2021.

Ioana Bica and Mihaela van der Schaar. Transfer learning on heterogeneous feature spaces for treatment effects estimation. Advances in Neural Information Processing Systems , 35:3718437198, 2022.

Léon Bottou, Jonas Peters, Joaquin Quiñonero-Candela, Denis X Charles, D Max Chickering, Elon Portugaly, Dipankar Ray, Patrice Simard, and Ed Snelson. Counterfactual reasoning and learning systems: The example of computational advertising. Journal of Machine Learning Research , 14(11), 2013.

Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd acm sigkdd international conference on knowledge discovery and data mining , pages 785-794, 2016.

Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiased machine learning for treatment and structural parameters, 2018.

Victor Chernozhukov, Mert Demirer, Esther Duflo, and Ivan Fernandez-Val. Generic machine learning inference on heterogeneous treatment effects in randomized experiments, with an application to immunization in india. Technical report, National Bureau of Economic Research, 2018.

Zhixuan Chu, Stephen L Rathbun, and Sheng Li. Graph infomax adversarial learning for treatment effect estimation with networked observational data. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining , pages 176-184, 2021.

Yifan Cui and EJ Tchetgen Tchetgen. Selective machine learning of doubly robust functionals. Biometrika , 111(2):517-535, 2024.

Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating cate? a critical look at ml benchmarking practices in treatment effect estimation. In Thirty-fifth conference on neural information processing systems datasets and benchmarks track (round 2) , 2021.

Alicia Curth and Mihaela van der Schaar. Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics , pages 1810-1818. PMLR, 2021.

Alicia Curth and Mihaela van der Schaar. On inductive biases for heterogeneous treatment effect estimation. Advances in Neural Information Processing Systems , 34:15883-15894, 2021.

Alicia Curth and Mihaela Van Der Schaar. In search of insights, not magic bullets: Towards demystification of the model selection dilemma in heterogeneous treatment effect estimation. In Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 6623-6642. PMLR, 23-29 Jul 2023.

Robert Donnelly, David Blei, Susan Athey, et al. Correction to: Counterfactual inference for consumer choice across many product categories. Quantitative Marketing and Economics , 19(3-4):409-409, 2021.

Vincent Dorie, Jennifer Hill, Uri Shalit, Marc Scott, and Dan Cervone. Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition. Statistical Science , 34(1):43-68, 2019.

Max H Farrell. Robust inference on average treatment effects with possibly more covariates than observations. Journal of Econometrics , 189(1):1-23, 2015.

Carlos Fernández-Loría, Foster Provost, Jesse Anderton, Benjamin Carterette, and Praveen Chandar. A comparison of methods for treatment assignment with an application to playlist generation. Information Systems Research , 34(2):786-803, 2023.

Aaron Fisher. Inverse-variance weighting for estimation of heterogeneous treatment effects. In Forty-first International Conference on Machine Learning .

Dylan J Foster and Vasilis Syrgkanis. Orthogonal statistical learning. The Annals of Statistics , 51(3):879-908, 2023.

Jared C Foster, Jeremy MG Taylor, and Stephen J Ruberg. Subgroup identification from randomized clinical trial data. Statistics in medicine , 30(24):2867-2880, 2011.

Ruocheng Guo, Lu Cheng, Jundong Li, P Richard Hahn, and Huan Liu. A survey of learning causality with data: Problems and methods. ACM Computing Surveys (CSUR) , 53(4):1-37, 2020.

Ruocheng Guo, Jundong Li, Yichuan Li, K Selçuk Candan, Adrienne Raglin, and Huan Liu. Ignite: A minimax game toward learning individual treatment effects from networked observational data. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence , pages 4534-4540, 2021.

P Richard Hahn, Jared S Murray, and Carlos M Carvalho. Bayesian regression tree models for causal inference: Regularization, confounding, and heterogeneous effects (with discussion). Bayesian Analysis , 15(3):965-1056, 2020.

Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. In International Conference on Learning Representations , 2019.

Jennifer L Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics , 20(1):217-240, 2011.

Paul W Holland. Statistics and causal inference. Journal of the American statistical Association , 81(396):945-960, 1986.

Zhaolin Hu and L Jeff Hong. Kullback-leibler divergence constrained distributionally robust optimization. Available at Optimization Online , 1(2):9, 2013.

Yiyan Huang, Cheuk Hang Leung, Shumin Ma, Zhiri Yuan, Qi Wu, Siyi Wang, Dongdong Wang, and Zhixiang Huang. Towards balanced representation learning for credit policy evaluation. In International Conference on Artificial Intelligence and Statistics , pages 3677-3692. PMLR, 2023.

Yiyan Huang, Cheuk Hang Leung, Qi Wu, Xing Yan, Shumin Ma, Zhiri Yuan, Dongdong Wang, and Zhixiang Huang. Robust causal learning for the estimation of average treatment effects. In 2022 International Joint Conference on Neural Networks (IJCNN) , pages 1-9. IEEE, 2022.

Yiyan Huang, Cheuk Hang Leung, Xing Yan, Qi Wu, Nanbo Peng, Dongdong Wang, and Zhixiang Huang. The causal learning of retail delinquency. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 204-212, 2021.

Yiyan Huang, WANG Siyi, Cheuk Hang Leung, WU Qi, WANG Dongdong, and Zhixiang Huang. Dignet: Learning decomposed patterns in representation balancing for treatment effect estimation. Transactions on Machine Learning Research , 2024.

Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning , pages 3020-3029. PMLR, 2016.

Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. The Journal of Machine Learning Research , 23(1):7489-7538, 2022.

Edward H Kennedy. Towards optimal doubly robust estimation of heterogeneous causal effects. Electronic Journal of Statistics , 17(2):3008-3049, 2023.

Newton Mwai Kinyanjui and Fredrik D Johansson. Adcb: An alzheimer's disease simulator for benchmarking observational estimators of causal effects. In Conference on Health, Inference, and Learning , pages 103-118. PMLR, 2022.

Toru Kitagawa and Aleksey Tetenov. Who should be treated? empirical welfare maximization methods for treatment choice. Econometrica , 86(2):591-616, 2018.

Kun Kuang, Peng Cui, Bo Li, Meng Jiang, Shiqiang Yang, and Fei Wang. Treatment effect estimation with data-driven variable decomposition. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 31, 2017.

Kun Kuang, Peng Cui, Hao Zou, Bo Li, Jianrong Tao, Fei Wu, and Shiqiang Yang. Data-driven variable decomposition for treatment effect estimation. IEEE Transactions on Knowledge and Data Engineering , 34(5):2120-2134, 2020.

Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh ShafieezadehAbadeh. Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations research & management science in the age of analytics , pages 130-166. Informs, 2019.

Sören R Künzel, Jasjeet S Sekhon, Peter J Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the national academy of sciences , 116(10):4156-4165, 2019.

Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems , 33:8847-8860, 2020.

Shuangning Li and Stefan Wager. Random graph asymptotics for treatment effect estimation under network interference. The Annals of Statistics , 50(4):2334-2358, 2022.

Yijun Li, Cheuk Hang Leung, Xiangqian Sun, Chaoqun Wang, Yiyan Huang, Xing Yan, Qi Wu, Dongdong Wang, and Zhixiang Huang. The causal impact of credit lines on spending distributions. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 38, pages 180-187, 2024.

Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. Advances in neural information processing systems , 30, 2017.

Shumin Ma, Cheuk Hang Leung, Qi Wu, Wei Liu, and Nanbo Peng. Understanding distributional ambiguity via non-robust chance constraint. In Proceedings of the First ACM International Conference on AI in Finance , pages 1-8, 2020.

Divyat Mahajan, Ioannis Mitliagkas, Brady Neal, and Vasilis Syrgkanis. Empirical analysis of model selection for heterogenous causal effect estimation. International Conference on Learning Representations , 2024.

Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: performance guarantees and tractable reformulations. Mathematical Programming , 171(1-2):115-166, 2018.

Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika , 108(2):299-319, 2021.

Ana Rita Nogueira, Andrea Pugnana, Salvatore Ruggieri, Dino Pedreschi, and João Gama. Methods and tools for causal discovery and causal inference. Wiley interdisciplinary reviews: data mining and knowledge discovery , 12(2):e1449, 2022.

Yung-Kyun Noh, Masashi Sugiyama, Song Liu, Marthinus C Plessis, Frank Chongwoo Park, and Daniel D Lee. Bias reduction and metric learning for nearest-neighbor estimation of kullback-leibler divergence. In Artificial Intelligence and Statistics , pages 669-677. PMLR, 2014.

Miruna Oprescu, Vasilis Syrgkanis, and Zhiwei Steven Wu. Orthogonal random forest for causal inference. In International Conference on Machine Learning , pages 4932-4941. PMLR, 2019.

Harsh Parikh, Carlos Varjao, Louise Xu, and Eric Tchetgen Tchetgen. Validating causal inference methods. In International Conference on Machine Learning , pages 17346-17358. PMLR, 2022.

Georg Ch Pflug. Multistage stochastic decision problems: Approximation by recursive structures and ambiguity modeling. European Journal of Operational Research , 306(3):1027-1039, 2023.

Zhaozhi Qian, Yao Zhang, Ioana Bica, Angela Wood, and Mihaela van der Schaar. Synctwin: Treatment effect estimation with longitudinal outcomes. Advances in Neural Information Processing Systems , 34:3178-3190, 2021.

Craig A Rolling and Yuhong Yang. Model selection for estimating treatment effects. Journal of the Royal Statistical Society Series B: Statistical Methodology , 76(4):749-769, 2014.

Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika , 70(1):41-55, 1983.

Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association , 100(469):322-331, 2005.

Yuta Saito and Shota Yasui. Counterfactual cross-validation: Stable model selection procedure for causal inference models. In International Conference on Machine Learning , pages 83988407. PMLR, 2020.

Alejandro Schuler, Michael Baiocchi, Robert Tibshirani, and Nigam Shah. A comparison of methods for model selection when estimating individual treatment effects. arXiv preprint arXiv:1804.05146 , 2018.

Alejandro Schuler, Ken Jung, Robert Tibshirani, Trevor Hastie, and Nigam Shah. Synthvalidation: Selecting the best causal inference method for a given dataset. arXiv preprint arXiv:1711.00083 , 2017.

Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International conference on machine learning , pages 3076-3085. PMLR, 2017.

Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems , 32, 2019.

Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the American Statistical Association , 113(523):1228-1242, 2018.

Chi Wang, Qingyun Wu, Markus Weimer, and Erkang Zhu. Flaml: A fast and lightweight automl library. Proceedings of Machine Learning and Systems , 3:434-447, 2021.

Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization. arXiv preprint arXiv:2109.11926 , 2021.

Qing Wang, Sanjeev R Kulkarni, and Sergio Verdú. A nearest-neighbor approach to estimating divergence between continuous random vectors. In 2006 IEEE International Symposium on Information Theory , pages 242-246. IEEE, 2006.

Anpeng Wu, Junkun Yuan, Kun Kuang, Bo Li, Runze Wu, Qiang Zhu, Yueting Zhuang, and Fei Wu. Learning decomposed representations for treatment effect estimation. IEEE Transactions on Knowledge and Data Engineering , 35(5):4989-5001, 2022.

Liuyi Yao, Zhixuan Chu, Sheng Li, Yaliang Li, Jing Gao, and Aidong Zhang. A survey on causal inference. ACM Transactions on Knowledge Discovery from Data (TKDD) , 15(5):1-46, 2021.

Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. Advances in neural information processing systems , 31, 2018.

Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International conference on learning representations , 2018.

Linying Zhang, Yixin Wang, Anna Ostropolets, Jami J Mulgrave, David M Blei, and George Hripcsak. The medical deconfounder: assessing treatment effects with electronic health records. In Machine Learning for Healthcare Conference , pages 490-512. PMLR, 2019.

Yao Zhang, Alexis Bellot, and Mihaela Schaar. Learning overlapping representations for the estimation of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics , pages 1005-1014. PMLR, 2020.

Appendix

A CATE Estimation Strategies

A.1 CATE Learners

We now detail how to construct CATE learners using the observed samples { ( X i , T i , Y i ) } n i =1 . Note that CATE learners are learned on the training set, so the sample size n here equals the training sample size. Denote n t by the sample size in the treat group, and n c by the sample size in the control group such that n = n t + n c .

S-learner: Let predictors= ( X,T ) , response= Y . Train a model ˆ µ ( X,T ) . Then we obtain ˆ τ S ( X ) :

T-learner: Let predictors= X T (covariates in the treat), response= Y T (outcome in the treat). Train a model ˆ µ 1 ( X ) . Let predictors= X C (covariates in the control), response= Y C (outcome in the control). Train a model ˆ µ 0 ( X ) . Then we obtain ˆ τ T ( X ) :

PS-learner: Fisrt-step: Train ˆ τ S ( X ) using the above-mentioned step in S-learner. Secondstep: Let predictors= X , response= ˆ τ S ( X ) . Train a model ˆ τ PS ( X ) from the following objective:

IPW-learner: First-step: let predictors= X , response= T . Train a propensity score model ˆ π ( X ) . Construct surrogate of CATE using pseudo-outcomes with inverse propensity weighting (IPW) formula: Y 1 , 0 IPW = Y 1 IPW -Y 0 IPW , where Y 1 IPW = TY ˆ π ( X ) and Y 0 IPW = (1 -T ) Y 1 -ˆ π ( X ) . Train a model ˆ τ IPW ( X ) from the following objective:

X-learner [47]: First-step: Train ˆ µ 1 ( X ) and ˆ µ 0 ( X ) using the the above-mentioned procedure in T-learner. Train a propensity score model ˆ π ( X ) using the the above-mentioned procedure in IPW-learner. Second-step: Let predictors= X T , response= ˆ µ 1 ( X T ) -Y T , and predictors= X C , response= ˆ µ 0 ( X C ) -Y C . Obtain a model ˆ τ X ( X ) by learning two separate functions ˆ τ 1 X ( X ) and ˆ τ 0 X ( X ) :

U-learner [25, 55]: First-step: Let predictors= X , response= Y . Train a model ˆ µ ( X ) to approximate the conditional mean outcome E [ Y | X ] . Train a propensity score model ˆ π ( X ) using the the above-mentioned procedure in IPW-learner. Second-step: Compute the outcome residual ξ = Y -ˆ µ ( X ) and treatment residual ν = T -ˆ π ( X ) . Train a model ˆ τ U ( X ) from the following objective:

DR-learner [41, 26]: First-step: Train ˆ µ 1 ( X ) and ˆ µ 0 ( X ) using the the above-mentioned procedure in T-learner. Train a propensity score model ˆ π ( X ) using the the abovementioned procedure in IPW-learner. Second-step: Construct surrogate of CATE using pseudo-outcomes with doubly robust (DR) formula: Y 1 , 0 DR = Y 1 DR -Y 0 DR , where Y 1 DR = ˆ µ 1 ( X ) + T ˆ π ( X ) ( Y -ˆ µ 1 ( X )) and Y 0 DR = ˆ µ 0 ( X ) + 1 -T 1 -ˆ π ( X ) ( Y -ˆ µ 0 ( X )) . Train a model ˆ τ DR ( X ) from the following objective:

R-learner [55]: First-step: Let predictors= X , response= Y . Train a model ˆ µ ( X ) to approximate the conditional mean outcome E [ Y | X ] . Train a propensity score model ˆ π ( X ) using the the above-mentioned procedure in IPW-learner. Second-step: Compute the outcome residual ξ = Y -ˆ µ ( X ) and treatment residual ν = T -ˆ π ( X ) . Train a model ˆ τ R ( X ) from the following objective:

RA-learner [18]: First-step: Train ˆ µ 1 ( X ) and ˆ µ 0 ( X ) using the the above-mentioned procedure in T-learner. Second-step: Construct surrogate of CATE using pseudo-outcomes with regression adjustment (RA) formula: Y RA = T ( Y -ˆ µ 0 ( X )) + (1 -T )(ˆ µ 1 ( X ) -Y ) . Train a model ˆ τ RA ( X ) from the following objective:

A.2 CATE Selectors

We now detail how to construct CATE selectors using the observed samples { ( X i , T i , Y i ) } n i =1 . Note that CATE selectors are constructed on the validation set, so the sample size n here equals the validation sample size.

Plug-in selector: Obtain any CATE learners ˜ τ using the observational validation data. Then plug-in ˜ τ into the following metric R plug ˜ τ (ˆ τ ) :

For each plug-in selector ˜ τ , the selected j ∗ -th CATE estimator is ˆ τ j ∗ , where j ∗ = arg min j ∈{ 1 ,...,J } R plug ˜ τ (ˆ τ j ) .

Pseudo-outcome selector:

Pseudo-DR: Utilize validation data to estimate nuisance parameters (˜ µ 1 , ˜ µ 0 , ˜ π ) , following the procedure described in Section A.1. ˜ Y DR = ˜ Y 1 DR -˜ Y 0 DR , where ˜ Y 1 DR = ˜ µ 1 ( X ) + T ˜ π ( X ) ( Y -˜ µ 1 ( X )) and ˜ Y 0 DR = ˜ µ 0 ( X ) + 1 -T 1 -˜ π ( X ) ( Y -˜ µ 0 ( X )) . Then the pseudo-DR metric is

For pseudo-DR selector, the selected j ∗ -th CATE estimator is ˆ τ j ∗ , where j ∗ = arg min j ∈{ 1 ,...,J } R pseudo DR (ˆ τ j ) .

Pseudo-R: Utilize validation data to estimate nuisance parameters (˜ µ, ˜ π ) , following the procedure described in Section A.1. Then the pseudo-R metric is

For pseudo-R selector, the selected j ∗ -th CATE estimator is ˆ τ j ∗ , where j ∗ = arg min j ∈{ 1 ,...,J } R pseudo R (ˆ τ j ) .

Pseudo-IF [5]: Utilize validation data to estimate nuisance parameters (˜ µ 1 , ˜ µ 0 , ˜ π ) , following the procedure described in Section A.1. Let ˜ τ ( X ) = (˜ µ 1 ( X ) -˜ µ 0 ( X )) . Then the pseudo-IF metric is

For pseudo-IF selector, the selected j ∗ -th CATE estimator is ˆ τ j ∗ , where j ∗ = arg min j ∈{ 1 ,...,J } R pseudo IF (ˆ τ j ) .

Other pseudo-outcome selector: By manipulating the formula of ˜ Y , it is possible to create additional pseudo-outcome selectors, such as the pseudo-IPW selector. In our paper, we choose pseudo-DR as the baseline because it is representative in the causal inference literature and it often demonstrates superior performance, owing to its doubly robust property.

B Proofs

B.1 Proof of Proposition 4.1

Proof.

B.2 Proof of Proposition 4.6

The following Proposition B.1 is useful in proving Proposition 4.6.

Proposition B.1. Assuming the random variable tuple ( X,T,Y 1 , Y 0 ) satisfies Assumption 2.1, we have

Proof.

Now we can prove Proposition 4.6.

Proof.

B.3 Proof of Theorem 4.4

Lemma B.2 (Theorem 1 in [34]) . Let f θ ( X ) denote the loss function of X and it is bounded almost surely. θ ∈ Θ represents the model parameters of the function f θ ( X ) . Let B ϵ ( P ) be the uncertainty ball centered at distribution P with ambiguity radius ϵ . Define κ as the mass of the distribution P on its essential supremum (Proposition 2 in [34]). Assume f θ ( X ) is bounded and log κ + ϵ < 0 , then we have

Our Theorem 4.4 follows by directly applying the above Lemma B.2.

B.4 Proof of Theorem 4.5

For notational simplicity, we denote W = ( X,T,Y ) ∈ W and Z = ˆ τ ( X ) Y . Assume Z is bounded within the range M and ¯ M . Define the following functions:

KL T C KL X X

Then we have the following lemma that guarantees the convergence for ˆ G 0 ( λ 0 ; W ) and ˆ G 1 ( λ 1 ; W ) . Lemma B.3. Assume 0 < ¯ λ ≤ λ 0 , λ 1 ≤ ¯ λ , and ˆ τ ( X ) Y is bounded within the range of ¯ M to ¯ M . Then with probability 1 -δ , we have

Proof. Denote h 0 ( W 1 , W 2 , . . . , W n ) = 1 n ∑ n i =1 g 0 ( λ 0 ; W i ) . We notice that h 0 ( W 1 , W 2 , . . . , W n ) satisfies the bounded difference inequality:

Note that | ˆ G 0 ( λ 0 ; W ) -G 0 ( λ 0 ; W ) | = | h 0 ( W 1 , W 2 , . . . , W n ) -E [ h 0 ( W 1 , W 2 , . . . , W n )] | . Then using McDiarmid's inequality, for any ϵ > 0 , we have

P

(∣

∣

=

∣

G

0

(

λ

P

(

h

|

0

;

0

(

W

W



2 exp

For some δ > 0 , we have

This solves ϵ such that



1

)

-

, W

-

2

0

0

G

(

λ

;

, . . . , W

2

n

(

ϵ

2

2 exp

(

M/λ

¯

n

0

ˆ

W

n

)

)

)

)

∣

∣

∣

-

E

ϵ

)

[



h

0

1

(

W

= 2exp

2

, W

2

(

2

, . . . , W

n

-

nϵ

(

exp

(

)]

2

¯

M/λ

0

))

ϵ

2

)

)

.

The above inequality should hold for any λ 0 such that 0 < λ ≤ λ 0 ≤ ¯ λ . Therefore, we have

Similarly, denote h 1 ( W 1 , W 2 , . . . , W n ) = 1 n ∑ n i =1 g 1 ( λ 1 ; W i ) . We note that h 1 ( W 1 , W 2 , . . . , W n ) satisfies the bounded difference inequality:

Then using McDiarmid's inequality, for any ϵ > 0 , we have

For some δ > 0 , we have

This solves ϵ such that

The above inequality should hold for any λ such that 0 < λ ≤ λ ≤ ¯ λ . Therefore, we have

In the following content, we will bound terms ∣ ∣ ∣ log( ˆ G 0 ( λ 0 ; W )) -log ( G 0 ( λ 0 ; W )) ∣ ∣ ∣ and ∣ ∣ ∣ log( ˆ G 1 ( λ 1 ; W )) -log ( G 1 ( λ 1 ; W )) ∣ ∣ ∣ . Lemma B.4 is useful for bounding these two terms.

Lemma B.4. Let c be a constant. For any x 1 , x 2 such that x 1 , x 2 ≥ c > 0 , we have

Proof. Without loss of generality, assume 0 < c ≤ x 1 ≤ x 2 . We then have

Taking the absolute value of both the left-hand side and the right-hand side, we have

Next, we introduce Lemma B.5 that bounds terms ∣ ∣ ∣ log( ˆ G 0 ( λ 0 ; W )) -log ( G 0 ( λ 0 ; W )) ∣ ∣ ∣ and ∣ ∣ ∣ log( ˆ G 1 ( λ 1 ; W )) -log ( G 1 ( λ 1 ; W )) ∣ ∣ ∣ .

Lemma B.5. Let u denote the probability of treat, i.e., u = P ( T = 1) . Assume that λ 0 , λ 1 ∈ Λ := [ ¯ λ, ¯ λ ] and ˆ τ ( X ) Y is bounded within ¯ M and ¯ M . Then for n ≥ max { 2 u 2 log ( 2 δ ) , 2 (1 -u ) 2 log ( 2 δ ) } , with probability 1 -δ , we have

If M ≤ ¯ M ≤ 0 :

Proof. First, we bound the term ∣ ∣ ∣ log( ˆ G 0 ( λ 0 ; W )) -log ( G 0 ( λ 0 ; W )) ∣ ∣ ∣ .

G 0 ( λ 0 ; W ) and ˆ G 0 ( λ 0 ; W ) are greater than 0 and bounded because Z = ˆ τ ( X ) Y is bounded within the range ¯ M and ¯ M . Therefore, applying Lemma B.4, we have

Moreover, for any λ 0 ∈ Λ , we have

Given ˆ u = 1 n ∑ n i =1 T i and u = E [ 1 n ∑ n i =1 T i ] , using Hoeffding's inequality, we have

We can solve n by

This indicates that (1 -ˆ u ) ≥ (1 -u ) / 2 with probability 1 -δ when n ≥ 2 (1 -u ) 2 log ( 2 δ ) . Combining this with equations (21) and (22), with probability 1 -δ , when n ≥ 2 (1 -u ) 2 log ( 2 δ ) , we have

Therefore, with probability 1 -δ , when n ≥ 2 (1 -u ) 2 log ( 2 δ ) , we have

Next, we bound the term ∣ ∣ ∣ log( ˆ G 1 ( λ 1 ; W )) -log ( G 1 ( λ 1 ; W )) ∣ ∣ ∣ . G 1 ( λ 1 ; W ) and ˆ G 1 ( λ 1 ; W ) are greater than 0 and bounded above. Therefore, applying Lemma B.4, we have

Moreover, for any λ 1 ∈ Λ , we have

Given ˆ u = 1 n ∑ n i =1 T i and u = E [ 1 n ∑ n i =1 T i ] , using Hoeffding's inequality, we have

We can solve n by

This indicates that hatu ≥ u/ 2 with probability 1 -δ when n ≥ 2 u 2 log ( 2 δ ) . Combining this with equations (23) and (24), with probability 1 -δ , when n ≥ 2 u 2 log ( 2 δ ) , we have

Therefore, with probability 1 -δ , when n ≥ 2 u 2 log ( 2 δ ) , we have

If M ≤ 0 : ∣ ∣ ∣ log( ˆ G 1 ( λ 1 ; W )) -log ( G 1 ( λ 1 ; W )) ∣ ∣ ∣ ≤ 2 exp( -¯ M/ ¯ λ ) u ∣ ∣ ∣ ˆ G 1 ( λ 1 ; W ) -G 1 ( λ 1 ; W ) ∣ ∣ ∣

This completes the proof of Lemma B.5.

Additionally, the following Lemma B.6 provides the bound of | log(ˆ u ) -log( u ) | .

Lemma B.6. Let ˆ u = 1 n ∑ n i =1 T i and u = E [ 1 n ∑ n i =1 T i ] . For n ≥ 2 u 2 log ( 2 δ ) , with probability 1 -δ , we have

Proof. Using Hoeffding's inequality, we have

Notably, using the results in the previous lemma, we know for n ≥ 2 u 2 log ( 2 δ ) , ˆ u ≥ u/ 2 . Therefore, we have

In the following, we will bound the term | ˆ V (ˆ τ ) -V (ˆ τ ) | using above lemmas. We first define functions F 0 ( λ 0 ) , ˆ F 0 ( λ 0 ) , F 1 ( λ 1 ) , and ˆ F 1 ( λ 1 ) :

The following Lemma B.7 bounds the term | ˆ F ( λ ) -F ( λ ) | .

Lemma B.7. Let u := P ( T = 1) . Assuming that 0 < ¯ λ ≤ λ ≤ ¯ λ and ˆ τ ( X ) Y is bounded within the range of ¯ M to ¯ M . Define C exp = 1 { ¯ M ≤ ¯ M ≤ 0 } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) + 1 { ¯ M ≤ 0 , ¯ M ≥ 0 } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) + 1 { 0 ≤ ¯ M ≤ ¯ M } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) . For n ≥ 2 /u 2 log(2 /δ ) , with probability 1 -δ , we have

Proof.

¯

Now, we can prove the result in Theorem 4.5.

Proof. Let ˆ λ 0 = arg min λ ˆ F 0 ( λ 0 ) , λ ∗ 0 = arg min λ 0 F 0 ( λ 0 ) , ˆ λ 1 = arg min λ ˆ F 1 ( λ 1 ) and λ ∗ 1 = arg min λ 1 F 1 ( λ 1 ) . Then we have

Therefore, we have

If M ≤ ¯ M ≤ 0 :

Finally, we have

. Note that u 1 = P ( T = 1) and u 0 = P ( T = 0) . C exp = 1 { ¯ M ≤ ¯ M ≤ 0 } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) + 1 { ¯ M ≤ 0 , ¯ M ≥ 0 } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) + 1 { 0 ≤ ¯ M ≤ ¯ M } exp ( ¯ M/ ¯ λ -¯ M/ ¯ λ ) .

C Additional Materials

C.1 Additional Explanations

Q1. Why DRM can select CATE estimators that are robust to the uncertainty in PEHE caused by selection bias and unobserved confounders? In Section 3.1 and Section 4.1, we have presented theoretical explanations for the reason why DRM can measure a CATE estimator's robustness against selection bias and unobserved confounding. Below we will explain it more specifically.

In causal inference, all the CATE estimators are constructed on the observational factual data. But how reliable the CATE estimator that learned on factual data is? This question can be never known unless we have the knowledge of the oracle PEHE. As shown in equation (6), we know that the PEHE is equal to two ˆ τ -dependent terms, E [ˆ τ ( X ) Y t | T = t ] and E [ˆ τ ( X ) Y t | T = 1 -t ] .

Unfortunately, E [ˆ τ ( X ) Y t | T = 1 -t ] is uncomputable empirically because we can only observe the factual distribution P F = P ( X,Y t | T = t ) but not the counterfactual distribution P CF = P ( X,Y t | T = 1 -t ) . The unobserved counterfactual distribution can be regarded as an uncertain distribution varying around the observed and certain factual distribution P F . If we could assume a "God's perspective" and observe P CF directly, the counterfactual distribution will be certain - like a quantum world! Such an uncertainty in P CF results in the uncertainty in PEHE. Now we will analyze the source of such uncertainty by analyzing the relationship between the uncertain distribution P CF and the certain distribution P F based on equation (2):

From above, we find the unobservable distribution P ( X,Y t | T = 1 -t ) is equal to the observable distribution P ( X,Y t | T = t ) multiplied with P ( Y t | T =1 -t,X ) P ( Y t | T = t,X ) P ( X | T =1 -t ) P ( X | T = t ) . In other words, p ( y t | T =1 -t,x ) p ( y t | T = t,x ) p ( x | T =1 -t ) p ( x | T = t ) controls the discrepancy between P F and P CF . Note that if there is no unmeasured confounders, then we have p ( y t | T = 1 -t, x ) = p ( y t | T = t, x ) ; and if there is no selection bias (covariate shift), then we have p ( x | T = 1 -t ) = p ( x | T = t ) . Now we understand the root cause of the discrepancy between P F and P CF (or between E [ τ ( X ) Y t | T = 1 -t ] and E [ τ ( X ) Y t | T = t ] ) lies at the unobserved confounders and selection bias (covariate shift). In the DRM method, the uncertainty caused by potential unobserved confounders and selection bias in PEHE can be further measured as the distributionally robust values ˆ V 1 and ˆ V 0 . Then the PEHE w.r.t. the CATE estimator ˆ τ will be at most R DRM (ˆ τ ) , as shown in equation (15). An estimator ˆ τ that attains smallest R DRM (ˆ τ ) by definition reflects the distributional robustness against potential unobserved confounders and selection bias.

Q2. How to set ϵ ∗ when there are unobserved confounders? When unobserved confounders are present, Proposition 3.6 can also provide guidance for setting ϵ ∗ . Taking ϵ ∗ 1 = D KL ( P C || P T ) as an example, we have

Therefore, when unobserved confounders present, we can set ϵ ∗ to a larger value than the one guided by Proposition 4.6. Simultaneously, as the empirical approximation of ϵ ∗ 1 = D KL ( P C X || P T X ) and ϵ ∗ 0 = D KL ( P T X || P C X ) can be biased, we also suggest set ϵ ∗ to a large value than the empiricallycomputed ones to ensure the ambiguity set is large enough to contain the target distribution. Therefore, we generally set ϵ ∗ 1 = D KL ( P C X || P T X ) + 5 . 2 and ϵ ∗ 0 = D KL ( P T X || P C X ) + 5 . 2 for all settings in our experiment. Theoretically, a larger ϵ ∗ should guarantee the DRM-selected estimator to be more robust, as it allows for a broader range of possible counterfactual distributions in the ambiguity set. However, setting ϵ ∗ too large can result in overly conservative estimator selection (similar to the well-known accuracy-robustness tradeoff). Therefore, how to determine a proper ambiguity radius still remains an open challenge in both our work and distributionally robust optimization literature.

C.2 Hyperparameters

For linear model, we use LogisticRegressionCV and RidgeCV (both are with 3-fold crossvalidation) from sklearn package to tune hyperparameters: Logistic regression: Cs ∈ {0.01, 0.1, 1, 10}; Ridge Regression: α ∈ {0.01, 0.1, 1, 10, 100}.

For Neural Net, we set the hidden layers as [200, 200, 200, 100, 100], each with the ReLU activation function. The model is trained using the Adam optimizer with a learning rate of 0.001, a batch size of 64, and 300 epochs.

For RF, XGBoost, and SVM model, we use AutoML [71, 53] (with 3-fold cross-validation) from flaml package to tune hyperparameters.

C.3 The Complementary Results

First, we would like to emphasize that the experimental results in this final version of the paper differ from those in the original version. We made several revisions based on the feedback from anonymous reviewers: 1) We add Neural Net model to the base ML models and add the U-learner to the meta-learners, increasing the number of CATE estimators from 24 to 36; 2) We adopted AutoML for hyperparameter tuning when training SVM, RF, and XGBoost. All code is available at https://github.com/yiyhuang3/CATE_estimator_selection .

Below, we prsent the complementary PEHE results for 36 candidate CATE estimators, where the candidate pool contains 4 ML models (LR, SVM, RF, and Neural Net) × 9 learners (U-, S-, T-, PS-, IPW-, X-, DR-, R-, RA-).

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: [Yes]

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a method for selecting a conditional average treatment effect estimator using the distributional robust optimization (DRO) technique. The proposed method does not require specifying models for nuisance functions.\n\nSTRENGTHS:\n1. The result is strong, since this paper provides a method of evaluating multiple CATE estimators without specifying other baseline CATE estimator. As a result, more reliable assessment of CATE estimators is possible.  \n2. The empirical comparison is extensive. This paper evaluates many existing CATE estimator and transparently exhibits the performance of the proposed method.   \n\nGiven that the paper contains several weaknesses discussed below, I think the paper’s result is interesting and powerful in practice.\n\nWEAKNESSES:\n**Misleading statements** \n\nThe following statement is generally not true. \n\n> According to Rubin Causal Model [56], the CATE is determined by comparing potential outcomes under different treatment assignments (i.e., treat and control) for a specific individual.\n> \n\nI think the authors are confusing the “unit-level treatment effect” with the “conditional treatment effect”. The unit-level treatment effect $Y_i(T=1) - Y_i(T=0)$ “*for a specific individual*” $i$ cannot be estimated. The CATE is a group-level treatment effect, where the group is specified by the conditioned covariates. Please explicitly distinguish between the term “CATE” and an individual treatment effect.\n\n**Noninformative Introduction**\n\nThe key research question in line 38, “Given multifarious options for CATE estimators, which should be chosen” are not well-defined in the Introduction section. As a reader, the minimum expectation for Introduction section is to see what were the limitations of previous methods with a simple example illustrating the limitations. However, by the design of the paper, until reaching to Section 3, Equation (2), it’s impossible to know (1) what are the plug-in and pseudo-outcome metrics and (2) what were their problems. The unclarity of the problem setting in Section 1 and 2, makes the first two sections non-informative. \n\n**Redundancy**\n\nI don’t see the point of categorizing previous methods as plug-in and pseudo-outcome. All previous methods rely on specifying nuisances, while the proposed method does not. I think such categorization is redundant and unnecessary for the literature review and explanation of the proposed method.\n\n**Subjective statements**\n\nI think this paper contains subjective evaluations of previous methods, which, as a reader, feels somewhat embarrassing. For example, in line 125, \"previous high-quality paper\" seems awkward and redundant — \"previous paper\" is enough. Similarly, \"standing on the shoulders of giants\" is also embarrassing. These kinds of subjective assessments of previous papers dilute the attractiveness of this paper and distract the authors from focusing on their own work.\n\nQUESTIONS:\n1. What is the relation with the paper [[1911.02029] Selective machine learning of doubly robust functionals (arxiv.org)](https://arxiv.org/abs/1911.02029)? \n2. Is the empirical result matched with previous literature [58, 16, 45]?\n\nLIMITATIONS:\nThe paper assumes the ignorability, which is pretty strong assumption in general.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the challenging problem of Counterfactual Average Treatment Effect (CATE) model selection, where the goal is to develop a model selection metric using only observed data, as counterfactual labels are not available. Previous work has focused on model selection metrics that involve learning nuisance parameters on the validation set, making the task of choosing the appropriate CATE model selection criteria complex. The authors propose a solution with a model selection criterion, termed DRM, that avoids the need for training additional nuisance models and is designed to be robust to distribution shifts between treatment and control groups. They first derive the optimization objective for DRM and then propose a relaxation that is tractable with observed data samples. The paper includes a finite sample analysis of the proposed DRM estimator and benchmarks it against existing model selection criteria using synthetic datasets.\n\nSTRENGTHS:\n* The paper makes a significant contribution to CATE model selection by introducing the DRM criterion, which is notable for not requiring the training of additional nuisance models. Prior work often involves specific inductive biases (e.g., S-Learner, T-Learner, DRLearner) and necessitates training extra models, which can lead to sub-optimal choices and affect model selection performance. The DRM criterion addresses these limitations by not being tied to any particular CATE estimation design. Instead, it derives an upper bound on the ideal PEHE metric that is tractable with observed data alone, avoiding the need for additional nuisance models.\n\n*  The theoretical formulation of DRM is compelling and enhances interpretability by linking the bound to the magnitude of the shift between control and treatment distributions. This directly addresses the challenge of distribution shift in effect estimation. Additionally, the authors derive an equivalent formulation in terms of the KL divergence between the control covariates ($P(X|T=0$) and treatment covariates ($P(X|T=1$), which can be computed using observed data. This approach provides a practical way to assess the impact of distribution shift on CATE model selection.\n\n* The derivation of the DRM criteria from the PEHE metric, along with the proposed estimation strategy, appears to be novel to the best of my knowledge.\n\n* The paper is well-written, with the methodology behind DRM presented in an accessible manner. The background section effectively motivates their approach, and the experimental results are clearly presented, accompanied by a thorough discussion of the performance compared to baselines.\n\nWEAKNESSES:\nWhile the authors have done a good job in the experiments by considering diverse synthetic datasets and interesting ablation studies, I think the empirical analysis can still be improved. I have summarized my concerns with the experiments below, which would help to provide a clearer understanding of the performance differences and strengthen the comparison of DRM with other model selection criteria.\n\n* The authors use only three base ML models to train the CATE estimators, which limits the diversity of models considered for model selection. To make the model selection problem more challenging and provide a more thorough evaluation, it would be beneficial to train a larger number of CATE estimators. This would offer a more comprehensive assessment of the DRM criterion's performance and robustness across a broader range of models.\n\n* Estimating the nuisance parameters for model selection metrics is crucial for their performance, as highlighted by Mahajan et al. [1], who showed that plug-in metrics based on T-Learner/X-Learner are effective for model selection. To ensure a fair analysis, the authors should consider using state-of-the-art techniques, such as AutoML, for estimating the nuisance parameters associated with the model selection metrics. Currently, the performance of many model selection metrics might be suboptimal due to potentially poor choices of nuisance parameters\n\n* The comparison of different model selection criteria in the results raises some concerns, particularly due to the high variance observed in the performance of several baselines. This variance might stem from poor results on a few datasets, while these baselines could be comparable to DRM on a majority of datasets. The authors should clarify this in the paper and offer more detailed insights, such as results broken down by dataset or other methods to highlight trends and performance variations.\n\nReferences:\n\n[1] Mahajan, Divyat, Ioannis Mitliagkas, Brady Neal, and Vasilis Syrgkanis. \"Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation.\" arXiv preprint arXiv:2211.01939 (2022)\n\nQUESTIONS:\n* In Table 1, I don't understand why DRM is more robust to selection bias as stated by the authors. Are authors comparing the performance gains with DRM over the best baselines across scenarios ($\\mathcal{E}= \\{0, 1, 2\\}$)? It doesn't seem like the performance gap becomes better with more selection bias. Also, the performance of Plug-R or other baselines across different scenarios ($\\mathcal{E}= \\{0, 1, 2\\}$) is statistically indistinguishable given the difference in means lies within the confidence interval. So I am not sure whether any of the prior metrics are susceptible to selection bias in this empirical study. \n\n* Why is the Plug-T metric much worse than the Plug-S metric? T-Learner should provide more flexibility in choosing the regression models than S-Learner, so I don't understand why would the estimates of the ground truth be worse with T-Learner (hence worse model selection).\n\nLIMITATIONS:\nYes, the authors have properly addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new metric for CATE model evaluation and selection. Specifically, they derive a distributionally robust metric (DRM) which is nuisance-free and robust against selection bias. They show and explain its robust performance in extensive benchmarking experiments against existing baselines.\n\nSTRENGTHS:\n- Robust and reliable CATE model evaluation is an important and often widely overlooked challenge in causal inference\n- The paper is well-structured and easy to follow\n- The authors provide proper theoretical derivation and extensive experimental validation for their method.\n\nWEAKNESSES:\n- The stated robustness against unobserved confounding sounds like a silver bullet. Here, the authors, provide experiments, but should more clearly discuss in theory why and to which degree of unobserved confounding their metric is more robust compared to others. Especially, since for selecting the optimal ambiguity radius they assume unconfoundedness. Hence, it is not clear if or why the metric is still robust, even though the experiments seem to support that.\n- The experimental results seem to be untransparent. The experimental setup for the baselines should be described in more detail, e.g. in the Appendix. It is not clear, if the models used for the baseline metrics were properly tuned. In particular, there is high standard variation in the Regret in Table 1, and large discrepancy compared to the observed performance in ranking in Table 2. Here, the setup should be described in full detail and possible reasons for this discrepancy should be discussed more clearly.\n\nQUESTIONS:\n- For selecting the ambiguity radius in Proposition 3.6, unconfoundedness is assumed. So why should the metric be still robust against unobserved confounding? Was the ambiguity radius in setting C selected differently?\n- Why is the ranking performance clearly worse than the average regret? The high standard deviation in the regret could imply that single runs with bad performance skew the evaluation which could be due to insufficient tuning of the used baseline models for the metrics. Clarification here would be appreciated. \n- What exactly do the authors mean in the limitations when they state “considering the ambiguity set constructed with other divergence such as Wasserstein may contain more diverse distributions”? What does more diverse mean here and what are possible shortcomings of KL divergence in this setting?\n\nLIMITATIONS:\nLimitations are mentioned sufficiently.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new model selection method for choosing an estimator of the conditional average treatment effect (CATE), namely, a Distributional Robust Metric (DRM). The proposed method, DRM, is nuisance-free: It does not require an additional estimation of the nuisance functions, unlike the majority of existing approaches. The DRM splits a precision in the estimation of the heterogeneous effects (PEHE) into two parts: the first term can be estimated nuisance-free and the second term is upper-bounded using the KL-divergence ambiguity set. Thus, the proposed method can only utilize the observational data.  Furthermore, the authors provided an estimation procedure for the DRM and the finite-sample convergence rates. Finally, multiple synthetic experiments were provided to prove the effectiveness of the method in comparison to the other existing approaches.\n\nSTRENGTHS:\nThe paper studies an important and heavily discussed problem in causal inference, CATE model selection. I find the idea of using the upper bound on the PEHE based on the distributional ambiguity set original and novel. The paper is well-structured and easy to follow.\n\nWEAKNESSES:\nTwo major weaknesses of the paper are, in my opinion, the following:\n1. A lack of understanding of the asymptotic behaviour of the DRM. Specifically, when the data size grows, the upper bound on the term (b) of Eq. (5) will still stay the same, as the epsilon-ball around $P_C$ or $P_T$ stays constant. Hence, there will be always a gap between the ground-truth PEHE and the DRM from Eq. (14). This renders the DRM not consistent. I encourage the authors to provide a discussion about this issue and some experiments with the varying data sizes, where they compare the DRM with other (consistent) nuisance-free selectors, e.g., nearest-neighbours matching.   \n2. It was not directly clear to me, why the optimal $\\epsilon^*$ is a forward KL-divergence but not, e.g., a reverse KL-divergence between treated and untreated covariates distributions. Note, that KL-divergence is not symmetric. Does it make sense to use the symmetric distributional distances? \n \nThere are also several minor suggestions for improvement:\n1. The hidden confounding setting needs to be considered more carefully. This setting requires different approaches, e.g., partial identification or sensitivity models, and, thus, fitting/selecting point CATE models is inadequate in this case.\n2. Some baselines were not mentioned in the related work or the experiments, e.g., U-learner [1]. \n\nI am open towards raising my score if the authors address the above-mentioned issues during the rebuttal.\n\nReferences:\n- [1] Fisher, Aaron. \"The Connection Between R-Learning and Inverse-Variance Weighting for Estimation of Heterogeneous Treatment Effects.\" arXiv preprint arXiv:2307.09700 (2023).\n\nQUESTIONS:\n- What ‘uncertainty in PEHE’ is meant in lines 164-165?\n\nLIMITATIONS:\nThe main limitation of the method, i.e., the lack of consistency, was not discussed by the authors.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-21"}