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: The paper gives theoretical explanation for improved generalization error of winning tickets for which only empirical results are known. Their results are based on teacher-student setting in which training samples are assumed to be generated from an unknown teacher network and student network is supposed to learn only from those samples. They give an accelerated gradient descent method for learning the pruned  network and convergence and sample complexity analysis for this algorithm. The empirical risk function is shown to have an enlarged convex region for a pruned network, which justifies the importance of the winning ticket. Learning on pruned network with the AGD algorithm gives a model closer to the teacher network with the same number of iterations, which implies better generalization of the trained pruned network. These findings are validated with experiments on synthetic and real datasets (MNIST, CIFAR-10).