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

## W1: 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? For this paper, we mainly consider the estimation of Average Treatment Effect (ATE). For estimation of ATE, there are numerous well-adopted methods developed based on different modeling techniques. In our paper, we consider the neural-network based methods (i.e. DragonNet, X-Learner) and methods in semiparametric literature (e.g., Double Selection Estimator (DSE), Approximate Balancing Residual Estimator (ARBE), and TMLE) as baselines. Since the proposed method is based on neural networks, comparing with other neural-network based methods is essential. For semiparametric methods, DSE and ARBE are state-of-the-art methods for high-dimensional causal inference, and TMLE can be easily extended to machine learning methods in terms of nuisance parameters estimations. We summarize the performance of these methods in simulation studies and ACIC dataset in Table 1. The methods listed in Table 2 are state-of-the art methods for estimation of Conditional Average Treatment Effect (CATE), most of which are also neural network based. Although CATE estimation is not our main focus, we still listed the comparison of CATE estimation just to highlight the existing neural network-based methods on CATE estimation.

## Q1: It would be better to say up front what dimension one hopes to achieve with the method (after variable reduction is done) In the rebuttal, we have tried a high-dimensional example with $p=1000$ and $n=800, 1500, 3000$. Please see the results in the global response. The proposed method worked for all cases and outperformed the baselines. ## Q2: Also, it’s important to say what density such models can attain. We added a simulated experiment for a function with nonlinear outcome function and nonlinear treatment effect function, for which the neural network can be dense as the true functions are not exact neural network functions. Please see the results in the global response. The proposed method worked for all cases and outperformed the baselines. ## Q3: Connection with causal graph discovery This is a thoughtful question. Indeed, the sparse DNN has been used in the literature for learning causal graphs via a double regression method, see Liang and Liang (2022)[1]. The method proposed in this paper can be equally applied there. [1] Liang, S. and Liang, F. (2022) A Double Regression Method for Graphical Modeling of High-dimensional Nonlinear and Non-Gaussian Data. Statistics and Its Interface, in press (see also arXiv:2212.04585 ## Q4: We will follow this suggestion to state the performance of the algorithm in revising the paper. ## Q5 & Q6: We are sorry for the misunderstanding. In this paper, we just assume the treatment variable is binary, while leaving the outcome and explanatory variables unrestricted. Furthermore, we have discussed the extensions of the binary treatment variable to multi-level or continuous treatment variables.

## W1: I believe the exposition can nonetheless be significantly improved. We will follow your suggestion to move the key assumptions and elements into the main text, ensuring a smooth flow. ## W2: What would be the performance if we further increase the covariate dimension p or decrease the sample size n? Please refer to the experiment results in the global response, where we considered the cases with $p=1000$ and $n=800$, 1500 and 3000. The proposed method worked for all cases and outperformed the baselines. ## W3: Missing references The references will be updated as suggested: Cite Robins et al. 1994 [1] for doubly-robustness, remove the arXiv paper Farrell et al. 2018, and add Chen et al. 2024 [2]. [1] Robins, J.M., Rotnitzky, A., and Zhao, L.P. (1994) Estimation of regression coefficients when some regressors are not always observed. J. Am. Stat. Assoc., 89, 846-866. [2] Chen, X., Liu, Y., Ma, S., and Zhang, Z. (2024) Causal inference of general treatment effects using neural networks with a diverging number of confounders, Journal of Econometrics, 238, 105555 ## W4: Theory-practice gap Tricks for efficiently training sparse DNNs have been discussed and developed in Sun et al. [1,2]. Specifically, Sun et al. [1] recommended starting with an overparameterized DNN, while Sun et al. [2] further proposed a prior-annealing method. This method involves gradually incorporating the prior onto an overparameterized DNN in an annealing manner and is potentially immune to local traps. In this paper, we adopt the approach suggested by Sun et al. [1], training the sparse StoNet by initially employing an overparameterized model. [1] Y. Sun, Q. Song, and F. Liang. Consistent sparse deep learning: Theory and computation. Journal of the American Statistical Association, 117(540):1981–1995, 2022. [2] Y. Sun, W. Xiong, and F. Liang. Sparse deep learning: A new framework immune to local traps and miscalibration. NeurIPS 2021, 2021. ## W5: Modeling philosophy Thank you for your thoughtful question. We address this query from two perspectives. In biological neural networks, as discussed in [1], a wealth of empirical evidence suggests that neurons represent information using sparse distributed patterns of activity. Concerning artificial DNNs, as shown in [2] based on the random matrix theory, the stochastic gradient-based training process results in an implicit heavy-tailed self-regularization in their weights, i.e., achieving an essentially sparse representation even without a sparse penalty used during the training process. The proposed work makes this sparse pattern of connection weights more explicit and validates the downstream inference. [1] S. Ahmad and J. Hawkins (2016) How do neurons operate on sparse distributed representations? A mathematical theory of sparsity, neurons and active dendrites, ArXiv:1601.00720. [2] C.H. Martin and M.W. Mahoney (2021). Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning, J. Mach. Learn. Res., 22, 165:1-73.

