Causal-StoNet: Causal Inference for High-Dimensional Complex Data
Authors: Yaxin Fang
Faming Liang
Venue: ICLR 2024 poster (2024)
Keywords: Causal Information
Stochastic Neural Network
Adaptive Stochastic Gradient MCMC
Sparsity
Missing Data
Primary area: causal reasoning
TL;DR: We propose an innovative stochastic deep learning approach for conducting causal inference with 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 stochastic deep learning approach for conducting causal inference with high-dimensional complex data. The proposed approach is based on some 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.

CAUSAI-STONET: CAUSAL INFERENCE FOR HIGH- NIGH-DIMENSIONAL COMPLEX DATA
Yaxin Fang
Faming Liang
Department of Statistics Purdue University
Department of Statistics
Purdue University
West Lafayette, IN 47907, USA
West Lafayette, IN 47907, USA
fang230@purdue. edu
fml ialiaai imlianguuudue... edu
ABSTRACT
With the advancement data science, the collection of increasingly complex datasets has become commonplace. In such datasets, the data dimension can be extremely high, and the undorlying data gencration process can be unknown and highly nonlinear. As a result, the task of making causal inference with high-dimensional complex data has become a fundamenta 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 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 coherent way. Furthermore, the proposed approach can also be used when missing values present in the datasets. Extensive numerical studies indicate that the proposed approach outperforms existing ones.
INTRODUCTION
Causal inference is Jundamental problem many disciplines such as medicine, econometrics and social science. The problem can be formulated under the potential outcomes framework by Rubin (1974). Let X E R2 denote a vector of p-dimensional confounders. In this paper, consider only the binary treatment A {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 ((a) that can be observed under the actual treatment. We are interested in estimating the average treatment effect (ATE) T = E[ (1) Y(O)]- It is known that ATE is identifiable if all confounders that influence both treatment and outcome are observed and the ignorubility and overlapping conditions (see Assumption A1 are sa satiifieed To estimate ATE, variety of methods, such as outcome regression, augmented /inenterdddiveeeeeeee 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. (Oulcome model) The parumelric model Pa(X Oa) is a correct specification for the oulcome. funclion px(X)= = E (a) IE X), (0,1); Le., Pa () pa(X, 0%) and O
is the true model parameter.
Assumption 2. (Treatment model) The parametric model p(X, 0,) is correct specification for the propensity score p(X) = P(A = 1[X); i.e., p(X) = p(X, 0:) and 0; is the true model parameter.

---

Let 02 and 0, be consistent estimators of 0 and 0;, respectively, For example, the AIPW estimator of the ATE is given by
A.Y. XX, 0.)
- 0.) (X., Q1)
(1 (1- A.). A. A,-((X,,
no (X, 0o)
p(X,0.)) 0.)
(X,
1 1- 1-p(X,,00 0.)
which is doubly robust [obiin et al. (1944) in the sense that in is consistent if either Assumption or 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ühhlaan 2013 Schwab et al. 2020). environmental and healthcare data (Antonclli et al.) 2019), and social media data (Sharmal 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 (iibshirani 1996) and other regularization methods have been used to select relevant covariates for figh-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 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 net works (Liang et al. 2022 Sun Liang 2022) and sparse deep learning (Liang al.) 2018b Sunn et al. 2022) ; propose a Causal Stochastic Neural Network, which is abbreviated as Causal-Sto.Net in what follows, for
addressing the above difficulties encountered in causal inference for high-dimensional complex data. The merits of Causal-StoNet are threc-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 mrodeidg framework for complex data generation processes. In Causal-StoNet, we replace a hidden neuron at an appropriate hidden layer by a visible trett variable. With its compositional regression structure, Causal-StoNet easily extends various causal inference scenarios, e-g. missing covariates, multi-level or continuous treatments, and mediation analysis, as discussed Section 5
2. Universal approximation ability. We prove that the StoNet possesses a valid approximation to deep neural network, thereby enabling it to possess the universal approximation ability to the outcome and propensity score functions.
3. Consistent sparse lcarning. By imposing an appropriate sparse penalty/prit on the structure of the StoNet, relevant variables to the outcome and propensity score can be identified along with the training the Causal-StoNel even under the setting of higl-dimennsionll covariates. As 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 holistic manner, providing 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 e.g et al. (2019) and Farrell et al. (2021). However, the
networks for causal inference, see e.g., (Shi consistency of the deep neural network estimator is not established in Shhi al. (2019). This property has been studied in Farrell et al.) (2021) y but under the low-dinensional scenario essentially. Otherwise, it requirs the underlying outcome and propensity score functions to be highly smooth with the smoothness degree even higher than the data dimension In addition, the methods in Shi et al. (2019) or Farrell et al. (2021) cannot perform covariate selection, and they hard to apply when the dataset contains missing values. Quite recently, Chen et al.] (2024) proposed some neural network- ATE estimators, where only the propensity score the outcome function is approximated using neural network.
2

---

There also numerous semi-paramette 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 20060 proposes flexible semi-parametric framework based on targeted regularization. These methods can also be combined with deep learning other machine learning models, but the flexibility of the StoNet in forward modeling of complex data gencration processes leads to the uniqueness of StoNee It can function effectively in various data scenarios, as discussed in Section
2
A BRIEF REVIEW OF STOCHASTIC NEURAL NETWORKS
The StoNet can be briefly described as follows. Consider DNN model with h hidden layers. For the sake of simplicity, we assume that the same activation function (i.) is used for all hidden units. By separating the leeding and aclivation operators of each hidden unit, we can rewrite the DNN model in the following form:
Y = b + wi X,
Y. b, + w. w. (Y. 1), i= 2..,.......
(2)
Y - + wh+1 Vr (Yr) + En+1,
where en+1 NO, (.1 Id++1) is Gaussian random error; Y, b, R4 for i 2, h; Y, b Rdo+i; V (YY = (w(Y,-1.)), (Y)-..))...., (Y,- ( and Vi is the jth element of Y-11; w, E Rd. for = 1,2,..., 1++1, 1, and do denotes the dimension of 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.
for i 2, ................................ + 1,
The StoNet is probabilistic deep learning model and constructed by adding auxiliary noise to Y's for = 1,2,... h in 2 Mathematically, the StoNet model is given by
Yı bj + wiX X + t1,
Y. b, + w. V (Y (Y,
+ i =, h
(3)
Y = b + wi+1 (Yr) Eh+11
as a composition of many simple regressions, where Y1, Y..... : Yn can be viewed as latent variables. Further, we assume that e; (N, l4)) for i 1,2, +11. For classification problems, 244 plays the rolc of temperature for the binomial or multinomial distribution formed at the output layer, and it works with (of.... together to control the variation of the latent variables (Y...... Yn).
It has been shown in Liang et al. (0222) 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 02 (w,, b:), let 0 = (01, ... .....))) denote the parameter vector of the StoNet, let dg denote the dimension of 0, and let e denote the space of 8. Let L e: R
denote the loss function of the DNN as defined in (2). which is given by
L( - 1 Σ log Y (x00))
(4)
where indexes the training samples. Under appropriate settings for oj's, the activation function vi, and the parameter space O 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.
sup oc
Σ log क (1) y X (8)
Σ log (Y (()) 0)
L 0,
n Co. 55)
n
where Y Ymis = (Y1. Y..... Y,) denotes the collection of all latent variables as defined in
The StoNet can work with wide range of Lipschitz continuous activation functions
3

---

such as tanh, sigmoid and ReLU. As explained in Liang et al. (002)). Assumption A2 also restricts the size of the noise added to each hidden unit by setting: 01 S 02 2 ... < h+1
0h+1 = 011), and d ( x for any & E (1,2.....h) h1, where the factor di+1 (h+ de can be understood as the amplification factor of the noise eN 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 0 for the DNN is unique up to some loss-invariant transfor- mations, such as reordering some hidden units and simultaneously changing the signs of some weights and biases, see Liang 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 (0) (log (Y| ()))))) 0)), see Assumption A3 Liang al. (2022) showed
0 0- 3 O,
as
(6)
where 0n arg max ( Σ log π (Y Y X 0)}, and 0 arg max (0).
That is, the DNN (2) can be trained by training the StoNet (3); they are asymptotically equivalent 00, thereby the universal approximation property also holds for the StoNet. is worth noting that in forward prediction, the StoNet ignores auxiliary noise added to hidden neurons and thus performs the DNN.
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 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 compulational 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
T TYY, Ymis, A A 0) Xm x XY (T) Oz) (A 02) (Y 0.),
(7)
where Ymus = (Y, Y2, Ya), 0 = (0,, 02 03, 04), and ( Yi, 02) corresponds to the propensity score. For the visible binary treatment unit, sigmoid activation function can be used for its probability interpretation.
(05)
(00)
InputXX
Figure 1: Nausal-StoNet Structure: the treatment is included as a visible unit (rectangle) in a middle layer, and Y2 denotes the latent variable of that layer but with the unit directly feeding to the treatment rectangle excluded; 'x' represents possible missing values.
4

---

