Reducing Balancing Error for Causal Inference via Optimal Transport
Yuguang Yan, Hao Zhou, Zeqin Yang, Weilin Chen, Ruichu Cai, Zhifeng Hao
Abstract
Most studies on causal inference tackle the issue of confounding bias by reducing the distribution shift between the control and treated groups. However, it remains an open question to adopt an appropriate metric for distribution shift in practice. In this paper, we define a generic balancing error on reweighted samples to characterize the confounding bias, and study the connection between the balancing error and the Wasserstein discrepancy derived from the theory of optimal transport. We not only regard the Wasserstein discrepancy as the metric of distribution shift, but also explore the association between the balancing error and the underlying cost function involved in the Wasserstein discrepancy. Motivated by this, we propose to reduce the balancing error under the framework of optimal transport with learnable marginal distributions and the cost function, which is implemented by jointly learning weights and representations associated with factual outcomes. The experiments on both synthetic and real-world datasets demonstrate the effectiveness of our proposed method.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext e844d38f-3b41-486c-bfc1-f042e5f7e8f8Cited by top-tier papers5
- Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal TransportYuguang Yan, Zongyu Li, Haolin Yang, Zeqin Yang et al.ICML 2025
- Adjusting Prediction Model Through Wasserstein Geodesic for Causal InferenceYuguang Yan, Haolin Yang, Zecong Chen, Weilin Chen et al.ICLR 2026
- PITE: Multi-Prototype Alignment for Individual Treatment Effect EstimationFuyuan Cao, Jiaxuan Zhang, Xiaoli LiAAAI 2026
- Matching without Group Barrier for Heterogeneous Treatment Effect EstimationYuguang Yan, Haolin Yang, Shihao Zhang, Weilin Chen et al.ICLR 2026
- Streaming Covariate Balancing via Discrepancy-Based Feature CoresetsYiXin Ren, Chenghou Jin, Yewei Xia, Zichuan Lin et al.ICML 2026
Builds on5
- Optimal Transport for Treatment Effect EstimationHao Wang, Jiajun Fan, Zhichao Chen, Haoxuan Li et al.NeurIPS 2023 · 71 citations
- In Search of Insights, Not Magic Bullets: Towards Demystification of the Model Selection Dilemma in Heterogeneous Treatment Effect EstimationAlicia Curth, Mihaela van der SchaarICML 2023 · 36 citations
- Empirical Analysis of Model Selection for Heterogeneous Causal Effect EstimationDivyat Mahajan, Ioannis Mitliagkas, Brady Neal, Vasilis SyrgkanisICLR 2024 · 28 citations
- fGOT: Graph Distances Based on Filters and Optimal TransportHermina Petric Maretic, Mireille El Gheche, Giovanni Chierchia, Pascal FrossardAAAI 2022 · 20 citations
- Covariate balancing using the integral probability metric for causal inferenceInsung Kong, Yuha Park, Joonhyuk Jung, Kwonsang Lee et al.ICML 2023 · 9 citations
Related papers
- Exploiting Geometry for Treatment Effect Estimation via Optimal TransportYuguang Yan, Zeqin Yang, Weilin Chen, Ruichu Cai et al.AAAI 2024 · 4 citations
- Learning to Re-weight Examples with Optimal Transport for Imbalanced ClassificationDandan Guo, Zhuo Li, Meixi Zheng, He Zhao et al.NeurIPS 2022 · 46 citations
- A Swiss Army Knife for Minimax Optimal TransportSofien Dhouib, Ievgen Redko, Tanguy Kerdoncuff, Rémi Emonet et al.ICML 2020 · 21 citations
- Generative multitask learning mitigates target-causing confoundingTaro Makino, Krzysztof J. Geras, Kyunghyun ChoNeurIPS 2022 · 9 citations
- DeepMatch: Balancing Deep Covariate Representations for Causal Inference Using Adversarial TrainingNathan KallusICML 2020 · 84 citations
