Counterfactual Contrastive Learning with Normalizing Flows for Robust Treatment Effect Estimation

We greatly appreciate your high recognition of our work, and are eager to share our thoughts with you.

$**Q1:**$ Relation to debiased ML.

$**R1:**$ Thanks for your insightful idea. We would like to clarify that our work primarily focuses on the challenges of deep learning-based causal effect estimation. As debiased machine learning (DML) is a traditional method, we did not explicitly discuss it in the literature review. We will include this discussion in the revision.

Then, we discuss the differences and connections between DML and our FCCL.

First, we explain the different mechanisms that DML and our proposed FCCL use to address confounding bias. DML handles confounding bias through orthogonal residual regression, i.e., $Y_{i}-M_{y}(X_{i}) =\tau(X_{i})(T_{i}-M_{t}(X_{i}))+\epsilon_{i}$, where $M_{y}(X_{i})$ eliminates the influence of the confounder $X$ on $Y$ and $M_{t}(X_{i})$ eliminates the influence of the confounder $X$ on $T$[1]. The way creates statistical independence between treatment and confounder to simulate randomization conditions. Instead, our method captures the characteristics of potential outcomes under different treatments, mitigating distribution shifts through sample-level alignment to emulate RCTs, which better captures individual-level heterogeneity.

Second, the DML way to addressing bias is highly instructive, particularly its unbiasedness. Therefore, we see great potential in incorporating DML into our future work, as R-learners draw inspiration from DML's approach to handling confounding bias, $\tau(\cdot)=argmin_{\tau}\frac{1}{n} \sum_{i=1}^{n} ((Y_{i} -M_{y}(X_{i}))-\tau(X_{i})(T_{i} -M_{t}(X_{i})) )^{2}$. Thanks again.

Besides, we add comparison experiments with DML (Table1) and the results show that FCCL achieves lower estimation errors.

Table1 ITE estimation errors (std) comparison with DML on IHDP.

$**Q2:**$ How to fix presentation shortcomings?

$**R2:**$ Thank you for pointing out the typos, we will correct them in next version.

( 1 ) We will add a formal definition of $k(x)$ in $**Definition 3.6**$.

( 2 ) We thank the reviewer for the suggestion. Indeed, $dis(\cdot ,\cdot )$ denotes the inner product of vectors in the representation space,i.e., $disc(\Phi^{t}(x),\Phi^{t}(\tilde{x}))=u \sum_{i} b_{i} \sqrt{2-2{{\Phi }^{t=1}({x}\_{i})}^{T}{\Phi }^{t=1}({\tilde{x} }\_{i})}+(1-u)\sum_{i} {a}_{i}\sqrt{2-2{{\Phi }^{t=0}({x}\_{i})}^{T}{\Phi }^{t=0}({\tilde{x} }\_{i})}$. Detailed proofs are provided in Appendix A. $**Theorem 4.2**$ shows that the estimation error $\epsilon _{PEHE}(h,\Phi )$ is upper bounded by the distance constraints in the representation space. Therefore, we implement the optimization via contrastive loss, specifically by leveraging the cosine similarity $sim(\cdot ,\cdot )$.

( 3 ) We apologize and will thoroughly go over this for the revision, specifically changing $[\rho \mathbf{1}{p} \mathbf{1}{p}^{\prime}+(1-\rho) \mathbf{I}{p}]$ to $[\rho \mathbf{1}\_{p} \mathbf{1}\_{p}^{\prime}+(1-\rho) \mathbf{I}\_{p}]$, where $\mathbf{1}\_{p}$ denotes the $p$-dimensional all-ones vector and $\mathbf{I}\_{p}$ denotes the identity matrix of size $p$.

$**Q3:**$ Experiments demonstrating the superiority of flow-based counter-factual generation. For example, one could use the MMD to quantify the difference between the distributions of real and synthetic (counter-factual) covariantes.

$**R3:**$ Following your valuable suggestion, we use the MMD to quantify the difference between the distributions of factual and counterfactual covariates (Table 2) and the result shows that the MMD is considerably smaller for our FCCL. Besides, we compare the effect of alternative counterfactual generation strategies on ITE estimation error in Table 3 in the main text. These results show the superiority of flow-based counterfactual generation.

Table2 MMD mean (std) comparison on IHDP.

[1] https://matheusfacure.github.io/python-causality-handbook/22-Debiased-Orthogonal-Machine-Learning.html#non-parametric-double-debiased-ml

We thank the reviewer for their time and for recognizing the importance of heterogeneous treatment effect estimation. Please see below answers to the questions.

$**Q1:**$ Provide a proper definition of boundary samples and discuss why handling them is challenging.

$**R1:**$ We thank the reviewer for consideration. We provide a formal definition of boundary samples in Appendix D.1, and will include this definition in the main text in the revision.

Furthermore, we analyze the challenges of handling boundary samples from two perspectives:

(1) Many scenarios demand flexible investigation of effect heterogeneity across individuals. However, boundary samples, due to their distinctive characteristics, reside in low-probability density regions distant from class centroids, which poses challenges for accurate ITE estimation.

(2) Existing methods, such as MMD, minimize the distribution discrepancy between treated and control groups by aligning their mean representations $\frac{1}{m\_1}\textstyle\sum_{i=1}^{m\_1}{\phi}(x\_{i}^{t})$ and $\frac{1}{m\_2}\textstyle\sum_{j=1}^{m\_2}{\phi}(x_{j}^{c})$. However, such methods overlook samples near the boundaries of the treatment and control distributions.

