# Causal-BALD: Deep Bayesian Active Learning of Outcomes to Infer Treatment-Effects from Observational Data

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

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 66 references that are in the map.
- First posted on arXiv 3 Nov 2021.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://mlsanity.com/paper/1511.07289/) (2015)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Deep Kernel Learning](https://mlsanity.com/paper/1511.02222/) (2015)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [Tune: A Research Platform for Distributed Model Selection and Training](https://mlsanity.com/paper/1807.05118/) (2018)
- [Regularisation of Neural Networks by Enforcing Lipschitz Continuity](https://mlsanity.com/paper/1804.04368/) (2018)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Scalable Variational Gaussian Process Classification](https://mlsanity.com/paper/1411.2005/) (2014)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Active Testing: Sample-Efficient Model Evaluation](https://mlsanity.com/paper/2103.05331/) (2021)
- [On Statistical Bias In Active Learning: How and When To Fix It](https://mlsanity.com/paper/2101.11665/) (2021)
- [Identifying Causal-Effect Inference Failure with Uncertainty-Aware Models](https://mlsanity.com/paper/2007.00163/) (2020)
- ["What We Can't Measure, We Can't Understand": Challenges to Demographic Data Procurement in the Pursuit of Fairness](https://mlsanity.com/paper/2011.02282/) (2020)
- [Interval Estimation of Individual-Level Causal Effects Under Unobserved Confounding](https://mlsanity.com/paper/1810.02894/) (2018)
- [Quantifying Ignorance in Individual-Level Causal-Effect Estimates under Hidden Confounding](https://mlsanity.com/paper/2103.04850/) (2021)
- [Bounds on the conditional and average treatment effect with unobserved confounding factors](https://mlsanity.com/paper/1808.09521/) (2018)
- [Stochastic Batch Acquisition: A Simple Baseline for Deep Active Learning](https://mlsanity.com/paper/2106.12059/) (2021)
- [Bayesian Nonparametric Causal Inference: Information Rates and Learning Algorithms](https://mlsanity.com/paper/1712.08914/) (2017)
- [Active Learning for Decision-Making from Imbalanced Observational Data](https://mlsanity.com/paper/1904.05268/) (2019)
- [Minimax Optimal Nonparametric Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/2002.06471/) (2020)
- [Deconfounding Scores: Feature Representations for Causal Effect Estimation with Weak Overlap](https://mlsanity.com/paper/2104.05762/) (2021)
- [PowerEvaluationBALD: Efficient Evaluation-Oriented Deep (Bayesian) Active Learning with Stochastic Acquisition Functions](https://mlsanity.com/paper/2101.03552/) (2021)

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

- [A Survey on Deep Active Learning: Recent Advances and New Frontiers](https://mlsanity.com/paper/2405.00334/) (2024)
- [Interventions, Where and How? Experimental Design for Causal Models at Scale](https://mlsanity.com/paper/2203.02016/) (2022)
- [Active Learning for Optimal Intervention Design in Causal Models](https://mlsanity.com/paper/2209.04744/) (2022)
- [Black-Box Batch Active Learning for Regression](https://mlsanity.com/paper/2302.08981/) (2023)
- [Sample Constrained Treatment Effect Estimation](https://mlsanity.com/paper/2210.06594/) (2022)
- [Learning Complementary Policies for Human-AI Teams](https://mlsanity.com/paper/2302.02944/) (2023)
- [Ensembled Prediction Intervals for Causal Outcomes Under Hidden Confounding](https://mlsanity.com/paper/2306.09520/) (2023)
- [Advancing Deep Active Learning & Data Subset Selection: Unifying Principles with Information-Theory Intuitions](https://mlsanity.com/paper/2401.04305/) (2024)
- [Improving Interpretability of Deep Active Learning for Flood Inundation Mapping Through Class Ambiguity Indices Using Multi-spectral Satellite Imagery](https://mlsanity.com/paper/2404.19043/) (2024)
- [Bayesian Active Learning for Semantic Segmentation](https://mlsanity.com/paper/2408.01694/) (2024)
- [Generalizing Off-Policy Evaluation From a Causal Perspective For Sequential Decision-Making](https://mlsanity.com/paper/2201.08262/) (2022)
- [Online Decision Mediation](https://mlsanity.com/paper/2310.18601/) (2023)
- [Progressive Generalization Risk Reduction for Data-Efficient Causal Effect Estimation](https://mlsanity.com/paper/2411.11256/) (2024)
- [Towards Regulatory-Confirmed Adaptive Clinical Trials: Machine Learning Opportunities and Solutions](https://mlsanity.com/paper/2503.09226/) (2025)
- [Enhancing Treatment Effect Estimation via Active Learning: A Counterfactual Covering Perspective](https://mlsanity.com/paper/2505.05242/) (2025)
- [X-Factor: Quality Is a Dataset-Intrinsic Property](https://mlsanity.com/paper/2505.22813/) (2025)
- [Task-specific experimental design for treatment effect estimation](https://mlsanity.com/paper/2306.05484/) (2023)
- [Can Active Sampling Reduce Causal Confusion in Offline Reinforcement Learning?](https://mlsanity.com/paper/2312.17168/) (2023)
- [Bayesian Active Learning for Censored Regression](https://mlsanity.com/paper/2402.11973/) (2024)
- [Is merging worth it? Securely evaluating the information gain for causal dataset acquisition](https://mlsanity.com/paper/2409.07215/) (2024)
- [Causal-EPIG: A Prediction-Oriented Active Learning Framework for CATE Estimation](https://mlsanity.com/paper/2509.21866/) (2025)
- [Learning Objective-Specific Active Learning Strategies with Attentive Neural Processes](https://mlsanity.com/paper/2309.05477/) (2023)
- [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)
- [Building Trust: Lessons from the Technion-Rambam Machine Learning in Healthcare Datathon Event](https://mlsanity.com/paper/2207.14638/) (2022)
- [Matched-Pair Experimental Design with Active Learning](https://mlsanity.com/paper/2509.10742/) (2025)
- [Treatment Effect Estimation with Differentiated Networked Effect on Graph Data](https://mlsanity.com/paper/2605.24358/) (2026)
- [Targeted Active Learning for Preference-Based Treatment Effects on Multivariate Outcomes](https://mlsanity.com/paper/2610.03824/) (2026)

## Abstract (arXiv)

Estimating personalized treatment effects from high-dimensional observational data is essential in situations where experimental designs are infeasible, unethical, or expensive. Existing approaches rely on fitting deep models on outcomes observed for treated and control populations. However, when measuring individual outcomes is costly, as is the case of a tumor biopsy, a sample-efficient strategy for acquiring each result is required. Deep Bayesian active learning provides a framework for efficient data acquisition by selecting points with high uncertainty. However, existing methods bias training data acquisition towards regions of non-overlapping support between the treated and control populations. These are not sample-efficient because the treatment effect is not identifiable in such regions. We introduce causal, Bayesian acquisition functions grounded in information theory that bias data acquisition towards regions with overlapping support to maximize sample efficiency for learning personalized treatment effects. We demonstrate the performance of the proposed acquisition strategies on synthetic and semi-synthetic datasets IHDP and CMNIST and their extensions, which aim to simulate common dataset biases and pathologies.