Paper: Statistical Inference for Deep Learning via Stochastic Modeling (ICLR-2024)

Reviewer comment (Novelty):
The main issues of this paper including **novelty and soundness of the results**.

Author response:
The work by Liang et al. (2022)[2] discussed the application of the StoNet in nonlinear sufficient dimension reduction. However, the major novelty of this paper are (i) to develop a sparse StoNet framework, showing that it can serve as a bridge for transferring the theory and methods developed for linear models to deep learning models; (ii) to develop a post-StoNet framework, showing that it provides a general method for faithfully quantifying the prediction uncertainty of large-scale deep learning models. 

Furthermore, we have provided rigorous theoretical justifications for the proposed frameworks and illustrated their performance using numerical examples.

Reviewer comment (Presentation):
In Figure 2, there is slight difference with regard to overall distribution between StoNet and DNN. The variance of StoNet is larger than DNN, which seems like that vanilla DNN performs better than StoNet. The result is confusing.

Author response:
In Lemma 1, it is shown that the StoNet and DNN are asymptotically equivalent when the noise $\sigma_i$ at each layer goes to 0. The goal of our theoretical results (Theorem 1 and Corollary 1) is to use StoNet as a tool to provide theoretical justification for the consistency of penalized DNN models. Therefore,  for the simulation example of Section 5, it is expected that the DNN and StoNet will produce similar results. Since StoNet introduces more noise during training, it is also expected that the lines in Figure 2(a) have a larger variability than those in Figure 2(b). 
Figure 2 matches with our theory.

Reviewer comment (Presentation):
Figure 3 lacks legend, what does lines of different color represent?

Author response:
In this figure, each line/color corresponds to a different feature, 
which shows how the gradient of the feature changes with the regularization parameter $\lambda$. 
Our major goal is to use this figure to illustrate that
the penalized DNN model can identify important features.

Reviewer comment (Evaluation):
In Table 2, the ACC results (mean and std.) of 'No Post Calibration' method and 'Temp. Scaling' method are exactly the same, which is counter-intuitive and is of low probability, what is the reason.

Author response:
The temperature scaling method (Guo et al. 2017 [1]) divides the pre-softmax output of the DNN model by a temperature, and then learns the optimal temperature in a separate calibration set. To be more precise, 
let $(z_1, \dots, z_K)$ be the neural network output for an $K$-class classification problem, and the temperature scaling method will produce $softmax(z_1 / T, z_2 / T, \dots, z_n / T)$ as the predicting probability vector. Since rescaling the values of $z_i$'s doesn't change their order, 
the prediction label 
$$
\arg\max  softmax(z_1 / T, z_2 / T, \dots, z_n / T),
$$
will be the same as the unscaled one. Therefore, temperature scaling won't change prediction labels, and  its prediction accuracy is exactly the same as the original model.

Reviewer comment (Experiments):
As for result in Table 3, the compared method (Vovk et al., 2005) is proposed too long ago. To be more convincing, further ablation study should be done, for example, compare with vanilla StoNet or other related works.

Author response:
The conformal prediction method has become popular in recent years as a general tool to apply on top of almost any ML model to provide prediction regions with marginal coverage guarantee. The quality of intervals depends on the performance of the underlying ML model and the choice of non-conformity score. Our intention is to emphasize that when the underling ML model is overfitting, our Post-StoNet method can mitigate the overfitting issue and provide shorter intervals with the correct coverage rate. In the revision, we will add another baseline Comformalized Quantile Regression[1] and vanilla StoNet without penalty. Please refer to our rebuttal to reviewer HASS(R2C3-1) for the results.

Reviewer comment (Evaluation):
In Table 1 and 2 of Liang et al. (2022), results of massive datasets and methods are listed, which is not discussed in this paper.

