Paper: Extracting Post-Treatment Covariates for Heterogeneous Treatment Effect Estimation (ICLR-2024)

Reviewer comment (Novelty):
In cases where the analyst genuinely has no prior knowledge of the causal structure, there are established methods in the causal inference literature for discovering causal relationships from observational data. For example, conditional independence tests can be applied to recover the causal structure, especially since post-treatment variables induce v-structures. This approach may offer a more direct and interpretable solution compared to encouraging conditional independence in an embedding space using mutual information regularizers as in the proposed method. Also, testing conditional independence in the original space seems to be a simpler problem compared to requiring conditional independence in an embedding space.

Author response:
Thank you for your insightful question regarding the use of causal discovery in the context of unknown causal structures in observational data. While it is true that conditional independence tests are a valuable tool in causal discovery, particularly in identifying v-structures induced by post-treatment variables, they also come with certain limitations, especially in complex data environments. For instance, when dealing with high-dimensional covariates, the computational cost of these methods can be quite substantial. Moreover, causal discovery based on conditional independence tests faces the challenge of Markov equivalence classes. This issue means that for some causal structures, it is not possible to determine the direction of causality solely through independence tests. Additionally, many of these methods tend to be linear and may struggle to model causal relationships in high-dimensional and complex environments effectively.

Our proposed method, PoNet, offers an alternative approach by focusing on identifying and distinguishing the distributions of three different types of covariates – confounders, collider post-treatment variables, and mediation post-treatment variables. As we demonstrate in Theorem 1, if we can successfully identify and separate these distributions from the covariates, the causal effect can be identified in the presence of post-treatment variables, without the need to fully understand the specific causal relationships between covariates.

But we deeply appreciate the value of your suggestion regarding the use of causal discovery. Leveraging causal discovery in the estimation of causal effects without predetermined causal structures represents an exciting and promising direction for future exploration in this field.Thank you once again for your insightful suggestion.

Reviewer comment (Theory):
The decoupling approach that combines the reconstruction loss and the mutual information regularizer seems to be a heuristic one since it does not really guarantee that the covariates are properly separated into pre-treatment and post-treatment ones, which could hurt the reliability of the proposed method.

Author response:
Thank you for your critical observation regarding our decoupling approach. We understand your concern about the potential heuristic nature of this approach. However, Our approach in PoNet is grounded in a hypothesized causal graph, from which we establish an inference policy. This policy is then operationalized using a data-driven model to decouple the covariates into three factors: confounders and two types of post-treatment variables. The particular structure of our model is designed as multi-channels based on the treatment assignment, to ensure that the learned representation of confounders is pre-treatment, while the other two factors are post-treatment. Moreover, through the supervised information from the data label and the backpropagation during training, we further ensure that the inferred confounders and mediation post-treatment variables are predictive  to outcome while collider post-treatment variables are irrelevant to the outcome. Then the reconstruction loss ensures that the learned representations retain the essential information from the original covariates, while the MIMR helps in encouraging the independence of the representations of confounders, mediation post-treatment, and collider post-treatment variables. This method aims to balance the need for information preservation and the separation of different types of variables, which is a challenging task in high-dimensional and complex datasets. Moreover, our extensive experimentation, detailed in the paper, demonstrates the effectiveness of our proposed model in evaluation. PoNet consistently outperforms state-of-the-art models across various scenarios, indicating its robustness and reliability in handling post-treatment bias. This empirical evidence supports the efficacy of our method, suggesting that, despite potential theoretical concerns, PoNet is a reliable tool for addressing post-treatment bias in practical applications.

Reviewer comment (Novelty):
The technical contribution is somewhat limited. The proposed algorithm feels like a band-aid solution with a structure very similar to TARNet (in terms of using two branches to handle the post-treatment covariates). Based on the experimental results reported in the manuscript, it seems that a generative approach with variable decomposition (that accommodates post-treatment covariates and confounding) can further improve the performance significantly.

