CATE Estimation With Potential Outcome Imputation From Local Regression

Ahmed Aloui, Juncheng Dong, Cat P. Le, Vahid Tarokh

Keywords: conditional average treatment effect, causal inference, local imputation, data augmentation

TL;DR: A model-agnostic data augmentation method for CATE estimation that leverages contrastive learning to impute missing potential outcomes, reducing treatment group discrepancies while ensuring minimal imputation error.

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 issue, we propose a model-agnostic data augmentation method for CATE estimation. First, we derive regret bounds for general data augmentation methods suggesting that a small imputation error may be necessary for accurate CATE estimation. Inspired by this idea, we propose a contrastive learning approach that reliably imputes missing potential outcomes for a selected subset of individuals formed using a similarity measure. We augment the original dataset with these reliable imputations to reduce the discrepancy between different treatment groups while inducing minimal imputation error. The augmented dataset can subsequently be employed to train standard CATE estimation models. We provide both theoretical guarantees and extensive numerical studies demonstrating the effectiveness of our approach in improving the accuracy and robustness of numerous CATE estimation models.

Readers: auai.org/UAI/2025/Conference, auai.org/UAI/2025/Conference/Area_Chairs, auai.org/UAI/2025/Conference/Reviewers, auai.org/UAI/2025/Conference/Submission616/Authors, auai.org/UAI/2025/Conference/Submission616/Reproducibility_Reviewers

Submission Number: 616