Why Lottery Ticket Wins? A Theoretical Perspective of Sample Complexity on Pruned Neural Networks
Year: 2022

Abstract:
The lottery ticket hypothesis (LTH) [20] states that learning on a properly pruned network (the winning ticket) improves test accuracy over the original unpruned network. Although LTH has been justified empirically in a broad range of deep neural network (DNN) involved applications like computer vision and natural language processing, the theoretical validation of the improved generalization of a winning ticket remains elusive. To the best of our knowledge, our work, for the first time, characterizes the performance of training a pruned neural network by analyzing the geometric structure of the objective function and the sample complexity to achieve zero generalization error. We show that the convex region near a desirable model with guaranteed generalization enlarges as the neural network model is pruned, indicating the structural importance of a winning ticket. Moreover, when the algorithm for training a pruned neural network is specified as an (accelerated) stochastic gradient descent algorithm, we theoretically show that the number of samples required for achieving zero generalization error is proportional to the number of the non-pruned weights in the hidden layer. With a fixed number of samples, training a pruned neural network enjoys a faster convergence rate to the desired model than training the original unpruned one, providing a formal justification of the improved generalization of the winning ticket. Our theoretical results are acquired from learning a pruned neural network of one hidden layer, while experimental results are further provided to justify the implications in pruning multi-layer neural networks.

heading: 1 Introduction
text: 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.

heading: 1.1 Related Work
text: 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.

heading: 2 Problem Formulation
text: 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?

heading: 3 Algorithm and Theoretical Results
text: Section 3.1 studies the geometric structure of (3), and the main results are in Section 3.2. Section 3.3 briefly introduces the proof sketch and technical novelty, and the limitations are in Section 3.4.

heading: 3.1 Local Geometric Structure
text: Theorem 1 characterizes the local convexity of f̂D in (3). It has two important implications.
1. Strictly locally convex near ground truth: f̂D is strictly convex nearW ∗P for some permutation matrix P , and the radius of the convex ball is negatively correlated with √ r̃, where r̃ is in the order of rave. Thus, the convex ball enlarges as any rj decreases.
2. Importance of the winning ticket architecture: Compared with training on the dense network directly, training on a properly pruned sub-network has a larger local convex region near W ∗P , which may lead to easier estimation of W ∗P . To some extent, this result can be viewed as a theoretical validation of the importance of the winning architecture (a good sub-network) in [20]. Formally, we have
Theorem 1 (Local Convexity). Assume the mask M of the learner network is accurate. Suppose constants ε0, ε1 ∈ (0, 1) and the number of samples satisfies
N = Ω ( ε−21 K 4r̃ log q ) , (4)
for some large constant q > 0, where
r̃ = 1
8K4
(∑K k=1 ∑K j=1(1 + δj,k)(rj + rk) 1 2 )2 , (5)
δj,k is 1 if the indices of non-pruned weights in the j-th and k-th neurons overlap and 0 otherwise. Then, there exists a permutation matrix P such that for anyW that satisfies
‖W −W ∗P ‖F = O ( ε0 K2 ) , andM W = W , (6)
its Hessian of f̂D, with probability at least 1−K · q−rmin , is bounded as:
Θ (1− ε0 − ε1
K2
) I ∇2f̂D(W ) Θ ( 1 K ) I. (7)
Remark 1.1 (Parameter r̃): Clearly r̃ is a monotonically increasing function of any rj from (5). Moreover, one can check that 18rave ≤ r̃ ≤ rave. Hence, r̃ is in the order of rave. Remark 1.2 (Local landscape): Theorem 1 shows that with enough samples as shown in (4), in a local region ofW ∗P as shown in (6), all the eigenvalues of the Hessian matrix of the empirical risk function are lower and upper bounded by two positive constants. This property is useful in designing efficient algorithms to recoverW ∗P , as shown in Section 3.2.
Remark 1.3 (Size of the convex region): When the number of samples N is fixed and r changes, ε1 can be Θ( √ r̃/N) while (4) is still met. ε0 in (7) can be arbitrarily close to but smaller than 1− ε1 so that the Hessian matrix is still positive definite. Then from (6), the radius of the convex ball is Θ(1) − Θ( √ r̃/N), indicating an enlarged region when r̃ decreases. The enlarged convex region serves as an important component in proving the faster convergence rate, summarized in Theorem 2. Besides this, as Figure 1 shown in [20], the authors claim that the learning is stable if the linear interpolation of the learned models with SGD noises still remain similar in performance, which is summarized as the concept “linearly connected region.” Intuitively, we conjecture that the winning ticket shows a better performance in the stability analysis because it has a larger convex region. In the other words, a larger convex region indicates that the learning is more likely to be stable in the linearly connected region.

