Optimal Transport for Treatment Effect Estimation

Summary

In this paper, authors have applied optimal transport to address selection-bias issue in binary treatment setting for individualized treatment effect estimation. They have proposed a relaxed mass-preserving regularizer to address the mis-alignment generated from outcome imbalance and outliers in non-ideal mini-batches. They also proposed a proximal factual outcome regularizer to relax the unconfoundedness assumption. They have done extensive evaluation of the proposed approach on IHDP and ACIC, two semi-synthetic datasets.

Strengths

- A novel method is proposed for treatment effect estimation in binary setting based on optimal transport (OT) that uses a relaxed mass-preserving regularizer to address the mis-alignment generated from outcome imbalance and outliers in non-ideal mini-batches, as well as a proximal factual outcome regularizer to relax the unconfoundedness assumption. - Extensive evaluations is performed to study the proposed model that provides necessary analysis and ablation studies. - Theoretical results are provided to bound the error in terms of the proposed components to show the convergence (I have not verified the proof). - Idea was clearly stated with sufficient details.

Weaknesses

- Baselines used in the paper seem old. There are very large number of methods proposed for treatment effect (TE) estimation. It would have been nice to see comparison with recent techniques. Hoever, even if authors don't add new baselines, I think paper has sufficient novelty to be above the acceptance threshold. - The paper solves (equation (1)) Conditional average treatment effect estimation also called as heterogeous treatment effects estimation, as called in machine learning literature, but paper did not even mention these names. Please update the paper so that readers don't have any confusion. There is already a lot of confusion in ML regarding the terminology. - Causal inference methods are typically based on some assumptions -- but I did not see any. Please clearly state the assumptions required to apply these methods.

Questions

- Training all networks for 400 epochs do not look fair to me. Are all the methods getting trained before 400 epochs or not? In my experience this affects the performance in causal inference. All the networks should be trained using early stopping and by setting large number of max-epochs so that every learner gets chance to train properly. - Have you used two layers each in representation learning and outcome mapping or total two layers for both? - First line of Introduction section does not look right to me. Doesn't RCT mostly answer population-level or heterogeneous-level questions but not the individual-level treatment effects? - In figures 1(a) and 2, please clearly state what are the different axis, e.g., for Figure 2, it is not clear what is x-axis or what is y-axis? - Why does title says 'a new perspective' when the idea is very widely used, including in causal inference, as authors have themselves discussed that?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Authors did not discuss the assumptions and limitations of thier work.

Summary

The authors propose to minimize a (relaxed) Sinkhorn distance between representations under the two treatment groups in order to balance the latent space and better predict the treatment effects. They highlight two problems with existing TARNet-based approaches that do not regularize with optimal-transport cost: mini-batch sampling effects (MSE) and unobserved confounder effects (UCE). Theorem 3.1 & Corollary 3.1 formally characterize how the proposed method suffers less from MSE. The authors also introduce additional regularization terms that are meant to deal with UCE. Finally, they show improved performance with two classical semi-synthetic benchmarks.

Strengths

* The problem of estimating heterogeneous treatment effects is important, the paper is written well, and the solution is creative. * The motivation for the relaxed Sinkhorn distance (Sec. 3.2) is great. * The benchmark results are solid. * The additional empirical explorations (Sec. 4.3 - 4.5) are interesting as well.

Weaknesses

My biggest problem is with the UCE part of the paper. * Sec. 3.3: the explanation for the hidden-confounding regularizer makes intuitive sense. However, I do not see a formal connection. Once the authors substitute the actual potential outcomes with the model's predictions (Eq. 11 -> Eq. 12), one begins to wonder if this can actually help with hidden confounding. * A major confounder with the issue above is that the authors do not appear to include hidden confounders in their semi-synthetic benchmarks. Therefore, it is impossible to tell if this regularization approach really helps with UCE. * If the authors still wish to make the case for UCE mitigation, they need to reference the growing literature on partial identification / sensitivity analysis for hidden confounders and how their regularization relates to sensitivity models. Besides that, I think the authors should spend more time addressing what the theoretical results mean for the MSE issue that they highlighted in the introduction. Small nitpicks: * Figure 1a, the orange annotation looks like it is referring to the X_1 part of the figure. * line 107-108: "Optimal transport is a preferred method due to its advantages over competitors." is overly vague. * line 108-109, "It accounts for the distribution’s geometry, making it effective [where the KL-divergence] fails." is a strange statement. Why does the geometry help here? No explanation is provided. * The beginning of the paper is a bit verbose.

Questions

How do you know that you address UCE? As the paper stands, I recommend that you refrain from discussing hidden confounders in your motivation for the method and stick with what really *can* be shown for your proposed regularization. I am willing to raise my score once this is addressed. Is MSE really an issue with previous approaches? Can you show that clearly?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations or potential societal impacts are not discussed.