Author response:
The work by Liang et al. (2022)[2] discussed the application of StoNet in nonlinear sufficient dimension reduction. However, the major goals of this paper are (i) to develop a sparse StoNet framework, showing that it can serve as a bridge for transferring the theory and methods developed for linear models to deep learning models; (ii) to develop a post-StoNet framework, showing that it provides a general method for faithfully quantifying the prediction uncertainty of large-scale deep learning models. Since the goals of the two papers are very different, we did not work on the same examples as those in Liang et al. (2022)[2].

Reviewer comment (Theory):
Some assumptions (already in StoNet) seem to be not realistic (see C1).

Author response:
In assumption A2, we essentially group the parameters that define equivalent neural network models into an equivalent class, and the assumption is actually imposed on the equivalent classes of parameters. In practice, the neural network models are usually trained with locally updating algorithms, and it is unlikely to visit all equivalent models. Therefore, practically, assumption A2 is more like a local assumption.

Reviewer comment (Theory):
It's unclear how the algorithm actually enforces sparsity (see C2).

Author response:
In practice, solving a model with LASSO penalty by a gradient descent type algorithm usually employs sub-gradients, where the lasso penalty term gives 0 gradients for the parameters at 0. This use of the sub-gradient in solving Lasso optimization problems has been justified in Nguyen and Yin (2021) [1]. 

We agree with you that the gradient descent or stochastic gradient descent will not lead to exact 0 parameters. To enforce sparsity in practice, we can put a threshold such as $\lambda$ on the learned parameters. Theoretically, to recover the sparse structure with LASSO penalty, 
one can put a $\beta$-min assumption. For example, in Liu and Yu (2013) [2], 
the authors assumed that 
$$
n^{\frac{1-c_3}{2}} \min_{1\leq i \leq s }|\beta_i^*| \geq M, \quad \lambda_n \propto n^{\frac{c_4 - 1}{2}},
$$
for some constants $M, c_1, c_2, c_3, c_4$ such that $c_1 + c_2 < c_3 \leq 1$, $c_2 < c_4 < c_3 - c_1$. Therefore, using $\lambda$ as a threshold will not miss true connections. 
In practice, similar to Figure 2, with an appropriate choice of $\lambda$, there usually exists a clear threshold to separate true variables and other variables.


### Reference

[1] Nguyen, N.N. and Yin, G. (2021). Stochastic Approximation with Discontinuous Dynamics,
Differential Inclusions, and Applications. 
arXiv:2108.12652v1.

 [2] Hanzhong Liu and Bin Yu (2013). Asymptotic properties of Lasso+mLS and Lasso+Ridge in sparse high-dimensional linear regression, Electronic Journal of Statistics, 7, 
 3124 - 3169.

Reviewer comment (Theory):
Theoretical guarantees for selection and subsequent inference are not perfectly clear (see Q1, Q2).

Author response:
Our results indeed rely on some assumptions on the eigenstructure of the data matrix, specifically, the  $m$-sparse minimum and maximum eigenvalues of the covariance matrix of the hidden variables at each hidden layer. Refer to assumption A3 and Lemma A1 (in the appendix) for the detail. "almost any training data" in the paper refers to the probability of the observed training data set $D_n$ from the data distribution. It doesn't mean that we don't need assumptions on the data distribution.

This is a thoughtful question. If we aim to make inference for the model parameters, a post-selection inference approach is indeed needed. However, in this paper, we aim to make inference for the uncertainty of the prediction. By Theorem 3 of Liu and Yu (2013) [1](select the model then refit selected parameters), the standard post-selection inference procedure (for the model parameters) is not needed, and the proposed inference procedure is valid. We can refit the parameter for the selected structure This point will be mentioned in the revision. 

### Reference
[1] Liu, Hanzhong, and Bin Yu. "Asymptotic properties of Lasso+ mLS and Lasso+ Ridge in sparse high-dimensional linear regression." (2013): 3124-3169.

Reviewer comment (Novelty):
It often does not get clear what the additional contribution is when compared to StoNet; see also first point on experiments in C3.