Author response:
Thank you for your comments regarding the technical aspects of our proposed algorithm. We acknowledge that in terms of the model architecture, we have not made extensive alterations from the TARNet framework, because architecture design is not our focus. However, we would like to emphasize that the structure design of the model is not the primary contribution of our work. Our significant contribution lies in unifying two different sources of post-treatment bias within a single framework and distinguishing between confounders and two types of post-treatment variables, which is rarely considered in the field of treatment effect estimation from observational data. In addition, our work importantly advances the theoretical understanding of causal effect identifiability under these complex conditions. We have rigorously proved that, given certain causal assumptions, causal effects are identifiable even in the presence of both confounders and two distinct types of post-treatment variables. Besides, in the current version we also enrich this work's theoretical contribution by proving the inferred confounders and post-treatment factors are minimally sufficient (It refers to a statistic that captures all the information in a sample that is relevant for estimating a parameter of the underlying probability distribution.) for the needed parameter $\theta$ for the true treatment effect, more details can be found in Appendix C.

Reviewer comment (Experiments):
What would be the performance of PoNet on datasets without post-treatment covariates? As in real-world scenarios practitioners may not be able to judge if all included covariates are all pre-treatment or not, it is useful if PoNet can be used in both scenarios and achieve state-of-the-art results.

Author response:
Thank you for raising the question about the applicability and performance of PoNet in scenarios where datasets may not have post-treatment covariates. In response to your suggestion, we have conducted additional experiments using a new semi-synthetic dataset generated from PeerRead, where post-treatment bias has been removed. The performance on error of PEHE, ATE and MSE on the test data is shown in the following table, we also have updated it in the Appendix:

|                     | $\sqrt{\epsilon_{PEHE}}$ | $\epsilon_{ATE}$ | $\epsilon_{MSE}$ |
|---------------------|--------------------------|------------------|------------------|
| LR                  | 2.718 $\pm$ 1.076        | 0.206 $\pm$ 0.089| 6.647 $\pm$ 4.495|
| OLS2                | 2.279 $\pm$ 0.852        | 0.102 $\pm$ 0.113| 6.403 $\pm$ 4.352|
| BART                | 2.736 $\pm$ 1.095        | 0.492 $\pm$ 0.261| 6.534 $\pm$ 5.020|
| Causal Forest       | 2.440 $\pm$ 1.011        | 0.264 $\pm$ 0.142| 7.082 $\pm$ 5.772|
| GANITE              | 2.471 $\pm$ 1.057        | 0.821 $\pm$ 0.294| 8.124 $\pm$ 5.663|
| CEVAE               | 2.677 $\pm$ 0.350        | 0.505 $\pm$ 0.118| 2.344 $\pm$ 0.204|
| TEDVAE              | 1.943 $\pm$ 1.023        | 0.302 $\pm$ 0.370| 5.318 $\pm$ 5.306|
| Tarnet              | 2.307 $\pm$ 0.586        | 0.145 $\pm$ 0.139| 5.503 $\pm$ 2.432|
| CFR                 | 2.117 $\pm$ 0.433        | 0.199 $\pm$ 0.115| 5.269 $\pm$ 1.858|
| PoNet               | **1.516 $\pm$ 0.463**    | **0.114 $\pm$ 0.045**| **2.750 $\pm$ 1.142**|

The results indicate that PoNet is not only adaptable to this situation but also maintains its efficiency, often achieving state-of-the-art results compared to other methods. Essentially, if the dataset only contains pre-treatment covariates, PoNet will focus on adjusting for  confounders to estimate the treatment effect accurately .Furthermore, our model’s  inferred​ $Z_c$,  which is excluded to predict outcome, in the absence of post-treatment bias, essentially acts to remove noise variables from the original covariates. This feature contributes to PoNet's superior performance in this scenario. The results from these experiments on the PeerRead dataset without post-treatment variables demonstrate that PoNet can effectively handle both post-treatment and confounding biases, showcasing its robustness and versatility.

Reviewer comment (Presentation):
The paper spends substantial time and effort on establishing the well-known problem of post-treatment bias (e.g., see Abstract and Section 2.2). While some of this is useful as motivation, the issue is already well-known, and less emphasis here would help focus more squarely on the contribution. Additionally, Theorem 1 could be better connected to the post-treatment bias problem, with the key point being that under a certain causal structure, knowledge of the conditional distribution of pure confounders and post-treatment variables given all variables could be relevant.