heading: 3.2 Convergence Analysis with Accelerated Gradient Descent
text: We propose to solve the non-convex problem (3) via the accelerated gradient descent (AGD) algorithm, summarized in Algorithm 1. Compared with the vanilla gradient descent (GD) algorithm, AGD has an additional momentum term, denoted by β(W (t) −W (t−1)), in each iteration. AGD enjoys a faster convergence rate than vanilla GD in solving optimization problems, including learning neural networks [65]. Vanilla GD can be viewed as a special case of AGD by letting β = 0. The initial point W (0) can be obtained through a tensor method, and the details are provided in Appendix B.
Algorithm 1 Accelerated Gradient Descent (AGD) Algorithm 1: Input: training data D = {(xn, yn)}Nn=1, gradient step size η, momentum parameter β, and an
initializationW (0) by the tensor initialization method; 2: Partition D into T = log(1/ε) disjoint subsets, denoted as {Dt}Tt=1; 3: for t = 1, 2, · · · , T do 4: W (t+1) = W (t) − η ·M ∇f̂Dt(W
(t)) + β(W (t) −W (t−1)) 5: end for 6: Return: W (T )
The theoretical analyses of our algorithm are summarized in Theorem 2 (convergence) and Lemma 1 (Initialization). The significance of these results can be interpreted from the following aspects.
1. Linear convergence to the oracle model: Theorem 2 implies that if initialized in the local convex region, the iterates generated by AGD converge linearly toW ∗P for some P when noiseless. When there is noise, they converge to a pointW (T ). The distance betweenW (T ) andW ∗P is proportional to the noise level and scales in terms of O( √ r̃/N). Moreover, when N is fixed, the convergence rate
of AGD is Θ( √ r̃/K). Recall that Algorithm 1 reduces to the vanilla GD by setting β = 0. The rate for the vanilla GD algorithm here is Θ( √ r̃/K) by setting β = 0 by Theorem 2, indicating a slower convergence than AGD. Lemma 1 shows the tensor initialization method indeed returns an initial point in the convex region.
2. Sample complexity for accurate estimation: We show that the required number of samples for successful estimation of the oracle model is Θ ( r̃ log q log(1/ε) ) for some large constant q and estimation accuracy ε. Our sample complexity is much less than the conventional bound of Θ(d log q log(1/ε)) for one-hidden-layer networks [69, 66, 67]. This is the first theoretical characterization of learning a pruned network from the perspective of sample complexity.
3. Improved generalization of winning tickets: We prove that with a fixed number of training samples, training on a properly pruned sub-network converges faster toW ∗P than training on the original dense network. Our theoretical analysis justifies that training on the winning ticket can meet or exceed the same test accuracy within the same number of iterations. To the best of our knowledge, our result here provides the first theoretical justification for this intriguing empirical finding of “improved generalization of winning tickets” by [20]. Theorem 2 (Convergence). Assume the maskM of the learner network is accurate. SupposeW (0) satisfies (6) and the number of samples satisfies
N = Ω ( ε−20 K 6r̃ log q log(1/ε) )
(8)
for some ε0 ∈ (0, 1/2). Let η = K/14 in Algorithm 1. Then, the iterates {W (t)}Tt=1 returned by Algorithm 1 converges linearly toW ∗ up to the noise level with probability at least 1−K2T · q−rmin
‖W (t) −W ∗P ‖F ≤ν(β)t‖W (0) −W ∗P ‖F +O (∑
j
√ rj log q
N
) · |ξ|, (9)
and ‖W (T ) −W ∗P ‖F ≤ε‖W ∗‖F +O (∑
j
√ rj log q
N
) · |ξ|, (10)
for a fixed permutation matrix P , where ν(β) is the rate of convergence that depends on β with ν(β∗) = 1−Θ ( 1−ε0√ K ) for some non-zero β∗ and ν(0) = 1−Θ ( 1−ε0 K ) . Lemma 1 (Initialization). Assume the noise |ξ| ≤ ‖W ∗‖2 and the number of samples N = Ω ( ε−20 K 5rmax log q )
for ε0 > 0 and large constant q, the tensor initialization method outputs W (0) such that (6) holds, i.e., ‖W (0) −W ∗‖F = O ( ε0σK K2 ) with probability at least 1− q−rmax .
Remark 2.1 (Faster convergence on pruned network): With a fixed number of samples, when r̃ decreases, ε0 can increase as Θ( √ r̃) while (8) is still met. Then ν(0) = Θ( √ r̃/K) and ν(β∗) =
Θ( √ r̃/K). Therefore, when r̃ decreases, both the stochastic and accelerated gradient descent
converge faster. Note that as long asW (0) is initialized in the local convex region, not necessarily by the tensor method, Theorem 2 guarantees the accurate recovery. [66, 67] analyze AGD on convolutional neural networks, while this paper focuses on network pruning.
Remark 2.2 (Sample complexity of initialization): From Lemma 1, the required number of samples for a proper initialization is Ω ( ε−20 K 5rmax log q ) . Because rmax ≤ Krave and r̃ = Ω(rave), this number is no greater than the sample complexity in (8). Thus, provided that (8) is met, Algorithm 1 can estimate the oracle network model accurately.
Remark 2.3 (Inaccurate mask): The above analyses are based on the assumption that the mask of the learner network is accurate. In practice, a mask can be obtained by an iterative pruning method such as [20] or a one-shot pruning method such as [55]. In Appendix E, we prove that the magnitude pruning method can obtain an accurate mask with enough training samples. Moreover, empirical experiments in Section 4.2 and 4.3 suggest that even if the mask is not accurate, the three properties (linear convergence, sample complexity with respect to the network size, and improved generalization of winning tickets) can still hold. Therefore, our theoretical results provide some insight into the empirical success of network pruning.