## Q1: Identifiability of neural network parameters. Indeed, it is known that the DNN model is generally nonidentifiable due to the symmetry of the network structure. For example, the approximation of the neural networks can be invariant if one permutes the orders of certain hidden nodes, simultaneously changes the signs of certain weights and biases if tanh is used as the activation function, or re-scales certain weights and bias if Relu is used as the activation function. In this paper, following Sun et al. (2022) [1], we consider a set of networks, denoted by $\boldsymbol{\Omega}$, for which each element can be viewed as a class of equivalent DNN models. Then we define an operator $\nu(\boldsymbol{\gamma},\boldsymbol{\theta}) \in \boldsymbol{\Omega}$ that maps any neural network to $\boldsymbol{\Omega}$ via appropriate transformations such as nodes permutation, sign changes, weight rescaling. To define $\boldsymbol{\Omega}$ usually comes with constraints on the neural network, and Liang et al. (2018) [2] gives an example of such constraints. However, even if we relax this identifiability assumption on network parameters, the validity of StoNet as an approximator to DNN can still be justified, since StoNet and DNN have asymptoticlaly equivalent loss functions (see equation (5)). [1] Y. Sun, Q. Song, and F. Liang. Consistent sparse deep learning: Theory and computation. Journal of the American Statistical Association, 117(540):1981–1995, 2022. [2] Liang, F., Li, Q., and Zhou, L. (2018), “Bayesian neural networks for selection of drug sensitive genes,” Journal of the American Statistical Association, 113, 955–972. ## Q2: Could the authors provide a heuristic explanation on the rate $O(n^{3/16})$? Again, this is a thoughtful question. A heuristic explanation is that the size of the true DNN model is allowed to increase with the sample size $n$ at the rate $O(n^{3/16})$, while ensuring the validity of inference for the treatment effect. We note that this rate is derived based on the Laplace approximation error for high-dimensional Bayesian sparse nonlinear models, as presented in Theorem 2.3 by Sun et al. 2022 [1]. Specifically, Sun et al. (2023) showed that the Laplace approximation error is of $O(r_n^4/n)$, a general result for the MAP-based posterior mean approximation in Bayesian sparse nonlinear models. We believe that this rate can be improved under the frequentist framework, but the proof appears to be challenging. Refer to our discussions about Farrell et al. (2021) [2] in the ``Related Works'' paragraph for further insights. [1] Sun, Q. Song, and F. Liang. Consistent sparse deep learning: Theory and computation. Journal of the American Statistical Association, 117(540):1981–1995, 2022. [2] M. Farrell, Tengyuan Liang, and S. Misra. Deep neural networks for estimation and inference. Econometrica, 89:181–213, 2021.