Author response:
We appreciate your suggestion to reduce the emphasis on the general problem of post-treatment bias, given its established recognition in the field. We agree that a more focused approach would indeed allow us to highlight the unique contributions of our work more effectively. But detailing the two sources of different post-treatment bias in the paper is crucial for the readability of the article and for the reader to understand the contribution of our work. Based on your suggestions, we have tried to shorten the description of the post-treatment bias in the revised version.

Regarding Theorem 1, we acknowledge the need for a clearer connection to the specific problem of post-treatment bias. Our intention with Theorem 1 is to establish that under certain causal structures, the knowledge of the conditional distribution of pure confounders and post-treatment variables given all variables is critical for identifying causal effect. In Theorem 1, we demonstrate a key insight: under the given conditions of covariates $X$ and treatment assignment $T$, the probability distribution of an outcome under an intervention $T$ is determined by the distribution of confounders and mediation post-treatment variables, rather than by collider post-treatment variables. This aligns directly with our analysis presented in Section 2.2. In that section, we discuss the importance of including mediation post-treatment variables in causal effect estimation to capture the indirect effects of treatment on outcome through the mediation variable $Z_m$. Conversely, we argue against conditioning on collider post-treatment variables $Z_c$, as this could introduce additional bias due to unobserved variables. This aligns perfectly with the conclusion of Theorem 1, which underscores the importance of accurately identifying and including the right variables in causal models to avoid post-treatment bias. And this theorem is pivotal in demonstrating how our proposed method effectively handles the complexities introduced by post-treatment variables, thereby mitigating post-treatment bias. Our work with Theorem 1 thus directly addresses the challenge of post-treatment bias by elucidating the specific types of variables that should be considered in the causal modeling process. This theoretical grounding strengthens the basis for our proposed model, PoNet, which is designed to effectively distinguish and utilize the appropriate post-treatment variables for accurate causal inference. Thank you for highlighting the need to more explicitly connect Theorem 1 to the issue of post-treatment bias. Your feedback has provided us with an opportunity to clarify this crucial aspect of our paper. We have clarified it in the updated version.

Reviewer comment (Writing):
Some aspects of the paper seem to distract from the main point. For example, the discussion of the Mutual Information Minimization Regularizer (MIMR) is framed as a contribution of the paper ("We propose a Mutual Information Minimization Regularizer (MIMR)"), but the problem the regularization is trying to solve could be more clearly emphasized, even though the empirical results on the MIMR are appreciated.

Author response:
We appreciate your appreciation of the empirical results regarding MIMR, and we agree that a more focused discussion on the specific problem it addresses would enhance the paper's clarity.

The primary purpose of the MIMR, as detailed in Section 3.3 of our paper, is to ensure the accurate separation of confounders and post-treatment variables, which is essential for unbiased treatment effect estimation. In situations where representations of confounders include information from mediation post-treatment variables, controlling for confounders might inadvertently introduce mediation post-treatment bias. Similarly, if mediation post-treatment variables contain confounder information, addressing confounding bias might not be fully effective.  if  mediation post-treatment variables $Z_m$ contain collider post-treatment information, conditioning on $Z_m$ to predict the outcome can lead to collider post-treatment bias as illustrated in Section 2.2. Precise differentiation between confounders and post-treatment variables is thus critical for reliable treatment effect estimation. To address this challenge, we utilize the MIMR based on kernel density estimation. This non-parametric method fits the distributions of the representations of these variables and measures their independence. By minimizing the mutual information between these variable representations, we effectively separate the underlying factors from each other, thus preventing the inadvertent introduction of bias into the treatment effect estimation. This separation is critical, as it allows for more precise differentiation between confounders and post-treatment variables, ensuring reliable treatment effect estimation. The MIMR thus plays a pivotal role in enhancing the robustness and accuracy of PoNet, our proposed model, by ensuring the independence of critical variables in the causal inference process. We hope this explanation clarifies the specific problem that MIMR addresses and its importance in the context of our work. We have emphasized this aspect more clearly in the revised version of our paper.

