L EARNING P RUNING -F RIENDLY N ETWORKS VIA F RANK -W OLFE : O NE -S HOT , A NY -S PARSITY , A ND N O R ETRAINING
A BSTRACT
We present a novel framework to train a large deep neural network (DNN) for only once , which can then be pruned to any sparsity ratio to preserve competitive accuracy without any re-training .
Conventional methods often require (iterative) pruning followed by re-training, which not only incurs large overhead beyond the original DNN training but also can be sensitive to retraining hyperparameters.
Our core idea is to re-cast the DNN training as an explicit pruning-aware process: that is formulated with an auxiliary K -sparse polytope constraint, to encourage network weights to lie in a convex hull spanned by K -sparse vectors, potentially resulting in more sparse weight matrices.
We then leverage a stochastic Frank-Wolfe (SFW) algorithm to solve this new constrained optimization, which naturally leads to sparse weight updates each time.
We further note an overlooked fact that existing DNN initializations were derived to enhance SGD training (e.g., avoid gradient explosion or collapse), but was unaligned with the challenges of training with SFW.
We hence also present the ﬁrst learning-based initialization scheme speciﬁcally for boosting SFW-based DNN training.
Experiments on CIFARand Tiny-ImageNet datasets demonstrate that our new framework named SFWpruning consistently achieves the state-of-the-art performance on various benchmark DNNs over a wide range of pruning
ratios. Moreover, SFW-pruning only needs to train once on the same model and dataset, for obtaining arbitrary ratios, while requiring neither iterative pruning nor retraining.
Codes are available in https://github.com/VITA-Group/SFW-Once-for-All-Pruning .
I NTRODUCTION
Deep neural networks (DNNs) achieve tremendous empirical success in various machine learning applications, but this usually requires a huge model size and high computational cost, challenging the usage of such models in various real-time and multi-platform applications.
For example (Cai et al., 2019), mobile applications on App stores have to support a diverse range of hardware devices, from high-end ﬂagships to low-end ones, which need different model capacities for different scenarios.
As a remedy, neural network compression (Han et al., 2015) tries to slim various DNNs to improve both the computation and memory efﬁciency.
One important way of network compression is network pruning, which prunes individual weights (unstructured pruning, Han et al.
or channels (structured pruning, Li et al.
(2016)) of the network, causing little degeneration in the test performance.
Usually, modern DNN pruning techniques require retraining or ﬁne-tuning of the compressed network, either in a one-shot manner or iterative manner (Frankle & Carbin, 2018).
When one aims to deploy the compressed neural network in different platforms, retraining the pruned model requires huge computational cost, resulting in excessive energy consumption.
To tackle such an efﬁciency problem, in this paper, we aim to answer the question of whether we can design an efﬁcient pruning method that does not need retrain the neural network.
Speciﬁcally, we hope to design one-shot unstructured pruning algorithms, which can guarantee consistent and competitive model performance under varying pruning ratios without retraining the neural network.
One-shot pruning is a natural choice since retraining is not allowed.
Intuitively, given a sparsity ratio, one can leave the most important individual weights (weight-magnitude unstructured pruning (Han et al., 2015)) untouched and the less important ones pruned according to their relative ”absolute values”.
In practice, different pruning ratio requirements correspond to different computational and memorial budgets.
However, previous pruning methods cannot address our question, since they mainly ﬁnd those important weights via standard gradient-based optimization methods during training, e.g., stochastic gradient descent (SGD), whose performance will degrade much if retraining is prohibited, especially at large pruning ratios (See Figure 1(a)).
The reason is that such optimization methods fail to take the desired sparse structures in the pruned model into consideration.
The process of ﬁnding those most important weights should be integrated with the optimization method.
In view of those, we innovate to re-cast DNN training as an explicit pruning-aware process, which is formulated with an auxiliary K -sparse polytope constraint to encourage the DNN weights to lie in the convex hull spanned by K -sparse vectors.
Correspondingly, we also propose to train the
DNN via solving the new constrained optimization problem using the stochastic Frank-Wolfe (SFW) algorithm (Reddi et al., 2016; Hazan & Luo, 2016) that can more organically handle constraints than
SGD-type algorithms.
SFW results in sparse weight updates and continually pushes less important weights to smaller magnitudes.
We observe DNNs trained in this way to have more small-value ( yet non-zero ) weights that can be gradually and smoothly removed when increasing the (one-shot) pruning ratio.
Applying one-shot weight magnitude pruning then yields consistently strong test performance across different pruning ratios without further retraining or ﬁne-tuning.
Moreover, we propose a new initialization scheme for SFW-based training.
We note that SFW algorithm is not a gradient-based optimization method, and the widely used initialization schemes like Kaiming initialization (He et al., 2015) for gradient-based methods do not ﬁt SFW algorithm since the latter does not use the gradient to update weights.
Our proposed initialization scheme is organically designed for SFW algorithm and is formulated as a meta learning problem, drawing motivations from Zhu et al.
It learns the layer-wise initialization scaling factors that lead to the largest loss reduction in the ﬁrst SFW training step.
We demonstrate that with the new initialization scheme, our proposed one-shot pruning algorithm can consistently achieve better test performance under different pruning ratios without retraining.
Now we summarize our main contributions.
• We explicitly re-cast DNN training as a pruning-aware process formulated under an auxiliary K -sparse polytope constraint, based on which we propose a new SFW-pruning framework that trains a DNN via solving a new constraint optimization problem using SFW and prunes the DNN in a one-shot fashion with no retraining.
• We customize a meta-learning-based initialization scheme for SFW-based DNN training, which leads to more consistent and competitive performance under varying pruning ratios.
• We empirically demonstrate that our proposed SFW-pruning framework is applicable across different architectures and datasets, achieving the state-of-the-art performance consistently over a wide range of pruning ratios without retraining.
R ELATED W ORK
Neural Network Pruning and Efﬁcient Deployment.
With larger, deeper, and more sophisticated models, DNNs achieved incredible success over the last decade.
However, large models introduce high computation and memory costs, making them difﬁcult in actual application and deployment.
Various methods, including knowledge distillation (Hinton et al., 2015; Romero et al., 2014), lowrank factorization (Denton et al., 2014; Yu et al., 2017), quantization (Courbariaux et al., 2016; Rastegari et al., 2016; Wu et al., 2018), and pruning (Han et al., 2015; Li et al., 2016; Liu et al., 2018), have been proposed to deploy large models in resource-constrained devices efﬁciently.
Among these methods, pruning has become a research hotspot for its ability to maintain high performance.
Typically, pruning can be roughly divided into two main branches: unstructured (Han et al., 2015) and structured (Hu et al., 2016).
Unstructured pruning is to remove individual parameters in the networks.
A common example, Iterative Magnitude Pruning (IMP) (Han et al., 2015; Frankle & Carbin, 2018; Chen et al., 2020b;a), repeats training and pruning cycles to attain a sparse network.
Although IMP works well, its high computation cost motivates more efﬁcient methods (Lee et al., 2018;
Wang et al., 2020a; Tanaka et al., 2020).
Another major branch is structured pruning.
It considers parameters in groups, removing entire neurons, ﬁlters, or channels (He et al., 2017; Li et al., 2016;
Guo et al., 2021) which can more directly save energy and memory (Luo & Wu, 2020).
Practical devices often vary in their on-board resource availability (Cai et al., 2019; Wang et al., 2020b).
To tackle this problem, pruning the network to different sizes and retraining is one naive way, yet costly.
(Han et al., 2015; He et al., 2017; Li et al., 2016).
Another approach is to design efﬁcient and scalable neural network architecture such as MobileNet (Howard et al., 2017) and ShufﬂeNets (Ma et al., 2018; Zhang et al., 2018).
However, these methods are computationally expensive or need human-based design.
To this end, Cai et al.
(2019) pioneer on a “once for all” scheme that can ﬂexibly obtain a large number of sub-networks to meet different resource constraints from a pre-trained “super network”, but their pre-training cost is gigantic and re-training is needed for each sub-network for restoring the optimal performance.
In our work, we view efﬁcient deployment from a novel aspect, i.e. training a neural network to naturally possess scalable weight sparsity, while keeping the training cost comparable with normal one-pass training.
Then with only a single-shot pruning, we can get the desired sub-network with any sparsity without retraining.
Stochastic Frank-Wolfe.
Frank–Wolfe (FW) (Frank et al., 1956) is a classical non-gradient-based method for convex optimization.
In recent years, FW has been applied in stochastic non-convex optimization (Reddi et al., 2016; Hazan & Luo, 2016), named as stochastic Frank–Wolfe (SFW) algorithm.
Hazan & Luo (2016) perform a theoretical analysis of the standard SFW algorithm with a convergence rate O (1 /t ) with Θ (cid:0) t 2 (cid:1) growing batch-sizes.
Several works further extend SFW-based algorithms in the deep learning ﬁeld and propose various variants of SFW
(Yurtsever et al., 2019; Shen et al., 2019; Xie et al., 2019; Zhang et al., 2020; Combettes et al., 2020).
Momentum is found to be a crucial part of SFW in training neural networks (Cutkosky & Orabona, 2019; Mokhtari et al., 2020;
Chen et al., 2018).
Recently, (Pokutta et al., 2020) apply a momentum based SFW algorithm in training NN and achieve high accuracy. Inspired from those prior arts, we propose SFW-based
NN training algorithm for one-shot pruning without retraining. We additionally investigate the previously overlooked problem of SFW initialization, and present a meta learning-based scheme.
Neural Networks Initialization.
An inappropriate initialization can lead to deep networks weights and gradients exploding or vanishing, challenging training deep neural networks.
Several standard initialization (Glorot & Bengio, 2010; He et al., 2015) methods are designed for gradient-based optimization methods to keep the variance per layer balanced.
Further, Glorot & Bengio (2010) and He et al.
(2015)’s
assumptions no longer hold for more complex architectures, motivating newer methods (Dauphin & Schoenholz, 2019; Zhang et al., 2019; Bachlechner et al., 2020).
However, most of these are tied to enhancing SGD-based training.
Since the K -sparse SFW algorithm is non-gradient based, new initialization methods need to be customized.
S TOCHASTIC F RANKE -W OLFE P RUNING F RAMEWORK
In this section, we formulate the DNN training process as an explicit pruning-aware process with an auxiliary K -sparse polytope constraint, and we solve the corresponding constrained optimization problem via a stochastic Franke-Wolfe (SFW) algorithm.
Speciﬁcally, we formulate the pruningaware process and review the basics of SFW algorithm in Section 3.1.
In Section 3.2, we present our proposed SFW-pruning framework based on the pruning-aware process and one-shot pruning, where we consider weight-magnitude unstructured pruning.
3. P RUNING -A WARE DNN T RAINING AND S TOCHASTIC F RANK -W OLFE A LGORITHM
Our core idea is to re-cast the DNN training process as an explicit pruning-aware process, which is achieved via an auxiliary K -sparse polytope constraint.
Formally, given a dataset D = { ( x i , y i ) } ni = where ( x i , y i ) ∈
X × Y and a loss function (cid:96) :
Y × Y (cid:55)→ R + , e.g., cross-entropy loss, we aim to train a deep neural network f
( θ ; · ) : X (cid:55)→ Y that minimizes the following pruning-aware objective :
[Equation]
To make the training process pruning-aware, we restrict the feasible parameter θ in a convex region
C , which potentially results in more sparse weight matrices and will be more pruning-friendly for only a small percentage of the weights are of large magnitudes.
Speciﬁcally, we choose C as a Ksparse polytope, and we solve (1) efﬁciently via a stochastic Frank–Wolfe (SFW) algorithm.
We review the K -sparse polytope and the basics of SFW below and further explain our motivations.
Stochastic Frank-Wolfe and K -Sparse Polytope Constraints.
SFW is a simple projection-free ﬁrst-order algorithm for solving convex constraint optimization problems (Reddi et al., 2016; Hazan & Luo, 2016;
Yurtsever et al., 2019; Shen et al., 2019; Xie et al., 2019; Zhang et al., 2020; Combettes et al., 2020).
Consider the constrained optimization objective (1) and denote the neural network weights learned by SFW by θ t .
At each iteration t , SFW ﬁrst calls a linear minimization oracle v t = LMO C ( (cid:98) ∇ θ L ( θ t ))
=
arg min v ∈C (cid:104) (cid:98) ∇ θ L ( θ t ) , v (cid:105) , where (cid:98) ∇ θ L ( θ t ) estimates the full gradient, e.g., gradient on a minibatch.
Given v t , SFW updates θ t along the direction of v t by a convex combination θ t +
=
θ t + α t ( v t − θ t ) =
α t v t + (1 − α t ) θ t with the learning rate α t ∈ [0 , 1] .
This keeps θ t always in C and saves any projection step.
In neural network training, to further improve the training performance and test accuracy, momentum is also introduced into SFW algorithm
(Xie et al., 2019; Pokutta et al., 2020).
We conclude the corresponding SFW algorithm in Algorithm 1.
A K -sparse polytope in R p of radius τ > 0 , p ∈ N , denoted by C ( K, τ ) , is the intersection of the
L 1 -ball B 1 ( τK ) and the L ∞ -ball B ∞ ( τ ) of R p .
One can obtain C ( K, τ ) by spanning all the vectors in R p , which have exactly K non-zero coordinates and the absolute value of the non-zero entries are τ .
It holds that (Pokutta et al., 2020) for any m ∈ R p the oracle v = LMO C ( K,τ ) ( m ) is given by ( v ) i = (cid:26) − τ · sign(( m ) i ) if ( m ) i is in the largest K coordinates of m , otherwise, ∀ 1 ≤ i ≤ p, (2) which is a vector with exactly K non-zero entries.
When applying the K -sparse polytope constraint in deep neural network training objective (1), we add the constraint on each layer of the network.
In other words, if we write out θ as θ =
( W 1 , b 1 , · · · , W L , b L ) for layer-wise weight parameter
W l and bias parameter b l , then each of the parameters is paired with a K -sparse polytope constraint
C ( K l , τ ) .
Here K l may vary between the different layers.
For notation simplicity, in the sequel, we denote such a layer-wise constraint as θ ∈ C
( { K l } Ll =1 , τ ) without making any confusion.
Why K -Sparse Polytope Constraints?
In each step of our optimization, the linear minimization oracle of a K -sparse polytope constraint returns a update vector with exactly K non-zero coordinates, which is then weighted-averaged with the current θ t according to the SFW algorithm (accumulated from the “ensemble” of K non-zero coordinate vectors, from all past update steps).
In other words, training with a K -sparse constraints amounts to a “voting” process, by all step updates, on which K elements “should be non-zero”, and the resulting weights could be viewed as the soft voting consensus.
By adding such a training constraint, each SFW step pushes less important weights smaller since they are averaged with zero, and those important weights are enhanced meanwhile.
As Figure 1(a) shows, compared to DNNs trained by SGD, those trained by SFW with K -sparse polytope constraints appear to have much more smaller weights ( but not exactly zero ), and less large ones.
Moreover, the amounts of weights, at different magnitude levels, change more “smoothly and “continually” in SFW-trained weights compared to that of SGD.
With such a weight distribution, when the pruning ratio increases, there will be no sudden “jump” to removing small values to impacting larger values) when the pruning ratio increases, yielding competitive test accuracies across the whole spectrum of pruning ratios, even without needing retraining.
We refer to Appendix B. for more computational details about K -Sparse Polytope Constraints.
3. S TOCHASTIC F RANK –W OLFE O NE - SHOT P RUNING WITHOUT R ETRAINING
Previously, we have motivated the usage of K -sparse polytope constraint in the pruning-aware DNN training objective (1) and using SFW algorithm in the training process.
In this subsection, we ﬁrst propose our SFW-pruning framework for deep neural network one-shot pruning without retraining based on SFW neural network training with K -sparse polytope constraints.
After, we provide some ﬁne-grained analysis of the hyperparameters involved in the newly proposed pruning framework.
Stochastic Frank–Wolfe Pruning Framework.
We now present our one-shot pruning framework, where a DNN is trained using the pruning-aware objective (1) with K -sparse constraints via SFW for only once and then undergoes a one-shot pruning (Algorithm 2).
The SFW training phase (Line 5 to 7) is pruning-aware, which aims to ﬁnd a proper DNN weight θ T that can minimize the training [Figure] loss while being pruning-friendly.
The following pruning procedure (Line 8 to 9) is adpoted with standard weight-magnitude unstructured pruning (Han et al., 2015; Frankle & Carbin, 2018) (prunes individual weights in the DNN).
For practical use, one can choose either of the pruning procedures.
We refer to technical details of the two pruning procedures to experiments in Section 5.
In Figure 1(b), we demonstrate the effectiveness of our framework using ResNet18 (He et al., 2016) on the CIFAR-10 (Krizhevsky
Despite the simplicity, the high-level idea of Algorithm 2 is intuitive: the discovery of sparse structures and pruning-aware DNN weights should be integrated with the optimization process , especially when retraining is prohibited.
We further demonstrate the effectiveness of Algorithm 2 via extensive experiments on different datasets and DNN architectures with unstructured pruning, achieving competitive test performance consistently across various pruning ratios, which is not achievable by random pruning or non-pruning-aware optimization methods.
Fine-grained Analysis of SFW Sparse Training.
Proper choices of several hyperparameters of
Algorithm 2 play an important in discovering pruning-friendly weights.
We note that intuitively it is a trade-off between the test performance and the one-shot pruning ratio of the DNN.
But given the same pruning ratio, our framework can ﬁnd more pruning-friend weights than random pruning or other optimization methods without sparsity awareness.
We refer to Appendix B.1 for a detailed ablation study of the hyperparameters in Algorithm 2.
K -sparse constraint parameters { K l } Ll =1 and τ : { K l } Ll =1 controls the sparsity of the linear minimization oracle solution, and τ controls the magnitude, which together inﬂuence the distribution of the learned parameters.
Experimental results show that relatively smaller choices of these hyperparameters can keep the test performance undamaged under higher one-shot pruning ratios, at the cost of relatively lower full model test performance.
Empirically, we ﬁnd K
= 5% and τ = 15 1 to be our default good choices that transfer well across DNNs and datasets.
Learning rate α : Recall that at step t , we update θ t by θ t +
=
θ t + α t ( v t − θ t ) =
α t v t +(1 − α t ) θ t which is the a convex combination of v t , the linear minimization orcale solution, and θ t , the current parameter.
Note that v t is K -sparse, which means that a larger learning rate (i.e., α t v t has larger weight in the convex combination) may result in sparser parameter θ t +1 .
This observation indicates that one can (dynamically) control the nature sparsity 2 of the neural network by controlling the scale of the learning rate α t .
We refer to Figure 1(c) for a demonstration of this phenomenon.
I NITIALIZATION S CHEME FOR S TOCHASTIC F RANK -W OLFE A LGORITHM
Most widely used initialization schemes for gradient-based methods, like Kaiming initialization (He et al., 2015), ensure that the mean of activations is zero, and the variance of the activations stays the same across layers to avoid gradient explosion or collapse.
However, the SFW algorithm does not use the gradient to update the weights, making such initialization schemes unaligned with the challenges of training with SFW.
A new initialization designed for SFW is desired.
To this end, we introduce a new initialization scheme by adapting the idea from Zhu et al.
The goal is to learn the best initialization θ 0 in the sense that it can allow the maximal loss reduction in the ﬁrst SFW step, which we hope can further bring more competitive and consistent pruning test performance in our SFW-pruning paradigm.
Speciﬁcally, the proposed method initializes each layer of the neural network with a uniform distribution, which is then multiplied by a layer-wise scaling parameter learned by the algorithm, resulting in a maximal one-SFW-step loss reduction.
Methodology.
We ﬁrst initialize the parameter θ 0 of weight matrices { W l } Ll =1 and bias vectors { b l } Ll =1 of the network with values drawn from independent zero-mean Gaussian distributions with the variance decided by the standard fan-in and fan-out of the layer (He et al., 2015).
l = 1 · · · L denotes the layer index.
For each layer l , we pair W l and b l with learnable non-negative scalars α i and β i that control the scaling of the layer at initialization, and we use β to denote the vector of scale factors ( ξ 1 , ζ 1 , · · · , ξ L , ζ L ) , and let θ β be the tuple ( ξ 1 W 1 , ζ 1 b 1 , · · · , ξ L W L , ζ L b L ) of rescaled parameters.
Also, we extend the loss deﬁnition as
[Equation]
which is the average loss of the model with θ on a minibatch of samples S .
Correspndingly, we denote one SFW optimization step obtained on minibatch S as SFW ( S, f, C
( { K l } Ll =1 , τ ) , θ 0 , 1 , α 0 , ρ ) .
Speciﬁcally, we optimize the scaling parameter β to obtain a good initialization θ β tailored for a stochastic Frank–Wolfe algorithm, i.e., SFW .
Consider the ﬁrst update step, starting from the rescaled initialization, i.e., θ 1 = SFW ( S, f, C ( { K l } Ll =1 , τ ) , θ 0 , 1 , α 0 , ρ ) .
We choose β so that the loss on θ 1 is as low as possible.
Formally, we optimize the scaling β through the following objective:
[Equation]
where S and ˜ S are two different random minibatches.
For his paper, we only consider K -sparse polytope constraints, which can be easily extended to other constraints like (cid:96) p -norm balls.
To make sure that the rescaled parameter θ β is still a feasible initialization, we make the constraint that θ β ∈ C ( K, τ ) .
We conclude the new initialization scheme for SFW in Algorithm 3.
Our algorithm features the following points, which differ from the work of Zhu et al.
• We do not apply constraints on the norm of the gradient ∇ θ L t ( θ β t ) , which is the case in Zhu et al.
This is because for K -sparse polytope constraints the linear minimization oracle LMO C
( { K l } Ll =1 ,τ ) ( ∇ θ L t ( θ β t )) only depends on the direction rather than norm.
• Since we need to ensure that the initialization is feasible, we project the obtained rescaled initialization θ β t back to the K -sparse polytope constraints at each training iteration.
This is done via a simple scaling of the current parameter to the K -sparse polytope C ( τ, K ) .
Besides, there are two other key points in the SFW-Init algorithm, motivated by Zhu et al.
The ﬁrst point is that we use different but overlapping minibatches S and ˜ S to calculate the current loss and the ﬁrst SFW step respectively.
The difference between two mini-batches is to prevent overﬁtting, while the overlapping of the minibatches is for better train performance.
See Zhu et al.
(2021) for a study of S and ˜ S .
The second point is that we detach the SFW algorithm step from the gradient ∇ β L t +1 (line 7) in order to facilitate the differentiation of the SFW algorithm (which involves sorting the coordinates of the gradient vector) and potential second-order derivatives.
The corresponding SFW algorithm modiﬁcation is marked in the grey color.
E XPERIMENTS 5. S ETTINGS
We summarize the key experiment setups, with hyperparameters of the implementation presented in Appendix A in detail.
We conduct experiments via two popular architectures, ResNet-18 (He et al., 2016) and VGG-16 (Simonyan & Zisserman, 2014), on two benchmark datasets, CIFAR-
(Krizhevsky et al., 2009) and Tiny-ImageNet (Wu et al., 2017).
Speciﬁcally, on all these setups, we demonstrate our SFW-pruning framework (Algorithm 2) by ﬁrst solving the K -sparse polytope constrained optimization problem via SFW with momentum and then applying one-shot pruning to different sparsity ratios.
We train each network for 180 epochs using a dynamic-changing learning
[Figure]
We ﬁrst evaluate our SFW-pruning framework in unstructured weight-magnitude pruning paradigm (Han et al., 2015; Frankle & Carbin, 2018) across a total of four combinations of datasets and architectures, i e., CIFAR-10 with { ResNet-18, VGG16 } and Tiny-ImageNet with { ResNet-18, VGG16 } . Note that for the fairness of comparison, we do not use our learned initialization (SFW-Init) here.
Speciﬁcally, after training the DNN with K -sparse polytope constraints via SFW, we prune it to various sparsity ratios by leaving the largest weights unchanged and the smaller ones pruned.
We consider neural network training without constraints via SGD and random one-shot pruning for comparison.
Figure 2 shows the achieved test performance of the pruned neural networks with different sparsities.
Several observations can be drawn from Figure 2: (i) without retraining, SFW-pruning signiﬁcantly outperforms magnitude-based and random pruning by SGD, across all architectures and datasets; (ii) over a wide range of sparsity ratios, SFW can keep pruning while maintaining a highly competitive performance without any retraining cost.
SFW-pruning is Competitive Over “SGD + Retraining”. We also compare our SFW-pruning to
SGD-pruning with retraining, which is to our disadvantage.
In Figure 3, SFW-pruning can achieve comparable accuracies against SGD-pruning with extra retraining cost, across most pruning ratios.
In contrast, without retraining, SGD-pruning can only tolerate pruning ratios above 30%.
This further validates that SFW-pruning is sparsity-aware.
Network weights of less importance have already converged to small magnitudes by SFW.
Unlike the performance gap induced by hard-thresholding by SGD, important weights found and trained by SFW-pruning are ready for pruning and inference.
5. A DDING THE N EW SFW I NITIALIZATION
We then evaluate the effectiveness of the initialization scheme (SFW-Init) designed speciﬁcally for SFW algorithm in Section 4.
We repeat the previous experiments for evaluating SFW-pruning framework in an unstructured paradigm, comparing the test performance with and without the new initialization scheme.
The corresponding results are shown in Figure 4.
We can see that the performance after adopting our customized SFW initialization is consistently better than Kaiming initialization. [Figure] [Figure]
C OMPARISON TO S TATE - OF - THE -A RT M ETHODS
Blalock We now compare our holistic solution: “ SFW-pruning + SFW-Init ”, with other state-of-the-art (SOTA) unstructured pruning methods (Lee et al., 2019; Blalock et al., 2020; Tanaka et al., 2020).
Note that most previous works require either iterative pruning or retraining for each sparsity ratio, which costs extremely heavy overhead and is to our disadvantage.
The full comparison of
VGG-16 on CIFAR-10 is demonstrated in
Figure 5.
Compared with both iterative pruning and pruning at initialization approaches, our method generates sparse networks at a range of pruning ratios in one-shot fashion that all maintain highly competitive performance.
We further compare with the more recent “oneshot no-retrain” method (Chen et al., 2021), pruning-during-training methods (You et al., 2020; Lin et al., 2020; Hubens et al., 2021), and groupsparsity inducing method (Deleu & Bengio, 2021).
As demonstrated in Table 1, our method outperforms all others with the only exception of DPF (Lin et al., 2020). However, DPF requires full re-training for each pruning ratio, whereas ours only trains once for all possible ratios. Pruning Ratios 50% 70% 80% 90%
[Table]  C ONCLUSION We proposed to consider Stochastic Frank-Wolfe (SFW) for train pruning-friendly networks.
Unlike previous pruning methods based on SGD that often require (iterative) pruning followed by retraining
[Table]
our method can prune a network to any sparsity by training only once
[Table]
which largely eliminates the pruning overhead.
Our core contribution is to train the network weights within in a convex hull spanned by K -sparse vectors, yielding much more smaller weights yet meanwhile “smoother” magnitude distribution, that is more amendable to gradual weight removal towards any sparsity ratio.
Our work is a pilot study to demonstrate the potential of non-gradient optimization methods in training deep networks, especially when additional weight structures like sparsity are desired
[Table]
A I MPLEMENTATION D ETAILS
Our experiments are conducted with ResNet-18 (He et al., 2016) and VGG-16 (Simonyan & Zisserman, 2014) on two benckmark datasets, CIFAR-10 (Krizhevsky et al., 2009) and Tiny-ImageNet (Wu et al., 2017).
We ﬁrst conclude all hyperparameters of SFW-pruning in the following table.
Hyperparameters
[Table]
Moreover, there are two other important technical points in choosing the learning rate α t during SFW training (Pokutta et al., 2020) that we adopt in our experiments.
We clarify them in the following.
• The ﬁrst point is the learning rate changing scheme.
We decrease the learning rate by 10 at epoch 61 and 121 .
Also, we dynamically change the learning rate (Pokutta et al., 2020): the learning rate is multiplied by 0.7 if the 5-epoch average loss is greater than the 10-epoch average loss, and is increased by a factor 1.06 if the opposite holds.
• The second point is to rescale the effective learning rate.
To be speciﬁc, in order to make tuning of the learning rate easier (Pokutta et al., 2020), one can decouple the learning rate from the size of the feasible region.
Here we adopt the following modiﬁed update rule,
[Equation]
Finally, we conclude the hyperparameters of SFWInit in the following table.
[Table]
B K -S PARSE P OLYTOPE C ONSTRAINTS AND DNN W EIGHT D ISTIRBUTIONS
In this section, we give a detailed description of the K -sparse polytope constraint that we use in our algorithm.
In Section B.1, we ﬁrst give two equivalent strict deﬁnitions of a K -sparse polytope and we further explain the motivation why it beneﬁts to equip such a constraint in our training objective.
Then in Section B.2, we analysis the choice of the parameters in the K -sparse polytope constraint via ablation studies on how they inﬂuence the pruning performance and the weight distributions.
The corresponding results further demonstrate our motivation of using such a constraint.
Finally, in
Section B.3, we show the weight distributions induced by SFW-training and SGD-training across different architectures and datasets which demonstrate the universal capability of encouraging sparse weights of SFW-training.
B. M ORE ABOUT K -S PARSE P OLYTOPE C ONSTRAINTS
A K -sparse polytope in R p of radius τ > 0 , p ∈ N , is a bounded subset of R p , which is denoted by
C ( K, τ ) .
There are two equivalent ways to deﬁne a K -sparse polytope C ( K, τ ) .
A K -sparse polytope C ( K, τ ) is the intersection of the L 1 -ball B 1 ( τK ) and the L ∞ -ball B ∞ ( τ ) of R p , i.e., C ( K, τ ) =
{ x ∈ R p : (cid:107) x (cid:107) 1 ≤ τK, (cid:107) x (cid:107) ∞ ≤ τ } .
A K -sparse polytope C ( K, τ ) can also be obtained by spanning all the vectors in R p which have exactly K non-zero coordinates and the absolute value of the non -zero entries are τ , i.e., C ( K, τ ) =
Span [0 , 1] ( { v ∈ R p : (cid:107) v (cid:107)
= K, ( v ) i ∈ { 0 , τ }} ) .
[Figure]
[Figure]
We consider τ ∈ { 5 , 10 , 15 } and K l ∈ { 5% , 10% } . It holds that (Pokutta et al., 2020) for any m ∈ R p the oracle v = LMO C ( K,τ ) ( m ) is given by ( v ) i = (cid:26) − τ · sign(( m ) i ) if ( m ) i is in the largest K coordinates of m , otherwise, ∀ 1 ≤ i ≤ p, (6) which is a vector with exactly K non-zero entries.
According to the SFW update formula θ t +1 =
θ t + α t ( v t − θ t ) =
α t v t + (1 − α t ) θ t , the DNN weight θ t is averaged with the K -sparse vector v t .
Thus each SFW step pushes those less important weights smaller since they are averaged with zero, and those more important weights are enhanced.
The motivation is that this can induce a weight distribution of DNN that has large amounts of small weights.
With such a weight distribution, oneshot pruning can more smoothly remove small values when the pruning ratio increases, yielding competitive test accuracy across the spectrum of pruning ratios, even when retraining is prohibited.
B. A BLATION S TUDY OF K -S PARSE P OLYTOPE C ONSTRAINTS P ARAMETERS
Ablation Study of K -sparse Polytope Constraint Hyperparameters.
We now conduct ablation studies on how the choices of τ and K in the K -sparse polytope constraints can inﬂuence the test performance of SFW-pruning algorithm and the weight distributions of the trained DNNs.
We run the SFW-pruning algorithm (Algorithm 2) using ResNet18 on CIFAR-10 with a range of different hyperparameters.
Speciﬁcally, we consider τ ∈ { 5 , 10 , 15 } and K l ∈ { 5% , 10% } .
Figure 6 shows the ablation study of τ , K choices on the pruning test performance.
Figure 7 shows the ablation study of τ , K choices on the weight distributions of the trained DNNs.
From (a) to (d) in Figure 6 and 7 we can ﬁnd that with smaller K , SFW-training can induce weight distributions with more smaller weights, which at the same time gives better pruning performance across different pruning ratios.
But this is at the cost of lower test performance of the full model .
This observation matches our motivation for using K -sparse polytope constraint.
When K is smaller, the update becomes sparser (recall the LMO solution (6)), which means that more weights are averaged with zero.
Consequently, the ﬁnal weight distributions will gain more smaller weights, which further beneﬁts the pruning procedure since we can more smoothly remove small values when the pruning ratio increases without harming the test performance.
From (e) to (g) in Figure 6 and 7, we can see that with smaller τ , SFW-training can induce weight distributions with more smaller weights.
But these weight distributions with too many smaller weights do not always generate better pruning performance consistently, e.g., see (f) in Figure 6 and 7.
Given the above ablation study we choose τ = 15 and K l = 5% as the hyperparameters in our experiments.
This can give both satisfactory full model performance and pruning performance across different pruning ratios.
We note that by choosing smaller K l , e.g., K l = 1% one can get even better pruning performance, but the full model performance also suffers.
Ablation Study of Initial Learning Rate.
Moreover, we study the inﬂuence of initial learning rate α 0 on the test performance of SFW-pruning algorithm and the weight distributions of the trained
DNNs.
Speciﬁcally we run the SFW-pruning algorithm (Algorithm 2) using VGG16 on CIFARwith a range of initial learning rate α 0 ∈ { 0 . 1 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 , 1 . 0 } .
Figure 8 shows the corresponding pruning performance and weight distributions.
One can see that larger learning rates encourage smaller weights and better performance (on both small and large pruning ratios).
Our conclusion still holds on more complex datasets like TinyImageNet.
As is shown in Figure 9, a weight distribution with more smaller values is still more friendly to pruning.
[Figure]
B
P ERFORMANCE ON M ORE A RCHITECTURES AND D ATASETS
Finally, we show the weight distributions induced by SFW-pruning across four different architecture and dataset combinations, and we compare them to the weight distributions induced by SGD training as well.
Speciﬁcally, we consider { ResNet18, VGG16 } and { CIFAR-10, CIFAR-100 } respectively.
The results are summarized in Figure 10.
[Figure]
[Figure]