Deconstructing Lottery Tickets: Zeros, Signs, and the Supermask
Abstract
The recent “Lottery Ticket Hypothesis” paper by Frankle & Carbin showed that a simple approach to creating sparse networks (keep the large weights) results in models that are trainable from scratch, but only when starting from the same initial weights.
The performance of these networks often exceeds the performance of the non-sparse base model, but for reasons that were not well understood.
In this paper we study the three critical components of the Lottery Ticket (LT) algorithm, showing that each may be varied signiﬁcantly without impacting the overall results.
Ablating these factors leads to new insights for why LT networks perform as well as they do.
We show why setting weights to zero is important, how signs are all you need to make the re-initialized network train, and why masking behaves like training.
Finally, we discover the existence of Supermasks, or masks that can be applied to an untrained, randomly initialized network to produce a model with performance far better than chance (86% on MNIST, 41% on CIFAR-10).
Introduction
Many neural networks are over-parameterized [1, 2], enabling compression of each layer [2, 15, 5] or of the entire network [10].
Some compression approaches enable more efﬁcient computation by pruning parameters, by factorizing matrices, or via other tricks [5, 6, 9, 11–17].
A recent work by Frankle & Carbin [3] presented a simple algorithm for ﬁnding sparse subnetworks within larger networks that can meet or exceed the performance of the original network.
Their approach is as follows: after training a network, set all weights smaller than some threshold to zero 3 , rewind the rest of the weights to their initial conﬁguration 4 , and then retrain the network from this starting conﬁguration but with the zero weights frozen (not trained).
See Section S1 for a more formal description of this algorithm.
In this paper we perform ablation studies along the above three dimensions of variability, considering alternate mask criteria (Section 2), alternate mask-1 actions (Section 3), and alternate mask-0 actions (Section 4).
These studies in aggregate reveal new insights for why lottery ticket networks work as they do.
Along the way we also discover the existence of Supermasks—masks that produce above-chance performance when applied to untrained networks (Section 5).
Mask criteria
We begin our investigation with a study of different Mask Criteria , or functions that decide which weights to keep vs. prune.
In this paper, we deﬁne the mask for each individual weight as a function of the weight’s values both at initialization and after training: M ( w i , w f ) .
We can visualize this function as a set of decision boundaries in a 2D space as shown in Figure S1.
In [3], the mask criterion simply keeps weights with large ﬁnal value; we refer to this as the large _ ﬁnal mask, M ( w i , w f ) =
| w f | .
In addition to large _ ﬁnal , we also experimented with the inverse, small _ ﬁnal , versions that evaluate weights based on their initial magnitudes instead, large _ init and small _ init , versions that select for both, large _ init _ large _ ﬁnal and small _ init _ small _ ﬁnal , two mask criteria that evaluate how far weights moved, magnitude _ increase and movement , and a control, random , that chooses masks randomly.
These nine masks are depicted along with their associated equations in Figure S2.
[Figure]
In this section and throughout the remainder of the paper
[Figure]
we follow the experimental framework from [3] and perform iterative pruning experiments on a fully-connected network (FC) trained on MNIST [8] and on three convolutional networks (Conv2, Conv
[Figure]
and Conv6) trained on CIFAR-10 [7].
For more achitecture and training details, see Section S2.
Results of these pruning experiments are shown in Figure 1.
Note that the ﬁrst six criteria out of the eight form three opposing pairs; in each case, we see when one member of the pair performs better than the random baseline, the opposing member performs worse than it.
We see that the mask criterion large _ ﬁnal , as employed by the LT algorithm, is indeed a competitive mask criteria.
However, the magnitude _ increase criterion turns out to work just as well as the large _ ﬁnal criterion, and in some cases signiﬁcantly better.
In Section 4, we provide an explanation to why some mask criteria work well while others don’t.
Mask-1 actions: show me a sign
Now that we have explored various ways of choosing which weights to keep and prune, we will consider what values to set for the kept weights.
In particular, we want to explore an interesting observation in [3] which showed that the pruned, skeletal LT networks train well when you rewind to its original initialization, but degrades in performance when you randomly reinitialize the network.
Why does reinitialization cause LT networks to train poorly?
Which components of the original initialization are important?
We evaluate a number of variants of reinitialization to investigate:
• “Reinit” experiments: reinitialize kept weights based on the original initialization distribution • “Reshufﬂe” experiments: reinitialize while respecting the original distribution of remaining weights in that layer by reshufﬂing the kept weights’ initial values • “Constant” experiments: reinitialize by setting mask-1 weight values to a positive or negative constant, with the constant set to be the standard deviation of each layer’s original initialization.
Thus every weight on a layer becomes one of three values: − α , 0 , or α .
We ﬁnd that none of these three variants alone are able to train as well as the original LT network, shown as dashed lines in Figure 2.
However, all three variants work better when we ensure that the new values of the kept weights are of the same sign as their original initial values.
These are shown as solid color lines in Figure 2.
Clearly, the common factor in all working variants including the original rewind action is the sign.
As long as you keep the sign, reinitialization is not a deal breaker; in fact, even setting all kept weights to a constant value consistently performs well!
[Figure]
[Figure] [Figure]
% Weights Remaining
[Figure]
T e s t
A cc u r a c y 100.0% 41.0% 16.8% 6.9% 2.8% 1.2% Weights Remaining 0.
T e s t
A cc u r a c
y rewind (large_final) random
[Figure]
reshuffle
[Figure]
constant,
[Figure]
constant
[Figure]
Mask-0 actions: masking is training
What should we do with weights that are pruned?
Typical network pruning procedures perform two actions on pruned weights: set them to zero, and freeze them in subsequent training.
However, it is unclear which of these two components leads to the increased performance in LT networks.
To separate the two factors, we run a simple experiment: we reproduce the LT iterative pruning experiments in which network weights are masked out in alternating train/mask/rewind cycles, but try an additional treatment: freeze masked weights at their initial values instead of at zero.
If pruned weights are truly unimportant, we would expect that setting them to any other sensible values, such as their original initializations, should lead to a similarly performing network.
Figure 3 shows the results for this experiment.
We ﬁnd that networks perform signiﬁcantly better when weights are frozen speciﬁcally at zero than at random initial values.
For these networks masked via the LT large _ ﬁnal criterion, zero would seem to be a particularly good value to set weights to when they had small ﬁnal values.
So why does zero work better than initial values?
One hypothesis is that the mask criterion we use tends to mask to zero those weights that were headed toward zero anyway .
To test out this hypothesis, we run another experiment interpolated between the previous two: for any weight to be frozen, we freeze it to zero if it moved toward zero over the course of training, and we freeze it at its random initial value if it moved away from zero.
Results are shown in Figure 3.
By setting only the selected subset of pruned weights to zero, we fully recover the performance of the original LT networks.
This supports our hypothesis that the beneﬁt derived from freezing values to zero comes from the fact that those values were moving toward zero anyway.
In fact, if we apply this treatment to all weights, including the kept weights, we can outperform even the original LT networks.
Supermasks
The hypothesis above suggests that for certain mask criteria, like large _ ﬁnal , that masking is training: the masking operation tends to move weights in the direction they would have moved during training .
If so, just how powerful is this training operation?
To answer this question, we can start all the way from the beginning—not training the network at all, but simply applying a mask to the randomly initialized network.
It turns out that with a well-chosen mask, an untrained network can already attain a test accuracy far better than chance.
Although it is not entirely implausible to have better-than-chance performance since the masks are derived from the training process, the large improvement in performance is still surprising because the only transmission of information from the training back to the initial network is via a zero-one mask based on a simple criterion.
We call these masks that can produce better-than-chance accuracy without training of the underlying weights “Supermasks”.
[Figure]
[Figure] [Figure]
We turn our attention to evaluating how various mask criteria perform as Supermasks.
In addition to evaluating the mask criteria from Section 2, we deﬁne a new large _ ﬁnal _ same _ sign mask criterion based on the demonstration in
Section 3 of the importance of signs and of keeping large weights.
The large _ ﬁnal _ same _ sign mask criterion selects for weights with large ﬁnal magnitudes that also maintained the same sign by the end of training.
This criterion is depicted in Figure S7.
Also included as a control is the large _ ﬁnal _ diﬀ _ sign .
Performances of Supermasks produced by allcriteria are included in Figure 4, compared with two baselines networks unmasked ( untrained _ baseline and networks fully trained ( trained _ baseline )
By using this simple mask criterion of large _ ﬁnal _ same _ sign , we can create networks that obtain a remarkable 80% test accuracy on MNIST and 24% on CIFAR-10 without training.
Another curious observation is that if we apply the mask to a signed constant (as described in Section 3) rather than the actual initial weights, we can produce even higher test accuracy of up to 86% on MNIST and 41% on CIFAR-10!
Detailed results across network architectures, pruning percentages, and these two treaments, are shown in Figure S8.
We ﬁnd it fascinating that these Supermasks exist and can be found via such simple criteria.
As an aside, they also present a method for network compression, since we only need to save a binary mask and a single random seed to reconstruct the full weights of the network.
S Lottery Ticket Pruning Algorithm
We describe the lottery ticket algorithm in more detail:
Initialize a mask m to all ones.
Randomly initialize the parameters w of a network f
( x ; w (cid:12) m ) 1.
Train the parameters w of the network f
( x ; w (cid:12) m ) to completion.
Denote the initial weights before training w i and the ﬁnal weights after training w f . 2.
Mask Criterion.
Use the mask criterion M ( w i , w f ) to produce a masking score for each currently unmasked weight.
Rank the weights in each layer by their scores, set the mask value for the top p % to 1, the bottom (100 − p )% to 0, breaking ties randomly.
Here p may vary by layer, and we follow the ratios chosen in [3], summarized in Table S1.
In [3] the mask selected weights with large ﬁnal value corresponding to M ( w i , w f ) =
| w f | .
In Section 2 we consider other mask criteria.
Mask-1 Action.
Take some action with the weights with mask value 1.
In [3] these weights were reset to their initial values and marked for training in the next round.
We consider this and other mask-1 actions in Section 3. 4.
Mask-0 Action.
Take some action with the weights with mask value 0.
In [3] these weights were pruned: set to 0 and frozen during any subsequent training.
We consider this and other mask-0 actions in Section 4. 5.
Repeat from 1 if performing iterative pruning.
S Architectures and training hyperparameters
Table S1 contains the architectures used in this study, together with relevant training hyperparameters, based off of experiments in [3].
[Table]
38K
[Table]
4M
[Table] [Table]
1M
[Table]
[Figure] [Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
[Equation]
[Figure]
S Further mask criteria details
Figure S1 shows the representation of large _ ﬁnal mask criterion on a 2D plane.
Figure S2 shows the formulation and visual depiction of all the mask criteria considered in Section 2.
Figure S3 shows the convergence speed and performance of all mask critera for FC on MNIST and Conv2, 4, 6 on CIFAR-10.
S Further mask-1 action details
Figure S4 shows the convergence speed and performance of various reinitialization methods for FC on MNIST and
Conv2, 4, 6 on CIFAR-10.
S Further mask-0 action details
Figure S5 shows the performance of various treatment of pruned weights for FC on MNIST and Conv2, 4, 6 on
CIFAR-10.
Figure S6 illustrates why the large _ ﬁnal criterion biases weights that were moving toward zero during training toward zero in the mask, effectively pushing them further in the direction they were headed.
It also illustrates why the large _ ﬁnal criterion creates Supermasks.
[Table]
S Further Supermask details
Figure S7 depicts the effect of Supermasks, as well as the two additional mask criteria considered only as Supermasks.
Figure S8 shows performance of various mask criteria on initial test accuracy for FC on MNIST and Conv 2, 4, 6 on
CIFAR-10.
S6.1 Learning to ﬁnd the best Supermasks
Now that we know Supermasks exist, and those derived from simple heuristics work remarkably well, we might wonder how far we can push the performance of Supermasks for a given network.
One can search, in the search space of all 2 n possible masks, where n is the number of parameters in that network.
We can also try learning it with regular optimizers.
To do that, we create a trainable mask variable for each layer while freezing all original parameters for that layer at their random initialization values.
For an original weight tensor w and a mask tensor m of the same shape, we have as the effective weight w (cid:48)
= w i (cid:12) g ( m ) , where w i denotes the initial values weights are frozen at, (cid:12) is element-wise multiplication and g is a point-wise function that transform a matrix of continuous values into binary values.
One example of g is (cid:98) ( S ( m )) (cid:101) , where S is the sigmoid function and (cid:98)(cid:101) means rounding.
Bias terms are added as usual to the product of w (cid:48) with the inputs as per the usual fully connected or convolutional kernels.
We train the masks with g ( m ) =
Bern( S ( m )) , where Bern( p ) is the bernoulli sampler with probability p .
It works slightly better than (cid:98) ( S ( m )) (cid:101) .
The bernoulli sampling adds some stochasticity that helps with training, mitigates the bias of all things starting at the same value, and uses in effect the expected value of S ( m ) , which is especially useful when they are close to 0.5.
By training the m matrix with SGD, we obtained up to 95.3% test accuracy on MNIST and 65.4% on CIFAR-10.
Results are shown in Figure S8, along with all the heuristic based, unlearned Supermasks.
Note that there is no straightforward way to control for the pruning percentage.
What we do is initializing m with a constant of different magnitudes, whose value nudges the network toward pruning more or less.
With this tactic we are able to produce masks with the amounts of pruning (percentages of zeros) ranging from 7% to 89%.
Further details about the training can be seen in Section S6.
Table S2 summarizes the best test accuracy obtained through different treatments.
The result shows striking improvement of learned Supermasks over heuristic based ones.
And learning Supermasks results in performance not too far from training the full network.
It suggests that a network upon initialization has already contained powerful subnetworks that work well.
Additionally, the learning of Supermask allows identifying a possibly optimal pruning rate for each layer, since each layer is free to learn the distribution of 0 s in m on their own.
For instance, in [3] the last layer of each network is designed to be pruned approximately half as much as the other layers, in our setting this ratio is automatically adjusted.
S6.2 Further training details for learning Supermasks
We train the networks with mask m for each layer (and all regular kernels and biases frozen) with SGD, 0.9 momentum.
The {FC, Conv2, Conv4, Conv6} networks respectively had {100, 100, 50, 20} for learning rates and trained for {2000, 2000, 1000, 800} iterations.
These hyperparameters may seem absurd, but a network of masks is quite different and cannot train well with typical learning rates.
Conv
and Conv6 showed signiﬁcant overﬁtting, thus we used early stopping as we are unable to use standard regularizing techniques.
For evaluation, we also use Bernoulli sampling, but average the accuracies over 10 independent samples.
For adjusting the amount pruned, we initialized m in every layer to be the same constant, which ranged from -5 to 5.
In the future it may be worth trying different initializations of m for each layer for more granular control over per-layer pruning rates.
A different method to try would be to add an L1 loss to inﬂuence layers to go toward certain values, which may alleviate the cold start problems of some networks not learning anything due to mask values starting too low (effectively having the entire network start at zero).
M N I S
T F C
C
I F
A
RC on v
C
I F
A
RC on v
C
I F
A
RC on v
[Figure]
M N I S
T F C
C
I F
A
RC on v
C
I F
A
RC on v
C
I F
A
R -
C on v
Test Accuracy
[Figure]
FC
[Figure]
Conv
[Figure]
Conv
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
Conv original set pruned to 0 (LT
[Figure]
set pruned to init set pruned to init orset all to init or
T e s t
A cc u r a c y a t E a r l y S t o pp i n g I t e r a
[Figure] i o n T e s t A cc u r a c y a t E a r l y
S t o
pp i n g I t e [Figure] n w i w f
[Figure] [Figure]