Sparse Deep Learning: A New Framework Immune to Local Traps and Miscalibration
Year: 2022

Abstract:
Deep learning has powered recent successes of artificial intelligence (AI). However, the deep neural network, as the basic model of deep learning, has suffered from issues such as local traps and miscalibration. In this paper, we provide a new framework for sparse deep learning, which has the above issues addressed in a coherent way. In particular, we lay down a theoretical foundation for sparse deep learning and propose prior annealing algorithms for learning sparse neural networks. The former has successfully tamed the sparse deep neural network into the framework of statistical modeling, enabling prediction uncertainty correctly quantified. The latter can be asymptotically guaranteed to converge to the global optimum, enabling the validity of the down-stream statistical inference. Numerical result indicates the superiority of the proposed method compared to the existing ones.

heading: 1 Introduction
text: During the past decade, deep neural networks (DNNs) have achieved the state-of-the-art performance in many machine learning tasks such as computer vision and natural language processing. However, the DNN suffers from a training-prediction dilemma from the perspective of statistical inference: A small DNN model can be well calibrated, but tends to get trapped into a local optimum; on the other hand, an over-parameterized DNN model can be easily trained to a global optimum (with zero training loss), but tends to be miscalibrated [14]. In consequence, it is often unclear whether a DNN is guaranteed to have a desired property after training instead of getting trapped into an arbitrarily poor local minimum, or whether its decision/prediction is reliable. This difficulty makes the trustworthiness of AI highly questionable.
To resolve this difficulty, researchers have attempted from two sides of the training-prediction dilemma. Towards understanding the optimization process of the DNN training, a line of researches have been done. For example, [13] and [27] studied the training loss surface of over-parameterized DNNs. They showed that for a fully connected DNN, almost all local minima are globally optimal, if the width of one layer of the DNN is no smaller than the training sample size and the network
∗To whom correspondence should be addressed: Faming Liang
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
structure from this layer on is pyramidal. Recently, [1, 8, 35] and [36] explored the convergence theory of the gradient-based algorithms in training over-parameterized DNNs. They showed that the gradient-based algorithms with random initialization can converge to global minima provided that the width of the DNN is polynomial in training sample size.
To improve calibration of the DNN, different methods have been developed, see e.g., Monte Carlo dropout [12] and deep ensemble [18]. However, these methods did not provide a rigorous study for the asymptotic distribution of the DNN prediction and thus could not correctly quantify its uncertainty. Recently, researchers have attempted to address this issue with sparse deep learning. For example, for Bayesian sparse neural networks, [19], [28] and [31] established the posterior consistency, and [33] further established the Bernstein-von Mises (BvM) theorem for linear and quadratic functionals. The latter guarantees in theory that the Bayesian credible region has a faithful frequentist coverage. However, since the theory by [33] does not cover the point evaluation functional, the uncertainty of the DNN prediction still cannot be correctly quantified. Moreover, their theory is developed with the spike-and-slab prior (i.e., each weight or bias of the DNN is subject to a spike-and-slab prior), whose discrete nature makes the resulting posterior distribution extremely hard to simulate. To facilitate computation, [31] employed a mixture Gaussian prior. However, due to nonconvexity of the loss function of the DNN, a direct MCMC simulation still cannot be guaranteed to converge to the right posterior distribution even with the mixture Gaussian prior.
In this paper, we provide a new framework for sparse deep learning, which successfully resolved the training-prediction dilemma. In particular, we propose two prior annealing algorithms, one from the frequentist perspective and one from the Bayesian perspective, for learning sparse neural networks. The algorithms start with an over-parameterized deep neural network and then have its structure gradually sparsified. We provide a theoretical guarantee that the training procedures are immune to local traps, the resulting sparse structures are consistent, and the predicted values are asymptotically normally distributed. The latter enables the prediction uncertainty correctly quantified. Our contribution in this paper is two-fold:
• We provide a new framework for sparse deep learning, which is immune to local traps and miscalibration.
• We lay down the theoretical foundation for how to make statistical inference with sparse deep neural networks.
The remaining part of the paper is organized as follows. Section 2 lays down the theoretical foundation for sparse deep learning. Section 3 describes the proposed prior annealing algorithms. Section 4 presents some numerical results. Section 5 concludes the paper.

