A U NIFIED F RAMEWORK OF S OFT T HRESHOLD P RUNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Soft threshold pruning is among the cutting-edge pruning methods with stateof-the-art performance 1 .
However, previous methods either aimlessly perform searching on the threshold scheduler or simply train the threshold, lacking theoretical explanation from a unified perspective.
In this work, we reformulate soft threshold pruning as an implicit optimization problem solved using the Iterative
Shrinkage-Thresholding Algorithm (ISTA), a classic method from the fields of sparse recovery and compressed sensing.
Under this theoretical framework, all threshold tuning strategies proposed in previous studies of soft threshold pruning are explained as a specific arrangement style of regularization term.
We further derive an optimal threshold scheduler through an in-depth study of threshold scheduling based on our framework.
This scheduler keeps L 1 -regularization in the equivalent optimization problem stable and corresponds to a consistent objective function and can be, in principle, applied to sparsify any mathematical model that includes parameters trained via SGD.
We conduct extensive experiments and verify its state-of-the-art performance on both Artificial Neural Networks (ResNetand MobileNet-V1) and Spiking Neural Networks (SEW ResNet-18) on ImageNet datasets.
It further evolves into a family of novel pruning methods, including sparsify-during-training, early pruning, and pruning at initialization via analysis based on our framework.
I NTRODUCTION
Pruning has been a thriving area of network compression.
Since the day deep neural networks stretch their tentacles to every corner of machine learning applications, the demand for shrinking the size of network parameters has never stopped growing.
Fewer parameters usually imply less computing burden on resource-constrained hardware such as embedding devices or neuromorphic chips.
Some pioneering studies have revealed considerable redundancies in both Artificial Neural
Networks (ANNs) (Han et al., 2015; 2016; Wen et al., 2016; Liu et al., 2017) and Spiking Neural
Networks (SNNs) (Qi et al., 2018; Chen et al., 2021; Yin et al., 2021; Deng et al., 2021; Kundu et al., 2021; Kim et al., 2022b).
In essence, pruning is an optimization problem under L 0 norm constraint on network parameters, which has been proved in Natarajan (1995) to be an NP-Hard problem.
Relaxing L 0 norm to L p (0 < p < 1) norm will not make it more tractable for it is still strongly NP-Hard (Ge et al., 2011).
Nowadays, research on pruning and sparse optimization is mainly focused on the L 1 -regularized problem, the tightest convex relaxation of L 0 norm, which dates back to a series of groundbreaking studies on compressed sensing (Donoho, 2006; Cand`es et al., 2006).
It technically allows us to solve
L 1 -regularized problem as an alternative or, sometimes even an equivalent option (Cand`es, 2008) to confront L 0 norm constraint.
A variety of modern methods such as magnitude-based pruning are still firmly rooted in solving the L 1 regularized optimization problem.
Be that as it may, L 1 regularization is mostly employed for shrinking the magnitude of weight before the hard thresholding step, which has started to be replaced by other sorts of novel regularization (Zhuang et al., 2020).
In the past few years, a new range of pruning methods based on soft threshold reparameterization of weights has been developing gradually.
Two representative ones are S oft T hreshold weight
R eparameterization ( STR ) (Kusupati et al., 2020) and S tate T ransition of D endritic S pines ( STDS )
(Chen et al., 2022).
These works adopt soft threshold mapping to map a set of hidden weights to the actual network weights.
They both achieve the best performance of that time.
STDS further demonstrates the analogy between soft threshold mapping and a structure in biological neural systems, i.e. , dendritic filopodia and mature dendritic spines.
However, few researchers notice that soft threshold mapping is inherently equivalent to the shrinkage operator in the orthonormal design case of
LASSO (Tibshirani, 1996), an L 1 -constrained optimization problem.
It further derives the Iterative
Shrinkage-Thresholding Algorithm (ISTA) (Daubechies et al., 2004; Elad, 2006), which used to be popularized in sparse recovery and compressed sensing.
ISTA has many variants (Bioucas-Dias &
Figueiredo, 2007; Beck & Teboulle, 2009b; Bayram & Selesnick, 2010) and has long been certified as an effective sparsification methods in all sorts of fields like deep learning (He et al., 2017; Zhang et al., 2018; Bai et al., 2020), computer vision (Beck & Teboulle, 2009a; Dong et al., 2013), medical imageology (Lustig et al., 2007; Otazo et al., 2015) and geophysics (Herrmann & Hennenfent, 2008).
Despite STDS undertaking an abecedarian analysis on the scheduling of threshold and providing a preliminary result that the threshold scheduler acts similarly to the iterative step of ISTA, many issues have not been fully addressed in previous soft threshold pruning literature, such as 1) the exact correspondence between ISTA and the manner of growing threshold, 2) the rationality of training threshold, 3) the underlying optimization problem, and 4) the way to improve existing methods without exhaustively trying different tricks for scheduling threshold.
In this work, we proposed a theoretical framework serving as a bridge between the underlying L 1regularized optimization problem and threshold scheduling by viewing soft threshold pruning as an implicit executed ISTA.
Specifically, we prove that the magnitude of the underlying L 1 regularized term is precisely equal to the quotient of the forward finite difference of threshold scheduler function and the learning rate.
To ease the burden, we use the derivative to replace the computation of finite differences and confirm it is a perfect substitute from the perspectives of sparsity and accuracy.
In this way, any kind of threshold tuning strategy is interpreted as a special scheme for changing Lpenalty.
Additionally, we induce a strategy of tuning L 1 penalty called continuation strategy (Xiao & Zhang, 2012), which was once all the rage in the field of sparse optimization, and derive broad categories of algorithms covering several tracks in the present taxonomy of pruning.
In summary, our contributions are as follows:
• Theoretical cornerstone of threshold tuning strategy.
To the best of our knowledge, this is the first work interpreting increasing threshold as an ever-changing regularized term.
Under theoretical analysis, we, for the first time, present a unified framework for the local equivalence of ISTA and soft threshold pruning using SGD as an optimizer.
This formulation enables us to break precursory works into parts and make a comprehensive study on threshold scheduling alone using the classical method in sparse optimization.
• Learning rate adapted threshold scheduler.
Through our proposed framework, we reveal the strong relation between the learning rate scheduler and the threshold scheduler.
Then we raise the viewpoint that a stable regularized term implies a coherent pace of the foregoing two schedulers.
We, therefore, proposed the Learning rate Adapted Threshold Scheduler (LATS), which achieves a state-of-the-art performance-sparsity tradeoff on both deep ANNs and SNNs.
• Sibling schedulers cover multiple tracks of pruning.
We propose an early pruning algorithm by translating the homotopy continuation algorithm into a pruning algorithm with our framework.
It achieves pretty close performance to LATS when used as a conventional early pruning method.
Moreover, the algorithm in the pruning-at-initialization setting erases some subsequent layers in ResNet and maintains the identity mapping, shrinking deep ResNet to a shallow one
R ELATED WORKS
There has been a deluge of pruning algorithms emerged since the term “deep compression” was invented.
These various studies emphasize different points like granularity (structured or unstructured) and stage of pruning (at initialization, during training, post training).
The difference in granularity is similar to that between LASSO and group LASSO.
Empirically, unstructured pruning tends to reach higher sparsity under the same accuracy degradation.
For the pruning phase, sparsify during training commonly lead to higher accuracy than early phase one, e.g. , pruning at initialization.
Moreover, pruning during training is cheaper than post-training pruning when the overhead of dense training is considered.
Some most relevant works are introduced as follows.
Sparsify during training.
The terminology sparsify-during-training (also called pruning-whilelearning, pruning-during-training) is mentioned in Hoefler et al.
(2021), which refers to pruning and training networks simultaneously including those iteratively pruned networks.
Recent works in this area includes STR (Kusupati et al., 2020), Top-KAST (Jayakumar et al., 2020), CS (Continuous
Sparsification) (Savarese et al., 2020), RigL
(Evci et al., 2020), WoodFisher (Singh & Alistarh, 2020), PSGD (Kim et al., 2020), GraNet (Liu et al., 2021a), Powerprop (Schwarz et al., 2021),
ProbMask (Zhou et al., 2021), GPO (Wang et al., 2022), OptG (Zhang et al., 2022) and STDS (Chen et al., 2022).
Many of these works are based on the reparameterization of weights using either binary mask or element-wise nonlinear mapping.
The former choose to confront L 0 constraint directly while the latter are committed to adjusting the landscape of loss function around zero.
Our method is based on soft threshold reparameterization, which is piecewise linear and has an intrinsic connection to the ISTA with L 1 regularization.
Early pruning.
We refer to a variant of sparsify-during-training as early pruning here, which only exerts pruning to network in the early stage of training.
It includes pruning at initialization, e.g. ,
GraSP (Wang et al., 2020), SynFlow (Tanaka et al., 2020), SBP-SR (Hayou et al., 2021), ProsPr (Alizadeh et al., 2022), and the conventional early pruning methods which stop pruning at a certain time (You et al., 2020; Liu et al., 2021b; Rachwan et al., 2022).
Most of these works are inspired by the discovery of Lottery Ticket Hypothesis (LTH) (Frankle & Carbin, 2019) or SNIP (Lee et al., 2019), if not both.
LTH suggests we can find a sparse subnetwork with comparable performance to the dense network after iterative retraining, while SNIP manages to train from an initial sparse network.
A wide array of criteria like connectivity sensitivity are taken for finding such sparse networks in the early stage with promising accuracy after training.
P RELIMINARIES
Notation.
We use | x | to denote the element-wise absolute value of x , and ∥ x ∥ p denotes the pnorm of x . x ⊙ y denotes the element-wise product of x and y .
If not otherwise specified, the superscript within parenthesis x ( i ) denotes x at i -th iteration.
The element-wise sigmoid function is denoted by σ ( x ) i :
= 1 / (1 + e − x i ) .
The soft threshold mapping is defined by the element-wise mapping S d ( x ) i :=
sign( x i ) · max {| x i | − d, 0 } with scalar threshold d .
3. S OFT THRESHOLD PRUNING
Basically, soft threshold pruning is a group of methods including two core steps:
Persistently parameterize network weight w through the soft threshold mapping w ( t ) ← S d ( θ ( t ) ) during training, where θ is the trainable hidden weight with the same shape as w . 2.
Growing threshold
d pushes the term max {| θ i | − d, 0 } in S d ( θ ) towards zero and thereby induces sparsity for w .
Generally, STR and STDS are built upon the above two designs.
A general form of soft threshold pruning is formulated as Algorithm 1.
In GPO, a convex combination of soft threshold mapping and identity mapping is introduced for reparameterization of weights.
Since GPO marginally improves performance with respect to STR, and will soon degenerate to STR as discussed in Appendix C, the following parts of this paper are focused on STR and STDS.
The differences between them are concluded from two aspects.
Gradient computing.
With the above reparameterization mapping, the forward step includes mapping via θ → w and evaluating loss with w as the actual network weight.
Hence, gradient is backpropagated via path L → w → θ .
The learning rule is thus used for updating θ instead of
←
Initialization 2:
3: ∆ θ t
← ∇ w ( L ( w t
)) ⊙
> d t ) weight 4: θ
+1)
← θ t
− η t ) (∆ θ ( t ) + λ θ t )
▷ weight 5: s
← s t
− η t )
∇ s w s
+ λ s t
6: d ( t
← σ g t
▷ Update threshold 7: w t
← S d θ w .
Note that S d is non-differentiable at ± d and has zero gradient in interval ( − d, d ) .
STR takes advantage of the subgradient at ± d and leaves zero gradients as it is.
STDS views S d as an identity mapping during backward and provides a convergence analysis by approximating S d using a smooth surrogate ˆ S d ( x )
= 1 α [ ζ ( α ( x − d )) − ζ ( − α ( x + d ))] , where ζ ( x ) :=
log(1+ e x ) denotes the softplus function.
Threshold tuning.
To enhance sparsity during training.
Recent studies focused on manipulating d throughout learning.
STR further parameterizes the threshold by another parameter s and assigns an independent threshold for each layer.
The mapping from s to threshold is individually designed for CNN and RNN.
It also set s trainable to train both the threshold and weights during learning .
STDS introduces the threshold scheduler as d ( t ) =
g ( t/T ) · D , wherein scheduler function g : [0 , 1] → [0 , 1] is increasing and satisfies g (0)
= 0 , g (1)
= 1 , T is the total training steps.
The final threshold D is the only adjustable hyperparameter for different sparsity levels when g is given.
Larger D always leads to higher sparsity in practice.
The above techniques are summarized in Table
[Table] [Table]
3. I TERATIVE SHRINKAGE - THRESHOLDING ALGORITHM
Consider a regularized optimization problem
[Equation]
) with assumptions listed below
Objective function f : R n → R is continuous differentiable, and L -smooth, i.e. , ∥∇ f
( x
−
∇ f ( y ) ∥ 2 ≤ L f ∥ x − y ∥ 2 , ∀ x ,
y ∈ R n , where L f > 0 is the Lipschitz constant of ∇ f . 2.
Regularization function r : R n → R is continuous convex and can be nonsmooth.
F
is bounded from below.
Leave the objective alone, applying vanilla SGD to f can be viewed as iteratively calculating proximal regularization (Martinet, 1970) of the linearized f at x .
In other words, for a given point x ,
F ( y ) can be approximated by expanding f to the quadratic term
[Equation]
where η is explained as “stepsize” in optimization or “learning rate” in the context of deep learning.
The above problem has a unique minimizer
[Equation]
Note that η = 1 L f gives a upperbound of f ( y ) , which is obtained by a property of L f -smoothness: f ( y ) ≤ f ( x )+ ⟨ y − x , ∇ f ( x ) ⟩ + L f 2 ∥
y − x ∥ 22 .
It implies with proper choice of η , we can optimize the upperbound of F ( y ) using minimizer arg min
y ˆ F η ( y ; x ) .
The general form of ISTA iteratively solves Eq.
3 as
[Equation]
which is also known as proximal gradient methods (Combettes & Wajs, 2005).
The detailed convergence analysis is presented in vast optimization literature like FISTA (Beck & Teboulle, 2009b) and GIST
(Gong et al., 2013).
For sparse learning, we are interested in L 1 regularization term r ( x ) = µ ∥ x ∥ 1 , µ > 0 .
Since L 1 norm is separable, Eq. 4 has a closed-form solution as
[Equation]
Eq.
5 gives the ISTA update rule under L 1 regularization.
A FRAMEWORK FOR SOFT THRESHOLD PRUNING
In this part, we will formulate the growing threshold under soft threshold reparameterization as an implicit ISTA.
For simplicity, we assume the threshold is global across all parameters, which is consistent with the setting of STDS and STR-GS, i.e. , the global threshold version of STR.
Hence, we use scalar d rather than vector d in Algorithm 1 to denote the global threshold.
To begin with, we investigate the update rule of nonzero components in actual weight
[Equation]
Assuming the sign of weight remains unchanged after an update, which happens when the gradient has the opposite sign of weight or the gradient magnitude is sufficiently small, we have the following
Lemma:
Lemma 1 (Local update rule) .
The update rule in Eq.
6 and Eq.
7 implies the following update of a nonzero component w ( θ, d ) in actual weight w
[Equation]
The formal version and proof of Lemma 1 is given in Appendix A. Note that the form in Eq.
8 is equivalent to the ISTA update rule with threshold equal to forward finite difference d
( t +1) − d ( t ) .
Recall Eq.
5, the corresponding optimization problem can be deduced as Theorem 1:
Theorem 1.
Let L ( w ) be the loss function depending on the network weight w , which is further reparamterized by hidden weight θ and threshold d as w ( θ , d ) =
S d ( θ ) .
When applying vanilla
SGD to reparameterized nonzero weight, the update rule is locally equivalent to solving the following problem using ISTA with penalty term min
w
[Equation]
where η ( t ) , d ( t ) denotes the learning rate and threshold at the t -th iteration, respectively.
The magnitude of L 1 -penalty is the quotient of forward finite difference of threshold scheduler divided by learning rate.
Theorem 1 serves as the glue holding the learning rate scheduler η ( t ) , threshold d ( t ) and penalty µ ( t ) together.
It also to some extent explains an increasing threshold is necessary for a positive penalty term.
Under this framework, we explain in Appendix B that the original STDS uses an improper threshold scheduler.
Furthermore, an analysis of training threshold is given in Appendix C.
F INDING OPTIMAL THRESHOLD SCHEDULER
In this part, we are devoted to evaluating some threshold schedulers based on our theoretical framework and literature on sparse optimization.
With the framework, any past strategy of tuning Lpenalty can be converted to a feasible pruning algorithm today.
5. L EARNING RATE ADAPTED THRESHOLD SCHEDULER
If the optimization problem in Theorem 1 is fixed, or in other words, the L 1 penalty is invariant during learning, we can deduce a unique threshold scheduler LATS .
In LATS, the change d ( t +1) − d ( t ) on threshold must keep paces with the learning rate η ( t ) to obtain stable regularized terms.
Corollary 1 ( L earning rate A dapted T hreshold S cheduler) .
For fixed L 1 regularized problem
[Equation]
where µ > 0 is the L 1 penalty coefficient, the threshold scheduler is governed by the learning rate scheduler, which is
[Equation]
For the most frequently used testbed, ResNet-50 (He et al., 2016) on ImageNet dataset (Deng et al., 2009), researchers usually use the cosine annealing learning rate scheduler with η min = 0 .
Assuming the initial threshold is zero, simple algebra shows the definition of corresponding LATS with a slight abuse of notation
[Equation]
where the threshold d ( n,b ) depends on epoch id n and batch id b .
Here n = 0 , 1 , . .
. , N − 1 denotes the current id of training epoch, and b = 1 , 2 , . . .
, B denotes the batch id in n -th epoch.
In effect, painstakingly coding LATS according to a learning rate scheduler is not inevitable.
To ease the computing burden and erase the accumulative error of learning rate, we turn to a simplified threshold scheduler by replacing the summation of learning rate with integration and increasing the update frequency to per-iteration.
Assuming the learning rate scheduler can be expressed by η ( n,b )
= h ( n/N ) , the simplified threshold scheduler is defined by
[Equation] [Equation]
Now we have the Simplified LATS ( S-LATS for short)
[Equation] [Equation]
It can be interpreted with scheduler form in STDS as d ( t ) =
g ( t/T ) · D , where g S-LATS ( x ) = sin( πx ) /π
+ x and D = µη max T/ 2 .
The final threshold D ∝ µ controls the overall strength of penalty.
Tuning D is thus akin to changing the magnitude of penalty.
In the following discussion, we employ the final threshold D instead of d ( N − 1 ,B ) to lighten the notation.
The threshold schedulers in the rest of this work will be expressed in a unified form of d ( t ) =
g ( t/T ) · D .
S-LATS has almost the same accuracy and sparsity performance as the original LATS under identical final thresholds D = 0 .
5 , 1 . 0 , 5 . 0 , which can be gleaned from Tab. 2.
Thus, we turn to a simplified threshold scheduler in the following discussion.
Detailed derivations in the above section are given in Appendix D.
[Table] [Table]
C ONTINUATION STRATEGY
Also known as “warm starting”, continuation strategy is designated for accelerating convergence.
Similar to annealing of learning rate, continuation refers to gradually reducing L 1 penalty during learning.
It is also explained in Hale et al.
as an analog to the homotopy algorithms in statistics.
Continuation method used to serve as a common trick in abundant classic literature concerning sparse optimization including GPSR (Figueiredo et al., 2007), fixed point continuation method (Hale et al., 2008), SpaRSA (Wright et al., 2009) and NESTA (Becker et al., 2011).
5.2. PGH SCHEDULER
In the series works of proximal gradient homotopy (PGH) (Xiao & Zhang, 2012; 2013; Lin & Xiao, 2014), the researchers provide proof of geometric convergence rate when inducing exponentially decaying regularization µ ( t ) =
β tT , where 0 < β < 1 .
Considering our formulation of soft threshold pruning in Theorem 1, PGH can be translated into the PGH scheduler (simplified using Eq.13), which can be written as
[Equation]
where the analytic form of g PGH is shown below
[Equation] [Equation]
5.2. L INK TO EARLY PRUNING
[Figure]
The union of scheduler functions are shown in blue area.
For the PGH scheduler, we interpolate β between 0 and 1 and get a series of PGH schedulers.
As shown in Fig.
1, the increasing in threshold slows over time, which is caused by decreasing regularization term.
For 0 < β < 1 , if we ignore pruning when regularization is below a preset threshold, we get a family of early pruning algorithms.
They stop pruning at different stages.
Unlike the stereotyped early pruning algorithm using static masks after the pruning stage, PGH scheduler uses a dynamic mask, enabling “dead” weights to be revived when the pruning stage ends.
It indicates that the sparsity may drop in the sheer training stage.
Recall that the regularized term is proportional to the forward finite difference and thus can be approximated by derivative, for conventional early pruning, we regard g ′ PGH ( t/T ) < 0 .
1 as the termination criterion of pruning.
Note that there are two limit cases when β approaches 0 or 1.
It is obvious that β → 1 leads to S-LATS for no decay is applied to regularization.
When β → 0 , the regularization is always zero.
In this case, PGH scheduler degenerates to the magnitude-based pruning after initialization followed by a normal training stage.
This is also referred to as sparse-tosparse training or pruning at initialization method.
E XPERIMENTS
In this section, we test our proposed threshold scheduler S-LATS on both deep ANNs and SNNs.
The favorable performances against the previous studies are confirmed.
We also tune hyperparameter β in PGH scheduler to maintain different phases of pruning algorithm.
Given the form of scheduler function, the diverse sparsity levels are tuned by switching the only hyperparameter D , which is directly associated with L 1 regularized term µ as discussed in section 5.1.
6. S-LATS
We test STDS integrated with S-LATS on various sparsity scenarios via tuning final threshold D .
It achieves state-of-the-art performances on both ANNs (ResNet-50, MobileNet-V1) and SNNs (SEW
ResNet-18 (Fang et al., 2021)).
The results on ResNet-like networks and MobileNet-V1 are illustrated in Fig.
2 and Appendix E, respectively.
In particular, it outperforms the default threshold scheduler of STDS, namely the sine scheduler, on SEW ResNet.
An ablation study on threshold scheduler is shown in Appendix F. Notably, our method surges ahead of all the other baselines under < 98% sparsity for ANN experiments.
It should be noted that the origin STDS excludes the last
FC layer from pruning in SNNs, while the results reported here are shown by rerunning it with the last layer pruned.
We also admit it fails to achieve comparable performance with a few baselines like OptG and in extreme sparsity ( ≥ 99%), which suggests the theory might be imperfect under such conditions.
[Figure] [Figure]
[Table]
Pruning using large batch size.
Some recent studies such as Powerprop (Schwarz et al., 2021) adopt a batch size of 4096 for pruning of ResNet-50, which is much larger than the standard training setting (256 batch size).
They both achieve fascinating performances under high sparsity.
Hence
[Table]
we verify the influence of large batch size.
Due to limited computing resources, we only increase it to 1024 and also enlarge the learning rate correspondingly.
The right of Fig.
2 shows the enlarged batch size indeed leads to higher performance for STDS + S-LATS, which outperforms all of SOTA studies.
Astonishingly, we observe that even though the performance of the dense network slightly degrades from 77.1 to 76.78, the accuracy of the sparse ones is improved overall.
Therefore, we attribute such performance improvement to the higher affinity to pruning when using a large batch size.
On the basis of the above finding, we believe applying an even larger batch size, like 4096 used in Powerprop, to our method will lead to top-notch accuracy under all levels of sparsity.
6. PGH SCHEDULER
Conventional early pruning Our experiments of early pruning includes β = 0 .
1 , 10 − 5 , 10 − 10 .
The corresponding ending criteria t/T = 0 . 743 , 0 .
382 , 0 . 231 are given by numerical solution.
It indicates pruning roughly stops at the 74th, the 38th and the 23th epoch.
The network using PGH scheduler with smaller β converges faster to a sparse one, which is illustrated in the left of Fig.
Surprisingly, we find in the middle of Fig.
3 that for different β , the datapoint of accuracy against
[Figure]
sparsity almost lies on the curve of S-LATS.
It suggests these schedulers have practically the same performances as S-LATS, but with faster convergence to sparse networks.
With the help of PGH scheduler, we are able to find sparse networks earlier with negligible performance degradation.
Pruning at initialization We also try β → 0 , which refers to increasing the threshold to its maximum at the first iteration.
Since our method is agnostic about the object of pruning, some layers are completely pruned as shown in the right of Fig.
3, wherein three consecutive layers within a residual block tend to be pruned simultaneously.
However, owing to the skip connection in ResNet, the feature can pass through shortcuts to the final FC layer, and thus most layers are still normally trained.
The aforementioned results are collected in Tab. 4.
[Table]
C ONCLUSION & D ISCUSSION
In this work, we present a framework interpreting increasing threshold as a constantly changing penalty term and reveal the underlying connection between soft threshold pruning and ISTA.
We also derive a couple of threshold schedulers, which achieve comparable performance to current
SOTA works and cover multiple tracks of pruning.
It is worth noting that our method is agnostic about the object of pruning.
This design endows our method with versatility while treating weight wheresoever equally, becoming totally ignorant of nowadays pruning researches like sparsity budget allocation, e.g. , Erd˝os-R´enyi (Mocanu et al., 2018) and Erd˝os-R´enyi-Kernel (ERK)
(Evci et al., 2020), or commonsense in this area like “Never prune the whole layer”.
We believe our method can for sure further benefit from knowledge in the prosperous field of network pruning.
A P ROOF OF THEOREMS AND LEMMAS
For clarity, we restate the theorem or lemma in the main text again here.
Lemma 1 (Local update rule) .
The update rule below
[Equation]
imply the following update of any nonzero component w ( θ, d ) in actual weight w
[Equation]
19) when | θ ( t +1) | > d ( t ) and the sign condition sign( θ ( t +1) )
=
sign( θ ( t ) ) are met.
Proof.
For any nonzero weight w ( t ) ̸ = 0 , θ ( t ) =
w ( t ) + d ( t ) sign( w ( t ) ) , using Eq.
17 we have
[Equation]
Let ¯ w ( t +1) :
= w ( t ) − η ( t ) ∇ w ( L ( w ( t ) )) be the target point of vanilla SGD without regularization.
Recall Eq.
20, we have
[Equation]
−
[Equation] [Equation]
21) which has the same sign as θ ( t +1) .
Now we have sign( w ( t +1) )
=
sign( θ ( t +1) )
=
sign( θ ( t ) )
=
sign( w ( t ) )
=
sign( ¯ w ( t +1) ) .
To evaluate the updated weight, by Eq.
18, Eq. 20, Eq. 35 and the definition of soft threshold mapping, we derive
[Equation]
B O RIGINAL THRESHOLD SCHEDULER IN STDS
In STDS, the authors propose the Sine scheduler d ( t ) = 12 (1+sin( π ( tT − 12 )))
D = 12 (1 − cos( tπT )) D .
With the form of simplified threshold scheduler in Eq. 13, by Theorem 1, the corresponding penalty µ ( t ) has the form
[Equation]
It is a function of training progress t/T with constant C .
Investigate the function tan
x/ 2) on interval (0 , π ) , we have µ ( x ) is increasing from 0 to + ∞ , which implies µ can be sufficiently large during training.
The loss is thus insignificant compared to the regularization term in the last stage of training and leads to performance degradation with respect to S-LATS.
C D ISCUSSION ABOUT TRAINING THRESHOLD
In the main text, we propose a framework explaining the style of growth in threshold as an everchanging optimization problem.
However, we only cover manually designed threshold schedulers.
We make a discussion here and show training threshold is not as easy as STR or GPO did.
Note that
GPO and STR share the same discussion of training threshold since GPO will degenerate to STR in a few training epochs.
Moreover, we suggest not simply setting threshold trainable if one really needs to investigate threshold tuning based on optimization of threshold.
C. L 2 PENALTY DOMINATES EARLY TRAINING OF STR.
In the official codebase of STR, we notice the trainable sparse threshold is also together with weight decay ( L 2 regularization λ
∥ s ∥ 22 ).
λ is of magnitude around 10 − 5 ∼ 10 − 4 .
The initial value s init is usually set to negative number with large magnitude around − 10 4 ∼ − 10 3 for CNN trials.
It is easy to see the magnitude of L 2 regularized term is around 0 .
01 ∼ 1 in the early stage of training.
Take CNN for instance.
For given loss function L and weight w l of the l -th layer, the gradient passed to threshold s l can be estimated as
[Equation]
wherein each term of sum has magnitude less than 14 ∇ w L ( w l ) since sigmoid function σ ( x ) is bounded in (0 , 1) .
Given that the stochastic gradient noise across parameters in w l admits to L ´ evy distribution (Xie et al., 2021), some negative and positive terms will balance each other, which makes us wonder whether the regularization term dominates the training of threshold.
Therefore, we compare the gradient to L 2 penalty by tracing the magnitude ratio of ∇ s l L to λ | s l | during training of
ResNet-50 on ImageNet, which is shown in Fig. 4.
[Figure]
It is evident that the L 2 penalty plays the leading role in the early stage rather than the gradient.
We even observe for several final layers, the penalty always dominates the training of threshold.
The update of s can thereby be rewritten to s ( t +1)
≈ s ( t ) (1 − η ( t ) λ ) , where η ( t ) is the learning rate at the t -th iteration.
Regardless of gradient, at the beginning of training, STR can be viewed as a special case of threshold scheduler as follows
[Equation]
By Theorem 1, the corresponding penalty can be derived by
[Equation]
where ¯ s lies between s ( t +1) and s ( t )
The existence of ¯ s is shown by Lagrange’s Mean Value
Theorem.
It is rather difficult to analytically give the explicit expression of µ ( t ) since it relies on the behavior of s , which has a dynamic decay rate 1 − η ( t ) λ .
Loosely speaking, we adopt the approximation ¯ s ≈ s ( t ) .
We still cannot show the explicit form of µ ( t ) , but now we can investigate the trend of penalty by analyzing function µ ( x ) = − σ ′ ( x ) x , which is increasing in ( −∞ , γ ) .
Here γ
≈ − 1 .
5434 is the unique negative root of e − x ( x +1)
− x +1 = 0 .
In the early stage, s is a negative number with a much larger magnitude than γ .
Based on the above, µ is increasing in the early stage.
C. GPO: FALL BACK TO STR SHORTLY
It concludes STR is majorly influenced by L 2 decay on s .
We will see that GPO has a similar behavior on k and GPO will degenerate to STR shortly after the training begins.
In GPO, the authors introduces another trainable parameter β in the reparameterization w = (1 − k ) S d ( θ ) + k θ with k = 10 − 6 | β | .
GPO starts from the identity mapping w = θ by initializing β
= 10 6 .
The gradient passed to β l in the l -th layer is given by
[Equation]
Notice that θ − S d ( θ ) =
(cid:26) θ, | θ | <
d d sign( θ ) , | θ | ≥ d has magnitude not greater than d , which derives that the gradient ∇ w L ( w l ) are added with coefficient whose magnitude is below 10 − 6 d .
From the codebase of GPO, we confirm the weight decay on β is 10 − 4 .
Since the initial β is 10 6 , the
L 2 regularized term is of magnitude around 10 − 4 × 10 6 = 100 at the beginning of training.
Recall that d = σ ( s ) < 1 , it is obvious that the gradient pass to β has a much smaller magnitude than the
L 2 regularization.
For this reason, the gradient can be ignored for k .
Furthermore, k will shrink exponentially to almost zero and the mapping in GPO will fall back to STR within a few epochs, which can be seen by w = (1 − k ) S d ( θ ) + k θ ≈
S d ( θ ) .
C. F OCUS ON THRESHOLD SCHEDULER INSTEAD OF THE FINAL THRESHOLD .
Conventional wisdom suggests that when a parameter is set trainable, it will be optimized automatically.
This idea should only work for those directly determining the performance, e.g. , weights in a dense network.
For pruning, authors of STDS find the differences in positions of performance versus sparsity curve should be ascribed to the evolving patterns of the threshold.
To verify this, we replace the threshold training mechanism in STR and default scheduler in STDS by several schedulers with the same final threshold D = 10 − 3 ( D = 0 .
5 for STDS) shown below
[Equation]
[Equation]
• Linear: d ( t ) =
tT D • Log2: d ( t ) = log 2 ( tT + 1) D
We conduct the training on ANN ResNet-50 for both methods, the results are shown in Tab. 5.
It is
[Table]
obvious that even though the final thresholds are set equally, the accuracy and overall sparsity vary with the scheduler.
The above results suggest the correct manner to manipulate threshold is pursuing a well-performed scheduler.
Even if one studies the learning of threshold, it should be concentrated on the optimization of threshold scheduler, which may require tools in discrete-time optimal control.
Due to the complex coupling between constantly changing threshold and final performance, discussion based on discrete-time optimal control is beyond the scope of this work.
D D ETAILED DERIVATION FOR LATS AND S-LATS
In this part, we provide detailed derivations for LATS and S-LATS, which covers Eq.
12 and Eq. 13.
D. LATS FOR COSINE ANNEALING SCHEDULER
In most of the deep learning applications, the training process includes the schedule of the learning rate, which is also known as learning rate scheduler .
Generally, the learning rate is updated at the end of an epoch.
Assuming there are N training epochs in total, each of which includes B training mini-batches.
The learning rate scheduler is defined as η ( n,b )
=
h ( n/N ) , which evaluates the learning rate at the b -th mini-batch in the n -th epoch.
Here n = 0 , 1 , . .
. , N − 1 denotes the current id of the training epoch, and b = 1 , 2 , . . .
, B denotes the batch id in n -th epoch.
We denote h : [0 , 1] → R + as the scheduler function for the learning rate.
For cosine annealing scheduler (Loshchilov & Hutter, 2017) with η min = 0 , we have
[Equation]
In Corollary 1, the threshold scheduler for LATS is obtained by d ( t ) =
d (0) + µ (cid:80) t − 1 i =
η ( i ) , where d ( t ) is the threshold after t mini-batches from the beginning.
Under the learning rate scheduler described in Eq.
28, the threshold d ( n,b ) is shown by accumulating all previous learning rates, which gives
[Equation]
To evaluate the sum of cos( iπ/N ) , we give the following results
[Equation]
Recall Eq. 29, we have
[Equation]
When d (0 , 0) = 0 , this gives LATS in the form of Eq.
D. T HE DETAILED MOTIVATION OF S-LATS
The implementation of LATS is rather complicated and requires meticulous coding.
Worse still, for some learning rate schedulers, e.g. , polynomial decay scheduler (Liu et al., 2015; Chen et al., 2018), the form of LATS cannot be reduced like cosine annealing scheduler in most cases, which means we may not get rid of the summation and it thus leads to accumulative error.
In brief, we cannot expect each learning rate scheduler corresponds to an analytical form of LATS without summation notation.
Returning to the first equation of Eq. 29, the form B (cid:80) n − 1 i =
h ( i/N ) is similar to the approximation of the Riemann integral BN (cid:82) ( n − 1) /Nh
( x )d x using rectangle methods.
Therefore, we replace the summation with integration while keeping the final threshold d ( T ) =
d ( N − 1 ,B ) by using a normalization term.
This lead to a simplified form of LATS, that is
[Equation]
E P RUNING EXPERIMENTS ON M OBILE N ET -V1 USING S-LATS
Besides the ResNet-like structures mentioned in the main text, we also conduct experiments on
MobileNet-V1 (Howard et al., 2017) to show the power of our proposed methods S-LATS on the lightweight network.
To make a fair comparison, we choose those SOTA studies using the standard training setting, i.e. , batch size of 256 and 100 training epochs.
They include STR (Kusupati et al., 2020), gradual pruning in WoodFisher (Singh & Alistarh, 2020), and a modern implementation (Gale et al., 2019) of Gradual Magnitude Pruning (GMP) (Zhu & Gupta, 2017).
Note that we do not compare S-LATS to OptG, since OptG adopts 1 . 8 × (180) training epochs, or the comparison would be unfair.
The results are shown in Tab. 6 and Fig.
Apparently, our proposed method towers over previous sparsify-during-training work.
We elaborate on the sparsity budgets and the corresponding final thresholds in Tab.
12 of Appendix I.
The hyperparameters for training MobileNet-V1 are stated in
Tab. 8 of Appendix H.
[Table]
[Figure]
F A BLATION STUDY OF THRESHOLD SCHEDULER ON R ES N ETF. R EMOVE L 2 DECAY FOR FAIR COMPARISON
To evaluate the performance gains brought by the threshold scheduler alone, the weight decay or Lpenalty must be removed.
To explain this, recall line 4 in Algorithm 1, we know the update rule of hidden weight θ is affected by both gradient and L 2 penalty λ
∥ θ ∥ 2 .
However, the analysis in
Lemma 1 is based on vanilla SGD without weight decay.
To account for this inconsistency, let’s first investigate the influence of weight decay on the equivalent optimization problem.
In the presence of weight decay, the update rule described by Eq.
17 has an additional penalty term as follows
[Equation]
Following derivation in Eq.
20, we have
[Equation]
[Equation]
Similarly, we denote ¯ w ( t +1) :
= w ( t ) − η ( t ) ∇ w ( L ( w ( t ) )) to be the target point of vanilla SGD without regularization.
With Eq.
34, we have
[Equation]
If ¯ w ( t +1) still satisfies the local relation that sign( w ( t +1) )
=
sign( θ ( t +1) )
=
sign( θ ( t ) )
=
sign( w ( t ) )
= sign( ¯ w ( t +1) ) ,
mimicking Eq.
22, we have
[Equation]
Apparently, the L 2 penalty of θ lies into the equivalent L 1 penalty term of w in the ISTA rule, making the analysis of the corresponding threshold scheduler intractable.
Accordingly, we decide to remove weight decay, i.e. , and set λ to zero in the ablation study to prevent an unpredictable threshold scheduler.
F. S INE SCHEDULER IN STDS VS S-LATS
After removing weight decay, we rerun the original STDS and our methods under several sparsity levels (through changing D ) while keeping the other hyperparameters and the batch size of 256.
As illustrated in Fig.
6, the results on ResNet-50 show our method clearly surpasses the original STDS on the ImageNet dataset.
A theoretical analysis is enclosed in Appendix B.
[Figure]
G S PARSITY VS FIRING RATE IN SNN S G. A N OVERVIEW OF SNN S
Spiking neural networks (SNNs) are honored as the third generation of neural network models (Maass, 1997), derived from biological neural network modeling.
SNNs are composed of spiking neurons, which release spikes in binary form, and connections between neurons.
The model of spiking neurons is a dynamical system described by one or more ordinary differential equations (ODE) and a firing threshold.
The dynamical system is also called “subthreshold dynamics” in the context of computational neuroscience.
A spike is generated and passed to all postsynaptic spiking neurons when the variable representing membrane potential exceeds the firing threshold.
Today, the most commonly used neuron model is the Leaky Integrate-and-Fire (LIF) model.
Specifically, LIF has the subthreshold dynamic as follows
[Equation]
where u ( t ) is the membrane potential at time t , u rest is the resting potential, τ m is the membrane constant, I and w denote input spikes and input weights respectively.
The firing behavior of LIF neurons is depicted as an instantaneous jump of membrane potential shown below
[Equation]
where u th , t f are the firing threshold and firing time respectively.
The ODE in Eq.37 can be discretized via the Euler method and transformed into an RNN-like iterative computing manner as follows
[Equation]
where u [ t − ] , u [ t ] are the membrane potential before and after firing at timestep t respectively, H ( · ) is the Heaviside step function modeling jump behavior when a spike is triggered.
However, training techniques in RNN, such as backpropagation through time (BPTT) (Werbos, 1990) cannot be directly applied to SNNs for the spiking behavior described by Heaviside is non-differentiable.
Thanks to the surrogate gradient method proposed in Wu et al.
; Neftci et al.
(2019), researchers can now incorporate BPTT into training of SNNs by switching to a “differentiable mode” of the Heaviside step function when computing gradient.
It refers to replacing Heaviside step with a differentiable surrogate function.
Surrogate gradient resembles the straight-through estimator (Bengio et al., 2013) closely in both computing style and ideology.
G. R EDUCING SNN S COST ON NEUROMORPHIC HARDWARE
SNNs are considered energy efficient when deployed on a series of dedicated hardware, also known as neuromorphic hardware or event-driven hardware .
On these chips, the computation is triggered only when there are incoming spikes and weights are nonzero (Merolla et al., 2014).
For this reason, there are three mainstream methods for alleviating the energy cost of a given SNN on neuromorphic chips including 1) unstructured pruning of weights, 2) reducing the number of spikes, and 3) searching for efficient SNN structures.
The NAS studies on SNNs (Na et al., 2022; Kim et al., 2022a) are not based on existing SNN structure, so we omit the discussion of NAS methods.
Many recent studies have made ample signs of progress on pruning and reducing spike counts.
They confirm there is a weak correlation between the number of spikes and weight sparsity (Deng et al., 2021;
Chen et al., 2022; Kim et al., 2022b).
We also evaluate the spike counts of pruned SNNs.
Compared to the number of spikes, a more frequently used metric is average firing rate , which is obtained by averaging the number of spikes across timesteps and spiking neurons.
We collect the average firing rate of each trial during inference using pruned SEW ResNet-18.
We further provide a plot of the average firing rate against the sparsity, which is shown in Fig.
The weak relationship
[Figure]
between the number of spikes and weight sparsity is manifested in the slightly decreased average firing rate.
Despite a downward trend in the average firing rate, the relative magnitude of the decline is trifling.
It is consistent with previous observations and suggests pruning is an inefficient means of reducing the number of spikes in SNNs.
In conclusion, pruning is an efficient way to induce weight sparsity and lower cost.
However, we should not expect the suppression of firing rates as a bonus.
H T RAINING HYPERPARAMETERS
We make the detailed setting in our experiments clear in Tab. 7, Tab. 8 and Tab.
[Table] [Table]
for ablation study in Appendix F)
[Table]
[Table] [Table]
Prune first and last layers? -
Yes
Simulation timestepsSEW functionADD
I S PARSITY BUDGETS
[Table] [Table]
[Table] [Table]
[Table] [Table]
model.10. model.10.
model.11. model.11.
model.12. model.12.
model.13. model.13.
fc 92.
[Table]
layer2.1.conv 76.
layer2.1.conv 63.
layer2.2.conv 68.
layer2.2.conv 70.
layer2.2.conv 63. layer2.3.conv 63.
layer2.3.conv 70.
layer2.3.conv 66.
layer3.0.conv 54.
layer3.0.conv 82.
layer3.0.conv 65.
layer3.0.downsample.
layer3.1.conv 78.
layer3.1.conv 82.
layer3.1.conv 69.
layer3.2.conv 74.
layer3.2.conv 81.
layer3.2.conv 68.
layer3.3.conv 66.
layer3.3.conv 78.
layer3.3.conv 65.
layer3.4.conv 57.
layer3.4.conv 79.
layer3.4.conv 62.
layer3.5.conv 53.
layer3.5.conv 79.
layer3.5.conv 58.
layer4.0.conv 90.
layer4.0.conv 99. layer4.0.conv 99.
layer4.0.downsample.
layer4.1.conv 97.
layer4.1.conv 99. layer4.1.conv 99.
layer4.2.conv 98.
layer4.2.conv 99. layer4.2.conv 99.
fc 83.
[Table]
[Table]
layer2.1.conv2.
layer3.0.conv1.
layer3.0.conv2.
layer3.0.downsample.0.
layer3.1.conv1.
layer3.1.conv2.
layer4.0.conv1.
layer4.0.conv2.
layer4.0.downsample.0.
layer4.1.conv1.
layer4.1.conv2.
fc 46.