Author response:
As one of the main contributions of this paper, we proposed the post-StoNet framework to quantify the prediction uncertainty for large-scale deep learning models, where the StoNet is used as a post hoc calibration model. 
In this framework, we don't need to change the standard DNN training procedure. A StoNet model can be applied on top of the learned representation of the last hidden layer of a "well-trained deep learning model''. This framework is simple while working better than the popular conformal method, especially when the "well-trained deep learning model'' is overfitted. As justified in the paper, the proposed post-StoNet method is able to correctly quantify the prediction uncertainty for the "well-trained deep learning model''.

In the revision, we will also add an ablation study to compare the post-StoNet procedure with and without penalties. In the following table, we can see that for the small data set, a StoNet without penalty failed to provide correct coverage rate, demonstrating the importance of the proposed sparse StoNet structure. And per request by Reviewer GYhy, we also include an additional baseline, Conformalized Quantile Regression(CQR) [1]

| Dataset | N | P | Model | Coverage Rate | Interval length |
|----------:|:-------:|:--------:|:-------:|:---------:|:--------:|
 | Wine |  1,599 | 11 | Post-StoNet($\lambda$ = 0.02) | 0.9042(0.0126) |  2.0553(0.0719) |
  | | | | Post-StoNet($\lambda$ = 0) | 0.2175(0.0209) | 0.2037(0.0304) |
  | | | | Split Conformal | 0.8958(0.0302) | 2.4534(0.1409) |
  | | | | Split CQR | 0.9237(0.0132) | 2.3534(0.0501) |

Reviewer comment (Presentation):
Quality of some graphics is really bad (Fig 2 and 3 almost not possible to read).

Author response:
We will add legends for Figures 2 and 3 in the revision. In Figure 2, we will change to 2 colors for relevant variables(5 lines that are far away from $y = 0$ line in the figure) and noise variables.

Reviewer comment (Theory):
There is no one "true parameter" in deep neural networks given their overparametrization-induced and hidden symmetries, and Assumption A2’s practical validity is questionable. While the authors recognize this and describe uniqueness up to loss-invariant transformations, this restriction does not account for scaling symmetries in ReLU networks or other hidden symmetries.

Author response:
Thank you for your valuable feedback. Here we would like to further clarify that the equivalent classes are defined implicitly, such that all equivalent models (after appropriate transformations) are mapped to a single model within an appropriately defined set of DNN models. 
We also agree with you that such an equivalent class can be data dependent, and the network weights 
and function distributions remain multimodal in nature. 

However, regarding statistical inference for such multimodal models, we want to elaborate by the problem 
of label switching in mixture models (see, for example, [1]). It is almost the same as the problem we are considering for DNNs but simpler. As shown in Liang and Wong (2001)[2], for statistical inference of such a label-switching mixture model, it is enough to conduct inference for one of its component. In particular, Liang and Wong (2001) developed an efficient Monte Carlo algorithm to sample all components of the model and tested the issue of using samples from all components or from a single component. They concluded that statistical inference based on all components of the model is equivalent to that based on a single component.  Based on this result, we state that our theory is developed based on a ``local/single'' component. 

Bearing in mind the label-switching mixture model as a prototype of the DNN model, algorithms that make inference based on a local/single model or a mixture of multiple models are all valid, as long as 
the models (optimized or sampled) are consistent. This is exactly the goal of this paper: we aim to 
find a consistent model to validate the downstream inference for deep learning. 
Note that our post-StoNet procedure is to find a consistent model for the data transformed via a well-trained DNN. 

Finally, we note that in terms of inference, the Laplace approximation method needs to approximate the Hessian matrix, which will be hard to do for large network models. On the other hand, deep ensemble model[3] lacks theoretical guarantee for its consistency in model estimation, and the resulting inference might not be reliable.

### Reference

[1] Stephens, Matthew. "Dealing with label switching in mixture models." Journal of the Royal Statistical Society: Series B (Statistical Methodology) 62.4 (2000): 795-809.

[2] Liang, Faming, and Wing Hung Wong. "Real-parameter evolutionary Monte Carlo with applications to Bayesian mixture models." Journal of the American Statistical Association 96.454 (2001): 653-666.