To ensure that sparse Causal-StoNet can be learned for high-dimensional data, where the mumber covariates can be much larger than the sample size n, we will follow Sun et al. 2022
2021 to impose a mixture Gaussian prior each component of 0, i..,
0 X0,, + (1 (O, oh,))
where 0 denotes a generic weight and bias of the Causal-StoNet, An (0, is the mixture proportion, oh is typically set to a very small number, while of,n is relatively large.
Let (x, A) denote the true outcome function, and let (*)) denote the true propensity score function. Tet p(æ, A; 0m) denote the DNN estimator of 1= (ar, A), and let f(x; 0.) denote the DNN estimator of p" (z:). For a given estimator 6n. both A(æ, A; Q.)) and p(2; 0,) are calculated as for the conventional DNN model (2) by ignoring the random errors e; Let Y = {%i i = 1.2,...,du d denote the true sparse structure of the Causal-StoNel,, which defined through a sparse DNN as in (A5). Here % is an indicator for the existence of connection c₁. Following Sunn et al. (2022), for each
{1,2.... d set fv 1
if the corresponding weight |0,1
Voi.a - fm log 1- din and 0 otherwise. Denote An da,
V200,n01.n
the estimated sparse Causal-StoNet structure by ((0,)) = : 1.2......o). Under appropriate conditions,
can show that the sparse Causal-StoNet leads to consistent estimates for μx (xx A); p (az), and y² 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 A$ hold, and I'n < 3/16
As n + co, the following results hold:
(a) (Propensity score function) With probability greater than - exxp(encf)) for some constant C,
(6(x; Q.) -p p=x))) (e + o((n-1/2)
(b) (Oulcome fanction) If 11 (ar) is bounded and the aclivation function ++)) E [-1,1], then, with probability greater than -xp(crnez) for some conslant c,
Em (x() Q,) n (x))3) O ((e - IT.) + o(n-1/2))
(c) (Structure selection) If Assumption A6 also holds, then P(S(O.) = 7.) 1 1.
By Theorem D the sparse Causal-StoNet provides consistent estimators for both the propen- sity score and outcome functions. Therefore, these estimators can be plugged into equation (1) to get double robust estimator for ATE. Morcover, 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 (015)..
Theorem 2. Suppose Assumptions AI A5 hold. Additionally, assume that the mixture Gaussian prior (8) is imposed on each connection of the StoNet, In < n3/16 and n 1+ x
wh < nn 2-5 and specify the network structure such that 0.5 { < < 1- and Ln = (n2)
for some S < 1/4. Then the following results hold:
(a) V V 12 T") 9, NN(0, 1), where V- is given in Supplement A.5 and T- denotes the true value of the ATE.
(b) V, vv where the estimator V, is given in Supplement
A.5
(c) (Uniformaly volid inference) Let P,r be set of data- nata-generating process satisfying Assumption A1 Then for ca = $ - a/2), we have
[
fn ea v (2//m/ - (1 -a) + 0.
sup P.EP.
Ppr
7" E
5

---

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 training the Causal-Stonet can be boiled down to solving a high- dimensional parameter estimation problem with latent variables present, i.e., maximizing
log π ( Y-U A| 0) + log (0).
To maximize (9) feasible method is adaptive stochastic gradient Markov chain Monte Carlo (SGMCMC), which, by the Bayesian version of Fisher's s identity (Song et al. 2020), converts the optimization problem to mean-field equation solving problem:
h(0) H(Y mise 1 0) mYmia X, Y A) 0,
(10)
where H(Ymiss 0) Vo log ग TYY Y miss A|X, 0) + Vo log T (0). The adaptive SGMCMC algo-
rithm 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 100 in the following equation:
h()) (2,)) / (z) =0
(11)
where Z a latent variable and π (== ()) probability density function parameterized by i E O The algorithm works by iterating between the following two steps:
(a) (Sampling) Simulate 2(*+1) (( ((k)) via a transition kernel induced by a SGM-
CMC algorithm such as stochastic gradient Langevin dynamics (Welling & Tch
2011 and stochastic gradient Hamilton Monte Carlo (SGHMC) (Chen et al. 2014].
(b) (Parameter updating) Set ( (()) + H (2(*k+1) (())), where 7k+1 denotes the step size used in the stochastic approximation procedure.
This algorithm is called adaptive SGMCMC the transition kernel used in step (a) changes along with the working estimate kkk) Applying the adaptive SGHMC algorithm to (10)
leads to Algorithm [ (in Appendix A.1. where SGHMC is used for simulation of the latent variables Y Yis at each iteration. The convergence of the adaptive SGHMC algorithm has
been studied in Liang et al.
(002)).
(022)) ) Suppose Assumplions A8 A13 hold. In Algoritio 7 if and 7R.i = Cn/(cg + ka) for some constants a E (0, 1), C. O, Ct2 0, Ce 2 0 and Cy 2 0, then there erists iteration ko and constont Ao > 0 such that for any Ko:
Lemma 1. [Linng el al.
we set ck.i = C./(co
ka)
B([ 0;11²) I AnTk.
(12)
where 0; denotes a solution to equation
10
In Liang et al. (2022). an explicit expression of Ag has been given. For simplicity, we have the expression omitted in this paper. Next, [iang et al. (002)) showed that h: -> co, the imputed latent variable 2(x) converges weakly to the desired posterior distribution x x=2(0*) in 2- 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 A11 hold. Then for any bounded function do- (+); ExØg- (z) f (z)g(z)dm( θr)= 2, 0 K. co, where (x)g-(2) = E Paco (2(0) and ((0()), ,()) : i = 1, 2. KJ denotes a set of parameter estimates and imputed latent variables that collected in 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) (Hart& & 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 ad adversarial networks (GANITE) YYonn et al. 2018), and
counterlactual regression net (CFRNet)) H
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) 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 800, 1600, 2400, 3200, 4000, respectively. For each scenario, 10 simulated datasets generated as described in Section A.7.1 Both the outcome and treatment effect functions in this experiment
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 I 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 otherwisc stated.
The results are depicted in Figure 2 demonstrating that Causal-StoNet maintains stable performance even with high-dimensione covariates and small sample sizes. Additionally, we conducted further simulations to investigate covariate selection accuracy and address missing value problems, as detailed Section A.7.2
12
CEVAE
CFANGSAND
DONUT
30-
In-Sample MWE
Sample
2.
1000
2400
4000
2400
Training samples
Training samples
Figure 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 caleulated over test set.
'The code of the experiments is available at: httpps///tthhhh.cc/aaaaa/alaaaa

---

4.3 ATLANTIC CAUSAL INFERENCE CONFERENCE 2019 DATA CHALLENGE
The Causal-StoNet is compared with baseline methods on 10 synthetic datasets with ho- mogeneous treatment effect from the Atlantic Causal Inference Conference (ACIC) 2019 Data Challenge. Each dataset contains 200 covariates with binary trealment variable, and the outcome variable is continuous. Since they both are synthetic, the true ATE is known. Results in Table @ 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 MAF.
Method
In-Sample
Out-of-Sample
Causal-StoNet
0.0501 (0.0118)
0.0542(0.0132)
DSE
0.0776(0.0193)
0.163)(0.0251)
ARBE
(..02)
00133
TML.F( (Tass))
0.0869(0.0164)
0.0867(0.0165)
TMLE (ensemble) DONUT
0.1140(0.0394)
0.1316(0.0429)
0.5294 (0.2640)
0.5290(0.2644)
In addition to ATE, we also consider the condi- tional average treatment effect (CATE), which measures the heterogeneous treatment effect for subpopulations or individuals based on their co- variates IR" and is defined by
Table 2: CATE estimation for an ACIC 2019 dataset.
Method
<PEHE
CATE
Causal-StoNet
0.0893
0.0118
CMDE
0.0823
0.0444
7 (ax) = ((y - - x z].
CMGP
0.2156
0.0258
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 metries: éPEНE (square root of the Precision in Estimation of Heterogeneous Effect (PEHE)) and CATE (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 éPEHE, it achieves the lowest CATE-
CEVAE
0.0867
0.0358
GANITE 0.1913
0.0485
X-Learner-RF
0.1877
0.0203
X- L-Lernee-BBRRT
0.0873
0.0720
SeRNet-Wass 0.1182
0.0421
CFRNet-MMD
0.1158
0.0849
4.4 TWINS DATA
We analyzed dataset of twin births from 1989 to 1991 in the United States. The treatment variable binary, with a denoting the heavier twin at birth; and the outcome variable is binary, with Y = 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. for data preprocessing steps. After data pre-procsssing, this final dataset, mortality rates for lighter
we obtained a dataset with 4,821 samples. and hcavier twins are 16.9% and 14.42%, respectively, resulting in true ATE of 22.4%%. 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 more stable ATE estimate than baseline methods.
Table 3: ATE estimates by different methods for twins data
Causal-StoNet
DSE
ARBE TMLE(Lasso) (MME(ensemble))
DONUT
-0.0232(0.0022) -..045500.0176) -0.0096 (0.0201) -0.1103(0.0500) (0.90()))779)) 0.073880.00128)
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.
8

---

