Estimating causal effects with the neural autoregressive density estimator

2.1 Non-linear causal effect estimation

The idea of estimating causal effects in non-linear settings is not new. For example, Hill [14] proposed a non-parametric Bayesian method – the Bayesian additive regression trees (BART) – to estimate average treatment effects (ATEs). More recently, the literature on non-linear causal effect estimation focused on estimating individual treatment effects, for example, by using representations of a feed-forward neural network [15,16], multi-task Gaussian processes [17], or a variant of a generative adversarial networks [18]. The approach proposed in this article can be viewed as a generalization of the idea in Hartford et al. [19], where deep neural networks are used to estimate causal effects using instrumental variables (deepIV), and further extended by Bennett et al. [20] to a generalized method of moments approach for instrumental variables. However, both papers are only concerned with the estimation of causal effects using instrumental variables, a special case of causal identification. All of this literature is mainly focused on the potential outcome framework literature, where the assumptions related to the causal graph are not stated as transparently as in Pearl’s framework.

There has been a recent surge in the literature combining the causal graph and formalism with neural estimation approaches. In particular, ref. [21] used a neural network that takes as input the adjacency matrix of a causal graph in order to predict parameters of a sum-product network that estimates an interventional distribution. This work differs from ours in two important ways. First, their proposed method relies on heterogeneous data, that is, both observational and interventional data; and second, their method does not require to make explicit assumptions about the distribution of the variables that conform the system, while our method does. The work by Pawlowski et al. [22] is very close to the method proposed in this article. There, the authors used deep generative models such as variational autoencoders and normalizing flows in order to estimate counterfactuals from data. Their approach requires assumptions similar to ours: both that the topology of the causal graph is known and some assumptions on the distribution of the variables of the system. They focus on the estimation of counterfactuals, which is the third level in Pearl’s ladder of causation.

2.2 Neural autoregressive density estimators

The literature related to autoregressive density estimators is comprehensive. As mentioned in Section 1, we are interested in approaches that use neural networks to learn useful representations from data. Here, we can highlight Bengio and Bengio [6] who introduced the idea of using neural networks to estimate the parameters of distributions of any variable given previously observed variables. Larochelle and Murray [4], Uria et al. [7], and Germain et al. [23] extended this idea in order to, for example, make the neural network training more efficient or model a richer family of distributions. To the best of our knowledge, NADEs have not been used for causal inference.

3.1 Assumptions

In order to use neural autoregressive density estimators to estimate causal effects we require the following assumptions:

• The underlying causal graph of the process is known. This graph is a directed acyclic graph (DAG) as presented in Section 3.2. Relaxations with respect to this assumption are explored in Section 4.6.

• The distribution family of every variable in the process is known. This is not so restrictive as we will see in Section 4.3.

These assumptions, either implicitly or explicitly, are in accordance with the common practice when estimating interventional distributions. However, it is worth stressing that in contrast to the majority of models in the literature, we do not make any assumptions about the functional form of the relationships between the variables and the corresponding parameters. These assumptions are relaxed exactly for the reason of being able to estimate more robust causal effect models for which these relationships are not known in advance.

3.2 Graphical models

Consider a set of J random variables X 1 , … , X J representing a system of interest. Furthermore, let lowercase letter variables x 1 , … , x J denote realizations of these random variables. We can represent such a system with a graph G composed of vertices V that represent the variable, and edges E that represent the conditional dependencies between the variables. In this article, we work with DAG, meaning that (1) edges are directed from one variable to another (for example, X i → X j ), and (2) there are no cycles in the graph; in other words, there are no directed paths of a variable to itself (there is no j with X j → ⋯ → X j ). The inclusion of both variables and conditional dependencies in a causal graph entails a joint probability distribution of the variables. The joint probability distribution is composed of the product of the distribution of each variable, conditioned on its parents on the graph P ( X ) = ∏ j J P ( X j ∣ PA ( X j ) ) . This is also known as the causal Markov condition.

