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

Authors: yiyan  huang
hang  cheuk
  leung
siyi  wang
yijun  li
qi  wu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-11 14:28:57

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.

Main paper: 
 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 Ỹ (τ ):EQUATION)One can establish a plug-in estimator τ or construct a pseudo-outcome estimator Ỹ 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 robust a 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:P F := P (X, Y t |T = t) = P (Y t |X, T = t)P (X|T = t);EQUATIONThe 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 Vt (τ ) 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. 
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 = T Y 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τ true (x) := E Y 1 -Y 0 | X = x = µ 1 (x) -µ 0 (x).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 . Unconfoundedness3 : 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 }:EQUATIONHere, 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 byEQUATIONNotably, 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, Ỹ 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] . 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] .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. 
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 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). 
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).Proposition 4.1. The PEHE w.r.t. the CATE estimator τ can be decomposed as follows:EQUATIONwhereζ = E[(µ 1 (X) -µ 0 (X)) 2 ]. The proof is deferred to Appendix B.1.Proposition 4.1 indicates that the PEHE is equal to four terms, whereE[τ (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:EQUATIONEquation (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 forE[τ (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 asEQUATIONHere, 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 asEQUATIONBy setting an adequately large ambiguity radius in Definition 4.2, the following inequalities hold forEQUATIONTo provide a clearer understanding, let us consider the example ofE 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 ofE 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 quantityE 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 asEQUATION)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 toEQUATION)The proof is deferred to Appendix B.3.In the finite-sample scenario, V 0 (τ ) and V 1 (τ ) can be empirically approximated as follows: Return Vt (τ j ) = arg min k∈{0,...,K-1}EQUATIONVt (τ j )[k] for t ∈ {0, 1}.Use V0 (τ j ) and V1 (τ j ) to compute R DRM (τ j ) by equation (15) . Output: τselect = arg min τ ∈{τ1,...,τ J } 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 T Y 1 = T Y , which aligns with the Consistency assumption in Assumption 2.1. We then provide a finite-sample analysis of the gap between Vt (τ ) 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 . DefineEQUATION)The proof is deferred to Appendix B.4.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 (τ ):EQUATIONEQUATIONHere, Z denotes τ (X)Y for notational simplicity. We then use the Newton-Raphson method to find the empirical solution for Vt (τ ), exploiting the convexity of Ft (λ t , ϵ t ; τ ) w.r.t. λ t . Based on the distributionally robust value of PEHE, i.e., V P EHE (τ ) in equation (10) , we finally obtain the selected estimator τselect = arg min τ ∈{τ1,...,τ J } R DRM (τ ) such thatEQUATION)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 haveEQUATION)The proof is deferred to Appendix B.2.Proposition 4.6 provides an important insight that the uncomputable termD 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. 
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 Nearest-Neighbor 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 withT i |X i ∼ Bern(1/(1 + exp(-ξ(β ′ T X i + 3)))), where Bern indicates the Bernoulli A (ρ = 0) A (ρ = 0.1) A (ρ = 0.3) B (ξ = 0) B (ξ = 2) C (m = 0.1) C (m = 0.5) C (m = 0.Y i = d j β ′ j X i;j + d j=1 d k=j β ′ j,k X i;j X i;k + d j=1 d k=j d l=k β ′ j,k,l X i;j X i;k X i;l + T i d j=1 γ j X i;j + ϵ i .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. 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:Regret = R oracle (τ select ) -R oracle (τ best ).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.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 Figure 1 : The stacked bar chart showing the distribution of the selected estimator's rank for each evaluation metric across rank intervals: [1-3], [4] [5] [6] [7] [8] [9] [10] [11] , [12] [13] [14] [15] [16] [17] [18] [19] , [20] [21] [22] [23] [24] [25] [26] [27] , and [28] [29] [30] [31] [32] [33] [34] [35] [36] . The greener (or redder) color indicates that the selected estimator ranks higher (or lower). For example, the dark red (or green) indicates the percentage of cases (out of 100 experiments) where the selected estimator ranks among the worst 9 estimators, specifically as ranks 28, 29, ..., or 36 (or among the best 3 estimators, specifically as ranks 1, 2, or 3).A (ρ = 0) A (ρ = 0.1) A (ρ = 0.3) B (ξ = 0) B (ξ = 2) C (m = 0.1) C (m = 0.5) C (m = 0.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. 
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 Vt (τ ) 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. 
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. 
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. • 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. 
 
 Appendix: A.1 CATE LearnersWe 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):τS (X) = μ(X, 1) -μ(X, 0).• 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):τT (X) = μ1 (X) -μ0 (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 τP S (X) from the following objective:τP S = arg min τ 1 n n i=1 (τ (X i ) -τS (X i )) 2 .• 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 IP W = Y 1 IP W -Y 0 IP W , where Y 1 IP W = T Y π(X) and Y 0 IP W = (1-T )Y 1-π(X). Train a model τIP W (X) from the following objective:τIP W = arg min τ 1 n n i=1 (τ (X i ) -Y 1,0 i,IP W ) 2 .• 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):τX (X) = (1 -π(X))τ 1 X (X) + π(X)τ 0 X (X), τ 1 X = arg min τ 1 n t nt i=1 (τ (X i ) -(Y i -μ0 (X i ))) 2 , τ 0 X = arg min τ 1 n c nc i=1 (τ (X i ) -(μ 1 (X i ) -Y i )) 2 .• U-learner [25, 55] ( ξ i ν i -τ (X i )) 2 .• 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:τDR = arg min τ 1 n n i=1 (τ (X i ) -Y 1,0 i,DR ) 2 .• R-learner [55] : First-step: Let predictors=X, response=Y . Train a model μ(X) to approximate the conditional mean outcomeE[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:τR = arg min τ 1 n n i=1 (ξ i -ν i τ (X i )) 2 .• 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:τRA = arg min τ 1 n n i=1 (τ (X i ) -Y i,RA ) 2 .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 τ (τ ):R plug τ (τ ) = 1 n n i=1 (τ (X i ) -τ (X i )) 2 .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: 1. Pseudo-DR: Utilize validation data to estimate nuisance parameters (μ 1 , μ0 , π), following the procedure described in Section A.1.ỸDR = Ỹ 1 DR -Ỹ 0 DR , where Ỹ 1 DR = μ1 (X) + T π(X) (Y -μ1 (X)) and Ỹ 0 DR = μ0 (X) + 1-T 1-π(X) (Y -μ0 (X)). Then the pseudo-DR metric is R pseudo DR (τ ) = 1 n n i=1 (τ (X i ) -Ỹi,DR ) 2 .For pseudo-DR selector, the selected j * -th CATE estimator is τj * , where j * = arg min j∈{1,...,J} R pseudo DR (τ j ). 2. Pseudo-R: Utilize validation data to estimate nuisance parameters (μ, π), following the procedure described in Section A.1. Then the pseudo-R metric isR pseudo R (τ ) = 1 n n i=1 ((Y i -μ(X i )) -τ (X i )(T i -π(X i ))) 2 .For pseudo-R selector, the selected j * -th CATE estimator is τj * , where j * = arg min j∈{1,...,J} R pseudo R (τ j ). 3. 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 isR pseudo IF (τ ) = 1 n n i=1 ((1 -B i )τ 2 (X i ) + B i Y i (τ (X i ) -τ (X i )) -A i (τ (X i ) -τ (X i )) 2 + τ 2 (X i )),where A i = T i -π(X i ), B i = 2T i (T i -π(X i ))C -1 i , C i = π(X i )(1 -π(X i )).For pseudo-IF selector, the selected j * -th CATE estimator is τj * , where j * = arg min j∈{1,...,J} R pseudo IF (τ j ).4. Other pseudo-outcome selector: By manipulating the formula of Ỹ , 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.1 Proof of Proposition 4.1Proof.(X) -τ true (X)) 2 ] = E[(τ (X) -(µ 1 (X) -µ 0 (X))) 2 ] = E[(τ (X) -µ 1 (X) + µ 0 (X)) 2 ] = E[(τ (X) -µ 1 (X)) 2 ] + E[µ 0 (X) 2 ] + 2E[(τ (X) -µ 1 (X))µ 0 (X)] = E[τ (X) 2 ] + E[µ 1 (X) 2 ] -2E[τ (X)µ 1 (X)] + E[µ 0 (X) 2 ] + 2E[τ (X)µ 0 (X)] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] -2E[τ (X)(µ 1 (X) -Y 1 + Y 1 )] + 2E[τ (X)(µ 0 (X) -Y 0 + Y 0 )] + E[µ 1 (X) 2 ] + E[µ 0 (X) 2 ] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] -2E[τ (X)Y 1 ] -2E[τ (X)(µ 1 (X) -Y 1 )] + 2E[τ (X)Y 0 ] + 2E[τ (X)(µ 0 (X) -Y 0 )] + E[µ 1 (X) 2 ] + E[µ 0 (X) 2 ] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] -2E[τ (X)Y 1 ] -2E[E[τ (X)µ 1 (X) -τ (X)Y 1 |X]] + 2E[τ (X)Y 0 ] + 2E[E[τ (X)µ 0 (X) -τ (X)Y 0 |X]] + E[µ 1 (X) 2 ] + E[µ 0 (X) 2 ] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] -2E[τ (X)Y 1 ] -2E[τ (X)µ 1 (X) -τ (X)E[Y 1 |X]] + 2E[τ (X)Y 0 ] + 2E[τ (X)µ 0 (X) -τ (X)E[Y 0 |X]] + E[µ 1 (X) 2 ] + E[µ 0 (X) 2 ] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] -2E[τ (X)Y 1 ] -2E[τ (X)µ 1 (X) -τ (X)µ 1 (X)] + 2E[τ (X)Y 0 ] + 2E[τ (X)µ 0 (X) -τ (X)µ 0 (X)] + E[µ 1 (X) 2 ] + E[µ 0 (X) 2 ] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] + 2E[τ (X)Y 0 ] -2E[τ (X)Y 1 ] + E[µ 1 (X) 2 ] + E[µ 0 (X) 2 ] -2E[µ 1 (X)µ 0 (X)] = E[τ (X) 2 ] + 2E[τ (X)Y 0 ] -2E[τ (X)Y 1 ] + ζ.The following Proposition B. EQUATIONProof.0 , Y 1 |T = 0) =p(Y 0 , Y 1 |X, T = 0)p(X|T = 0) =p(Y 0 , Y 1 |X)p(X|T = 0). (Unconfoundedness) p(X, Y 0 , Y 1 |T = 1) =p(Y 0 , Y 1 |X, T = 1)p(X|T = 1) =p(Y 0 , Y 1 |X)p(X|T = 1). (Unconfoundedness)Now we can prove Proposition 4.6.Proof. 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 V := supEQUATIONQ∈Bϵ(P ) E Q [f θ (X)] = min λ>0 λϵ + λ log E P [exp(f θ (X)/λ)].Our Theorem 4.4 follows by directly applying the above Lemma B.2.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:G 0 (λ 0 ; W ) = E[g 0 (λ 0 ; W )], Ĝ0 (λ 0 ; W ) = 1 n n i=1 g 0 (λ 0 ; W i ),where g 0 (λ 0 ; W ) = (1 -T ) exp (Z/λ 0 ) ;G 1 (λ 1 ; W ) = E[g 1 (λ 1 ; W )], Ĝ1 (λ 1 ; W ) = 1 n n i=1 g 1 (λ 1 ; W i ),where g 1 (λ 1 ; W ) = T exp (-Z/λ 1 ) .Then we have the following lemma that guarantees the convergence for Ĝ0 (λ 0 ; W ) and Ĝ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 haveEQUATION)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:sup W1,...,Wn,W ′ i ∈W |h 0 (W 1 , . . . , W i , • • • , W n ) -h 0 (W 1 , . . . , W ′ i , • • • , W n )| = sup Wi,W ′ i ∈W |g 0 (λ 0 ; W i ) -g 0 (λ 0 ; W ′ i )| n ≤ 2 sup Wi∈W |g 0 (λ 0 ; W i )| n ≤ 2 exp M/λ 0 n . Note that | Ĝ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 haveP Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) ≥ ϵ = P (|h 0 (W 1 , W 2 , . . . , W n ) -E[h 0 (W 1 , W 2 , . . . , W n )]| ≥ ϵ) ≤ 2 exp   - 2ϵ 2 n( 2 exp( M/λ0) n ) 2   = 2 exp -nϵ 2 2 exp M/λ 0 2 .For some δ > 0, we haveP Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) ≥ ϵ ≤ 2 exp -nϵ 2 2 exp M/λ 0 2 ≤ δ.This solves ϵ such thatϵ ≥ 2 log 2 δ exp M/λ 0 2 n .The above inequality should hold for any λ 0 such that 0 < λ ≤ λ 0 ≤ λ. Therefore, we haveIf M ≥ 0 : ϵ ≥ 2 log 2 δ exp M/ λ 2 n ; If M ≤ 0 : ϵ ≥ 2 log 2 δ exp M/ λ 2 n . 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: sup W1,...,Wn,W ′ i ∈W |h 1 (W 1 , . . . , W i , • • • , W n ) -h 1 (W 1 , . . . , W ′ i , • • • , W n )| = sup Wi,W ′ i ∈W |g 1 (λ 1 ; W i ) -g 1 (λ 1 ; W ′ i )| n ≤ 2 sup Wi∈W |g 1 (λ 1 ; W i )| n ≤ 2 exp (-M/λ 1 ) n .Then using McDiarmid's inequality, for any ϵ > 0, we haveP Ĝ1 (λ 1 ; W ) -G 1 (λ 1 ; W ) ≥ ϵ = P (|h 1 (W 1 , W 2 , . . . , W n ) -E[h 1 (W 1 , W 2 , . . . , W n )]| ≥ ϵ) ≤ 2 exp - 2ϵ 2 n( 2 exp(- M/λ1) n ) 2 = 2 exp -nϵ 2 2 (exp (-M/λ 1 )) 2 .For some δ > 0, we haveP Ĝ1 (λ 1 ; W ) -G 1 (λ 1 ; W ) ≥ ϵ ≤ 2 exp -nϵ 2 2 (exp (-M/λ 1 )) 2 ≤ δ.This solves ϵ such thatϵ ≥ 2 log 2 δ (exp (-M/λ 1 )) 2 n .The above inequality should hold for any λ 1 such that 0 < λ ≤ λ 1 ≤ λ. Therefore, we haveIf M ≥ 0 : ϵ ≥ 2 log 2 δ exp -M/ λ 2 n ; If M ≤ 0 : ϵ ≥ 2 log 2 δ (exp (-M/ λ)) 2 n .In the following content, we will bound terms log( Ĝ0 (λ 0 ; W )) -log (G 0 (λ 0 ; W )) andlog( Ĝ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| log(x 1 ) -log(x 2 )| ≤ 1 c |x 1 -x 2 | (19)Proof. Without loss of generality, assume 0 < c ≤ x 1 ≤ x 2 . We then havelog(x 2 ) -log(x 1 ) = log( x 2 x 1 ) = log(1 + x 2 x 1 -1) ≤ x 2 x 1 -1 = x 2 -x 1 x 1 ≤ x 2 -x 1 c .Taking the absolute value of both the left-hand side and the right-hand side, we have| log(x 1 ) -log(x 2 )| ≤ 1 c |x 1 -x 2 |.Next, we introduce Lemma B.5 that bounds terms log( Ĝ0 (λ 0 ; W )) -log (G 0 (λ 0 ; W )) andlog( Ĝ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 haveEQUATION)Proof. First, we bound the term log( Ĝ0 (λ 0 ; W )) -log (G 0 (λ 0 ; W )) .G 0 (λ 0 ; W ) and Ĝ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 havelog( Ĝ0 (λ 0 ; W )) -log (G 0 (λ 0 ; W )) ≤ 1 c Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) , where c = min inf λ0∈Λ,W ∈W Ĝ0 (λ 0 ; W ), inf λ0∈Λ,W ∈W G 0 (λ 0 ; W ) .Moreover, for any λ 0 ∈ Λ, we haveEQUATIONEQUATION) Given û = 1 n n i=1 T i and u = E[ 1 n n i=1T i ], using Hoeffding's inequality, we haveP 1 n n i=1 (1 -T i ) -E[ 1 n n i=1 (1 -T i )] ≥ E[ 1 n n i=1 (1 -T i )] 2 ≤ 2 exp - 2( 1-u 2 ) 2 n( 1 n ) 2 ≤ δ.We can solve n by2 exp - n(1 -u) 2 2 ≤ δ ⇒ n ≥ 2 (1 -u) 2 log 2 δ .This indicates that (1 -û) ≥ (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 haveIf M ≥ 0 : inf λ0∈Λ,W ∈W G 0 (λ 0 ; W ) ≥ exp( M/ λ)(1 -u); inf λ0∈Λ,W ∈W Ĝ0 (λ 0 ; W ) ≥ exp( M/ λ)(1 -û) ≥ exp( M/ λ)(1 -u)/2. If M ≤ 0 : inf λ0∈Λ,W ∈W G 0 (λ 0 ; W ) ≥ exp( M/ λ)(1 -u); inf λ0∈Λ,W ∈W Ĝ0 (λ 0 ; W ) ≥ exp( M/ λ)(1 -û) ≥ exp( M/ λ)(1 -u)/2.Therefore, with probability 1 -δ, when n ≥ 2 (1-u) 2 log 2 δ , we haveIf M ≥ 0 : log( Ĝ0 (λ 0 ; W )) -log (G 0 (λ 0 ; W )) ≤ 2 exp( M/ λ)(1 -u) Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) ; If M ≤ 0 : log( Ĝ0 (λ 0 ; W )) -log (G 0 (λ 0 ; W )) ≤ 2 exp( M/ λ)(1 -u) Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) .Next, we bound the term log( Ĝ1 (λ 1 ; W )) -log (G 1 (λ 1 ; W )) . G 1 (λ 1 ; W ) and Ĝ1 (λ 1 ; W ) are greater than 0 and bounded above. Therefore, applying Lemma B.4, we havelog( Ĝ1 (λ 1 ; W )) -log (G 1 (λ 1 ; W )) ≤ 1 c Ĝ1 (λ 1 ; W ) -G 1 (λ 1 ; W ) , where c = min inf λ1∈Λ,W ∈W Ĝ1 (λ 1 ; W ), inf λ1∈Λ,W ∈W G 1 (λ 1 ; W ) .Moreover, for any λ 1 ∈ Λ, we haveEQUATIONIfEQUATIONGiven û =1 n n i=1 T i and u = E[ 1 n n i=1T i ], using Hoeffding's inequality, we haveP 1 n n i=1 T i -E[ 1 n n i=1 T i ] ≥ E[ 1 n n i=1 T i ] 2 ≤ 2 exp - 2( u 2 ) 2 n( 1 n ) 2 ≤ δ.We can solve n by2 exp - nu 2 2 ≤ δ ⇒ n ≥ 2 u 2 log 2 δ .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 haveIf M ≥ 0 : inf λ1∈Λ,W ∈W G 1 (λ 1 ; W ) ≥ exp(-M/ λ)u; inf λ1∈Λ,W ∈W Ĝ1 (λ 1 ; W ) ≥ exp(-M/ λ)û ≥ exp(-M/ λ)u/2. If M ≤ 0 : inf λ1∈Λ,W ∈W G 1 (λ 1 ; W ) ≥ exp(-M/ λ)u; inf λ1∈Λ,W ∈W Ĝ1 (λ 1 ; W ) ≥ exp(-M/ λ)û ≥ exp( M/ λ)u/2.Therefore, with probability 1 -δ, when n ≥ 2 u 2 log 2 δ , we haveIf M ≥ 0 : log( Ĝ1 (λ 1 ; W )) -log (G 1 (λ 1 ; W )) ≤ 2 exp(-M/ λ)u Ĝ1 (λ 1 ; W ) -G 1 (λ 1 ; W ) ; If M ≤ 0 : log( Ĝ1 (λ 1 ; W )) -log (G 1 (λ 1 ; W ≤ 2 exp(-M/ λ)u Ĝ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(û) -log(u)|.EQUATIONProof. Using Hoeffding's inequality, we haveP (|û -u| ≥ ϵ) = P 1 n n i=1 T i -E[ 1 n n i=1 T i ] ≥ ϵ ≤ 2 exp -2nϵ 2 , 2 exp -2nϵ 2 ≤ δ solves ϵ ≥ log( 2 δ ) 2n .Notably, using the results in the previous lemma, we know for n ≥ 2 u 2 log 2 δ , û ≥ u/2. Therefore, we have| log(û) -log(u)| ≤ 1 min{û, u} |û -u|. (By Lemma B.4) ≤ 2 u |û -u| ≤ 2 u O   log( 2 δ ) 2n   = O   2 log( 2 δ ) nu 2   .In the following, we will bound the term | V(τ ) -V(τ )| using above lemmas. We first define functions F 0 (λ 0 ), F0 (λ 0 ), F 1 (λ 1 ), and F1 (λ 1 ):F 0 (λ 0 ) = λ 0 ϵ 0 + λ 0 log(E P C [exp(τ (X)Y /λ 0 )]) = λ 0 ϵ 0 + λ 0 log 1 1 -u E[(1 -T ) exp(τ (X)Y /λ 0 )] ; F0 (λ 0 ) = λ 0 ϵ 0 + λ 0 log( 1 n c n i=1 (1 -T i ) exp(τ (X i )Y i /λ 0 )) = λ 0 ϵ 0 + λ 0 log 1 n(1 -û) n i=1 (1 -T i ) exp(τ (X i )Y /λ 0 ) . F 1 (λ 1 ) = λ 1 ϵ 1 + λ 1 log(E P T [exp(-τ (X)Y /λ 1 )]) = λ 1 ϵ 1 + λ 1 log 1 u E[T exp(-τ (X)Y /λ 1 )] ; F1 (λ 1 ) = λ 1 ϵ 1 + λ 1 log( 1 n t n i=1 T i exp(-τ (X i )Y i /λ 1 )) = λ 1 ϵ 1 + λ 1 log 1 nû n i=1 T i exp(-τ (X i )Y i /λ 1 ) .The following Lemma B.7 bounds the term Proof. Let λ0 = arg min λ F0 (λ 0 ), λ * 0 = arg min λ0 F 0 (λ 0 ), λ1 = arg min λ F1 (λ 1 ) and λ * 1 = arg min λ1 F 1 (λ 1 ). Then we have| 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 | F0 (λ 0 ) -F 0 (λ 0 )| ≤ O   8λ 2 0 log 2 δ n(1 -u) 2 C 2 exp   + O   2λ 2 0 log( 2 δ ) n(1 -u) 2   ; | F1 (λ 1 ) -F 1 (λ 1 )| ≤ O   8λ 2 1 log 2 δ nu 2 C 2 exp   + O   2λ 2 1 log( 2 δ ) nu 2   . (26) Proof. | F0 (λ 0 ) -F 0 (λ 0 )| = λ 0 log 1 1 -u E[(1 -T ) exp(τ (X)Y /λ 0 )] -log 1 n(1 -û) n i=1 (1 -T i ) exp(τ (X i )Y i /λ 0 ) = λ 0 log (E[(1 -T ) exp(τ (X)Y /λ 0 )]) -log 1 n n i=1 (1 -T i ) exp(τ (X i )Y i /λ 0 ) + log(1 -û) -log(1 -u) ≤ λ 0 log (E[(1 -T ) exp(τ (X)Y /λ 0 )]) -log 1 n n i=1 (1 -T i ) exp(τ (X i )Y i /λ 0 ) + λ 0 |log(1 -û) -log(1 -u)| . If M ≤ M ≤ 0 : | F0 (λ 0 ) -F 0 (λ 0 )| ≤ 2λ 0 exp( M/ λ)(1 -u) Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) + λ 0 |log(1 -û) -log(1 -u)| (By Lemma B.5) ≤ O   8λ 2 0 log 2 δ n(1 -u) 2 exp M/ λ -M/ λ 2   + O   2λ 2 0 log( 2 δ ) n(1 -u) 2   (By Lemma B.3 and Lemma B.6) If M ≤ 0, M ≥ 0 : | F0 (λ 0 ) -F 0 (λ 0 )| ≤ 2λ 0 exp( M/ λ)(1 -u) Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) + λ 0 |log(1 -û) -log(1 -u)| (By Lemma B.5) ≤ O   8λ 2 0 log 2 δ n(1 -u) 2 exp M/ λ -M/ λ 2   + O   2λ 2 0 log( 2 δ ) n(1 -u) 2   (By Lemma B.3 and Lemma B.6) If 0 ≤ M ≤ M : | F0 (λ 0 ) -F 0 (λ 0 )| ≤ 2λ 0 exp( M/ λ)(1 -u) Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) + λ 0 |log(1 -û) -log(1 -u)| (By Lemma B.5) ≤ O   8λ 0 log 2 δ n(1 -u) 2 exp M/ λ -M/ λ 2   + O   2λ 2 0 log( 2 δ ) n(1 -u) 2   (By Lemma B.3 and Lemma B.6) | F1 (λ 1 ) -F 1 (λ 1 )| = λ 1 log 1 u E[T exp(-τ (X)Y /λ 1 )] -log 1 nû n i=1 T i exp(τ (X i )Y i /λ 0 ) = λ 1 log (E[T exp(τ (X)Y /λ 1 )]) -log 1 n n i=1 T i exp(τ (X i )Y i /λ 1 ) + log(û) -log(u) ≤ λ 1 log (E[T exp(τ (X)Y /λ 1 )]) -log 1 n n i=1 T i exp(τ (X i )Y i /λ 1 ) + λ 1 |log(û) -log(u)| . If M ≤ M ≤ 0 : | F1 (λ 1 ) -F 1 (λ 1 )| ≤ 2λ 1 exp(-M/ λ)u Ĝ0 (λ 0 ; W ) -G 0 (λ 0 ; W ) + λ 1 |log(û) -log(u)| (By Lemma B.5) ≤ O   8λ 2 1 log 2 δ nu 2 exp M/ λ -M/ λ 2   + O   2λ 2 1 log( 2 δ ) nu 2   (By Lemma B.3 and Lemma B.6) If M ≤ 0, M ≥ 0 : | F1 (λ 1 ) -F 1 (λ 1 )| ≤ 2λ 1 exp(-M/ λ)u Ĝ1 (λ 1 ; W ) -G 1 (λ 1 ; W ) + λ 1 |log(û) -log(u)| (By Lemma B.5) ≤ O   8λ 2 1 log 2 δ nu 2 exp M/ λ -M/ λ 2   + O   2λ 2 1 log( 2 δ ) nu 2   (By Lemma B.3 and Lemma B.6) If 0 ≤ M ≤ M : | F1 (λ 1 ) -F 1 (λ 1 )| ≤ 2λ 1 exp(-M/ λ)u Ĝ1 (λ 1 ; W ) -G 1 (λ 1 ; W ) + λ 1 |log(û) -log(u)| (By Lemma B.5) ≤ O   8λ 2 1 log 2 δ nu 2 exp M/ λ -M/ λ 2   + O   2λ 2 1 log( 2 δ ) nu 2   (By Lemma B.V 0 (τ ) -V0 (τ ) = F 0 (λ * 0 ) -F0 ( λ0 ) = F 0 (λ * 0 ) -F0 ( λ0 ) + F 0 ( λ0 ) -F 0 ( λ0 ) = F 0 ( λ0 ) -F0 ( λ0 ) + F 0 (λ * 0 ) -F 0 ( λ0 ) ≤ |F 0 ( λ0 ) -F0 ( λ0 )| + 0 ≤ sup λ0 |F 0 (λ 0 ) -F0 (λ 0 )|. V0 (τ ) -V 0 (τ ) = F0 ( λ0 ) -F 0 (λ * 0 ) = F0 ( λ0 ) -F 0 (λ * 0 ) + F0 (λ * 0 ) -F0 (λ * 0 ) = F0 (λ * 0 ) -F 0 (λ * 0 ) + F0 ( λ0 ) -F0 (λ * 0 ) ≤ | F0 (λ * 0 ) -F 0 (λ * 0 )| + 0 ≤ sup λ0 | F0 (λ 0 ) -F 0 (λ 0 )|. V 1 (τ ) -V1 (τ ) = F 1 (λ * 1 ) -F1 ( λ1 ) = F 1 (λ * 1 ) -F1 ( λ1 ) + F 1 ( λ1 ) -F 1 ( λ1 ) = F 1 ( λ1 ) -F1 ( λ1 ) + F 1 (λ * 1 ) -F 1 ( λ1 ) ≤ |F 1 ( λ1 ) -F1 ( λ1 )| + 0 ≤ sup λ1 |F 1 (λ 1 ) -F1 (λ 1 )|. V1 (τ ) -V 1 (τ ) = F1 ( λ1 ) -F 1 (λ * 1 ) = F1 ( λ1 ) -F 1 (λ * 1 ) + F1 (λ * 1 ) -F1 (λ * 1 ) = F1 (λ * 1 ) -F 1 (λ * 1 ) + F1 ( λ1 ) -F1 (λ * 1 ) ≤ | F1 (λ * 1 ) -F 1 (λ * 1 )| + 0 ≤ sup λ1 | F1 (λ 1 ) -F 1 (λ 1 )|.Therefore, we haveIf M ≤ M ≤ 0 : | V0 (τ ) -V 0 (τ )| ≤ sup λ | F (λ) -F (λ)| ≤ O   8 λ2 log 2 δ n(1 -u) 2 exp M/ λ -M/ λ 2   + O   2 λ2 log( 2 δ ) n(1 -u) 2   ; | V1 (τ ) -V 1 (τ )| ≤ sup λ | F (λ) -F (λ)| ≤ O   8 λ2 log 2 δ nu 2 exp M/ λ -M/ λ 2   + O   2 λ2 log( 2 δ ) nu 2   . If M ≤ 0, M ≥ 0 : | V0 (τ ) -V 0 (τ )| ≤ sup λ | F (λ) -F (λ)| ≤ O   8 λ2 log 2 δ n(1 -u) 2 exp M/ λ -M/ λ 2   + O   2 λ2 log( 2 δ ) n(1 -u) 2   ; | V1 (τ ) -V 1 (τ )| ≤ sup λ | F (λ) -F (λ)| ≤ O   8 λ2 log 2 δ nu 2 exp M/ λ -M/ λ 2   + O   2 λ2 log( 2 δ ) nu 2   . If 0 ≤ M ≤ M : | V0 (τ ) -V 0 (τ )| ≤ sup λ | F (λ) -F (λ)| ≤ O   8 λ2 log 2 δ n(1 -u) 2 exp M/ λ -M/ λ 2   + O   2 λ2 log( 2 δ ) n(1 -u) 2   ; | V1 (τ ) -V 1 (τ )| ≤ sup λ | F (λ) -F (λ)| ≤ O   8 λ2 log 2 δ nu 2 exp M/ λ -M/ λ 2   + O   2 λ2 log( 2 δ ) nu 2   .Finally, we have| Vt (τ ) -V t (τ )| ≤ O   8 λ2 log 2 δ nu 2 t C 2 exp   + O   2 λ2 log( 2 δ ) nu 2 t   .Note that u 1 = P (T = 1) and u 0 = P (T = 0). • 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 exp = 1 { M ≤ M ≤0} exp M/ λ -M/ λ + 1 { M ≤0, M ≥0} exp M/ λ -M/ λ + 1 {0≤ M ≤ M } exp M/ λ -M/ λ .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-). 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.Note that in the setting C of our experiments, the unconfoundedness assumption is violated, leading to misspecified nuisance parameters in CATE estimators and baseline selectors. 


