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

Dear Reviewer 1tzr,

We sincerely appreciate your positive recognition of the soundness of our theoretical results and the extensiveness and transparency of our empirical experiments. We are also grateful for your valuable suggestions in improving our presentation and bringing new insights to future work. Below we will respond to your comments one by one.

W1. Distinguishing Conditional Average Treatment Effect (CATE) and Individualized Treatment Effect (ITE)

R1: We agree with your argument about the difference between CATE and ITE. In the revised version, we will clarify this in the Introduction, explicitly defining ITE as $Y^1-Y^0$ (unit-level) and CATE as $\mathbb{E}[Y^1-Y^0|X]$ (subgroup-level).

W2. The presentation of Introduction

R2: We will restructure it as follows:

• Background on causal inference and CATE.
• Motivation for CATE estimator selection and existing methods for CATE estimator selection.
• Challenges faced by current methods (moving content from Section 3.1 here)
• Introducing our DRM method.

This restructuring will allow readers to quickly grasp existing CATE model selection metrics and their limitations before introducing our method, thus providing clearer motivation for our work. Thank you for your suggestion!

W3. Redundancy of categorizing previous methods as plug-in and pseudo-outcome

R3: Thank you for this comment, while we still think it is proper to keep such categorization for the following reasons.

• This maintains consistency with [16, 45].

• The plug-in and pseudo-outcome approaches differ fundamentally in their construction form, as detailed in Section A.2:

1. Plug-in metrics approximate the ground-truth CATE function using $\tilde{\tau}(X)$, which relies on off-the-shelf nuisance functions and only covariates variable $X$.

2. Pseudo-outcome metrics approximate the ground-truth CATE function using $\tilde{Y}(X,T,Y)$, incorporating off-the-shelf nuisance functions and variables $(X,T,Y)$.

• Both DR-based and R-based objectives can be used to construct plug-in and pseudo-outcome metrics (see Section A.2). This categorization helps prevent potential confusion between Plug-R/Plug-DR and Pseudo-R/Pseudo-DR.

We will emphasize their distinction in lines 115-123 of the revised paper.

W4. Subjective statements

R4: We appreciate this helpful suggestion for our presentation. We have removed the subjective statements and will ensure our tone remains objective and rigorous in future paper writing.

Q1. What is the relation with the paper [arxiv 1911.02029] Selective machine learning of doubly robust functionals?

R1:

(a) Difference: While both papers address model selection in causal inference, they differ significantly in the interested quantity and selection objects:

Interested quantity:

• Our paper: Conditional Average Treatment Effect (CATE), i.e., $E[Y^1 - Y^0|X]$.

• Their paper: Average Treatment Effect (ATE), i.e., $E[Y^1 - Y^0]$.

Selection objects:

• Our paper: Selects candidate CATE functions. The specific details are stated in Section 3, Background of CATE Estimator Selection.

• Their paper: Selects candidate nuisance functions in the doubly robust (DR) estimator of ATE. The specific details are as follows:

The DR estimator of ATE is $\theta_{DR}=\frac{1}{n}\sum_{i=1}^{n} [\mu(X_i,1)-\mu(X_i,0)+\frac{T_i}{\pi(X_i)}(Y-\mu(X_i,1))-\frac{1-T_i}{1-\pi(X_i)}(Y-\mu(X_i,0))]$. In $\theta_{DR}$, the nuisance function $\mu(X,T)$ can be fitted with different machine learners, e.g., Ridge regression, Support Vector Machine, Random Forest, etc, and the nuisance function $\pi(X)$ can be fitted with different machine learners, e.g., Logistic Regression, Support Vector Machine, Random Forest, etc. This learner-fitting necessity is the same as plug-in and pseudo-outcome methods in CATE estimator selection. But differently, the main goal of their paper is to select nuisance parameters $\mu$ and $\pi$ for the DR estimator of ATE, from $J_{\mu}$ candidate outcome nuisance functions $\mu_1(X,T), ..., \mu_{J_{\mu}}(X,T)$ and $J_{\pi}$ candidate propensity score nuisance functions $\pi_1(X), ..., \pi_{J_{\pi}}(X)$.