Working with graphs allows researchers to include causal semantics in their models, that is, it allows them to include a hypothesized true data generating process. A direct implication from such a hypothesized model is that researchers can simulate interventions to their system of interest. In addition, if the hypothesized model is correct and the mechanisms involved in the process are known, the outcome of such an intervention should correspond to an intervention in the real world. Mathematically, the interventional distribution of an outcome variable X o under an intervention X j = x j can be defined as P ( X o ∣ do ( X j = x j ) ) [2]. In this equation, the d o operator, “ do ( X j = x j ) ,” expresses that X j is forced to be x j , and any incoming edge into X j is deleted.

3.3 Do-calculus

However, in many cases, the mechanisms are not known exactly, and performing an intervention in the real world is not possible. Therefore, it is required to estimate the effect of an intervention from observational data. In such cases, the interventional distributions can be derived as operations of observational distributions using the rules of do-calculus [2]. We define G X j ¯ as the graph G after deleting all the incoming arrows to X j . In a similar way, we define G X j ̲ as the graph G after deleting all the outcoming arrows from X j . Finally, we define the conditional independence relation of X j , X i conditioned on X k as X j ⊥ ⊥ X i ∣ X k .

Theorem 3.4.1 in ref. [2], which outlines the rules of do-calculus, allowing us to identify interventional distributions from observational ones reads as follows:

Let G be the DAG associated with a causal model, and let P ( ⋅ ) be the probability distribution induced by that model. For any disjoint subsets of variables X , Y , Z , and W , we have the following rules:

Rule 1 (Insertion/deletion of observations):

Rule 2 (Action/observation exchange):

Rule 3 (Insertion/deletion of actions):

where Z ( W ) is the set of Z -nodes that are not ancestors of any W -node in G X ¯ .

The interventional distribution allows us to estimate other quantities of interest, for example, the ATE of treatment x j relative to treatment x ˆ j , on X o as

3.4 Autoregressive density estimators

Before introducing the relationship with causal models, we briefly explain the concept of NADEs. These models are a family of generative models that attempt to estimate the joint distribution of the variables by factorizing the joint distribution into a product of conditional distributions [4,6,23]. By using the chain rule of probability, any joint probability distribution P ( X ) can be factorized as a conditional distribution where the probability of X j is conditioned on a function f j of its “parents:” PA ( X j ) . Note, however, that in the machine learning literature these parents have no causal semantics as explained in the previous section. This applies to any arbitrary graph that contains all the variables of interest:

Even though ref. [6] suggested that these models can be constructed on the basis of probabilistic causal models, most of the machine learning literature has moved away from this approach. Instead, the machine learning literature has focused on proving that the densities estimated by these models are robust with respect to changes in the joint probability factorization. There is a body of evidence from the deep generative model literature that suggests this is indeed possible [7,23,24]. Some of these models are valid as an approximation of associational inference. However, as we are interested in causal inference, and more specifically, interventional type of inference, further assumptions with respect to the factorization of the distribution are required. In other words, we cannot just impose an arbitrary factorization of conditional distributions as in the body of neural autoregressive estimator literature. Instead, we must apply a “causal factorization” as discussed in ref. [25]. The individual conditional distributions are also called independent causal mechanisms or Markov Kernels [26].

We propose to parametrize the functions f j as independent fully connected feed-forward neural networks and use them as independent causal mechanisms in order to estimate the effects of interventions. For every variable X j in our system, there is a neural network. Each of the neural networks takes the parents of the variable of interest PA ( X j ) and outputs parameters that define the distribution of X j . For example, if X j is a binary variable, the network outputs a probability, p , of X j = 1 . Likewise, if X j is continuous and approximately Gaussian distributed, the neural network would output the mean μ and the standard deviation σ of X j . These networks are trained jointly, using the negative log-likelihood in equation (2) as the loss function. Figure 1 illustrates an example of such a model for a causal graph of the “kidney stone” example [26,27]. This graph contains assumptions related to the conditional dependencies between variables, for example, that R (recovery) is dependent on both KS (size of kidney stones) and T (treatments). Each network takes parents of each variable as input and outputs the parameters of the respective variable. In this case, the output is a single neuron, since each random variable follows a Bernoulli distribution with parameter p .

Figure 1

Example of an NADE model. (Left) The causal graph we want to represent with neural networks. (Right) The NADE representing the causal graph. All the variables are parameterized as Bernoulli distributed random variables.