heading: 2 Theoretical Foundation for Sparse Deep Learning
text: As mentioned previously, sparse deep learning has received much attention as a promising way for addressing the miscalibration issue of the DNN. Theoretically, the approximation power of the sparse DNN has been studied for various classes of functions [30, 3]. Under the Bayesian setting, posterior consistency has been established in [19, 28, 31]. In particular, the work [31] has achieved important progress toward taming sparse DNNs into the framework of statistical modeling. They provide a neural network approximation theory fundamentally different from the existing ones. In the existing theory, no data is involved and a small network can potentially achieve an arbitrarily small approximation error by allowing connection weights to take values in an unbounded space[24]. In contrast, the theory by [31] links the network approximation error, the network size, and the bound of connection weights to the training sample size. They prove that for a given training sample size n, a sparse DNN of size O(n/ log(n)) has been large enough to approximate many types of functions, such as affine functions and piecewise smooth functions, arbitrarily well as n→∞. Moreover, they prove that the sparse DNN possesses many theoretical guarantees. For example, its structure is more interpretable, from which the relevant variables can be consistently identified for high-dimensional nonlinear systems; and its generalization error bound is asymptotically optimal.
From the perspective of statistical inference, some gaps remain toward taming sparse DNNs into the framework of statistical modeling. This paper bridges the gap by establishing (i) asymptotic normality of the connection weights, and (ii) asymptotic normality of the prediction.

