Potential Outcome Imputation for CATE Estimation

Ahmed Aloui, Juncheng Dong, Cat Phuoc Le, Vahid Tarokh

Keywords: causal inference, treatment effects, data augmentation

TL;DR: In this work, we propose a model-agnostic data augmentation method to improve treatment effect estimation.

Abstract: One of the most significant challenges in Conditional Average Treatment Effect (CATE) estimation is the statistical discrepancy between distinct treatment groups. To address this, we propose a model-agnostic data augmentation method for CATE estimation. We first derive regret bounds for general data augmentation methods, indicating that reduced group discrepancy and low imputation error enhance CATE estimation. Inspired by this, we introduce a contrastive learning approach that reliably imputes missing potential outcomes for a selected subset of individuals based on a similarity measure. These reliable imputations augment the original dataset, reducing the discrepancy between treatment groups while inducing minimal imputation error. The augmented dataset can then be used to train standard CATE estimation models. We provide theoretical guarantees and extensive numerical studies, demonstrating our approach's effectiveness in improving the accuracy and robustness of various CATE estimation models.

Primary Area: causal reasoning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 8946