## W1: 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. The technical innovations of this paper lie in two aspects. First, we propose using the stochastic neural network as a general tool for modeling natural systems by incorporating middle-level observations as visible units in some hidden layers. While we acknowledge that the stochastic neural network has been proposed in the literature, our utilization of it is innovative. This approach enables the modeling of many complex systems by deep neural networks (DNNs) in a natural way, facilitating downstream inference. Second, to ensure that Causal-StoNet can handle high-dimensional data and that the resulting statistical inference is valid, we impose a mixture Gaussian prior on each connection of the Causal-StoNet (see Equation (8)). Under some regularity conditions, we establish consistency for structure selection, outcome function estimation, and propensity score function estimation. Consequently, we establish the validity of statistical inference for the treatment effect. Without the sparsity prior, statistical inference with Causal-StoNet cannot be valid. Here, we would like to emphasize that for many neural network-based methods, justifying the validity of downstream inference can be challenging unless the dataset size is sufficiently large and the network size is relatively small. On the other hand, we are aware that deep neural networks trained with stochastic gradient-based methods possess the property of implicit heavy-tailed self-regularization (see Ref. [1]). However, since the self-regularization phenomenon cannot be rigorously measured, at least for now, establishing the validity of downstream inference can still be challenging. Our method provides a theoretical guarantee for the validity of the neural network-based causal inference. [1] C.H. Martin and M.W. Mahoney (2021). Implicit Self-Regularization in Deep Neural Networks: Evidence from Random Matrix Theory and Implications for Learning, J. Mach. Learn. Res., 22, 165:1-73. ## W2: 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? As mentioned above, the sparsity prior enables Causal-StoNet to perform effectively with high-dimensional confounders. Moreover, the resulting sparse structure ensures consistency of the estimation and the validity of the subsequent causal inference. In contrast, both CEVAE and CFRNet utilize fully-connected neural networks, where the consistency of estimation cannot be ensured. Consequently, validating the subsequent causal inference becomes challenging. For comparison, we have applied CEVAE and CFRNet to the ACIC example, see Table 2. The comparison indicates that both methods are inferior to the proposed one for the example. References: Shalit, Uri, Fredrik D. Johansson, and David Sontag. “Estimating individual treatment effect: generalization bounds and algorithms." International Conference on Machine Learning. PMLR, 2017. Johansson, Fredrik, Uri Shalit, and David Sontag. “Learning representations for counterfactual inference.” International conference on machine learning. PMLR, 2016. ## W3: 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}^n$ is relatively large? Is there any rationale for this? What is $\pi(\theta)$ in Eq.~9? Is it the prior? We will double-check and clarify the notations in the revision. The restriction on the noise is needed to ensure that the Causal StoNet can have asymptotically equivalent loss function as DNN (see equation (5)), which is crucial in justifying that the Causal-SoNet is a good approximator of the underlying DNN and hence a valid universal learner for the treatment model and the outcome model. The $\pi(\theta)$ in equation (9) refers to the mixture Gaussian prior. ## W4: Since is $\boldsymbol{Y_{\text{miss}}}^i$ unobserved, how do you minimise Eq.~9? Equation (9) can be minimized by solving a mean-field equation, as given in (10), using an adaptive SGMCMC algorithm. The Adaptive SGMCMC algorithm operates by iterating over the following two steps: (1) Sampling $\boldsymbol{Y_{\text{miss}}}^i$ from $\pi(\boldsymbol{Y_{\text{miss}}}^i | \boldsymbol{X}, A, \boldsymbol{Y}, \boldsymbol{\theta})$; (2) update $\boldsymbol{\theta}$ based on the sampled latent variables. For more details of the algorithm, please refer to section 3.2. Additionally, Appendix A1 provides the pseudo-code for the training algorithm of Causal-StoNet, and Appendix A6 includes a brief discussion of the convergence analysis of the algorithm.