heading: 3.3 The Sketch of Proofs and Technical Novelty
text: Our proof outline is inspired by [69] on fully connected neural networks, however, major technical changes are made in this paper to generalize the analysis to an arbitrarily pruned network. To characterize the local convex region of f̂D (Theorem 1), the idea is to bound the Hessian matrix of the population risk function, which is the expectation of the empirical risk function, locally and then characterize the distance between the empirical and population risk functions through the concentration bounds. Then, the convergence of AGD (Theorem 2) is established based on the desired local curvature, which in turn determines the sample complexity. Finally, to initialize in the local convex region (Lemma 1), we construct tensors that contain the weights information and apply a decomposition method to estimate the weights.
Our technical novelties are as follows. First, a direct application of the results in [69] leads to a sample complexity bound that is linear in the feature dimension d. We develop new techniques to tighten the sample complexity bound to be linear in r̃, which can be significantly smaller than d for a sufficiently pruned network. Specifically, we develop new concentration bounds (Lemmas 4 and 5 in Appendix) to bound the distance between the population and empirical risk functions rather than using the bound in [69]. Second, instead of restricting the acitivation to be smooth for convergence analysis, we study the case of ReLU function which is non-smooth. Third, new tensors are constructed for pruned networks (see (21)-(23) in Appendix) in computing the initialization, and our new concentration bounds are employed to reduce the required number of samples for a proper initialization. Last, Algorithm 1 employs AGD and is proved to converge faster than the GD algorithm in [69].