Reviewer comment (Theory):
There are practical applicability concerns about the proposed DAG. Since post-treatment variables are presumably omnipresent and difficult to identify, it is unclear when or whether investigators would be willing to assume the proposed graph in Figure 1c. More discussion of the limitations of the proposed approach is needed, such as in a dedicated limitations section. For example, if only a single covariate is available and it is unknown whether it is pre- or post-treatment, the proposed method would not work (according to the reviewer’s understanding). Clarifying how many covariates are needed and similar practical constraints would help contextualize the strengths and limits of the contribution.

Author response:
Thank you for your insightful query about graph structure as depicted in Figure 1c. Our study acknowledges and addresses the omnipresence of post-treatment variables in real-world scenarios and the complexities they introduce in causal inference. The proposed DAG in our paper is designed to illustrate a **general** framework for handling these complexities, but we understand that its applicability might vary based on the specific context and knowledge of the covariates involved. 

The question of how many covariates are "needed" for the effective application of the proposed model is also an important consideration. Our experiments, which encompassed synthetic, semi-synthetic, real world datasets with dimensions ranging from tens to thousands of covariates, have demonstrated that PoNet is effective even in scenarios where the number of covariates is relatively small. This effectiveness is particularly notable in comparison to methods that only consider confounders. The strength of PoNet lies in its data-driven decomposition to identify and distinguish between pre-treatment and post-treatment covariates when the pre-post nature of covariates is unknown. We acknowledge, however, that in some extreme scenarios, such as when there is only a single covariate, PoNet might not perform optimally. This limitation arises from PoNet's design, which separates out three distinct types of factors from the covariates. In such extreme cases, our model might not be as effective, but it is important to note that this limitation is not unique to PoNet and applies to other popular causal estimation methods as well. We appreciate your insightful perspective on this matter. It highlights an interesting point about the relationship between the choice of estimation method and the contextual background knowledge of the application scenario. Practitioners should carefully consider the nature of their data and the specific requirements of their research context when choosing a causal estimation method.

Another limitation of this study is our assumption that the mediation and collider post-treatment variables are entirely independent. Our theoretical analysis suggests that to mitigate post-treatment bias, we should condition on mediation post-treatment variables while excluding collider post-treatment variables in the inference policy. However, it is possible for certain post-treatment variables to act as both mediators and colliders simultaneously. This duality introduces the risk of collider post-treatment bias when they are included as conditioning factors in inference, and mediation post-treatment bias when they are not. Consequently, further exploration is warranted in this area. It is crucial to investigate whether there are more theoretically sound strategies to effectively eliminate such biases or to determine the prevalence of dual-type post-treatment variables in real-world settings and carefully select the scenarios for application, thereby offering more insightful guidance to practitioners. 

We have updated the above limitation discussion in the updated Appendix. Thank you very much for your valuable comments!

Reviewer comment (Presentation):
The figures and tables legends are not to be read on their own. It is a very nice feature to be able to understand the figures without reading the main text (at least for variables definitions, etc). Additionally, the text on figures is generally too small, in particular for Figure 2 which is barely readable.

Author response:
Thank you for pointing this out. We are sorry that the legends in the figures and tables are not clear for reading. In the revised version, we have revised all figures and tables to include comprehensive legends. These will detail all variables and key aspects, ensuring they can be understood independently of the main text. We appreciate your attention to these details.

Reviewer comment (Reproducibility):
The source code is not provided. The absence of a ready-to-use code for practitioners is a severe limit to the usefulness of this work by other researchers in the community, and mostly practitioners.

Author response:
Thank you for your insightful feedback regarding the availability of the source code for our work. We would like to kindly point out that the source code for our model is indeed available and the link has been provided in the supplementary material. Specifically, the link to the code can be found in Appendix Section E, which discusses additional experiments. We understand how crucial it is for the code to be easily discoverable and accessible, and we apologize if its location in the supplementary material was not immediately apparent. To address this, we have put the code link at the beginning of the experiment section in the main body of the paper.

Reviewer comment (Theory):
The theoretical part is a little limited, it could be enriched for instance with a bound on the reconstruction from the obtained representations.

Author response:
Thank you for your valuable feedback regarding the theoretical contributions of our paper. We appreciate your suggestion to enrich the theoretical framework.

