Causal Optimal Transport for Treatment Effect Estimation

Treatment effect estimation helps answer questions, such as whether a specific treatment affects the outcome of interest. One fundamental issue in this research is to alleviate the treatment assignment bias among those treated units and controlled units. Classical causal inference methods resort to the propensity score estimation, which unfortunately tends to be misspecified when only limited overlapping exists between the treated and the controlled units. Moreover, existing supervised methods mainly consider the treatment assignment information underlying the factual space, and thus, their performance of counterfactual inference may be degraded due to overfitting of the factual results. To alleviate those issues, we build on the optimal transport theory and propose a novel causal optimal transport (CausalOT) model to estimate an individual treatment effect (ITE). With the proposed propensity measure, CausalOT can infer the counterfactual outcome by solving a novel regularized optimal transport problem, which allows the utilization of global information on observational covariates to alleviate the issue of limited overlapping. In addition, a novel counterfactual loss is designed for CausalOT to align the factual outcome distribution with the counterfactual outcome distribution. Most importantly, we prove the theoretical generalization bound for the counterfactual error of CausalOT. Empirical studies on benchmark datasets confirm that the proposed CausalOT outperforms state-of-the-art causal inference methods.