heading: 2.1 Posterior consistency and structure selection consistency
text: This subsection provides a brief review of the sparse DNN theory developed in [31] and gives the conditions that we will use in the followed theoretical developments. Without loss of generality, we letDn = (x(i), y(i))i=1,...,n denote a dataset of n i.i.d observations, where x(i) ∈ Rpn and y(i) ∈ R. Consider a generalized linear model with the distribution of y given by
f(y|µ∗(x)) = exp{A(µ∗(x))y +B(µ∗(x)) + C(y)},
where µ∗(x) is a nonlinear function of x and A(·), B(·) and C(·) are appropriately defined functions. For example, for normal regression, we have A(µ∗) = µ∗/σ2, B(µ∗) = −µ∗2/2σ2, C(y) = −y2/2σ2− log(2πσ2)/2, and σ2 is a constant. We approximate µ∗(x) using a fully connected DNN with Hn − 1 hidden layers. Let Lh denote the number of hidden units at layer h with LHn = 1 for the output layer and L0 = pn for the input layer. Let wh ∈ RLh×Lh−1 and bh ∈ RLh×1, h ∈ {1, 2, ...,Hn} denote the weights and bias of layer h, and let ψh : RLh×1 → RLh×1 denote a coordinate-wise and piecewise differentiable activation function of layer h. The DNN forms a nonlinear mapping
µ(β,x) = wHnψHn−1 [ · · ·ψ1 [ w1x+ b1 ] · · · ] + bHn , (1)
where β = (w, b) = { whij , b h k : h ∈ {1, . . . ,Hn}, i, k ∈ {1, . . . , Lh}, j ∈ {1, . . . , Lh−1} } denotes
the collection of all weights and biases, consisting of Kn = ∑Hn h=1 (Lh−1 × Lh + Lh) elements. For convenience, we treat bias as a special connection and call each element in β a connection weight. In order to represent the structure for a sparse DNN, we introduce an indicator variable for each connection weight. Let γw h and γb h
denote the indicator variables associated with wh and bh, respectively. Let γ = {γwhij ,γb h
k : h ∈ {1, . . . ,Hn}, i, k ∈ {1, . . . , Lh} , j ∈ {1, . . . , Lh−1}}, which specifies the structure of the sparse DNN. With slight abuse of notation, we will also write µ(β,x) as µ(β,x,γ) to include the information of the network structure. We assume µ∗(x) can be well approximated by a parsimonious neural network with relevant variables, and call this parsimonious network as the true DNN model. More precisely, we define the true DNN model as
(β∗,γ∗) = arg min (β,γ)∈Gn, ‖µ(β,γ,x)−µ∗(x)‖L2(Ω)≤$n |γ|, (2)
where Gn := G(C0, C1, ε, pn, Hn, L1, L2, . . . , LHn) denotes the space of valid sparse networks satisfying condition A.2 (given below) for the given values of Hn, pn, and Lh’s, and $n is some sequence converging to 0 as n → ∞. For any given DNN (β,γ), the error µ(β,γ,x) − µ∗(x) can be generally decomposed as the network approximation error µ(β∗,γ∗,x) − µ∗(x) and the network estimation error µ(β,γ,x)− µ(β∗,γ∗,x). The L2 norm of the former is bounded by $n, and the order of the latter will be given in Lemma 2.1. For the sparse DNN, we make the following assumptions:
A.1 The input x is bounded by 1 entry-wisely, i.e. x ∈ Ω = [−1, 1]pn , and the density of x is bounded in its support Ω uniformly with respect to n.
A.2 The true sparse DNN model satisfies the following conditions:
A.2.1 The network structure satisfies: rnHn log n +rn logL+ sn log pn ≤ C0n1−ε, where 0 < ε < 1 is a small constant, rn = |γ∗| denotes the connectivity of γ∗, L = max1≤j≤Hn−1 Lj denotes the maximum hidden layer width, sn denotes the input dimension of γ∗. A.2.2 The network weights are polynomially bounded: ‖β∗‖∞ ≤ En, where En = nC1 for some constant C1 > 0.
A.3 The activation function ψ is Lipschitz continuous with a Lipschitz constant of 1.
Refer to [31] for explanations and discussions on these assumptions. We let each connection weight and bias be subject to a mixture Gaussian prior, i.e.,
whij ∼ λnN(0, σ21,n) + (1− λn)N(0, σ20,n), bhk ∼ λnN(0, σ21,n) + (1− λn)N(0, σ20,n), (3)
where λn ∈ (0, 1) is the mixture proportion, σ20,n is typically set to a very small number, while σ21,n is relatively large.
Posterior Consistency Let P ∗ and E∗ denote the respective probability measure and expectation with respect to data Dn. Let d(p1, p2) denote the Hellinger distance between two densities p1(x, y) and p2(x, y). Let π(A | Dn) be the posterior probability of an event A. Lemma 2.1. (Theorem 2.1 of [31]) Suppose Assumptions A.1-A.3 hold. If the mixture Gaussian prior (3) satisfies the conditions: λn = O(1/{Kn[nHn(Lpn)]τ}) for some constant τ > 0, En/{Hn log n + logL}1/2 . σ1,n . nα for some constant α > 0, and σ0,n . min { 1/{ √ nKn(n 3/2σ1,0/Hn) Hn}, 1/{ √ nKn(nEn/Hn) Hn} }
, then there exists an error sequence 2n = O($ 2 n) +O(ζ 2 n) such that limn→∞ n = 0 and limn→∞ n 2 n =∞, and the posterior distribution satisfies
P ∗ { π[d(pβ, pµ∗) > 4 n|Dn] ≥ 2e−cn 2 n } ≤ 2e−cn 2 n ,
E∗Dnπ[d(pβ, pµ∗) > 4 n|Dn] ≤ 4e −2cn 2n ,
(4)
for sufficiently large n, where c denotes a constant, ζ2n = [rnHn log n + rn logL + sn log pn]/n, pµ∗ denotes the underlying true data distribution, and pβ denotes the data distribution reconstructed by the Bayesian DNN based on its posterior samples.
Structure Selection Consistency The DNN is generally nonidentifiable due to the symmetry of network structure. For example, µ(β,γ,x) can be invariant if one permutes certain hidden nodes or simultaneously changes the signs or scales of certain weights. As in [31], we define a set of DNNs by Θ such that any possible DNN can be represented by one and only one DNN in Θ via nodes permutation, sign changes, weight rescaling, etc. Let ν(γ,β) ∈ Θ be an operator that maps any DNN to Θ via appropriate weight transformations. To serve the purpose of structure selection in Θ, we consider the marginal inclusion posterior probability (MIPP) approach proposed in [20]. For each connection, we define its MIPP by qi = ∫ ∑ γ ei|ν(γ,β)π(γ|β, Dn)π(β|Dn)dβ for i = 1, 2, . . . ,Kn, where ei|ν(γ,β) is the indicator of connection i. The MIPP approach is to choose the connections whose MIPPs are greater than a threshold q̂, i.e., setting γ̂ q̂ = {i : qi > q̂, i = 1, 2, . . . ,Kn} as an estimator of γ∗ ∈ Θ. Let A( n) = {β : d(pβ, pµ∗) ≥ n} and define ρ( n) = max1≤i≤Kn ∫ A( n)c ∑ γ |ei|ν(γ,β) − ei|ν(γ∗,β∗)|π(γ|β, Dn)π(β|Dn)dβ, which measures the structure difference between the true and sampled models on A( n)c. Then we have: Lemma 2.2. (Theorem 2.2 of [31]) If the conditions of Lemma 2.1 hold and ρ( n)→ 0 as n→∞ and n → 0, then (i) max1≤i≤Kn{|qi − ei|ν(γ∗,β∗)|} p→ 0; (ii) (sure screening) P (γ∗ ⊂ γ̂ q̂) p→ 1 for any pre-specified q̂ ∈ (0, 1); (iii) (consistency) P (γ∗ = γ̂0.5) p→ 1.
Lemma 2.2 implies consistency of variable selection for the true DNN model as defined in (2).