Firstly, we would like to highlight that our paper already proposes a significant theoretical contribution regarding the identification of causal treatment in the presence of both confounders and post-treatment variables, as detailed in Theorem 1. We have provided extensive proof for this in the appendix.

Regarding the bounds of reconstruction loss, we have thoroughly considered this aspect. The reconstruction loss based on mean square error, a common regularization term in deep learning, serves dual purposes in our work. It ensures learning about collider post-treatment variables during backpropagation and retains the information in the original covariates. However, the challenge lies in the fact that the reconstructed covariates are processed through a nonlinear decoder based on a multi-layer perceptron. This creates a non-convex problem for solving the bound of this reconstruction loss, presenting a complex challenge in optimization.

But we still totally understand your concerns about the theoretical depth, thus we have further enriched our paper from another angle in the revised version. Noting that Theorem 1 indicates that the treatment effect can be identified by recovering $C$ and $Z_m$, we further provide a new theorem:

The joint set of inferred factors for $C$, $Z_m$ is minimally sufficient for the optimal parameters $\theta$ which estimation of unbiased treatment effects needs.

 It refers to the inferred confounders and mediation post-treatment underlying factors are minimally sufficient for recovering the true treatment effect's optimal parameters. This theory asserts that the inferred post-treatment variables and confounders are the most parsimonious statistics, yet they retain all necessary information for estimating the optimal parameters required for a sufficient estimation of the true treatment effect. The theory has been updated in Section 2 of the main body and detailed proof of this new theorem is also provided in the appendix C.

Reviewer comment (Experiments):
Can you include the method proposed by Li et al. 2022 in the benchmark? As you mention it partially solves the problem of post-treatment bias you address.

Author response:
Thank you for your suggestion to include the method proposed by Li et al. 2022 in our benchmark comparisons. We acknowledge the relevance of this method, as it addresses one aspect of post-treatment bias, which is a key focus of our work. we indeed intended to include Li et al. 2022's method in our comparative analysis. However,  the source code for their method is not publicly available. The absence of this code posed a barrier to implementing their method accurately for a fair and rigorous comparison in our experiments. we are open to including this comparison in future work, should the source code become available or if we can collaborate with the authors for an accurate implementation. In the meantime, we have made efforts to ensure that our benchmark comparisons are comprehensive and include a range of relevant methods that are available for comparison. We appreciate your understanding in this matter and are grateful for your suggestion.

Reviewer comment (Evaluation):
It would be interesting to consider the error on the ATE, in addition to the CATE, as an evaluation metric.

Author response:
We are immensely grateful for your constructive suggestion to consider the error on ATE, in addition to the CATE, as an evaluation metric in our study. Acting on your valuable feedback, we have re-conducted our experiments to include the error metrics of ATE on the synthetic and semi-synthetic datasets in our evaluation.  For example, the updated results regarding to the ATE metric on semi-synthetic PeerRead  are as follows:

|    PeerRead   |        d=50        |        d=100       |       d=200       |
|:-------------:|:------------------:|:------------------:|:-----------------:|
|       LR      |  0.405 $\pm$ 0.164 | 0.317 $\pm$ 0.343  | 1.140 $\pm$ 0.785 |
|      OLS2     | 0.147 $\pm$ 0.067  | 0.363 $\pm$ 0.074  | 0.723 $\pm$ 0.108 |
|      BART     | 0.209 $\pm$ 0.089  | 0.445 $\pm$ 0.241  | 1.253 $\pm$ 0.613 |
| Causal Forest | 0.172 $\pm$ 0.066  | 0.386 $\pm$ 0.183  | 0.938 $\pm$ 0.390 |
|     CEVAE     | 0.415$ \pm$ 0.079  | 0.544 $\pm$ 0.125  | 0.761$\pm$0.357   |
|     GANITE    | 1.740 $\pm$ 0.390  | 1.748 $\pm$ 0.243  | 0.506 $\pm$ 0.396 |
|     TEDVAE    | 0.339 $\pm$ 0.169  | 0.399 $\pm$ 0.314  | 0.599 $\pm$ 0.639 |
|     Tarnet    | 1.018 $\pm$ 0.210  | 0.498 $\pm$ 0.198  | 0.757 $\pm$ 0.246 |
|      CFR      | 1.045 $\pm$ 0.181  |  0.544 $\pm$ 0.220 | 0.706 $\pm$ 0.233 |
|      PoNet     | 0.125 $\pm$ 0.058  | 0.131 $\pm$ 0.077  | 0.138 $\pm$ 0.137 |