heading: 3.4 Limitations
text: Like most theoretical works based on the oracle-learner setup, limitations of this work include (1) one hidden layer only; and (2) the input follows the Gaussian distribution. Extension to multi-layer might be possible if the following technical challenges are addressed. First, when characterizing the local convex region, one needs to show that the Hessian matrix is positive definite. In multi-layer networks, the Hessian matrix is more complicated to compute. Second, new concentration bounds need to be developed because the input feature distributions to the second and third layers depend on the weights in previous layers. Third, the initialization approach needs to be revised. The team is also investigating the other input distributions such as Gaussian mixture models.

heading: 4 Numerical Experiments
text: The theoretical results are first verified on synthetic data, and we then analyze the pruning performance on both synthetic and real datasets. In Section 4.1, Algorithm 1 is implemented with minor modification, such that, the initial point is randomly selected as ‖W (0) −W ∗‖F /‖W ∗‖F < λ for some λ > 0 to reduce the computation. Algorithm 1 terminates when ‖W (t+1)−W (t)‖F /‖W (t)‖F is smaller than 10−8 or reaching 10000 iterations. In Sections 4.2 and 4.3, the Gradient Signal Preservation (GraSP) algorithm [55] and IMP algorithm [10, 20]2 are implemented to prune the neural networks. As many works like [11, 10, 20] have already verified the faster convergence and better generalization accuracy of the winning tickets empirically, we only include the results of some representative experiments, such as training MNIST and CIFAR-10 on Lenet-5 [32] and Resnet-50 [27] networks, to verify our theoretical findings.
The synthetic data are generated using a oracle model in Figure 1. The input xn’s are randomly generated from Gaussian distribution N (0, Id×d) independently, and indices of non-pruned weights of the j-th neuron are obtained by randomly selecting rj numbers without replacement from [d]. For the convenience of generating specific r̃, the indices of non-pruned weights are almost overlapped ( ∑ j ∑ k δjδk > 0.95K
2) except for Figure 5. In Figures 2 and 4, rj is selected uniformly from [0.9r̃, 1.1r̃] for a given r̃, and rj are the same in value for all j in other figures. Each non-zero entry ofW ∗ is randomly selected from [−0.5, 0.5] independently. The noise ξn’s are i.i.d. from N (0, σ2), and the noise level is measured by σ/Ey , where Ey is the root mean square of the noiseless outputs.
2The source codes used are downloaded from https://github.com/VITA-Group/CV_LTH_Pre-training.

heading: 4.1 Evaluation of theoretical findings on synthetic data
text: Local convexity near W ∗. We set the number of neurons K = 10, the dimension of the data d = 500 and the sample size N = 5000. Figure 2 illustrates the success rate of Algorithm 1 when r̃ changes. The y-axis is the relative distance of the initializationW (0) to the ground-truth. For each pair of r̃ and the initial distance, 100 trails are constructed with the network weights, training data and the initializationW (0) are all generated independently in each trail. Each trail is called successful if the relative error of the solutionW returned by Algorithm 1, measured by ‖W −W ∗‖F /‖W ∗‖F , is less than 10−4. A black block means Algorithm 1 fails in estimatingW ∗ in all trails while a white block indicates all success. As Algorithm 1 succeeds ifW (0) is in the local convex region nearW ∗, we can see that the radius of convex region is indeed linear in −r̃ 12 , as predicted by Theorem 1. Convergence rate. Figure 3 shows the convergence rate of Algorithm 1 when r̃ changes. N = 5000, d = 300, K = 10, η = 0.5, and β = 0.2. Figure 3(a) shows that the relative error decreases exponentially as the number of iterations increases, indicating the linear convergence of Algorithm 1. As shown in Figure 3(b), the results are averaged over 20 trials with different initial points, and the areas in low transparency represent the standard deviation errors. We can see that the convergence rate is almost linear in √ r̃, as predicted by Theorem 2. We also compare with GD by setting β as 0. One can see that AGD has a smaller convergence rate than GD, indicating faster convergence.
10 12 14 16 18 20 2 6
10 14 18 22 26 30 34 38
Figure 2: The radius of the local convex region against r̃ 1 2
Sample complexity. Figures 4 and 5 show the success rate of Algorithm 1 when varying N and r̃. d is fixed as 100. In Figure 4, we construct 100 independent trails for each pair of N and r̃, where the ground-truth model and training data are generated independently in each trail. One can see that the required number of samples for successful estimation is linear in r̃, as predicted by (8). In Figure 5, rj is fixed as 20 for all neurons, but different network architectures after pruning are considered. One can see that although the number of remaining weights is the same, r̃ can be different in different architectures, and the sample complexity increases as r̃ increases, as predicted by (8).
r̃
r̃ 1 2 at different noise level
Performance in noisy case. Figure 6 shows the relative error of the learned model by Algorithm 1 from noisy measurements when r̃ changes. N = 1000, K = 10, and d = 300. The results are averaged over 100 independent trials, and standard deviation is around 2% to 8% of the corresponding relative errors. The relative error is linear in r̃ 1 2 , as predicted by (9). Moreover, the relative error is proportional to the noise level |ξ|.