[3] Lakshminarayanan, Balaji, Alexander Pritzel, and Charles Blundell. "Simple and scalable predictive uncertainty estimation using deep ensembles." Advances in neural information processing systems 30 (2017).

Reviewer comment (Theory):
It is unclear how the model is optimized with a Lasso penalty given that this penalty is non-smooth and stochastic variants hence do not yield exact-zero solutions.

Author response:
Thank you for your valuable feedback. To address this question, let's clarify the structure of our theory. 

First, the proof for Theorem 1, i.e., Lasso theory transfer from linear models to StoNet, is based on the theoretical framework provided by the imputation-regularized optimization (IRO)
algorithm [1](Algorithm 1 in the appendix, it is a stochastic EM type algorithm). 
In the proof, the regularized regression problem is assumed to be exactly 
solved at each iteration.
Given the imputed value for the hidden units, it corresponds to solving a series of linear regression with LASSO penalty,
which is actually doable. 
Next, we pointed out that the objective function involved in Theorem 1 can also be solved using an adaptive stochastic gradient MCMC algorithm. 
Therefore, the proof for Theorem 1 is valid based on the IRO algorithm, which does not involve any issue 
on optimizing a non-smooth penalty with stochastic gradients. 

Regarding the issue on optimizing a non-smooth penalty with SGD,  particularly for the Lasso penalty, 
it has been shown in Nguyen and Yin (2021) [2] that SGD works. Therefore, our use of the adaptive SGMCMC for training the StoNet is valid. 

### Reference

[1] Liang, Faming, et al. "An imputation–regularized optimization algorithm for high dimensional missing data problems and beyond." Journal of the Royal Statistical Society Series B: Statistical Methodology 80.5 (2018): 899-926.

[2] Nguyen, N.N. and Yin, G. (2021). Stochastic Approximation with Discontinuous Dynamics,
Differential Inclusions, and Applications. 
arXiv:2108.12652v1.

Reviewer comment (Experiments):
Experiments mainly present coverage rates and calibration results, and it is unclear how much of this performance comes from StoNet itself (missing ablation study).

Author response:
As one of the main contributions of this paper, we proposed the post-StoNet framework to quantify the prediction uncertainty for large-scale deep learning models, where the StoNet is used as a post hoc calibration model. 
In this framework, we don't need to change the standard DNN training procedure. A StoNet model can be applied on top of the learned representation of the last hidden layer of a "well-trained deep learning model''. This framework is simple while working better than the popular conformal method, especially when the "well-trained deep learning model'' is overfitted. As justified in the paper, the proposed post-StoNet method is able to correctly quantify the prediction uncertainty for the "well-trained deep learning model''.

In the revision, we will also add an ablation study to compare the post-StoNet procedure with and without penalties. In the following table, we can see that for the small data set, a StoNet without penalty failed to provide correct coverage rate, demonstrating the importance of the proposed sparse StoNet structure. And per request by Reviewer GYhy, we also include an additional baseline, Conformalized Quantile Regression(CQR) [1]

| Dataset | N | P | Model | Coverage Rate | Interval length |
|----------:|:-------:|:--------:|:-------:|:---------:|:--------:|
 | Wine |  1,599 | 11 | Post-StoNet($\lambda$ = 0.02) | 0.9042(0.0126) |  2.0553(0.0719) |
  | | | | Post-StoNet($\lambda$ = 0) | 0.2175(0.0209) | 0.2037(0.0304) |
  | | | | Split Conformal | 0.8958(0.0302) | 2.4534(0.1409) |
  | | | | Split CQR | 0.9237(0.0132) | 2.3534(0.0501) |

Reviewer comment (Experiments):
Experiments only contain one simulation study with a fixed setup, which is restrictive.

Author response:
In the rebuttal, we have conducted an additional variable selection 
experiment, which will be added to the paper in the revision. The experiment can be described as follows.