Our method achieves robust performance via sample-level alignment, especially in scenarios where individual differences are significant. Besides, we also provide boundary sample analysis by listing the estimation errors of the first five samples from Table 4 in Appendix D, as shown in Table 1. More detailed results are presented in Appendix D.

Table1 The ITE estimation error $|\epsilon_{ITE}|$ of boundary samples on IHDP.

$**Q2:**$ Is optimization end-to-end? or are diffeomorphic counterfactuals found first and then the representation and prediction modules are trained?

$**R2:**$ Our approach falls into the latter: we first generate the diffeomorphic counterfactuals, and then train the representation and prediction modules.

$**Q3:**$ Why the smaller evaluation metric “dis” is desirable?

$**R3:**$ A smaller "dis" value indicates that dispersion of samples in the latent representation space is reduced, which further reflects a smaller distance between boundary samples and their corresponding class centers. This shows that there are fewer boundary samples, enabling better model fitting and exhibiting lower ITE estimation bias, as demonstrated by our results. We will make this point clearer in the next version.

$**Q4 (Weaknesses2):**$ The evaluation could benefit from inclusion of more challenging datasets (e.g., ACIC).

$**R4:**$ Thank you for your valuable suggestions. For validation, we additionally evaluate our method against several representative baselines on ACIC (Table 2). We observe that our method still outperforms the baselines.

Table2 Within-sample and out-of-sample mean (std) for the metrics on ACIC.

Thank you for your insightful suggestions. We have addressed the comments related to the counterfactual generation and model robustness evaluation. Please see our responses below.

$**Q1:**$ How does FCCL compare to diffusion-based counterfactual generators?

$**R1:**$ We thank the reviewer for the question. We add experiments on diffusion-based counterfactual generation (Table1), which shows performance comparable to FCCL, but with slightly inferior result on $\epsilon _{PEHE}^{out-of}$. This may be due to the noise-driven mechanism of the diffusion model [1], which causes counterfactual to deviate from the sample semantic space, especially in the out-sample cases. In contrast, the flow-based model ensures that counterfactuals reside on the same manifold as the original instances, which generates both meaningful and reliable counterfactuals, making FCCL more reliable for individual-level treatment effect predictions. We will include this discussion in the revision.

Table1 ITE estimation errors (std) with different generation methods on IHDP.

$**Q2:**$ How does the contrastive loss temperature $\tau$ affect performance?

$**R2:**$ We perform sensitivity analysis focusing on the temperature coefficient $\tau$ (see Table 2). We observe minimal variation in performance with respect to $\tau$, which demonstrates our model's general robustness.

Table2 ITE estimation errors (std) with different values of $\tau$ on IHDP.

$**Q3:**$ Can FCCL generalize to continuous treatments?

$**R3:**$ This is a valuable point. FCCL currently focuses on binary treatments, we can extend it to continuous treatments in future work. The most direct approach would involve discretizing the continuous treatment variable[2]. Specifically, we can split treatments into $E$ heads, each assigned a dose level from the range$[a, b]$, which is divided into $E$ equal intervals of width $(b-a)/E$. We will explore better approaches to handle continuous treatments directly.

$**Q4 (Claim2):**$ Additional ablation studies isolating the contrastive loss would be beneficial.

$**R4:**$ We discussed the impact of contrastive loss, thank you for going into Figure 4 in the main text.

$**Q5 (Claim3):**$ Empirical validation of whether this ITE bound holds in practice is not explicitly tested.

$**R5:**$ Our empirical analysis demonstrates that the proposed ITE bound effectively guides ITE model training and becomes tighter than the bound proposed in CFR as iterations increase under identical conditions[3] (Table 3).

Table3 Generalization-error bound comparison on IHDP.

$**Q6 (Methods And Evaluation Criteria):**$ Further quantitative measures of alignment would reinforce the claims (e.g., KL divergence).

$**R6:**$ Thank you for your suggestion. We add KL divergence comparisons between treatment and control distributions for two typical methods (see Table 4), which demonstrates that our FCCL can better address distribution shifts. We will include this metric in Figure 3 in the next version.

Table4 KL divergence mean (std) comparison on IHDP.

[1] Kotelnikov Akim, et al. Tabddpm: Modelling tabular data with diffusion models. In International Conference on Machine Learning, pp. 17564–17579. PMLR, 2023.

[2] Schwab et al. Learning counterfactual representations for estimating individual dose-response curves. Proceedings of the AAAI Conference on Artificial Intelligence, pp. 5612-5619, 2020.

[3] Shalit, U., et al. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076–3085. PMLR, 2017.

We appreciate your time and your thoughtful and encouraging comments. We hope our responses can resolve your concern.

$**Q1:**$ I suggest authors give some theoretical or empirical evidence to show that the bound is tight (at least tighter than the bound proposed in CFR).

$**R1:**$ Following your helpful suggestion, we conduct an experimental analysis to compare the generalization error bounds of our proposed FCCL (Equation (14) in Appendix A) with the CFR bound (Equation (8) in Appendix A.2, [1]) under the same conditions.

As shown in Table 1, we track the generalization error bounds at different iterations. The results demonstrate that as the number of iterations increases, FCCL achieves a significantly tighter bound than that of CFR.

Table1 Generalization-error bound comparison on IHDP.

In addition, we would like to clarify that our FCCL provides a different theoretical perspective from CFR. Our generalization error bound links ITE estimation error to the generalization error of factual predictions and representation distances, motivating our focus on minimizing these distances via sample-level alignment. Instead, the bound of CFR shows ITE estimation error can be reduced by the difference between the treated and control distributions.

[1] Shalit, U., et al. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076–3085. PMLR, 2017.