# Causal-EPIG: A Prediction-Oriented Active Learning Framework for CATE Estimation

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 60 references that are in the map.
- First posted on arXiv 26 Sep 2025.

## Builds on (most cited first, 30 of 33)

- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Discussion on "Bayesian Regression Tree Models for Causal Inference: Regularization, Confounding, and Heterogeneous Effects" by Hahn, Murray and Carvalho](https://mlsanity.com/paper/2108.02836/) (2021)
- [External Validity: From Do-Calculus to Transportability Across Populations](https://mlsanity.com/paper/1503.01603/) (2015)
- [Removing Hidden Confounding by Experimental Grounding](https://mlsanity.com/paper/1810.11646/) (2018)
- [Prediction-Oriented Bayesian Active Learning](https://mlsanity.com/paper/2304.08151/) (2023)
- [Interval Estimation of Individual-Level Causal Effects Under Unobserved Confounding](https://mlsanity.com/paper/1810.02894/) (2018)
- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [Causal Inference in Recommender Systems: A Survey and Future Directions](https://mlsanity.com/paper/2208.12397/) (2022)
- [On the role of surrogates in the efficient estimation of treatment effects with limited outcome data](https://mlsanity.com/paper/2003.12408/) (2020)
- [Stochastic Batch Acquisition: A Simple Baseline for Deep Active Learning](https://mlsanity.com/paper/2106.12059/) (2021)
- [Active Bayesian Causal Inference](https://mlsanity.com/paper/2206.02063/) (2022)
- [Applied Causal Inference Powered by ML and AI](https://mlsanity.com/paper/2403.02467/) (2024)
- [Causal-BALD: Deep Bayesian Active Learning of Outcomes to Infer Treatment-Effects from Observational Data](https://mlsanity.com/paper/2111.02275/) (2021)
- [Efficient Adaptive Experimental Design for Average Treatment Effect Estimation](https://mlsanity.com/paper/2002.05308/) (2020)
- [Semiparametric Efficient Inference in Adaptive Experiments](https://mlsanity.com/paper/2311.18274/) (2023)
- [Transductive Active Learning: Theory and Applications](https://mlsanity.com/paper/2402.15898/) (2024)
- [Bayesian Causal Inference: A Critical Review](https://mlsanity.com/paper/2206.15460/) (2022)
- [A General Algorithm for Deciding Transportability of Experimental Results](https://mlsanity.com/paper/1312.7485/) (2013)
- [Active Learning for Decision-Making from Imbalanced Observational Data](https://mlsanity.com/paper/1904.05268/) (2019)
- [Stratification Trees for Adaptive Randomization in Randomized Controlled Trials](https://mlsanity.com/paper/1806.05127/) (2018)
- [Balancing Covariates in Randomized Experiments with the Gram-Schmidt Walk Design](https://mlsanity.com/paper/1911.03071/) (2019)
- [Black-Box Batch Active Learning for Regression](https://mlsanity.com/paper/2302.08981/) (2023)
- [Active Adaptive Experimental Design for Treatment Effect Estimation with Covariate Choices](https://mlsanity.com/paper/2403.03589/) (2024)
- [Stochastic Tree Ensembles for Estimating Heterogeneous Effects](https://mlsanity.com/paper/2209.06998/) (2022)
- [Sample Constrained Treatment Effect Estimation](https://mlsanity.com/paper/2210.06594/) (2022)
- [Designing Randomized Experiments to Predict Unit-Specific Treatment Effects](https://mlsanity.com/paper/2310.18500/) (2023)
- [Optimal Causal Annotations: An Application to Casenotes in Social Services](https://mlsanity.com/paper/2502.10605/) (2025)
- [Towards Regulatory-Confirmed Adaptive Clinical Trials: Machine Learning Opportunities and Solutions](https://mlsanity.com/paper/2503.09226/) (2025)

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

- [ActiveCQ: Active Estimation of Causal Quantities](https://mlsanity.com/paper/2509.24293/) (2025)
- [Observationally Informed Adaptive Causal Experimental Design](https://mlsanity.com/paper/2603.03785/) (2026)

## Abstract (arXiv)

Estimating the Conditional Average Treatment Effect (CATE) is often constrained by the high cost of obtaining outcome measurements, making active learning essential. However, conventional active learning strategies suffer from a fundamental objective mismatch. They are designed to reduce uncertainty in model parameters or in observable factual outcomes, failing to directly target the unobservable causal quantities that are the true objects of interest. To address this misalignment, we introduce the principle of causal objective alignment, which posits that acquisition functions should target unobservable causal quantities, such as the potential outcomes and the CATE, rather than indirect proxies. We operationalize this principle through the Causal-EPIG framework, which adapts the information-theoretic criterion of Expected Predictive Information Gain (EPIG) to explicitly quantify the value of a query in terms of reducing uncertainty about unobservable causal quantities. From this unified framework, we derive two distinct strategies that embody a fundamental trade-off: a comprehensive approach that robustly models the full causal mechanisms via the joint potential outcomes, and a focused approach that directly targets the CATE estimand for maximum sample efficiency. Extensive experiments demonstrate that our strategies consistently outperform standard baselines, and crucially, reveal that the optimal strategy is context-dependent, contingent on the base estimator and data complexity. Our framework thus provides a principled guide for sample-efficient CATE estimation in practice.