P ICKING W INNING T ICKETS B EFORE T RAINING BY P RESERVING G RADIENT F LOW
Anonymous authors
Paper under double-blind review
A BSTRACT
Overparameterization has been shown to beneﬁt both the optimization and generalization of neural networks, but large networks are resource hungry at both training and test time.
Network pruning can reduce test-time resource requirements, but is typically applied to trained networks and therefore cannot avoid the expensive training process.
We aim to prune networks at initialization, thereby saving resources at training time as well.
Speciﬁcally, we argue that efﬁcient training requires preserving the gradient ﬂow through the network.
This leads to a simple but effective pruning criterion we term Gradient Signal Preservation (GraSP).
We empirically investigate the effectiveness of the proposed method with extensive experiments on
CIFAR-10, CIFAR-100, Tiny-ImageNet and ImageNet, using VGGNet and ResNet architectures.
Our method can prune 80% of the weights of a VGG-16 network on ImageNet at initialization, with only a 1.6% drop in top-1 accuracy.
Moreover, our method achieves signiﬁcantly better performance than the baseline at extreme sparsity levels.
I NTRODUCTION
Deep neural networks exhibit good optimization and generalization performance in the overparameterized regime (Zhang et al., 2016; Neyshabur et al., 2019; Arora et al., 2019; Zhang et al., 2019b), but both training and inference for large networks are computationally expensive.
Network pruning (LeCun et al., 1990; Hassibi et al., 1993; Han et al., 2015b; Dong et al., 2017; Zeng &
Urtasun, 2019; Wang et al., 2019) has been shown to reduce the test-time resource requirements with minimal performance degradation.
However, as the pruning is typically done to a trained network, these methods don’t save resources at training time .
Moreover, it has been argued that it is hard to train sparse architectures from scratch while maintaining comparable performance to their dense counterparts (Han et al., 2015a; Li et al., 2016).
Therefore, we ask: can we prune a network prior to training, so that we can improve computational efﬁciency at training time?
Recently, Frankle & Carbin (2019) shed light on this problem by proposing the Lottery Ticket
Hypothesis (LTH), namely that there exist sparse, trainable sub-networks (called “winning tickets”) within the larger network.
They identify the winning tickets by taking a pre-trained network and removing connections with weights smaller than a pre-speciﬁced threshold.
They then reset the remaining weights to their initial values, and retrain the sub-network from scratch.
Hence, they showed that the pre-trained weights are not necessary, only the pruned architecture and the corresponding initial weight values.
Nevertheless, like traditional pruning methods, the LTH approach still requires training the full-sized network for identifying the sparse sub-networks.
Can we identify sparse, trainable sub-networks at initialization?
This would allow us to exploit sparse computation with speciﬁed hardware for saving computation cost.
As shown by Dey et al.
(2019), with pre-speciﬁed sparsity one can achieve 5x efﬁciency for both training and testing.
At the ﬁrst glance, a randomly initialized network seems to provide little information that we can use to judge the importance of individual connections, since the choice would seem to depend on complicated training dynamics.
However, recent work suggests this goal is attainable.
Lee et al.
(2018) proposed the ﬁrst algorithm for pruning at initialization time: Single-shot Network Pruning (SNIP), which uses a connection sensitivity criterion to prune weights with both small magnitude and small gradients.
Their empirical results are promising in the sense that they can ﬁnd sparse, trainable sub-networks at the initialization.
However, connection sensitivity is sub-optimal as a criterion because the gradient of
[Figure]
each weight might change dramatically after pruning due to complicated interactions between weights.
Since SNIP only considers the gradient for one weight in isolation, it could remove connections that are important to the ﬂow of information through the network.
This effect is illustrated in Figure 1.
Practically, we ﬁnd that this blocking of information ﬂow manifests as a reduction in the norm of the gradient ( e.g. In Figure 1, the network pruned by SNIP has a zero gradient).
Therefore, we aim to prune connections in a way that accounts for their role in the network’s gradient ﬂow.
Speciﬁcally, we take the gradient norm after pruning as our criterion, and prune those weights whose removal will result in least decrease in the gradient norm after pruning.
Because we rely on preserving the gradient ﬂow to prune the network, we name our method Gradient Signal
Preservation (GraSP).
Our approach is easy to implement and conceptually simple.
Moreover, the recently introduced Neural Tangent Kernel (NTK) (Jacot et al., 2018) provides tools for studying the learning dynamics in the output space.
Building on the analysis of Arora et al.
(2019), we show that our pruning criterion tends to keep those weights which will be beneﬁcial for optimization.
We evaluate GraSP on CIFAR-10, CIFAR-100 (Krizhevsky, 2009), Tiny-ImageNet and ImageNet (Deng et al., 2009) with modern neural networks, such as VGGNet (Simonyan & Zisserman, 2014) and
ResNet (He et al., 2016).
GraSP signiﬁcantly outperforms SNIP in the extreme sparsity regime.
R ELATED W ORK AND B ACKGROUND
In this section, we review the literature on neural network pruning including pruning after training, during training and before training.
Lastly, we also review the Neural Tangent Kernel (NTK)
(Jacot et al., 2018), which builds up the foundation for justifying our method in Section 4.
2. N ETWORK P RUNING
After training.
Most pruning algorithms (LeCun et al., 1990; Hassibi et al., 1993; Dong et al., 2017;
Han et al., 2015b; Li et al., 2016; Molchanov et al., 2016) operate on a pre-trained network.
The main idea is to identify those weights which are most redundant, and whose removal will therefore least degrade the performance.
Magnitude based pruning algorithms (Han et al., 2015b;a) remove those weights which are smaller than a threshold, which may incorrectly measure the importance of each weight.
In contrast, Hessian-based pruning algorithms (LeCun et al., 1990; Hassibi et al., 1993) compute the importance of each weight by measuring how its removal will affect the loss.
More recently, Wang et al.
(2019) proposed a novel network reparameterization based on the Kroneckerfactored Eigenbasis for further boosting the performance of Hessian-based methods.
However, all the aforementioned methods require a pre-trained network, and therefore aren’t applicable at initialization.
During training.
There are also some works which attempt to incorporate pruning into the training procedure itself.
Srinivas & Babu (2016) proposed generalized dropout, allowing for tuning the individual dropout rates during training, which can result in a sparse network after training.
Louizos et al.
(2018) proposed a method for dealing with discontinuity in training L 0 norm regularized networks in order to obtain sparse networks.
Both methods require roughly the same computational cost as training the full network.
Other works, such as Bellec et al.
; Mocanu et al.
Mostafa & Wang (2019); Dettmers & Zettlemoyer (2019), follow the prune-redistribute-regrowth cycle for pruning.
Among them, Dettmers & Zettlemoyer (2019) proposed sparse momentum algorithm, which dynamically determines the sparse mask based on the mean momentum magnitude during training.
However, their method requires to maintain the momentum of all the weights during training, and thus does not save the memory.
These works not only change the standard training procedure, but also hard to enjoy the practical accelerations from hardware (Dey et al., 2019).
Before training.
Pruning at initialization is more challenging because we need to account for the effect on the training dynamics when removing each weight.
There have been several attempts to conduct pruning before training.
Frankle & Carbin (2019); Frankle et al.
(2019) proposed and validated the Lottery Ticket Hypothesis (LTH) that the network structure found by traditional pruning algorithms and the corresponding initialization are enough for training the sub-network from scratch.
Lee et al.
(2018) proposed the SNIP algorithm, which was the ﬁrst attempt to directly identify trainable and sparse sub-networks at initialization time.
Their method was based on connection sensitivity , which aims to preserve the loss after pruning, and achieved impressive results.
Concurrently to our work, Lee et al.
(2019b) studied the pruning problem from a signal propagation perspective, and proposed to use an orthogonal initialization for ensuring faithful signal propagations.
Though sharing the same spirit as our GraSP algorithm, they focus on initialization scheme which is orthogonal to speciﬁc pruning method.
2. N EURAL T ANGENT K ERNEL AND C ONVERGENCE A NALYSIS
Jacot et al.
(2018) analyzed the dynamics of neural net training by directly analyzing the evolution of the network’s predictions in output space.
Let L denote the cost function, X the set of all training samples, Z = f
( X ; θ ) ∈ R nk × 1 the outputs of the neural network, and k and n the output space dimension and the number of training examples.
For a step of gradient descent, the change to the network’s predictions can be approximated with a ﬁrst-order Taylor approximation:
[Equation]
where the matrix Θ t ( X , X ) is the Neural Tangent Kernel (NTK) at time step t :
[Equation]
where ∇ θ f ( X ; θ ) denotes the network Jacobian over the whole training set.
Jacot et al.
(2018) showed that for inﬁnitely wide networks, with proper initialization, the NTK exactly captures the output space dynamics throughout training, which means NTK reveals the ﬁnal state of the network at initialization without any training.
In particular, Θ t ( X , X ) is constant.
Arora et al.
(2019) used the NTK to analyze optimization and generalization phenomena, showing that under the assumptions of constant NTK and squared error loss, the training dynamics can be analyzed in closed form:
[Equation]
where Y ∈ R nk × 1 is all the targets, Θ = UΛU (cid:62)
= (cid:80) ni
=
λ i u i u (cid:62) i is the eigendecomposition, and (cid:15) is a bounded error term.
Although the constant NTK assumption holds only in the inﬁnite width limit,
Lee et al. (2019a) found close empirical agreement between the NTK dynamics and the true dynamics for wide but practical networks, such as wide ResNet architectures (Zagoruyko & Komodakis, 2016).
Therefore, if running the training procedure itself is too expensive (as in our setting), one ought to instead be able to achieve a rough approximation to the training dynamics using the NTK.
R EVISITING S INGLE -S HOT N ETWORK P RUNING (SNIP)
Single-shot network pruning was introduced by Lee et al.
(2018), who used the term to refer both to the general problem setting and to their speciﬁc algorithm.
To avoid ambiguity, we refer to the general problem of pruning before training as foresight pruning .
For completeness, we ﬁrst revisit the formulation of foresight pruning, and then point out issues of SNIP for motivating our method.
Problem Formulation.
Suppose we have a neural network f parameterized by θ ∈ R d , and our objective is to minimize the empirical risk L ( θ )
= 1 N (cid:80) ( x ,y )
∼D [ (cid:96) ( f ( x ; θ ) , y )] given a training set
D = { ( x i , y i ) } ni
=1 .
Then, the foresight pruning problem can be formulated as:
[Equation]
[Equation]
where (cid:100) p · d (cid:101) is the number of weights to be removed, and A is a known training algorithm (e.g. SGD), which takes the mask m
(here we marginalize out the initial weights θ 0 for simplicity), and returns the trained weights.
Since globally minimizing Eqn. 4 is intractable, we are instead interested in heuristics that result in good practical performance.
Revisiting SNIP.
SNIP (Lee et al., 2018) was the ﬁrst algorithm proposed for foresight pruning, and it leverages the notion of connection sensitivity to remove unimportant connections (weights).
They deﬁne this in terms of how removing a single weight θ q in isolation will affect the loss:
[Equation]
δ q ) (cid:15) (cid:12)(cid:12)(cid:12)(cid:12) =
(cid:12)(cid:12)(cid:12)(cid:12)
θ q ∂ L ∂θ q (cid:12)(cid:12)(cid:12)(cid:12) (5) where θ q is the q th element of θ 0 , and δ q is a one-hot vector with only q th element equals to θ q .
Essentially, SNIP aims at preserving the loss of the pruned network almost the same as that of the original network.
Preserving the loss value motivated several classic methods for pruning a trained network , such as optimal brain damage (LeCun et al., 1990) and optimal brain surgery (Hassibi et al., 1993).
While the motivation for loss preservation of a trained network is clear, it is less clear why this is a good criterion for foresight pruning.
After all, at initialization, the loss is no better than chance.
We argue that at the beginning of training, it is more important to preserve the training dynamics than the loss itself.
SNIP does not do this automatically, because even if removing a particular connection doesn’t affect the loss, it could still block the ﬂow of information through the network.
For instance, we noticed in our experiments that SNIP with a high pruning ratio (e.g. 99%) tends to eliminate nearly all the weights in a particular layer, creating a bottleneck in the network.
Therefore, we would prefer a pruning criterion which accounts for how the presence or absence of one connection inﬂuences the training of the rest of the network.
G RADIENT S IGNAL P RESERVATION
We now introduce and motivate our foresight pruning criterion, Gradient Signal Preservation (GraSP).
To understand the problem we are trying to address, observe that the network after pruning will have fewer parameters and sparse connectivity, hindering the ﬂow of gradients through the network and potentially slowing the optimization.
This is reﬂected in Figure 3, which shows the reduction in gradient norm for random pruning with various pruning ratios.
Moreover, the performance of the pruned networks is correspondingly worse (see Table 1).
Mathematically, a larger gradient norm indicates that, to the ﬁrst order, each gradient update achieves a greater loss reduction, as characterized by the following directional derivative:
[Equation]
Since we only care about the performance of the pruned network, our goal is to preserve or even increase (if possible) the gradient ﬂow after pruning ( i.e. , the gradient ﬂow of the pruned network).
[Table]
In order to achieve this
[Table]
as adopted in LeCun et al.
[Table]
we cast the pruning operation as adding a perturbation δ on the initial weights for characterizing how removing one weight will affect the gradient ﬂow after pruning :
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where S ( δ ) is a function of δ , and it denotes the change in (6) to a δ -perturbation on the initial weight θ 0 .
The Hessian matrix H captures the correlations between each weight, and thus modulates the effect of pruning on the remained weights.
When H is identity, the above criteria recovers SNIP up to the absolute value (recall the criteria of SNIP is | δ (cid:62) g | ).
However, it has been observed that different weights are highly coupled (Hassibi et al., 1993), indicating that the Hessian H is far from identity.
In practice, we adopt the equation deﬁned above ( i.e. , eqn.
as the measure of the importance of each weight.
Speciﬁcally, if S ( δ ) is negative, then removing the corresponding weights will reduce the gradient ﬂow, otherwise it will not.
Hence, we will prefer removing those weights whose removal will not decrease the gradient ﬂow ﬁrst.
For each weight, the importance can be computed in the following way:
[Equation]
For a given pruning ratio p , we can obtain the resulting pruning mask by computing the pruning criteria in a single pass, sort them accordingly and then remove the top p % of the weights (see
Algorithm 1).
Roughly speaking, GraSP takes change in gradient ﬂow after pruning into account, whereas SNIP only tries to preserve the loss after pruning, which may not perserve the gradient ﬂow.
Intuitively speaking, GraSP is accounting for how removing one weight will affect the gradient of each other weight.
Moreover, our criteria is easy to implement, and can be computed by backward twice without explicitly computing the Hessian (Pearlmutter, 1994; Schraudolph, 2002).
4. U NDERSTANDING G RA SP THROUGH L INEARIZED T RAINING D YNAMICS
The above discussion concerns only the training dynamics at initialization time.
To understand the longer-term dynamics, we leverage the recently proposed Neural Tangent Kernel (NTK), which has been shown to be able to capture the training dynamics throughout training for practical networks (Lee et al., 2019a).
Speciﬁcally, as we introduced in section 2.2, eqn.
(3) characterizes how the training error changes throughout the training process, which only depends on time step t , training targets Y and the NTK Θ .
Since the NTK stays almost constant for wide but practical networks (Lee et al., 2019a), e.g. , wide ResNet (Zagoruyko & Komodakis, 2016), we can regard that eqn.
3 only depends on t once the network is initialized.
This shows that NTK captures the training dynamics throughout the training process.
Therefore, if we decompose eqn.
(6) in the following form:
[Equation]
By relating it to eqn.
(3), we can see that our objective is implicitly encouraging the eigenspace of Θ to align well with the target Y by pruning, which controls how each gradient step changes the training error.
Formally speaking, when we are able to maximize eqn.
(9), it will encourage the eigenspace of the NTK distributing large eigenvalues in the direction of Y , which in turn accelerates the decrease of the loss (Arora et al., 2019) and beneﬁts the optimization in A .
[Table]
E XPERIMENTS
In this section, we conduct various experiments to validate the effectiveness of our proposed pruning algorithm in terms of the test accuracy vs. pruning ratios.
We also include two traditional pruning algorithms (LeCun et al., 1990; Zeng & Urtasun, 2019), which operate on the pre-trained networks, for serving as an upper bound for foresight pruning.
Besides, we study the convergence performance of the sub-networks obtained by different pruning methods for investigating the relationship between gradient norm and ﬁnal performance.
Lastly, we study the role of initialization and batch sizes in terms of the performance of GraSP for ablation study.
5. P RUNING R ESULTS ON M ODERN C ONV N ETS
To evaluate the effectiveness of GraSP on real world tasks, we test GraSP on four image classiﬁcation datasets, CIFAR-10/100, Tiny-ImageNet and ImageNet, with two modern network architectures,
VGGNet and ResNet 1 .
For the experiments on CIFAR-10/100 and Tiny-ImageNet, we use a minibatch with ten times of the number of classes for both GraSP and SNIP 2 according to Lee et al.
The pruned network is trained with Kaiming initialization (He et al., 2015) using SGD for 160 epochs for CIFAR-10/100, and 300 epochs for Tiny-ImageNet, with an initial learning rate of 0 . 1 and batch size 128 .
The learning rate is decayed by a factor of 0 .
1 at 1 / 2 and 3 / 4 of the total number of epochs.
Moreover, we run each experiment for 3 trials for obtaining more stable results.
For ImageNet, we adopt the Pytorch (Paszke et al., 2017) ofﬁcial implementation, but we used more epochs for training according to Liu et al.
Speciﬁcally, we train the pruned networks with
SGD for 150 epochs, and decay the learning rate by a factor of 0.1 every 50 epochs.
We ﬁrst compre GraSP against random pruning, which generates the mask randomly for a given pruning ratio.
The test accuracy is reported in Table 1.
We can observe GraSP clearly outperforms random pruning, and the performance gap can be up to more than 30% in terms of test accuracy.
We further compare GraSP with more competitive baselines on CIFAR-10/100 and Tiny-ImageNet for pruning ratios { 85% , 90% , 95% , 98% } , and the results can be referred in Table 2 and 4.
When the pruning ratio is low, i.e. 85% , 90% , both SNIP and GraSP can achieve very close results to the baselines, though still do not outperform pruning algorithms that operate on trained networks, as expected.
However, GraSP achieves signiﬁcant better results for higher pruning ratios with more complicated networks (e.g. ResNet) and datasets (e.g. CIFAR-100 and Tiny-ImageNet), showing the advantages of directly relating the pruning criteria with the gradient ﬂow after pruning.
Moreover, in most cases, we notice that either SNIP or GraSP can match or even slightly outperform LT with original initialization, which indicates that both methods can identify meaningful structures.
We further experiment with the late resetting as suggested in Frankle et al.
by resetting the weights
[Table]
[Figure] [Figure]
of the winning tickets to their values
at 5 epoch.
By doing so, we observe a boost in the performance of the LT across all the settings, which is consistent with (Frankle et al., 2019).
[Table]
However, the above three datasets are overly simple and small-scale for reaching robust conclusions, and thus we conduct further experiments on ImageNet using ResNet-50 and VGG16 with pruning ratios { 60% , 80% , 90% } to further validate the effectiveness of GraSP.
The results are shown in Table 3, we can see that when the pruning ratio is 60% , both SNIP and
GraSP can achieve very close performance to the original one, and these two methods perform almost the same.
However, as we increase the pruning ratio, we observe that GraSP surpasses
SNIP more and more.
When the pruning ratio is 90% , GraSP can beat SNIP by 6 . 2% for ResNetand 4 .
7% for VGG16 in top-1 accuracy, showing the advantages of GraSP at larger pruning ratios.
We seek to investigate the reasons behind in Section 5.2 for obtaining a better understanding on
GraSP.
5. A NALYSIS ON C ONVERGENCE P ERFORMANCE AND G RADIENT N ORM
Based on our analysis in Section 4.1 and the observations in the previous subsection, we conduct further experiments on CIFAR-100 with ResNet32 to investigate where the performance gains come from in the high sparsity regions.
We present the training statistics in terms of the training and test loss in Figure 2.
We observe that the main bottleneck of pruned neural networks is underﬁtting, and thus support our optimization considerations when deriving the pruning criteria.
As a result, the network pruned with GraSP can achieve much lower loss for both training and testing and the decrease in training loss is also much faster than SNIP.
Besides, we also plot the gradient norm of pruned networks at initialization for a ResNet32 on
CIFAR-100.
The gradient norm is computed as the average of the gradients over the entire dataset, and we normalize the gradient of the original network to be 1 .
We run each experiment for 10 trials in order to obtain more stable results.
We observe that both SNIP and GraSP result in a lower gradient norm at high sparsity (e.g. 98%), but GraSP can better preserve the gradient norm after pruning, and also yield better results than SNIP at the high sparsity regions.
Moreover, in these high sparsity regions, the pruned network usually underﬁts the training data, and thus the optimization will be a
[Figure]
problem.
The randomly pruned network has a much lower gradient norm and performs the worst.
In contrast, the network pruned by GraSP can start with a higher gradient norm and thus hopefully more training progress can be made as evidenced by our results in Figure 2.
5. V ISUALIZING THE N ETWORK AFTER P RUNING
In order to probe the difference of the network pruned by SNIP and GraSP, we present the portion of the remaining weights at each layer of the sparse network obtained by SNIP and GraSP in Figure.
We observe that these two pruning methods result in different pruning strategies.
Speciﬁcally, GraSP aims for preserving the gradient ﬂow after pruning, and thus will not prune too aggressively for each layer.
Moreover, it has been known that those convolutional layers at the top usually learn highly sparse features and thus more weights can be pruned (Zhang et al., 2019a).
As a result, both methods prune most of the weights at the top layers, but GraSP will preserve more weights than SNIP due to the consideration of preserving the gradient ﬂow.
In contrast, SNIP is more likely to prune nearly all the weights in top layers, and thus those layers will be the bottleneck of blocking the information ﬂow from the output layer to the input layer
. 5. E FFECT OF B ATCH S IZE AND I NITIALIZATION
[Table]
[Table] [Table]
[Table] [Table]
We also study how the batch size and initialization will affect GraSP on CIFAR-10 and CIFAR100 with ResNet32 for showing the robustness to different hyperparameters.
Speciﬁcally, we test
GraSP with three different initialization methods: Kaiming normal (He et al., 2015), Normal N (0 , 0 . 1) , and Xavier normal (Glorot &
Bengio, 2010), as well as different mini-batch sizes.
We present the mean and standard variance of the test accuracy obtained with different initialization methods and by varying the batch sizes in Table 5.
For CIFAR-10 and CIFAR-100, we use batch sizes { 100 , 400 , · · · , 25600 , 50000 } and { 1000 , 4000 , 16000 , 32000 , 50000 } respectively.
We observe that GraSP can achieve reasonable performance with different initialization methods, and also the effect of batch size is minimal for networks pruned with Kaiming initialization, which is one of the most commonly adopted initialization techniques in training neural networks.
D ISCUSSION AND C ONCLUSION
We propose Gradient Signal Preservation (GraSP), a pruning criterion motivated by preserving the gradient ﬂow through the network after pruning.
It can also be interpreted as aligning the large eigenvalues of the Neural Tangent Kernel with the targets.
Empirically, GraSP is able to prune the weights of a network at initialization , while still performing competitively to traditional pruning algorithms, which requires ﬁrst training the network.
More broadly, foresight pruning can be a way for enabling training super large models that no existing GPUs can ﬁt in, and also reduce the training cost by adopting sparse matrices operations.
Readers may notice that there is still a performance gap between GraSP and traditional pruning algorithms, and also the LT with late resetting performs better than LT with original initialization.
However, this does not negate the possibility that foresight pruning can match the performance of traditional pruning algorithms while still enjoying cheaper training cost.
As evidence, Evci et al.
(2019) show that there exists a linear and monotonically decreasing path from the sparse initialization to the solution found by pruning the fully-trained dense network, but current optimizers fail to achieve this.
Therefore, designing better gradient-based optimizers that can exploit such paths will also be a good direction to explore.