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

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.

Paper

References (79)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer 1tzr8/10 · confidence 4/52024-07-10

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.

Strengths

1. 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. 2. The empirical comparison is extensive. This paper evaluates many existing CATE estimator and transparently exhibits the performance of the proposed method. Given that the paper contains several weaknesses discussed below, I think the paper’s result is interesting and powerful in practice.

Weaknesses

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

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]?

Rating

8

Confidence

4

Soundness

4

Presentation

2

Contribution

4

Limitations

The paper assumes the ignorability, which is pretty strong assumption in general.

Authorsrebuttal2024-08-10

Thank you for your reply

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

Reviewer 9qbS8/10 · confidence 4/52024-07-13

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.

Strengths

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

Weaknesses

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

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

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, the authors have properly addressed the limitations of their work.

Authorsrebuttal2024-08-10

Thank you for your reply

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

Reviewer e7g26/10 · confidence 3/52024-07-15

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.

Strengths

- Robust and reliable CATE model evaluation is an important and often widely overlooked challenge in causal inference - The paper is well-structured and easy to follow - The authors provide proper theoretical derivation and extensive experimental validation for their method.

Weaknesses

- 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

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

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations are mentioned sufficiently.

Authorsrebuttal2024-08-13

Response to Reviewer e7g2

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

Reviewer T6Zv6/10 · confidence 4/52024-07-21

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.

Strengths

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

Weaknesses

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

- What ‘uncertainty in PEHE’ is meant in lines 164-165?

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The main limitation of the method, i.e., the lack of consistency, was not discussed by the authors.

Reviewer T6Zv2024-08-12

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.

Authorsrebuttal2024-08-12

Response to Reviewer T6Zv

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

Reviewer T6Zv2024-08-12

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.

Authorsrebuttal2024-08-13

Response to Reviewer T6Zv

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

Reviewer T6Zv2024-08-13

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

Authorsrebuttal2024-08-13

Response to Reviewer T6Zv

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

Reviewer 1tzr2024-08-10

Response

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.

Reviewer 9qbS2024-08-10

Thanks for the detailed response to my questions! My concerns have been addressed and I have updated my rating accordingly.

Reviewer e7g22024-08-13

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.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC