Causal-StoNet: Causal Inference for High-Dimensional Complex Data

Abstract

With the advancement of data science, the collection of increasingly complex datasets has become commonplace. In such datasets, the data dimension can be extremely high, and the underlying data generation process can be unknown and highly nonlinear. As a result, the task of making causal inference with high-dimensional complex data has become a fundamental problem in many disciplines, such as medicine, econometrics, and social science. However, the existing methods for causal inference are frequently developed under the assumption that the data dimension is low or that the underlying data generation process is linear or approximately linear. To address these challenges, this paper proposes a novel causal inference approach for dealing with high-dimensional complex data. The proposed approach is based on deep learning techniques, including sparse deep learning theory and stochastic neural networks, that have been developed in recent literature. By using these techniques, the proposed approach can address both the high dimensionality and unknown data generation process in a coherent way. Furthermore, the proposed approach can also be used when missing values are present in the datasets. Extensive numerical studies indicate that the proposed approach outperforms existing ones.

1. Introduction

Causal inference is a fundamental problem in many disciplines such as medicine, econometrics and social science. The problem can be formulated under the potential outcomes framework by Rubin (1974). Let $\mathit{X}\in {ℝ}^p$ denote a vector of $p$-dimensional confounders. In this paper, we consider only the binary treatment $A\in \{0,1\}$, but discuss extensions to multiple-level treatments or continuous treatments later. For each subject at each treatment level $a$, we assume there exists a potential outcome $\mathit{Y}(a)$ that can be observed under the actual treatment. We are interested in estimating the average treatment effect (ATE) $\tau =\mathbb{E}[Y(1)-Y(0)]$. It is known that ATE is identifiable if all confounders that influence both treatment and outcome are observed and the ignorability and overlapping conditions (see Assumption A1) are satisfied. To estimate ATE, a variety of methods, such as outcome regression, augmented/inverse probability weighting (AIPW/IPW) and matching, have been developed. See Imbens (2004) and Rosenbaum (2002) for overviews. These methods often work under the assumptions:

Assumption 1.

(Outcome model) The parametric model ${\mu }_a\left(\mathit{X},{\mathit{\theta }}_a\right)$ is a correct specification for the outcome function ${\mu }_a(\mathit{X})=\mathbb{E}[\mathit{Y}(a)\mid \mathit{X}]$, $a\in \{0,1\}$; i.e., ${\mu }_a(\mathit{X})={\mu }_a\left(\mathit{X},{\mathit{\theta }}_a^{\ast }\right)$ and ${\mathit{\theta }}_a^{\ast }$ is the true model parameter.

Assumption 2.

(Treatment model) The parametric model $p\left(\mathit{X},{\mathit{\theta }}_s\right)$ is a correct specification for the propensity score $p(\mathit{X})=P(A=1\mid \mathit{X})$; i.e., $p(\mathit{X})=p\left(\mathit{X},{\mathit{\theta }}_s^{\ast }\right)$ and ${\mathit{\theta }}_s^{\ast }$ is the true model parameter.

Let ${\widehat{\mathit{\theta }}}_a$ and ${\widehat{\mathit{\theta }}}_s$ be consistent estimators of ${\mathit{\theta }}_a^{\ast }$ and ${\mathit{\theta }}_s^{\ast }$, respectively. For example, the AIPW estimator of the ATE is given by

$${\widehat{\tau }}_n=\frac{1}{n}\sum _{i=1}^n\left[\frac{A_i{\mathit{Y}}_i}{p\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_s\right)}-\frac{A_i-p\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_s\right)}{p\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_s\right)}{\mu }_1\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_1\right)-\frac{\left(1-A_i\right)Y_i}{1-p\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_s\right)}-\frac{A_i-p\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_s\right)}{1-p\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_s\right)}{\mu }_0\left({\mathit{X}}_i,{\widehat{\mathit{\theta }}}_0\right)\right],$$ (1)

which is doubly robust (Robins et al. (1994)) in the sense that ${\widehat{\tau }}_n$ is consistent if either Assumption 1 or 2 holds and locally efficient if both assumptions hold.

In practice, estimating the ATE often poses two main difficulties: (i) high dimensionality of covariates, which is common in genomic data (Bühlmann, 2013; Schwab et al., 2020), environmental and healthcare data (Antonelli et al., 2019), and social media data (Sharma et al., 2020); and (ii) unknown functional forms of the outcome and propensity score. Various methods have been proposed to address these challenges but in a separate manner. For example, Lasso (Tibshirani, 1996) and other regularization methods have been used to select relevant covariates for high-dimensional problems under the linear model framework (Belloni et al., 2014; Farrell, 2015). On the other hand, deep neural networks have been employed to estimate the outcome and propensity score functions for low-dimensional data (Shi et al., 2019; Farrell et al., 2021). However, none of these methods address both difficulties in a unified manner. Furthermore, missing values are often present in the datasets, which further complicates the causal inference problem (Yang et al., 2019; Guan & Yang, 2019).

Building on existing works on stochastic neural networks (Liang et al., 2022; Sun & Liang, 2022) and sparse deep learning (Liang et al., 2018b; Sun et al., 2022), we propose a Causal Stochastic Neural Network, which is abbreviated as Causal-StoNet in what follows, for addressing the above difficulties encountered in causal inference for high-dimensional complex data. The merits of Causal-StoNet are three-folds:

1. A natural forward-modeling framework. As described in Section 2, the StoNet has been formulated as a composition of multiple simple linear and logistic regressions, providing a natural forward-modeling framework for complex data generation processes. In Causal-StoNet, we replace a hidden neuron at an appropriate hidden layer by a visible treatment variable. With its compositional regression structure, Causal-StoNet easily extends to various causal inference scenarios, e.g. missing covariates, multi-level or continuous treatments, and mediation analysis, as discussed in Section 5.

2. Universal approximation ability. We prove that the StoNet possesses a valid approximation to a deep neural network, thereby enabling it to possess the universal approximation ability to the outcome and propensity score functions.

3. Consistent sparse learning. By imposing an appropriate sparse penalty/prior on the structure of the StoNet, relevant variables to the outcome and propensity score can be identified along with the training of the Causal-StoNet even under the setting of high-dimensional covariates. As a result, the outcome and propensity score can be properly estimated even when their exact functional forms are unknown.

In summary, the Causal-StoNet has successfully tackled the issues of high-dimensional covariates, unknown functional forms, and missing data in a holistic manner, providing a robust and reliable approach of causal inference for high-dimensional complex data.

Related Works

In the literature, there are quite a few works employing deep neural networks for causal inference, see e.g., Shi et al. (2019) and Farrell et al. (2021). However, the consistency of the deep neural network estimator is not established in Shi et al. (2019). This property has been studied in Farrell et al. (2021) but under the low-dimensional scenario essentially. Otherwise, it requires the underlying outcome and propensity score functions to be highly smooth with the smoothness degree even higher than the data dimension $p$. In addition, the methods in Shi et al. (2019) or Farrell et al. (2021) cannot perform covariate selection, and they are hard to apply when the dataset contains missing values. Quite recently, Chen et al. (2024) proposed some neural network-based ATE estimators, where only the propensity score or the outcome function is approximated using a neural network. There are also numerous semi-parametric casual estimation methods in the literature. Causal trees (Athey & Imbens, 2015; Li et al., 2015) develops data-driven approach to estimate heterogeneous treatment effect for subpopulations. Super learner (van der Laan et al., 2007) utilizes an ensemble of different models to enhance the causal effect estimation. Targeted maximum likelihood estimation (van der Laan & Rubin, 2006) proposes a flexible semi-parametric framework based on targeted regularization. These methods can also be combined with deep learning or other machine learning models, but the flexibility of the StoNet in forward modeling of complex data generation processes leads to the uniqueness of Causal-StoNet. It can function effectively in various data scenarios, as discussed in Section 5.

2. A Brief Review of Stochastic Neural Networks

The StoNet can be briefly described as follows. Consider a DNN model with $h$ hidden layers. For the sake of simplicity, we assume that the same activation function $\psi (\cdot )$ is used for all hidden units. By separating the feeding and activation operators of each hidden unit, we can rewrite the DNN model in the following form:

$$\begin{gathered} {\tilde{\mathit{Y}}}_1={\mathit{b}}_1+{\mathit{w}}_1\mathit{X}, \\ {\tilde{\mathit{Y}}}_i={\mathit{b}}_i+{\mathit{w}}_i\text{Ψ}\left({\tilde{\mathit{Y}}}_{i-1}\right),i=2,3,\dots ,h, \\ \mathit{Y}={\mathit{b}}_{h+1}+{\mathit{w}}_{h+1}\text{Ψ}\left({\tilde{\mathit{Y}}}_h\right)+{\mathit{e}}_{h+1}, \end{gathered}$$ (2)

where ${\mathit{e}}_{h+1}\sim N\left(0,{\sigma }_{h+1}^2{I}_{d_{h+1}}\right)$ is Gaussian random error; ${\tilde{\mathit{Y}}}_i$, ${\mathit{b}}_i\in {ℝ}^{d_i}$ for $i=1,2,\dots ,h$; $\mathit{Y}$, ${\mathit{b}}_{h+1}\in {ℝ}^{d_{h+1}}$; $\text{Ψ}\left({\tilde{\mathit{Y}}}_{i-1}\right)={\left(\psi \left({\tilde{Y}}_{i-1,1}\right),\psi \left({\tilde{Y}}_{i-1,2}\right),\dots ,\psi \left({\tilde{Y}}_{i-1,d_{i-1}}\right)\right)}^T$ for $i=2,3,\dots ,h+1$, and ${\tilde{Y}}_{i-1,j}$ is the $j\text{th}$ element of ${\tilde{\mathit{Y}}}_{i-1}$; ${\mathit{w}}_i\in {ℝ}^{d_i\times d_{i-1}}$ for $i=1,2,\dots ,h+1$, and $d_0=p$ denotes the dimension of $\mathit{X}$. For simplicity, we consider only the regression problems in (2). By replacing the third equation in (2) with a logit model, the DNN model can be extended to classification problems.

The StoNet is a probabilistic deep learning model and constructed by adding auxiliary noise to ${\tilde{\mathit{Y}}}_i$’s for $i=1,2,\dots ,h$ in (2). Mathematically, the StoNet model is given by

$$\begin{array}{l}{\mathit{Y}}_1={\mathit{b}}_1+{\mathit{w}}_1\mathit{X}+{\mathit{e}}_1,\\ {\mathit{Y}}_i={\mathit{b}}_i+{\mathit{w}}_i\text{Ψ}\left({\mathit{Y}}_{i-1}\right)+{\mathit{e}}_i,i=2,3,\dots ,h,\\ \mathit{Y}={\mathit{b}}_{h+1}+{\mathit{w}}_{h+1}\text{Ψ}\left({\mathit{Y}}_h\right)+{\mathit{e}}_{h+1},\end{array}$$ (3)

as a composition of many simple regressions, where ${\mathit{Y}}_1,{\mathit{Y}}_2,\dots ,{\mathit{Y}}_h$ can be viewed as latent variables. Further, we assume that ${\mathit{e}}_i\sim N\left(0,{\sigma }_i^2{I}_{d_i}\right)$ for $i=1,2,\dots ,h+1$. For classification problems, ${\sigma }_{h+1}^2$ plays the role of temperature for the binomial or multinomial distribution formed at the output layer, and it works with $\left\{{\sigma }_1^2,\dots ,{\sigma }_h^2\right\}$ together to control the variation of the latent variables $\left\{{\mathit{Y}}_1,\dots ,{\mathit{Y}}_h\right\}$.

It has been shown in Liang et al. (2022) that the StoNet is a valid approximator to the DNN, i.e., asymptotically they have the same loss function as the training sample size $n$ becomes large. Let ${\mathit{\theta }}_i=\left({\mathit{w}}_i,{\mathit{b}}_i\right)$, let $\mathit{\theta }=\left({\mathit{\theta }}_1,{\mathit{\theta }}_2,\cdots ,{\mathit{\theta }}_{h+1}\right)$ denote the parameter vector of the StoNet, let $d_{\mathit{\theta }}$ denote the dimension of $\mathit{\theta }$, and let $\text{Θ}$ denote the space of $\mathit{\theta }$. Let $L:\text{Θ}\to ℝ$ denote the loss function of the DNN as defined in (2), which is given by

$$L(\mathit{\theta })=-\frac{1}{n}\sum _{i=1}^n\log \pi \left({\mathit{Y}}^{(i)}\mid {\mathit{X}}^{(i)},\mathit{\theta }\right),$$ (4)

where $i$ indexes the training samples. Under appropriate settings for ${\sigma }_i$’s, the activation function $\psi$, and the parameter space $\text{Θ}$, see Assumption A2 (in Appendix), Liang et al. (2022) showed that the StoNet and the DNN have asymptotically the same training loss function, i.e.,

$$\underset{\mathit{\theta }\in \text{Θ}}{\sup }\left|\frac{1}{n}\sum _{i=1}^n\log \pi \left({\mathit{Y}}^{(i)},{\mathit{Y}}_{mis}^{(i)}\mid {\mathit{X}}^{(i)},\mathit{\theta }\right)-\frac{1}{n}\sum _{i=1}^n\log \pi \left({\mathit{Y}}^{(i)}\mid {\mathit{X}}^{(i)},\mathit{\theta }\right)\right|\stackrel{p\text{}}{\to }0,\mathit{\text{as}}n\to \infty ,$$ (5)

where ${\mathit{Y}}_{mis}=\left({\mathit{Y}}_1,{\mathit{Y}}_2,\dots ,{\mathit{Y}}_h\right)$ denotes the collection of all latent variables as defined in (3). The StoNet can work with a wide range of Lipschitz continuous activation functions such as tanh, sigmoid and ReLU. As explained in Liang et al. (2022), Assumption A2 also restricts the size of the noise added to each hidden unit by setting: ${\sigma }_1\le {\sigma }_2\le \cdots \le {\sigma }_{h+1}$, ${\sigma }_{h+1}=O(1)$, and $d_{h+1}\left({\prod }_{i=k+1}^h{d}_i^2\right)d_k{\sigma }_k^2\prec \frac{1}{h}$ for any $k\in \{1,2,\dots ,h\}$, where the factor $d_{h+1}\left({\prod }_{i=k+1}^h{d}_i^2\right)d_k$ can be understood as the amplification factor of the noise ${\mathit{e}}_k$ at the output layer. In general, the noise added to the first few hidden layers should be small to prevent large random errors propagated to the output layer.

Further, it is assumed that each $\mathit{\theta }$ for the DNN is unique up to some loss-invariant transformations, such as reordering some hidden units and simultaneously changing the signs of some weights and biases, see Liang et al. (2018b) and Sun et al. (2022) for similar assumptions used in the study. Then, under some regularity assumptions for the population negative energy function $Q^{\ast }(\mathit{\theta })=\mathbb{E}(\log \pi (\mathit{Y}\mid \mathit{X},\mathit{\theta }))$, see Assumption A3, Liang et al. (2022) showed

$$‖{\widehat{\mathit{\theta }}}_n-{\mathit{\theta }}^{\ast }‖\stackrel{p\text{}}{\to }0,\text{as}n\to \infty ,$$ (6)

where ${\widehat{\mathit{\theta }}}_n=\arg {\max }_{\mathit{\theta }\in \text{Θ}}\left\{\frac{1}{n}{\sum }_{i=1}^n\log \pi \left({\mathit{Y}}^{(i)},{\mathit{Y}}_{mis}^{(i)}\mid {\mathit{X}}^{(i)},\mathit{\theta }\right)\right\}$, and ${\theta }^{\ast }=\arg {\max }_{\mathit{\theta }\in \text{Θ}}{Q}^{\ast }(\mathit{\theta })$.

That is, the DNN (2) can be trained by training the StoNet (3); they are asymptotically equivalent as $n\to \infty$, thereby the universal approximation property also holds for the StoNet. It is worth noting that in forward prediction, the StoNet ignores auxiliary noise added to hidden neurons and thus performs as the DNN.

3. Causal-StoNet

3.1. The Structure

The StoNet provides a unified solution for the challenges faced in causal inference for high-dimensional complex data. In this section, we address the challenges of high-dimensional covariates and unknown functional forms of the outcome and propensity score, leaving the treatment of missing data to Section 5.

Figure 1 illustrates the structure of the Causal-StoNet, where the treatment variable is encompassed as a visible unit in an intermediate hidden layer. The compositional regression architecture of the StoNet ensures seamless computational handling of this setup without introducing any computational complexities. Let $A$ denote the treatment variable. The Causal-StoNet is to learn a decomposition of the joint distribution

$$\pi \left(\mathit{Y},{\mathit{Y}}_{mis},A\mid \mathit{X},\mathit{\theta }\right)\propto \pi \left({\mathit{Y}}_1\mid \mathit{X},{\mathit{\theta }}_1\right)\pi \left({\mathit{Y}}_2\mid {\mathit{Y}}_1,{\mathit{\theta }}_2\right)\pi \left(A\mid {\mathit{Y}}_1,{\mathit{\theta }}_2\right)\pi \left({\mathit{Y}}_3\mid {\mathit{Y}}_2,A,{\mathit{\theta }}_3\right)\pi \left(\mathit{Y}\mid {\mathit{Y}}_3,{\mathit{\theta }}_4\right),$$ (7)

where ${\mathit{Y}}_{mis}=\left({\mathit{Y}}_1,{\mathit{Y}}_2,{\mathit{Y}}_3\right)$, $\mathit{\theta }=\left({\mathit{\theta }}_1,{\mathit{\theta }}_2,{\mathit{\theta }}_3,{\mathit{\theta }}_4\right)$, and $\pi \left(A\mid {\mathit{Y}}_1,{\mathit{\theta }}_2\right)$ corresponds to the propensity score. For the visible binary treatment unit, a sigmoid activation function can be used for its probability interpretation.

Figure 1:

Causal-StoNet Structure: the treatment is included as a visible unit (rectangle) in a middle layer, and ${\mathit{Y}}_2$ denotes the latent variable of that layer but with the unit directly feeding to the treatment rectangle excluded; ‘x’ represents possible missing values.

To ensure that a sparse Causal-StoNet can be learned for high-dimensional data, where the number of covariates $p$ can be much larger than the sample size $n$, we will follow Sun et al. (2022; 2021) to impose a mixture Gaussian prior on each component of $\mathit{\theta }$, i.e.,

$$\theta \sim {\lambda }_{n}N\left(0,{\sigma }_{1,n}^2\right)+\left(1-{\lambda }_n\right)N\left(0,{\sigma }_{0,n}^2\right),$$ (8)

where $\theta$ denotes a generic weight and bias of the Causal-StoNet, ${\lambda }_n\in (0,1)$ is the mixture proportion, ${\sigma }_{0,n}^2$ is typically set to a very small number, while ${\sigma }_{1,n}^2$ is relatively large.

Let ${\mu }^{*}(\mathit{x},A)$ denote the true outcome function, and let $p^{*}(\mathit{x})$ denote the true propensity score function. Let $\widehat{\mu }\left(\mathit{x},A;{\widehat{\mathit{\theta }}}_n\right)$ denote the DNN estimator of ${\mu }^{*}(\mathit{x},A)$, and let $\widehat{p}\left(\mathit{x};{\widehat{\mathit{\theta }}}_n\right)$ denote the DNN estimator of $p^{*}(\mathit{x})$. For a given estimator ${\widehat{\mathit{\theta }}}_n$, both $\widehat{\mu }\left(\mathit{x},A;{\widehat{\mathit{\theta }}}_n\right)$ and $\widehat{p}\left(\mathit{x};{\widehat{\mathit{\theta }}}_n\right)$ are calculated as for the conventional DNN model (2) by ignoring the random errors ${\mathit{e}}_i$’s. Let ${\mathit{\gamma }}^{*}=\left\{{\gamma }_i:i=1,2,\dots ,d_{\mathit{\theta }}\right\}$ denote the true sparse structure of the Causal-StoNet, which is defined through a sparse DNN as in (A5). Here ${\gamma }_i$ is an indicator for the existence of connection $c_i$. Following Sun et al. (2022), for each $i\in \left\{1,2,\dots ,d_{\mathit{\theta }}\right\}$, we set ${\widehat{\gamma }}_i=1$ if the corresponding weight $\left|{\widehat{\mathit{\theta }}}_i\right|>\frac{\sqrt{2}{\sigma }_{0,n}{\sigma }_{1,n}}{\sqrt{{\sigma }_{1,n}^2-{\sigma }_{0,n}^2}}\sqrt{\log \left(\frac{1-{\lambda }_n}{{\lambda }_n}\frac{{\sigma }_{1,n}}{{\sigma }_{0,n}}\right)}$ and 0 otherwise. Denote the estimated sparse Causal-StoNet structure $\widehat{\mathit{\gamma }}\left({\widehat{\mathit{\theta }}}_n\right)=\left\{{\widehat{\gamma }}_i:i=1,2,\dots ,d_{\mathit{\theta }}\right\}$. Under appropriate conditions, we can show that the sparse Causal-StoNet leads to consistent estimates for ${\mu }^{\ast }(\mathit{x},A)$, $p^{\ast }(\mathit{x})$, and ${\mathit{\gamma }}^{\ast }$. This can be summarized as the following theorem, whose proof is given in the appendix.

Theorem 1.

Assume that the mixture Gaussian prior (8) is imposed on each connection of the StoNet, Assumptions A2–A5 hold, and $r_n\prec n^{3/16}$. As $n\to \infty$, the following results hold:

1. (Propensity score function) With probability greater than $1-\exp \left\{cn{ϵ}_n^2\right\}$ for some constant $c$,

   $${\mathbb{E}}_{\mathit{x}}\left[{\left(\widehat{p}\left(\mathit{x};{\widehat{\mathit{\theta }}}_n\right)-p^{\ast }(\mathit{x})\right)}^2\right]=O\left({ϵ}_n^2+e^{-cn{ϵ}_n^2/16}\right)+o\left(n^{-1/2}\right).$$
2. (Outcome function) If ${\mu }^{\ast }(\mathit{x})$ is bounded and the activation function $\psi (\cdot )\in [-1,1]$, then, with probability greater than $1-\exp \left\{cn{ϵ}_n^2\right\}$ for some constant $c$,

   $${\mathbb{E}}_{\mathit{x}}\left({\left|\widehat{\mu }\left(\mathit{x},A;{\widehat{\mathit{\theta }}}_n\right)-{\mu }^{\ast }(\mathit{x})\right|}^2\right)=O\left(\left({ϵ}_n^2+e^{-cn{ϵ}_n^2}\right){\overline{L}}_n^2\right)+o\left(n^{-1/2}\right).$$

3. (Structure selection) If Assumption A6 also holds, then $P\left(\widehat{\mathit{\gamma }}\left({\widehat{\mathit{\theta }}}_n\right)={\mathit{\gamma }}_{\ast }\right)\stackrel{p\text{}}{\to }1$.

By Theorem 1, the sparse Causal-StoNet provides consistent estimators for both the propensity score and outcome functions. Therefore, these estimators can be plugged into equation (1) to get a double robust estimator for ATE. Moreover, the sparse Causal-StoNet also provides consistent identification for the covariates relevant to the treatment and outcome variables, which ensures that the covariates selected for the propensity score function are contained in those selected for the outcome function.

Regarding theoretical properties of the ATE estimator, we have Theorem 2 by following the theory established in Farrell (2015).

Theorem 2.

Suppose Assumptions A1–A5 hold. Additionally, assume that the mixture Gaussian prior (8) is imposed on each connection of the StoNet, $r_n\prec n^{3/16}$, and $n^{-1+\xi }\prec {\varpi }_n^2\prec n^{-\frac{1}{2}-\xi }$, and specify the network structure such that $0.5+\xi <\epsilon <1-\xi$ and ${\overline{L}}_n=O\left(n^{\xi }\right)$ for some $0\le \xi <1/4$. Then the following results hold:

1. $V_{\tau }^{-1/2}\sqrt{n}\left({\widehat{\tau }}_n-{\tau }^{\ast }\right)\stackrel{d\text{}}{\to }\mathcal{N}(0,1)$, where $V_{\tau }$ is given in Supplement A.5, and ${\tau }^{\ast }$ denotes the true value of the ATE.

2. ${\widehat{V}}_{\tau }-V_{\tau }\stackrel{p\text{}}{\to }1$, where the estimator ${\widehat{V}}_{\tau }$ is given in Supplement A.5.

3. (Uniformaly valid inference) Let ${\mathcal{P}}_n$ be a set of data-generating process satisfying Assumption A1. Then for $c_{\alpha }={\text{Φ}}^{-1}(1-\alpha /2)$, we have

   $$\underset{P_n\in {\mathcal{P}}_n}{\sup }\left|{\mathbb{P}}_{P_n}\left[{\tau }^{\ast }\in \left\{{\widehat{\tau }}_n\pm c_{\alpha }\sqrt{{\widehat{V}}_{\tau }/n}\right\}\right]-(1-\alpha )\right|\to 0.$$

We note that by the theory of the StoNet, we can also estimate the propensity score and outcome functions separately by running two sparse DNNs in the way of double machine learning (Chernozhukov et al., 2018). However, in this double machine learning implementation, the covariates selected for the propensity score function might not be a subset of those selected for the outcome function, leading to ambiguity in interpretation for the role that certain covariates play in the causal system. The Causal-StoNet avoids this issue by jointly estimating the propensity score and outcome functions.

3.2. Adaptive Stochastic Gradient MCMC for Training Causal-StoNet

As implied by Theorem 1, training the Causal-Stonet can be boiled down to solving a high-dimensional parameter estimation problem with latent variables present, i.e., maximizing

$$\sum _{i=1}^n\log \pi \left({\mathit{Y}}^{(i)},{\mathit{Y}}_{mis}^{(i)},A\mid {\mathit{X}}^{(i)},\mathit{\theta }\right)+\log \pi (\mathit{\theta }).$$ (9)

To maximize (9), a feasible method is adaptive stochastic gradient Markov chain Monte Carlo (SGMCMC), which, by the Bayesian version of Fisher’s identity (Song et al., 2020), converts the optimization problem to a mean-field equation solving problem:

$$h(\mathit{\theta }):=\int H\left({\mathit{Y}}_{mis},\mathit{\theta }\right)d\pi \left({\mathit{Y}}_{mis}\mid \mathit{X},\mathit{Y},A,\mathit{\theta }\right)=0,$$ (10)

where $H\left({\mathit{Y}}_{mis},\mathit{\theta }\right)={\nabla }_{\mathit{\theta }}\log \pi \left(\mathit{Y},{\mathit{Y}}_{mis},A\mid \mathit{X},\mathit{\theta }\right)+{\nabla }_{\mathit{\theta }}\log \pi (\mathit{\theta })$. The adaptive SGMCMC algorithm works under the framework of stochastic approximation MCMC (Benveniste et al., 1990; Liang et al., 2007). It can be briefly described as follows.

For simplicity of notation, we rewrite (10) in the following equation:

$$h(\mathit{\theta })=\mathbb{E}[H(\mathit{Z},\mathit{\theta })]=\int H(z,\mathit{\theta })\pi (z\mid \mathit{\theta })dz=0,$$ (11)

where $\mathit{Z}$ is a latent variable and $\pi (z\mid \mathit{\theta })$ is a probability density function parameterized by $\mathit{\theta }\in \text{Θ}$. The algorithm works by iterating between the following two steps:

1. (Sampling) Simulate $z^{(k+1)}\sim \pi \left(z\mid {\mathit{\theta }}^{(k)}\right)$ via a transition kernel induced by a SGMCMC algorithm such as stochastic gradient Langevin dynamics (Welling & Teh, 2011) and stochastic gradient Hamilton Monte Carlo (SGHMC) (Chen et al., 2014).

2. (Parameter updating) Set ${\mathit{\theta }}^{(k+1)}={\mathit{\theta }}^{(k)}+{\gamma }_{k+1}H\left(z^{(k+1)},{\mathit{\theta }}^{(k)}\right)$, where ${\gamma }_{k+1}$ denotes the step size used in the stochastic approximation procedure.

This algorithm is called adaptive SGMCMC as the transition kernel used in step (a) changes along with the working estimate ${\mathit{\theta }}^{(k)}$. Applying the adaptive SGHMC algorithm to (10) leads to Algorithm 1 (in Appendix A.1), where SGHMC is used for simulation of the latent variables ${\mathit{Y}}_{mis}$ at each iteration. The convergence of the adaptive SGHMC algorithm has been studied in Liang et al. (2022).

Lemma 1.

(Liang et al. (2022)) Suppose Assumptions A8-A13 hold. In Algorithm 1, if we set ${ϵ}_{k,i}=C_{ϵ}/\left(c_e+k^{\alpha }\right)$ and ${\gamma }_{k,i}=C_{\gamma }/\left(c_g+k^{\alpha }\right)$ for some constants $\alpha \in (0,1)$, $C_{ϵ}>0$, $C_{\gamma }>0$, $c_e\ge 0$ and $c_g\ge 0$, then there exists an iteration $k_0$ and a constant ${\lambda }_0>0$ such that for $k>k_0$,

$$\mathbb{E}\left({‖{\widehat{\mathit{\theta }}}^{(k)}-{\widehat{\mathit{\theta }}}_n^{\ast }‖}^2\right)\le {\lambda }_0{\gamma }_k,$$ (12)

where ${\widehat{\mathit{\theta }}}_n^{\ast }$ denotes a solution to equation (10).

In Liang et al. (2022), an explicit expression of ${\lambda }_0$ has been given. For simplicity, we have the expression omitted in this paper. Next, Liang et al. (2022) showed that as $k\to \infty$, the imputed latent variable $z^{(k)}$ converges weakly to the desired posterior distribution $\pi \left(z\mid {\mathit{\theta }}^{\ast }\right)$ in 2-Wasserstein distance. Similarly, we establish Lemma 2 which can be used in statistical inference for the problems with missing data being involved.

Lemma 2.

Suppose Assumptions A8-A13 hold. Then for any bounded function ${\varphi }_{{\mathit{\theta }}^{\ast }}(\cdot )$, ${\mathbb{E}}_{\mathcal{K}}\widehat{{\varphi }_{{\widehat{\mathit{\theta }}}^{\ast }}}(z)-\int {\varphi }_{{\widehat{\mathit{\theta }}}^{\ast }}(z)d\pi \left(z\mid {\widehat{\mathit{\theta }}}^{\ast }\right)\stackrel{p\text{}}{\to }0$ as $\mathcal{K}\to \infty$, where $\widehat{{\mathbb{E}}_{\mathcal{K}}{\varphi }_{{\widehat{\mathit{\theta }}}^{\ast }}}(z)=\frac{1}{\mathcal{K}}{\sum }_{i=1}^{\mathcal{K}}{\varphi }_{{\widehat{\mathit{\theta }}}^{(i)}}\left(z^{(i)}\right)$ and $\left\{\left({\widehat{\mathit{\theta }}}^{(i)},z^{(i)}\right):i=1,2,\dots ,\mathcal{K}\right\}$ denotes a set of parameter estimates and imputed latent variables that are collected in a run of the SGHMC algorithm.

4. Numerical Examples

4.1. Setup

Baselines.

We compare Causal-StoNet with the following baselines: (1) Designed for average treatment effect: double selection estimator (DSE)(Belloni et al., 2014), approximate residual balancing estimator (ARBE) (Athey et al., 2018), targeted maximum likelihood estimator (TMLE) (van der Laan & Rubin, 2006), and deep orthogonal networks for unconfounded treatments (DONUT) (Hatt & Feuerriegel, 2021); (2) Designed for heterogeneous treatment effect: X-learner (Künzel et al., 2017), Dragonnet(Shi et al., 2019), causal multi-task deep ensemble (CMDE) (Jiang et al., 2023)), causal effect variational autoencoder (CEVAE) (Louizos et al., 2017), generative adversarial networks (GANITE) (Yoon et al., 2018), and counterfactual regression net (CFRNet) ¹.

Performance metrics.

We consider these metrics: (a) estimation accuracy of ATE, which is measured by the mean absolute error (MAE) of the ATE estimates; (ii) estimation accuracy of CATE, which is measured by precision in estimation of heterogeneous effect (PEHE); and (iii) covariate selection accuracy for the treatment and outcome models, which is measured by false and negative selection rates (FSR and NSR) as defined in Section A.7.2.

4.2. Simulation with Varying Sample Size

We ran simulated experiment with covariate dimension p = 1000 and training sample size n = 800,1600,2400,3200,4000, respectively. For each scenario, 10 simulated datasets are generated as described in Section A.7.1. Both the outcome and treatment effect functions in this experiment are nonlinear. As DSE and ARBE are formulated under linear assumptions, we didn’t include them as baselines to ensure a fair comparison. In each experiment, Algorithm 1 was executed 10 times, and the best model was selected based on BIC, as suggested by Sun et al. (2022). This setting will be default for all the experiments unless otherwise stated.

The results are depicted in Figure 2, demonstrating that Causal-StoNet maintains stable performance even with high-dimensional covariates and small sample sizes. Additionally, we conducted further simulations to investigate covariate selection accuracy and address missing value problems, as detailed in Section A.7.2.

Figure 2:

In-sample MAE and Out-of-Sample MAE of ATE estimation with varying training sample sizes. In-sample MAE is calculated over training and validation sets, Out-of-Sample MAE is calculated over test set

4.3. Atlantic Causal Inference Conference 2019 Data Challenge

The Causal-StoNet is compared with baseline methods on 10 synthetic datasets with homogeneous treatment effect from the Atlantic Causal Inference Conference (ACIC) 2019 Data Challenge. Each dataset contains 200 covariates with binary treatment variable, and the outcome variable is continuous. Since they both are synthetic, the true ATE is known. Results in Table 1 demonstrate that Causal-StoNet consistently provides more accurate estimates than the competitive methods.

Table 1:

ATE estimation across 10 ACIC 2019 datasets, where the number in the parentheses is the standard deviation of the MAE.

Method	In-Sample	Out-of-Sample
Causal-StoNet	0.0501(0.0118)	0.0542(0.0132)
DSE	0.0776(0.0193)	0.1632(0.0251)
ARBE	0.0729(0.0166)	0.1335(0.0179)
TMLE(Lasso)	0.0869(0.0164)	0.0867(0.0165)
TMLE(ensemble)	0.1140(0.0394)	0.1316(0.0429)
DONUT	0.5294 (0.2640)	0.5290(0.2642)

In addition to ATE, we also consider the conditional average treatment effect (CATE), which measures the heterogeneous treatment effect for subpopulations or individuals based on their covariates $\mathit{x}\in {ℝ}^p$ and is defined by

$$\tau (\mathit{x})=\mathbb{E}[\mathit{Y}(1)-\mathit{Y}(0)\mid \mathit{X}=\mathit{x}].$$

We evaluated Causal-StoNet’s performance in CATE estimation on an ACIC 2019 dataset with heterogeneous treatment effect, where both the treatment and the outcome are binary, using two metrics: ${ϵ}_{\text{PEHE}}$ (square root of the Precision in Estimation of Heterogeneous Effect (PEHE)) and ${ϵ}_{\text{ATE}}$ (absolute error of estimated ATE). The results in Table 2 show that while Causal-StoNet slightly lags behind CMDE, CEVAE, and X-Learner-BART in ${ϵ}_{\text{PEHE}}$, it achieves the lowest ${ϵ}_{\text{ATE}}$.

Table 2:

CATE estimation for an ACIC 2019 dataset.

Method	${ϵ}_{\text{PEHE}}$	${ϵ}_{\text{ATE}}$
Causal-StoNet	0.0893	0.0118
CMDE	0.0823	0.0444
CMGP	0.2156	0.0258
CEVAE	0.0867	0.0358
GANITE	0.1913	0.0485
X-Learner-RF	0.1877	0.0203
X-Learner-BART	0.0873	0.0720
CFRNet-Wass	0.1182	0.0421
CFRNet-MMD	0.1158	0.0849

4.4. Twins Data

We analyzed a dataset of twin births from 1989 to 1991 in the United States. The treatment variable is binary, with $a=1$ denoting the heavier twin at birth; and the outcome variable is binary, with $Y=1$ indicating twin mortality within the first year. We regard each twin-pair’s records as potential outcomes, allowing us to find the true ATE. The dataset includes 46 covariates. Refer to A.7.3 for data preprocessing steps. After data pre-processing, we obtained a dataset with 4,821 samples. In this final dataset, mortality rates for lighter and heavier twins are 16.9% and 14.42%, respectively, resulting in a true ATE of −2.48%. We conducted the experiment in three-fold cross validation, where we partitioned the dataset into three subsets, trained the model using two subsets and estimated the ATE using the remaining one. Table 3 reports the averaged ATE over three folds and the standard deviation of the average. Causal-StoNet yields a more stable ATE estimate than baseline methods.

Table 3:

ATE estimates by different methods for twins data

Causal-StoNet	DSE	ARBE	TMLE(Lasso)	TMLE(ensemble)	DONUT
−0.0232(0.0042)	−0.0405(0.0176)	−0.0096 (0.0201)	−0.1103(0.0599)	−0.1290(0.0779)	−0.0738(0.0128)

Table A2 shows the covariates selected by Causal-StoNet for the propensity score and outcome models in the three-fold cross-validation experiments. As expected, some covariates that are known to be relevant to the outcome, such as gestat10, have been selected for both the treatment and outcome models.

5. Some Variants of Causal-StoNet

The proposed Causal-StoNet can be easily extended to various scenarios of causal inference, such as covariates with missing values, multi-level or continuous treatments, and the presence of mediation variables. The extensions can be briefly described as follows.

Missing at Random (MAR)

Let ${\mathit{X}}_{obs}$ denote the observed covariates, let ${\mathit{X}}_{mis}$ denote the missed covariate values, and let $R$ denote the missing pattern represented as a binary vector. Under the mechanism of missingness at random, i.e., ${\mathit{X}}_{mis}⫫R\mid \left({\mathit{X}}_{obs},A,\mathit{Y}\right)$, the Causal-StoNet as depicted in Figure 1 is to learn a decomposition of the joint distribution

$$\begin{array}{l}\pi \left(\mathit{Y},{\mathit{Y}}_{mis},{\mathit{X}}_{mis},A\mid {\mathit{X}}_{obs},R,\mathit{\theta }\right)\propto \pi \left({\mathit{X}}_{mis}\mid {\mathit{X}}_{obs}\right)\pi \left({\mathit{Y}}_1\mid {\mathit{X}}_{obs},{\mathit{X}}_{mis},{\theta }_1\right)\pi \left({\mathit{Y}}_2\mid {\mathit{Y}}_1,{\theta }_2\right)\\ \times \pi \left(A\mid {\mathit{Y}}_1,{\theta }_2\right)\pi \left({\mathit{Y}}_3\mid {\mathit{Y}}_2,A,{\theta }_3\right)\pi \left(\mathit{Y}\mid {\mathit{Y}}_3,{\theta }_4\right),\end{array}$$ (13)

where ${\mathit{Y}}_{mis}=\left({\mathit{Y}}_1,{\mathit{Y}}_2,{\mathit{Y}}_3\right)$, $\mathit{\theta }=\left({\theta }_1,{\theta }_2,{\theta }_3,{\theta }_4\right)$, $\pi \left(A\mid {\mathit{Y}}_1,{\theta }_2\right)$ corresponds to the propensity score, and $\pi \left({\mathit{X}}_{mis}\mid {\mathit{X}}_{obs}\right)$ can be formulated in graphical models (see e.g. Liang et al. (2018a)) and will not be detailed here. It is easy to see that in this scenario, the Causal-StoNet can still be trained using Algorithm 1 by treating ${\mathit{X}}_{mis}$ as part of the latent variables. Statistical inference with imputed missing data can then be made based on Lemma 2.

Missing not at Random (MNAR)

The Causal-StoNet can also be extended to the scenario of MNAR, where the missing pattern depends on the missing values themselves even after controlling for observed data. To make the full data distribution identifiable, following Yang et al. (2019), we will assume that the missing pattern $R$ is independent of the outcome given the treatment and confounders, i.e., $\mathit{Y}⫫R\mid \left(A,{\mathit{X}}_{obs},{\mathit{X}}_{mis}\right)$. Under this assumption,

$$\begin{array}{c}\pi \left(\mathit{Y},{\mathit{Y}}_{mis},{\mathit{X}}_{mis},A\mid {\mathit{X}}_{obs},R,\mathit{\theta }\right)\propto \pi \left({\mathit{X}}_{mis}\mid {\mathit{X}}_{obs}\right)\pi \left(A\mid {\mathit{X}}_{obs},{\mathit{X}}_{mis},\mathit{\theta }\right)\pi \left(R\mid {\mathit{X}}_{obs},{\mathit{X}}_{mis},A,\mathit{\theta }\right)\\ \times \pi \left(\mathit{Y}\mid {\mathit{X}}_{obs},{\mathit{X}}_{mis},A,\mathit{\theta }\right).\end{array}$$

To accommodate the term $\pi \left(R\mid {\mathit{X}}_{obs},{\mathit{X}}_{mis},A,\mathit{\theta }\right)$ in the decomposition, we can include some extra visible units for $R$ at some layer between the treatment layer and the output layer. Note that the $R$ units will not be forwardly connected to the output layer.

Multilevel or Continuous Treatment Variables

The extension of the Causal-StoNet to this scenario is straightforward. For continuous treatment variable, the Causal-StoNet as depicted in Figure 1 can be directly applied with an appropriate modification of the activation function for the treatment neuron. For multilevel treatment variable, we can simply include multiple visible treatment neurons in the sample hidden layer, with a softmax activation function being used for them.

Causal Mediation Analysis

In this scenario, we aim to measure how the treatment effect is affected by intermediate/mediation variables. For example, Pearl (2001) gave an example where the side effect of a drug may cause patients to take aspirin, and the latter has a separate effect on the disease that the drug was originally prescribed for. The mediation analysis can be easily conducted with the Causal-StoNet by including some extra visible units for mediation variables at some layer between the treatment layer and the output layer. The mediation units was fed by the treatment unit and other hidden units of the same layer, and then feeds forward to cast its effect on the outcome layer.

6. Conclusion

We have developed an effective method for causal inference with high-dimensional complex data, which addresses the difficulties, including high-dimensional covariates, unknown treatment and outcome functional forms, and missing data, that are frequently encountered in the practice of modern data science. The proposed method does not only possess attractive theoretical properties, but also numerically outperforms the existing methods as demonstrated by our extensive examples.

The Causal-StoNet introduces an innovative deep neural network structure, incorporating visible neurons in its middle layers. Its stochastic deep learning nature renders Causal-StoNet essentially a universal tool for causal inference. It can model complex data generation processes in a forward manner, consistently identify relevant features, and provide accurate approximation to the underlying functions. Furthermore, the flexibility of adaptive SGMCMC algorithms, which impute latent variables (and handle missing data) while consistently estimating model parameters, greatly facilitates the computation of Causal-StoNet.