The results of these additional experiments are now incorporated in our paper. Specifically, the results on the synthetic and  semi-synthetic datasets are presented in Table 6 and 7 in Appendix. Encouragingly, these additional experimental results demonstrate that our proposed method continues to outperform other baseline methods in terms of the ATE error metrics. This further validates the effectiveness of our approach and strengthens our contributions to the field.

Reviewer comment (Evaluation):
Regarding the experiment on real data, is it possible to also report the different CATE (and ATE) estimates? To have an idea of the impact of the different methods on the actual estimand of interest, even if we do not know the true value.

Author response:
Thank you for your question about reporting the CATE  and ATE estimates on real-world data. 

Regarding CATE estimates, due to the nature of real-world data, which lacks counterfactual outcomes for each sample, it is unfortunately not feasible to calculate CATE as it is defined by the average root mean square difference between the actual and predicted CATE for each sample: ${\epsilon_{PEHE}}=\sqrt{\frac{1}{N}\sum_{i=1}(\tau_i-\hat{\tau}_i)^2}$,  

where  $\tau_i=y_{i}^{t_{i}=1}-y_{i}^{t_{i}=0}, \hat{\tau_i}= \hat{y}_i^{t_i=1}-\hat{y}_i^{t_i=0}$ are the ground truth CATE and the estimated CATE, respectively.

For the ATE error metric, while the standard definition involves the absolute difference between the true ATE and the predicted ATE: $ |Error_{ate}| = |ATE -\hat{ATE}| = \left| \frac{1}{n} \sum_{i=1}^{n}(y_i(1) - y_i(0)) - \sum_{i=1}^{n}(\hat{y}_i(1) - \hat{y}_i(0)) \right|$, the actual calculation requires knowledge of each sample's true counterfactual outcome, which is not available in real-world data.

 However, understanding the importance of your inquiry, we approximated the ATE calculation as follows:  $ATE = \frac{1}{m}\sum_{i=1}^{m} y_i(1) - \frac{1}{k}\sum_{i=1}^{k} y_i(0)$, where m and k are the number of samples in treated and control group. Then we conducted experiments on real-world data using this approximation. In our study, we focused on our proposed method PoNet, along with CFR and TarNet (excluding others like OLS1, OLS2, TEDVAE, BART, Causal Forest, GANITE, CEVAE due to their significantly larger error values). We recorded the ATE error metrics for these methods every 10 epochs during 100 training epochs:

|        | epoch=1 | epoch=10 | epoch=20 | epoch=30 | epoch=40 | epoch=50 | epoch=60 | epoch=70 | epoch=80 | epoch=100 |
|:------:|:-------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:---------:|
|  PoNet |  0.023  |   0.18   |   0.153  |   0.071  |   0.03   |   0.051  |   0.016  |   0.033  |   0.098  |   0.151   |
|   CFR  |  0.282  |   0.095  |   0.028  |   0.088  |   0.096  |   0.092  |   0.117  |   0.345  |   0.14   |   0.062   |
| TARnet |  0.346  |   0.121  |   0.015  |   0.076  |   0.084  |   0.095  |   0.263  |   0.187  |   0.059  |   0.033   |

The results  indicate that the ATE error metrics are highly unstable, appearing somewhat random. This instability is likely due to the unavailability of true ATE values in real-world data, suggesting that calculating ATE error metrics without counterfactual outcomes may not be a reliable evaluation metric in real-world settings. We hope this explanation clarifies our approach and findings, and we appreciate your understanding of the inherent limitations when working with real-world data.

Reviewer comment (Writing):
A detail, the first sentence of section 3.3 is incomplete.

Author response:
Thank you for pointing this out. The sentence should be: Separating confounders and post-treatment variables is essential for unbiased treatment effect estimation. We have corrected it in the revised version.