# CATE Estimation With Potential Outcome Imputation From Local Regression

arXiv: https://arxiv.org/abs/2311.03630 · HTML: https://mlsanity.com/paper/2311.03630/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 49 references that are in the map.
- First posted on arXiv 7 Nov 2023.

## Builds on (most cited first, 20 of 20)

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [What Makes for Good Views for Contrastive Learning?](https://mlsanity.com/paper/2005.10243/) (2020)
- [Self-Supervised Learning of Pretext-Invariant Representations](https://mlsanity.com/paper/1912.01991/) (2019)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [A Survey of Data Augmentation Approaches for NLP](https://mlsanity.com/paper/2105.03075/) (2021)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Discriminative Unsupervised Feature Learning with Exemplar Convolutional Neural Networks](https://mlsanity.com/paper/1406.6909/) (2014)
- [Recursive Partitioning for Heterogeneous Causal Effects](https://mlsanity.com/paper/1504.01132/) (2015)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [GridMask Data Augmentation](https://mlsanity.com/paper/2001.04086/) (2020)
- [Unsupervised Learning by Predicting Noise](https://mlsanity.com/paper/1704.05310/) (2017)
- [Uniform Error Bounds for Gaussian Process Regression with Application to Safe Control](https://mlsanity.com/paper/1906.01376/) (2019)
- [Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks](https://mlsanity.com/paper/1810.00656/) (2018)
- [Selfie: Self-supervised Pretraining for Image Embedding](https://mlsanity.com/paper/1906.02940/) (2019)
- [RieszNet and ForestRiesz: Automatic Debiased Machine Learning with Neural Nets and Random Forests](https://mlsanity.com/paper/2110.03031/) (2021)
- [Estimating Causal Effects using a Multi-task Deep Ensemble](https://mlsanity.com/paper/2301.11351/) (2023)

## Built on by (most cited first, 2 of 2)

- [OASIS: Conditional Distribution Shaping for Offline Safe Reinforcement Learning](https://mlsanity.com/paper/2407.14653/) (2024)
- [An Analysis of Causal Effect Estimation using Outcome Invariant Data Augmentation](https://mlsanity.com/paper/2510.25128/) (2025)

## Abstract (arXiv)

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.