3.5 Connection to Pearl’s do-calculus

The way to connect the model (described in the previous subsection) to the do-calculus framework is by realizing that equation (2) is closely related to the truncated factorization formula [2]. For an intervention where X j takes a single value, do ( X j = x j ′ ) , this corresponds to

The neural networks will approximate each of these conditional distributions or independent causal mechanisms [26] that are fundamental to the estimation of the causal effects. In order to compute causal effects, it is required that effects can be identified and that we use the rules of the do-calculus as explained in more detail in refs. [2,3] and [28].

4.1 Experiment I: binary variables

To estimate causal effects with binary outcomes from Algorithm 1, we need to estimate the distributions in equation (4). In order to do this, we used neural networks as shown in Figure 1. There is one neural network for each variable in the system. Each neural network takes the parents of the corresponding variable as input and produces the parameters of a distribution as output. In this case, all the variables were modelled as random variables with Bernoulli distributions. This means that, for all neural networks in our system, the output of each one of them was a single parameter p .

Table 2 presents the estimated conditionals. The fact that we are able to recover the probabilities from the generated data is not surprising, since knowing the structure of the causal graph reduces the estimation problem to a conditional density estimation problem. For a linear model, the resulting ATE is 6.43%. This value is close to the analytical one since, in this model, the probability of the outcome does not include non-linear relationships between the covariates.

4.2 Experiment II: continuous response variables

In the second simulation experiment, we generated the confounding variable KS, and the treatment variable T as before, but now changed the recovery variable R to be normally distributed with R ∼ N ( R ∣ μ = T ∗ 4 + exp ( KS ) , σ = 2 ) . That is, we changed Line 4 of Algorithm 1 to follow a normal distribution conditional on treatment and kidney size. The analytical causal effect of the treatment variable is 4. In this case, the output of the neural network that corresponds to the recovery variable has two units, one for the mean and one for the variance of a normal distribution. In that way, we are able to estimate the likelihood of the parameters, conditioned on our data, given the assumption of a normally distributed recovery variable. Table 3 shows the results of the proposed method. The treatment effect from the generated data is equivalent to the effect derived analytically. The difference between the analytical ATE and the neural approach is 0.11 points, while the difference with the linear model is 0.27 points. Neither of these differences is large enough to invalidate one of the approaches. However, this phenomenon does tell that the neural approach is competitive also in linear settings.

4.3 Experiment III: continuous confounding variables

In the third simulation experiment, the confounding variable KS was generated as a continuous variable. In this experiment, the aim was to test the impact of the second assumption in Section 3.1 that the distribution family of every variable in the process is known. To do this, two different data sets with different distributions of the confounding variable – gamma and lognormal – were generated. That is, we changed Line 2 in Algorithm 1 with a gamma distribution (left column of Table 4) or a lognormal distribution (right column of the same table). Due to the non-differentiable nature of its parameters, gamma distributions cannot be learnt in a straightforward way using neural networks. As a result, we used a lognormal parameterization for both generated data sets. That is, in the first neural network of Figure 1, the output is now a mean and a variance with which we can evaluate the likelihood using the kidney stone data.

The treatment variable was still a binary variable. A cutoff equal of 10 (the mean of the confounder) was applied to decide whether the kidney stone is small or large T = 1 { KS > 10 } . After deciding whether the kidney stone is large or small, we assigned the treatment using the same probabilities as in Table 1. Finally, the response variable was defined as a normal distribution R ∼ N ( R ∣ μ = T ∗ 4 + KS , σ = 2 ) . In the second data set, the confounding variable was drawn from a lognormal distribution with μ = 2.5 and σ = 0.25 . The rest of the DGP was the same.

The DGP for this experiment can be summarized as follows:

Since we have a continuous confounding variable, the sum in equation (4) turns into an integral. If the integral is mathematically intractable, numerical approximations can be used. It is proposed here to use Monte Carlo integration in order to obtain an approximation of the interventional distribution and the ATE. To do this, we sample from the inferred distribution and propagate forward through the neural network to get different estimates of the response variable.

