Statistical Inference for Deep Learning via Stochastic Modeling

Thank you for your constructive reviews and comments, please find our point-by-point response below

R3W1 The main issues of this paper including novelty and soundness of the results.

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.

R3W2 Abstract description is not clear, and writing needs to be improved

Thank you for your comment. We will enhance the abstract by precisely highlighting the contributions of this paper.

R3W3 Comparison with other method is too few and experiment performance improvement is marginal.

Thank you for your comment. In the rebuttal, we have conducted more numerical experiments, including a large-scale model for CIFAR 100, an ablation study for the regularization parameter $\lambda$, comparison with split Conformalized quantile regression(CQR),
and one more simulation example for high-dimensional nonlinear variable selection. Please refer to our reply to Reviewer 9U3v and Reviewer HASS for the details (R1W2, R2C3-1, R2C3-3). And in our results, we have included the mean and standard deviation of the metric, which shows that our improvements are significant(see R1W2 of our rebuttal to reviewer 9U3v)

R3Q1 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.

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.

R3Q2 Figure 3 lacks legend, what does lines of different color represent?

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.

R3Q3 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.

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.

Reference

[1] Guo, Chuan, et al. "On calibration of modern neural networks." International conference on machine learning. PMLR, 2017.

[2] Liang, Siqi, Yan Sun, and Faming Liang. "Nonlinear Sufficient Dimension Reduction with a Stochastic Neural Network." Advances in Neural Information Processing Systems 35 (2022): 27360-27373.

R3Q4 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.

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.

R3Q5 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.

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].

Reference

[1]Romano, Yaniv, Evan Patterson, and Emmanuel Candes. "Conformalized quantile regression." Advances in neural information processing systems 32 (2019).

[2] Liang, Siqi, Yan Sun, and Faming Liang. "Nonlinear Sufficient Dimension Reduction with a Stochastic Neural Network." Advances in Neural Information Processing Systems 35 (2022): 27360-27373.

R2W1-3 The empirical experiments are limited (see C3)

R2W2 Originality / Novelty: It often does not get clear what the additional contribution is when compared to StoNet; see also first point on experiments in C3

R2C3-1 The experiments mainly present coverage rates and calibration results and it is unclear how much of this performance comes from StoNet itself (missing ablation study)

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]

R2C3-2 Only contain one simulation study with a fixed setup (restrictive)

R2C3-3 do not elaborate on the selection quality except for the small simulation study (missing empirical evidence)

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$):

$

y = \frac{5x_2}{1 + x_1^2} +5\sin (x_3 x_4) + 2x_5 + 0x_6 + \dots 0x_{2000} + \epsilon

$

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.

Reference

[1]Romano, Yaniv, Evan Patterson, and Emmanuel Candes. "Conformalized quantile regression." Advances in neural information processing systems 32 (2019).

Thank you for your constructive reviews and comments, please find our point-by-point response below

R2W1-1 Some assumptions (already in StoNet) seem to be not realistic (see C1)

R2C1 There is no one "true parameter" in deep neural networks given their overparametrization-induced and hidden symmetries and I don't see how Assumption A2 holds in practice. While the authors recognize this, writing "Given nonidentifiability of the neural network model, Assumption A2 has implicitly assumed that each $\theta$ 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", it does not get clear what this restriction of the Assumption implies and does certainly not account for scaling symmetries present in ReLU networks or hidden symmetries as mentioned above.

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.

R2W1-2 It's unclear, how the algorithm actually enforces sparsity

R2C2 How is the model optimized with a Lasso penalty given that this penalty is non-smooth and stochastic variants hence do not yield exact-zero solutions?

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.

Thanks for your detailed answer, clarifications and additional experiments. A couple of remaining questions (in case time allows to answer):

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.

