Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
103,458 characters · 32 sections · 116 citation commands
A Primer on Deep Learning for Causal Inference
\singlespacing
\tcblistof[\section*]{boxstyle}{Boxes}
\doublespacing
This primer aims to introduce social science readers to an exciting literature exploring how deep neural networks can be used to estimate causal effects. In recent years, both causal inference frameworks and deep learning have seen rapid adoption across science, industry, and medicine. Causal inference has a long tradition in the social sciences, and social scientists are increasingly exploring the use of machine learning for causal inference Athey2016,wager2018, chernozhukov2016double. Nevertheless, deep learning remains conspicuously underutilized by social scientists compared to other ML approaches, both for causal inference and more generally.
The deep learning revolution has been spurred by the flexibility and expressiveness of these models. Neural networks are nearly non-parametric and can theoretically approximate any continuous function Cybenko1989ApproximationBS, making them well suited for both classification and regression tasks. Furthermore, they can be configured with different architectures and objectives to learn from a variety of quantitative data as well as text, images, video, networks, and speech. These advantages allow them to learn vector “representations" of complex data with emergent properties. Simple examples of representation learning include the Word2Vec algorithm that discovers semantic relationship between words in texts, or face classification models that learn vectors describing facial features mikolov2013. More recently, generative models like DALL-E, Stable Diffusion, and ChatGPT have shown how coherent text passages and life-like images can be reconstructed from learned representations.
Here we explore the potential for leveraging these advantages to estimate causal effects. Causal inference frameworks are non-parametric, but the linear models traditionally used to estimate causal effects require strong parametric assumptions. In contrast, the nearly non-parametric nature of neural networks allows us to estimate smooth response surfaces that capture heterogeneous treatment effects for individual units with low bias.\footnote{Neural networks can have hundreds to billions of parameters making them effectively non-parametric. The risks of overparameterization of neural networks are discussed in Section (ref).} The ability of these models to learn from complex data means we can extend causal inference to new settings where confounding is complicated, time-varying, or even encoded in texts, graphs, or images (see Box (ref) for hypothetical examples). Lastly given the right objectives, neural networks promise to learn deconfounded representations of data, presenting a new strategy for treatment modeling.
This primer synthesizes existing literature on deep causal estimators, but it is not a review; its goals are fundamentally pedagogical and prospective rather than retrospective. In Section (ref), we introduce social scientists to the fundamental concepts of deep learning, and the basic workflow for building and training their own deep neural networks within a supervised learning framework. For readers unfamiliar with causal inference, Section (ref) introduces the assumptions of causal identification and three fundamental estimation strategies within the selection on observables design: matching, outcome modeling, and inverse propensity score weighting. Machine learning models often perform poorly in both theory and practice when only one of these strategies is employed, so we also introduce the concept of double robustness.
Section (ref) is the main body of the article. Here we introduce four related deep learning models for the estimation of heterogeneous treatment effects: the S-learner, T-learner, TARNet and Dragonnet Shalit2017,Shi2019. Although this literature is rapidly evolving, these four models are sufficient to illustrate how traditional estimation strategies can be used in creative ways that leverage the key strengths of neural networks (i.e., deconfounding through representation learning, semi-parametric inference). Section (ref) deals with the practical considerations of building confidence intervals and interpreting neural networks. These guidelines are concretized in the \href{https://github.com/kochbj/Deep-Learning-for-Causal-Inference}{companion online tutorials}, which show readers how to implement and interpret the models described in Section (ref) in Tensorflow 2 and PyTorch.
In Section (ref), we focus on the future of deep causal information: estimators that can disentange counfounding relationships embedded within texts, images, graphs, or time-varying data. In the interest of clarity, we give hypothetical examples of the types of questions social scientists might answer with these models, and briefly describe ongoing research on each of these modalities. For fuller treatments of some of these models, see the appendix. We conclude with a discussion of how neural networks fit into the broader literature on machine learning for causal inference (Section (ref)).
The primer makes multiple contributions. First, it is one of the first pieces in the sociological literature to introduce the fundamentals of deep learning not only at a conceptual level (e.g, backpropagation, representation learning), but at a practical one (e.g., validation, hyperparameter tuning). Our recommendations for training and interpreting neural networks are supported by heavily annotated tutorials that teach readers without prior familiarity with deep learning how to build their own custom models in Tensorflow 2 and PyTorch. Second, we use this foundation and select examples to build intuition on how the core strengths of deep learning can be leveraged for causal inference. Finally, we highlight future directions for this literature and argue why the future of causal estimation runs through deep learning.
Artificial neural networks (ANN) are statistical models inspired by the human brain brand_koch_xu,Goodfellow-et-al-2016. In an ANN, each “neuron" in the network takes the weighted sum of its inputs (the outputs of other neurons) and transforms them using a differentiable, non-linear function (e.g. sigmoid, rectified linear unit) that outputs a value between 0 and 1 if the transformed value is above some threshold. Neurons are arrayed in layers where an input layer takes the raw data, and neurons in subsequent layers take the weighted sum of outputs in previous layers as input. An “output" layer contains a neuron for each of the predicted outcomes with transformation functions appropriate to those outcomes. For example, a regression network that predicts one outcome will have a single output neuron without a transformation function so that it produces a real number. A regression network without any hidden layers corresponds exactly to a generalized linear model (Fig. (ref)A). When additional “hidden" layers are added between the input and output layers, the architecture is called a feed-forward network or multi-layer perceptron (Fig. (ref)B). A neural network with multiple hidden layers is called a “deep" network, hence the name “deep learning" lecun2015deep. A neural network with a single, large enough hidden layer can theoretically approximate any continuous function Cybenko1989ApproximationBS.
Neural networks are trained to predict their outcomes by optimizing a loss function (also called an objective or cost function). During training, the backpropagation algorithm uses calculus's chain rule to assign portions of the total error in the loss function to each neuron in the network. An optimizer, such as the stochastic gradient descent algorithm or the currently popular ADAM algorithm Kingma2014AdamAM, then moves each parameter in the opposite direction of this error gradient. Neural networks first rose to popularity in the 1980s but fell out of favor compared to other machine learning model families (e.g., support vector machines) due to their expense of training. By the late 2000s, improvements to backpropagation, advances in computing power (i.e., graphic cards), and access to larger datasets collectively enabled a deep learning revolution where ANNs began to significantly outperform other model families. Today, deep learning is the hegemonic machine learning approach in industries and fields other than social science.
This section focuses on the practice of training neural networks within a supervised learning framework. While the principles behind supervised machine learning are universal, the workflow for neural networks differs substantially from other ML approaches (e.g., random forests, support vector machines) in practice. Figure (ref) presents this workflow in four different parts: Set Up, Training, Model Evaluation, and Interpretation. We delve into each of these topics in more detail below. Box (ref) contains a basic introduction to supervised learning for unfamiliar readers.
The first step in training a neural network, as in other types of supervised machine learning, is to split your dataset into training, validation, and testing datasets (Fig. (ref)A). If the network is being used for statistical inference, as here, the testing dataset is optional, and inference may be conducted on just the validation set or the full dataset.
While the computational graph and loss function define a deep learning architecture (Box (ref)), actual implementations can vary significantly due to the choice of hyperparameters. In supervised machine learning, hyperparameters are parameters that are not learned automatically when training the model, but must be specified by the analyst. In deep learning, architectural hyperparameters include the number of layers to use for each section of the computational graph, the number of neurons to use in each layer, and the activation functions to be used by neurons. While some basic rules of thumb apply (e.g., use fewer layers than neurons), these choices remain poorly understood theoretically\footnote{For some interesting work on understanding neural networks theoretically from a statistical physics perspective see PDLT-2022.}; Decisions are generally made by comparing empirical performance on the validation set, a practice called hyperparameter tuning.
Neural networks are trained by repeatedly making predictions from the training set, calculating error gradients for each parameter, and backpropagating small fractions of those error gradients. (Fig. (ref) B). A full pass through examples in the training set is called a training loop or epoch. At the beginning of each epoch, the training set is divided into mini-batches of 2 to 1000 units, randomly sampled without replacement. This practice not only aids in memory management, it also improves optimization. Using small random samples reduces the risk of large “exploding” error gradients, particularly early in the training, that could cause the model to overshoot optimal solutions and instead get stuck in local minima.
The size of mini-batches can be considered a hyperparameter.\footnote{In the specific context of causal inference, we recommend not having mini-batches that are too small such that the model can learn from both treated and control units with sufficient overlap.} Because a mini-batch of data is only a sample of a sample (the training dataset), the optimizer only adjusts weight parameters by a fraction of the error gradient (the learning rate) to avoid overfitting. The learning rate is also a hyperparameter, that typically varies between 0.0001 and 0.01.
The non-convex nature of most loss functions\footnote{In convex functions (e.g. the OLS loss), there is a single minimum, so optimizing the function means that you will always converge at the same parameter weights. This is not the case for non-convex functions which may have many local minima.} means that optimization often requires hundreds to potentially millions of epochs of training. Moreover, neural networks are highly susceptible to overfitting because it is easy to overparameterize them with excessive neurons/layers. To ward against overfitting, error metrics on the complete validation set are computed at the end of every epoch. In a regularization practice called “early stopping,” analysts usually stop training once validation metrics stop improving. Other common regularization techniques include weight decay (i.e., $\ell^2$ norm, ridge, or Tikhonov) penalties on the parameters, dropout of neurons during training, and batch normalization.
Dropout is a regularization technique in deep learning where certain nodes are randomly silenced from training during a given epoch Srivastava2014DropoutAS. The general idea of dropout is to force two neurons in the same layer to learn different aspects of the covariate/feature space and reduce overfitting. Batch normalization is another regularization technique applied to a layer of neurons Ioffe2015BatchNA. By standardizing (i.e. z-scoring) the inputs to a layer on a per-batch basis and then rescaling them using trainable parameters, batch normalization smooths the optimization of the loss function. The addition and extent of each of these regularization techniques can be treated as hyperparemeters.
(Tutorial 2 \href{https://colab.research.google.com/drive/1MBUkQO7rh89JrlAV0p1aLNoF9CIB3rZZ?usp=sharing})
After the model has been trained, the analyst compares models assembled with different hyperparameterizations or initial parameter values (Fig. (ref)C). Hyperparameterizations can be chosen using random search, an exhaustive grid search of all possible combinations, or strategic search algorithms like Bayesian hyperparameter optimization or evolutionary optimization snoek2012. Validation loss metrics on the final epoch are commonly used for these comparisons.
Model selection for causal estimators is complicated by the fundamental problem of causal inference: we are not actually interested in the observed “factual” outcomes and propensity scores, but the CATE and ATE. In the case of algorithms like Dragonnet (ref) where the validation loss explicitly targets a causal quantity, we use that as the model selection criterion. In cases where the algorithm is only trained for outcome modeling or propensity modeling, other solutions are needed. In the Appendix, we describe johannson2020's proposal to use matching on a nearest neighbor approximation of the Precision in Estimated Heterogeneous Effects (PEHE), a measure of CATE bias, as an alternative model selection metric (Appendix A.(ref)).
The development of more sophisticated methods for model selection of causal estimators through data simulation is an active area of research within this literature.\footnote{We note that crossfitting zivich2021machine, another approach that has emerged for model selection of other types of machine learning causal estimators may work for the models discussed here, but is likely data-inefficient.} For example, parikh2022 use deep generative models to approximate the data generating distribution under weak, non-parametric assumptions. alaa2019validating independently model each outcome and the propensity score before using influence functions to assess model error.
One comparative advantage of deep learning over other machine learning approaches has been the ability of ANNs to encode and automatically compress informative features from complex data into flexible, relevant “representations" or “embeddings" that make downstream supervised learning tasks easier Goodfellow-et-al-2016, bengio2013deep. While other machine learning approaches may also encode representations, they often require extensive pre-processing to create useful features for the algorithm (i.e., feature engineering). Through the lens of representation learning, a geometric interpretation of the role of each layer in a supervised neural network is to transform its inputs (either raw data or output of previous layers) into a typically lower (but possibly higher) dimensional vector space. As a means to share statistical power, encoded representations can also be jointly learned for two tasks at once in multi-task learning.
The simplest example of a representation might be the final layer in a feed-forward network, where the early layers of the network can be understood as non-linearly encoding the inputs into an array of latent linear features for the output neuron Goodfellow-et-al-2016 (Fig. (ref)B). A famous example of representation learning is the use of neural networks for face detection. Examining the representations produced by each layer of these networks shows that each subsequent layer seems to capture increasingly abstract features of a face (first edges, then noses and eyes, and finally whole faces) lecun2015deep. A more familiar example of representation learning to social scientists might be word vector models like Word2Vec mikolov2013. Word2Vec is a neural network with one hidden layer and one output layer where words that are semantically similar are closer together in the representation space created by the hidden layer of the network.
The novel contribution of deep learning to causal estimation is the proposal that a neural network can learn a function $\Phi$ that produces representations of the covariates decorrelated from the treatment. Fundamentally, the idea is that $\Phi$ can transform the treated and control covariate distributions into a representation space such that they are indistinguishable (Fig. (ref)). To ensure that these representations are also still predictive of the outcome (multi-task learning), multiple loss functions are generally applied simultaneously to balance these objectives. This approach is applied in a majority of the algorithms presented in section (ref).
The papers described in this primer are primarily framed within the Potential Outcomes causal framework (Neyman-Rubin causal model) rubin1974estimating,imbens2015causal. This framework is concerned with identifying the “potential outcomes" of each unit $i$ in the sample, had it received treatment ($Y(1)$) or not received treatment ($Y(0)$). However, because each unit can only receive one treatment regime in reality (being treated or remaining untreated), it is not possible to observe both potential outcomes for each individual (often termed “the fundamental problem of causal inference") holland1986statistics. While we cannot thus identify individual treatment effects $\tau_i=Y_i(1)-Y_i(0)$ for each unit, causal inference frameworks allow us to probabilistically estimate average treatment effects ($ATE$) and average treatment effects conditional on select covariates ($CATE$) across samples of treated and control units. Within this literature, the motivation of many papers is to present algorithms that can both infer CATEs from observational data, but also predict them for out-of-sample units where treatment status is unknown. For readers unfamiliar with causal inference, a short introduction is glossed in Box (ref) with a concrete example, used in the tutorials, in Box (ref).
The ATE is defined as:
$$ATE =\mathbb{E}[Y_i(1)-Y_i(0)] = \mathbb{E}[\tau_i]$$
where $Y_i(1)$ and $Y_i(0)$ are the potential outcomes had the unit $i$ received or not received the treatment, respectively. The CATE is defined as,
$$CATE =\mathbb{E}[Y_i(1)-Y_i(0)|X_i=x] = \mathbb{E}[\tau_i|X_i=x] $$
where $X$ is the set of selected, observable covariates, and $x \in X$.
Within the machine learning literature on causal inference treated here, the primary strategy for causal identification is selection on observables. A challenge to identifying causal effects is the presence of confounding relationships between covariates associated with both the treatment and the outcome.
The key assumptions allowing the identification of causal effects in the presence of confounding is:
1. Conditional Ignorability/Exchangability The potential outcomes $Y(0)$, $Y(1)$ and the treatment $T$ are conditionally independent given $X$, $$Y(0),Y(1)\perp \!\!\! \perp T|X $$ Conditional Ignorability specifies that there are no unmeasured confounders that affect both treatment and outcome outside of those in the observed covariates/features $X$. Additionally $X$ may contain predictors of the outcome (helping precision), but should not contain instrumental variables (hurting precision and potentially amplifying residual bias) or colliders within the conditioning set.\footnote{A variable is a collider if it is caused by two other variables. Controlling for colliding variables, or descendants of colliding variables, will induce a spurious correlation between the parents. In the case of adjusting for confounding, controlling for a collider variable can (re-)open a confounding path that would otherwise be closed, introducing additional bias.}
Other standard assumptions invoked to justify causal identification are:
2. Consistency/Stable Unit Treatment Value Assumption (SUTVA). Consistency specifies that when a unit receives treatment, their observed outcome is exactly the corresponding potential outcome (and the same goes for the outcomes under the control condition). Moreover, the response of any unit does not vary with the treatment assignment to other units (i.e., no network or spillover effects), and the form/level of treatment is homogeneous and consistent across units (no multiple versions of the treatment). Note that this is an identification assumption, based on our understanding of the data generating process, and independent of the model chosen for estimation. More formally, $$T=t \rightarrow Y=Y(T)$$
3. Overlap. For all $x \in X$ (i.e., any observed covariate value), all treatments $t\in \{0,1\}$ have a non-zero probability of being observed in the data, within the “strata" defined by such covariates, $$1 > p(T=t|X=x) > 0$$
4. An additional assumption sometimes invoked at the interface of identification and estimation using neural networks is:
Invertability $$\Phi^{-1}(\Phi(X)) = X$$ In words, there must exist an inverse function of the representation function $\Phi$ encoded by a neural network that can reproduce $X$ from representation space. This is required for the Conditional Ignorability assumption to hold when using representation learning. From a practical perspective, it also means that the representation we created is rich enough to capture the causal relationships we are interested in.
For reference, we describe the full notation used within the review in Box (ref).
Once a strategy for identifying causal effects from available data has been developed (arguably the harder and more important part of causal inference), statistical methods can be used to estimate causal effects by controlling for confounding bias, selection bias, and/or measurement error. There are two fundamental approaches to estimation: treatment modeling to control for correlations between the covariates $X$ and the treatment $T$, and outcome modeling to control for correlations between the treatment $X$ and the outcome $Y$ (Fig. (ref)). Below we briefly review three traditional techniques for removing confounding bias to motivate our systematization of deep learning models. First, we discuss outcome modeling through regression. Next, we consider treatment modeling through non-parametric matching. Finally, we discuss treatment modeling through inverse propensity score weighting (IPW) and introduce the concept of double robustness.
Assuming the treatment effect is constant across covariates/features or the probability of treatment is constant across all covariates/features (both improbable assumptions), the simplest consistent approach to estimating the $ATE$ is to regress the outcome on the treatment indicator and covariates using a linear model.\footnote{Another outcome modeling approach that could be used to estimate the outcome, not discussed here, is g-computation robins1986new,hernancausal2020.} The ATE is then the coefficient of the treatment indicator. Without loss of generality, we call outcome models of this nature, linear or non-linear, $h$:
$$\hat{Y}_i(T)=h(X_i,T)$$
A slightly more sophisticated semi-parametric approach to outcome modeling, used widely in the application of machine learning to causal inference, is to use $h(X,T)$ to impute $\hat{Y}(1)$ and $\hat{Y}(0)$, and calculate the CATE for each unit as a plug-in estimator:
$$\widehat{CATE_i}=\hat{\tau_i}=\hat{Y_i(1)}-\hat{Y_i(0)}=h(X_i,1)-h(X_i,0)$$
and the ATE as:
$$\widehat{ATE}=\frac{1}{N}\sum_{i=1}^N{\hat{\tau_i}}$$
A common treatment-modeling strategy is balancing the treated and control covariate distributions through matching. Matching requires the analyst to select a distance measure that captures the difference in observed covariate distributons between a treated and untreated unit austin2011introduction. Units with treatment status $T$ can then be matched with one or more counterparts with treatment status $1-T$ using a variety of algorithms stuart2010matching. In a one-to-one matching scenario where each treated unit has an otherwise identical untreated counterpart, the covariate distribution of treated and control units is indistinguishable.
Another common approach is inverse propensity score weighting (IPW). In IPW, units are weighted on their inverse propensity to receive treatment. Without loss of generality, we call the propensity function $\pi$. The propensity score is calculated as the probability of receiving treatment conditional on covariates:
$$\pi(X,T)=P(T|X)$$
The simplest IPW estimator of the ATE is then:
Note that only one of the two terms is active for any given unit. Furthermore, this presentation looks different than how the IPW is generally presented because we use $\pi$ as a function with different outputs depending on the value of $T$ rather than a scalar (Box (ref)).\footnote{To de-emphasize the contribution of units with extreme weights due to sparse data, sometimes a “stabilized" IPW is used glynn2010introduction.}
IPW weighting is attractive because if the propensity score $\pi$ is specified correctly, it is an unbiased estimator of the ATE. Moreover, the IPW is consistent if $\pi$ is estimated consistently rosenbaum1983central,glynn2010introduction.
Because different models make different assumptions, it is not uncommon to combine outcome modeling with propensity modeling or matching estimators to create doubly-robust estimators. For example, one of the most widely used doubly-robust estimators is the Augmented-IPW (AIPW) estimator.
The first term is the difference in prediction from two outcome models, one for treated and one for control units, while the last terms is a “corrected” IPW estimator replacing the raw outcome by the residuals from the regression models. As expected, this estimator is unbiased if the IPW and regression estimators are consistently estimated. However, the model is attractive because it will be consistent if either the propensity score $\pi(X,T)$ is correctly specified or the regression model $h$ is consistently specified glynn2010introduction. The model also provide efficiency gains with respect to the use of each model separately, and especially with respect to weighting alone.
Doubly robust estimation is especially important for causal estimation using machine learning. When using simple outcome plug-in estimators, bias is directly dependent on estimation error, which may be different for each potential outcome depending on the modeling strategy kennedy2020. Machine learning estimation of the propensity score can also rely heavily on non-confounding predictors, giving rise to extreme weights schnitzer2016variable. More generally, there are no asymptotic linearity guarantees for machine learning estimators which may converge at a slow rate, leading to misleading confidence intervals naimi2021,zivich2021machine. For these reasons, plug-in machine learning estimation often has poor empirical performance when not using double robust estimators benkeser2017,kennedy2020,zivich2021machine.
The growth of machine learning for causal inference literature has thus been largely driven by the introduction of semi-parametric frameworks. Semi-parametric frameworks address these issues by using machine learning only to estimate the nuissance parameters (i.e., potential outcomes and propensity score) of influence functions for causal parameters like the ATE and CATE chernozhukov2016double,kennedy2016semiparametric,van2011targeted. In these approaches, the estimation of causal parameters is only-second order dependent on machine learning error, there is double-robustness against inconsistent estimation, and guarantees of fast convergence and asymptotically-valid confidence intervals even if the machine learning models converge slowly benkeser2017,kennedy2020,naimi2021,zivich2021machine. We use the final algorithm introduced below, Dragonnet, as an opportunity to provide an intuitive introduction to semi-parametric theory and how it can be used for doubly robust estimation Shi2019.
The architectures proposed in the deep learning literature for causal estimation build upon the core idea discussed above. First, we introduce “S-Learners" and “T-Learners" to show how neural networks can be used to estimate non-linearities in potential outcomes. Second, given the right objectives, a neural network can learn representations of the treated and control distributions that are deconfounded (Fig. (ref)). This approach, which can be related theoretically to non-parametric matching, is illustrated by the foundational TARNet algorithm in section (ref) Shalit2017. Finally, the machine learning for causal inference literature has been largely driven by the introduction of semi-parametric frameworks that allow predictive machine learning models to be plugged-in to doubly robust estimation equations van2011targeted,chernozhukov2016double,chernozhukov2021automatic. In section (ref), we introduce the concept of influence functions and the targeted maximum likelihood estimator to explain the Dragonnet algorithm. For clarity the algorithms presented here all share a familial resemblence to the TARNet algorithm. However, we note that there are many other approaches to using deep learning for causal inference (e.g., the generative models described in Appendix A.(ref)).
S-Learners and T-Learners (Tutorial 1 \href{https://colab.research.google.com/drive/1hjnyfJjFm0wWM3BcZMi0cpW0uBRd5c5f?usp=sharing})
Because at most one potential outcome is unobserved, it is not possible to apply supervised models to directly learn treatment effects. Across econometrics, biostatistics, and machine learning, a common approach to this challenge has been to instead use machine learning to model each potential outcome separately and use plug-in estimators for treatment effects chernozhukov2016double,van2011targeted,wager2018. As with linear models, a single neural model can be trained to learn both potential outcomes (“S[ingle]-learner") (Fig. (ref)B), or two independent models can be trained to learn each potential outcome (a “T-learner") johannson2020 (Fig. (ref)A). In both cases, the neural network estimators would be feed-forward networks tasked with minimizing the MSE in the prediction of observed outcomes. In a slight abuse of notation, the joint loss function for a T-learner can be written as:
where $h_1$ and $h_0$ represent separate networks for each potential outcome.
After training, inputting the same unit into both networks of a T-learner will produce predictions for both potential outcomes: $\hat{Y}(T)$ and $\hat{Y}(1-T)$. We can plug-in these predictions to estimate the $CATE$ for each unit, $$\hat{\tau_i}=(1-2T_i)(\hat{Y_i}(1-T_i)-\hat{Y_i}(T_i))$$ where the first term is a switch to make sure the treated potential outcome comes first. The average treatment effect as, $$\widehat{ATE}=\frac{1}{N}\sum_{i=1}^N\hat{\tau_i}$$
Nearly all of the models described below combine this plug-in outcome modeling approach with other forms of treatment adjustment.
TARNet (Tutorial 1 \href{https://colab.research.google.com/drive/1hjnyfJjFm0wWM3BcZMi0cpW0uBRd5c5f?usp=sharing})
Balancing is a treatment adjustment strategy that aims to deconfound the treatment from outcome by forcing the treated and control covariate distributions closer together Johansson2016. The novel contribution of deep learning to the selection on observables literature is the proposal that a neural network can transform the covariates into a representation space $\Phi$ such that the treated and control covariate distributions are indistinguishable (Fig. (ref)).
To encourage a neural network to learn balanced representations, the seminal paper in this literature, Shalit2017, proposes a simple two-headed neural network called Treatment Agnostic Regression Network (TARNet) that extends the outcome modeling T-learner with shared representation layers (Fig. (ref)B). Each head models a separate potential outcome: one head learns the function $\hat{Y}(1)=h_1(\Phi(X),1)$, and the other head learns the function $\hat{Y}(0)=h_0(\Phi(X),0)$. During training, only one head will receive error gradients at a time (the one predicting the observed outcome). However, both heads backpropagate their gradients to shared representation layers that learn $\Phi(X)$. The idea is that these representation layers must learn to balance the data because they are tasked with predicting both outcomes. The authors of this algorithm have subsequently extended TARNet with additional losses in an algorithm called CFRNET that explicitly encourage balancing by minimizing a statistical distance between the two covariate distributions in representation space (see Appendix A.(ref) for details) Johansson2018,johannson2020.
The complete objective for the network is to fit the parameters of $h$ and $\Phi$ for all $n$ units in the training sample such that,
or more compactly,
where $\mathcal{R}(h)$ is a model complexity term (e.g., for $L_2$ regularization) and $\lambda$ is a hyperparameter chosen through model selection. For coded versions of TARNet in Tensorflow and Pytorch, see Box (ref).
Rather than applying losses directly to the representation function, IPW methods estimate propensity scores from representations using the function $\pi(\Phi(X),T)=P(T|\Phi(X))$. As in traditional IPW estimators, these methods exploit the sufficiency of correctly-specified propensity scores to reweight the plugged-in outcome predictions and provide unbiased estimates of the ATE rosenbaum1983central. Because these models combine outcome modeling with IPW, they retain the attractive statistical properties of doubly robust estimators discussed in section (ref) Atan2018. In this section we focus on Shi2019's Dragonnet model, which adapts semi-parametric estimation theory for batch-wise neural network training in a procedure they call “Targeted Regularization" (TarReg) kennedy2016semiparametric. Given the increasing importance of semi-parametric theory and “double machine learning" across the causal estimation literature, we include a brief introduction to semi-parametric theory and targeted maximum likelihood estimation (TMLE) before diving into the details of the Dragonnet algorithm van2011targeted,chernozhukov2016double.
Dragonnet (Tutorial 3 \href{https://colab.research.google.com/drive/1XzyOINgdSr78_KT7HJRJn07AwkHh-fG8?usp=sharing} / Tutorial 4 \href{https://colab.research.google.com/drive/1NHYTbvGq-cWyy-mm0TrBH2rAmMqtcgPJ?usp=sharing})
A trivial extension to TARNet is to add a third head to predict the propensity score. This third head could use multiple neural network layers or just a single neuron, as proposed in Dragonnet (Fig. (ref)C) Shi2019. Dragonnet uses this additional head to develop a training procedure called “Targeted Regularization" for semi-parametric causal estimation, inspired by “Targeted Maxmimum Likelihood Estimation" (TMLE)van2011targeted.
With three heads, the basic loss function for this network looks like:
with $\alpha$ being a hyperparameter to balance the two objectives. The mean squared error and binary cross-entropy are standard objective functions in machine learning for regression and binary classification, respectively. Note that the first term is simply an expansion of the first term in equation (ref)
Below, we explore how the authors add a second loss on top of this one to allow for semi-parametric estimation.
In recent years, semi-parametric theory has emerged as a dominant theoretical framework for applying machine learning algorithms, including neural networks, to causal estimation chernozhukov2016double,chernozhukov2021automatic,chernozhukov2022riesznet,farrell2021deep,kennedy2016semiparametric,nie2021quasi,van2011targeted,wager2018. The great appeal of these frameworks is that they allow for machine learning algorithms to be plugged-in for non-linear estimates of outcomes and propensity score, while still providing attractive statistical guarantees (e.g., consistency, efficiency, asymptotically-valid confidence intervals).
At a very intuitive level, semi-parametric causal estimation is focused on estimating a target parameter of a distribution $P$ (the $ATE$) of treatment effects $T(P)$ fisher2021visually. While we do not know the true distribution of treatment effects because we lack counterfactuals, we do know some parameters of this distribution (e.g., the treatment assignment mechanism). We can encode these constraints in the form of a likelihood that parametrically defines a set of possible approximate distributions $\mathcal{P}$ from our existing data $P$. Within this set there is a sample-inferred distribution $\tilde{P}\in\mathcal{P}$, that can be used to estimate $T(P)$ using $T(\tilde{P})$.
Regardless of $\tilde{P}$ chosen, $\tilde{P}\neq P \rightarrow T(\tilde{P})\neq T(P)$. We do not know how to pick $\tilde{P}$ with finite data to get the best estimate $T(\tilde{P})$. We can maximize a likelihood function to pick $\tilde{P}$, but there may be “nuisance" parameters in the likelihood that are not the target and we do not care about estimating accurately. Maximum likelihood optimization may provide lower-biased estimates of these nuissance terms at the cost of better estimates of $T(P)$.
To sharpen the likelihood's focus on $T(P)$, we define a “nudge" parameter $\epsilon$ that moves $\tilde{P}$ closer to $P$ (thus moving $T(\tilde{P})$ closer to $T(P)$). An influence curve of $T(P)$ tells us how changes in $\epsilon$ will induce changes in $T(P+\epsilon(\tilde{P}-P))$. We'll use this influence curve to fit $\epsilon$ to get a better approximation of $T(P)$ within the likelihood framework. In particular, there is a specific efficient influence curve (EIC) that provides us with the lowest variance estimates of $T(P)$. In causal estimation, solving the EIC for the ATE yields estimates that are asymptotically unbiased, efficient, and have confidence intervals with (asymptotically) correct coverage.
The EIC for the ATE is,
Setting $EIC_{ATE}$ to it's mean of 0,
The underbraces illustrate how $EIC_{ATE}$ resembles a doubly robust estimator. When the EIC is minimized (set to 0) as in equation (ref), the $ATE$ is equal to the outcome modeling estimate plus a treatment modeling estimate proportional to the residual error.
Targeted Regularization (TarReg) is closely modeled after “Targeted Maxmimum Likelihood Estimation" (TMLE) van2011targeted. TMLE is an iterative procedure where a nuissance parameter $\epsilon$ is used to nudge the outcome models towards sharper estimates of the ATE when minimizing the EIC as in Equation (ref).\footnote{For a deeper dive on targeted learning, we recommend chambaz2020.}
Targeted Regularization takes TMLE and adapts it for a neural network loss function. The main difference is that steps 1 and 2 above are done concurrently by Dragonnet, and that the loss functions for the first three steps are combined into a single loss applied to the whole network at the end of each batch. It requires adding a single free parameter to the Dragonnet network for $\epsilon$.
At a very intuitive level, Targeted Regularization is appealing because it introduces a loss function to TARNet that explicitly encourages the network to learn the mean of the treatment effect distribution, and not just the outcome distribution. The Targeted Regularization procedure proceeds as follows:
In each epoch:
Step 3 of Targeted Regularization is exactly equivalent to minimizing the EIC up to a constant $\beta$.
At the end of training, we can thus estimate the targeted regularization estimate of the ATE $\hat{ATE_{TR}}$ as in TMLE: $$\hat{ATE_{TR}}=\frac{1}{N}\sum_{i=1}^N{ \underbrace{h^*(\Phi(X_i),1)}_{Y^*_i(1)}-\underbrace{ h^*(\Phi(X_i),0)}_{Y^*_i(0)}}$$
Compared to S-learners, T-learners, and TARNet, the Dragonnet algorithm is particularly attractive because of the statistical guarantees afforded by its semiparametric framework. It is doubly robust, unbiased, converges at a rate of $\frac{1}{\sqrt{n}}$, and the sampling distribution is asymptotically normal. Below we describe how to create assymptotically-valid confidence intervals for this estimator.
In this section, we move from theory to practice, and treat best practices for building confidence intervals and interpreting heterogeneous treatment effects. Both of these topics are active areas of development, not only within the causal inference literature, but across machine learning research. Here we specifically focus on recommendations that can be easily implemented by analysts.
(Tutorial 4 \href{https://colab.research.google.com/drive/1NHYTbvGq-cWyy-mm0TrBH2rAmMqtcgPJ?usp=sharing})
In this paper, we feature Dragonnet over other approaches because of its attractive statistical properties. Because the Targeted Regularization procedure in Dragonnet is essentially a variant of TMLE, an asymptotically valid standard error can be calculated as the sample corrected variance of the efficient influence curve $\sigma_{\hat{ATE}}$, where
and,
(van2011targeted, pp. 96)
In Tutorial 5, we show how $\sigma_{\hat{ATE}}$ can be used to calculate a Wald confidence interval for Dragonnet. While not featured in this review, asymptotically valid conference intervals can also be calculated using RieszNet, a variant of Dragonnet introduced in chernozhukov2022riesznet that connects neural network estimation to the automatically debiased machine learning literature currently popular in causal econometrics chernozhukov2016double,chernozhukov2021automatic.
(Tutorial 4 \href{https://colab.research.google.com/drive/1NHYTbvGq-cWyy-mm0TrBH2rAmMqtcgPJ?usp=sharing})
A lack of interpretability has been a barrier to the adoption of machine learning methods like neural networks and random forests in social science settings. However, the literature on post-hoc interpretability techniques has matured considerably over the past five years, and several techniques for identifying important features/covariates such as permutation importance, LIME scores, SHAP scores, Individual Conditional Expectation plots etc... are in widespread usage today altmann2010,goldstein2017,lundberg2017,ribeiro2016should. For a broad and accessible treatment on interpreting machine learning models, see molnar2022interpretable.
Building on criteria used to evaluate other explainable AI methods, crabbe2022benchmarking note four desirable properties of a feature importance technique for the interpretation of deep causal estimators: sensitivity, completeness, linearity, and implementation invariance sundararajan2017axiomatic. A method that is 'sensitive' can distinguish between features that are simply predictive of the outcome, and those that actually influence CATE heterogeneity. A method that is 'complete' identifies all features that, together, explain all effect heterogeneity compared to a baseline. A 'linear' method is one where the feature importance scores additively describe the prediction. Lastly, the approach should be agnostic to both the model architecture (e.g., TARNet, Dragonnet) and different architectural hyperparameterizations (i.e., invariant to implementation). Of the feature importance methods surveyed, they identify two that manifest all four of these qualities: SHAP scores, and integrated gradients.
SHAP (SHapley Additive exPlanations) scores have emerged as one of the most popular methods for evaluating machine learning models in recent years lundberg2017. SHAP is what is called a “local” interpretability method: it provides feature importance estimates for each individual datum. Theoretically, SHAP frames feature importance estimation as a cooperative (game-theoretic) game between covariates to predict a specific outcome. Under the hood, the algorithm exhaustively compares all possible “coalitions” of covariates and their ability to predict the outcome (win the game). Predictions from this powerset of coalitions are used to calculate the additive marginal contributions of each feature in prediction using Shapley values. The disadvantage of SHAP is that, even with computational tricks, calculating scores for every unit can become computationally intractable in high dimensional datasets. SHAP scores are interpreted in comparison to a causal baseline of the ATE.
Because of the computational expense of SHAP scores, crabbe2022benchmarking also recommend another local-interpretability method called “Integrated Gradients” sundararajan2017axiomatic. Intuitively, this algorithm draws a straight-line, linear path in feature space between the target input (individual unit) and a baseline (i.e., a hypothetical unit who is exactly average on all covariates). A feature importance score can then be constructed by calculating the gradient in prediction error along this path with respect to the feature of interest. Note that SHAP scores can also be understood theoretically within the path framework. From this perspective, coalitions are paths in which each feature is turned on sequentially, and the SHAP score is the expectation across these paths. This interpretation leads to a gradient-based algorithm for calculating SHAP scores specifically for neural networks, which is also in the SHAP package. In practice, we recommend that analysts experiment with both integrated gradients and SHAP scores.
To move from theory to empirics, the \href{https://github.com/kochbj/Deep-Learning-for-Causal-Inference/tree/main}{online tutorials} show how to implement many of the ideas presented throughout this primer. The tutorials are hosted in notebooks in the Google Colaboratory environment. When users open a Colab notebook, Google immediately provides a free virtual machine with standard Python machine learning packages available. This means that readers need not install anything on their own computers to experiment with these models. The tutorials are written in the Python programming language and provide examples in both Tensorflow2 and Pytorch, the two most popular deep learning frameworks. We note that both Tensorflow2 and Pytorch have implementations in R. However, we strongly recommend that readers interested in getting into deep learning work in Python, which has a much richer ecosystem of third-party packages for machine learning.
Currently there are five tutorials:
As exciting as neural networks are for heterogeneous treatment effect estimation from quantitative data, a great promise of deep causal estimation is inference when treatments, confounders, and mediators are encoded in high-dimensional data (e.g., text, images, social networks, speech, and video) or are time-varying. This is a strong advantage of neural networks over other machine learning approaches, which do not generalize competitively to non-quantitative data. In these scenarios, multi-task objectives and tailored architectures can be used to learn representations that are simultaneously rich, capture information about causal quantities, and disentangle their relationships. Moreover, the inherent flexibility of neural networks means that, in many cases, the TARNet-style models presented above can serve as the foundations to inference on text and graphs with some architectural modifications, additional losses, and new identification assumptions.
This literature is rapidly evolving, so readers should treat this section of the primer as fundamentally prospective. To maintain accessibility, our primary goal here is to introduce readers to hypothetical scenarios where they might perform causal inference on text, network, or image data. Second, we selectively review contemporary, theoretically-motivated literature on deep causal estimation in these settings. The identification assumptions for different data types differ substantially, so we generally leave those to the interested reader. Finally, we briefly discuss approaches for dealing with time-varying confounding. We also take this section as an opportunity to introduce the Transformer or Graph Neural Network, an architecture now used in most contemporary deep learning models to learn from complex data (Box (ref)).
In recent years, an interdisciplinary community across both social science and computer science has coalesced around causal inference from text (see keith-etal-2020-text and feder2021causal for exhaustive reviews). Broadly speaking, texts may capture information about any causal quantity (treatments, outcomes, confounders, mediators) we might be interested in. For example, in an exit-polling experiment, analysts might want to measure toxicity ($Y$) in text responses to political prompts. In an observational study of e-mail response times ($Y$), analysts might want to measure the effects of the tone of the email ($T$). In this scenario, the analyst might also want to control for confounders like subject matter ($X$). Each of these scenarios presents distinct identification challenges feder2021causal. But in all cases, we can use low-dimensional representations of the high dimensional text to extract, quantify, and disentangle relationships between nuanced qualities like tone and subject matter.
The ability of neural networks to automatically extract features makes them particularly suited for the last scenario when both treatment information and confounding covariates are encoded in text. In many cases, we may not have explicitly identified, quantified, or labeled all of the confounders in text (e.g., subject matter and tone of emails), but we would still like to control for them. pryzant2020causal,Veitch2019UsingTE, and gui2022causal address this problem by prepending Transformer-layers (Box (ref)) for reading text to the beginning of TARNet or Dragonnet. Veitch2019UsingTE demonstrate the viability of this approach on a Science of Science question testing the causal effect of equations on getting papers accepted to computer science conferences. pryzant2020causal,gui2022causal explore the more complicated scenario not in which the treatment is explicitly known (e.g., equations in papers, gender of authors), but is instead externally perceived upon reading (e.g., politeness/rudeness of an email or toxicity of a social media post). In these models, an additional loss function is also added for learning text representations concurrently with the causal inference losses discussed above.
A smaller literature has leveraged relational data for causal inference in two distinct scenarios. In the first traditional selection on observable settings, we wish to control for information about unobserved confounding inferable from homophilous ties. For example, age or gender might be unmeasured in our data, but we might expect people to develop friendship ties with those of the same gender identity or age cohort.
This scenario suggests estimation strategies similar to those when confounders are encoded in text. Much like Transformer layers can be prepended to TARNet-style estimators to learn from text, graph neural networks (an analog of the Transformer) can be preprended to learn from graphs. guo2019counterfactual provides a first pass at this problem by adding GNN layers to CFRNet Shalit2017. Veitch2019UsingET instead adapt Dragonnet in a semi-parametric framework to allow for consistent estimates of the treatment and outcome, assuming the network representation encodes significant information about confounders.
The second, more challenging scenario is estimating the causal effect of social influence on outcomes from observational data. For example, Cristali2022UsingEF introduce the problem of measuring the effects of vaccination ($T$) on peer vaccination choice ($Y$). This is a hard problem because a) SUTVA is a fundamental assumption of all causal inference frameworks and b) it is hard to disentangle whether changes in the outcome result from the treatment via peer effects (e.g, person A pressuring person B to vaccinate), or from homophily (e.g., person A and person B having similar political leanings). In other words, contagion and homophily are generically confounded shalizithomas2011. mcfowland2021 are the first to tackle this problem by making strong parametric assumptions about the generation of network ties and the outcome model. Cristali2022UsingEF instead propose an approach using neural network-learned representations of the graph.
While ideas from causal inference have been leveraged extensively to improve image classification, to our knowledge there are no papers that explore causal inference where treatments, confounders, mediators, or predictors are encoded in images.\footnote{jesson2021quantifying introduce a simulation where the MNIST digit dataset serves as covariates $X$ as toy example of high-dimensional confounding, but not a possible application.} That being said, some scenarios proposed for causal text analysis should apply here as well. For example, consider the conjoint experiment by todorov2005inferences where both the treatment (e.g. incumbency of a politician) and potential latent confounders (e.g., party, age, gender, race) are encoded in an image. In this setting, a TARNet-like model adapted to learn and condition on image representations could improve treatment effect estimation by controlling for confounders such as the politician's age. Causal inference on images is an area ripe for exploration, and we hope to see more work here in the future.
One natural extension of deep causal estimation is to scenarios where treatments are administered over time and confounding may be time-varying. While “g-methods” developed by Robins et al. for estimating effects with time-varying treatments and confounding have existed for decades, the statistical assumptions encoded in these models are quite strong robins1994correcting,robins2000marginal,robins2009longitudinal. Due to their reliance on generalized linear models to define the “structural” component, they assume that the outcome is a linear function of all covariates and treatment. Second, for identification, they make strong assumptions about which previous timesteps confound the current one. Third, they require different coefficients to be estimated at each time steps. Transformers (Box (ref)) and recurrent neural networks, a simpler model for sequential data (Appendix A.(ref)), should be able to capture long-term dependencies and non-linearities in ways that marginal structural models and g-computation cannot.
Several papers have begun to explore these possibilities in the context of personalized medicine. lim2018forecasting build a marginal structural model using a recurrent neural network, and bica2020estimating extend this framework with an additional loss to more explicitly deal with time varying confounding by forcing the model to “unlearn” information about the previous time steps. melnychuk2022 go one step further by adapting bica2020estimating’s approach with a transformer. Inspired by longitudinal targeted maximum likelihood, frauen2022estimating add a semi-parametric targeting layer to their RNN to create a g-computation algorithm that is doubly robust and asymptotically efficient. lietal2021 instead propose an RNN framework for g-computation that allows for dynamic treatment regimes. All of these papers use simulations of tumor growth dynamics, naturalistic simulations based on vital signs from intensive care unit visits, or factual datasets exploring treatment response to physical therapy for back pain.
In this primer we introduce social scientists to the emerging machine learning literature on deep learning for causal inference. To set the stage, we first provide both an intuitive introduction to fundamental deep learning concepts like representation and multi-task learning, as well as practical guidelines for training neural networks. In the main body of the article, we show how ML researchers have adapted core treatment and outcome modeling strategies to leverage the particular strengths of neural networks for heterogeneous treatment effect estimation. We follow with a discussion on inference (e.g., model selection, confidence intervals, interpretation), and closed with a prospective look at algorithms for inference from text, social networks, images, and time varying data.
Deep learning is not the only potential tool for heterogeneous treatment effect inference, and there are robust literatures exploring the usage of other methods in both the econometrics and biostatistics communities van2011targeted,chernozhukov2016double,wager2018. While these literatures are certainly more mature, below we discuss reasons why we think the use gap between neural networks and other machine learning methods will continue to narrow, a change that we must prepare for.
First, neural networks are better at modeling non-linear heterogeneity (e.g., in treatment responses) than other machine learning methods. In extensive simulations, curth2021 found that when the data-generating process for treatment heterogeneity includes exponential relationships, neural networks outperformed random forests, but tree-based methods are robust when the data-generating process is built on linear functions. Neural networks were also consistently better at predicting outlier treatment effects than forests. These differences result from how the two methods model functions. While neural networks can approximate any continuous function with enough neurons, random forests must build non-linear or non-orthogonal decision boundaries using piecewise functions and average predictions. Consistent with these differences, curth2021 also find that neural networks do better when variables are constructed as continuous covariates, and vice versa when they are dichotomized.
From a statistical perspective, the rise of semi-parametric and double machine learning frameworks has also narrowed the gap between neural networks and other types of machine learning in terms of theoretical guarantees. For example, the TMLE-inspired Dragonnet algorithm featured here is unbiased, plausibly consistent, and converges to the target estimand at a fast rate of $\frac{1}{\sqrt{n}}$. The closely-related Riezsnet double machine learning model (not featured) boasts similar guarantees chernozhukov2022riesznet. Beyond these algorithms, there is a growing adjacent literature of model-agnostic plug-in learners (e.g., X-learner, R-learner) that can leverage the strengths of neural networks nie2021quasi,kunzel2019metalearners.
Third, folk beliefs about the data-hungriness and uninterpretability of neural networks are overstated. Neural networks are data-hungry when over-parameterized or learning from high-dimensional data like images, but we show in the tutorials that modest-sized, well-regularized neural networks can successfully infer heterogeneous treatment effects in a naturalistic simulation of quantitative data with less than 800 units. In Section (ref), we also highlight the considerable progress in machine learning interpretability over the past five years, much of which has been on model-agnostic approaches that benefit all black-box algorithms equally.\footnote{Critics often point to out-of-bag feature importances as a particular strength of random forests, but this approach has been shown to be less accurate than model-agnostic permutation importances anyways altmann2010.}
In our opinion, the most pressing limitation of current deep learning approaches is the difficulty of optimizing neural networks. Theoretically, this stems from a) the complexity of the loss functions which are often non-convex, and b) the ease of over-parameterizing these models to fit these functions. If neural networks are to be used as statistical estimators, statistical guarantees must be backed by optimization guarantees and/or more rigorous methods for model selection. Outside of statistical estimation, this limitation has largely been addressed through empirical testing on test data and strategic model selection. Within the statistical estimation context, this gap will likely need to be addressed by simulation-based sensitivity analyses and, in the short term, comparisons to other model families.
Moreover, there has been a lack of mature tools and empirical applications of these models. A major goal of this primer, and the tutorials in particular, is to synthesize the theoretical literature, practical training and interpretation guidelines, and annotated code so that social scientists in one place can start using these models. Deep learning frameworks like Tensorflow and Pytorch are becoming more accessible every year, but we note that canned Python packages like Uber's causalML exist for interested readers who just want to experiment with a few of these models chen2020causalml.
Despite current limitations, we believe the future of causal estimation runs through deep learning. As causal inference ventures into new settings, the flexibility of neural networks will become essential for learning from text, graph, image, video, and speech data. For time-varying settings, we believe the ability of neural networks to model non-linearities and long-range temporal dependencies will ultimately lead to solutions with net weaker assumptions than current approaches. Overall, we are optimistic and excited to see where deep causal estimation heads over the next few years.
The accompanying tutorials are available at \href{https://github.com/kochbj/Deep-Learning-for-Causal-Inference}{https://github.com/kochbj/Deep-Learning-for-Causal-Inference}. The tutorials use the IHDP naturalistic simulation introduced in Hill2011 as an example. The 25 covariates/features for the 747 units (139 treated) in the dataset were taken from an experiment, but Hill simulated the outcomes to create known counterfactuals. The data are available from Fredrik Johansson's website \href{https://www.fredjo.com}{https://www.fredjo.com}.