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
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.
In this connection, we thenfocus on the two concepts flatness and the distance from the initial weights of the winning tickets.
Wenext apply the PAC-Bayesian theory to LTH on the basis of the flatness motivation and show that itcan explain the relationship between LTH and generalization behavior.
We use the PAC-Bayes boundfor a spike-and-slab distribution to analyze winning tickets, which is based on our experimentalfindings that reducing the expected sharpness restricted to an unpruned parameter space and addingthe regularization of distance from the initial weights can enhance the test performance of winningtickets.
Finally, we revisit existing algorithms such as IMP, continuous sparsification [34] from thepoint of view of the PAC-Bayes bound optimization.
This consideration gives an interpretation ofthese methods as an approximation of bound optimization.
To sum up, our contributions are as follows.
• We experimentally show that IMP with a small learning rate finds relatively sharper minima and 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-and- slab 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.
Liu et al.
[26] discussed the relationship between winningtickets and the learning rate considering the similarity between initial and trained weights in termsof pruning mask overlap instead of the distance from the initial weights.
There is a recent study onfinding a mask 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 winning tickets generalization based on sample complexity.
Althoughit is limited to the case of a one-hidden-layer neural network, they provided insight into why thesparsity increases the generalization ability.
For a PAC-Bayesian theory, Hayou et al.
[15] also usedspike-and-slab prior and posterior distributions.
However, their motivation and purpose are completelydifferent from our work. Their aim was to obtain a sophisticated pruning mask by optimizing the PAC-Bayes bound, and this is mainly in the context of network pruning rather than LTH. Our workdiffers in that we use it to analyze the generalization behavior of a given winning ticket on the basis
[Figure]
[Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] [Figure] [Figure] [Figure]
[Figure] Clean Noisy Learningrate
T e s t A cc u r a c y [Figure] of our empirical findings on the learning rate, the flatness, and the distance from the 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.
Thesmall learning rate is good for finding the winning ticket, which is contrary to the intuition that alarge learning rate is good in terms of generalization ability.
First, we show the small learning rateis actually disadvantageous for generalization ability of subnetworks through experiments underlabel noise. Next, as a reason for this, we show that the small learning rate finds the relativelysharper minima, and it is possible to find the winning tickets in the flatter minima using sharpnessregularization. We also focused on the distance from the initial weights as a reason for why the smalllearning rate is important for IMP.
On the basis of these findings, further discussion is developed totheoretically support them in the next section.
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
[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]
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.
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 {+
+}
[Figure]
shows that SGD with the small learning rate finds relatively a sharper minimum and that IMP can finda 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 thatthe 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.2.
Add noise to the parameters including the pruned weights.3.
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 Gaussian distribution 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 boundbut 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 with a peak at zero-weight δ { 0 } ,in the PAC-Bayesian bound.
Let σ p and σ q be vectors whose i th element is a variance of Gaussiandistribution, λ p and λ q be vectors that represents the mixture ratio of prior and posterior, respectively.θ represents the network parameter, 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
[Equation]
where
[Equation]
Here, λ p,i is set to the target sparsity, so if we want 90% sparsity winning tickets, then λ p,i is setto 0 . 1 .
Since the structure of the pruned network is given, we set the element of λ q,i to 0 or 1asymptotically according to the pruning mask m and obtain the following kl divergence about λ .
This operation has been 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.
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.
[Table]
Training accuracy gets lower after reverting; If we assume that retraining can reach weights thatshow the same or better accuracy because weights achieving good accuracy with the same structureexist, it seems to make sense to use large _ final to decrease the KL term while suppressing theincrease in training loss.
The results in Table 2 confirm this assumption empirically.
We can alsoroughly understand why IMP succeeds by pruning small magnitude weights using the simple Taylorexpansion discussion (see Appendix B.6).
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]
(6) 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. 7.
[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.