Why Lottery Ticket Wins? A Theoretical Perspective of Sample Complexity on Sparse Neural Networks

1 Introduction . Neural network pruning can reduce the computational cost of model training and inference significantly and potentially lessen the chance of overfitting [ 33 , 26 , 15 , 25 , 28 , 51 , 58 , 41 ] . The recent Lottery Ticket Hypothesis ( LTH ) [ 20 ] claims that a randomly initialized dense neural network al- 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ways contains a so-called “ winning ticket , ” which is a sub-network bundled with the corresponding initialization , such that when trained in isolation , this winning ticket can achieve at least the same testing accuracy as that of the original network by running at most the same amount of training time . This so-called “ improved generalization of winning tickets ” is verified empirically in [ 20 ] . LTH has attracted a significant amount of recent research interests [ 45 , 70 , 39 ] . Despite the empirical success [ 19 , 63 , 55 , 11 ] , the theoretical justification of winning tickets remains elusive except for a few recent works . [ 39 ] provides the first theoretical evidence that within a randomly initialized neural network , there exists a good sub-network that can achieve the same test performance as the original network . Meanwhile , recent work [ 42 ] trains neural network by adding the ` 1 regularization term to obtain a relatively sparse neural network , which has a better performance numerically . However , the theoretical foundation of network pruning is limited . The existing theoretical works usually focus on finding a sub-network that achieves a tolerable loss in either expressive power or training accuracy , compared with the original dense network [ 2 , 71 , 61 , 43 , 4 , 3 , 35 , 5 , 59 ] . To the best of our knowledge , there exists no theoretical support for the improved generalization achieved by winning tickets , i.e. , pruned networks with faster convergence and better test accuracy . Contributions : This paper provides the first systematic analysis of learning pruned neural networks with a finite number of training samples in the oracle-learner setup , where the training data are generated by a unknown neural network , the oracle , and another network , the learner , is trained on the dataset . Our analytical results also provide a justification of the LTH from the perspective of the sample complexity . In particular , we provide the first theoretical justification of the improved generalization of winning tickets . Specific contributions include : 1 . Pruned neural network learning via accelerated gradient descent ( AGD ) : We propose an AGD algorithm with tensor initialization to learn the pruned model from training samples . Our algorithm converges to the oracle model linearly , which has guaranteed generalization . 2 . First sample complexity analysis for pruned networks : We characterize the required number of samples for successful convergence , termed as the sample complexity . Our sample complexity bound depends linearly on the number of the non-pruned weights and is a significant reduction from directly applying conventional complexity bounds in [ 69 , 66 , 67 ] . 3 . Characterization of the benign optimization landscape of pruned networks : We show analytically that the empirical risk function has an enlarged convex region for a pruned network , justifying the importance of a good sub-network ( i.e. , the winning ticket ) . 4 . Characterization of the improved generalization of winning tickets : We show that gradientdescent methods converge faster to the oracle model when the neural network is properly pruned , or equivalently , learning on a pruned network returns a model closer to the oracle model with the same number of iterations , indicating the improved generalization of winning tickets . Notations . Vectors are bold lowercase , matrices and tensors are bold uppercase . Scalars are in normal font , and sets are in calligraphy and blackboard bold font . I denote the identity matrix . N and R denote the sets of nature number and real number , respectively . ‖z‖ denotes the ` 2-norm of a vector z , and ‖Z‖2 , ‖Z‖F and ‖Z‖∞ denote the spectral norm , Frobenius norm and the maximum value of matrix Z , respectively . [ Z ] stands for the set of { 1 , 2 , · · · , Z } for any number Z ∈ N. In addition , f ( r ) = O ( g ( r ) ) ( or f ( r ) = Ω ( g ( r ) ) ) if f ≤ C · g ( or f ≥ C · g ) for some constant C > 0 when r is large enough . f ( r ) = Θ ( g ( r ) ) if both f ( r ) = O ( g ( r ) ) and f ( r ) = Ω ( g ( r ) ) holds , where c · g ≤ f ≤ C · g for some constant 0 ≤ c ≤ C when r is large enough . 1.1 Related Work . Network pruning . Network pruning methods seek a compressed model while maintaining the expressive power . Numerical experiments have shown that over 90 % of the parameters can be pruned without a significant performance loss [ 10 ] . Examples of pruning methods include irregular weight pruning [ 25 ] , structured weight pruning [ 57 ] , neuron-based pruning [ 28 ] , and projecting the weights to a low-rank subspace [ 13 ] . Winning tickets . [ 20 ] employs an Iterative Magnitude Pruning ( IMP ) algorithm to obtain the proper sub-network and initialization . IMP and its variations [ 22 , 46 ] succeed in deeper networks like Residual Networks ( Resnet ) -50 and Bidirectional Encoder Representations from Transformers ( BERT ) network [ 11 ] . [ 21 ] shows that IMP succeeds in finding the “ winning ticket ” if the ticket is stable to stochastic gradient descent noise . In parallel , [ 36 ] shows numerically that the “ winning ticket ” initialization does not improve over a random initialization once the correct sub-networks are found , suggesting that the benefit of “ winning ticket ” mainly comes from the sub-network structures . [ 18 ] analyzes the sample complexity of IMP from the perspective of recovering a sparse vector in a linear model rather than learning neural networks . Feature sparsity . High-dimensional data often contains redundant features , and only a subset of the features is used in training [ 6 , 14 , 27 , 60 , 68 ] . Conventional approaches like wrapper and filter methods score the importance of each feature in a certain way and select the ones with highest scores [ 24 ] . Optimization-based methods add variants of the ` 0 norm as a regularization to promote feature sparsity [ 68 ] . Different from network pruning where the feature dimension still remains high during training , the feature dimension is significantly reduced in training when promoting feature sparsity . Over-parameterized model . When the number of weights in a neural network is much larger than the number of training samples , the landscape of the objective function of the learning problem has no spurious local minima , and first-order algorithms converge to one of the global optima [ 37 , 44 , 64 , 50 , 9 , 49 , 38 ] . However , the global optima is not guaranteed to generalize well on testing data [ 62 , 64 ] . Generalization analyses . The existing generalization analyses mostly fall within three categories . One line of research employs the Mean Field approach to model the training process by a differential equation assuming infinite network width and infinitesimal training step size [ 12 , 40 , 56 ] . Another approach is the neural tangent kernel ( NTK ) [ 30 ] , which requires strong and probably unpractical over-parameterization such that the nonlinear neural network model behaves as its linearization around the initialization [ 1 , 17 , 72 , 73 ] . The third line of works follow the oracle-learner setup , where the data are generated by an unknown oracle model , and the learning objective is to estimate the oracle model , which has a generalization guarantee on testing data . However , the objective function has intractably many spurious local minima even for one-hidden-layer neural networks [ 48 , 47 , 64 ] . Assuming an infinite number of training samples , [ 8 , 16 , 52 ] develop learning methods to estimate the oracle model . [ 23 , 69 , 66 , 67 ] extend to the practical case of a finite number of samples and characterize the sample complexity for recovering the oracle model . Because the analysis complexity explodes when the number of hidden layers increases , all the analytical results about estimating the oracle model are limited to one-hidden-layer neural networks , and the input distribution is often assumed to be the standard Gaussian distribution . 2 Problem Formulation . In an oracle-learner model , given any input x ∈ Rd , the corresponding output y is generated by a pruned one-hidden-layer neural network , called oracle , as shown in Figure 1 . The oracle network is equipped with K neurons where the j-th neuron is connected to any arbitrary r∗j ( r ∗ j ≤ d ) input features . LetW ∗ = [ w∗1 , · · · , w∗K ] ∈ Rd×K denotes all the weights ( pruned ones are represented by zero ) . The number of non-zero entries in w∗j is at most r ∗ j . The oracle network is not unique because permuting neurons together with the corresponding weights does not change the output . Therefore , the output label y obtained by the oracle network satisfies 1 y = 1 K K∑ j=1 φ ( w∗Tj x ) + ξ : = g ( x ; W ∗ ) + ξ = g ( x ; W ∗P ) + ξ , ( 1 ) where ξ is arbitrary unknown additive noise bounded by some constant |ξ| , φ is the rectified linear unit ( ReLU ) activation function with φ ( z ) = max { z , 0 } , and P ∈ { 0 , 1 } K×K is any permutation matrix . M∗ is a mask matrix for the oracle network , such that M∗j , i equals to 1 if the weight w∗j , i is not pruned , and 0 otherwise . Then , M∗ is an indicator matrix for the non-zero entries ofW ∗ with M∗ W ∗ = W ∗ , where is entry-wise multiplication . Based on N pairs of training samples D = { xn , yn } Nn=1 generated by the oracle , we train on a learner network equipped with the same number of neurons in the oracle network . However , the j-th neuron in the learner network is connected to rj input features rather than r∗j . Let rmin , rmax , and rave denote the minimum , maximum , and average value of { rj } Kj=1 , respectively . LetM denote the 1It is extendable to binary classification , and the output is generated by Prob ( yn = 1|xn ) = g ( xn ; W ∗ ) . mask matrix with respect to the learner network , and wj is the j-th column ofW . The empirical risk function is defined as f̂D ( W ) = 1 2N N∑ n=1 ( 1 K K∑ j=1 φ ( wTj xn ) − yn ) 2 . ( 2 ) When the maskM is given , the learning objective is to estimate a proper weight matrixW for the learner network from the training samples D via solving minW∈Rd×K f̂D ( W ) s.t . M W = W . ( 3 ) M is called an accurate mask if the support ofM covers the support of a permutation ofM∗ , i.e. , there exists a permutation matrix P such that ( M∗P ) M = M∗ . When M is accurate , and ξ = 0 , there exists a permutation matrix P such that W ∗P is a global optimizer to ( 3 ) . Hence , if W ∗P can be estimated by solving ( 3 ) , one can learn the oracle network accurately , which has guaranteed generalization performance on the testing data . We assume xn is independent and identically distributed from the standard Gaussian distribution N ( 0 , Id×d ) . The Gaussian assumption is motivated by the data whitening [ 34 ] and batch normalization techniques [ 29 ] that are commonly used in practice to improve learning performance . Moreover , training one-hidden-layer neural network with multiple neurons has intractable many fake minima [ 47 ] without any input distribution assumption . In addition , the theoretical results in Section 3 assume an accurate mask , and inaccurate mask is evaluated empirically in Section 4 . The questions that this paper addresses include : 1. what algorithm to solve ( 3 ) ? 2. what is the sample complexity for the accurate estimate of the weights in the oracle network ? 3. what is the impact of the network pruning on the difficulty of the learning problem and the performance of the learned model ?

Summary: This paper presents a theoretical analysis of the advantages of learning pruned neural networks . This analysis considers a teacher-student setup with a finite number of training samples. The theoretical results presented in the paper show that pruned networks can have multiple advantages, such as faster training convergence, a lower number of samples required for successful convergence and an enlarged convex region.