heading: 2.2 Asymptotic Normality of Connection Weights
text: In this section, we establish the asymptotic normality of the network parameters and predictions. Let nln(β) = ∑n i=1 log(pβ(xi, yi)) denote the log-likelihood function, and let π(β) denote the density of the mixture Gaussian prior (3). Let hi1,i2,...,id(β) denote the d-th order partial derivatives ∂dln(β) ∂βi1∂βi2 ···∂βid . Let Hn(β) denote the Hessian matrix of ln(β). Let hij(β) and hij(β) denote the (i, j)-th component ofHn(β) andH−1n (β), respectively. Let λ̄n(β) and λn(β) denotes the maximum and minimum eigenvalue of the Hessian matrix Hn(β), respectively. Let Bλ,n = λ̄ 1/2 n (β ∗)/λn(β ∗)
and bλ,n = √ rn/nBλ,n, where rn is the connectivity of γ∗. For a DNN parameterized by β, we define the weight truncation at the true model γ∗: (βγ∗)i = βi for i ∈ γ∗ and (βγ∗)i = 0 otherwise. For the mixture Gaussian prior (3), let Bδn(β ∗) = {β : |βi − β ∗ i | < δn,∀i ∈ γ∗, |βi − β ∗ i | < 2σ0,n log( σ1,n λnσ0,n
),∀i /∈ γ∗}. We follow the definition of asymptotic normality in [4] and [33]: Definition 2.1. Denote by dβ the bounded Lipschitz metric for weak convergence and by φn the mapping φn : β → √ n(g(β)− g∗). We say that the posterior distribution of the functional g(β) is asymptotically normal with the center g∗ and variance G if dβ(π[· | Dn] ◦ φ−1n , N(0, G)) → 0 in P ∗-probability as n→∞. We will write this more compactly as π[· | Dn] ◦ φ−1n N(0, G).
Theorem 2.1 establishes the asymptotic normality of ν̃(β), where ν̃(β) denotes a transformation of β which is invariant with respect to µ(β,γ,x) while minimizing ‖ν̃(β)− β∗‖∞.
Theorem 2.1. Assume the conditions of Lemma 2.2 hold with ρ( n) = o( 1Kn ) and C1 > 2 3 in Condition A.2.2. For some δn s.t. rn√n . δn . 1 3 √ nrn , let A( n, δn) = {β : maxi∈γ∗ |βi − β ∗ i | > δn, d(pβ, pµ∗) ≤ n}, where n is the posterior contraction rate as defined in Lemma 2.1. Assume there exists some constants C > 2 and M > 0 such that
C.1 β∗ = (β∗1,β ∗ 2, . . . ,β ∗ Kn) is generic [11, 10], mini∈γ∗ |β ∗ i | > Cδn and π(A( n, δn) |
Dn)→ 0 as n→∞.
C.2 |hi(β∗)| < M , |hj,k(β∗)| < M , |hj,k(β∗)| < M , |hi,j,k(β)| < M , |hl(β)| < M hold for any i, j, k ∈ γ∗, l /∈ γ∗ and β ∈ B2δn(β ∗).
C.3 sup { |Eβ(aTU)3| : ‖βγ∗ − β ∗‖ ≤ 1.2bλ,n, ‖a‖ = 1 } ≤ 0.1 √ n/rnλ 2 n(β ∗)/λ̄ 1/2 n (β
∗) and Bλ,n = O(1), where U = Z − Eβγ∗ (Z), Z denotes a random variable drawn from a neural network model parameterized by βγ∗ , and Eβγ∗ (Z) denotes the mean of Z.
Then π[ √ n(ν̃(β) − β∗) | Dn] N(0,V ) in P ∗-probability as n → ∞, where V = (vij), and vi,j = E(h i,j(β∗)) if i, j ∈ γ∗ and 0 otherwise.
Condition C.1 is essentially an identifiability condition, i.e., when n is sufficiently large, the DNN weights cannot be too far away from the true weights if the DNN produces approximately the same distribution as the true data. Condition C.2 gives typical conditions on derivatives of the DNN. Condition C.3 ensures consistency of the MLE of β∗ for the given structure γ∗ [29].

