# Individual Treatment Effect Estimation Through Controlled Neural Network Training in Two Stages

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 10 of its 42 references that are in the map.
- First posted on arXiv 21 Jan 2022.

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

- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [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)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Deep Neural Networks for Estimation and Inference](https://mlsanity.com/paper/1809.09953/) (2018)
- [Orthogonal Random Forest for Causal Inference](https://mlsanity.com/paper/1806.03467/) (2018)
- [Double/Debiased Machine Learning for Treatment and Causal Parameters](https://mlsanity.com/paper/1608.00060/) (2016)

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

- [Heterogeneous Treatment Effect with Trained Kernels of the Nadaraya-Watson Regression](https://mlsanity.com/paper/2207.09139/) (2022)
- [BENK: The Beran Estimator with Neural Kernels for Estimating the Heterogeneous Treatment Effect](https://mlsanity.com/paper/2211.10793/) (2022)
- [Attention-Based Estimation of the Individual Treatment Benefit Probability under Dose Variation](https://mlsanity.com/paper/2606.13821/) (2026)

## Abstract (arXiv)

We develop a Causal-Deep Neural Network (CDNN) model trained in two stages to infer causal impact estimates at an individual unit level. Using only the pre-treatment features in stage 1 in the absence of any treatment information, we learn an encoding for the covariates that best represents the outcome. In the 2nd stage we further seek to predict the unexplained outcome from stage 1, by introducing the treatment indicator variables alongside the encoded covariates. We prove that even without explicitly computing the treatment residual, our method still satisfies the desirable local Neyman orthogonality, making it robust to small perturbations in the nuisance parameters. Furthermore, by establishing connections with the representation learning approaches, we create a framework from which multiple variants of our algorithm can be derived. We perform initial experiments on the publicly available data sets to compare these variants and get guidance in selecting the best variant of our CDNN method. On evaluating CDNN against the state-of-the-art approaches on three benchmarking datasets, we observe that CDNN is highly competitive and often yields the most accurate individual treatment effect estimates. We highlight the strong merits of CDNN in terms of its extensibility to multiple use cases.