heading: 4.2 Performance with inaccurate mask on synthetic data
text: The performance of Algorithm 1 is evaluated when the maskM of the learner network is inaccurate. The number of neurons K is 5. The dimension of inputs d is 100. r∗j of the oracle model is 20 for
all j ∈ [K]. GraSP algorithm [55] is employed to find masks based only on early-trained weights in 20 iterations of AGD. The mask accuracy is measured by ‖M∗ M‖0/‖M∗‖0, whereM∗ is the mask of the oracle model. The pruning ratio is defined as (1− rave/d)× 100%. The number of training samples N is 200. The model returned by Algorithm 1 is evaluated on Ntest = 105 samples, and the test error is measured by √∑ n |yn − ŷn|2/Ntest, where ŷn is the output of the learned model with the input xn, and (xn, yn) is the n-th testing sample generated by the oracle network.
Improved generalization by GraSP. Figure 7 shows the test error with different pruning ratios. For each pruning ratio, we randomly generate 1000 independent trials. Because the mask of the learner network in each trail is generated independently, we compute the average test error of the learned models in all the trails with same mask accuracy. If there are less than 10 trails for certain mask accuracy, the result of that mask accuracy is not reported as it is statistically meaningless. The test error decreases as the mask accuracy increases. More importantly, at fixed mask accuracy, the test error decreases as the pruning ratio increases. That means the generalization performance improves when r̃ deceases, even if the mask is not accurate.
Linear convergence. Figure 8 shows the convergence rate of Algorithm 1 with different pruning ratios. We show the smallest number of iterations required to achieve a certain test error of the learned model, and the results are averaged over the independent trials with mask accuracy between 0.85 and 0.90. Even with inaccurate mask, the test error converges linearly. Moreover, as the pruning ratio increases, Algorithm 1 converges faster.
Sample complexity with respect to the pruning ratio. Figure 9 shows the test error when the number of training samples N changes. All the other parameters except N remain the same. The results are averaged over the trials with mask accuracy between 0.85 and 0.90. We can see the test error decreases when N increases. More importantly, as the pruning ratio increases, the required number of samples to achieve the same test error (no less than 10−3) decreases dramatically. That means the sample complexity decreases as r̃ decreases even if the mask is inaccurate.