How exactly are these equivalent classes defined? Because as far as I understand, symmetries can also be data-dependent (https://arxiv.org/pdf/2210.17216.pdf), exist after accounting for permutation and scaling symmetries (e.g., for hidden symmetries https://arxiv.org/abs/2306.06179), and even if all of these could be accounted for, network weight and function distributions apparently remain multimodal in nature (see, e.g., https://arxiv.org/pdf/2304.02902.pdf, E.2 / why would methods like deep ensemble work so much better than local methods otherwise).

As you mention that your method is "local", this also begs the question what a comparison with local methods like Laplace approximation (https://arxiv.org/abs/2106.14806) and non-local methods like deep ensembles (https://arxiv.org/abs/1612.01474) would look like.

To enforce sparsity in practice, we can put a threshold such as …

Is it actually possible to simply transfer the arguments from the Liu and Yu paper to your much more complicated setup? A lot of its theory is heavily based on the linear model structure.

And in the first place: Can you actually be sure that optimizing a non-smooth penalty with an optimizer not aware of the non-smoothness will do something meaningful? I thought most of the guarantees require a certain smoothness notion of the underlying function.

In the following table, we can see that for the small data set …

Just as a side note: I think the Wine dataset is ordinal and would require something like a proportional odds model.

Thank you for your valuable feedback, please find our response to your questions below

R2R1 How exactly are these equivalent classes defined? Because as far as I understand, symmetries can also be data-dependent (https://arxiv.org/pdf/2210.17216.pdf), exist after accounting for permutation and scaling symmetries (e.g., for hidden symmetries https://arxiv.org/abs/2306.06179), and even if all of these could be accounted for, network weight and function distributions apparently remain multimodal in nature (see, e.g., https://arxiv.org/pdf/2304.02902.pdf, E.2 / why would methods like deep ensemble work so much better than local methods otherwise). As you mention that your method is "local", this also begs the question what a comparison with local methods like Laplace approximation (https://arxiv.org/abs/2106.14806) and non-local methods like deep ensembles (https://arxiv.org/abs/1612.01474) would look like.

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).

R1W4 The proposed MAP learning does not really utilise the power of the probabilistic model.

In section 3.2, we estimate the parameter $\theta$ via solving an integration equation (i.e., the equation appeared at the bottom of page 5) using an adaptive stochastic gradient MCMC algorithm. The resulting estimator can be viewed as a maximum a posteriori (MAP) estimator or a regularized frequentist estimator. The power of the StoNet, which enables uncertainty quantification for its prediction, comes from its latent variable structure. A fully Bayesian treatment for the parameters of StoNet is not necessary at least for the goal of this paper.

R1W5 The authors should compare in more detail to existing approaches introducing noise to the network, e.g. Gaussian dropout. Overall, I think the work should be more linked to the existing research.

In the revision, we will expand the 'related work' paragraph to include more discussions on other types of stochastic neural networks (that introduce noise to the network). Specifically, the dropout-type methods primarily focus on reducing overfitting, and many other stochastic neural networks are mostly based on heuristic arguments or frame the problem within a framework of variational inference. In Table 4 of Sun and Liang (2022) [2], a kernel-expanded StoNet (similar to the one studied in the paper) and some other stochastic methods (including dropout, variational inference, probabilistic back-propagation) have been compared in terms of prediction. The numerical results indicate that StoNet generally outperforms the others.

As mentioned in the paper, one of our major goals is to provide a mathematically rigorous approach, post-StoNet, for quantifying the prediction uncertainty of large-scale deep learning models. However, the existing stochastic neural networks are generally invalid for prediction uncertainty quantification. In particular, it is unclear how to make statistical inferences with the dropout method.

R1W6 There are works considering similar treatment as section 4 that should be at least cited, e.g.[1]

Thank you for pointing out the reference. It also discussed the approximate distribution of the hidden layers, but it is under the Variational Bayes framework, where the parameter of each layer follows the variational distribution, and the goal is to compute Evidence Lower Bound (ELBO). In the revision, we will cite and discuss it as a piece of related works.

R1W7 The developed theory mostly relies on the convergence of MAP/MLE, which happens to be very slow in practice.

The proposed algorithm decomposes the nonconvex high-dimensional DNN training problem into a series of low-dimensional convex optimization problems, each corresponding to solving a Lasso linear/logistic regression problem. Since the dimension of parameters involved in each of the Lasso optimization problems is low, the convergence of the algorithm should not be a concern. Furthermore, in section 6.2, the StoNet model we used in the post-StoNet approach is usually small. For example, for the CIFAR10 example, the post-StoNet model we used has only one hidden layer with 100 hidden units. Again, the computation for such a small model should not be a concern.

Reference

[1] Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E. Turner, Jos'e Miguel Hern'andez-Lobato, Alexander L. Gaunt Deterministic Variational Inference for Robust Bayesian Neural Networks.

[2] Sun, Yan, and Faming Liang. "A kernel-expanded stochastic neural network." Journal of the Royal Statistical Society Series B: Statistical Methodology 84.2 (2022): 547-578.

Thank you for your constructive reviews and comments, please find our point-by-point response below

R1W(eakness)1 The differences in Table 1 look very small and I wonder if they are statistically significant at all?

Table 1 gives the coverage rate of StoNet with different choices of $\sigma$. Please note that through this example, we intend to show the trend that the coverage rate becomes closer to the target level 95% as the value of $\sigma$ is decreased. Therefore, the results with different values of $\sigma$ are not expected to be significantly different. Moreover, the similar coverage rates reported in the table imply that the performance of the StoNet is not very sensitive to the choice of $\sigma$. This is just as what we expected.

R1W2 I think more empirical evidence would make the paper stronger. The authors discuss scalability, but there's only one experiment with bigger networks (Table 2) where the gains are very marginal?

We will add an experiment on CIFAR100 in the revision. The following table shows the prediction accuracy, negative log likelihood loss (NLL), and expected calibration error (ECE) obtained by the ResNet110 model for the dataset with different posthoc calibration methods. The training setting for the model is the same as that used for the CIFAR10 dataset (see the paper for the details) and we use $\lambda = 5e-5$ for the penalty. The comparison shows that the proposed Post-StoNet approach provides better ECE, demonstrating improvement in model calibration.

Our method indeed provides significant improvement over the baselines. For example, in Table 2 of the paper, for the ResNet110 model, we reported the mean and standard deviation of each metric. Post-StoNet achieves $0.70\\%(0.14\\%)$ ECE, while temperature scaling achieves $1.22\\%(0.16\\%)$ ECE. A two-sample t-test with pooled variance will have a $p$-value of $3.95e-7$. Hence, the difference is highly significant. In the table of CIFAR100 results below, we also include the $p$-value of the two-sample t-tests comparing the ECE of our methods with other baselines, which again shows our method provides significant improvements.

We also conducted another experiment with ImageNet data, where we used a pre-trained ResNet152 model from PyTorch. The output of the last hidden layer is 2048. We use 20% of the validation set(10000 images) as the calibration set for posthoc calibration. The size of the calibration set is relatively small, in this case, we found that training a logistic regression with LASSO penalty gives better results, where we use the parameter of the last hidden layer of the pre-trained model to initialize the parameter of the logistic regression and train the parameter with SGD for 500 epochs with learning rate 0.001, momentum 0.9 and batch size 128. This model can be viewed as a StoNet without hidden layers, or introducing shortcut connections in StoNet. As we discussed in section 6.2, as long as the model gives a consistent estimation, it is still valid. The result is given in the following table. Because of the use of the pre-trained model, we only report one value for each metric, the Post-StoNet still makes clear improvement in terms of NLL and ECE compared to the original model.

R1W3 I don't think overparameterization of the NNs is necessarily a bad thing, it seems more like an open research question?

We agree with you that overparameterization is not necessarily bad. In terms of optimization, many works have shown good properties of the energy landscape of over-parameterized models. However, how to quantify the prediction uncertainty of such an over-parameterized model is largely an open problem.

The split conformal method provides a possible way to address this problem, but the resulting prediction confidence interval can be overly wide when the model is over-fitted (typically happens for over-parameterized neural networks). We address this problem using the proposed post-StoNet method. It adapts the learned model to be
consistent in prediction, leading to faithful prediction intervals.