Decision: This paper proposes a distributionally-robust method for selecting CATE estimators. Reviewers were generally very positive about the paper's contributions and emphasized many of its strength. Some limitations (e.g., the non-symmetric nature of D_KL) were discussed in length in a way that seems to have both grounded the author's choices and results, as well as improved the paper. Others should be implemented into the next revision (e.g., details and clarity regarding experimental setup and evaluation, inaccuracies and redundancies of some claims). One point that was brought up multiple times is the extended use of subjective statements - these should be reduced considerably (or removed altogether) in the final version. Nonetheless, overall this is a good paper that is likely to advance the current state of research on CATE model selection and promote the importance of robust models in this literature.

Reviews: review: paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: Two major weaknesses of the paper are, in my opinion, the following:
1. 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.   
2. 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? 
 
There are also several minor suggestions for improvement:
1. 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.
2. Some baselines were not mentioned in the related work or the experiments, e.g., U-learner [1]. 

I am open towards raising my score if the authors address the above-mentioned issues during the rebuttal.

References:
- [1] Fisher, Aaron. "The Connection Between R-Learning and Inverse-Variance Weighting for Estimation of Heterogeneous Treatment Effects." arXiv preprint arXiv:2307.09700 (2023).
questions: questions: - What ‘uncertainty in PEHE’ is meant in lines 164-165?
limitations: limitations: The main limitation of the method, i.e., the lack of consistency, was not discussed by the authors.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: - 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.
- 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.
questions: questions: - 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?
- 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. 
- 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?
limitations: limitations: Limitations are mentioned sufficiently.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: While 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.