(b) Connection: Interestingly, the nuisance function selection method in their paper offers insights for improving the baseline CATE selection methods Plug-DR and Pseudo-DR. Both Plug-DR and Pseudo-DR metrics involve the process of constructing $\tilde{Y}_i=\mu(X_i,1)-\mu(X_i,0)+\frac{T_i}{\pi(X_i)}(Y-\mu(X_i,1))-\frac{1-T_i}{1-\pi(X_i)}(Y-\mu(X_i,0))$ (see Sections A.1 and A.2). Note that similar to the mentioned paper, one also needs to select a machine learner to fit the nuisance functions $\mu(X,T)$ and $\pi(X)$ for constructing $\tilde{Y}$. Their approach could be helpful in choosing appropriate machine learning learners for fitting nuisance functions $\mu(X,T)$ and $\pi(X)$ in constructing $\tilde{Y}$. This opens an interesting avenue for future research, i.e., exploring nuisance function selection methods for Plug-DR and Pseudo-DR metrics in CATE model selection. We will discuss this interesting research direction in Related Work.

Q2. Is the empirical result matched with [58, 16, 45]?

R2: Our results align with some key findings from previous literature. For example, in lines 288-289, we observe the excellence of the R-objective in many scenarios, corroborating the findings in [58]. In lines 290-292, the results confirm the outperformance of R-based selectors as CATE complexity decreases, consistent with the conclusions in [16]. Notably, our analysis extends beyond previous work by examining CATE estimator performance against the level of unobserved confounders and selection bias, providing new insights into CATE estimator selection.

Finally, we would like to thank you again for your time and effort in reviewing our paper. Your feedback is helpful and valuable in improving our work! We remain very open to discussing any additional comments you may have.

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

Dear Reviewer 9qbs,

Thank you for your thorough review of our paper. We are delighted that you recognize the novelty and significance of our DRM method in advancing CATE model selection, particularly its nuisance-free and robustness to distribution shift. Below we will address your comments.

• Suggestions:

S1. Considering a larger set of base models for CATE estimators.

