Learning Pruning-Friendly Networks via Frank-Wolfe: One-Shot, Any-Sparsity, And No Retraining
Authors: Miao Lu
Xiaolong Luo
Tianlong Chen
Wuyang Chen
Dong Liu
Zhangyang Wang
Venue: ICLR 2022 Spotlight (2022)
Keywords: Pruning
Frank-Wolfe

Abstract: We present a novel framework to train a large deep neural network (DNN) for only $\textit{once}$, which can then be pruned to $\textit{any sparsity ratio}$ to preserve competitive accuracy $\textit{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 $\textit{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 first learning-based initialization scheme specifically for boosting SFW-based DNN training. Experiments on CIFAR-10 and Tiny-ImageNet datasets demonstrate that our new framework named $\textbf{SFW-pruning}$ 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. All codes will be released to the public. 

Published as a conference paper at ICLR 2022
LEARNING PRUNING-FRIENDLY NETWORKS VIA FRANK-WOLFE: ONE-SHOT, ANY-SPARSITY, AND
No RETRAINING
Miao Lu Xiaolong Luo"", Tianlong Chen', Wuyang Chen', Dong Liu', Zhangyang Wang' University of Science and Technology of China, University of Texas at Austin {I [Lumiac, Lx121338maaaa www.uusc......... edu en, dongeliu@ustc. edu. en {t stianleng.o chen, wuyang, chen, at www.saaageettteeeecco
ABSTRACT
We present novel framework to train large deep neural network (DNN) for only once, which can then be pruned to any sparsity ratio to preserve competitive accuracy without any re-Iraining. 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 retraining hyperparameters. Our core idea is to re-cast the DNN training as an explicit pruning-aware process: that is formulated with an auxiliary -ssaare polytope constraint, to encourage net- work weights to lie in a convex hull spanned by K K-sparse vectors, potentially re- sulting 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 ex- isting DNN initializations were derived to enhance SGD training (e.g., avoid gra- dient explosion or collapse). but was unaligned with the challenges of training with SFW. We hence also present the first learning-based initialization scheme specifically for boosting SFW-based DNN training. Experiments on CIFAR-10 and Tiny-ImageNet datasets demonstrate that our new framework named SFW- pruning consistently achieves the state-of-the-ar performance on various bench- mark DNNs over 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 https: //g thubb com/ VI TA TTAA-AAor/S---A Once- for-All-Prunn
INTRODUCTION
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-platfion et al., 2019), mobile applications on App stores have to support a diverse range of hardware devices, from high-end flagships to low-end ones, which need different model capacities for different sce- narios. As a remedy, neural network compression (Han et al., 2015) tries to slim various DNNs to improve both the computation and memory efficiency. One important way of network compression is network pruning, which prunes individual weights (unstructured pruning, Han et al. (2015)) 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 fine-tuning of the com- pressed 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.
applications. For example (Cai
To tackle such an efficiency problem, in this paper, aim answer the question of whether can design an efficient pruning method that does not need retrain the neural network. Specifically, 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 natural choice since retraining is not allowed. Intuitively, given sparsity

---

Published as a conference paper at ICLR 2022
ratio, one can leave the most important individual weights (weight-magnitudl unstructured pruning (Han et al., 2015)) untouched and the less important ones pruned according to their relative "abso- lute 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 find those important weights via standard gradient- optimization methods during train- ing, e.g., stochastic gradient descent (SGD), whose performance will degrade much if retraining is prohibited, especially at large pruning ratios (See Figure ((a)). The reason is that such optimiza- tion methods fail to take the desired sparse structures in the pruned model into consideration. The process of finding 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 formulated with an auxiliary K-sparse polytope constraint to encourage the DNN weights to lie in the convex hull spanned by K 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-valus (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 fine-tuning.
Moreover, we propose new initialization scheme for SFW-based training. We note that SFW al- gorithm 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 fit 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 mo- tivations from Zhu et al. (2021). It learns the layer-wise initialization scaling factors that lead to the largest loss reduction in the first 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-awar process formulated under an auxil- iary K-sparse polytope constraint, based on which we propose a new SFW-pruning frame- work 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- -easeing--baeed 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 applicable across different architectures and datasets, achieving the state-of-the-art performance sistently over a wide range of pruning ratios without retraining.
perfformcc con-
2 RELATED WORK
Neural Network Pruning and Efficient 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 difficult in actual application and deployment. Various methods, including knowledge distillation (Hinton et al., 2015; Romero et al., 2014), low- rank factorization (Denton et al., 2014; Yu et al., 2017), quantization (Courbariaux Rastegari et al., 2016; Wu et al., 2018), and pruning (Han et al., 201 5; Li et al., 2016; I. Liu et al., 201 8), have been proposed to deploy large models in resource-constraingg
al., 2016;
devices efficiently. Among these methods, pruning has become 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 sparse network. Although IMP works well, its high computation cost motivates more efficient methods (Lee et al., 2018; Wang et al., 2020a; Tanaka et al., 2020). Another major branch is structured pruning. It
2

---

Published as a conference paper at ICLR 2022
considers parameters in groups, removing entire neurons, filters, 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 efficient and scalable neural network architecture such as MobileNet (Howard et al., 2017) and Shuf- fleNets (Ma 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 flexibly obtain a large number of sub-networks meet different resource constraints from a pre-trained "super network", but their pre- pre-training cost is gigantic and re- is needed for each sub-network for restoring the optimal performance. our work, we view efficient deployment from 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 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) a classical non- graadretdtetteeeddddd 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) al- gorithm. Hazan & Luo (2016) perform theoretical analysis of the standard SFW algorithm with convergence rate (1/tt) with θ (1²) growing batch-sizes. Several works further extend SFW-based algorithms in the deep learning field 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 be a crucial part of SFW in training neural networks (Cutkosky & Orabona, 2019; Mokhtari 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 pre- viously overlooked problem of SFW initialization, and present meta learning-
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 longer hold for more complex architectures, motivating newer
He et al. (2015)'s assumptions methods (Dauphin & Schoenholz, 2019; Zhang et al., 2019; Bachlechner et al., 2020). However, most of these tied to enhancing SGD-based training. Since the K -sspare SFW algorithm is non-gradient based, new initialization methods need to be customized.
3 STOCHASTIC FRANKE-WOLd PRUNING FRAMEWORK
In this section, we formulate the DNN training process as an explicit pruning-aware process with an auxiliary K K-sparse polytope constraint, and we solve the corresponding constrained optimization problem via a stochastic Franke- Wolfe (SFW) algorithm. Specifically, we formulate the pruning- aware process and review the basics of SFW algorithm in Section 3.1. In Section 3.2. we present our proposed SFW-pruming framework based on the pruning-aware process and one-shot pruning, where we consider weight-magnitude unstructured pruning.
3.1 PRUNING-AWARE DNN TRAINING AND STOCHASTIC FRANK- WOLFE ALGORITHM
Our core idea is to re-cast the DNN training process as an explicit pruning-aware process, which is achieved via auxiliary K K-sparse polytope constraint. Formally, given dataset DD {(x;, v:) 1 where (x,, y:) E X' x y and a loss function f y x y + R_, e.g., cross-entropp loss, we aim to train deep neural network ((:-)) .): X' y that minimizes the following pruning-awe -aware objective:
min [(((0x)) = min L( OEC n
DEC
To make the training process awaning-aware, we restrict the feasible parameter 0 in a convex region C, which potentially results in more sparse weight matrices and will be more pruning-friendl for
3

---

Published as a conference paper at ICLR 2022
only small percentage of the weights of large magnitudes. Specifically, we choose C as K- sparse polytope, and we solve (1) efficiently via a stochastic Frank Wolle (SFW) algorithm. We review the K -sparse polytope and the basics of SFW below and further explain our motivations.
Stochastic Trank-Wolfe and K K-Sparse Polytope Constraints. SFW is simple projection-free first-order algorithm for solving convex constraint optimization problems (Reddi et al., 2016; Hazan & Luo, 2016; Yurtsever 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 Oi. At each iteration t. SFW first calls a linear minimization oracle v, = LMOc( L (O)) = arg milleec( V (O,), v), where VoL(O,) estimates the full gradient, e.g., gradient a minibatch. Given Ut. SFW updates O, along the direction of v. by convex combination 0,+1 = 0₂ + (n - 0,) = a + (1 - at)0, with the learning rate a E [0. 1]. This keeps 0, 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 -sspase polytope in R of radius 7 > O, P E N, denoted by C(K, r). is the intersection of the - B, (rK) and the -ball (T) of RR. One can obtain C(K, r) by spanning all the vectors in RP, which have exactly K non-zero coordinates and the absolute value of the non-zero entries are T, It holds that (Pokutta et al., 2020) for any m IR" the oracle LMOc(K.r)(M) is given by
{o- T. sign (mm))) if (m)) in the largest K coordinates of m. (v), 10
A1 < i P:
(2)
otherwise,
which is a vector with exactly K non-zero entries. When applying the K K-sparse polytope constraint in deep neural network training objective (1), we add the constraint In other words, if we write out 0 as = (W1. ... WL, by) for layer-wise weight parameter Wi and bias parameter b,, then each of the parameters is paired with K K-sparse polytope constraint C(K,,7). Here K2 may vary between the different layers. For notation simplicity, in the sequel, we denote such a layer-wise wise constraint as 0 E C((KI) without making any confusion.
each layer of the network.
Why K- K-Sparse Polytope Constraints? In each step of our optimization, the linear minimization oracle of a K K-sparee polytope constraint returns update vector with exactly K non-zero coordi- nates, which then reighted-averagee with the current 0, according to the SFW algorithm (aceu- mulated from the "ensemble" of K non-zero coordinate vectors, from all past update steps). In other words, training with K K-sparse constraints amounts to a "vottng" process, by all step updates, on which K elements "should be non-zero", and the resulting weights could be viewed as the soft vot- ing 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 weight distribution, when the pruning ratio increases, there will be no sudden "jump" to removing small values to im- pacting 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- -SSparse Polytope Constraints.
3.2 STOCHASTIC FRANK- WOLFE ONE-SHOT PRUNING WITHOUT RETRAINING
Previously, we have motivated the usage of K K-sparse polytope constraint in the pruning-aware DNN training objective (1) and using SFW algorithm in the training process. this subsection, we first propose our SFW W-pruning framework for deep neural network one-shot pruning without retraining based on SFW neural network training with K K-sparse polytope constraints. After, we provide some fine-grained analysis of the hyperparameters involved in the newly proposed pruning framework.
Stochastic Frank-Wole Pruning Framework. We now present our one-shot pruning framework, where a DNN is trained using the pruning-awarr objective (1) with K 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 pwaning-aware, which aims to find proper DNN weight 0, that can minimize the training
4

---

Published as a conference paper at ICLR 2022
Algorithm 1: Stochastic Frank-Wolfe with Momentum for Deep Neural Network Training
1: Input: Dataset D = {xx, y)) deep neural network f(0:-). convex constraint region C, initialization point 0g C, linear minimization oracle LMOc, number of steps T, initial learning rate og E [0, 1). learning rate scheme 1 scheme, momentum E [0, 1].
2: Output: final point r = (D, LMOc, 0, T, agg 1 Ar_scheme,p) p).
3: for 1..=11.... T do
4: Update learning rate a₂ + 1r.scheme(t) and gradient estimator VoL(O,-1).
5: Update momentum vector m, (1 p)m 1oomm-+pVoLLO--) (V.(())-1).
6: Solve linear minimization oracle 22 t LM)c(
7: Update neural network weights 0, t- P.-1 az 0.).
8: end for
Algorithm 2: Stochastic Frank-Wolfe Pruning Framework (SFW-Pruning)
1: Input: Dataset D, deep neural network f(On;)) -), diameter т, sparsity {Ki)= initial weight
0o E C({K))= r), linear minimization oracle LM training epoch T, initial
learning rate ao E j0, 1J. learning rate scheme Ir.scheme, momentum E [0, 1]. desired pruning ratio s, Init
(bool), and Pruning.Proceduss
(weight-magnitude unstructured pruning).
2: Output: a pruned neural network f(0):)) with pruning sparsity ratio
3: # Sparse training phase using SFW with K-sparse constraints
4: if Initial iiaaiinnSchheme True then
5: 0o SF 0g).
6: end if
7: Train the deep neural network f on D via SFW, i.c., set
0r < -FWWW(D, f. ff (((K))), IMOc(1))) 0, T, 7, 1r..cheeeeee
8: # One-shot magnitude pruning phase
9: Pruning the deep neural network (0)::) via Pruning.Procedure and get (07;
0.30
10-
Pruning Tast Accurscy
80
0.25
10a
SGD
60
SGD-magniude
0.20
VEG16
10"
Polytopc
Nature
ResNet18
40
10a
0.10
10
20-
107
0.05
Weight Absolute Values
intlal Rate
(a) Weight distributions
(b) Test accuracy
(c) Nature sparsity
Figure 1: (a) shows the weight magnitude distributions of DNNs trained by SGD and SFW respec- tively, averaged over architectures (ResNet18, VGG16) and 2 datasets (CIFAR-10, 10, CIFAR-100), See Appendix B.3 for weight distributions for each combination. (b) shows the test accuracy under three different settings: SFW-pruning framework; SGD training and one-shot magnitude pruning; and SGD training and random pruning. The r-axis the desired pruning ratio. (c) shows the na- ture sparsity of two architectures constraints on CIFAR-11 dataset with ou {0.1, 0.2, 0.4, 0.6, 1.0).
{ [ResNet18, VGG16) trained by SFW with K-sparse polytope
loss while being pruning-friendly. The following pruning procedure (Line 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 1(b), we demonstrate the effectiveness of our framework using ResNet18 (He et al., 2016) on the CIFAR-10 (Krizhevsky al., 2009) dataset for unstructured pruning.
experiments in Section 5. In Figure
5

---

Published as a conference paper at ICLR 2022
Despite the simplicity, the high-level idea of Algorithm 2 is intuitive: the discovery of sparse struc- tures and pruning-award DNN weights should be integrated with the optimization process, especially when retraining is prohibited. We further demonstrate the effectiveness of Algorithm 2 via exten- sive 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-awarr 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 trade-off between the test performance and the one-shot pruning ratio of the DNN. But given the same pruning ratio, our framework can find more pruning-friend weights than random pruning or other optimization methods without sparsity awareness. We refer to Appendix B.1 for detailed ablation study of the hyperparameters in Algorithm 2.
K K-sparse constraint parameters { and r: {KI)_ controls the sparsity of the linear mini- mization oracle solution, and controls the magnitude, which together influence the distribution of the learned parameters. Experimental results show that relatively smaller choices of these hyperpa- rameters can keep the test performance undamaged under higher one-shot pruning ratios, at the cost of relatively lower full model test performance. Empirically, we find K = 5% and T = 15 to be our default good choices that transfer well across DNNs and datasets.
Learning rate or: Recall that at step I, we update O, by 0.+1 = 0, (r - 0.) 02 + (1 11--a (1-0,)) which is the convex combination of v,, the linear minimization orcale solution, and θ,, the current parameter. Note that v, is K K-sparse, which means that a larger learning rate (i.c., ar4, has larger weight in the convex combination) may result in sparser parameter 0 This observation indicates that one can (dynamically) control the nature sparsity² of the neural network by controlling the seale of the learning rate a₁- We refer to Figure (cc) for demonstration of this phenomenon.
4 INITIALIZATION SCHEME FOR STOCHASTIC FRANK-WOLFE ALGORITHM
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 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. (2021). The goal is to learn the best initialization 0, in the sense that it can allow the maximal loss reduction in the first SFW step. which we hope can further bring more competitive and consistent pruning test performance in our SFW-pruning paradigm. Specifically, the proposed method initializes each layer of the neural network with a uniform distribution, which is then multiplied by layer-wise scaling parameter learned by the algorithm, resulting maximal one-SFW-step loss reduction.
Methodology. We first initialize the parameter On of weight matrices (Wi)f= and bias vectors (br) 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). 1 L denotes the layer index. For each layer I, we pair W; and b, with learnable non-negative scalars a and B, that control the scaling of the layer at initialization, and we use ,3 denote the vector of scale factors (E,, S,, ... EL, (L), and let 0 be the tuple (& W1, Wi,, WL, (bb)) of rescaled parameters. Also, we extend the loss definition as
1 S
L(S:0)
Σ {(f(:x)).).
(3)
(x.))SS
which is the average loss of the model with 0 on a minibatch of samples S. Correspndingly, we de- note one SFW optimization step obtained on minibatch S as SFW(S, (,C((K)) 7), 0p. 1, ao p).
Here K2 5% refers 5% of the weights in layer I. Also, let be rescaled by the expected initializa- tion norm. See Section for detailed implementation setup. We refer to the nature sparsity of neural network as the largest pruning ratio that keeps the test perfor- mance undamaged.

---

Published as a conference paper at ICLR 2022
Algorithm 3: Stochastic Frank- Initialization Scheme (SPW-Init)
Input: Dataset D, SFW parameters (See Algorithm 1), K sparse polytope constraint
C( C((Ki)= learning rate K, total iterations I', lower bound of weight and bias scales e, E.
2: Output: A new initialization of SFWInit(((06:)))
3: Set B1 f-
4: for = , do
5: Draw random samples St from training set ID.
6: Draw /2 samples to replace [S,] /2 samples in S, and let this new minibatch be S,
7: Set L(+1 f- 1 E(x.)es, (afx))eS,((((SFNO))) f.
08" a, a.. p):x), v).
[5.1
8: Set Be+1 = Be -VBI++1
9: Clamp 3.41 using the lower bound c and E.
10: Project B.+1 back to C(T, K) as B.+1 = Cl,,, (3++1)-
end for
Specifically, we optimize the scaling parameter B to obtain good initialization 0 tailored for stochastic Frank-Wolfe algorithm, i.e., SFW. Consider the first update step, starting from the rescaled initialization, ie., 01 = SFW(S, f. ((((()))= 7), 0o. 1, ao p). We choose & so that the loss on 02 is as low possible. Formally, optimize the scaling B through the following objective:
min B s.t. EC(K.T)
L S.W(S 0..,
(4)
where and are two different random minibatches. For his paper, we only consider K-sparse polytope constraints, which can be easily extended to other constraints like (?-norm balls. To make sure that the rescaled parameter 0 still feasible initialization, we make the constraint that 0 C(K,1) 7). 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. (2021):
We do not apply constraints on the norm of the gradient (0%) which the case in Zhu et al. (2021). This is because for K K-sparse polytope constraints the linear minimization oracle (() (V, (+%)) only depends on the direction rather than norm.
Since we need to ensure that the initialization is feasible, we project the obtained rescaled initialization 0º. back to the K K-sparse polytope constraints each training iteration. This is done via simple scaling of the current parameter to the K K-sparse polytope C(T, K).
Besides, there are two other key points in the SFW-Init algorithm, motivated by Zhu et al. (2021). The first point is that we use different but overlapping minibatches and Š to calculate the current loss and the first SFW step respectively. The difference between overfitting, while the overlapping of the minibatches for better train performance. See Zhu et al. (2021) for study of 5 and The second point is that we detach the SFW algorithm step from the gradient V Vala+1 (line 7) in order facilitate the differentiation of the SFW algorithm (which involves sorting the coordinates of the gradient vector) and potential second-orde derivatives. The corresponding SFW algorithm modification marked in the grey color.
mini-batches is prevent
EXPERIMENTS
5.1 SETTINGS
We summarize the key experiment setups, with hyperparameters of the implementation presented in Appendix A 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-10 (Krizhevsky et al., 2009) and ImageetteeNet (Wu et al., 2017). Specifically, on all these setups, demonstrate our SFW-pruning framework (Algorithm 2) by first solving the K 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 dynamic-changing learning

---

Published as a conference paper at ICLR 2022
stot www.tactuurd ACCuRDO Woddet CINAR0
Sie Stat LNNDOA Lare Accrddy V0AC UIdee
22 Stic MRSTTUCtar ACCRCCY AOSRUCT Tay 350
1313 S104 fadd ACCKOOS Wiith Maago
Puatigg Test Accurey
4
ScP-rage Lade
SPW-anning
EPW-nuiing
40
Pureg Heee
Phinng Hale
Phiiing lialo
Punggibd
(a) CIFAR-I0 R (b) CIFAR-IO VGG-16
(c) Tiny-Image ResNet--8 (d) Tiny-Image VGG-16
Figure 2: Test performance of unstructured pruning networks with different sparsity ratios without retraining. We compare three settings: SFW- Wwrrninnng. SGD training with one-shot weight-mane magnitude pruning. SGD-training with one-shot random pruning. All initial learning rate fixed to ag 1.0.
rate schedule used by Pokutta et al. (2020). We evaluate the neural networks by testing accuracy after pruning under different sparsity ratios and without retraining.
5.2 ONE-SHOT WEIGHT-MAGNITUDE UNSTRUCTURED PRUNING
We first evaluate our SFW-pruning framework in unstructured weight-magnitude pruning paradigm (Han et al., 2015; Frankle & Carbin, 2018) across total of four combinations of datasets and archi- tectures, i.e., CIFAR-10 with (ResNet-18, VGG- 16) and Tiny-ImageNet with (ResNet-18, VGG-
Comparison Retraining
80
Test Accuracy
60
SGD-magnitude with retraining
SPW-pruning
16). Note that for the fairness of comparison, we do not use our learned initialization (SFW-Init) here.
SGD-magniude without retraining
40
Specifically, after training the DNN with K K-sparse polytope constraints via SFW, we prune it to var- ious sparsity ratios by leaving the largest weights unchanged and the smaller ones pruned. We con- sider neural network training without constraints via SGD and random one-shot pruning for compar- ison. Figure 2 shows the achieved test performance of the pruned neural networks with different sparsi- ties. Several observations can be drawn from Figure 2: (i) without retraining, SFW- SFW-pruming significantly outperforms magnitude- and random prun- ing 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.
20
0.64
0.4
0.38
0.17
0.07
0.03
0.01
Pruning Patio
Figure Our one-shot SFW-pruning (no retrain- ing) can even achieve competitive performance against "SGDD-puning with extra retraining costs" (ResNet-18,, CIFAR-10).
SFW-pruning 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 fur- ther validates that SFW 3FW-pruning is sparsity-aware. Network weights of less importance have already converged to small magnitudes by SFW. Unlike the performance gap induced by surd-thresholding by SGD, important weights found and trained by SFW SFW-pruning are ready for pruning and inference.
5.3 ADDING THE NEW SFW INITIALIZATION
We then evaluate the effectiveness of the initialization scheme (SFW-Init) designed specifically for SFW algorithm in Section 4. We repeat the previous experiments for evaluating SFW-pruning frame- work in unstructured paradigm, comparing the test performance with and without the initial- ization scheme. The corresponding results are shown in Figure 4. We can see that the performance after adopting customized SFW initialization is consistently better than Kaiming initialization.
8

---

Published as a conference paper at ICLR 2022
100
intaibaaton Comper Rerretto CPAR30
130 195a malee Compabees VG236 pharto
64
Fiba ostte Compariics feawdd Trp.
19ttl Comparies VoGld Try-erapehrr
80
S/W peiring with Sier22
Staiprasing SPitin Srippratig
2g W SFAt:: graagra
Paiing
Pruning
PuniiRote
(a) CIFAR-IO ResNet-18 (b) CIFAR-I0 VGG-16 (c) Tiny-Image ResNet-18 (d) Tiny-Image VGG-16
Figure 4: We study SFW- LFW-innning with and without SFWInit. Test accuracies of the pruned DNNs with different sparsity ratios are obtained without retraining. All have initial learning rate ag = 1.0.
5.4 COMPARISON TO STATE- ARE-ARR METHODS
holistic solution: "SFW-pruning + SFW-Init", with other tfth--tfth (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
We now compare
100
80
Test Accuracy
magnitude
60
SFW-pruning
SGD-Global Weight
40
SGD-Layer Welght
costs extremely heavy overhead and is to our disadvantage. The full comparison of VGG-16 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.
SGD Gol-Golbal Gradient
SGD-Layer Gradlent
20
GraSP
SynFlow
1.0
0.1
0.06
Pruning Ratio
Figure 5: Comparison with SOTA unstructured
pruning methods on VGG-16 and CIFAR-10. Solid lines are reported by Blalock et (2020).
We further compare with the more recent "one- shot no-retrain" method (Chen et al., 2021).
muning-during-trainiin methods (You al., 2020; Lin al., 2020; Hubens et al., 2021), and group- sparsity inducing method (Deleu & Bengio, 2021). As demonstrated in Table 1, our method out- performs 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%
95%
SFW-Pruning SFW-Initt (ours)
93.10
93.10
93.10
93.10
92.00
One-Cycle Pruning (Hubens et al., 2021)
90.87
90.72
90.67
Early Bird (You et al., 2020)
93.21
92.80
OTO (Chen et al., 2021)
90.35
90.35
90.35
90.35
90.35
DPF (Lin et al., 2020)
93.87
Group MDP (Deleu & Bengio, 2021)
89.38
Table 1: Comparisons to more state-of-the-art methods on VGG-16 and CIFAR-10.
6 CONCLUSION
We proposed to consider Stochastic Frank- Wolfe (SFW) for train pruning-friene networks. Unlike previous pruning methods based on SGD that often require (iterative) pruning followed by retraining, our method can prune a network to any sparsity by training only once, which largely eliminates the pruning overhead. Our core contribution is to train the network weights within in convex hull spanned by K 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.
9

---

Published as a conference paper at ICLR 2022
REFERENCES
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. arXiv preprint 0488: 2020.
Davis Blalock, Jose Javier Gonzalez Ortiz, Jonathan Frankle, and John Guttag. What is the state of neural network pruning? arXiv preprint arXiv:2003.03033 2020.
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once-for-all: Train one network and specialize it for efficient deployment. arXiv preprint arXiv:1908.09791, 2019.
Lin Chen, Christopher Harshaw, Hamed Hassani, and Amin Karbasi. Projection-free online opti- mization with stochastic gradient: From convexity to submodularity. In International Conference on Machine Learning. pp. 814-823, PMLR, 2018
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Michael Carbin, and Zhangyang Wang. The lottery tickets hypothesis for supervised and self-supervised pre-training in computer vision models. arXiv preprini arXiv:2012.06908 2020a.
Tianlong Chen, Jonathan Frankle, Shiyu Chang, Sijia Liu, Yang Zhang, Zhangyang Wang, and Michael Carbin. The lottery ticket hypothesis for pre-trained bert networks. arXiv, abs/2007. 12223, 2020b.
Tianyi Chen, Bo Ji, Tianyu Ding, Biyi Fang, Guanyi Wang. Zhihui Zhu, Luming Liang, Yixin Shi, Sheng Yi, and Xiao Tu. Only train once: A one-shot neural network training and pruning framework. Advances in Neural Information Processing Systems, 34, 2021.
Cyrille w Combettes, Christoph Spiegel, and Sebastian Pokutta. Projection adaptive gradients for large-scale optimization. arXiv preprint 2rXiv:2009.144 14114, 2020.
Matthieu Courbariaux, Itay Hubara, Daniel Soudry, Ran Yaniv, and Yoshua Bengio. Binarized neural networks: Training deep neural networks with weights and activations constrained to+ or-1. arXiv preprint 02220.1602.02880, 2016.
Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. arXiv preprint
2019.
Yann Dauphin and Samuel S Schoenholz. Metainit: Initializing learning by learning to initialize. 2019.
Tristan Deleu and Yoshua Bengio. Structured sparsity inducing adaptive optimizers for deep learn- ing. arXiv preprint arXiv:2102.03866 2021.
Emily L Denton, Wojciech Zaremba, Joan Bruna, Yann LeCun, and Rob Fergus. Exploiting linear structure within convolutional networks for efficient evaluation. In Advances in neural informa- tion processing systems. pp. 1269-1277, 2014.
Marguerite Frank, Philip Wolfe, et al. An algorithm for quadratic programming. Naval research logistics quarterly, 3(1-2):95-110, 1956.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: 1803.08803.3335 2018.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249-256. JMLR Workshop and Conference Proceedings, 2010.
Yi Guo, Huan Yuan, Jianchao Tan, Zhangyang Wang, Sen Yang, and Ji Liu. Gdp: Stabilized neural network pruning via gates with differentiable polarization. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pp. 5239-52500 2021.
Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv: /5 15.000149 2015.
10

---

Published as a conference paper at ICLR 2022
Elad Hazan and Haipeng Luo. Variance-reduccd and projection-free stochastic optimization. In International Conference on Machine Learning, pp. 1263-1271. PMLR, 2016
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural net- works. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1389-1397, 2017.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge neural network. arXiv preprint arXiv:/503.025.: 2015.
Andrew G Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: /704. 2017.
Hu, Peng, Y Tai, C Tang, and N Trimming. A data-driven neuron pruning approach towards efficient deep architectures. arXiv preprint arXivv1607.032550 46, 2016.
Nathan Hubens, Matei Mancas, Bernard Gosselin, Marius Preda, and Titus Zaharia. One-cycle pruning: Pruning convnets under tight training budget. arXiv preprint urXiv:2 1e0.2107700000. 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shov shot network pruning based on connection sensitivity. arXiv preprint arXiv:/810.02344 2018.
Namhoon Lee, Thalaiyasingam Ajanthan, Stephen Gould, and Philip HS Torr A signal propagation perspective for pruming neural networks at initialization. arXiv preprint arXiv:1906. 000000667777. 2019.
Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. arXiv preprint 16Xiv:1608.087/11 2016.
Tao Lin, Sebastian U Stich, Luis Barba, Daniil Dmitriev, and Martin Jaggi. Dynamic model pruning with feedback. arXiv preprint (xXi::2006607253, 2020.
Zhuang Liu, Mingjic Sun, Tinghui Zhou, Gao Huang, and Trevor Darrell. Rethinking the value of network pruning. arXiv preprint arXiv: arXiv:1810.05270, 2018.
Jian-Hao Luo and Jianxin Wu. Neural network pruning with residual-connection and limited-data. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 1468-1467, 2020. 2020.
Ningning Ma, Xiangyu Zhang, Hai-Tao Zheng, and Jian Sun. Shufflenet v2: Practical guidelines for efficient enn architecture design. In Proceedings of the European conference on computer vision (ECCV), pp. 116-131, 2018.
Aryan Mokhtari, Hamed Hassani, and Amin Karbasi. Stochastic conditional gradient methods: From convex minimization to submodular maximization. Journal of machine learning research, 2020.
Sebastian Pokutta, Christoph Spiegel, and Max Zimmer. Deep neural network training with frank- wolfe. arXiv preprint arXivv20100.02233, 2020.
Mohammad Rastegari, Vicente Ordonez, Joseph Redmon, and Ali Farhadi. Xnor-net: Imagenet classification using binary convolutional neural networks. In European conference on computer vision, pp. 525-542. Springer, 2016.
11

---

Published as a conference paper at ICLR 2022
Sashank Reddi, Suvrit Sra, Barnabás Póczos, and Alex Smola. Stochastic frank wolfe methods for nonconvex optimization. In 2016 54th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 1244-1251. IEEE, 2016.
Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. arXiv preprint arXiv:/412.6550, 2014.
Zebang Shen, Cong Fang, Peilin Zhao, Junzhou Huang, and Hui Qian. Complexities in projection- free stochastic non-convex minimization. In The 22nd International Conference on Artificial Intelligence and Statisties, pp. 2868-2876. PMLR, 2019.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint ArXiv 1556, 2014.
Hidenori Tanaka, Daniel Kunin, Daniel LK Yamins, and Surya Ganguli. Pruning neural networks without any data by iteratively conserving synaptic flow. arXiv preprint wwXiv:2006.05467. 2020.
Chaoqi Wang. Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. arXiv preprint arXiv:2002.07376 2020a.
Haotao Wang, Tianlong Chen, Shupeng Gui. TingKuei Hu, Ji Liu, and Zhangyang Wang. Once-for- all adversarial training: In-situ tradeoff between robustness and accuracy for free. Advances in Neural Information Processing Systems, 33:7449 7461, 2020b.
Jiayu Wu, Qixiang Zhang, and Guoxi Xu. Tiny imagenet challenge. Technical Report, 2017.
Junru Wu, Yue Wang, Zhenyu Wu, Zhangyang Wang, Ashok Veeraraghavan, and Yingyan Lin. Deep k-means: Re-training and parameter sharing with harder cluster assignments for compressing deep convolutions. In International Conference on Machine Learning, pp. 2018.
PMLR,
Jiahao Xie, Zebang Shen, Chao Zhang, Hui Qian, and Boyu Wang. Stochastic recursive gradient- based methods for projection-free online learning. 2019.
Haoran You, Chaojian Li, Pengfei Xu, Yonggan Fu, Yue Wang, Xiaohan Chen, Richard G Baraniuk, Zhangyang Wang, and Yingyan Lin. Drawing carly-bird tickets: Toward more efficient training of deep networks. In International Conference on Learning Representations, 2020.
Xiyu Yu, Tongliang Liu, Xinchao Wang, and Dacheng Tao. On compressing deep models by low rank and sparse decomposition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 7370-7379, 2017.
Alp Yurtsever, Suvrit Sra, and Volkan Cevher. Conditional gradient methods via stochastic path- integrated differential estimator. In International Conference on Machine Learning, pp. 7282- 7291. PMLR. 2019.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv: 1x0xv 009311, 2019.
Mingrui Zhang, Zebang Shen, Aryan Mokhtari, Hamed Hassani, and Amin Karbasi. One sample stochastic frank-wolfe. In International Conference on Artificial Intelligence and Statistics, pp. 4012-4023, PMLR, 2020.
Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun. Shufflenet: An extremely efficient convolutional neural network for mobile devices. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6848-6856, 2018
Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong, W Ronny Huang, and Tom Goldstein. Gra- dinit: Learning to initialize neural networks for stable and efficient training. arXiv preprint 0020092002.08098. 2021.
12

---

Published as a conference paper at ICLR 2022
A IMPLEMENTATION DETAILS
Our experiments are conducted with ResNet-18 (He et al., 2016) and VGG-16 (Simonyan & Zis- serman, 2014) two benckmark datasets, CIFAR-10 (Krizhevsky et al., 2009) and Tiny-ImageNet (Wu et al., 2017). We first conclude all hyperparameters of SFrniruning in the following table.
Hyperparameters
Value
Hyperparameters
Value
Hyperparameters
Value
Initial learning rate og
1.0
Training batchsize
128
Test batchsize
100
Radius 7
15
K-frac {Ki]im
5%
Training epoch T
180
Momentum P
0.9
Table 2: Values of hyperparameters in SFW-pruning (Algorithm 2).
Moreover, there are two other important technical points in choosing the learning rate or during SFW training (Pokutta et al., 2020) that we adopt in our experiments. We clarify them in the following.
The first point is the learning rate changing scheme. We decrease the learning rate by 10 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 factor 1.06 if the opposite holds.
The second point 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 modified update rule,
rescale the effective learning rate. To be specific, in order to make
8,41 0g minfar (L))) / 8.112. 1} (uz 0.)
Finally, we conclude the hyperparameters of SFWInit in the following table.
Hyperparameters
Value
Hyperparameters
Value
Hyperparameters
Value
Learning rate x
0.001
Training iterations T
390
Minimal scaling e,
0.01
Table 3: Values of hyperparameters in SFWIint (Algorithm 3).
B K-SPARSE POLYTOPE CONSTRAINTS AND DNN WEIGHT DISTIRBUTIONS
In this section, we give detailed description of the K-sparse polytope constraint that we use in our algorithm. In Section B.I, we first give two equivalent strict definitions a K K-sparse polytope and further explain the motivation why it benefits to equip such a constraint in our training objective. Then in Section B.2, we analysis the choice of the parameters in the K K-sparse polytope constraint via ablation studies how they influence the pruning performance and the weight distributions. The corresponding results further demonstrate our motivation of using such a constraint. Finally, in Section B.3, 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.1 MORE ABOUT K-SPARSE POLYTOPE CONSTRAINTS
A K -paases polytope in R2 of radius > 0, p E N. is bounded subset of IR, which is denoted by C(K, r). There are two equivalent ways to define K-sparse polytope C( (K,
1. A K K-sparse polytope C(K, r) is the intersection of the L' -ball (T(K) and the ball Be (T) of RR, i.e., C(K, )) = {ar E TRP || S 7 TK, ||0||| S T)).
2. A K K-sparse polytope C(K, 7) can also be obtained by spanning all the vectors in R which have exactly non-zen zoro coordinates and the absolute value of the non-zero entries are т, i.e., C(K,T) Span ( RP l|ol|o K, (w); (0,+), )))-
13

---

Published as a conference paper at ICLR 2022
Capare Natian Al Naa Dan3
Sapare puldapo No Win Daald
3nd 2000 E-0 Desede
24 puttgggeo kod Faania
40-
5
(b) T= 10
(c) T==11 15
50% poatigg 5
5rikgreeigg 6-0L 1
.0A www.ood
5/00
(d) K2 = 5%
(e) K₁ 10%
(l) Conclusion
Figure 6: Ablation study of the influence T and K in K-sparse polytope constraint on the pruning performance (ResNet18 on CIFAR-IO). We consider 7 E 10, 15) and K₁ E 15%, 10%).
Speerse hobttree Ab Study with 1agi-
X Spern Polyttas 4D atte Study with 10
X Souns Polytope Ablater Stide tt tap-15
10'
105
100
12
207
107
Wosth Aoodat
AlsoMie VARKO
(a) T
(b) T 10
(c) T 15
Soarre Tolytree Sh
CSperse Ph atoen
Scaape Noiteps Hatter
10P
10a
107
107
107
Westl 1005
Wayge Aiwirio Valees
Absukie Vokee
(d) Ki 5%
(e) Ki 10%
(I) Conclusion
Figure 7: Ablation study of the influence of 7 and K in K K-sparse polytope constraint on the weight distributions (ResNet18 on CIF. We consider 7 E (5, 10, 15) and K₁ E (5%, 10%).
It holds that (Pokutta et al., 2020) for any m E IR the oracle v = LMOc(K,r) (m) is given by
(v); { {- T. (ig)((m):) if (m); is in the largest K coordinates of m.
visisp. P.
(6)
otherwise,
which is a vector with exactly K non-zero entries. According to the SFW update formula 0.+1
0, + (v - 0.) = a,,,, + (1-a, 0, the DNN weight O, is averaged with the K-sparse vector V,. 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 weight distribution, one- shot 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.
14

---

Published as a conference paper at ICLR 2022
B.2 ABLATION STUDY OF K-SPARSE POLYTOPE CONSTRAINTS PARAMETERS
Ablation Study of K K-sparse Polytope Constraint Hyperparameters. We now conduct ablation studies on how the choices of T and K in the K-sparse polytope constraints can influence 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. Specifically, we consider T E {5, 10, 15) and K₁ E (5%, 10%]. Figure 6 shows the ablation study of т, K choices on the pruning test performance. Figure 7 shows the ablation study of T. K choices on the weight distributions of the trained DNNs.
From (a) to (d) in Figure 6 and 7 we can find 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 motivation for using K 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 final weight distributions will gain more smaller weights, which further benefits 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 r, FFW-training can induce weight dis- tributions 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 choose 15 and K2 5% the hyperparameters in our periments. This can give both satisfactory full model performance and pruning performance across different pruning ratios. We note that by choosing smaller Kr, e.g., K₁ 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 influence of initial learning rate ag on the test performance of SFW-pruning algorithm and the weight distributions of the trained DNNs. Specifically we run the SFW SFW-pruning algorithm (Algorithm 2) using VGG16 on CIFAR-10 10 with a range of initial learning rate og (0.1. 0.2, 0.4, 0.6, 0.8, 1.0]. Figure 8 shows the cor- responding 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 TinylmageNet. As is shown Figure 9, weight distribution with more smaller values is still more friendly to pruning.
347
1.22
Woyge Alviltt Vilco
(a) Pruning performance
(b) Weight Distributions
Figure 8: Comparison of pruning performance and weight distributions induced by SFW-pruning for VGG16 on CIFAR-10 with different initial learning rates og E 0.. 0.4, 0.6, 0.8, 1.881..0. 1.0}
B.3 PERFORMANCE ON MORE ARCHITECTURES AND DATASETS
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. Specifically, we consider (ResNet18, VGG16) and (CIFAR-1O, CIFAR-I00) 100) respectively. The results are summarized in Figure 10.
15

---

Published as a conference paper at ICLR 2022
/uuiing Tixt Thiieapee
38"
knd Taae20
382
38
Moogh Aisesle Vakes
Pruning performance
(b) Weight distributions
Figure 9: Comparison of pruning performance and weight distributions induced by SFW-pruning for ResNet18 on SinylmageNet The results show that even on complex dataset like TinyImageNet, appropriately more smaller weights can still benefit the pruning procedure.
127
107
10h
10
107
197
107
30'
107
591
Wedgh
(a) CIFAR-10 R8 (b) CIFAR-10 VGG-16 (c) CIFAR-100 ResNet-18 (d) CIFAR-100 VGG-16
Figure 10: Comparison of weight distributions induced by SFW and SGD over different architec- tures and datasets. For SFW-pruning. choose 15 and K; 5%, with og 1.0.
16


Reviews:
[{"review_id": "rT-Z_1uLk2P", "reviewer": "Reviewer_22uh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The topic studied by this paper is of general interest. The authors manage to organize the paper in a good way.\n\nStrength:\n\n(1) Formulate the pruning strategy with an optimization problem and solve it through Frank-Wolfe\n\n(2) The learning-based initialization helps to achieve better performance.\n\nQuestions:\n\n(1) Which step in the proposed algorithm other than backpropagation has the largest cost per iteration? \n\n(2) Does the algorithm model-agnostic? How is the performance in various vision models\n\n(3) Does the proposed algorithms scales to larger dataset such as ImageNet?", "summary_of_the_paper": "This paper proposes a Frank-Wolfe based approach for efficient pruning. \nThe highlights of this paper are:\n\n(1)  Propose a learning-based initialization scheme for SFW-based DNN training\n\n(2)  Achieves promising performance on DNN benchmark when we set different pruning ratios", "summary_of_the_review": "The paper proposes an interesting algorithm for training large DNNs. Particularly, the proposed method could handle different pruning ratios. \nThere still exist concerns regarding the scalability of the algorithm. The authors are encouraged to address it with more experiments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nuvyw4mEOUR", "reviewer": "Reviewer_9FPN", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think the proposed method here is actually quite interesting but this work fails to discuss a whole body of relevant literature that also aims to train sparse neural networks. As a result the chosen baselines for the experiments section fail to show how effective this method actually is.\n\nHere are few examples I was expecting to see discussed and/or evaluated:\n\n- Directly related to the \"one-shot no retrain\" claim of this paper:\n    - [Only Train Once: A One-Shot Neural Network Training And Pruning Framework by Chen et al](https://arxiv.org/abs/2107.07467)\n- Pruning during training:\n    - [Sparse Networks from Scratch: Faster Training without Losing Performance by Dettmers et al.](https://arxiv.org/abs/1907.04840)\n    - [The State of Sparsity in Deep Neural Networks by Gale et al.](https://arxiv.org/abs/1902.09574)\n    - [Dynamic Model Pruning with Feedback by Lin et al.](https://arxiv.org/abs/2006.07253)\n- there are also lots of works on inducing group sparsity on neural networks that's relevant to this paper. Some examples are:\n    - [Towards Compact ConvNets via Structure-Sparsity Regularized Filter Pruning by Lin et al.](https://arxiv.org/abs/1901.07827)\n    - [Learning Structured Sparsity in Deep Neural Networks by Wen et al.](https://arxiv.org/abs/1608.03665)\n    - [Structured Sparsity Inducing Adaptive Optimizers for Deep Learning by Deleu et al.](https://arxiv.org/abs/2102.03869)\n    - [Neuron-level Structured Pruning using Polarization\nRegularizer by Zhuang et al.](https://papers.nips.cc/paper/2020/file/703957b6dd9e3a7980e040bee50ded65-Paper.pdf)\n- and of course numerous papers that traine sparse neural networks with Lasso (or directly L0) regularization\n\nResults in Figure 2:\n- while it's true that SFW maintains performance in low pruning ratios but it seems to degrade a lot above 0.9 pruning ratio which is what typical pruning methods usually aim for. So doesn't it make more sense to compare SFW against other pruning methods at _their_ \"nature sparsity\"? I guess my point here is: why should I choose SFW over another classic pruning method that can achieve better performance at higher pruning ratio? \n- In Fig2. the performance of SGD on TinyImageNet is lower than what it should be (and of course lower than Franke-Wolfe training). I was wondering if authors have an explanation for this?\n\nResults in Figure 5:\n- Section 5.4 is supposed to compare against SOTA pruning methods but chooses pruning-at-initialisation baselines such as GraSP and SynFlow. This makes very little sense to me. Why not compare against pruning-during-training methods or even pruning-after-training methods?\n\nRandom:\n- Fig1a averages over multiple architectures and multiple datasets. Why does it make sense to do so? Is the shaded area the standard deviation over these differences or over multiple runs with different seed?\n- Isn't Fig1.b exactly the same as Fig2.a? \n\n", "summary_of_the_paper": "This paper swaps gradient-based optimisation for training neural networks with Frank-Wolfe algorithm with a constrain that pushes less important weights towards smaller values. The resulting trained models can be pruned to different sparsity targets without needing to be retrained.", "summary_of_the_review": "This paper proposes an interesting alternative to gradient descent that can be used to train models with a weight distribution suitable for one-shot pruning. I find the method section well-written and convincing however I was expecting to see many other relevant work here and the experiments section fails to compare to these relevant works, so it's difficult for me to evaluate the effectiveness of the method compared to simpler/other alternatives.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JYxuxQFGqLQ", "reviewer": "Reviewer_urBu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The core strength of this paper is to train a deep neural network so that it is more resilient to pruning. The author achieves this by using the K-sparse polytope constraint to enforce the model has more small weights. And they empirically show that such weight distribution can help weight pruning. They also use stochastic Frank-Wolfe for optimizing the problem. A gradient-based initialization scheme is also used to improve the trade-off between sparsity and performance.\n\nThe weakness of this paper can be summarized into the following points:\n1. The core benefit of the proposed method is not clear.  In 'Why K-Sparse Polytope Constraints?', authors argue that 'each SFW step pushes those less important weights smaller' and such weight distribution 'can more smoothly remove small values' and 'yielding competitive test accuracies'. On the other hand, in 'Learning rate $\\alpha$', authors argue that a larger learning rate can 'result in sparser parameter $\\theta_{t+1}$' and higher natural sparsity. The latter argument suggests that the benefit of K-sparse constraints is due to they can learn sparser solutions for the deep neural network. The former argument suggests that the benefit of K-sparse constraints is a weight distribution that can make pruning easier. The latter argument is a quite trivial result. The former argument is interesting but requires more analysis. In addition, it seems that the K-sparse constraint is not designed to obtain a sparse solution. The property of the Frank-Wolf algorithm ensures $\\theta$ always attains the K-sparse constraint. If the K-sparse constraint is really useful, then it is not clear why $\\alpha$ is important to the performance.\n\n2. The authors should provide a more detailed analysis of K and $\\tau$ for K-sparse polytope constraints, and how they affect the weight distributions. I also suggest the authors include a more detailed definition of K-sparse polytope constraints and their properties, to make the paper self-contain. \n\n3. Authors use the Frank-Wolfe framework, which is drastically different from regular techniques for training deep neural networks. As a result, ImageNet experiments with standard baseline models, like ResNet-50, are required, and numerical results should also be reported. The reason is that scalability is crucial to weight pruning algorithms. Moreover, on ImageNet, whether Frank-Wolfe algorithms can obtain similar performance as popular stochastic optimizers is not clear. \n\n4. In Figure.2, the advantage of the proposed method becomes much smaller when the dataset is more complex. When applying large models on small datasets, it is plausible to say small weights are not important. However, when the dataset becomes complex, small weights may also become important as suggested by previous works [1]. From this perspective, the effects of pushing more weights smaller may vanish when the dataset is large enough.\n\n[1] Ye, Jianbo, et al. \"Rethinking the smaller-norm-less-informative assumption in channel pruning of convolution layers.\" ICLR, 2018.\n\n----------------------------------------------------\n\nAfter the authors' response, most of my questions were addressed, and I increased my score to 6.\n", "summary_of_the_paper": "In this paper, the authors proposed to use stochastic Frank-Wolfe with K-sparse polytope constraints to train deep neural networks and make them pruning friendly. ", "summary_of_the_review": "In summary, I think this paper is below the acceptance threshold since the main claim lacks enough support. And current results may be due to sparsity incurred by the large learning rate $\\alpha$. In addition, experimental results are not enough.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-qWGW0BZaE", "reviewer": "Reviewer_BXpo", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n\n1. The problem of pruning-aware training is interesting and has a lot of potential.\n2. The formulation and approach to optimize it using SFW is novel.\n3. In addition an initialization scheme that is suited for SFW is also presented.\n4. The results demonstrate the value of pruning-aware training using the proposed approach.\n\n# Weaknesses\n\n1. The motivation for K-sparse constraint set is not clear. More clarity is required on how the authors chose to use such a constraint set.\n\n2. If I understand correctly, the \\tau is a constant for all the weights in the network and therefore, the SFW update at any iteration is basically based on the sign of the gradient. Is this understanding correct? If yes, it further raises questions about the choice of the K-sparse constraint set. What is the need to use such a fixed \\tau? How to decide the value of \\tau? It is intriguing that even with such a restrictive update the final weights are competitive to SGD based training.\n\n3. The performance is still worse than SGD with retraining at extreme sparsities (fig. 5) but considering the cost of retraining the proposed method is useful.", "summary_of_the_paper": "The paper presents an algorithm to train a DNN in a pruning-aware manner such that the trained network can be pruned at various sparisities without any fine-tuning. The approach is based on stochastic frank-wolfe and the results are impressive for a wide-range of sparsities.", "summary_of_the_review": "The idea is interesting and the method is technically sound. The choice of the constraint set should be motivated and clarified to improve the quality of the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This work considers one-shot pruning in deep neural networks. The main departure from previous work is to consider stochastic Frank-Wolfe. The reported results are convincing although a number of baselines were missing from the initial submission. The authors provide a balanced account of the strengths and weaknesses of the proposed approach.

The authors adequately addressed the concerns of the reviewers. For instance they ran additional experiments to compare to missing pruning baselines. I would encourage the authors to revise the manuscript by including the missing related work, the additional clarification discussions (e.g., motivation for K-sparse constraints, follow-up analysis, and cost per iteration) and to include the additional experiments that were conducted (e.g., pruning with training).
Decision: Accept (Spotlight)