Causal-EPIG: Causally Aligned Active CATE Estimation

Erdun Gao, Jake Fawkes, Dino Sejdinovic
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:33151-33183, 2026.

Abstract

Estimating the Conditional Average Treatment Effect (CATE) is constrained by the high cost of obtaining outcome measurements, making active learning valuable. However, conventional strategies suffer from a fundamental objective mismatch: they reduce uncertainty in model parameters or observable outcomes rather than the unobservable causal quantities of interest. We address this via the principle of causal objective alignment, positing that acquisition functions should target potential outcomes or CATE directly. We operationalize this through Causal-EPIG, a framework adapting Expected Predictive Information Gain to quantify uncertainty reduction in causal quantities. We derive two distinct strategies: a comprehensive approach that targets the joint potential-outcome structure, and a focused approach that directly targets the CATE estimand for sample efficiency. We provide theoretical justification for our framework, establishing a formal link between CATE estimation error and posterior uncertainty in causal quantities. Extensive experiments demonstrate that our strategies improve sample efficiency over standard baselines, and crucially, reveal that the preferred strategy is context-dependent, contingent on the base estimator and treatment-effect structure. Our framework thus provides a principled guide for sample-efficient CATE estimation in practice.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-gao26i, title = {Causal-{EPIG}: Causally Aligned Active {CATE} Estimation}, author = {Gao, Erdun and Fawkes, Jake and Sejdinovic, Dino}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {33151--33183}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/gao26i/gao26i.pdf}, url = {https://proceedings.mlr.press/v306/gao26i.html}, abstract = {Estimating the Conditional Average Treatment Effect (CATE) is constrained by the high cost of obtaining outcome measurements, making active learning valuable. However, conventional strategies suffer from a fundamental objective mismatch: they reduce uncertainty in model parameters or observable outcomes rather than the unobservable causal quantities of interest. We address this via the principle of causal objective alignment, positing that acquisition functions should target potential outcomes or CATE directly. We operationalize this through Causal-EPIG, a framework adapting Expected Predictive Information Gain to quantify uncertainty reduction in causal quantities. We derive two distinct strategies: a comprehensive approach that targets the joint potential-outcome structure, and a focused approach that directly targets the CATE estimand for sample efficiency. We provide theoretical justification for our framework, establishing a formal link between CATE estimation error and posterior uncertainty in causal quantities. Extensive experiments demonstrate that our strategies improve sample efficiency over standard baselines, and crucially, reveal that the preferred strategy is context-dependent, contingent on the base estimator and treatment-effect structure. Our framework thus provides a principled guide for sample-efficient CATE estimation in practice.} }
Endnote
%0 Conference Paper %T Causal-EPIG: Causally Aligned Active CATE Estimation %A Erdun Gao %A Jake Fawkes %A Dino Sejdinovic %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-gao26i %I PMLR %P 33151--33183 %U https://proceedings.mlr.press/v306/gao26i.html %V 306 %X Estimating the Conditional Average Treatment Effect (CATE) is constrained by the high cost of obtaining outcome measurements, making active learning valuable. However, conventional strategies suffer from a fundamental objective mismatch: they reduce uncertainty in model parameters or observable outcomes rather than the unobservable causal quantities of interest. We address this via the principle of causal objective alignment, positing that acquisition functions should target potential outcomes or CATE directly. We operationalize this through Causal-EPIG, a framework adapting Expected Predictive Information Gain to quantify uncertainty reduction in causal quantities. We derive two distinct strategies: a comprehensive approach that targets the joint potential-outcome structure, and a focused approach that directly targets the CATE estimand for sample efficiency. We provide theoretical justification for our framework, establishing a formal link between CATE estimation error and posterior uncertainty in causal quantities. Extensive experiments demonstrate that our strategies improve sample efficiency over standard baselines, and crucially, reveal that the preferred strategy is context-dependent, contingent on the base estimator and treatment-effect structure. Our framework thus provides a principled guide for sample-efficient CATE estimation in practice.
APA
Gao, E., Fawkes, J. & Sejdinovic, D.. (2026). Causal-EPIG: Causally Aligned Active CATE Estimation. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:33151-33183 Available from https://proceedings.mlr.press/v306/gao26i.html.

Related Material