R1. As you may note, prior research has conducted comprehensive empirical investigations conducted a thorough and comprehensive empirical investigation, providing valuable insights for baseline CATE selectors. In our work, our primary focus has been on proposing a new metric for CATE estimator selection, and we believe the 24 CATE estimators with 8 widely-used meta-learners and 3 base ML models were adequate for the current examination. However, we do agree that considering a larger set of base models could offer a more comprehensive analysis. As such, we have expanded our experiments to include an additional Neural Net model, and the results are reported in Table 4 of the rebuttal pdf, and the updated results will be added into the revised paper. Specifically, the Net has 5 hidden layers [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.

S2. Using AutoML for estimating nuisance parameters.

R2. In the original paper, we have tuned hyperparameters for each model in line with [16]. The details are discussed in GR3 of the top-page rebuttal. Thank you for recommending AutoML. A recent ICLR paper [45] has shown the benefits of using AutoML to search a broader grid of hyperparameters, ensuring well-trained nuisance functions. We will use AutoML to update our results and also in future research for hyperparameter tuning.

S3. The high variance phenomenon should be further discussed.

R3. Thank you for this insightful suggestion. The high variance phenomenon is discussed in GR2 of the top-page rebuttal and the discussion will be added into the revised paper.

• Questions

Q1. Not sure whether any of the prior metrics are susceptible to selection bias. Why DRM is robust to selection bias compared to baselines?

R1. It is widely acknowledged in the causal inference literature that increasing selection bias can lead to larger biases in estimators and selectors (e.g., [8, 30, 67, 68]). Accordingly, for most selectors, the regret is expected to increase with higher selection bias - a trend clearly evident in the results presented in Table 1. Therefore, we would not expect a regret-decreasing trend with higher selection bias for most selectors, including our DRM method.

DRM metric is robust to selection bias when comparing its performance to other baseline selectors, such as Plug-R. Specifically, as the level of selection bias increases from $\xi=1$ to $\xi=2$, the average regret of Plug-R increases significantly, from 1.27 to 4.38 - a jump of 3.11. In contrast, the average regret of DRM only increases from 0.11 to 1.21, a much smaller change of 1.10. However, it is important to note that in the case of no selection bias ($\xi=0$), DRM does not show a clear advantage over some baselines. This aligns with our expectations because DRM is designed to select estimators that are robust to distribution shifts.

Q2. Why is Plug-T worse than Plug-S?

R2. The key factor affecting the performance of Plug-T relative to Plug-S appears to be the level of heterogeneity. Heterogeneity refers to the complexity of the CATE function relative to the potential outcome functions. In our paper, the parameter $\rho$ controls the degree of heterogeneity. According to Table 1, the average regret gap between Plug-T and Plug-S decreases as $\rho$ increases from 0 to 0.3. Specifically, the regret gap decreases from 38.12 - 3.40 = 34.72 to 34.38 - 2.34 = 32.04 as heterogeneity increases. This suggests a trend that Plug-T would perform on par with (or better than) Plug-S as the level of heterogeneity increases.

This observation aligns with insights from prior literature. For example, the authors in [16] found that the T-learner outperforms the S-learner in terms of PEHE when the heterogeneity is much larger. Additionally, [40] noted that the S-learner performs better when the CATE complexity is lower than the potential outcome functions, while the T-learner is preferred when the CATE is more complex. Furthermore, the authors in [45] used the original ACIC data but discarded instances with low-variance CATE to ensure sufficient heterogeneity, where they found that T-learner outperforms S-learner.

Therefore, based on our results and the insights from previous studies, we believe the level of heterogeneity is a key factor that affects the relative performance of Plug-T and Plug-S. It would be interesting to further investigate other possible factors that may influence their comparative performance, as they are widely used meta-learners in real-world applications.

Finally, we would like to express our sincere gratitude for your time and effort in reviewing our paper and providing valuable suggestions and feedback! We are very open to addressing and discussing any further comments or questions you may have.

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

Dear Reviewer e7g2,

We greatly appreciate your thoughtful feedback and insightful comments. Thank you for your recognition of our theoretical and experimental analysis! We will address each of your comments below.

Q1. Should more clearly discuss why and to which degree of unobserved confounding their metric is more robust compared to others.

R1. Regarding "why": Please kindly refer to GR1 of the top-page rebuttal.

Regarding "to what degree", the level of robustness depends on the choice of the ambiguity radius $\epsilon$. Theoretically, a larger $\epsilon$ should guarantee the DRM-selected estimator to be more robust to hidden confounding, as it allows for a broader range of possible counterfactual distributions in the ambiguity set. However, setting $\epsilon$ too large can result in overly conservative estimator selection (similar to the well-known accuracy-robustness tradeoff). Therefore, as lines 216-220 explain, determining $\epsilon$ involves a careful balance between ensuring the counterfactual distribution is contained in the ambiguity set (i.e., robustness) and maintaining a tight upper bound (i.e., accuracy).

2. In Proposition 3.6, unconfoundedness is assumed. Why should the metric still be robust against unobserved confounding? Was the ambiguity radius in setting C selected differently?

R2. As discussed above, the DRM-selected estimators should be robust against unobserved confounders provided that $\epsilon$ is set appropriately. However, determining the proper value of $\epsilon$ remains an open challenge in the distributionally robust optimization literature [29, 46, 39, 41, 63]. In our work, Proposition 3.6 can guide us to set $\epsilon$ when unconfounded. When unobserved confounders are present, we can set $\epsilon$ to a larger value than the one guided by Proposition 3.6. The reason is:

$D_{KL}(P_C||P_T)$

$=\int_{\mathcal{X}}\int_{\mathcal{Y}^0}\int_{\mathcal{Y}^1}p(y^0,y^1|x,T=0)p(x|T=0) \log \frac{p(y^0,y^1|x,T=0)p(x|T=0)}{p(y^0,y^1|x,T=1)p(x|T=1)}dy^1dy^0dx$

$=D_{KL}(P^C_X || P^T_X)+\int_{\mathcal{X}}\int_{\mathcal{Y}^0}\int_{\mathcal{Y}^1}p(y^0,y^1|x,T=0)p(x|T=0) \log \frac{p(y^0,y^1|x,T=0)}{p(y^0,y^1|x,T=1)}dy^1dy^0dx$

$>D_{KL}(P^C_X || P^T_X)$

In setting C, yes, the ambiguity radius is set differently. We set $\epsilon_1=D_{KL}(P^C_X || P^T_X)+5$ and $\epsilon_0=D_{KL}(P^T_X || P^C_X)+5$. This ensures that the ambiguity set is sufficiently large to contain the possible uncertain counterfactual distributions. Based on your insightful question, we will further emphasize the ambiguity radius for unmeasured confounding in the revised paper.

Q3. The high variance for baselines in regret and worse ranking performance of DRM.

R3. The high variance phenomenon is discussed in GR2 of the top-page rebuttal and the discussion will be added in the revised paper. Regarding the ranking performance of DRM, we have provided some explanation in lines 310-317 of the original paper. Now we provide some additional discussion.

• First, high rank correlation does not necessarily imply low regret. E.g., in Table 3 of the rebuttal pdf, we have 24 CATE estimators with true performance rank in ascending order: [1, 2, 3, ..., 24]. A plug-X selector may give a surrogate PEHE value, 11.59, then it selects the 6th estimator and gives a rank order [6, 7, 8, 5, 4, 3, 9, 10, 11, 2, 1, 12, 13, 14, ..., 24]. The Spearman correlation between the true rank and Plug-X rank is 0.89. However, the regret of the selected 6th-ranked estimator is 11.63, much larger than 2.00 produced by the top-ranked estimator.

• Second, the DRM approach is designed to select estimators based on their distributionally robust (worst-case) performance. Consider that a coach selects athletes to attend the Olympics based on their scores over 100 games. The first athlete consistently scores between 90-95, with an average of 93, while the second athlete's scores fluctuate between 85-96, with an average of 94. The coach may prefer to select the first athlete, as the worst-case performance is much higher, even though the second athlete has a higher average score (94) and best score (96).

Q4. Should discuss baseline setups: whether models for baseline metrics were properly tuned.

R4. The baseline setups and hyperparameter tuning method in our paper are in line with [16]. Details are provided in GR3 of the top-page rebuttal. We will add these in the revised paper.

Q5. What does “more diverse mean” and what are possible shortcomings of KL divergence?

R5. By "more diverse", we mean the ability to consider a broader range of distribution types and supports within the ambiguity set. The KL-divergence defined in Eqn (6) is a non-symmetric measure and requires P and Q to have the same distribution type and support.

The limitations of the KL divergence:

• Non-symmetry: The lack of symmetry means the KL divergence does not satisfy the triangle inequality, which can limit its applicability in certain areas, such as domain adaptation and transfer learning. Additionally, it makes the distributional distance hard to interpret.
• Requirement for the same distribution type and support: This restriction can reduce the flexibility of the KL divergence-based ambiguity set. E.g., it prevents comparisons between continuous and discrete distributions.

Therefore, We recognize that using a Wasserstein-based ambiguity set would be a promising research direction. It can be used to compare distributions of different types and supports, capturing a wider range of possible counterfactual distributions. However, as mentioned in the paper, the dual problem of the Wasserstein version involves an intractable "sup" problem, which presents additional challenges that require further investigation.

Thank you again for your time and effort in providing these constructive insights that help improve the clarity and rigor of our work! We are very welcome to any additional feedback you may have.

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

Dear Reviewer T6Zv,

We are grateful for your thorough summary and greatly appreciate your recognition of the motivation and novelty underlying our proposed method. Thank you for your time and effort in providing constructive and helpful feedback. We will respond to each of your comments below.

Q1. DRM is not consistent with ground-truth PEHE. Better discuss this issue and do some experiments with the varying data sizes, comparing with other (consistent) nuisance-free selectors, e.g., nearest-neighbours matching.

R1: Thank you for this suggestion. First, we would like to emphasize that we did not claim the proposed DRM metric converges to the ground-truth PEHE. Our theoretical results only state that $\hat{\mathcal{V}}^t$ should converge to the $\mathcal{V}^t$ at a rate of $1/\sqrt{n}$. Consequently, together with the law of large numbers for other terms in $\mathcal{R}^{DRM}(\hat{\tau})$ in Eqn. (14) should converge to $\mathcal{V}_{PEHE}(\hat{\tau})$ in Eqn. (9) at a rate of $1/\sqrt{n}$. Additionally, we believe it is not a weakness that DRM is not consistent with PEHE for the following two reasons:

• By definition, DRM measures the robustness of PEHE w.r.t. $\hat{\tau}$. We can never observe the ground-truth PEHE due to the uncertainty incurred by selection bias and hidden confounders (as discussed in GR_1 of the top-page rebuttal). Therefore, instead of pursuing the unavailable PEHE, we aim to quantify its uncertainty using a distributionally robust value of PEHE: If some $\hat{\tau}$ has a small $\mathcal{R}^{DRM}(\hat{\tau})$, then $\hat{\tau}$ is robust to the PEHE uncertainty caused by the uncertainty in counterfactual distribution.
• To the best of our knowledge, no existing nuisance-free estimator/selector (e.g., the mentioned matching) is assured to be a consistent estimator of the ground-truth CATE/PEHE, but some nuisance-free estimators are consistent estimators of ATE (not CATE) given additional assumptions. E.g., (Abadie & Imbens, 2006) show the matching estimator converges to ATE at a rate slower than $1/\sqrt{n}$, given the assumption that the density of covariates is bounded.

Per your suggestion, we have compared DRM and nearest-neighbor matching with varying sample sizes in Table 2 of the attached rebuttal pdf reports the details and results. The findings show that DRM and matching exhibit a decreasing trend in regret with increasing sample size. However, when the sample size exceeds 10,000, the average regret for matching fluctuates between 6.5 and 7, while the average regret for DRM becomes steady. These results suggest that neither DRM nor matching converges to the ground-truth PEHE because a consistent metric should have a regret that tends to zero. Nevertheless, DRM demonstrates more stable and effective performance than matching as the sample size varies.

Ref: (Abadie & Imbens, 2006) Large sample properties of matching estimators for average treatment effects, Econometrica, 2006

Q2. Why the $\epsilon^*$ is determined as $D_{KL}(P_C||P_T)$ instead of $D_{KL}(P_T||P_C)$? Does it make sense to use the symmetric distributional distances?

R2. Indeed, we set $\epsilon_1^*=D_{KL}(P^C_X||P^T_X)$ and $\epsilon_0^*=D_{KL}(P^T_X||P^C_X)$ in the original paper. In lines 222-237, we intended to explain the choice of $\epsilon^*$ with $\epsilon_1^*=D_{KL}(P^C_X||P^T_X)$ as an example, but forgot to explain $\epsilon_0^*$, which led to a misleading presentation. We are grateful for your pointing this out, and we will revise the incomplete expression to make it clearer.

Regarding your second question, we believe using symmetric distances would be more meaningful. The non-symmetric nature of KL divergence makes it fail to satisfy the triangle inequality and can make the divergence values difficult to interpret and compare between distributions. This restricts its practical use in fields like domain adaptation and transfer learning. Exploring the use of symmetric divergences, such as the Wasserstein distance, would be very useful in DRM (as discussed in line 343). However, solving the distributionally robust optimization problem with a Wasserstein-based ambiguity set can be computationally challenging due to the intractable dual problem involving an infinite "sup" objective. Therefore, considering alternative symmetric divergences that may provide a more tractable optimization framework for DRM would be a promising direction.

Minor suggestions:

S1. More baseline CATE estimators aiming at the hidden confounding setting should be considered.

R1. At the current stage, one key focus is developing a CATE selection method that can robustly handle the distributional shifts. Therefore, the paper is mainly around “selection” rather than “estimation”. However, we strongly agree that considering more CATE estimators aiming at hidden confounder issues would provide more insights for testing different selectors under the hidden confounder setting, which will be further investigated. We will mention this interesting suggestion and leave it for future research.

S2. U-learner can be included as a baseline CATE estimator.

R2. We have incorporated the U-learner [1] into our experimental evaluation, and the updated results are reported in Table 1 of the top-page rebuttal pdf. We will add a discussion of the U-learner and incorporate the updated experimental results in the final version of the paper.

Q. What 'uncertainty in PEHE' is meant in lines 164-165?

R. We hope GR_1 of the top-page rebuttal has addressed this question well, and we will add these details in the Appendix to provide clearer explanations for the proposed method.

Finally, thank you again for providing a thorough review and constructive suggestions! Your comments are valuable in helping us improve our work, and we will carefully incorporate your suggestions when finalizing the paper. We are very welcome any further questions or comments you may have.

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

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