Analyzing Lottery Ticket Hypothesis from PAC-Bayesian Theory Perspective
Abstract
The lottery ticket hypothesis (LTH) has attracted attention because it can explain why over-parameterized models often show high generalization ability.
It is known that when we use iterative magnitude pruning (IMP), which is an algorithm to find sparse networks with high generalization ability that can be trained from the initial weights independently, called winning tickets , the initial large learning rate does not work well in deep neural networks such as ResNet.
However, since the initial large learning rate generally helps the optimizer to converge to flatter minima, we hypothesize that the winning tickets have relatively sharp minima, which is considered a disadvantage in terms of generalization ability.
In this paper, we confirm this hypothesis and show that the PAC-Bayesian theory can provide an explicit understanding of the relationship between LTH and generalization behavior.
On the basis of our experimental findings that IMP with a small learning rate finds relatively sharp minima and that the distance from the initial weights is deeply involved in winning tickets, we offer the PAC-Bayes bound using a spike-and-slab distribution to analyze winning tickets.
Finally, we revisit existing algorithms for finding winning tickets from a PAC-Bayesian perspective and provide new insights into 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 parameters leads 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 of parameters, the pruned network can often work well with little or no accuracy loss.
However, training this sparse subnetwork independently from the initial weights often does not work and we can only obtain the sparse subnetwork through pruning after training the whole network.
Frankle and Carbin [10] presented “Lottery Ticket Hypothesis (LTH)” that states the existence of winning tickets: small but critical subnetworks which can be trained independently from scratch.
They proposed an algorithm called iterative magnitude pruning (IMP) to obtain a winning ticket.
They pointed out that, for deeper networks such as VGG [35] and ResNet [16], small learning rate is required to obtain a winning ticket.
However, since a large learning rate helps the generalization ability of neural networks [25], the learning rate should be well controlled to find a winning ticket that has a higher test accuracy.
Frankle et al.
[11] found a correlation between the stability to SGD noise and the ability to find the winning ticket and empirically showed that the large learning rate moves 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 noise setting compared to the subnetwork created with the large learning rate; that is, the generalization ability of winning tickets is degraded due to the learning rate constraint.
In this connection, we then focus on the two concepts flatness and the distance from the initial weights of the winning tickets.
We next apply the PAC-Bayesian theory to LTH on the basis of the flatness motivation and show that it can explain the relationship between LTH and generalization behavior.
We use the PAC-Bayes bound for a spike-and-slab distribution to analyze winning tickets, which is based on our experimental findings that reducing the expected sharpness restricted to an unpruned parameter space and adding the regularization of distance from the initial weights can enhance the test performance of winning tickets.
Finally, we revisit existing algorithms such as IMP, continuous sparsification [34] from the point of view of the PAC-Bayes bound optimization.
This consideration gives an interpretation of these 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 distance and 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 their behavior.
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.
They proposed IMP with rewinding to an early epoch to avoid the very early training process because training 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 from the initial weights on LTH.
Bain [2] plotted the loss landscape of winning tickets visually and found that 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 generalization of pruned networks considering flatness.
They measured flatness by using the trace of Hessian and gradient covariance matrix.
He et al.
[17] refers the distance from the initial weights to analyze the label noise robustness of winning tickets; they stated that the influence of the label noise is lessened if the distance is suppressed.
Liu et al.
[26] discussed the relationship between winning tickets and the learning rate considering the similarity between initial and trained weights in terms of pruning mask overlap instead of the distance from the initial weights.
There is a recent study on finding 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 comparable accuracy 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.
Although it is limited to the case of a one-hidden-layer neural network, they provided insight into why 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 are completely different 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 work differs in that we use it to analyze the generalization behavior of a given winning ticket on the basis 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.
[Figure]
[Figure]
Clean
Noisy
[Figure]
cc u r a c y
[Figure]
Empirical Analysis
We first show some empirical results because our new findings about winning tickets in this empirical analysis motivate the PAC-Bayesian analysis for LTH; thus, we interpret the results on the basis of the PAC-Bayesian perspective in the next section.
We empirically investigated the properties of winning tickets mainly related to the learning rate.
The small learning rate is good for finding the winning ticket, which is contrary to the intuition that a large learning rate is good in terms of generalization ability.
First, we show the small learning rate is actually disadvantageous for generalization ability of subnetworks through experiments under label noise.
Next, as a reason for this, we show that the small learning rate finds the relatively sharper minima, and it is possible to find the winning tickets in the flatter minima using sharpness regularization.
We also focused on the distance from the initial weights as a reason for why the small learning rate is important for IMP.
On the basis of these findings, we will develop a discussion from a PAC-Bayesian perspective in the next section. We will show this perspective will theoretically support the findings in this 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 randomly flipped to see whether the generalization behavior of winning tickets is degraded by the constraint that the learning rate must be small.
Figure 1 shows the test accuracy on clean and label noise datasets of sparse 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.
We added the original unpruned baseline (dashed green line) to discuss if the subnetwork is a winning ticket, and this baseline shows no such accuracy drop when increasing the learning rate.
The subnetwork eventually performs worse than the original unpruned network, which means that IMP ultimately fails to find the winning ticket.
In contrast, the test accuracy generally continues to improve without decreasing as the learning rate is increased in the high label noise setting (orange line).
The test accuracy increases even when it is not a winning ticket for no label setting at the same learning rate.
To sum up, the large learning rate is not suitable for finding winning tickets under a clean dataset but is advantageous in the high label noise setting. [Figure]
3.2 Flatness
Now that we have seen that winning tickets have undesirable properties due to the small learning rate
[Figure]
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 on pruned 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 the 1-d loss landscape of the subnetwork that has a high sparsity.
This landscape shows the training loss around the trained weights adding perturbation restricted to unpruned parameters.
We can see that the parameters of the winning ticket (small learning rate) are in the sharper minimum compared with the large learning rate.
The large learning rate can find a flatter minimum; however, the training loss is higher and the test accuracy is worse than the small learning rate.
This graph is a fixed sparsity loss landscape; therefore, it is not possible to discuss whether or not the subnetwork is a winning ticket from this graph.
There is actually an accuracy drop in the large learning rate setting when sparsity is changed, which means this is not a winning ticket.
Given that IMP requires a small learning rate rather than a large one, these results suggest two possible interpretations; 1) sharp minimum is essential to find winning tickets, therefore the large learning rate fails in IMP, 2) sharp minimum is simply the result of small learning rate training, and flat minimum is better for winning tickets if possible.
To investigate these possibilities, we used sharpness-aware minimization (SAM) [8] and neural variable risk minimization (NVRM) [37] to search for parameters that lie in neighborhoods having uniformly low loss.
They differ in that SAM minimizes the maximum loss in the neighborhood, whereas NVRM minimizes the expected loss in the neighborhood.
We used both of them to compare with normal SGD as a baseline.
These optimizers are based on SGD with the same setting as the original LTH paper [10], and the noise considered in these methods is limited to the unpruned parameters.
Figure 2 also shows the loss landscape when SAM is used; It can actually reach a flatter minimum.
The loss landscape is as flat as the large learning rate setting; however, unlike this, the test accuracy is higher than that of the small learning rate.
Next, we investigated the trace of Hessian to analyze the flatness of winning tickets created by these optimizers.
The trace of Hessian is used as a measure of flatness in the prior work [4, 21, 39].
[-Table 
 [Figure] shows-]Tableshows that SGD with the small learning rate finds relatively a sharper minimum and that IMP can find a flatter one by using SAM or NVRM.