heading: 2.2.1 Asymptotic Normality of Prediction
text: Theorem 2.2 establishes asymptotic normality of the prediction µ(β,x0) for a test data point x0, which implies that a faithful prediction interval can be constructed for the learnt sparse neural network. Refer to Appendix A.4 for how to construct the prediction interval based on the theorem. Let µi1,i2,...,id(β,x0) denote the d-th order partial derivative
∂dµ(β,x0) ∂βi1∂βi2 ···∂βid .
Theorem 2.2. Assume the conditions of Theorem 2.1 and the following condition hold: |µi(β∗,x0)| < M , |µi,j(β,x0)| < M , |µk(β,x0)| < M hold for any i, j ∈ γ∗, k /∈ γ∗ and β ∈ B2δn(β ∗), where M is as defined in Theorem 2.1. Then π[ √ n(µ(β,x0)−µ(β∗,x0)) | Dn] N(0,Σ), where Σ = ∇γ∗µ(β∗,x0)TH−1∇γ∗µ(β∗,x0) and H = E(−∇2γ∗ ln(β ∗)) is the Fisher information matrix.
The asymptotic normality for general smooth functional has been established in [4]. For linear and quadratic functional of deep ReLU network with a spike-and-slab prior, the asymptotic normality has been established in [33]. The DNN prediction µ(β,x0) can be viewed as a point evaluation functional over the neural network function space. However, in general, this functional is not smooth with respect to the locally asymptotic normal (LAN) norm. The results of [4] and [33] are not directly applicable for the asymptotic normality of µ(β,x0).

heading: 3 Prior Annealing Algorithms for Sparse DNN Computation
text: As implied by Theorems 2.1 and 2.2, a consistent estimator of (γ∗,β∗) is essential for statistical inference of the sparse DNN. Toward this goal, [31] proved that the marginal inclusion probabilities qi’s can be estimated using Laplace approximation at the mode of the log-posterior. Based on this result, they proposed a multiple-run procedure. In each run, they first maximize the log-posterior by an optimization algorithm, such as SGD or Adam; then sparsify the DNN structure by truncating the weights less than a threshold to zero, where the threshold is calculated from the prior (3) based on the Laplace approximation theory; and then refine the weights of the sparsified DNN by running an optimization algorithm for a few iterations. Finally, they select a sparse DNN model from those obtained in the multiple runs according to their Bayesian evidence or BIC values. The BIC is suggested when the size of the sparse DNN is large.
Although the multiple-run procedure works well for many problems, it is hard to justify that it will lead to a consistent estimator of the true model (γ∗,β∗). To tackle this issue, we propose two prior annealing algorithms, one from the frequentist perspective and one from the Bayesian perspective.

heading: 3.1 Prior Annealing: Frequentist Computation
text: It has been shown in [27, 13] that the loss of an over-parameterized DNN exhibits good properties:
(S∗) For a fully connected DNN with an analytic activation function and a convex loss function at the output layer, if the number of hidden units of one layer is larger than the number of training points and the network structure from this layer on is pyramidal, then almost all local minima are globally optimal.
Motivated by this result, we propose a prior annealing algorithm, which is immune to local traps and aims to find a consistent estimate of (β∗,γ∗) as defined in (2). The detailed procedure of the algorithm is given in Algorithm 1.
Algorithm 1 Prior annealing: Frequentist
(i) (Initial training) Train a DNN satisfying condition (S*) such that a global optimal solution β0 = arg maxβ ln(β) is reached, which can be accomplished using SGD or Adam [16].
(ii) (Prior annealing) Initialize β at β0 and simulate from a sequence of distributions
π(β|Dn, τ, η(k), σ(k)0,n) ∝ enln(β)/τπ η(k)/τ k (β) for k = 1, 2, . . . ,m, where 0 < η
(1) ≤ η(2) ≤ · · · ≤ η(m) = 1, πk = λnN(0, σ21,n) + (1−λn)N(0, (σ (k) 0,n) 2), and σinit0,n = σ (1) 0,n ≥ σ (2) 0,n ≥ · · · ≥ σ (m) 0,n = σ end 0,n . The simulation can be done in an annealing manner using a stochastic gradient MCMC algorithm [34, 6, 22, 26]. After the stage m has been reached, continue to run the simulated annealing algorithm by gradually decreasing the temperature τ to a very small value. Denote the resulting DNN by β̂ = (β̂1, β̂2, . . . , β̂Kn).
(iii) (Structure sparsification) For each connection i ∈ {1, 2, . . . ,Kn}, set γ̃i = 1 if |β̂i| > √ 2σ0,nσ1,n√ σ21,n−σ20,n
√ log (
1−λn λn σ1,n σ0,n ) and 0 otherwise, where the threshold value of |β̂i| is ob-
tained by solving π(γi = 1|βi) > 0.5 based on the mixture Gaussian prior as in [31]. Denote the yielded sparse DNN structure by γ̃.
(iv) (Nonzero-weights refining) Refine the nonzero weights of the sparsified DNN by maximizing ln(β). Denote the resulting estimate by β̃γ̃ , which represents the MLE of β ∗.
For Algorithm 1, the consistency of (γ̃, β̃γ̃) as an estimator of (γ ∗,β∗) can be proved based on Theorem 3.4 of [27] for global convergence of β0, the property of simulated annealing (by choosing an appropriate sequence of ηk and a cooling schedule of τ ), Theorem 2.2 for consistency of structure selection, Theorem 2.3 of [31] for consistency of structure sparsification, and Theorem 2.1 of [29] for consistency of MLE under the scenario of dimension diverging. Then we can construct the confidence intervals for neural network predictions using (γ̃, β̃γ̃). The detailed procedure is given in supplementary material.
Intuitively, the initial training phase can reach the global optimum of the likelihood function. In the prior annealing phase, as we slowly add the effect of the prior, the landscape of the target distribution is gradually changed and the MCMC algorithm is likely to hit the region around the optimum of the target distribution. More explanations on the effect of the prior can be found in the supplementary material. In practice, let t denote the step index, a simple implementation of the initial training and prior annealing phases of Algorithm 1 can be given as follows: (i) for 0 < t < T1, run initial training; (ii) for T1 ≤ t ≤ T2, fix σ(t)0,n = σinit0,n and linearly increase ηt by
setting η(t) = t−T1T2−T1 ; (iii) for T2 ≤ t ≤ T3, fix η (t) = 1 and linearly decrease ( σ (t) 0,n )2 by setting(
σ (t) 0,n )2 = T3−tT3−T2 ( σinit0,n )2 + t−T2T3−T2 ( σend0,n )2 ; (iv) for t > T3, fix η(t) = 1 and σ (t) 0,n = σ end 0,n and
gradually decrease the temperature τ , e.g., setting τt = ct−T3 for some constant c.