* 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.

* 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

* 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.

References:

[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)
questions: questions: * 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. 

* 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).
limitations: limitations: Yes, the authors have properly addressed the limitations of their work.
paper_summary: paper_summary: This 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.
strength_weakness: strength_weakness: **Misleading statements** 

The following statement is generally not true. 

> 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.
> 

I 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.

**Noninformative Introduction**

The 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. 

**Redundancy**

I 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.

**Subjective statements**

I 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.
questions: questions: 1. What is the relation with the paper [[1911.02029] Selective machine learning of doubly robust functionals (arxiv.org)](https://arxiv.org/abs/1911.02029)? 
2. Is the empirical result matched with previous literature [58, 16, 45]?
limitations: limitations: The paper assumes the ignorability, which is pretty strong assumption in general.

Comments: title: Response to Reviewer T6Zv
Response to Reviewer e7g2
Response to Reviewer T6Zv
Response to Reviewer T6Zv
Thank you for your reply
Thank you for your reply
Response
comment: Dear Reviewer T6Zv,

Thank you for your thorough and insightful suggestions! We find that the paper [1] you mentioned is very useful for both CATE estimation and model selection. Their findings, which reveal that the minimax rate of PEHE depends on smoothness and sparsity rather than selection bias, align well with the key investigation in [16] regarding how CATE complexity affects model selection metrics. This has provided us with new insights: How can we select estimators that achieve the optimal PEHE minimax rates? We believe this is a very interesting and important avenue, and it may be closely related to the smoothness and sparsity of CATE and response surfaces.

We are also grateful for your second suggestion. The Integral Probability Metric (IPM), e.g., the Wasserstein distance used in [3], was briefly discussed in line 343 of our original paper. Together with your original question about "symmetric distribution discrepancy" and our rebuttal R2, we will incorporate more detailed discussions on the possible usage of Wasserstein distance in the proposed DRM method. Regarding the paper [2] you mentioned, we find their proposed upper bound for potential outcome error very interesting. They claim that the bound will be tight given a proper selection of the hyperparameter λ. Such a tightness property suggests a potential for f-divergence in the DRM method. In light of this, we intend to conduct a comparative analysis of the advantages and disadvantages of using IPM (such as Wasserstein or MMD) and f-divergence versus KL-divergence.

Finally, we want to express our gratitude again for your quick replies and insightful discussions. The suggested references are valuable and insightful for our future research. Thank you for your time, expertise, and comments on helping us improve the study!

Best regards,

Authors of 4584
Thank you for the quick response!

Yet, I still think that the k-NN is a consistent estimator (an instantiation of the T-learner). The reason is that the distribution shift (or selection bias) only matters for CATE estimation in the low-sample regime [1]. In large-sample regimes, any possible universal function approximator can serve as a consistent estimator of $\mu_0(x)$ and $\mu_1(x)$. Thus, the distribution shift does not hinder the consistency of estimation.

Nevertheless, I consider this paper very interesting and thought-provoking. Therefore, I will raise my score. 

P.S. Here is something that I realized after submitting the review for this paper (therefore the following does not influence my valuation of the current work). There was a recent concurrent paper from ICML 2024, which tackles the same problem of CATE model selection/validation [2]. There, the authors used a _total variation distance_ to bound the counterfactual term of the PEHE risk. Also, they compared their work with the _integral probability metrics_ bounds from [3]. In the final version of the manuscript, I encourage the authors of this work to incorporate the discussion on the above-mentioned works [2, 3] (with alternatives to KL-divergence).  

References:
- [1]  Alaa, Ahmed, and Mihaela Schaar. "Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design." International Conference on Machine Learning. PMLR, 2018.
- [2] Csillag, Daniel, Claudio José Struchiner, and Guilherme Tegoni Goedert. "Generalization Bounds for Causal Regression: Insights, Guarantees and Sensitivity Analysis." arXiv preprint arXiv:2405.09516 (2024).
- [3] 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.514
Dear Reviewer e7g2,

Thank you for your thoughtful reply. We are delighted our rebuttal has addressed most of your concerns, and we greatly appreciate your suggestions and will implement them as follows:

Regarding the experimental aspects, we will include more details, such as settings, model training procedures, hyperparameters, and the choice of ambiguous radius under the hidden confounder setting. We believe these additional materials will enhance the clarity and transparency of our work.

As for your suggestion concerning the ambiguity radius, we will incorporate a sensitivity analysis for the proposed DRM under the hidden confounders setting. We believe this analysis will provide valuable guidance for practitioners and researchers in selecting appropriate ambiguity radius when unobserved confounders exist.

We will carefully revise our paper as we respond in rebuttal. Once again, thank you for taking the time to review our paper, and we appreciate your recognition of our work!

Best regards,

Authors 4584
Thank you for your rebuttal, it addressed most of my concerns.

In general, I think that the writing of the paper could be improved at some points, e.g., less subjective statements and more clarity and transparency regarding the implementation.
Therefore, please include the additional explanations regarding the empirical evaluation, especially the exact tuning settings and the considered ambiguity radius, in the paper or appendix.

Further, I would strongly encourage the authors to include an additional sensitivity study over varying values of the ambiguity radius under unobserved confounding and recommendations how to select the ambiguity radius in practice.

However, overall I think this is still an interesting and valuable paper. Hence, I will increase my score by 1.
Dear Reviewer T6Zv,

Thank you for sharing this reference and giving us a chance to discuss it with you. We strongly agree with you that the kNN regressor can be a consistent estimator of target functions in some machine learning contexts. **However, such consistency property may not necessarily hold when fitting potential outcome surfaces due to the distribution shift between factual and counterfactual distributions.**

The paper you mentioned and Chapters 6 and 11 of [1] prove the consistency of the kNN regression estimator under certain conditions, such as $k/n \rightarrow 0$. However, note that **the consistency holds with the assumption that data are independent and identically distributed (i.i.d.)**, as discussed in the Introduction of [1]. Though the i.i.d. assumption is common in most machine learning scenarios, it doesn't necessarily hold when estimating potential outcome surfaces $\mu_t$ for $t \in $ {0, 1}. Please allow us to explain in detail below.

In the Rubin Causal Model framework, observational (factual) data are i.i.d. tuples $(x_i,t_i,y_i)_{i=1}^{n}$, following the factual distribution $P^{F}:=P(X,T,Y)$. For each pair $(x_i,t_i)$, there also exists an unobserved counterfactual outcome $y_i^{CF}$. The unobserved (counterfactual) data are i.i.d. tuples $(x_i,t_i,y_i^{CF}) _{i=1}^{n}$, following the counterfactual distribution $P^{CF}:=P(X,T,Y^{CF})$. As explained in Section 3.1 and GR1 of our top-page rebuttal, **$P^{F}$ is not identical to $P^{CF}$ in general**. Consequently, a model trained on $P^{F}$ may not predict well on $P^{CF}$.

Let's consider estimating $\mu_0(x)$ for further explanations. To infer the potential outcome $Y^0$ for treated ($T=1$) samples, the kNN regressor first approximates $\hat{\mu}_ {0}(X)$ using controlled samples $(X_i,Y^0_i) _ {i=1}^{n_ {control} }$, then uses $\hat{\mu} _ {0}(X)$ to predict $Y^0$ for treated samples, with the predicted samples being $(X_i, \hat{Y}^0_i) _ {i=1}^{n_ {treat} }$. However, $\hat{\mu} _ {0}(X)$ does not necessarily generalize well on $Y^0$ predictions for treated samples, because the source training data (controlled samples) do not have the same distribution as the target prediction data (treated samples), meaning that there exists a distribution shift problem $P(X,Y^0|T=0) \neq P(X,Y^0|T=1)$. Such a distribution shift (discussed in Section 3.1 and GR1 of our top-page rebuttal) renders kNN-based learner $\hat{\mu}_{0}(X)$ to be a consistent estimator of $\mu_0(x)$. 

Therefore, the above analysis demonstrates that, if kNN-based learner $\hat{\mu}_ {0}(X)$ would be a consistent estimator of $\mu_ {0}(X)$, one of the following two conditions must be met:

* The factual distribution for controlled samples, $P(X_i,Y^0_i|T_i=0)$, should be identical to the counterfactual distribution for treated samples, $P(X_i,Y^0_i|T_i=1)$. This necessitates the absence of distribution shift between factual and counterfactual distributions, which in turn requires two key conditions: the unconfoundedness and access to infinite data from Randomized Controlled Trials (RCTs).

* The training samples should be $(X_i,Y^0_i)$ (including the whole samples) instead of $(X_i,Y^0_i)|T_i=0$ (only including controlled samples). This requires counterfactual knowledge, because the potential outcome $Y^0_i$ remains unknown for treated ($T_i=1$) data.

We hope this clarifies our perspective on the consistency of kNN in causal inference contexts. Note that the above analysis and conditions for consistency property are not exclusive to kNN but also extend to other machine learning methods employed in estimating potential outcomes or CATE.  We will further incorporate the above discussions in Appendix, as they further underscore the importance and significance of our method in CATE estimator selection. The DRM method offers distinct advantages compared with previous selection metrics: it does not require counterfactual data and is capable of selecting estimators that are robust to the distribution shift between factual and counterfactual samples.

Thank you again for your active engagement and effort in helping us improve the paper! We are very welcome any further thoughts or questions you may have.

Reference: [1] Devroye, Luc, et al. A Probabilistic Theory of Pattern Recognition. Stochastic Modelling and Applied Probability.

Best regards,

Authors 4584
Thank you for the additional details.

Regarding the consistency of the knn estimator: What about the following work?
-  Devroye, Luc, et al. "On the strong universal consistency of nearest neighbor regression function estimates." The annals of Statistics 22.3 (1994): 1371-1385.
Dear Reviewer T6Zv,

Thank you for your reply. We are pleased to hear that we have successfully addressed most of your concerns!

In both our original experiment (e.g., Table 1 in original paper) and new experiments, we set $k=1$ for the k-nearest neighbours estimator, which directly follows the code and strategy of the k-nearest neighbours selection metric in [16] to maintain consistency with literature.

Aditionally, if possible, we would appreciate any additional information or references about the claim that "the k-nearest neighbours estimator would consistently estimate both of the potential outcomes surfaces if k is chosen in a data-driven way." To the best of our knowledge, we have not found any evidence supporting this conclusion. We sincerely hope to address this and improve the current comparison. Your suggestions would be valuable in improving the kNN baseline, and further strengthen the experiment quality of our paper.

Thank you once again for your time in reviewing our paper and responding to our rebuttal. We really look forward to your further reply!


Best regards,

Authors 4584
Thank you for clarifying most of the concerns.

Still, it seems to me that the $k$-nearest neighbours estimator would consistently estimate both of the potential outcomes surfaces, $\mathbb{E}(Y \mid X =x, T = t)$, if $k$ is chosen in a data-driven way. I wonder how the authors chose the hyperparameter $k$ for the new synthetic experiments.
Dear Reviewer 9qbS,

Thank you so much for your recognition of our paper and rebuttal! We will carefully incorporate your suggestions in finalizing our paper. We would like to express our gratitude again for your time and effort in the review process!

Best regards,

Authors of paper 4584
Thanks for the detailed response to my questions! My concerns have been addressed and I have updated my rating accordingly.
Dear Reviewer 1tzr,

Thank you so much for your recognition and appreciation! We will definitely address all your comments in the final version, as we also believe the suggestions are very important in improving our paper quality. Thank you again for your prompt reply and we truly appreciate your time and help during the whole review process!

Best regards,

Authors of paper 4584
Thank you for addressing my concerns and questions. I sincerely hope that my concerns will be fully addressed as promised, because I believe this paper is strong and deserves acceptance. Thank you for the additional experiments and detailed explanations. I will raise my score.