The data are generated from the following high-dimensional nonlinear model ($p=2000$): 
\begin{equation*}
y = \frac{5x_2}{1 + x_1^2} +5\sin (x_3 x_4) + 2x_5 + 0x_6 + \dots 0x_{2000} + \epsilon
\end{equation*}
where $\epsilon \sim N(0,1)$, $x_i = \frac{e + z_i }{\sqrt{2}}$, $e, z_1, \dots, z_{2000} \sim N(0,1)$ are independent. The explanatory $x_i$ are mutually correlated, which makes the problem more challenging. We generate $n=10000$ data for training. We use a StoNet with 1 hidden layer, and 500 hidden units, trained with LASSO penalty $\lambda = 8e-3$. We use $\lambda$ as the threshold to select variables. We repeated the experiment 10 times, the model selects $\\{6,5,5,6,5,5,5,5,5,6\\}$ variables, and all 5 relevant variables are selected. 
Formally, to quantify the selection quality, we can define false selection rate:  $$FSR = \frac{\sum_{i=1}^{10} |\hat{S_i} \setminus S| }{\sum_{i=1}^{10}|\hat{S}_i|}, $$

and negative selection rate: $$NSR = \frac{\sum_{i=1}^{10}|S \setminus \hat{S_i}|}{\sum_{i=1}^{10}|S|},$$  where $S$ is the set of true variables and $\hat{S}_i$ is the set of selected variables in the $i$-th data set. 
Then the corresponding false selection rate is 0.056, and the corresponding negative selection rate is 0.
This example demonstrates that the StoNet model with a LASSO penalty can correctly identify relevant variables for general nonlinear systems.

Reviewer comment (Theory):
The paper proves consistent structure selection without any requirements on the feature matrix (“almost any training data”), whereas such results typically require rather restrictive assumptions even in simpler models; clarification is needed if this condition is overlooked in the Appendix.

Author response:
Our results indeed rely on some assumptions on the eigenstructure of the data matrix, specifically, the  $m$-sparse minimum and maximum eigenvalues of the covariance matrix of the hidden variables at each hidden layer. Refer to assumption A3 and Lemma A1 (in the appendix) for the detail. "almost any training data" in the paper refers to the probability of the observed training data set $D_n$ from the data distribution. It doesn't mean that we don't need assumptions on the data distribution.

Reviewer comment (Reproducibility):
How the model is optimized with a Lasso penalty is questioned again (see C2).

Author response:
Thank you for your valuable feedback. To address this question, let's clarify the structure of our theory. 

First, the proof for Theorem 1, i.e., Lasso theory transfer from linear models to StoNet, is based on the theoretical framework provided by the imputation-regularized optimization (IRO)
algorithm [1](Algorithm 1 in the appendix, it is a stochastic EM type algorithm). 
In the proof, the regularized regression problem is assumed to be exactly 
solved at each iteration.
Given the imputed value for the hidden units, it corresponds to solving a series of linear regression with LASSO penalty,
which is actually doable. 
Next, we pointed out that the objective function involved in Theorem 1 can also be solved using an adaptive stochastic gradient MCMC algorithm. 
Therefore, the proof for Theorem 1 is valid based on the IRO algorithm, which does not involve any issue 
on optimizing a non-smooth penalty with stochastic gradients. 

Regarding the issue on optimizing a non-smooth penalty with SGD,  particularly for the Lasso penalty, 
it has been shown in Nguyen and Yin (2021) [2] that SGD works. Therefore, our use of the adaptive SGMCMC for training the StoNet is valid. 

### Reference

[1] Liang, Faming, et al. "An imputation–regularized optimization algorithm for high dimensional missing data problems and beyond." Journal of the Royal Statistical Society Series B: Statistical Methodology 80.5 (2018): 899-926.

[2] Nguyen, N.N. and Yin, G. (2021). Stochastic Approximation with Discontinuous Dynamics,
Differential Inclusions, and Applications. 
arXiv:2108.12652v1.