These optimizers can find relatively flatter minima, which improves the test 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 SAM instead of SGD (see Appendix B.2). The large learning rate has already found relatively flatter [Table]
[Figure] [Table] [Figure] minima as shown in Figure 2; therefore, it is considered that the results do not change by using SAM.
This fact implies that the flatness around the parameters found by SGD with a large learning rate and by SAM with a small learning rate are similar, however winning tickets cannot be found with the large learning rate because they find different solutions in terms of some other 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 be found by IMP within a range not far from the initial weights.
In order to confirm this, we ran IMP suppressing this distance and compare it with the usual regularization by the L 2 norm regularization.
In addition, we also discussed the pruning mask structure in Appendix B.3 in relation to the distance from the initial weights.
We empirically confirmed that the winning ticket can be obtained via the regularization of this distance even 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 ∈ { 0 , 1 } | θ | be a pruning mask. In this experiment, the parameters were regularized only for unpruned weights.
Specifically, we designed the 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 80 %, 90 % and 95 % sparsity subnetworks and the original unpruned network as a baseline.
If the subnetwork accuracy is close to the whole network accuracy, it is considered to be successful in finding the winning ticket.
As discussed previously, the large learning rate fails to find winning tickets unlike the small learning rate setting; however adding the regularization from the initial weights changes this trend.
In Figure 3, λ around 2 × 10 − 4 shows that a winning ticket is found since the 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 the distance from the initial weights even with the large learning rate.
This finding is shown more clearly in Appendix B.4.
We confirmed it by examining the test accuracy when sparsity is changed for problem settings other than ResNet20 + CIFAR10.
We can also obtain interesting results when the learning rate is small.
The accuracy drop is small while λ 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 same for large learning rate setting), but in the case of l 2 _ norm , the gap widens significantly even though generalization ability increases because of the regularization.
The strong norm regularization makes IMP fail 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 the training 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. Therefore, it is important to ensure a balance between them.
As for IMP, there is a specific situation where suppressing the parameter norm makes the training loss large due to the failure to find a winning ticket.
If we take the initial weights as a prior mean instead of 0 , the training loss can decrease by suppressing the KL term when the trained weights are far from the initial weights.
This experiment corresponds to what to take as a prior mean in terms of minimizing PAC-Bayes bound, and this result indicates that setting the initial weights to a prior mean 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 bound of a spike-and-slab formulation based on our experiments.
Next, we show that this bound captures the generalization behavior of winning tickets and revisit existing algorithms from the perspective of this bound.
4.1 Flatness on winning tickets
While we could simply consider the neighborhood of the trained weights in an unpruned neural network, it is not trivial to define the flatness of pruned networks depending on how the pruned weights 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 parameter space can no longer be in the local minimum; thus, it is uncertain whether flatness has any meaning in such a setting.
He et al.
[17] conducted re-dense training and showed that solutions at high sparsity are no longer minimizers in high dimensions.
They also found that the winning ticket has higher sharpness than the original network based on Measure 3 and concluded that highly sparse solutions do not 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 Measure 1.
We will also compare this formulation with the normal Gaussian formulation (Measure 2) through numerical experiments and show that the spike-and-slab formulation captures the generalization behavior of winning tickets better.
This supports our findings that using SAM and NVRM on pruned networks 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 for analyzing winning ticket properties.
As discussed in Section 3.3, the distance from the initial weights of the winning ticket is expected to be small.
We set the prior mean as the initial weights to take advantage 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 initial weights 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 optimized. This is because the more sparse the subnetwork is, the larger the bound becomes.
In addition, noise is inevitably added to the pruned weights when considering expected sharpness since the variance of the pruned part cannot be set to zero.
In order to limit the distance from the initial weights and noise added in expected sharpness only to the unpruned weights, we use a spike-and-slab distribution, which is the mixture of the Gaussian distribution N and Dirac delta distribution with a peak at zero-weight δ { 0 } , in the PAC-Bayes bound.
Let σ p and σ q be vectors whose i th element is a variance of Gaussian distribution, λ 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
[Equation]
λ p
[Equation] i is set to the target sparsity [Equation] so if we want 90% sparsity winning tickets, then λ p,i is set to 0 .
Since the structure of the pruned network is given, we set the element of λ q,i to 0 orasymptotically according to the pruning mask m and obtain the following kl divergence about λ
[Equation] This operation has been conventionally done in the entropy discussion [28]. [Equation]
The expected sharpness is as follows.
[Equation]
where
[Equation]
This means that flatness of pruned networks can be discussed adding noise to only unpruned weights.
The advantage of this definition compared to the Gaussian distribution setting, where we cannot avoid adding noise to pruned weights, will be discussed in the next subsection.
4.3 Numerical Experiment
We conducted numerical experiments to confirm that our PAC-Bayes formulation can adequately explain the behavior of winning tickets.
We optimized the posterior variance to minimize the PACBayes bound, and plot KL term and the training loss on the posterior distribution to investigate that the bound can capture the test accuracy of the subnetworks produced by IMP.
As a comparison, we also experiment with the Gaussian distribution setting using the zero-mean prior.
The PAC-Bayesian bound used here is the variational KL bound [6] (Theorem D.2, and see also Appendix B.7), and we implemented the code following Dziugaite et al.
Figure 4 shows the distribution of the training risk term and KL term when optimizing the PAC-Bayes bound, and the actual test accuracy is colored.
We show that the bound with the spike-and-slab formulation successfully explains the behavior of winning tickets, dividing the point cloud into three groups: A) winning tickets (moderate learning rate), B) subnetworks that failed to find winning tickets (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 and the 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 does not capture the test accuracy well because B should have higher test accuracy considering its KL term and training risk term.
This means that, as the learning rate is increased and winning tickets can 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 the winning ticket fails to be found (see Appendix B.5).
In terms of not only the bound optimization but also the analysis of the existing winning ticket, it is preferable to use the spike-and-slab distribution to 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 optimizing the PAC-Bayes bound.
Although we focus only on IMP and continuous sparsification, this view could be helpful 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 IMP behaves in the sense of a PAC-Bayes bound with our formulation instead of viewing IMP as a direct bound optimization problem.
The risk term and KL term in the PAC-Bayes bound are basically in a trade-off relationship; choosing a complex model to fit the training data will increase the KL term, while choosing a simple model [Table] [Table] may not have a high accuracy on the training data.
We point out that the two steps of IMP: 1) train the subnetwork, 2) prune the subnetwork and revert its weights, optimize the overall bound by alternately reducing 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.
This process reduces the training risk, and the increase in KL term is not expected to be so large in our formulation because the prior mean is set to the initial weights and the trained weights of IMP are not far 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 reverts the trained weights to the initial state.
Reverting θ to θ init and changing part of λ q from 1 tomake the KL term small.
The number of Gaussian KL summations decreases and the distance from initialization gets to 0 , and the λ KL part is also reduced if the prior mixture ratio λ p is set to the final target sparsity.
This KL reduction is not dependent on the pruning criterion.
The problem here is how to minimize the increase in training loss, which is related to what heuristic pruning criterion we choose and why pruning weights with a small absolute value works well.
Table 2 lists the training accuracy drop when we use three different pruning criteria; large _ final leaves a large absolute value of weights and corresponds to IMP, small _ final conversely leaves small weights, and random prunes randomly.
This notation of criteria follows that of Zhou et al.
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 show the same or better accuracy because weights achieving good accuracy with the same structure exist, it seems to make sense to use large _ final to decrease the KL term while suppressing the increase in training loss.
The results in Table 2 confirm this assumption empirically.
We can also discuss the reason why large _ final pruning criterion is good for IMP by considering the following simple 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 the assumption 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 parameters continuously instead of alternating between training and pruning.
This target function is as follows,
[Equation] [Equation]
where η > 0 is a hyperparameter.
Continuous sparsification is formulated as the training loss minimization with the L 0 regularization of weights, and a sigmoid function σ is used for the continuous 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 by first-order Taylor expansion on the training risk, 2) simplify the second term to the L 1 norm of λ q because the second term can be viewed as a weighted summation of λ q,i , and 3) remove the third term, which can be regarded as a regularization to the target sparsity.
This yields the following, which is similar to Eq. 8.
[Equation]
Since the Gaussian KL ϕ i is approximated, the distance from the initial weights is not taken into account in this setting.
The authors adopt a problem setting where the weights trained a few epochs ahead instead of the initial weights are used for ticket search following Frankle et al.
[11]; therefore their work does not have to consider suppressing the learning rate, i.e., the distance from the initial weights.
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 relationship between continuous sparsification and PBP.
Conclusion
In this work, we explored the fact that a small learning rate is required to find winning tickets, and we provided empirical analysis related to flatness and the distance from the initial weights.
On the basis of these findings, we used the PAC-Bayesian framework to analyze winning tickets and experimentally showed that it captures the generalization behavior.
Finally, we reconsidered IMP and continuous sparsification from a PAC-Bayesian perspective.
In this study, we do not analyze the case where no solution exists near the initial weights, which needs IMP with rewinding to early epoch. A Background
In this section, we explain the background knowledge for IMP and the PAC-Bayesian theory.
IMP Iterative pruning is a method of obtaining a subnetwork by repeating training and pruning in stages.
Frankle and Carbin [10] showed that iterative pruning (Algorithm 1) could find the winning ticket by adding an operation to restore the weights to the initial weights after pruning.
In Frankle and Carbin [10], the mask criterion is simply to keep the weights with a large final magnitude, | θ T | ; This is called Iterative Magnitude Pruning (IMP).
This paper follows their setting: p is set to 0 .
2 and the models are pruned globally.
Frankle et al.
[11] proposed an iterative pruning with rewinding to avoid lower-stability phase to
SGD noise in the early training in each IMP step.
The algorithm does not return to the exact initial weights, but instead returns the weights trained slightly in advance as the initial weights.
It can find the winning ticket even with the initial large learning rate or in the harder settings such as ImageNet; however it revises the original LTH setting.
Our paper focuses on analyzing the properties of winning tickets rather than improving the accuracy or robustness of winning tickets; thus, we will not consider this problem setting.
Liu et al.
[27] claimed that the large learning rate performs better in the larger model settings such as
ResNet56 and ResNet110, contrary to Frankle and Carbin [10], which stated that IMP requires the small learning rate to find winning tickets.
From our findings in Section 3.3 and 4.3, we suppose that there is no good solution near the initial weights in such a setting; therefore the small learning rate cannot find the winning tickets, and that IMP with rewinding is effective in such a setting because it shifts the initial weights closer to the final trained weights by pretraining a little and searches for the winning ticket in a better lottery.
PAC-Bayesian Theory
First, we provide the notations used in this paper.
Denote the training sample S =
{ ( x i , y i ) } mi =1 ∈ ( X × Y ) m , which is randomly sampled from an underlying data distribution
D .
Let H be a set of hypotheses, and ℓ : H × X × Y → [0 , 1] be a loss function.
Given f ∈ H , we formulate the empirical risk on S and the generalization error on D as
[Equation]
The PAC-Bayesian framework gives a bound on the generalization error of a posterior distribution Q over the hypothesis H ; we denote it L D ( Q ) .
It assumes that we have a prior distribution P on
H which does not depend on training data, and we update it to Q through the learning process.
Optimizing the PAC-Bayes bound controls the balance between the empirical risk and the closeness to the prior (small complexity of the model).
Although there are several types of PAC-Bayes bound, we consider the following well-used form of the PAC-Bayes bound.
Theorem 1 (Alquier et al.
bound [1])
Given a real number δ ∈ (0 , 1] , a non-negative real number η , and a prior distribution P on H defined before seeing any training sample ( X, Y ) ∈ S , with probability at least 1 − δ , for all Q on H
[Equation]
12) where
[Equation]
It is known that flatness is closely related to the generalization ability of neural networks [21] [20].
As shown in previous studies ([5], [31] [32]), it can be viewed in the expected sense from a PAC-Bayesian perspective.
We decompose the PAC-Bayes bound as follows.
[Equation]
where f ∈ H is a solution obtained by a training.
The expected sharpness term represents the amount of change in empirical risk around the trained weights, and the solutions in flatter minima are expected to have a relatively smaller value of this term.
In the PAC-Bayesian framework, the role flatness plays in generalization behavior can be understood in this way.
For example, we consider the case where the Gaussian distribution is used for the prior and posterior.
Let N ( µ , Σ ) be the Gaussian distribution, where µ is the mean and Σ , is the covariance matrix and θ be the parameters of a neural network.
We set a prior P to be N ( 0 , σ 2 I ) and a posterior Q to be
N ( θ , σ 2 I ) , where σ > 0 , and the KL term is calculated as ∥ θ ∥ 22 / (2 σ 2 ) .
This is consistent with the conventional understanding: solutions in flat minima obtained with norm regularization can achieve good generalization accuracy.
When considering this PAC-Bayesian framework in the pruned networks, we should note that prior P has to be defined without depending on the training samples S .
It can be thought that if the parameter space is restricted to an unpruned weight subspace, we do not have to consider the pruned network case differently.
However this is not valid because this prior depends on the structure of the pruning mask m , which is found after seeing the training samples S .
Target sparsity β does not depend on S , so we can use it in the prior.
B Other Experiments B.1 The test accuracy when SAM and NVRM are used
[Table]
In Section 3.
[Table]
we observe that the winning ticket is in a relatively sharper minima due to a small learning rate and that IMP can find a flatter minimum by using SAM or NVRM.
Table 3 shows the test accuracy when a flatter solution is obtained by using SAM and NVRM.
They can achieve a test accuracy the same as or even better than the SGD with a small learning rate
[Table]
and the improvement in test accuracy can be seen especially in VGG16.
It is considered that since VGG16 has a larger number of parameters than ResNet20, a flatter solution with a relatively small training loss could be found.
We found no significant difference in test accuracy between SAM and NVRM
[Table]
B.2 SAM with the large learning rate
Figure 5 shows the test accuracy when we use SAM with a large learning rate; we experimented only
ResNet20 + CIFAR10 due to computational resource limitations.
This figure shows that SAM does not improve test accuracy from that of IMP with a large learning rate and that SAM does not help to find winning tickets because the test accuracy continues to decline as sparsity increases.
Since
IMP with a large learning rate already produces relatively flatter solutions, SAM will not change the
[Figure]
[Figure]
results as expected.
This also confirms that the improved generalization accuracy when we use SAM instead of SGD for IMP with a small learning rate is because of finding flatter solution, not some other side effects of SAM.
B.3 Mask structure with different learning rates
We conducted the following experiment related to distance from initial weights.
In Figure 6, we found that IMP with different learning rates produce the pruning masks with different properties; learning with a large learning rate on a sub-network obtained with a small learning rate also does not provide winning tickets.
Retraining with a large learning rate for a given mask greatly improves test accuracy at a low sparsity.
In contrast, as sparsity is increased, this improvement decreases or the test accuracy worsens; it is generally the same trend as IMP with the large learning rate (orange line).
We found that the mask obtained by IMP with a small learning rate has a structure that performs well with weights around the initialization, and it does not work well when trained directly of the subnetwork with a large learning rate.
[Figure]
B.4 Regularization of the distance from the initial weights
Figure 7 shows the test accuracy on CIFAR10 and CIFAR100 when we trained ResNet20 and VGG with a regularization from the initial weights.
For a large learning rate setting (orange line), increasing sparsity significantly reduces the test accuracy, which means that IMP cannot find the winning ticket.
By adding regularization from the initialization, this decrease in the test accuracy can be reduced (red line), showing a similar trend for a small learning setting (green line), which is successful in finding winning tickets.
B.5 The distribution of parameter changing the learning rate
Figure 8 shows the parameter distribution of ResNet20 and VGG16 changing the learning rate.
We trained them on CIFAR10 and plot the unpruned weights in order from the smallest to largest: 20% , 40% , 60% , 80% .
The same sparsity where the winning tickets cannot be found in Figure 1 shows a change in the trend of the distribution.
Although the difference is not apparent when we measure
L2 norm, we confirmed that each parameter, which was near the initial weights originally, becomes distributed in a wider range when the learning rate is increased.
B.6 Proof of Eq.
We estimate the deviation of f when the ∆ θ moves from the trained weights θ using Taylor’s theorem.
Let H be a Hessian matrix, then we have
[Equation]
[Figure]
where λ θ + γ∆ θ max is a top eigenvalue of H ( θ + ∆ θ ) .
First equation comes from quadratic Taylor’s theorem, second equation comes from the fact that θ is a trained weights and ∇ f ( θ ) = 0 , and third inequality holds because of sub-multiplicativity of matrix norm.
B.7 Variational KL bound
In Section 4.3, we optimized the following variational KL bound [6], [15].
Given a real number δ ∈ (0 , 1] , with probability 1 − δ over the training sample S ,
[Equation]
(16) where
[Equation]
We used this bound in our experiment because it can avoid selecting variables that appear in the
PAC-Bayes bound, such as η in Eq.