heading: 3.2 Prior Annealing: Bayesian Computation
text: For certain problems the size (or #nonzero elements) of γ∗ is large, calculation of the Fisher information matrix is difficult. In this case, the prediction uncertainty can be quantified via posterior simulations. The simulation can be started with a DNN satisfying condition (S*) and performed using a SGMCMC algorithm [22, 26] with an annealed prior as defined in step (ii) of Algorithm 1 (For Bayesian approach, we may fix the temperature τ = 1). The over-parameterized structure and annealed prior make the simulations immune to local traps.
To justify the Bayesian estimator for the prediction mean and variance, we study the deviation of the path averaging estimator 1T ∑T t=1 φ(β (t)) and the posterior mean ∫ φ(β)π(β|Dn, η∗, σ∗0,n)dβ for some test function φ(β). For simplicity, we will focus on SGLD with prior annealing. Our analysis can be easily generalized to other SGMCMC algorithms [5].
For a test function φ(·), the difference between φ(β) and ∫ φ(β)π(β|Dn, η∗, σ∗0,n)dβ can be characterized by the Poisson equation:
Lψ(β) = φ(β)− ∫ φ(β)π(β|Dn, η∗, σ∗0,n)dβ,
where ψ(·) is the solution of the Poisson equation and L is the infinitesimal generator of the Langevin diffusion. i.e. for the following Langevin diffusion dβ(t) = ∇ log(π(β|Dn, η∗, σ∗0,n))dt+ √ 2IdWt, where I is identity matrix and Wt is Brownian motion, we have Lψ(β) := 〈∇ψ(β),∇ log(π(β|Dn, η∗, σ∗0,n)) + tr(∇2ψ(β)).
Let Dkψ denote the kth-order derivatives of ψ. To control the perturbation of φ(β), we need the following assumption about the function ψ(β): Assumption 3.1. For k ∈ {0, 1, 2, 3}, Dkψ exists and there exists a function V , s.t. ||Dkψ|| . Vpk for some constant pk > 0. In addition, V is smooth and the expectation of Vp on β(t) is bounded for some p ≤ 2 maxk{pk}, i.e. supt E(Vp(β (t))) < ∞, ∑ s∈(0,1) Vp(sβ1 + (1 − s)β2) . Vp(β1) + Vp(β2).
In step t of the SGLD algorithm, the drift term is replaced by∇β log π(β(t)|D(t)m,n, η(t), σ(t)0,n), where D (t) m,n is used to represent the mini-batch data used in step t. Let Lt be the corresponding infinitesimal generator. Let δt = Lt − L. To quantify the effect of δt, we introduce the following assumption: Assumption 3.2. β(t) has bounded expectation and the expectation of log-prior is Lipschitz continuous with respect to σ0,n, i.e. there exists some constant M s.t. supt E(|β
(t)|) ≤M <∞. For all t, |E log(π(β(t)|λn, σ(t1)0,n , σ1,n))− E log(π(β (t)|λn, σ(t2)0,n , σ1,n))| ≤M |σ (t1) 0,n − σ (t2) 0,n |.
Then we have the following theorem: Theorem 3.1. Suppose the model satisfy assumption 3.2, and a constant learning rate of is used. For a test function φ(·), if the solution of the Poisson equation ψ(·) satisfy assumption 3.1, then
E
( 1
T T−1∑ t=1 φ(β(t))− ∫ φ(β)π(β|Dn, η∗, σ∗0,n)dβ ) = O ( 1 T + ∑T−1 t=0 (|η (t) − η∗|+ |σ(t)0,n − σ∗0,n|) T + ) ,
(5) where σ∗0,n is treated as a fixed constant.
Theorem 3.1 shows that with prior annealing, the path averaging estimator can still be used for estimating the mean and variance of the prediction and constructing the confidence interval. The detailed procedure is given in supplementary material. For the case that a decaying learning rate is used, a similar theorem can be developed as in [5].

heading: 4 Numerical Experiments
text: This section illustrates the performance of the proposed method on synthetic and real data examples.2 For the synthetic example, the frequentist algorithm is employed to construct prediction intervals.
2The code for running these experiments can be found in https://github.com/sylydya/ Sparse-Deep-Learning-A-New-Framework-Immuneto-Local-Traps-and-Miscalibration
The real data example involves a large network, so both the frequentist and Bayesian algorithms are employed along with comparisons with some existing network pruning methods.

heading: 4.1 Synthetic Example
text: We consider a high-dimensional nonlinear regression problem, which shows that our method can identify the sparse network structure and relevant features as well as produce prediction intervals with correct coverage rates. The datasets were generated as in [31], where the explanatory variables x1, . . . , xpn were simulated by independently generating e, z1, . . . , zpn from N(0, 1) and setting xi =
e+zi√ 2 . The response variable was generated from a nonlinear regression model:
y = 5x2
1 + x21 + 5 sin(x3x4) + 2x5 + 0x6 + · · ·+ 0x2000 + ,
where ∼ N(0, 1). Ten datasets were generated, each consisting of 10000 samples for training and 1000 samples for testing. This example was taken from [31], through which we show that the prior annealing method can achieve similar results with the multiple-run method proposed in [31].
We modeled the data by a DNN of structure 2000-10000-100-10-1 with tanh activation function. Here we intentionally made the network very wide in one hidden layer to satisfy the condition (S*). Algorithm 1 was employed to learn the model. The detailed setup for the experiments were given in the supplementary material. The variable selection performance were measured using the false
selection rate FSR = ∑10 i=1 |Ŝi\S|∑10 i=1 |Ŝi| and negative selection rate NSR = ∑10 i=1 |S\Ŝi|∑10 i=1 |S| , where S is the set of true variables, Ŝi is the set of selected variables from dataset i and |Ŝi| is the size of Ŝi. The predictive performance is measured by mean square prediction error (MSPE) and mean square fitting error (MSFE). We compare our method with the multiple-run method (BNN_evidence) [31] as well as other existing variable selection methods including Sparse input neural network(Spinn) [11], Bayesian adaptive regression tree (BART) [2], linear model with lasso penalty (LASSO) [32], and sure independence screening with SCAD penalty (SIS)[9]. To demonstrate the importance of selecting correct variables, we also compare our method with two dense model with the same network structure: DNN trained with dropout(Dropout) and DNN trained with no regularization(DNN). Detailed setups for these methods were given in the supplementary material as well. The results were summarized in Table 1. With a single run, our method BNN_anneal achieves similar result with the multiple-run method. The latter trained the model for 10 times and selected the best one using Bayesian evidence. While for Spinn (with LASSO penalty), even with over-parametrized structure, it performs worse than the sparse BNN model.
To quantify the uncertainty of the prediction, we conducted 100 experiments over different training sets as generated previously. We constructed 95% prediction intervals over 1000 test points. Over the 1000 test points, the average coverage rate of the prediction intervals is 94.72%(0.61%), where (0.61%) denote the standard deviation. Figure 1 shows the prediction intervals constructed for 20 of the testing points. Refer to the supplementary material for the detail of the computation.

heading: 4.2 Real Data Example
text: As a different type of applications of the proposed method, we conducted unstructured network pruning experiments on CIFAR10 dataset[17]. Following the setup in [21], we train the residual network[15] with different networks size and pruned the network to different sparsity levels. The detailed experimental setup can be found in the supplementary material.
We compared the proposed methods, BNN_anneal (Algorithm 1) and BNN_average (averaged over last 75 networks simulated by the Bayesian version of the prior annealing algorithm), with several state-of-the-art unstructured pruning methods, including Consistent Sparse Deep Learning (BNN_BIC) [31], Dynamic pruning with feedback (DPF) [21], Dynamic Sparse Reparameterization (DSR) [25] and Sparse Momentum (SM) [7]. The results of the baseline methods were taken from [21] and [31]. The results of prediction accuracy for different models and target sparsity levels were summarized in Table 2. Due to the threshold used in step (iii) of Algorithm 1, it is hard for our method to make the pruning ratio exactly the same as the targeted one. We intentionally make the pruning ratio smaller than the target ratio, while our method still achieve better test accuracy. Compared to BNN_BIC, the test accuracy is very close, but the result of BNN_BIC is obtained by running the experiment 10 times while our method only run once. To further demonstrate that the proposed method result in better model calibration, we followed the setup of [23] and compared the proposed method with DPF on several metrics designed for model calibration, including negtive log likelihood (NLL), symmetrized, discretized KL distance between in and out of sample entropy distributions (JS-Distance), and expected calibration error (ECE). For JS-Distance, we used the test data of SVHN data set3 as out-of-distribution samples. The results were summarized in Table 3. As discussed in [23, 14], a well calibrated model tends to have smaller NLL, larger JS-Distance and smaller ECE. The comparison shows that the proposed method outperforms DPF in most cases. In addition to the network pruning method, we also train a dense model with the standard training set up. Compared to the dense model, the sparse network has worse accuracy, but it tends to outperform the dense network in terms of ECE and JS-Distance, which indicates that sparsification is also a useful way for improving calibration of the DNN.

heading: 5 Conclusion
text: This work, together with [31], has built a solid theoretical foundation for sparse deep learning, which has successfully tamed the sparse deep neural network into the framework of statistical modeling. As implied by Lemma 2.1, Lemma 2.2, Theorem 2.1, and Theorem 2.2, the sparse DNN can be simply viewed as a nonlinear statistical model which, like a traditional statistical model, possesses many nice properties such as posterior consistency, variable selection consistency, and asymptotic
3The Street View House Numbers (SVHN) Dataset: http://ufldl.stanford.edu/housenumbers/
normality. We have shown how the prediction uncertainty of the sparse DNN can be quantified based on the asymptotic normality theory, and provided algorithms for training sparse DNNs with theoretical guarantees for its convergence to the global optimum. The latter ensures the validity of the down-stream statistical inference.

heading: Acknowledgment
text: Funding in direct support of this work: NSF grant DMS-2015498, NIH grants R01-GM117597 and R01-GM126089, and Liang’s startup fund at Purdue University.

Review summaries:
The paper proposes an approach to learning sparse deep neural networks (neural networks where most of the weights are set to zero) using annealing: starting by training a large overparameterized neural network to fit the training data perfectly, then adding weight penalties according to a prespecified schedule. A theoretical argument is provided to argue why this is reasonable, with two new theory proofs provided which extend previous work presented in Sun et al. 2021. Experiments on synthetic data show that the proposed algorithms perform feature selection better than other existing algorithms. Experiments on CIFAR10 show the algorithms are competitive in terms of the accuracy vs. pruning tradeoff.

Overall, the work presents an interesting direction into NN pruning. However, the paper is quite hard to follow and overly complex in notation, with little effort to assist the reader. Further, the paper severely lacks useful devices, such as figures or in-depth description, to explain the methods to readers, making the work challenging to follow.  Last, some of the experimental evaluation is somewhat questionable, as specified in the details comments.  Though the direction seems useful, additional work is needed in communicating the results and in experiments in order to warrant acceptance. 

The submission addresses the issue of sparsifying trained deep neural networks in an optimal way, preserving performance and improving calibration. It builds on the formalism and results of Sun et al. (2021), proves additional results and proposes two prior-annealing algorithms, which they demonstrate on synthetic and real data, comparing favorably to existing sparsifying algorithms. 

1. Theoretical contributions: the authors extend the theory established by Sun et al. [2021] and study the asymptotic behaviors of sparse deep learning. 2. Methodological contributions: the authors propose prior annealing algorithms, one from the frequentist perspective and one from the Bayesian perspective, to implement the corresponding sparse neural nets. They show the performance of the prior annealing algorithms on one synthetic example and the CIFAR10 datasets.