Summary

The objective of this research paper is to tackle two primary concerns: (1) the effects of mini-batch sampling (MSE), which lead to misalignment in non-ideal mini-batches characterized by outcome imbalance and outliers, and (2) the impacts of unobserved confounders (UCE), resulting in inaccurate discrepancy calculation due to the neglect of these unobserved factors. To address these challenges, this paper presents a novel approach based on optimal transport in the context of causality. More specifically, the proposed approach builds upon the stochastic optimal transport framework. It introduces a relaxed mass preserving regularizer to mitigate the MSE issue and devises a proximal factual outcome regularizer to address the UCE problem. By incorporating these techniques, the proposed method achieves a principled solution for treatment selection bias. The effectiveness of the proposed method is extensively evaluated through a series of experiments. The results demonstrate that the approach successfully overcomes treatment selection bias and outperforms existing state-of-the-art methods by a significant margin.

Strengths

The paper is clearly written and well-presented. It delves into the important question of improving the estimation of individualized treatment effects, a crucial topic within the field of causal inference. Optimal transport, an emerging field in machine learning, is effectively employed to measure treatment selection bias. Furthermore, the two proposed regularizers are technically intriguing and efficient, as supported by both theoretical and empirical evidence.

Weaknesses

I have two main concerns: 1. The potential applicability of the proposed relaxed mass-preserving regularizer extends beyond causal effect estimation and encompasses a wide range of application settings. It appears to be independent of the estimation of causal effects. The authors acknowledge that the OT discrepancy can be easily affected by various sampling cases, highlighting that it is not solely limited to the distribution discrepancy across treatment groups. Consequently, the main contribution of this regularizer may be seen as an incremental addition to the field of causal inference, raising concerns about the originality and novelty of the paper. 2. The existing literature on the subject appears to be considerably understudied. For example, important references like "Optimal Transport for Counterfactual Estimation: A Method for Causal Inference" by Arthur Charpentier, Emmanuel Flachaire, and Ewen Gallic are missing and remain unaddressed in the related work section. It would be beneficial for the authors to conduct a more comprehensive investigation of the related literature and consider incorporating these references into their study.

Questions

See above

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

See above

Summary

The paper introduces a new approach called Entire Space CounterFactual Regression (ESCFR) to estimate individual treatment effects from observational data, addressing two critical problems that existing methods fail to solve. The first problem is mini-batch sampling effects (MSE), which lead to misalignment in non-ideal mini-batches with outcome imbalance and outliers. The second problem is unobserved confounder effects (UCE), which result in inaccurate discrepancy calculation due to the neglect of unobserved confounders. ESCFR combines the principles of optimal transport and causality to overcome these issues.

Strengths

- The theoretical evidence is solid and convincing. - The topic, mentioned in the paper as the mini-batch misalignment and UCE error, are interesting and practical issues of ITE estimation. - Experimental results are solid.

Weaknesses

Is there some theoretical guarantee for the UCE regularizer proposed in this paper? If so, I think the quality can be further improved.

Questions

See in Weaknesses

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

See in Weaknesses

Response to Authors

I appreciate the detailed response to my questions from the authors. Given that both my evaluation and the author's rebuttal are in accord about the significance and potential of the crossover between OT and causal inference as a field, as well as the concerns surrounding the proposed method's contribution, I believe that my rating can adequately convey my standpoint on this paper. As a result, I intend to keep my score.

Raising my score to 6

Thank you for your response. After reading your rebuttals as well as Section D.3 more carefully, I think I understand the corresponding sensitivity assumption for hidden confounders that lies behind your regularization approach. I trust that you will clarify this in the main text. You also addressed some of my concerns surrounding the benchmarks. Still, there are many ways to induce hidden confounding in semi-synthetic benchmark, for instance by hiding some of the covariates. This could have been a valuable addition to the study. Overall, I have updated my score to reflect my new assessment of this work.

Concluding Response

Dear Reviewers, Many thanks for your engagement in the discussion phase. For convenience, we have prepared a concise summary on the final updates according to your suggestions. In the final version, we will - provide a theoretical analysis of PFOR and its relationship with partial identification [VwsR, MG3J]; - provide a more comprehensive description of related works and include all suggested citations [hjjj]; - revise the paper according to the detailed feedbacks and clarify the strengths of PFOR in the main text, with limitation discussion and toy examples; provide more detailed analysis of MSE issue [MG3J]; - revise the typos and details that are kindly pointed out (e.g., CATE v.s. ITE), reconsider the scope of our title [jCP6]. We are thankful for the feedback and your suggestions on the paper improvement.