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

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 $\ell_{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.

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 $\beta({\boldsymbol{W}}^{(t)}-{\boldsymbol{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 $\beta=0$. The initial point ${\boldsymbol{W}}^{(0)}$ can be obtained through a tensor method, and the details are provided in Appendix.

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.

Remark 2.1 (Faster convergence on pruned network): With a fixed number of samples, when $\widetilde{r}$ decreases, $\varepsilon_{0}$ can increase as $\Theta(\sqrt{\widetilde{r}})$ while (8) is still met. Then $\nu(0)=\Theta(\sqrt{\widetilde{r}}/K)$ and $\nu(\beta^{*})=\Theta(\sqrt{\widetilde{r}/K})$. Therefore, when $\widetilde{r}$ decreases, both the stochastic and accelerated gradient descent converge faster. Note that as long as ${\boldsymbol{W}}^{(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 $\Omega\big{(}\varepsilon_{0}^{-2}K^{5}r_{\max}\log q\big{)}$. Because $r_{\max}\leq Kr_{\text{ave}}$ and $\widetilde{r}=\Omega(r_{\text{ave}})$, 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, 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.

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 $\hat{f}_{\mathcal{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 $\tilde{r}$, which can be significantly smaller than $d$ for a sufficiently pruned network. Specifically, we develop new concentration bounds 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 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].

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.

Local convexity near ${\boldsymbol{W}}^{*}$. We set the number of neurons $K=10$, the dimension of the data $d=500$ and the sample size $N=5000$. Figure 3 illustrates the success rate of Algorithm 1 when $\widetilde{r}$ changes. The $y$-axis is the relative distance of the initialization ${\boldsymbol{W}}^{(0)}$ to the ground-truth. For each pair of $\widetilde{r}$ and the initial distance, $100$ trails are constructed with the network weights, training data and the initialization ${\boldsymbol{W}}^{(0)}$ are all generated independently in each trail. Each trail is called successful if the relative error of the solution ${\boldsymbol{W}}$ returned by Algorithm 1, measured by $\|{\boldsymbol{W}}-{\boldsymbol{W}}^{*}\|_{F}/\|{\boldsymbol{W}}^{*}\|_{F}$, is less than $10^{-4}$. A black block means Algorithm 1 fails in estimating ${\boldsymbol{W}}^{*}$ in all trails while a white block indicates all success. As Algorithm 1 succeeds if ${\boldsymbol{W}}^{(0)}$ is in the local convex region near ${\boldsymbol{W}}^{*}$, we can see that the radius of convex region is indeed linear in $-\widetilde{r}^{\frac{1}{2}}$, as predicted by Theorem 1.

Convergence rate. Figure 3 shows the convergence rate of Algorithm 1 when $\widetilde{r}$ changes. $N=5000$, $d=300$, $K=10$, $\eta=0.5$, and $\beta=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 $1/\sqrt{\widetilde{r}}$, as predicted by Theorem 2. We also compare with GD by setting $\beta$ as 0. One can see that AGD has a smaller convergence rate than GD, indicating faster convergence.

Sample complexity. Figures 6 and 6 show the success rate of Algorithm 1 when varying $N$ and $\widetilde{r}$. $d$ is fixed as $100$. In Figure 6, we construct $100$ independent trails for each pair of $N$ and $\widetilde{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 $\widetilde{r}$, as predicted by (8). In Figure 6, $r_{j}$ 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, $\widetilde{r}$ can be different in different architectures, and the sample complexity increases as $\widetilde{r}$ increases, as predicted by (8).

Performance in noisy case. Figure 6 shows the relative error of the learned model by Algorithm 1 from noisy measurements when $\widetilde{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 $\widetilde{r}^{\frac{1}{2}}$, as predicted by (9). Moreover, the relative error is proportional to the noise level $|\xi|$.

The performance of Algorithm 1 is evaluated when the mask ${\boldsymbol{M}}$ 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\in[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 $\|{\boldsymbol{M}}^{*}\odot{\boldsymbol{M}}\|_{0}/\|{\boldsymbol{M}}^{*}\|_{0}$, where ${\boldsymbol{M}}^{*}$ is the mask of the oracle model. The pruning ratio is defined as $(1-{r_{\textrm{ave}}}/{d})\times 100\%$. The number of training samples $N$ is $200$. The model returned by Algorithm 1 is evaluated on $N_{\text{test}}=10^{5}$ samples, and the test error is measured by $\sqrt{\sum_{n}|y_{n}-\hat{y}_{n}|^{2}/N_{\text{test}}}$, where $\hat{y}_{n}$ is the output of the learned model with the input ${\boldsymbol{x}}_{n}$, and $({\boldsymbol{x}}_{n},y_{n})$ is the $n$-th testing sample generated by the oracle network.

Improved generalization by GraSP. Figure 9 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 $\widetilde{r}$ deceases, even if the mask is not accurate.

Linear convergence. Figure 9 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 $\widetilde{r}$ decreases even if the mask is inaccurate.

We implement the IMP algorithm to obtain pruned networks on synthetic, MNIST and CIFAR-10 datasets. Figure 12 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\in[K]$. The noise level $\sigma/E_{y}=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 $r_{j}$ 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 12 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., $P_{1}$ and $P_{3}$ in Figures 12 and 12) that has better testing performance than a dense model (e.g., $P_{2}$ and $P_{4}$ in Figures 12 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.