5 SOME VARIANTS OF CAUSAL-STONET
The proposed Causal-Sto.eet can be casily 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 Xohs denote the observed covariates, let X mis denote the missed covariate values, and let R denote the missing pattern represented as binary vector. Under the mechanism of missingness at random, i.e., X ( A, Y), the Cousal-StoNet as depicted in Figure is to learn decomposition of the joint distribution
T(Y, Ymis, X mis, A|X obs: obs, R, 0) XT Xia X ohs ) Xbba, X mis, 0 (Y2
(13)
x π ( 82) Yg Y2 A, (01, 22 0z. 04): π (A( corresponds to the propensity score, and (Xmis Xohs) 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[]by treating Xmis part of the latent variables. Statistical inference with imputed missing data can then be made based on Lemma P
04).
where Yiis (Y, Y3),
Missing not at Random (MNAR) The Causal-StoNet can also be extended to the sccnario 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 al. (2019), we will assume that the missing pattern R is independent of the outcome given the treatment and confounders, i.e., (L,R(A, oobs: Xmis). Under this assumption, T (Y, Ymiss X mis,; AlXobs; R, 0) OT (Xmis X obs ) (T| abs: X mis, O) (R]Xabs: X miss A, 0)
x T ( X Xobs Xmis: A, 0)
To accommodate the term x XRRoosss obs: X mis: A, 0) 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 [ can be directly applicd 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 A Analysis In this scenario, we aim to measure how the treat effect is affected by intermediate/ meediaaiin variables. For example, Pearl (001) gave an example where the side effect of a drug may cause patients take aspirin, and the latter has a separate effect on the disease that the drug was originally prescribed for. The mediation analysis can be casily 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.
CONCLUSION
We have developed an effective method for causal inference with high-dimensional complex data, which addresses the difficulties, including high-dimensional covariates, unknown treat- ment 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 the- oretical 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 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.
9

---

REFERENCES
Christophe Andrieu, Eric Moulines, and Pierre Priouret. Stability of stochastic approximation under verifiable conditions. STAM Journal on Control and Oplimization, 44(1):283 312, 2005.
Joseph Antonelli, Giovanni Parmigiani, and Francesca Dominici. High-dimensional confound- ing adjustment using continuous spike and slab priors. Bayesian analysis, 14 3:805- 828, 2019.
Susan Athey and Guido Imbens. Recursive partitioning for heterogeneous causal effects. Proceedings of the National Academy of Sciences, 113:7353 7360, 2015. URL https
//api.ssmanticccchol..e httpi/ioo..rrrrrrrroooooo (
Susan Athey, Guido W. Imbens, and Stefan Wager. Approximate residual balancing: De- biased inference of average treatment effects in high dimensions. Journal of the Royal Statistical Society. Series B, Statistical methodology, 80(4):597-623, 2018.
Alexandre Belloni, Victor Chernozhukov, and Christian Hansen. Inference on treatment effects after selection amongst high-dimensional controls. Review of Economic Studies, 81: 608 650, 2014.
Albert Benveniste, Michael Métivier, and Pierre Priouret. Adaplive Algorithams and Stochastic Approzimations. Berlin: Springer, 1990.
Peter Bühlmann. Causal statistical inference in high dimensions. Matheraatcal Methods of Operations Research, 77:357 370, 2013.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pp. 1683-1691 2014.
Xiaohong Chen, Ying Liu, Shujie Ma, and Zheng Zhang. Causal inference of general treatment effects using neural networks with a diverging number of confounders. Journal of Econometrics, 238(1)) (1):105555,, 2024. ISSN 0304-4076. doi: https:/ http://///ooooo1101111111 aeconom.20253 105555. 503044006220027113
URL 10555.
https///www.sceeeeeeeettttt het.a//eeeeece//eeeeeeeettttttaa
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiasee machine learning for treatment and structural parameters. The Economeriies Journal, 21, 2018.
Hengjian Cui, Runze Li, and Wei Zhong. Model-free fcature screening for ultrahigh dimen- sional discriminant analysis. Journal of the American Statistical Association, 110(510): 630 641, 2015.
Wei Deng, Xiao Zhang, Faming Liang, and Guang Lin. An adaptive empirical bayesian method for sparse doep learning. Advances in neural information processing systems, 2019: 5563-5573, 2019.
M. Farrell. Robust inference on average treatment effects with possibly more covariates than observations. Journal of Econometries, 189:1 23, 2015.
M. Farrell, Tengyuan Liang, and S. Misra. Deep neural networks for estimation and inference. Econometrica, 89:181-213, 2021.
Xuefeng Gao, Mert Gürbiizbalaban, and Lingjiong Zhu. Global convergence of stochastic gradient hamiltonian monte carlo for nonconvex stochastic optimization: Nonasymptotic performance bounds and momentum-basee acceleration. Operations Research, 2021.
Qian Guan and Shu Yang. A unified framework for causal inference with multiple imputation using martingale. Xrxi:: Methodology, 2019.
Tobias Hatt and Stefan Feuerriegel. Estimating average treatment effects via orthogonal regularization. In Proceedings of the 30th ACM International Conference on Information & Knowledge Management, pp. 680 689, 2021.
10

---

Guido Imbens. Nonparametric estimation of average treatment effects under exogencity: A review. The Review of Economics and Statistics, 86:4-29, 2004.
Wenxin Jiang. Bayesian variable selection for high dimensional generalized lincar models: convergence rates of the fitted densities. The Annals of Stalislics, 35( 35(4):1487 1511, 2007.
Ziyang Jiang, Zhuoran Hou, Yiling Liu, Yiman Ren, Keyu Li, and David Carlson. Estimating causal effects using a multi-task deep ensemble. In Proceedings of the 40 th International Conference on Machine Learning (ICML), PMLR, pp. 680-689, 2023.
Sören R. Künzel, Jasject S. Sckhon, Peter J. Bickel, and Bin Yu. Metalearners for estimating heterogencous treatment effects using machine learning. Proceedings of the National Academy of Sciences of the Uniled States of America, 116:4156 4165, 2017. URL
Stttps://app semanticscholas
Serg/orrpusID. 73455742
Lihua Lei and Emmanuel J. Candès. Conformal Inference of Counterfactuals and Individual Treatment Effects. Journal of the Royal Statistical Society Series B: Statistical Methodology,
83(5):911-995 10 2021. ISSN 1369-7412. doi: 10.1111/rssb. hal1l1/rssb.12455. 12445. URL pttpss//doo org/10. 111.1/ssbsbsbb..e 12445.
Jiuyong Li, Saisai Ma, Thuc Duy Le, Lin Liu, and Jixue Liu. Causal decision trees. IEEE Transactions on Knowledge and Data Engineering, 29:257 271, 2015. URL https
/appi wewanticccholoaar org/ (CorppssID::7188772
F. Liang, C. Liu, and R.J. Carroll. stochastic approximation in monte carlo computation. Journal of the American Statistical Association, 002:305-320, 2007.
F. Liang, Q. Song, and K. Yu. Bayesian subset modeling for high dimensional generalized linear models. Journal of the American Stalislical Associalion, 108:589 606, 2013.
F. Liang, Y. Cheng, and G. Lin. Simulated stochastic approximation annealing for global optimization with a square-root cooling schedule. Journal of the American Statistical Association, 109:847-988 863, 2014.
F. Liang. B. Jia, Xue, Q. Li, and Y. Luo. An imputation--eeguliizzd optimization algorithm for high-dimensional missing data problems and beyond. Journal of the Royal Statistical Sociely, Series B, 80(5):899 926, 2018a.
F. Liang, Q Li, and L. Zhou. Bayesian neural networks for selection of drug sensitive genes. Journal of the American Statistical Association, 113:955 972, 2018b.
Siqi Liang, Yan Sun, and Faming Liang. Nonlinear sufficient dimension reduction with stochastic neural network. NeurIPS, 2022.
Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling, Causal effect inference with deep latent-variable models. 31st Conference on Neural Information Processing Systems (NIPS 2017), 2017.
Judea Pearl. Direct and indirect effects. In Proceedings of the 17th Conference on Uncertainty in Artificial Intelligence (UAI), pp. 411 420, 2001.
Maxim Raginsky, Alexander Rakhlin, and Matus Telgarsky. Non-convex learning via stochastic gradient langevin dynamics: nonasymptotic analysis. In Conference on Learning Theory, pp. 1674-1703. PMLR, 2017.
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathematical statistics, pp. 400 407, 1951.
James M. Robins, Andrea Rotnitzky, and Lue Ping Zhao. Estimation of regression coefficients when some regressors are not always observed. Journal of the American Statistical Associa- tion, 89:846-866, 1994. URL httph::t///pppat www.sticschoiilaaa.t.t owgg/Corpuuseeeee :2: 120769390
P. R. Rosenbaum. Observational Studies (2nd edition). Springer, New York, 2002.
11

---

Donald B. Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of Educational Psychology, 66:688-7011 701, 1974.
Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M. Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. 2020.
ArXiv, abs/
Ankit Sharma, Garima Gupla, Ranjitha Prasad, Arnab Chatterjee, Lovekesh Vig, and Gautam M. ShrofT. Hi-ci: Deep causal inference high dimensions. Journal of Machine Learning Research, 2020.
Claudia Shi, David M. Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. In NewrIPS, 2019.
Q. Song, Y. Sun, M. Ye, and F. Liang, Extended stochastic gradient meme for large-scale bayesian variable selection. Biometrika, 107 107(4))977-100) 2020.
Y. Sun, Q Song, and F. Liang. Consistent sparse deep learning: Theory and computation. Journal of the American Statistical Association, 117((40):198111195, 2022.
Yan Sun and Farning Liang. A kernel-expanded stochastic neural network. Journal of the Royal Statistical Society Series B, 84(2):547-577 577, 2022.
Yan Sun, Wenjun Xiong, and Faming Liang. Sparse deep learning: new framework immune to local traps and miscalibration. NearIPS 2021, 2021.
R. Tibshirani. Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society, Series B, 58:267-2288 1996.
Mark J. van der Laan and Daniel Rubin. Targeted maximum likelihood learning. The International Journal of Biostatistics, 2(1):11, 2006.
Mark J. van der Laan, Eric C. Polley, and Alan E. Hubbard. Super learner. Slalistical Appli- calions Genelies and Molecular Biology, 6, 2007.
org/CorpusID: 7162180
Max Welling and Yee Whye Tch. Bayesian learning via stochastic gradient Langevin dynamics. In ICML, 2011.
Teh Whye, H ThieryAlexandre, and J VollmerSebastian. Consistency and fluctuations for stochastic gradient langevin dynamics. Journal of Machine Learning Research, 2016.
Li Xie, Kun Wang, Haiquan Chen, Yanlong Shi, Yuanqi Zhang, Hao yu Lin, Yuan ke Liang, Yingsheng Xiao, Zhi-Yong Wu, Zhongyu Yuan, and Si qi Qiu. Markers associated with tumor recurrence in patients with breast cancer achieving pathologic complete response after neoadjuvant chemotherapy. Frontiers in Oncology, 12, 2022. URL https
//api..emaatticccholaa... wttp.oooooo/oos./222522
Shu Yang, Linbo Wang, and Peng Ding. Causal inference with confounders missing not at random. Biometrika, 106:875-888, -888, 2019.
Jinsung Yoon, James Jordon, and Mihaela van der Schaar. Ganite: Estimation of individual- ized treatment effects using generative adversarial nets. In International Conference on Learning Representations, 2018. URL https://api. swwanticscholar..c.mm wtt.://cccooaaaaooooorrprppss.ee Irrg/orrusID: 65516833
A APPENDIX
A.1 ADAPTIVE STOCHASTIC GRADIENT HAMILTONIAN MONTE CARLO ALGORITIM
Let (x(())) y((t)) = (x()), y denote a training sample 8, and let ((s.k) (x)(a)) ..... y((.))) denote the latent variables imputed for the training sample s at iteration 1.a
k. where the subscript a indicates that the imputed values are affected by the treatment variable A.
12

---

Algorithm 1: An Adaptive SGHMC algorithm for training StoNet
Input: total iteration number K, Monte Carlo step number tM the learning rate sequence {fA,e : k = 1. ...... K; 2....... h + 1}, and the step size sequence (Tk. : k = ........ K = 1,2, ...,...... .,.,. 1); and the step size sequence {TR. k == 1.2..... K; 1,2,..., 1}; Initialization: Randomly initialize the network parameters O(0) = (01
for k= f=_,,,...,4 do
STEP Subsampling: Draw a mini-batch of data and denote it by Sk:
STEP 1: Backward Sampling
For each observation E Sk, sample y( s, in the order from layer h to layer 1, from
(V ( yleik) a (5.b) @(k-1) 0 1+1,a yiaiaa a 141a ++1 )m 10
((,k)
1-1.a)
by rumning SGHMC in EMC steps: Initialize 8. (2.0) 0, and initialize y:(ak0)) by the corresponding Y, calculated in 22. for 1= 2....... tMe do for i = h I...., do
a (o.k.)) =(1 Ck (1 (1) + CR. V to log (vaaa (y(akii-)) @(t-1), y )((abal11)
(((k111)) 0/4411 y/aabil-) + 2ek, :ee
te,s V (( log i+1,a
(s,,)) -y(,,,1111) + c e (a.1.-1)
(A1)
where en N(0, Ia,), Enis is the learning rate, and ry is the friction
coelficient. end end
((a0) (s.kitte)) for i = 1,2,.... Set 1. t.a
STEP Parameter Update
Update ((k) = (0()) ((()) ,...,, ((h) by stochastic approximation
Robbins & Monrro
1951]: :
((h) = 0(h-11 1) + Yk,i
Σ sES
Vo. log x
(1( y((a)t)
+
SA Vo, log T (0)
1.
n
(A2)
for 1,2,..., h 1, where Ya. is the step size used for updating O,.
end
A.2 ASSUMPTIONS ON DATA GENERATING PROCESS
Suppose that the dataset (1) ai, 2:) :i= 1,2,.... n) is generated from a process Pn, where a 10, (() = pa (X) x)+U0 ++U, U, and U denotes random error. P., obeys Assumption A1 with bounds uniformly in n:
Assumption A1.
(a) ((),, az 2.) :1 .........) an i.i.d sample from (Y.AX X).
(b) (Ignorability) (0) Y
(c) (Overlap) 0<6 &1 < P(A 1|X) S C2 almost surely for some Çi and Çz.
(d) XX < U for some U 0.
(e) For some d > 0, Ellua X) (X) and [E[ [+8] are bounded.
13

---

A.3 ASSUMPTIONS FOR STONETS
The property of the StoNet as an approximator to the DNN, i.c., asymptotically they have the same loss fuuction the training sample size n 00; has been studied in Liang et al.) (2022). A briel review for their theory is provided as follows, which form the basis for this work.
Let 0 = (wi, bi..... Wh+1, bx+1) denote the collection of all weights of the StoNet (3) let ④ denote the space of 0, let Yiss (Yi, .2...., .., Yx) denote the collection of all la- tent variables, let T(Y, Yuis (xX)) denote the likelihood function of the StoNet, and let TYNN(Y 0) denote the likelihood function of the DNN model (2). Regarding the network structure, activation function and the variance of the latent variables, they made the following assumption:
Assumption A2. (i) + is compact, i.e., O is contained in da-ball centered at 0 with rudius r; (ti) E(log π (Y Ymis X, 00))² < oo for any 0 E O; (ti)) the activation function (()) is c'-Lipschilz continuous for some constant d; (iv) the network' ' depth h and widths di's are both allowed to increase with n; (V) 01 S 02 S ... oht1, oht1 = O(1), and di (iie+1d)ddd d x h for any k (1.2......
Assumption A2 (iii) allows the StoNet to work with a wide range of Lipschitz continuous activation functions such as tanh, sigmoid and ReLU. Assumption A2 (v) constrains the size of noise added to each hidden neuron, where the factor dx+1 k= df )d can be understood as the amplification factor of the noise ex at the output layer. In general, the noise added to the first few hidden layers should be small to prevent large random crrors propagated to the output layer. Under Assumption A2 they proved the result (
Further, regarding the equivalence between training the StoNet and the DNN, they made the following assumption regarding the energy surface of the DNN. Let (* E(log TDNN TDNN(Y)X 0)), where the expectation is taken with respect to the joint distribution T (X,Y) By Assumption A2 (ii) and the law of large numbers,
log DDN (Y (0,00 Q' 20
(A3)
holds uniformly over θ They assumed (0) satisfies the following regularity conditions:
Assumption A3. (O) is continuous in 0 and uniquely marimized at 0 ((t for any
e > 0, supacey, B()) (0) erists, where B(e) = {0 1|0 8- < <c,, and = (0')
(upp > 0.
A3 restricts the shape of (0) around the global maximizer, which cannot be discontimuous or too flat. Given nonidentifabiility of the neural network model, Assumption A3 has implicitly assumed that cach 0 is unique up to the loss-invariant transformations, e.g., reordering the hidden neurons of the same hidden layer and simultaneously changing the signs of some weights and biases. Under Assumptions A2 and A3 they proved the result 6 In summary, we have the following lemuna for the Sto
Assumption
Lemma A1. Liang et al., (022) Suppose Assumptions and A3 hold, and π TYY, Ymis| (X,) is continuous in 0. Then 5 and
hold.
A.4 PROOF OF THEOREM
To prove Theorem I we first give a brief review for the theory of sparse deep learning Sun et al. (2022). The theory under the context of Bayesian deep neural networks where each parameter, including the bias and connection weights, is subject to a mixture Gaussian prior distribution.
Assume that the distribution of y given x given in generalized linear model as ply. x; 0) = exp(A(p (x))y + 3(p"(x)) + C(y)),
where (x) denotes a nonlinear function of x, and A(-), B(-) and C(-) are appropriately defined functions.
14

---

Motivated by the universal approximation ability of the DNN, they proposed to approximate 4² (x) using a DNN with Hn, hidden layers and L, hidden units at layer h, where LH, = for the output layer and Lo Pn, for the imput layer. Let w" R xn. and ba E R h E (1,2,.., Hn} denote the weights and bias of layer h, and let 1/ R4ax1 + TRLax1 denote coordinate-wise and piecewise differentiable activation function of layer h. The DNN forms a nonlinear mapping
p(0, x) w a/ (...z)) 1w0 x + b'] 1..
(A4)
where 0 -w b) = wulj b/ E (1, ..... Hnj.ikk E {l, Lob.j E {1, (I, (I,...,A-11} denotes
the collection of all weights and biases, consisting of Kn = = (Ln-1 Ln + Ln) elements
in total. To facilitate representation of the sparse DNN, they further introduced an indicator variable for each weight and bias of the DNN, which indicates the existence of the connection in the network. Let y" and 70t denote the matrix and vector of the indicator variables associated with wht and b", respectively. Let 7 = (ng 7 E {1, .... Hah, i, k (1,-., T.) J E (1,..., [n-1]} and 0, (wfz. U/ : 7 ~,, =1 ,h E (1,2, Hn), ) i, & (1..., j {I,..., L, which specify, respectively, the structure and associated parameters for a sparse DNN.
Among many DNNs that can provide good approximation to n ('(x), they define the Irue DNN model as
(0" x-))==
arg min (0,7)S.. -00
hl.
(A5)
(o)
where Gn = g(Co, Cr, Pr. Hn. L,, L LH.) denotes the space of valid sparse DNNs
satisfying condition Assumption A4 (A.2) for the given values of Hn, Pu, and L, is and wa is some sequence converging to 0 as n - oc. For any given DNN (0,7), the error 110,, 7, u (x) can be generally decomposed as the network approximation error p(0","", - H (x) and the network estimation error p00.7,v) - p(0", ",",""") x). They generally treated wn as the network approximation error. In addition, they made the following two assumptions.
A Assumption A4. A.1 The imput x is bounded by 1 entry-wisely, ke. x E x [-1, 1)"",
and the density of x is bounded in its support X uniformly with respect n.
A.2 The true sporse DNN model satisfies the following conditions:
A.2.1 The network structure satisfies: rn Hn log n +rn log Ln + Sn log Pr. I Con' where 0 < E < 1 is small constant, = l denotes the connectivity of r T. = maxisjsH,,-1 I.j denotes the marimum hidden layer width, 8n denotes the input dimension of 7Y"
A.2.2 The nelwork weights are polynomially bounded:: || S En, where E, = Cr for some constant C1 0.
A.3 The activation function vp is Lipschitz continuous with Lipschitz constant of 1.
Assumption A5. The miature Gaussian prior (6) satisfies the conditions: Au
(n
for some constant T > 0, En/(H, logn + log Ln Lnj1f² 11/2 S
Jin S na for some conslant a > 0, and Jon S min (1/(VKK, (13/2 Hn)Ha)) }, (I..
Ha
Based on the two assumptions, Sun et al. consistency of the sparse DNN.
2022 proved the following result on posterior
Lemma A2. (Theorem 2. 1; Sun et al. (022) Suppose Assumptions A4 A5 hold. Then there
cxists an error sequence 2 = ((w,)) + O(C)) such that limm-+oo En = and lim ne, 0o, and the posterior distribution satisfies
IP p [(d(po.P/-)) > 4e, Dn] 2 20 ene?] } S 2 -ene
(A6)
ED, n xld(po, pp > Dal < 4e 2one?"
for sufficiently large n, where e denotes a constant, Dn denotes a dataset of n i. obser- vations, C2 [r,, Hr log n rn log Ton + 8n log logpn]/n, Py- denotes the underlying true data
15

---

distribution, and po denotes the data distribution reconstructed by the Bayesian DNN based on its posterior samples.
is known that the DNN model is generally nonidentifiable due the symmetry of the network structure. For example, the approximation the orders of certain hidden nodes, simultaneously changes the signs of certain weights and biases if tanh is used as the activation function, re-scales certain weights and bias if Relu used as the activation function. To address this issue, Sun et al. (2022) considered a set of DNNs, denoted by S2 for which each clement can be viewed models. Let 21(y, 0) E Ω be an operator that maps any neural network to S via appropriate transformations such nodes permulation, sign changes, weight. rescaling, etc. To the purpose of structure selection in the space It they adopted the marginal posterior inclusion probability approach (Liang al. 2013).
can be invariant if one permutes
equivalent class of DNN
For each connection G₁, they defined its marginal posterior inclusion probability by
qi = //E Σ (11)) = 1,2,.... Kn:
(A7)
where ()) is the indicator for the existence of connection c; in the network ( 0). Similarly, 6 (iv(r.)))) denotes the indicator for the connection in the true model "(>",, 0*). Let Py- denote the underlying true data distribution, and let po denote the data distribution reconstructed by the Bayesian DNN based on its posterior samples. Let A(e,) (0
d(pe. P 2 6x), where d(pI.P2) P2) denotes the Hellinger distance between two distributions. Define
/
p(G,) max 13isK, Alen)
Eeeqpu(r.))
- Gijpir (..))
which measures the structure difference between the true model and the sampled models on the set A(c,)). Further, they made the assumption:
Assumption A6. ple,) + 0, as n + co and En + 0.
That is, when n is sufficiently large, if a DNN has approximatoly the same probability distribution as the true DNN, then the structure of the DNN, after mapping into the parameter space Ω must coincide with that of the true DNN.
1.2,..., Let R/ {i qi 1,2,1,2.... Kn}
estimator of % (i Gijerr .0)) 1, 1...., Knh, where Y can be viewed as the uniquenized true model. In summary, they have the following result regarding network structure selection:
Lemma A3. (Theorem 2.2; Sun et al. (2022) Suppose Assumptions A4 A6 hold. Then, as + 00, we have
(a) (sure screening) P(r, C 9c) & for any pre-specified Ç (0,
(b) (Consistency) P(7-=fo1 $
As shown by Sun et al. 2022 Lemma
A3 implies consistency of covariate selection.
For binary classification problems, the method leads to predictive distribution p(æ))
S p(æ; (O) as the Bayesian estimator of the true classification distribution p (x) ==
P(y 1/x = 1). Let (x)()) = (d()))x(dA) = =(x), and v' = (x)()
(z))
Lemma A4. Suppose Assumptions
A4 A5 probability greater than [xx
hold. Then the following inequality holds with
Em (())(x) - n ( (x)))²
< 4 16E and /1 /6/ 2
where 8 < (pp 0.5 D.) denotes lower bound of the selection probability of the selection
rule > 0.5, and [. [.] denotes expectation with respect to V i.e., the probability measure of
16

---

Proof. The proof of Lemma A4 follows from the arguments of Jiang 2007 (23)-(25)), i.e. for binary classification, we have
(around equations
Ex (((A(()) (x))²) 1 4d (p(x), Pp-) z 4e 8T(d(pe Pμ*) cn Dn.//
The proof can then be completed by applying the first inequality of (A6) to x[d(po: PH-) G, (n]
Note that the overlapping condition generally assumed for the treatments ensures that Ç is bounded away from zero and thus the approximation 4e, + 16e cne en holds.
Regarding the generalization error for regression problems, they proved the following result:
Lemma A5. (Theorem 2.6; Sun et al. (022)) ) Suppose Assumptions A4 A5 hold. If a is compact, the activation function ((-) E - 1, and p (æ) is bounded, then the following inequality holds with probability greater than 2expeene
Eig ((00 (x))² x (Cg t enee I
where Cn is defined in Lemma A2 and (-) denotes an expectation with respect to v(x), the probability measure of x.
To accelcrate computation, Sun et al. 2022) ) suggested to replace the Bayesian estimators involved in Lemmas A5 by their Laplace approximators. That is, instead of sampling from the posterior distribution, they conducted optimization to maximize the objective function
a Σ log 2og(p((). x,; + log(T(0)),
(0)
(A8)
n
where (0) denotes the mixture Gaussian prior specified in Denote the resulting maximum posteriori (MAP) estimator by
arg max I., (0). OE
(A9)
Assumption A7. Assume Tn = |y* grows with the sample size n at a rate of o4n)/4,
i.e.
Tn < 1/4
By invoking the Laplace approximation theorem, they show that the consistency results established in Lemmas A3 A5 still hold for O,, with the approximation error decaying at rate of O ) This leads to the following corollary:
Corollary 1. For sparse DNNs, the consistency results established in Lemmas A3 A5 also
hold for the maximum a posteriori estimator A9 provided that Assumption A7 also holds.
Moreover, address the local trap issue possibly encountered in maximizing AS Sun et al. (2022) suggested to run optimization procedure, such as Adam or SGD, mu multiple times and select the solution according to the BIC criterion.
Proof of Thcorem
Proof First, we note that these results hold for the sparse DNN following from Lemmas A3 and the property of Laplace approximation. Specifically, the term o(n-1/2) in parts (a) and (b) follows from Theorem 2.3 of Sun et al. (2022), where accuracy of the Laplace approximation for Bayesian sparse DNNs is given as O(r2/n). Therefore, for the mean squared errors in part (a), we have
Fe (f(x; (x; 0m)) (x)) < 2 (f(x) -
-p" p='(((())) (x)) + 2 (f(x;; Qm) - ((x))
0
2 + =cne + +(11./n²)
17

---

where p(x) p(x; (o),,, defined in Lemma A4 denotes the Bayesian estimator of (x=). Therefore, the result follows under the assumption In < 3/16 For part (b), the result can be justified in similar way.
Then, by equation (5), these results can also be achieved by the Causal-StoNet which is
trained by maximizing the penalized log-likelihodd function p This completes the proof of the theorem.
A.5 PROOF OF THEOREM 2
For convenience, for any a, a' E {0, 1}, we define the following notations:
(a E(P(A a|X =x)), Pa (pa(2)),
(x(x) = (E A = a, = =x), Ho = (aY())),
(x) E(U²) A a, =x).
With the above notation, we have
o
( (x()
V, E
+E ((11(X) m1) -mmo(X) -mo)
Pi (X)
For estimation, we define
n
Ba = ix(X x. A a, 0n), Pa(2r) p(x -x, a On), Pa Σ
I(A: a),
T(A: = a) - Ma(33))
+ Aelaai)
Aa =
(x(x)))
i=1
With the notation, we have
D, = E, [(A -1 1(A, 0) - io (x:)))
+Bn ((p) - (.)) - -(o((=)) - 20))
(x;)²
where B[:] denotes the empirical mean over n samples.
Proof of Theorem
Proof. The results directly follow from Theorem and Corollary 2 of Farrell (0155) Tt. is casy to verify that the conditions n- -1+E x w ensure that 2 e xn - 1/2 and ( te L x xn 1/2
x n - 0.5 < E < 1 1- E, and I,, O
and thus Assumption 3 of Farrell (2015) holds. Assumptions 1 and 2 of Farrell (2015 are given in Assumption A1 of this paper. Therefore, Theorem 3 and Corollary 2 of Farrell 2015) hold, and the statement of Theorem 2 is implied.
A.6 CONVERCENCE OF THE SGHMC ALGORITHM
Notations: We let D denote a datasct of n observations, and let D; denote the i-th observation of D. For StoNet, D; has included both the input and output variables of the observation. re-denote the latent variable corresponding to D, by Z,, and
For simplicity of notation, denote by fD. (zi, 0) = log क (2; D., (D,0) 0) the negative log-density function of Z,. Let z = (21, 22...., 2n) be realization of Z (Z1, Zz..... Za), and let FD(2,0) Σ fp. (2, (),
To study the convergence of Algorithm I we need the following assumptions:
Assumption A8. The function FD takes nonnegative real values, and there crist constonts A, B 2 0, such that [F, 0*) < A, (V 0*) s B, [VVoD((, < B, and (0, 0*) < B.
18

---

Assumption A9. (Smoothness) FD(:.-) is M-smooth and :(:,-) is M M-Lipschitz: there erists some constant M > such that for any Z, R and any 0' E e,
VV z, 0) VzFD ( 0l) - M Z Z'Il + M 0'1l
(VFD((Z,0) VeFD(Z' 00) | - M| Z Z'Il M 0 0'1l
() ( 0'1l S M Z Z MM10 -0'|- -
Assumption A10. (Dissipativity) For any O the function FD 0*) is (m. b) disssipattee there erist some constants m > 2 and 2 0 such that (Z, V F (Z, 0²)) 2 ml 2 b.
The smotthhess and dissipativity conditions are regular for studying the convergence of stochastic gradient MCMC algorithms. and they have been used in many papers such as Raginsky et al. (2017) and Gao et al. (2021). As implied by the definition of Fz 0), the values of MZ and b increase linearly with the sample size Therefore, we impose nonzero lower bound on m to facilitate related proofs.
Assumption A11. (Gradient noise) There erists a constant Ç E [0,1) such that for any Z and 0, E 22 0) F2 0)| = 2<(M² + M 10 -0 B²).
Introduction of the extra constant facilitates our study. For the full data case, we have = 0, i.e., the gradient VzFD(Z,0) 0) can be evaluated accurately.
Assumption A12. The step size {7x JEN is a positive decreasing sequence such that " + O and _k=1 7k oc. In addition, lel h(()
[Z 0)), then there exists 5 >> 0 such that for 0 112,, and lim in
any O (0 0", (h))) 2 8
26 2 7++1 Tre O.
7641
As shown by Benveniste et al. 1990) (p.244), Assumption A12 can be satisfied by setting
Tkk = aā ha) for some constants ā >> 0, b 2 O, and a E (0, A 20ō). By (A2). 5 increases linearly with the sample size n. Therefore, if we set ((1/n) then 20ā
be satisfied,
where (2(-) denotes the order of the lower bound of function. In this paper, we simply
choose E (0, 1) by assuming that has been appropriately with 28a 2 held.
Assumption A13. (Solution of Poisson equation) For any O E θ 2 3, and function V(z) = ||=||, there exists a function no on 3 that solves the Poisson equation po(z) Topo 2) = (0, h(0), where To denotes probability transition kernel with Topoa(z) 13 (To xz))))))" 'z'd such that
H(Kk, 2k+1) h(0)) + moe (zK+1) To (#++1); k= 1,2,.....
(A10)
Moreover, for all 0, 0' E θ and z E 3, we have (1))(2) uo (z)1l S s1 - 0 -0][(())) and [[))(=)] (z) for some constants Çi > O and 52 >
This assumption is also regular for studying the convergence of stochastic gradient MCMC algorithms, see e.g., Whye et al. (2016) and Deng et all (2019) Alternatively, one can assume that the MCMC algorithms satisly the drilt condition, and then Assumption A13 be verified, see e.g., Andrieu et. al. (2005).
Outline of the Proof of Lemma Lemma can be proved in a similar way to Theorem 3.1 of Liang et al. (0222 Note that in the proof of Lemma I the boundedness of θ is not assumed.
Lemma can be proved in a similar way to Theorem 3.3 of Liang et al. (2014) by ignoring the paramcter 7 used in the proof there and modifying some notations appropriately.
Outline of the Proof of Lemma 2
A.7 MORE NUMERICAL RESULTS
A.7.1 AN ILLUSTRATIVE EXAMPLE WITH VARYING SAMPLE SIZE
Data Generation Procedure 10 simulation datasets are generated in the following procedure, which is inspired by Lei & Candès (2011
19

---

Generate 21.1.. 81000 independently from a truncated standard normal distri- bution on the interval [-10, 10]. Set I; = etz for covariates highly correlated.
= = -,.. 1000, making the
v2
The propensity score e(æ) = [1 + (2,4 (2(1)) + (23) + (((z))))) where $2,1 is the CDF of the beta distribution with shape parameters (2, 4), and & denotes the CDF of the standard normal distribution. This ensures that e(x) E [0.25, 0.5], thereby providing sufficient overlap. Treatment A, is hence generated by Bernoulli distibution with the probablity of success being e(x;), and resampling from the treatment and control groups has been performed for ensuring that the dataset contains balanced samples for treatment group and control group.
For simulation of observed outcome, we consider
y; c(x;) + TA; + * A: - azi, i 112...., n,
503
<(x:) 1 I+ v²
+ 2xg
where r(2:)) = ((1) ((((1))((22)) and f(x) T+exp(-2+0.0) In other words, we set treatment effect T(x,) T Th. We generated the data under the setting т = 3 and = 0.25 with Merain E {800, 1600, 2400, 3200, 4000), neal E (200, 400, 600, 800, 1000), Aetst E (200, 400, 600, 800, 1000).
2
Resulsts For TMLE, we use the ensemble of lasso and XGBoost to estimate the nuisance functions. For X-Learner, we consider the model with Random Forest (X-Learner-RF) and the model with Bayesian Additive Regression Trees result of X-Learner-RR since the performance of X-Learner-RE that of X-Learner-BART on our simulation dalaset. For CFRNet, consider the model with Wasserstein distance (CFRNet-Wass (CFRNet-MMD), and only presents the result of MMM,, following similiar logic. For most of the benchmark models, we refer to the code implementation of (Jiang et al. 2023 The results summarzed in 2
Secarner-BARTT) but only presents the
consistently better than
and the model with Maximum Mean Discrepancy
A.7.2 AN ILLUSTRATIVE EXAMPLE FOR MISSING DATA PROBLEMS
Data Generation Procedure With the following procedure, we simulated 10 datascts for three scenarios: a) complete data, b) missing random (MAR), and c) missing not at random (MNAR). Each dataset consists of 10,000 training samples, 1000 validation samples, and 1000 Lest samples. In our setting, only training set contains missing values.
1. Generate $1, $100 from an auto-regressive process of order 2 with the concentra- tion matrix given by
0.5,
j- = 1, 2.... .9. .99.
0.25,
if - = 2, - =, ,9.. 98. if - -,i-=1
C5
1,
1000
0,
otherwise
2. Generate the binary treatment variable A E [0,1] from a Bernoulli distribution with the success probability given by p(r1, 33, 23) 1/(1 + -s(a (2. where 22g) - taantan22a
s(21, 32, F3, F5) = tanhtannhaa
2xx),
and resampling from the treatment and control groups has been performed for ensuring that the dataset contains equal numbers of treatment and control samples.
3. Generate outcome variable Y by
Y 4 tanntt + +xa) + tanh(2x2 - 22ry) 2AA) 2xa) taanh-22 -2x4 + ms)))
+ 2 tanh( -A A + 2 tanh(tanh
Oze + 1z
where e N(O0 1) and is independent of at's.
20

---

4. For MAR, randomly deleted 10% of the observations in 21 and 2ą missing values. For MNAR, we first generate missing pattern R1 and RĄ which are binary vectors with 1 representing observed and 0 representing missing, from Bernoulli distribution with the success probability given by 1/(1 + 0- (210)) and Pa = 1/(1 + e--ee -a( (="( m100), where
$1(21, ... T100) = 4 2A +
22j-1, where j -,,.. ,...,,000 50
and (4( =100 4 2-4--4+ + (-0.1) where j 1, ,5, 50
then delete the observations based Rı and R4- The missing rate for MNAR scenario is roughly 10%.
Definition of the False Selection Rate (FSR) and Negative Selection Rate (NSR) In this paper, the FSR and NSR are defined based on 10 datasets:
FSR Σ S \
> 1S $.1 NSR =
E S
Σ 5
where S is the set of true variables, S the set of selected variables for dataset i, and [S|]] is the size of S,.
Results For this experiment, we assumee that already know the neighborhood structure of the Gaussian graphical model that represents the correlations between covariates, rendering the distribution x mmmmis Xobs) be correctly modcled. For case where the structure unknown, estimating the structure is necessary before rumning the Causal-StoNel. The covariate selection accuracy and Out-of-Sample MAE of the estimated ATE are summarized in Table A1 compared with the result for complete dataset.. The result shows that even with missing values present in the dataset, Causal-StoNeet can still correctly identify the true covariates for the outcome and propensity score models and achieve reasonably accurate estimates for the ATE.
Table A1: Covariate selection accuracy and MAE of the ATE estimates for the illustrative
example of missing data.
Out-of Sample MAE
FSRY
NSRY
FSRA
NSRA
Complete MAR 0.0797(0.0134) MNAR (..68)))
0.0103(0.0024)
0
0
0
0
0 0
A.7.3 TWINS DATA
Data Preproccssing In data preproccssing, we focused on the samc-sex twin pairs with born-wiights less than 2 kg and assigned treatments according to t,|2,, 2 Bernoull( (" x wn (z/100 0.1))). Here, (-)) is sigmoid, w. NO0,0.1 I), wx ~ (5, 0.1), # represents the covariate gestat10 (gestation weeks before birth), and æ encompasses the other 45 covariates.
A.7.4 TCGA-BRCA DATA
The Breast Cancer dataset from the TCGA database collects clinical data and gene ex- pression data. for breast cancer patients. The gene expression data contains the expression measurements of 20531 genes. Our goal is to investigate the causal effect radiation therapy paticnts' vital status, while accounting for relevant clinical features and genetic covariates.
The treatment variable A binary (1 for radiation therapy, otherwise), and the outcome variable Y is also binary (1 for death, 0 otherwise). For genetic covariates, applied the sure independence variable screening method (Cui et al. 2015) to reduce the number of genes from 20531 to 100. After preprocessing, the dataset. contains 110 covariates with
21

---

Table A2: Covariates selected for the treatment and outcome models by the Causal-StoNet in a three-fold cross-validation experiment for the twins data, where the number in the parentheses indicates the times that the covariate was elected.
[Id birattnd()), mager8(3), ormoth(3), mrace(3)
m3) dmar(1); mpre5(3), adequacy(3), frace(3), birmon(3) g0stat10(3), csex(1), incervix(2), cigar6(3), crace(3), data _ear__)
Treatment Model
(3), df) foduc6(3), dlivord_min(3), dtotord_ mmn(33,
brstate reg(3), stocefipb reg(3), mplbir reg(3), bord(3)
p3d (3), mager8(3), ormoth(3), m3ace(3)
(3) d2nr((2), mpre5(3), adequacy (3), [race(3), birmon(3), 10s3, osex (3), hydra(2), incervix(3), cigar6(3), crace(3), data year(3)
Outcome Model
(3)evistq()), (3) feduc6(3), dlivord min(3), dtotord min
brstate _eatt,, stoccfipb (tocclipb_reg(3 mplbir_reg(3)) bord(3)
Table A3: ATE estimates by different methods for BRCA-TCGA data
Causal-StoNet.
DSE
ARBE TML.F(asss) TML.(ememmble)
DONUT
-0.1243 (0.0058) -0.0378(0.0243)
-0.1497 (0.1021) -0.0509 (0.0172) -0.0468 (0.0153) -0.0129 (0.0088)
sample size of 845. Table A3 reports the ATE estimates by different methods, with the Causal-StoNel estimate having a smaller standard error. Table A4 summarizes the most frequently selected covariates for both the treat and oulcome models across the three cross- valisaviiiationn folds. Our method identified pathology NN stage and gene ALDH3A2 as significant variables. It is exciting to note that clinical N stage and the tumor expression of ALDH3A2 have been reported by Xie et al. (2022) as potential markers for predicting tumor recurrence in breast cancer patients who achieve pathologic complcte response after neoadjuvant chemotherapy.
Table A4: Most frequently selected covariates for both the treatment and outcome models by the Causal-StoNet in a three-fold cross-validation experiment for the -RGAA data
Clinical
years to birth, date of initial pathologic diagnosis, number of lymph nodes, pathology N stage"
Features
Genetic ACMSD, AKAP2, ALDH3A2I, ALGI, ALMS1P, ANKLE1, ANKRD54, AQP7P1, AZII, B4GALTI, B4GALT2, BCL6B, C10orf47, C10orf72
Features
A.8 PARAMETERS SETTING USED IN THE EXPERIMENTS
A.8.1 CAUSAL-STONET
Illustrated Example with Varying Sample Size The network has hidden layers with structure 32-16-8-1, where Tanh is used as the activation function. The treatment variable is set at the second hidden node of the second hidden layer. Variances of latent variables are of = 10-3 o = 10-5, o = 10-7, oh 10-9. For the mixture Gaussian prior, of = 10-5 of 0.01, and X, = 10 10-6 The epochs for pre-training, training, and refining after sparsification are 50, 200, and 200, respectively.
For SGHMC imputation, tuMe = 1,
0.1. Initial imputation learning rate is set as é1 = 3 x 10-3, <2 3 x 10-4, , 63 5 x 10-7, and decays at training stage and refining stage. For epoch l in these stages, the imputation learning rate is Ek. I
For parameter update, Initial learning rate for training stage are 71 = 10 10-3, 722 = 10 1 10 74 5x 13 and the initial learning rate for refining stage after the sparsification are 71 = 10 722 = 10 -7 73- = 10 9 74 = x 10-14, Learning rate decays for training and refining stage, and for epoch A in these stages, the learning rate is 78, I+
22

---

Illustrated Example with Missing Value The network has 3 hidden layers with structure 8-6-5-1, where Tanh is used as the activation function. The treatment variable is set at the second hidden node of the second hidden layer. Variances of latent variables are oh = 10-3, oh = 10-5, 2h = 10-7, oh = 10-9. For the mixture Gaussian prior, of 3 x 10-3, of = 0.3, and A = 10 The epochs for pro-training, training, and refining after sparsification are 100, 1500, and 200, respectively.
For SGHMC imputation, tHNC 1. = 0.1. Initial imputation learning rate for the network is set as 61 = 10 62 = x 10-1,
2 63 = 10-6, and imputation learning rate for missing covariates imputation miis = 3 x 10-1. Both imputation learning rates decay at training stage and refining stage. For epoch &t in these stages, the imputation learning rate is ER, T+cikki
For parameter update, Initial learning rate for training stage are 71 10-3, 72 = 33x 10 7s = 10-7 74 = =00 10-12, 12 and the initial learning rate for refining stage after the sparsification are 71 = 10 - 122 = 10-7, 7/ 10- 74 10-13, Learning rate decays for training and refining stage, and for epoch k in these stages, the learning rate is 7h,i I+7
%
ACIC The network has 2 hidden layers with structure 200-64-32, where Tanh is used as the activation function. The treatment variable is set the sccond hidden node of the first hidden layer. Variances of latent variables are of = 10-5, o 10-7, oř, 10-9. For the mixture Gaussian prior, of = 2 x 10-4, of = 0.1, and An 10 6, The number of epochs for pre-training is 100. The parameters were pruned twice, first after 500 training epochs, and second after 1000 training epochs. The first pruning be deemed as another initialization for the model. After the second pruning, the model was trained for another 200 epochs for refining the network parameters.
For SGHMC imputation, LuvC = 1, a = 1. Initial imputation learning rate is set as 6 3 x 10-1, <2 10-6, and decays at training stage and refining stage. For epoch k in these stages, the imputation learning rate is EK I+e:x
For parameter update, Initial learning rate for training stage are 712 = 5 10 72 5x 10 Y = 5 x 10-13 13 and the initial learning rate for refining stage after the sparsification are ' = 5 10-7, 72 5 10-9, 7 10-13 Learning rate decays for training and refining stage, and for epoch l in these stages, the learning rate is Th = 147xx
%
Twins The network has 3 hidden layers with structure 46-16-8-4, where Tanh is used as the activation function. The treatment variable is set at the second hidden node of the second hidden layer. Variances of latent variables are o2 = 10-3, o² = 10-4 oh = 10-5, o² 10 6 For the mixture Ganssian prior, oî = x 10-4, of = 0.1, and An = 10 The epochs for pre-training, training, and refining after sparsification are 100, 1500, and 200, respectively.
For SGHMC imputation, tuMe 1, = 1. Initial imputation learning rate set as 61 x x10-2, c = 9 x 10-3, = 9 x 10-5, and decays at training stage and refining stage. For epoch in these stages, the imputation learning rate is fx = I++,x
For parameter update, Initial learning rate for training stage are
10-3, = 10-4,
a = 10-5 74 1 and the initial learning rate for refining stage after the sparsification are = 10-4, ?/2 = 10 5, 73 = 10-6, 10-10, Learning rate decays for training and refining stage, and for epoch in these stages, the learning rate 7, 7
1+TExR1
TCGA-BRCA The network has 3 hidden layers with structure 110-64-16-4, where Tanh is used as the activation function. The treatment variable is set at the second hidden node of the second hidden layer. Variances of latent variables are o
10 3 oh 10-5 on. = 10-7 2 = 10-9. For the mixture Gaussian prior, of = 1 of = 0.01, and A, = 10 The epochs for pre-training, training, and refining after sparsification are 100, 1500, and 200, respectively.
23

---

For SGHMC imputation, taaC 1, 1. Initial imputation learning rate is set as 6 = 10-3, 3, 62 10 <3 = 10 6, and decays at training stage and refining stage. For epoch à in these stages, the imputation learning rate is Ck I+
For parameter update, Initial learning rale for training stage are 71 10 3 72 = 5 10-5, 5 73 = 5 74 10-12 and the initial learning rate for refining stage after the sparsification are 71 = 10 7/2 = 10 73 = 10 8, 71 10010, 13, Learning rate decays for training and refining stage, and for epoch k in these stages, the learning rate is k,
1
24


Reviews:
[{"review_id": "8KfgSfX583", "reviewer": "Reviewer_7o25", "summary": "This paper introduces an approach to causal inference tailored for high-dimensional complex data. This approach draws upon recent developments in deep learning techniques, including sparse deep learning and stochastic neural networks. By leveraging these techniques, the proposed approach effectively tackles both the high dimensionality and the complexity of the underlying data generation process. Moreover, it can handle scenarios with missing values in the datasets.", "questions": "Please see section weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The proposed method Causal-StoNet is proven to have the universal approximation ability, making it a versatile tool for modeling outcome and propensity score functions.\n\n- The paper provides a strong theoretical foundation for its approach, including proofs and mathematical support, enhancing its reliability.", "weaknesses": "1/ It seems to me that the proposed method is a direct application of stochastic neural networks. Please clearly highlight technical innovation of the proposed method.\n\n2/ Can the authors explain which component of their method make it works well for high-dimensional confounder X? Why the existing methods such as: CEVAE, CFR Net are not possible to deal with high dimensions of X?\n\n3/ It seems to me that many concepts and notations in the paper are unexplained. For example, why do we choose $\\sigma_{0,n}^n$ to be a very small number while $\\sigma_{1,n}^2$ is relatively large? Is there any rationale for this? What is $\\pi(\\theta)$ in Eq.~9? Is it the prior?\n\n4/ Since $Y_{mis}^i$ is unobserved, how do you minimise Eq.~9?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Qd5B2EP4fp", "reviewer": "Reviewer_GhoU", "summary": "This paper proposes a deep-learning-based causal inference method, that assumes the nuisance parameters are nonlinear high-dimensional models, fit by a stochastic neural network with sparsity-pursuing properties. This paper also gives theoretical guarantees, building upon a series of papers by Liang and colleagues on Bayesian neural nets.", "questions": "1. Do the theoretical results rely on equation (6)? If so, then the theoretical results seem to go against the conventional wisdom in the deep learning literature, that is, the neural parameters themselves are not scientifically meaningful so it is not that important to learn the neural net parameters.\n\n2. In theory, the propensity model and outcome model are both sparse models with sparsity levels lower than n^{3/16}. In linear models, the sparsity allowed is n^{1 / 2} up to log factors. Could the authors provide a heuristic explanation on the rate n^{3 / 16}?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper tries to address an important problem in causal inference, by taking the nuisance parameters to be high-dimensional sparse nonlinear models. State-of-the-art deep learning methods are used, e.g. stochastic/Bayesian sparsity-pursuing neural nets, to address this challenge.\n\nThe theoretical results look sound, building upon a series of earlier works on statistical guarantees of Bayesian neural nets by Liang and colleagues.", "weaknesses": "1. Even though I selected ``good'' in Presentation, I believe the exposition can nonetheless be significantly improved. For example, some of the key assumptions and elements are delayed to Appendix so the flow seems a bit broken, in particular in the first several sections. I strongly recommend that the authors consider revising their manuscript to make the flow smoother.\n\n2. In the simulation, this paper considered scenarios with p=100/200 and n=10000, which do not seem to be a very high-dimensional regime. What would be the performance if we further increase the covariate dimension p or decrease the sample size n?\n\n3. Missing references: doubly-robustness should be traced back to Robins et al. 1994 JASA. Also, Farrell et al. was cited twice in the paper (one arxiv version, one joe published version). A recent paper led by Xiaohong Chen and colleagues (Chen, Liu, Ma, Zhang, to appear in JoE) also addresses a similar problem, though their nuisance models are slightly different (Barron space). This paper should also be cited.\n\n4. Theory-practice gap: As we know, almost all theoretical works on deep learning do not reflect practice. For instance, sparse neural nets are generally difficult to fit, as persuasively argued in Farrell, Liang, and Misra ECTA 2021. The authors are recommended to comment on this, for readers to better understand what are the key elements that allow Causal-Stonet to learn the sparse neural nets.\n\n5. In causal inference, often times $x$ has clear scientific meaning. For sparse models, I would imagine that the input layer sparsity seems to be more important. If putting sparsity in output layer, this is essentially saying that there is some sparse nonlinear representation of the input. By assuming sparse neural nets, however, this is saying that the nonlinear representation itself is also in some sense sparse. Is it really aligned with our view of the real world or is it more like a contrived modeling assumption? I hope that the authors could further comment on the modeling philosophy adopted in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KcNygeAI5k", "reviewer": "Reviewer_vAHH", "summary": "The authors propose an algorithm, CausalStoNet, which aims to give accurate causal inferences for systems of up to 100 variables (after variable reduction) under fairly relaxed nonlinearity conditions. The method is based on deep learning ideas.", "questions": "1.\tI guess, in my mind, it does very little good to say a method can deal with high dimensions without saying how high. There are two examples, one for 43 variables and another for 100 (genome reduced to this). Some may not consider this to be high-dimensional, so it would be better to say up front what dimension one hopes to achieve with the method (after variable reduction is done).\n2.\tAlso, it’s important to say what density such models can attain. It’s possible with some very high-dimensional data that the models may, in fact, be very dense, a situation that can be addressed currently in the linear Gaussian or non-Gaussian case. However, this hasn’t been addressed to my knowledge for models with more general connection functions, a possible advantage of the method in this paper. \n3.\tThere are recent papers that address a dense searches for the linear, Gaussian case and linear, non-Gaussian cases. The secret of some of these papers is to relax the Faithfulness condition, something which does allow some nonlinear functions to be addressed in a linear framework. The secret of others (like DirectLiNGAM) is to move to the linear, non-Gaussian regime. It would be really wonderful if an approach like the one in this paper could be shown to improve these algorithms (which look to be state-of-the-art) for some choices of nonlinear functions.\n4.\tAlong these lines, I think it’s important when saying you’re outperforming methods to include in this the relaxation of assumptions that you’ve done and to compare only to other methods that relax assumptions in similar ways or that use stronger assumptions, explicitly noting this and show that with the stronger assumptions worse results are achieved. (The latter is not always the case.)\n5.\tThe restriction to binary data for some of the theory is somewhat severe since very few real datasets consist entirely of binary variables, though perhaps I misunderstand.\n6.\tI’m not sure how “MNR” abbreviates “Missing At Random”; I think it must be “MAR.”\n\n\nNOTE: I read the response; I thought it was convincing, so I'm raising my rating. Thanks.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This was a good paper, I thought. The applications were fruitful, and the explanation of the theory was clear. It was well-situated in the literature, and useful and correct extensions of the method were proposed.", "weaknesses": "Several methods are compared, but I don't see a discussion of looking for the best available methods. Are these methods in Table 2 the best available methods? See also my questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)