Estimating Heterogeneous Treatment Effects: Mutual Information Bounds and Learning Algorithms
Xingzhuo Guo, Yuchen Zhang, Jianmin Wang, Mingsheng Long
Abstract
Estimating heterogeneous treatment effects (HTE) from observational studies is rising in importance due to the widespread accumulation of data in many fields. Due to the selection bias behind the inaccessibility of counterfactual data, the problem differs fundamentally from supervised learning in a challenging way. However, existing works on modeling selection bias and corresponding algorithms do not naturally generalize to non-binary treatment spaces. To address this limitation, we propose to use mutual information to describe selection bias in estimating HTE and derive a novel error bound using the mutual information between the covariates and the treatments, which is the first error bound to cover general treatment schemes including multinoulli or continuous spaces. We then bring forth theoretically justified algorithms, the Mutual Information Treatment Network (MitNet), using adversarial optimization to reduce selection bias and obtain more accurate HTE estimations. Our algorithm reaches remarkable performance in both simulation study and empirical evaluation.
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 4d2342e4-3021-473a-bd62-339fdb95efceCited by top-tier papers6
- Bounds on Representation-Induced Confounding Bias for Treatment Effect EstimationValentyn Melnychuk, Dennis Frauen, Stefan FeuerriegelICLR 2024 · 23 citations
- Removing Biases from Molecular Representations via Information MaximizationChenyu Wang, Sharut Gupta, Caroline Uhler, Tommi S. JaakkolaICLR 2024 · 11 citations
- Estimating Distributional Treatment Effects in Randomized Experiments: Machine Learning for Variance ReductionUndral Byambadalai, Tatsushi Oka, Shota YasuiICML 2024 · 7 citations
- On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive RandomizationUndral Byambadalai, Tomu Hirata, Tatsushi Oka, Shota YasuiICML 2025
- Treatment Effect Estimation with Differentiated Networked Effect on Graph DataXiaofeng Lin, Han Bao, Hisashi KashimaKDD 2026
Builds on5
- Learning Counterfactual Representations for Estimating Individual Dose-Response CurvesPatrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M. Buhmann et al.AAAI 2020 · 159 citations
- Identifying Causal-Effect Inference Failure with Uncertainty-Aware ModelsAndrew Jesson, Sören Mindermann, Uri Shalit, Yarin GalNeurIPS 2020 · 85 citations
- VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous TreatmentsLizhen Nie, Mao Ye, Qiang Liu, Dan NicolaeICLR 2021 · 81 citations
- Quantifying Ignorance in Individual-Level Causal-Effect Estimates under Hidden ConfoundingAndrew Jesson, Sören Mindermann, Yarin Gal, Uri ShalitICML 2021 · 66 citations
- Causal Effect Inference for Structured TreatmentsJean Kaddour, Yuchen Zhu, Qi Liu, Matt J. Kusner et al.NeurIPS 2021 · 62 citations
Related papers
- Comparison of meta-learners for estimating multi-valued treatment heterogeneous effectsNaoufal Acharki, Ramiro Lugo, Antoine Bertoncello, Josselin GarnierICML 2023 · 18 citations
- Budgeted Heterogeneous Treatment Effect EstimationTian Qin, Tian-Zuo Wang, Zhi-Hua ZhouICML 2021 · 18 citations
- A Relative Error-Based Evaluation Framework of Heterogeneous Treatment Effect EstimatorsJiayi Guo, Haoxuan Li, Ye Tian, Peng WuICLR 2026 · 2 citations
- Improving Data-driven Heterogeneous Treatment Effect Estimation Under Structure UncertaintyChristopher Tran, Elena ZhelevaKDD 2022 · 4 citations
- Proximity Matters: Local Proximity Enhanced Balancing for Treatment Effect EstimationHao Wang, Zhichao Chen, Zhaoran Liu, Xu Chen et al.KDD 2025 · 4 citations