The results of this experiment are shown in Table 4. Even though the analytical treatment effects are not recovered perfectly, they are approximated quite accurately in this way. The ATE is the same for any number of samples, as the functional relation between the mean of the outcome and the confounder is linear.

These results also reveal that the assumption of knowing the exact parameterization of the distribution is not critical. This is supported by the possibility of increasing the complexity of any modeled variable by using a mixture of densities [29] or normalizing flows [30,31,32]. In other words, we can increase the number of output units of any neural network to parametrize a variable with a highly complex distribution. Since the DGP of the outcome variable is linear, the linear benchmark is able to recover the treatment effects in the same way as the proposed approach.

4.4 Experiment IV: non-linear relation between all the variables in the system

In the fourth simulation study, data for both the confounding and the recovery value were generated as continuous variables. In this case, the effect of the confounding and treatment variable on the recovery is non-linear. In fact, this setting allows us to compute treatment effects conditioned on the confounding variable.

The confounding variable was simulated from a Log-normal distribution, KS ∼ Lognormal ( μ = 2.5 , σ = 0.25 ) . The treatment, as in the previous case, was modeled using a Bernoulli distribution with probability p = 1 1 + exp ( ( KS − μ ) / 10 ) . p was modeled using a transformation of KS in order to avoid extreme probabilities. Finally, the response variable was simulated from a normal distribution, R ∼ N 50 T KS + 3 , 1 .

We summarize the DGP in the following algorithm:

These definitions make treatment effects non-linear, as illustrated in Figure 2, in which the means and 90% confidence bands using the non-parametric bootstrap [33] are plotted for both the neural approach (left) and the linear approach (right). The results are not surprising for the linear model, as the linear model is only able to retrieve a pooled estimate of the treatment effects. On the other hand, the neural network model is able to estimate the individual treatment precisely.

Figure 2

Comparison between the linear and neural network models using the back-door adjustment formula. The solid blue line represents the mean of the bootstrap, while the bands around the mean represent the 90% confidence bands. A histogram of KS is represented in the background of the plots.

Nevertheless, Figure 2 indicates some weakness in our approach. As we move away from the support of the data, the estimated treatment effects become less precise. For values below KS = 8 , the estimated treatment effect deviates drastically from the analytical one. The ability to accurately estimate causal effects out of the support of the observed data using neural networks depends on the predictive power of the neural network in those areas of the data. In other words, the estimates of the out of distribution causal effects depend on the generalizability of the neural networks.

In case of a linear model, the ATE is constant. That is, regardless of the value of KS, a unique value for the causal effect is found. The problem with this estimation is that the non-linear relation between the confounding variable and the treatment is aggregated and thereby give rise to aggregation bias [34]. This bias can be of any order of magnitude. For example, in our simulation study, with R ∼ N ( R ∣ 4 ∗ T ∗ exp ( − 99 l ) + KS ) , the aggregated treatment effect would be close to 2 and underestimate the effect for small KS and overestimating it for large values.

4.5 Experiment V: the front door adjustment

In the previous simulation experiments, all the variables in the causal graph were observed. In fact, because of the simplicity of the graph, the causal effects collapsed to a conditional estimation as noted in Section 4.4. In this simulation study, the aim is to highlight that our approach works in conditions where there are unobserved confounders and, as a result, the back-door criterion does not hold. In Figure 3, a causal graph with those conditions is presented. The white background of the KS circle denotes an unobserved variable and, hence, equation (4) cannot be used.

Figure 3

Graphical model of an extended kidney stone example where the confounder KS is not observed.

In this case, KS was simulated from a standard normal distribution, and T was simulated from a normal distribution T ∼ N ( sin ( KS ) , 0.1 ) . Moreover, Mg, which is a mediator variable representing the amount of a certain chemical (in this case magnesium) in the treatment, was simulated from a normal distribution Mg ∼ N ( 1 + T 2 , 0.1 ) , while R was simulated from a normal distribution R ∼ N sin ( KS 2 ) + 5 Mg , 0.1 . The graphical model and the generative process are shown in Figure 3.

Given the conditions of the causal graph, the back-door criterion cannot be used, and we must rely on other methods to estimate the causal effect of T on R . The method to identify the causal effect in this graph is called the front-door adjustment [2]

