Analyzing Lottery Ticket Hypothesis from PAC-Bayesian Theory Perspective
Abstract
The lottery ticket hypothesis (LTH) has attracted attention because it can explainwhy over-parameterized models often show high generalization ability.
It is knownthat when we use iterative magnitude pruning (IMP), which is an algorithm to findsparse networks with high generalization ability that can be trained from the initialweights independently, called winning tickets , the initial large learning rate doesnot work well in deep neural networks such as ResNet.
However, since the initiallarge learning rate generally helps the optimizer to converge to flatter minima,we hypothesize that the winning tickets have relatively sharp minima, which isconsidered a disadvantage in terms of generalization ability.
In this paper, weconfirm this hypothesis and show that the PAC-Bayesian theory can provide anexplicit understanding of the relationship between LTH and generalization behavior.
On the basis of our experimental findings that IMP with a small learning rate findsrelatively sharp minima and that the distance from the initial weights is deeplyinvolved in winning tickets, we offer the PAC-Bayes bound using a spike-and-slabdistribution to analyze winning tickets.
Finally, we revisit existing algorithms forfinding winning tickets from a PAC-Bayesian perspective and provide new insightsinto these methods.
Introduction
The high generalization ability of modern neural networks can be attributed to the heavier overparameterization and effective learning algorithms [22, 30, 41].
This increase in the number of parametersleads to high computational cost and high memory usage and network pruning is one of the effective techniques for addressing these problems [12, 14, 23].
After pruning a significant number ofparameters, the pruned network can often work well with little or no accuracy loss.
However, trainingthis sparse subnetwork independently from the initial weights often does not work and we can onlyobtain the sparse subnetwork through pruning after training the whole network.
Frankle and Carbin [10] presented “Lottery Ticket Hypothesis (LTH)” that states the existence ofwinning tickets: small but critical noise and the ability to find the winning ticket and empirically showed that the large learning ratemoves weights too much under the low-stability learning process.
In this paper, we first empirically show that winning tickets are actually more vulnerable to label noisesetting compared to the subnetwork created with the large learning rate; that is, the generalizationability of winning tickets is degraded due to the learning rate constraint.
To explain this, we nextapply the PAC-Bayesian theory to LTH and show that it can explain the relationship between LTHand generalization behavior.
We use the PAC-Bayes bound for a spike-and-slab distribution toanalyze winning tickets, which is based on our experimental findings that reducing the expectedsharpness restricted to an unpruned parameter space and adding the regularization of distance fromthe initial weights can enhance the test performance of winning tickets.
Finally, we revisit existingalgorithms such as IMP, continuous sparsification [34] from the point of view of the PAC-Bayesbound optimization.
This consideration gives an interpretation of these methods as an approximationof bound optimization.
To sum up, our contributions are as follows.
• We experimentally show that IMP with a small learning rate finds relatively sharper minimaand that the distance from the initial weights is critical for IMP,
i
e., balancing the distanceand the training error helps to find them.•
On the basis of our findings, we reveal that the PAC-Bayesian formulation for a spike-andslab distribution effectively captures the winning tickets behavior.•
We revisit the existing algorithms from the PAC-Bayesian perspective and explain theirbehavior.
Related Work
Learning Rate
The initial large learning rate often improves generalization of deep neural networks[19, 25].
The large learning rate generally helps an optimizer to converge to flatter minima [24, 38],and these flatter minima are advantageous for generalization ability rather than sharper minima[18, 20, 21].
The relationship between flatness and generalization can be considered from a PACBayesian perspective [5, 31, 32] (see Appendix A).
Frankle et al.
[11] provided the insights in why
IMP with the large learning rate fails into find winning tickets in some problem settings.
Theyproposed IMP with rewinding to an early epoch to avoid the very early training process becausetraining is not stable to SGD noise in the early training regime (see Appendix A).
Empirical Results on LTH Some studies have also investigated the flatness and the distance fromthe initial weights on LTH.
Bain [2] plotted the loss landscape of winning tickets visually and foundthat IMP produces more convex and sharper minimum relative to random pruning.
Bartoldson et al.[3] found that pruning regularizes similarly to noise injection, and they discussed the generalizationof pruned networks considering flatness.
They measured flatness by using the trace of Hessian andgradient covariance matrix.
He et al.
[17] refers the distance from the initial weights to analyze thelabel noise robustness of winning tickets; they stated that the influence of the label noise is lessenedif the distance is suppressed.
Liu et al.
[26] discussed the relationship between winning tickets andthe learning rate considering the similarity between initial and trained weights in terms of pruningmask overlap instead of the distance from the initial weights.
There is a recent study on finding amask that shows good accuracy without training weights at all [33, 43].
Theoretical Results on LTH Malach et al.
[29] demonstrated a subnetwork with a comparableaccuracy to the original network in a sufficiently over-parameterized network, without any training.
Zhang et al.
[42] analyzed the generalization capability of winning tickets based on sample complexity.
Although it is limited to the case of a one-hidden-layer neural network, they provided insight intowhy the sparsity increases the generalization ability.
For a PAC-Bayesian theory, Hayou et al.
[15]also used spike-and-slab prior and posterior distributions.
However, their motivation and purpose
[Figure]
are completely different from our work.
Their aim was to obtain a sophisticated pruning mask byoptimizing the PAC-Bayes bound, and this is mainly in the context of network pruning rather than
LTH.
Our work differs in that we use it to analyze the generalization behavior of a given winningticket on the basis of our empirical findings on the learning rate, the flatness, and the distance fromthe initial weights.
In addition to these, we also discuss the relationship with existing algorithms.
Empirical Analysis
We first show some empirical results because our new findings about winning tickets in this empiricalanalysis motivate the PAC-Bayesian analysis for LTH; thus, we interpret the results on the basis ofthe PAC-Bayesian perspective in the next section.
We empirically investigated the properties of winning tickets mainly related to the learning rate.
First,we show that winning tickets are vulnerable to label noise due to a small learning rate and that IMPwith a small learning rate finds relatively sharp minima. We also focused on the distance from theinitial weights and show that the regularization of this distance, rather than the regularization of thenorm, is important in the discovery of winning tickets.
We followed the experimental setting of
Frankle and Carbin [10] and used a modified version of OpenLTH repository [9].
3.1 Vulnerability to label noise
We examined the test accuracy when some fractions of the labels in the training set are randomlyflipped to see whether the generalization behavior of winning tickets is degraded by the constraint thatthe learning rate must be small.
Figure 1 shows the test accuracy on clean and label noise datasets ofsparse subnetworks produced by IMP with different learning rates.
As for the no label noise setting(green line), there is an accuracy drop at some point as the learning rate is increased.
The subnetworkeventually performs worse than the original unpruned network, which means that IMP ultimatelyfails to find the winning ticket.
In contrast, the test accuracy generally increases as the learning rateincreases in the high label noise setting (orange line).
The test accuracy continues to improve evenwhen it is not the winning ticket.
To sum up, the large learning rate is not suitable for finding winningtickets under a clean dataset but is advantageous in the high label noise setting.
3.2 Flatness
Now that we have seen that winning tickets have undesirable properties due to the small learning rate,we investigated whether this result comes from the difference in flatness around the found solution.
There are many previous studies that discuss the relationship between the learning rate and flatness[24, 38].
In this experiment, perturbations are added only to unpruned weights to consider flatness onpruned networks.
We will discuss this justification in detail in Section 4.
[Figure]
T r a i n i n g
L o ss
The trace of Hessian is calculated by PyHessian [40].
We averaged over three different subnetworks.
ResNet
VGG
[Table]
First, we visualized the loss landscape shape of subnetworks produced by IMP.
Figure 2 shows the1-d loss landscape of the subnetwork that has a high sparsity.
This landscape shows the test accuracyaround the trained weights adding perturbation restricted to unpruned parameters.
We can see that theparameters of the winning ticket (small learning rate) are in the sharper minimum compared with thelarge learning rate.
The large learning rate can find a flatter minimum; however, it is not the winningticket.
This graph is a fixed sparsity loss landscape; therefore, it is not possible to discuss whether ornot the subnetwork is a winning ticket from this graph.
There is actually a accuracy drop in the largelearning rate setting when sparsity is changed.
Given that IMP requires a small learning rate rather than a large one, these results suggest twopossible interpretations; 1) sharp minimum is essential to find winning tickets, therefore the largelearning rate fails in IMP, 2) sharp minimum is simply the result of small learning rate training, andflat minimum is better for winning tickets if possible.
To investigate these possibilities, we usedsharpness-aware minimization (SAM) [8] and neural variable risk minimization (NVRM) [37] tosearch for parameters that lie in neighborhoods having uniformly low loss.
They differ in that SAMminimizes the maximum loss in the neighborhood, whereas NVRM minimizes the expected lossin the neighborhood.
We used both of them to compare with normal SGD.
These optimizers arebased on SGD with the same setting as the original LTH paper, and the noise considered in them islimited to the unpruned parameters.
Figure 2 also shows the loss landscape when SAM is used; Itcan actually reach a flatter minimum.
The loss landscape is as flat as the large learning rate setting;however the test accuracy is higher than that of the large learning rate.
Next, we investigated the trace of Hessian to analyze the flatness of winning tickets created by theseoptimizers.
The trace of Hessian is used as a measure of flatness in the prior work [4, 21, 39].
Table [-
-]1shows that SGD with the small learning rate finds relatively a sharper minimum and that IMP can find
[Figure]
a flatter one by using SAM or NVRM.
These optimizer can find flatter minima, which improves thetest accuracy to some extent compared with the SGD with the small learning rate (see Appendix B.1).
We also found that the large learning rate still cannot find winning tickets even though we use SAMinstead of SGD (see Appendix B.2).
The large learning rate has already found relatively flat minima;therefore, the results do not change by using SAM.
This fact implies that the flatness of solutionsfound by SGD with a large learning rate and by SAM with a small learning rate are similar, butwinning tickets cannot be found with large learning rate because they find different solutions in someother properties.
Next, we analyze this difference by focusing on the distance from the initial weights.
3.3 Distance from initial weights
As a reason for the small learning rate constraint, we hypothesized that winning tickets can only befound in IMP within a range not far from the initial weights.
In order to confirm this, we run IMPsuppressing this distance and compare it with the usual regularization by the L 2 norm regularization.
We empirically analyzed that the winning ticket can be obtained via the regularization of this distanceeven under the large learning rate setting.
Let L S be the empirical risk on the training samples S , λbe a regularization hyperparameter, θ be network weights, and m be a pruning mask.
We designedthe loss function with l 2 _ init and l 2 _ norm , respectively, as follows.
[Equation]
Figure 3 shows the test accuracy with different regularizations by changing the hyperparameter λ .
We plot 90 % and 95 % sparsity subnetworks and the original unpruned network as a baseline.
Ifthe subnetwork accuracy is close to the whole network accuracy, it is considered to be successfulin finding the winning ticket.
As discussed previously, the large learning rate fails to find winningtickets unlike the small learning rate setting; however adding the regularization from the initialweights changes this trend.
In Figure 3, λ around 5 × 10 − 4 shows that a winning ticket is found sincethe accuracy drop from the whole network is suddenly reduced, and there is no such a trend when
L 2 norm regularization is added.
This means that IMP can obtain winning tickets suppressing thedistance from the initial weights even with the large learning rate.
(Other settings in Appendix B.4.)
We can also obtain interesting results when the learning rate is small.
The accuracy drop is smallwhile λ is small; however it becomes large increasing λ .
In the case of l 2 _ init , it is possible that the gap widens because sparse networks are more affected by strong regularization (it is the samefor large learning rate setting), but in the case of l 2 _ norm , the gap widens significantly even thoughgeneralization ability increases because of regularization.
The strong norm regularization makes itfail to find a winning ticket even if a winning ticket could be found originally.
These results are related to the prior-mean selection from the PAC-Bayesian perspective.
In thetraining of a normal network, there is a trade-off between suppressing the parameter norm, i.e.,reducing the KL term, and reducing the training loss, and it is important to ensure a balance betweenthem.
As for IMP, there is a specific situation where suppressing the parameter norm makes thetraining loss large due to the failure to find a winning ticket.
If we take the initial weights as a priormean instead of 0 , the training loss can decrease by suppressing the KL term when the trained weightsare far from the initial weights.
This experiment corresponds to what to take as a prior mean in termsof minimizing PAC-Bayes bound, and this result indicates that setting the initial weights to a priormean seems to be compatible with minimizing the PAC-Bayes bound in the case of winning tickets.
PAC Bayesian Analysis on LTH
First, we present some possible definitions of subnetwork flatness and consider a PAC-Bayes boundof a spike-and-slab formulation based on our experiments.
Next, we show that this bound capturesthe generalization behavior of winning tickets and revisit existing algorithms from the perspective ofthis bound.
4.1 Flatness on winning tickets
While we could simply consider the neighborhood of the trained weights in an unpruned neuralnetwork, it is not trivial to define the flatness of pruned networks depending on how the prunedweights are taken into consideration.
The possible measure of the expected sharpness are as follows.
Add noise to the parameters restricted on unpruned parameter space.
Add noise to the parameters including the pruned weights.
Recover pruned weights and train the whole network to convergence (re-dense training [13])and measure its flatness.
Measure 2 is the same as the unpruned neural network, but the sparse weights in the whole parameterspace can no longer be in the local minimum; thus, it is uncertain whether flatness has any meaningin such a setting.
He et al.
[17] conducted re-dense training and showed that solutions at high sparsityare no longer minimizers in high dimensions.
They also found that the winning ticket has highersharpness than the original network based on Measure 3 and concluded that highly sparse solutions donot stick around the flat basins of minimizers.
However, none of these metrics has any justification.
It is known that flatness can be viewed from a PAC-Bayesian perspective [5, 31, 32].
We use a
PAC-Bayes bound of a spike-and-slab formulation, where expected sharpness corresponds to Measure1.
We will also compare this formulation with the normal Gaussian formulation (Measure 2) throughnumerical experiments and show that the spike-and-slab formulation captures the generalizationbehavior of winning tickets better.
This supports our findings that using SAM and NVRM on prunednetworks can enhance the generalization accuracy of winning tickets.
4.2 Spike-and-slab formulation
There are several problems when it comes to using the Gaussian distribution as prior and posterior foranalyzing winning ticket properties.
As discussed in Section 3.3, the distance from the initial weightsof the winning ticket is expected to be small.
We set the prior mean as the initial weights to takeadvantage of this property; however the original PAC Bayesian formulation based on the Gaussiandistribution has the disadvantage that the pruned weights have weight 0 and the norm of the initialweights corresponding to these weights remain in the KL part.
This not only results in a large bound but also behaves contrary to the purpose of getting sparse subnetworks when the bound is optimizedbecause the more sparse the subnetwork are, the larger the bound gets.
In addition, noise is inevitablyadded to the pruned weights when considering expected sharpness since the variance of the prunedpart cannot be set to zero.
In order to limit the distance from the initial weights and noise added inexpected sharpness only to the unpruned weights, we use a spike-and-slab distribution, which is themixture of the Gaussian distribution N and Dirac delta distribution δ , in the PAC-Bayesian bound.
Let σ p and σ q be vectors whose i th element is a variance of Gaussian distribution, λ p and λ q bevectors that represents the mixture ratio of prior and posterior, respectively. θ represents the networkparameter, and θ init is the initial weights and ¯ θ is the trained weights.
We design the prior P and posterior Q as follows.
[Equation]
The KL divergence of the spike-and-slab distribution [36] can be calculated as KL [Equation] [Equation]
P
[Equation]
where
[Equation]
Since the structure of the pruned network is given, we set the element of λ q,i to 0 or 1 asymptoticallyaccording to the pruning mask m and obtain the following kl divergence about λ .
This operation hasbeen conventionally done in the entropy discussion [28].
[Equation]
4.3 Numerical Experiment
We conducted numerical experiments to confirm that our PAC-Bayes formulation can adequatelyexplain the behavior of winning tickets.
We optimized the posterior variance to minimize the PACBayes bound, and plot the training loss above the posterior and KL term to investigate that thebound can capture the test accuracy of the subnetworks produced by IMP.
As a comparison, we alsoexperiment with the Gaussian distribution setting using the zero-mean prior.
The PAC-Bayesianbound used here is the variational KL bound [6] (Theorem D.2, and see also Appendix B.7), and weimplemented the code following Dziugaite et al.
Figure 4 shows the distribution of the training risk term and KL term when optimizing the PAC-Bayesbound, and the actual test accuracy is colored.
We show that the bound with the spike-and-slabformulation successfully explains the behavior of winning tickets, dividing the point cloud into threegroups:
A) winning tickets (moderate learning rate), B) subnetworks that failed to find winningtickets (too high learning rate) and C) not much trained subnetworks (too small learning rate).
In the left figure, as the learning rate is increased, the distance from the initial weights increases andthe training risk gradually decreases from C to A; The same trend also appears on the right figure.
On the other hand, the distribution of B differs greatly between left and right.
The right figure doesnot capture the test accuracy well because B should have higher test accuracy considering its KLterm and training risk term.
This means that, as the learning rate is increased and winning ticketscan no longer be found, the KL term when prior mean is zero becomes much smaller than that of A.
This is because many parameters become close to zero rather than near their initial weights when thewinning ticket fails to be found (see Appendix B.5).
In terms of not only the bound optimization butalso the analysis of the existing winning ticket, it is preferable to use the spike-and-slab distributionto set the prior mean as the initial weights. [Table]
4.4 Revisiting existing algorithms
We reconsider the existing algorithms for winning ticket searching from the perspective of optimizingthe PAC-Bayes bound.
Although we focus on IMP and continuous sparsification, this view could behelpful for other methods as well.
4.4.1 IMP
IMP is a heuristic method and does not have an explicit target function.
Here, we explain how IMPbehaves in the sense of a PAC-Bayes bound with our formulation instead of viewing IMP as a directbound optimization problem.
The risk term and KL term in the PAC-Bayes bound are basically in a trade-off relationship; choosinga complex model to fit the training data will increase the KL term, while choosing a simple modelmay not have high accuracy on the training data.
We point out that the two steps of IMP: 1) train thesubnetwork, 2) prune the subnetwork and revert its weights, optimize the overall bound by alternatelyreducing one term while suppressing the increase in the other term.
In the first step , IMP trains the subnetwork from initial weights under a given pruning mask.
Thisprocess reduces the training risk, and the increase in KL term is not expected to be so large in ourformulation because the prior mean is set to the initial weights and the trained weights of IMP are notfar from the initial weights as shown in Sections 3.3 and 4.3.
In the second step , IMP prunes a certain percentage of the smallest magnitude weights and revertsthe trained weights to the initial state.
Reverting θ to θ init and changing part of λ q from 1 to 0make the KL term small.
The number of Gaussian KL summations decreases and the distance frominitialization gets to 0 , and the λ KL part is also reduced if the prior mixture ratio λ p is set to the finaltarget sparsity.
This KL reduction is not dependent on the pruning criterion.
The problem here ishow to minimize the increase in training loss, which is related to what heuristic pruning criterion wechoose and why pruning weights with a small absolute value work well.
Table 2 lists the training accuracy drop when we use three different pruning criteria; large _ finalleaves a large absolute value of weights and corresponds to IMP, small _ final conversely leavessmall weights, and random prunes randomly.
This notation of criteria follows that of Zhou et al.[43].
As expected, large _ final has a smaller drop in training accuracy after pruning than the others.
Training accuracy gets lower after reverting; If we assume that retraining can reach weights that showthe same or better accuracy because weights achieving good accuracy with the same structure exist, itseems to make sense to use large _ final to decrease the KL term while suppressing the increase intraining loss.
The results in Table 2 confirm this assumption empirically.
We can also discuss thiswith the following consideration using Taylor expansion (see Appendix B.6). [Equation] where λ θ + γ∆ θ max is a top eigenvalue of Hessian H ( θ + ∆ θ ) . This provides a brief insight into why IMP succeeds by pruning small magnitude weights under theassumption that the maximum eigenvalues are not very different.
4.4.2 Continuous Sparsification
Continuous sparsification [34] is a method to find winning tickets by removing the parameterscontinuously instead of alternating between training and pruning.
This target function is as follows,
[Equation]
(7) where η > 0 is a hyperparameter.
Continuous sparsification is formulated as the training lossminimization with the L 0 regularization of weights, and a sigmoid function σ is used for thecontinuous relaxation of the regularization term as follows.
[Equation]
We can regard this function as an approximation of the PAC-Bayes bound of our formulation.
Letϕ i ≥ 0 be the Gaussian KL part in 3, the summation of training risk and KL is as follows.
[Equation]
We make three approximations: 1) replace Q with E
[ θ ] over the spike-and-slab distribution byfirst-order Taylor expansion on the training risk, 2) simplify the second term to the L 1 norm of λ qbecause the second term can be viewed as a weighted summation of λ q,i , and 3) remove the thirdterm, which can be regarded as a regularization to the target sparsity.
This yields the following, whichis similar to Eq. 8.
[Equation]
Since the Gaussian KL ϕ i is approximated, the distance from the initial weights is not taken intoaccount in this setting.
The authors adopt a problem setting where the weights trained a few epochsahead instead of the initial weights are used for ticket search following Frankle et al.
[11]; thereforetheir work does not have to consider suppressing the learning rate, i.e., the distance from the initialweights.
Note that Hayou et al.
[15] proposed PAC-Bayes pruning (PBP) by optimizing the PACBayes bound.
However, the limitation of our analysis is that we cannot reveal an explicit relationshipbetween continuous sparcification and PBP.
Conclusion
In this work, we explored the fact that a small learning rate is required to find winning tickets, andwe provided empirical analysis related to flatness and the distance from the initial weights.
Onthe basis of these findings, we used the PAC-Bayesian framework to analyze winning tickets andexperimentally showed that it captures the generalization behavior.
Finally, we reconsider IMP andcontinuous sparsification from a PAC-Bayesian perspective.
In this study, we do not analyze the casewhere no solution exists near the initial weights, which needs IMP with rewinding to early epoch.