Paper: Causal-StoNet: Causal Inference for High-Dimensional Complex Data (ICLR-2024)

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Experiments):
There are recent papers that address 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.

Author response:
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

Reviewer comment (Evaluation):
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 showing that with the stronger assumptions worse results are achieved. (The latter is not always the case.)

Author response:
We will follow this suggestion to state the performance of the algorithm in revising the paper.

Reviewer comment (Theory):
The 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.

Author response:
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.

Reviewer comment (Presentation):
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.

Author response:
We will follow your suggestion to move the key assumptions and elements into the main text, ensuring a smooth flow.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Theory):
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}?

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Writing):
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?

Author response:
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.