# Covariate-Balancing-Aware Interpretable Deep Learning models for Treatment Effect Estimation

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 44 references that are in the map.
- First posted on arXiv 7 Mar 2022.

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

- [Gradient Descent Provably Optimizes Over-parameterized Neural Networks](https://mlsanity.com/paper/1810.02054/) (2018)
- [Causal Effect Inference with Deep Latent-Variable Models](https://mlsanity.com/paper/1705.08821/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Neural Additive Models: Interpretable Machine Learning with Neural Nets](https://mlsanity.com/paper/2004.13912/) (2020)
- [Generalized Random Forests](https://mlsanity.com/paper/1610.01271/) (2016)
- [Adapting Neural Networks for the Estimation of Treatment Effects](https://mlsanity.com/paper/1906.02120/) (2019)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Estimating the Effects of Continuous-valued Interventions using Generative Adversarial Networks](https://mlsanity.com/paper/2002.12326/) (2020)
- [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)
- [Double/Debiased/Neyman Machine Learning of Treatment Effects](https://mlsanity.com/paper/1701.08687/) (2017)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)
- [Deep Neural Networks for Estimation and Inference](https://mlsanity.com/paper/1809.09953/) (2018)
- [Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks](https://mlsanity.com/paper/1810.00656/) (2018)
- [Support points](https://mlsanity.com/paper/1609.01811/) (2016)
- [Causal Effect Inference for Structured Treatments](https://mlsanity.com/paper/2106.01939/) (2021)
- [Deep Counterfactual Networks with Propensity-Dropout](https://mlsanity.com/paper/1706.05966/) (2017)
- [A Dynamical View on Optimization Algorithms of Overparameterized Neural Networks](https://mlsanity.com/paper/2010.13165/) (2020)
- [Energy Balancing of Covariate Distributions](https://mlsanity.com/paper/2004.13962/) (2020)

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

- [Benchmarking Heterogeneous Treatment Effect Models through the Lens of Interpretability](https://mlsanity.com/paper/2206.08363/) (2022)
- [NAISR: A 3D Neural Additive Model for Interpretable Shape Representation](https://mlsanity.com/paper/2303.09234/) (2023)
- [DISCRET: Synthesizing Faithful Explanations For Treatment Effect Estimation](https://mlsanity.com/paper/2406.00611/) (2024)
- [CausalKANs: interpretable treatment effect estimation with Kolmogorov-Arnold networks](https://mlsanity.com/paper/2509.22467/) (2025)
- [LucidAtlas: Learning Uncertainty-Aware, Covariate-Disentangled, Individualized Atlas Representations](https://mlsanity.com/paper/2502.08445/) (2025)

## Abstract (arXiv)

Estimating treatment effects is of great importance for many biomedical applications with observational data. Particularly, interpretability of the treatment effects is preferable for many biomedical researchers. In this paper, we first provide a theoretical analysis and derive an upper bound for the bias of average treatment effect (ATE) estimation under the strong ignorability assumption. Derived by leveraging appealing properties of the Weighted Energy Distance, our upper bound is tighter than what has been reported in the literature. Motivated by the theoretical analysis, we propose a novel objective function for estimating the ATE that uses the energy distance balancing score and hence does not require correct specification of the propensity score model. We also leverage recently developed neural additive models to improve interpretability of deep learning models used for potential outcome prediction. We further enhance our proposed model with an energy distance balancing score weighted regularization. The superiority of our proposed model over current state-of-the-art methods is demonstrated in semi-synthetic experiments using two benchmark datasets, namely, IHDP and ACIC.