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

1 Introduction . 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 can not 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 can not 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 . 2 Theoretical Foundation for Sparse Deep Learning . 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 . 2.1 Posterior consistency and structure selection consistency . 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 ) .

Summary: 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.