4.5.2 Results

Since the analytical solution of this causal graph is not readily available, the true causal effects were simulated from binned conditional distributions using the true generative process. We compared the distributions using two values of T = 0 and T = 0.5 ( T is constrained between − 1 and 1). Furthermore, the integrals in equation (5) are approximated using Monte Carlo integration as explained in Section 4.3.

In Figure 4, we compare the estimated distribution of R for three models with the respective true distribution. On the left, the linear model finds a single distribution for the different values of T , as expected. In the middle, our distribution is seen to match the true simulated distribution. The discrepancies between the true distribution and the real distribution arise from the different elements of stochasticity, e.g., the sampling of the true effects, the stochastic optimization procedure, and the Monte Carlo approximation of the integral. In the plot to the right of Figure 4 it is observed that pure conditioning with respect to T and integrating over Mg from the conditional network P ( R ∣ X = x , Mg = m g ) render a distribution that is very different from the true distribution. We also included the Wasserstein Distance (WD) in every plot as a metric of comparison between the different models and the true distribution. The proposed approach performs better for the interventional distributions when compared to the other approaches.

Figure 4

Comparison of three different ways to estimate causal effects. (Left) Linear model. (Center) Our proposed approach. (Right) Simple conditioning. (Above) Comparing against the true do ( X = 0 ) , and (Below) against the true do ( X = 0.5 ) . The WD between the distributions is specified on each plot.

4.6 Experiment VI: unobserved confounder

In this section, we explore the robustness of our model when the assumption of knowing the causal graph is not satisfied. To examine this, three additional experiments were carried out. Two experiments in which the unobserved confounder has a linear relationship with the outcome variable with different “degrees” of confounding, and one experiment in which the relationship is non-linear. In all cases, it was assumed that there is a single unobserved confounder U , which affects both treatment and outcome variables. The factor multiplying this term in the “mild” case is 0.3 and 3 for the “strong” case. The true causal graph and DGP can be summarized as follows (Figure 5):

Figure 5

Causal graph for the unobserved confounder experiment.

The assumption of a known graph structure is violated by ignoring the existence of U . In other words, the causal effect is estimated using the same architecture as in Case IV, using the data from Algorithm 5. The results of the best performing models can be found in Figures 6 and 7.

Figure 6

Comparison between the linear and neural network models using the back-door adjustment formula in the mild unobserved confounding scenario. The solid blue line represents the mean of the bootstrap, while the bands around the mean represent the 90% confidence bands. A histogram of KS is represented in the background of the plots.

Figure 7

Comparison between the linear and neural network models using the back-door adjustment formula in the strong unobserved confounding scenario. The solid blue line represents the mean of the bootstrap, while the bands around the mean represent the 90% confidence bands. A histogram of KS is represented in the background of the plots.

Both figures show a deviation from the true treatment effect as expected. The results for both the neural and the linear model are relatively similar to those without unobserved confounding, in the sense that the neural model learns a non-linear treatment effect and the confidence bands are relatively tight. Likewise, for the mild confounding case, the linear model learns a linear treatment effect with big confidence bands. It is interesting, however, that in the strong confounding case, the confidence bands of the linear model are considerably tighter than the rest of the models. This, we believe, is due to the fact that the confounding captures most of the variance of the outcome variable and, consequently, the best fit linear model is the one closest to the value of the confounding.

The second type of model with which we tested the assumption of a known causal graph is a model where the relationship between the unobserved confounder and the outcome variable is non-linear. The DGP described in Algorithm 6 was used to simulate the data.

The estimated causal effects using this simulation are presented in Figure 8.

Figure 8

Comparison between the linear and neural network models using the back-door adjustment formula in the non-linear unobserved confounding scenario. The solid blue line represents the mean of the bootstrap, while the bands around the mean represent the 90% confidence bands. A histogram of KS is represented in the background of the plots.

As expected from all the simulations where the causal graph is not completely known, the estimated causal effects are far from the actual effects regardless of whether the estimator is linear or non-linear. The main lesson of this study cases is that, regardless of the estimator being used, if the true causal graph is not known, the estimators are not going to be particular strong predictors for the outcome.