Advancing Model Pruning via Bi-level Optimization
Anonymous Author(s)
AffiliationAddress email
Abstract
The deployment constraints in practical applications necessitate the pruning oflarge-scale deep learning models, i.e. , promoting their weight sparsity.
As illustrated by the Lottery Ticket Hypothesis (LTH), pruning also has the potential ofimproving their generalization ability.
At the core of LTH, iterative magnitudepruning (IMP) is the predominant pruning method to successfully find ‘winningtickets’, i.e. , pruned sparse models with better generalization than the originaldense models.
Yet, the computation cost of IMP grows prohibitively as the targetedpruning ratio increases.
To reduce the computation overhead, various efficient‘one-shot’ pruning methods have been developed but these schemes are usuallyunable to find winning tickets as good as IMP.
This raises the question of how toclose the gap between pruning accuracy and pruning efficiency?
To tackle it, wepursue the algorithmic advancement of model pruning.
Specifically, we formulatethe pruning problem from a fresh and novel viewpoint, bi-level optimization (BLO).
We show that the BLO interpretation provides a technically-grounded optimizationbase for an efficient implementation of the pruning-retraining learning paradigmused in IMP.
We also show that the proposed bi-level optimization-oriented pruningmethod (termed B I P) is a special class of BLO problems with a bi-linear problemstructure.
By leveraging such bi-linearity, we theoretically show that B I P canbe solved as easily as first-order optimization, thus inheriting the computationefficiency.
Through extensive experiments on both structured and unstructuredpruning with 5 model architectures and 4 data sets, we demonstrate that B I P canfind better winning tickets than IMP in most cases, and is computationally asefficient as the one-shot pruning schemes, demonstrating 2-7 × speedup over IMPfor the same level of model accuracy and sparsity.
Introduction
Deep neural networks (DNNs) have dramatically improved the state-of-the-art in various fields [1, 2].
While over-parameterized structures are key to the improved generalization of DNNs [3], they createnew problems – the millions or even billions of parameters not only increase computational costsduring inference, but also pose serious deployment challenges on resource-limited devices [4].
As aresult, model pruning has seen a lot of research interest in recent years, focusing on reducing modelsizes by removing (or pruning) redundant parameters [4–8].
Among various proposed model pruning algorithms [5, 9–24], the heuristics-based Iterative Magnitude
Pruning (IMP) is the current dominant approach to achieving model sparsity without sufferingperformance loss, as suggested and empirically justified by the Lottery Ticket Hypothesis (LTH)[12].
The LTH hypothesizes the existence of a sub-network (the so-called ‘winning ticket’) whentrained in isolation ( e.g. , from either a random initialization [12] or an early-rewinding point ofdense model training [14]), can match the performance of the original dense model [5, 12–16].
The core idea of IMP is to iteratively prune and retrain the model while progressively pruninga small ratio of the remaining weights in each iteration, and continuing till the desired pruning ratio has been reached.
While IMP often finds the winning ticket, it incurs the cost of repeatedmodel retraining from scratch, making it prohibitively expensive for large datasets or large modelarchitectures [25, 26].
To improve the efficiency of model pruning, numerous heuristics-basedone-shot pruning methods [12, 17–21], e.g. , one-shot magnitude pruning (OMP), have been proposed.
These schemes directly prune the model to the target sparsity and are significantly more efficient than
IMP.
Yet, the promise of current one-shot pruning methods is dataset/model-specific [18, 27] andmostly lies in the low pruning ratio regime [21].
As systematically studied in [18, 28], there existsa clear performance gap between one-shot pruning and IMP.
As an alternative to heuristics-basedschemes, optimization-based pruning methods [10, 11, 22–24] still follow the pruning-retrainingparadigm and adopt sparsity regularization [29, 30] or parameterized masking [11, 22, 24] to prunemodels efficiently.
However, these methods do not always match the accuracy of IMP and thus havenot been widely used to find winning tickets [12, 17, 19–21, 26].
The empirical results showing thatoptimization underperforms heuristics motivate us to revisit the algorithmic fundamentals of thisresearch: Can we reformulate optimization-based pruning to match and improve over IMP in termsof model accuracy?
[Figure]
To this end, we put forth a novel perspective of modelpruning as a bi-level optimization (BLO) problem.
Inthis new formulation, we show that BLO provides atechnically-grounded optimization basis for an efficientimplementation of the pruning-retraining paradigm, thekey algorithmic component used in IMP.
To the bestof our knowledge, we make the first rigorous connection between model pruning and BLO.
Technically, wepropose a novel bi -level optimization-enabled p runingmethod (termed B I P).
We further show how B I P takesadvantage of the bi-linearity of the pruning problem toavoid the computational challenges of common BLOmethods, and is as efficient as any first-order alternatingoptimization scheme.
Practically, we demonstrate thesuperiority of the proposed B I P in terms of accuracy,sparsity, and computation efficiency through extensiveexperiments.
B I P finds the best winning tickets nearly in all settings while taking time comparableto the one-shot OMP.
In Fig.
1, we present a snapshot of our empirical results for CIFAR-10 with
ResNet architectures at 74% pruning ratio.
In all cases, B I P ( ⋆ ) finds winning tickets, improvingaccuracy over the dense model ( • ) and matching IMP ( ▲ ), while being upto 5 × faster than IMP.
Our contributions can be summarized as follows:
• (Formulation)
We rethink the algorithmic foundation of model pruning through the lens of BLO(bi-level optimization).
The new BLO-oriented formulation provides an optimization basis of thewidely-accepted pruning-retraining paradigm.
Moreover, the formulation disentangles pruning andretraining variables, providing the flexibility to design the interface between pruning and retraining.• (Algorithm)
We propose the new bi-level pruning (B I P) algorithm, which is built upon the aforementioned BLO formulation and the implicit gradient-based optimization theory.
Unlike computationallyintensive standard BLO solvers, we theoretically show that B
I P is as efficient as any first-orderoptimization by taking advantage of the bi-linear nature of the pruning variables.• (Experiments)
We conduct extensive experiments across 4 datasets (CIFAR-10/100, Tiny-ImageNetand ImageNet), 5 model architectures, and 2 pruning settings (unstructured and structured pruning).We show that B I P enjoys the high accuracy of IMP and the high efficiency of the one-shot pruningschemes. Specifically, we show that (i) B I P achieves higher test accuracy than IMP and finds the bestwinning tickets nearly in all settings and (ii) B I P is highly efficient (comparable to one-shot pruningschemes), achieving 2-7 × speedup over IMP for the same level of model accuracy and sparsity
Related Work and Open Question
Neural network pruning.
As neural networks become deeper and more sophisticated, modelpruning technology has gained increasing attention over the last decade since pruned models arenecessary for the deployment of deep networks in practical applications [4, 31].
With the goalof finding highly-sparse and highly-accurate subnetworks from original dense models, a variety of pruning methods have been developed such as heuristics-based pruning [12, 17, 19–21, 26]and optimization-based pruning [11, 22–24].
The former identifies redundant model weights byleveraging heuristics-based metrics such as weight magnitudes [6, 12, 14, 16, 18, 32, 28], gradientmagnitudes [17, 19, 20, 33], and Hessian statistics [34–38].The latter is typically built on: 1) sparsitypromoting optimization [10, 30, 39–41], where model weights are trained by penalizing their sparsityinducing norms, such as ℓ 1 norm for irregular weight pruning and ℓ 2 norm for structured pruning;2) parameterized masking [11, 22, 42], where model weight scores are optimized to filter the mostimportant weights and achieve better performance.
Iterative vs. one-shot pruning, and motivation.
Existing schemes can be further categorized intoone-shot or iterative pruning based on the pruning schedule employed for achieving the targeted modelsparsity.
Among the iterative schemes, the IMP (Iterative Magnitude Pruning scheme) [12, 15, 43–52] has played a significant role in identifying high-quality sub-networks or ‘winning tickets’, aspostulated by LTH (Lottery Ticket Hypothesis) [13, 14].
For example, the current state-of-the-art(SOTA) implementation of IMP in [18] can lead to a pruned ResNet-20 on CIFAR-10 with 74%sparsity and 92 .
12% test accuracy, matching the performance of the original dense model (see red ⋆in Fig.
2- Left ).
The IMP algorithm typically contains two key ingredients: (i) a temporally-evolvingpruning schedule to progressively increase model sparsity over pruning iterations, and (ii) the pruningretraining learning mechanism applied at each pruning iteration.
With a target pruning ratio of p %with T pruning iterations, an example pruning schedule in (i) could be as follows – each iterationprunes ( p %) 1 / T of the currently unpruned model weights, progressively pruning fewer weights ineach iteration.
For (ii), the unpruned weights in each pruning iteration are re-set to the weightsat initialization or at an early-training epoch [13], and re-trained till convergence.
In brief, IMPrepeatedly prunes, resets, and trains the network over multiple iterations.
[Figure]
However, winning tickets found by IMP incur significant computationalcosts.
The sparsest winning ticketfound by IMP in Fig.
2- Left (red⋆ ) utilizes T = 6 pruning iterations.
As shown in Fig.
2- Right , this takes3 × more time than the original training of the dense model.
To avoidthe computational cost of IMP, different kinds of ‘accelerated’ pruningmethods were developed [12, 17, 19–21, 26], and many fall into the oneshot pruning category: The networkis directly pruned to the target sparsity and retrained once.
In particular, OMP (one-shot magnitude pruning) is an important baselinethat simplifies IMP [12].
It follows the pruning-retraining paradigm, but prunes the model to thetarget ratio with a single pruning iteration.
Although one-shot pruning schemes are computationallycheap (Fig.
2- Right ), they incur a significant accuracy drop compared to IMP (Fig.
2- Left ).
Even if IMP is customized with a reduced number of training epochs per pruning round, the pruning accuracyalso drops largely (see Fig. A3).
Hence, there is a need for advanced model pruning techniques tofind winning tickets like IMP, while being efficient like one-shot pruning.
Different from unstructured weight pruning described above, structured pruning takes into consideration the sparse patterns of the model, such as filter and channel-wise pruning [40, 42, 53–57].
Structured pruning is desirable for deep model deployment in the presence of hardware constraints [4, 31].
However, compared to unstructured pruning, it is usually more challenging to maintain the performance and find structure-aware winning tickets [25, 26].
143 Open question.
As discussed above, one-shot pruning methods are unable to match the predictiveperformance of IMP, and structure-aware winning tickets are hard to find.
One possible cause is therequirement to prune in a single iteration or in a structure-constrained space mainly through heuristicsbased strategies, while the best optimization foundation of model pruning is underdeveloped.
Thus,we consider the following open question:
Is there an optimization basis for a successful pruning algorithm that can attain high pruned model accuracy (like IMP) and high computational efficiency (like one-shot pruning)?
The model pruning problem has a natural hierarchical structure – we need to find the best mask toprune model parameters, and then , given the mask, find the best model weights for the unprunedmodel parameters.
Given this hierarchical structure, we believe that the bi-level optimization (BLO)framework is one promising optimization basis for a successful pruning algorithm.
We briefly reviewthe BLO-related work below.
Bi-level optimization (BLO).
BLO is a general hierarchical optimization framework, where theupper-level problem depends on the solution to the lower-level problem.
Such a nested structuremakes the BLO in its most generic form very difficult to solve, since it is hard to characterize theinfluence of the lower-level optimization on the upper-level problem.
Various existing literaturefocuses on the design and analysis of algorithms for various special cases of BLO.
Applications rangefrom the classical approximate descent methods [58, 59], penalty-based methods [60, 61], to recent
BigSAM [62] and its extensions [63, 64].
It is also studied in the area of stochastic algorithms [65–67]and back-propagation based algorithms [68–70].
BLO has also advanced adversarial training [71],meta-learning [72], data poisoning attack generation [73], automated machine learning [74] as wellas reinforcement learning [75].
However, to the best of our knowledge, the BLO framework has notbeen rigorously considered for neural network pruning.
Our work, for the first time, advances thealgorithmic foundation of neural network pruning through the lens of BLO.
B I P: Model Pruning via Bi-level Optimization
In this section, we re-investigate the problem of model pruning through the lens of BLO and developthe bi-level pruning (B I P) algorithm.
We can theoretically show that B
I P can be solved as easily asthe first-order alternating optimization by taking advantage of the bi-linearity of pruning variables.
A BLO viewpoint on model pruning As described in the previous section, the pruning-retraininglearning paradigm covers two kinds of tasks: ❶ pruning that determines the sparse pattern of modelweights, and ❷ training remaining non-zero weights to recover the model accuracy.
In existingoptimization-based pruning methods [76–79], the tasks ❶ - ❷ are typically achieved by optimizingmodel weights, together with penalizing their sparsity-inducing norms, e.g. , the ℓ 1 and ℓ 2 norms [80].
Different from the above formulation, we propose to separate optimization variables involved in thepruning tasks ❶ and ❷ .
This leads to the (binary) pruning mask variable m ∈ { 0 , 1 } n and the modelweight variable θ ∈ R n .
Here n denotes the total number of model parameters.
Accordingly, thepruned model is given by ( m ⊙ θ ) , where ⊙ denotes the element-wise multiplication.
As will beevident later, this form of variable disentanglement enables us to explicitly depict how the pruningand retraining process co-evolve, and helps customize the pruning task with high flexibility.
We next elaborate on how BLO can be established to co-optimize the pruning mask m and theretrained sparse model weights θ .
Given the pruning ratio p % , the sparsity constraint is given bym ∈ S , where S =
{ m | m ∈ { 0 , 1 } n , 1 T m ≤ k } and k
=
(1 − p %) n .
Our goal is to prunethe original model directly to the targeted pruning ratio p % ( i.e. , without calling for the IMP-likesparsity schedule as described in Sec. 2) and obtain the optimized sparse model ( m ⊙ θ ) .
To thisend, we interpret the pruning task ( i.e. , ❶ ) and the model retraining task ( i.e. , ❷ ) as two optimizationlevels , where the former is formulated as an upper-level optimization problem, and it relies on theoptimization of the lower-level retraining task.
We thus cast the model pruning problem as thefollowing BLO problem (with ❷ being nested inside ❶ ):
[Equation]
(cid:123)(cid:122) (cid:125) ❶ : Pruningtask ; subject to θ ∗ ( m ) = arg min θ ∈ R n := g ( m , θ ) (cid:122) (cid:125)(cid:124) (cid:123) ℓ ( m ⊙ θ )
+ γ 2 ∥ θ ∥(cid:124) (cid:123)(cid:122) (cid:125) ❷ : Sparsity-fixedmodelretraining , (1) where ℓ denotes the training loss ( e.g. , cross-entropy), m and θ are the upper-level and lower-leveloptimization variables respectively, θ ∗ ( m ) signifies the lower-level solution obtained by minimizingthe objective function g given the pruning mask m , and γ > 0 is a regularization parameter introducedto convexify the lower-level optimization so as to stabilize the gradient flow from θ ∗ ( m ) to m andthus the convergence of BLO [67, 81].
In a sharp contrast to existing single-level optimization-basedmodel pruning methods [76–79], the BLO formulation (1) brings in two advantages.
First, BLO has the flexibility to use mismatched pruning and retraining objectives at the upper andlower optimization levels, respectively.
This flexibility allows us to regularize the lower-level trainingobjective function in (1) and customize the implemented optimization methods at both levels.
To be more specific, one can update the upper-level pruning mask m using a data batch (called B 2 ) distinctfrom the one (called B 1 ) used for obtaining the lower-level solution θ ∗ ( m ) .
The resulting BLOprocedure can then mimic the idea of meta learning to improve model generalization [82], where thelower-level problem fine-tunes θ using B 1 , and the upper-level problem validates the generalizationof the sparsity-aware finetuned model ( m ⊙ θ ∗ ( m )) using B 2 .
Second, BLO enables us to explicitly model and optimize the coupling between the retrained modelweights θ ∗ ( m ) and the pruning mask m through the implicit gradient (IG)-based optimizationroutine.
Here IG refers to the gradient of the lower-level solution θ ∗ ( m ) with respect to (w.r.t.)the upper-level variable m , and its derivation calls the implicit function theory [61].
The use of IGmakes our proposed BLO-oriented pruning (1) significantly different from the greedy alternatingminimization [83] that learns the upper-level and lower-level variables independently ( i.e. , minimizesone variable by fixing the other).
We refer readers to the following section for the detailed IG theory.
We will also show in Sec. 4 that the pruning strategy from (1) can outperform IMP in many pruningscenarios but is much more efficient as it does not call for the scheduler of iterative pruning ratios.
Optimization foundation of B I
P.
The key optimization challenge of solving the B
I
P problem (1)lies in the computation of IG (implicit gradient).
Prior to developing an effective solution, we firstelaborate on the IG challenge , the unique characteristic of BLO.
In the context of gradient descent,the gradient of the objective function in (1) yields
[Equation]
where ∇ m and ∇ θ denote the partial derivatives of the bi-variate function ℓ ( m ⊙ θ )
w.r.t.
thevariable m and θ respectively, d θ ⊤ /d m ∈ R n × n denotes the vector-wise full derivative , and for easeof notation, we will omit the transpose ⊤ when the context is clear.
In (2), the IG challenge refers tothe demand for computing the full gradient of the implicit function θ ∗ ( m ) =
arg min θ g ( m , θ ) w.r.t.m , where recall from (1) that g ( m , θ ) :=
ℓ ( m ⊙ θ ) + γ 2 ∥ θ ∥ 22 .
Next, we derive the IG formula following the more rigorous implicit function theory [61, 67, 72].
Based on the fact that θ ∗ ( m ) satisfies the stationarity condition for the lower-level objective functionin (2), it is not difficult to obtain that (see derivation in Appendix A)
[Equation]
where ∇ 2 m θ ℓ and ∇ 2 θ ℓ denote the second-order partial derivatives of ℓ respectively, and ( · ) − 1 denotesthe matrix inversion operation.
Yet, the exact IG formula (3) remains difficult to calculate due to the presence of matrix inversion andsecond-order partial derivatives.
To simplify it, we impose the Hessian-free assumption, ∇ 2 θ ℓ = 0 ,which is mild in general; For example, the decision boundaries of neural networks with ReLUactivations are piece-wise linear in a tropical hyper-surface [84], and this assumption has been widelyused in BLO-involved applications such as meta learning [85] and adversarial learning [71].
Given∇ 2 θ ℓ = 0 , the matrix inversion in (3) can be then mitigated, leading to the IG formula
[Equation]
234 At the first glance, the computation of the simplified IG (4) still requires the mixed (second-order)partial derivative ∇ 2 m θ ℓ .
However, B I P is a special class of BLO problems with bi-linear variables( m ⊙ θ ) .
Based on this bi-linearity, we can prove that IG in (4) can be analytically expressed usingonly first-order derivatives; see the following theorem.
Proposition 1 Assuming ∇ 2 θ ℓ = 0 and defining ∇ z ℓ ( z ) := ∇ z ℓ ( z ) |
z = m ⊙ θ ∗ , the implicit gradient(4) is then given by
[Equation]
Further, the gradient of the objective function given by (2) becomes
[Equation]
where ⊙ denotes the element-wise multiplication.
Proof : Using chain-rule, we can obtain that
[Equation]
where diag( · ) represents a diagonal matrix with · being the main diagonal vector.
Further, we canconvert (4) to
[Equation]
where the last equality holds due to the Hessian-free assumption.
With (9) and (4) we can prove (5).
Next, substituting the IG (5) to the upper-level gradient (2), we obtain thatdℓ [Equation] ⊙
[Equation]
which leads to (6).
The proof is now complete.
□
The key insight drawn from Prop. 1 is that the bi-linearity of pruning variables ( i.e. , m ⊙ θ ∗ ) makesthe IG-involved gradient (2) easily solvable, and the computational complexity is almost the same asthat of computing the first-order gradient ∇ z ℓ ( z ) just once, as supported by (6)
B I P algorithm and implementation.
We next formalize the B
I P algorithm based on Prop. 1 andthe alternating gradient descent based BLO solver [67].
At iteration t , there are two main steps.
★ Lower-level SGD for model retraining :
Given m ( t − 1) , θ ( t − 1) , and z ( t − 1) :
= m ( t − 1) ⊙ θ ( t − 1) , weupdate θ ( t ) by randomly selecting a data batch with the learning rate α and applying SGD (stochasticgradient descent) to the lower-level problem of (1),
[Equation]
★ Upper-level SPGD for pruning :
Given m ( t − 1) , θ ( t ) , and z ( t +1 / 2) :
= m ( t − 1) ⊙ θ ( t ) , we updatem using SPGD (stochastic projected gradient descent) along the IG-enhanced descent direction (2),
[Equation]
where β > 0 is the upper-level learning rate, and P S ( · ) denotes the Euclidean projection onto theconstraint set S given by S
= { m | m ∈ { 0 , 1 } n , 1 T m ≤ k } in (1) and is achieved by the top- khard-thresholding operation as will be detailed later.
In B I P, the ( θ -step) and ( m -step) steps execute iteratively.
In practice, we find that B I P convergesvery
well; see Fig.
A7 for an ablation study.
It is also worth noting that both ( θ -step) and ( m -step) onlyrequire the first-order information ∇ z ℓ ( z ) , demonstrating that B I P can be conducted as efficiently asfirst-order optimization.
We present more implementation details of B I P below and refer readers to
Appendix B for a detailed algorithm description.
✦ Discrete optimization over m : We follow the ‘convex relaxation + hard thresholding’ mechanismused in [11, 22].
Specifically, we relax the binary masking variables to continuous masking scoresm ∈ [ 0 , 1 ] .
We then acquire loss gradients at the backward pass based on the relaxed m .
At theforward pass, we project it onto the discrete constraint set S using the hard thresholding operator,where the top k elements are set to 1 s and the others to 0 s. See Appendix B for more discussion.
✦ Data batch selection for lower-level and upper-level optimization: We adopt different data batches(with the same batch size) when implementing ( θ -step) and ( m -step).
This is one of the advantagesof the BLO formulation, which enables the flexibility to customize the lower-level and upper-levelproblems.
The use of diverse data batches is beneficial to generalization as shown in [82].
✦ Hyperparameter tuning: As described in ( θ -step)-( m -step), B I P needs to set two learning rates αand β for lower-level and upper-level optimization, respectively.
We choose α = 0 .
01 and β = 0 .
1 in all experiments, where we adopt the mask learning rate β from Hydra [22] and set a smallerlower-level learning rate α , as θ is initialized by a pre-trained dense model.
We show ablation studyon α in Fig. 6(c).
BLO also brings in the low-level convexification parameter γ .
We set γ = 1 .
inexperiments and refer readers to Fig.
6(b) for a sanity check.
✦ One-step vs. multi-step SGD: In ( θ -step), the one-step SGD is used and helps reduce the computation overhead.
In practice, we also find that the one-step SGD is sufficient: The use of multi-step
SGD in B I P does not yield much significant improvement over the one-step version; see Fig. 6(a).
✦ Extension to structured pruning: We formulate and solve the B
I
P problem in the context ofunstructured (element-wise) weight pruning.
However, if define the pruning mask m w.r.t.
model’sstructural units ( e.g. , filters), then it is easy to apply B I P to structured weight pruning; see Fig.
4 and
Fig.
A6 for empirical results.
Experiments
In this section, we present extensive experimental results to show the effectiveness of B I P acrossmultiple model architectures, various datasets, and different pruning setups.
Compared to IMP,one-shot pruning, and optimization-based pruning baselines, we find that B
I P can find better winningtickets in most cases and is computationally efficient.
4.1 Experiment Setup
Datasets and models.
Following the pruning benchmark in [18], we consider 4 datasets including
CIFAR-10 [86], CIFAR-100 [86], Tiny-ImageNet [87], ImageNet [88], and 5 architecture types including ResNet-20/56/18/50 and VGG-16 [89, 90].
Tab. A1 summarizes these datasets and modelconfigurations and setups.
Baselines, training, and evaluation.
As baselines, we mainly focus on 4 SOTA pruning methods, 299 ①
IMP [12], ② OMP [12], ③ G RASP [91] (a one-shot pruning method by analyzing gradientflow at initialization), and ④ H YDRA [22] (an optimization-based pruning method that optimizesmasking scores).
It is worth noting that there exist various implementations of IMP, e.g. , specifiedby different learning rates and model initialization or ‘rewinding’ strategies [13].
To make a faircomparison, we follow the recent IMP benchmark in [18], which can find the best winning tickets overcurrent heuristics-based pruning baselines.
We also remark that H YDRA is originally proposed forimproving the adversarial robustness of a pruned model, but it can be easily customized for standardpruning when setting the adversary’s strength as 0 [22].
We choose H YDRA as a baseline becauseit can be regarded as a single-level variant of B I P with post-optimization weight retraining.
Whenimplementing B I P, unless specified otherwise, we use the 1 -step SGD in ( θ -step), and set the learningrates ( α, β ) and the lower-level regularization parameter γ as described in the previous section.
Whenimplementing baselines, we follow their official repository setups.
We evaluate the performanceof all methods mainly from two perspectives: (1) the test accuracy of the sub-network, and (2) theruntime of pruning to reach the desired sparsity.
We refer readers to Tab. A3 and Appendix C.1 formore training and evaluation details, such as training epochs and learning rate schedules.
4.2 Experiment Results
B I P identifies high-accuracy winning tickets.
In what follows, we look at the quality of winningtickets identified by B I P. Two key observations can be drawn from our results:
(1) B I P finds winningtickets of higher accuracy and/or higher sparsity than the baselines in most cases (as shown in Fig.
3and Tab. 1); (2) The superiority of B I P holds for both unstructured pruning and structured pruning (asshown in Fig.
4 and Fig.
A6 ).
Fig.
3 shows the unstructured pruning trajectory (given by test accuracy vs. pruning ratio) of B I Pand baseline methods in 8 model-dataset setups.
For comparison, we also present the performance ofthe original dense model.
As we can see, the proposed B I P approach finds the best winning tickets(in terms of the highest accuracy) compared to the baselines across all the pruning setups.
Note that awinning ticket refers to a sparse model with test accuracy that is at least as good as the dense model(see the marked green dash line).
Among the baseline methods, IMP is the pruning method mostcompetitive to ours.
However, the improvement brought by B I P is significant with respect to thevariance of IMP, except for the 60% - 80% sparsity regime in (CIFAR-10, ResNet-20).
In the case of(CIFAR-100, ResNet-20) where IMP can not find any winning tickets (as confirmed by [18]), B I P
[Table]
[Figure]
We observe that B
I P consistently outperforms the other baselines.
Note in the (ImageNet, ResNet- setting, we only compare B I P with our strongest baseline IMP due to computational resource constraints.
still manages to find winning tickets with around 0 .
6% improvement over the dense model. In Tab. 1,we summarize the sparsest winning tickets along the pruning trajectory identified by different pruningmethods. As we can see, B I P can identify the winning tickets with higher sparsity levels than theother methods, except in the case of (CIFAR-10, ResNet-20).
We also refer readers to Appendix C.3for more experimental results to justify the high quality of winning tickets identified by B I P.
Fig.
4 demonstrates the structured pruning trajectory on the CIFAR-10/100 datasets.
Here we focuson filter pruning, where the filter is regarded as a masking unit in (1).
We refer readers to Appendix C.3for channel-wise pruning results.
Compared to Fig.
3, Fig.
4 shows that it becomes more difficultto find winning tickets of high accuracy and sparsity in the structured pruning, and the gap amongdifferent methods decreases.
This is not surprising, since filter pruning imposes much stricter pruningstructure constraints than irregular pruning.
However, B I P still outperforms all the baselines.
Mostimportantly, it identifies clear winning tickets in the low sparse regime even when IMP fails.
[Figure]
B I P is computationally efficient.
In our experiments, another key observation is that B I P yields sparsity-agnostic runtime complexity while IMP leads to runtime exponential to thetarget sparsity.
Fig.
5 shows the computation cost of differentmethods versus pruning ratios on (CIFAR-10, ResNet-18).
Forexample, B I P takes 86 mins to find the sparsest winning ticket(with 89% sparsity in Tab. 1).
This yields 7 × less runtime than
IMP, which consumes 620 mins to find a comparable winningticket with 87% sparsity.
Compared to the optimization-basedbaseline H YDRA , B I P is more efficient as it does not relyon the extra post-optimization retraining; see
Tab. A2 for adetailed summary of runtime and number of training epochsrequired by different pruning methods.
Further, B I P takes about 1.25 × more computation time than [Figure] [Figure] [Figure]
[Figure] [Figure] Ratio [Figure]
T e s t
A cc u r ac y ( % ) 92.
Dense Model w/o IG w/ IG
[Figure]
G RASP and OMP.
However, the latter methods lead to worse pruned model accuracy, as demonstratedby their failure to find winning tickets in Tab. 1, Fig. 3, and Fig. 4.
Rewinding effect on B I P. Another advantage of B I P is that it is not sensitive to model rewinding.
Recall that rewinding is a strategy used in LTH [14] to determine what model initialization should beused for re-training a pruned sparse model.
As shown in [18], an IMP-identified winning ticket couldbe sensitive to the choice of a rewinding point.
Fig 6(a) shows the test accuracy of the B I P-prunedmodel when it is retrained at different rewinding epochs, where ‘N/A’ in the x-axis represents the caseof no retraining (and thus no rewinding).
As we can see, a carefully-tuned rewinding scheme does notlead to a significant improvement over B I P without retraining.
This suggests that the subnetworksfound by B I P are already of high quality, and justifies the improved computation efficiency of B
I Pover H YDRA (which requires post-optimization model retraining), as shown in Fig.
Sanity check of B I P on specialized hyperparameters.
In Fig.6(b), we show the sensitivity of
B I P to its specialized hyperparameters at lower-level optimization, including the number of SGDiterations ( N ) in ( θ -step), and the regularization parameter γ in (1).
Fig 6(b1) shows the test accuracyof B I P-pruned models versus the choice of N .
As we can see, more SGD iterations for the lower-leveloptimization do not improve the performance of B I
P. This is because in B I P, the θ -step is initializedby a pre-trained model which does not ask for aggressive weight updating.
The best performance of
B I P is achieved at N ≤ 3 .
We choose N = 1 throughout the experiments as it is computationallylightest.
Fig 6(b2) shows the performance of B I P by varying γ .
As we can see, the choice ofγ ∈ { 0 . 5 , 1 } yields the best pruning accuracy across all pruning ratios.
If γ is too small, the lackof strong convexity of lower-level optimization would hamper the convergence.
If γ is too large,the lower-level optimization would depart from the true model objective and causes a performancedrop.
Fig.
6(b3) demonstrates the necessity of the IG enhancement in B I
P. We compare B I P with its
IG-free version by dropping the IG term in (2).
We observe that B I P without IG (marked in blue)leads to a significant performance drop ( > 1% ) at various sparsities.
This highlights that the IG in the( m -step) plays a critical role in the performance improvements obtained by B I P, justifying our novel
BLO formulation for the pruning problem.
See Fig.
A5 for more results on hyperparameter tuning.
Conclusion
We proposed the B I P method to find sparse networks through the lens of BLO.
Our work advanced thealgorithmic foundation of model pruning by characterizing its pruning-retraining hierarchy using BLO.
We theoretically showed that B
I P can be solved as easily as first-order optimization by exploiting thebi-linearity of pruning variables.
We also empirically showed that B
I P can find high-quality winningtickets very efficiently compared to the predominant iterative pruning method.
In the future, we willexplore ways to investigate the effect of Hessian-free assumption on the performance of B I P, andinvestigate the performance of our proposal for actual hardware acceleration.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section .
• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Appendix C.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Appendix C. (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 3.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Section 3 andAppendix A.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[Yes] See supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Sec.
(c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes] See Sec.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Table A24.
If you are using existing assets (e g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[No] The licences of used datasets are
[Table]
provided in the cited references
[Table] (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
[Table]
Our code is included in the supplementary material
[Table] (d) Did you discuss whether and how consent was obtained from people whose data you’re [Table]
using/curating?
[N/A]
[Table] (e) Did you discuss whether the data you are using/curating contains personally identifiable [Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table] (a) Did you include the full text of instructions given to participants and screenshots, if [Table] applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
[Table]