# Adapting Neural Networks for the Estimation of Treatment Effects

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

## Key facts

- Cited by 287 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 5 Jun 2019.

## Built on by (most cited first, 30 of 287)

- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [The Causal-Neural Connection: Expressiveness, Learnability, and Inference](https://mlsanity.com/paper/2107.00793/) (2021)
- [Diffusion Causal Models for Counterfactual Estimation](https://mlsanity.com/paper/2202.10166/) (2022)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Causal machine learning for predicting treatment outcomes](https://mlsanity.com/paper/2410.08770/) (2024)
- [VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments](https://mlsanity.com/paper/2103.07861/) (2021)
- [Neural Causal Models for Counterfactual Identification and Estimation](https://mlsanity.com/paper/2210.00035/) (2022)
- [Do-PFN: In-Context Learning for Causal Effect Estimation](https://mlsanity.com/paper/2506.06039/) (2025)
- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [RieszNet and ForestRiesz: Automatic Debiased Machine Learning with Neural Nets and Random Forests](https://mlsanity.com/paper/2110.03031/) (2021)
- [Conditional Distributional Treatment Effect with Kernel Conditional Mean Embeddings and U-Statistic Regression](https://mlsanity.com/paper/2102.08208/) (2021)
- [Causal Effect Inference for Structured Treatments](https://mlsanity.com/paper/2106.01939/) (2021)
- [Measuring axiomatic soundness of counterfactual image models](https://mlsanity.com/paper/2303.01274/) (2023)
- [Optimal Transport for Treatment Effect Estimation](https://mlsanity.com/paper/2310.18286/) (2023)
- [Foundation Models for Causal Inference via Prior-Data Fitted Networks](https://mlsanity.com/paper/2506.10914/) (2025)
- [A Critical Look at the Consistency of Causal Estimation With Deep Latent Variable Models](https://mlsanity.com/paper/2102.06648/) (2021)
- [Estimating average causal effects from patient trajectories](https://mlsanity.com/paper/2203.01228/) (2022)
- [Causal-BALD: Deep Bayesian Active Learning of Outcomes to Infer Treatment-Effects from Observational Data](https://mlsanity.com/paper/2111.02275/) (2021)
- [Estimating Average Treatment Effects via Orthogonal Regularization](https://mlsanity.com/paper/2101.08490/) (2021)
- [SurvITE: Learning Heterogeneous Treatment Effects from Time-to-Event Data](https://mlsanity.com/paper/2110.14001/) (2021)
- [Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2202.01336/) (2022)
- [Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation](https://mlsanity.com/paper/2211.01939/) (2022)
- [Explicit Feature Interaction-aware Uplift Network for Online Marketing](https://mlsanity.com/paper/2306.00315/) (2023)
- [Sharp Bounds for Generalized Causal Sensitivity Analysis](https://mlsanity.com/paper/2305.16988/) (2023)
- [Algorithm Fairness in AI for Medicine and Healthcare](https://mlsanity.com/paper/2110.00603/) (2021)
- [Normalizing Flows for Interventional Density Estimation](https://mlsanity.com/paper/2209.06203/) (2022)
- [Unbiased Scene Graph Generation via Two-stage Causal Modeling](https://mlsanity.com/paper/2307.05276/) (2023)
- [Scalable Sensitivity and Uncertainty Analysis for Causal-Effect Estimates of Continuous-Valued Interventions](https://mlsanity.com/paper/2204.10022/) (2022)
- [Estimating Categorical Counterfactuals via Deep Twin Networks](https://mlsanity.com/paper/2109.01904/) (2021)
- [Doubly Robust Causal Effect Estimation under Networked Interference via Targeted Learning](https://mlsanity.com/paper/2405.03342/) (2024)

## Abstract (arXiv)

This paper addresses the use of neural networks for the estimation of treatment effects from observational data. Generally, estimation proceeds in two stages. First, we fit models for the expected outcome and the probability of treatment (propensity score) for each unit. Second, we plug these fitted models into a downstream estimator of the effect. Neural networks are a natural choice for the models in the first step. The question we address is: how can we adapt the design and training of the neural networks used in the first step in order to improve the quality of the final estimate of the treatment effect? We propose two adaptations based on insights from the statistical literature on the estimation of treatment effects. The first is a new architecture, the Dragonnet, that exploits the sufficiency of the propensity score for estimation adjustment. The second is a regularization procedure, targeted regularization, that induces a bias towards models that have non-parametrically optimal asymptotic properties `out-of-the-box`. Studies on benchmark datasets for causal inference show these adaptations outperform existing methods. Code is available at github.com/claudiashi57/dragonnet.