heading: 4.3 Performance of IMP on synthetic, MNIST and CIFAR-10 datasets
text: We implement the IMP algorithm to obtain pruned networks on synthetic, MNIST and CIFAR-10 datasets. Figure 10 shows the test performance of a pruned network on synthetic data with different sample sizes. Here in the oracle network model, K = 5, d = 100, and r∗j = 20 for all j ∈ [K]. The noise level σ/Ey = 10−3. One observation is that for a fixed sample size N greater than 100, the test error decreases as the pruning ratio increases. This verifies that the IMP algorithm indeed prunes the network properly. It also shows that the learned model improves as the pruning progresses, verifying our theoretical result in Theorem 2 that the difference of the learned model from the oracle model decreases as rj decreases. The second observation is that the test error decreases as N increases for any fixed pruning ratio. This verifies our result in Theorem 2 that the difference of the learned model from the oracle model decreases as the number of training samples increases. When the pruning ratio is too large (greater than 80%), the pruned network cannot explain the data properly, and thus the test error is large for all N . When the number of samples is too small, like N = 100, the test error is always large, because it does not meet the sample complexity requirement for estimating the oracle model even though the network is properly pruned.
Figures 11 and 12 show the test performance of learned models by implementing the IMP algorithm on MNIST and CIFAR-10 using Lenet-5 [32] and Resnet-50 [27] architecture, respectively. The
experiments follow the standard setup in [10] except for the size of the training sets. To demonstrate the effect of sample complexity, we randomly selected N samples from the original training set without replacement. As we can see, a properly pruned network (i.e., winning ticket) helps reduce the sample complexity required to reach the test accuracy of the original dense model. For example, training on a pruned network returns a model (e.g., P1 and P3 in Figures 11 and 12) that has better testing performance than a dense model (e.g., P2 and P4 in Figures 11 and 12) trained on a larger data set. Given the number of samples, we consistently find the characteristic behavior of winning tickets: That is, the test accuracy could increase when the pruning ratio increases, indicating the effectiveness of pruning. The test accuracy then drops when the network is overly pruned. The results show that our theoretical characterization of sample complexity is well aligned with the empirical performance of pruned neural networks and explains the improved generalization observed in LTH.

heading: 5 Conclusions
text: This paper provides the first theoretical analysis of learning one-hidden-layer pruned neural networks, which offers formal justification of the improved generalization of winning ticket observed from empirical findings in LTH. We characterize analytically the impact of the number of remaining weights in a pruned network on the required number of samples for training, the convergence rate of the learning algorithm, and the accuracy of the learned model. We also provide extensive numerical validations of our theoretical findings.
Broader impacts
We see no ethical or immediate societal consequence of our work. This paper contributes to the theoretical foundation of both network pruning and generalization guarantee. The former encourages the development of learning method to reduce the computational cost. The latter increases the public trust in incorporating AI technology in critical domains.

heading: Acknowledgement
text: This work was supported by AFOSR FA9550-20-1-0122, ARO W911NF-21-1-0255, NSF 1932196 and the Rensselaer-IBM AI Research Collaboration (http://airc.rpi.edu), part of the IBM AI Horizons Network (http://ibm.biz/AIHorizons). We thank Tianlong Chen at University of Texas at Austin, Haolin Xiong at Rensselaer Polytechnic Institute and Yihua Zhang at Michigan State University for the help in formulating numerical experiments. We thank all anonymous reviewers for their constructive comments.

Review summaries:
The paper analyzes the possibility of recovering the weight of a teacher network using a student network which is as dense as the original one. The authors prove that with enough samples drawn from the teacher network, the recovery can be done on one-hidden-layer neural networks with linear convergence speed. The authors also discuss the correlation between the needed samples and the sparsity of the teacher network.  To sum up, the well-designed structure makes the workﬂow clear and easy to follow, but the further analysis and discussion are expected to clarify some contributions in the techniques as well as in the evaluation section.

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.

This work takes a theoretical view of LTH, leveraging the geometric structure of the objective function to analyze the generalization error of a pruned network trained in a teacher-student fashion. In particular, they prove that, as a model is pruned, the desirable (convex) region with high-guaranteed generalization performance enlarges, providing explanation for improved performance of winning tickets. Then, the sample complexity of training the pruned network to achieve zero generalization error is analyzed, finding that the number of samples required is proportional to the number of un-pruned weights. Interestingly, the work finds that pruned models enjoy faster convergence to high performance, providing another possible explanation of why winning tickets outperform dense networks. 

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