A U NIFIED F RAMEWORK FOR S OFT T HRESHOLD P RUNING
A BSTRACT
Soft threshold pruning is among the cutting-edge pruning methods with state-ofthe-art performance 1 .
However, previous methods either perform aimless searching on the threshold scheduler or simply set the threshold trainable, 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 prun- ing are concluded as different styles of tuning L 1 -regularization term
We fur - ther derive an optimal threshold scheduler through an in-depth study of threshold scheduling based on our framework.
This scheduler keeps L regularization coef ficient stable, implying a time-invariant objective function from the perspective of optimization. In principle the derived pruning algorithm could sparsify any math ematical model trained via SGD. We conduct extensive experiments and verify its state
-of- the art performance on both Artificial Neural Networks (ResNet 50 and MobileNet V and Spiking Neural Networks (SEW ResNet-
) on ImageNet datasets. On the basis of this framework we derive a family of pruning meth ods including sparsify during training early pruning, and pruning at initializa- tion. The code is available at https://github com/Yanqi-Chen/LATS .
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 embedded 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 can be formulated as an optimization problem under constraint on L 0 norm, the number of nonzero components in network parameters. Assuming L is the loss function of vectorized network weight w , we expect lower L 0 norm ∥ w ∥ 0 along with lower loss L ( w ) . Despite different formulations like hard constraints [Equation]
or soft constraints (penalized)
[Equation] all these forms Nguyen 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).
These researches 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.
The term “reparameterization” here refers to a specific mapping to network weights w from a latent space of hidden parameters θ . The “geometry” of latent space could be designed for guiding actual weights w towards sparsity . In soft threshold pruning, the mapping is an element-wise soft threshold function with time-variant threshold. Among these studies, 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).
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 also appear as the shrinkage operator in the solution of LASSO (Tibshirani, 1996) when the design matrix is orthonormal.
The studies on LASSO 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 an abecedarian analysis on the similarity between STDS and ISTA, many issues remains to be addressed, such as 1) the exact equivalence between ISTA and the growing threshold in soft threshold pruning, 2) the necessity of setting threshold trainable in STR, and 3) 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. The bridge is built upon the key finding that soft threshold pruning is an implicit ISTA for nonzero weights .
Specifically, we prove that the L 1 coefficient in the underlying optimization problem is determined by both threshold and learning rate.
In this way, any threshold tuning strategy can now be interpreted as a scheme for tuning Lpenalty.
We find that a time-invariant L 1 coefficient lead to performance towering over previous pruning studies.
Moreover, we bring 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, to the field of pruning. It derives broad categories of algorithms covering several tracks in the present taxonomy of pruning.
In brief, our contributions are summarized 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 present a unified framework for the local equivalence of ISTA and soft threshold pruning.
It enables us to make a comprehensive study on threshold tuning using the classic 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 show that an time-invariant L 1 coefficient requires the changing of threshold being proportional to the learning rate.
The Learning rate Adapted Threshold Scheduler (LATS) built upon L 1 coefficient 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 indistinguishable performance to LATS 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 after several epochs of training (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 of gradient descent.
The element-wise sigmoid function is denoted by σ ( x ) i :
= 1 / (1 + e − x i ) .
The soft threshold mapping is also an element-wise mapping defined by 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 will iteratively execute following three core steps:
(i) Mapping hidden weight to actual weight w through the soft threshold mapping w ( t ) ← S d ( θ ( t ) ) during training, where θ is a trainable hidden weight with the same shape as w .
(ii) Training hidden weight θ through backpropagation in latent space (iii)
Growing threshold d pushes the term max {| θ i | − d, 0 } in S d ( θ ) towards zero and thereby enforces sparsity for w .
[Table]
We provide a general form of soft threshold pruning as Algorithm 1, which is a prototype of both
STR and STDS.
GPO change the mapping in line 7 to a convex combination of soft threshold and identity mapping, leading to a slight difference.
However, GPO only obtain marginal performance improvement with respect to STR, and will soon degenerate to STR as discussed in Appendix C. Therefore, the following of this paper are focused on STR and STDS.
The differences between them are concluded from two aspects.
Gradient computing.
With the reparameterization mapping w ( t ) ← S d ( θ ( t ) ) , 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 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.
These studies spontaneously try manipulating threshold d for different sparsity.
STR assigns an independent threshold for each layer, resulting into a threshold vector d . Besides, STR further parameterizes the threshold by another trainable parameter s .
The mapping from s to d is individually designed for CNN and RNN.
Compared to STR, STDS set threshold manually by introducing 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 formulation is based on the idea that increasing threshold from 0 to D could follow different paths.
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 1.
[Table]
3. I TERATIVE SHRINKAGE - THRESHOLDING ALGORITHM ISTA
is initially derived from solving linear inverse problem with regularization min x ∈ R n (cid:8) ∥
Ax − b ∥ 22 + r ( x ) (cid:9) , where A ∈ R m × n and b ∈ R m .
ISTA is later extended to general objective as min x ∈ R n { F ( x ) : = f ( x ) + r ( x ) } with assumptions listed below
(i)
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 .
(ii) Regularization function r : R n → R is continuous convex and can be nonsmooth.
(iii) F is bounded from below.
Leave the regularization r ( x ) alone, applying vanilla SGD to f can be viewed as iteratively calculating proximal regularization (Martinet, 1970) of the linearized f at x , which is suggested by the following fact: x − η ∇ f ( x ) =
arg min y { f ( x ) + ⟨ y − x , ∇ f ( x ) ⟩ + 12 η ∥ y − x ∥ 22 } , where η is explained as “stepsize” in optimization or “learning rate” in the context of deep learning.
For a given point x , F ( y ) can be approximated by expanding f to the quadratic term in a similar vein
[Equation]
The above problem admits a unique minimizer
[Equation]
Note that η = 1 L f gives a upperbound of f ( y ) , which is obtained by the descent lemma f ( y ) ≤ f ( x ) + ⟨ y − x , ∇ f ( x ) ⟩ + L f 2 ∥
y − x ∥ 22 Beck (2017).
It implies with proper choice of η , we are virtually optimizing the upperbound of F ( y ) using minimizer in Eq.
The general form of ISTA iteratively solves Eq.
5 as x (
[Equation]
) = [Equation] x ; x ( t ) ) , 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 sparsity, we are interested in L 1 regularization term r ( x ) = µ ∥ x ∥ 1 , µ > 0 .
Since L 1 norm is separable, we have a closed-form solution with element-wise soft threshold operation as
[Equation]
Eq. 6 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]
Line 4 in Algorithm 1 ,
[Equation]
[Equation]
Line 7 in Algorithm 1 .
[Equation] [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.
7 and Eq. 8 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.
9 is equivalent to the ISTA update rule with threshold equal to forward finite difference d
( t +1) − d ( t ) .
Recall Eq. 6, 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
[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 justify the increasing threshold for a positive penalty term.
Under this framework, we further explain in Appendix B that the original STDS uses an improper threshold scheduler.
Moreover, the validity of training threshold is discussed 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 µ ( t ) ≡ µ during learning, we can deduce a unique threshold scheduler LATS .
In LATS, the change on threshold d ( t +1) − d ( t ) must be proportional to the learning rate η ( t ) during training.
Corollary 1 ( L earning rate A dapted T hreshold S cheduler) .
For fixed L 1 regularized problem
[Equation]
where µ > 0 is the time-independent L 1 penalty coefficient, the threshold scheduler is governed by the learning rate scheduler as d ( t ) =
d (0) + µ t − 1 (cid:80)
[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 (Loshchilov & Hutter, 2017) with η min = 0 .
Assuming the initial threshold is zero, simple algebra shows the definition of corresponding LATS with a slight abuse of notation
[Equation]
12) 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.
We elaborate the detailed derivation of Eq.
12 in Appendix D.
5. S IMPLIFIED THRESHOLD SCHEDULER
Computation such as Eq.
12 is intricate for implementation.
In effect, painstakingly coding LATS according to given learning rate scheduler is not inevitable.
To ease the computing burden, we turn to replacing sum of learning rate with integration of learning rate function.
To be specific, assuming the learning rate scheduler can be expressed by η ( n,b ) = h ( n/N ) , the simplified threshold scheduler is defined by
[Equation]
The simplification is loyal to the idea that the value of Riemann integral could be approximated by rectangle method.
Eq.
13 can be interpreted with scheduler form in STDS as d ( t ) =
g ( t/T ) · D with scheduler function g ( x ) = (cid:82)
tT 0 h ( x )d x/ (cid:82) 1 0 h ( x )d x and final threshold D = d ( N − 1 ,B ) . The detailed derivation of Eq. 13 is given in Appendix D. Now we have the Simplified LATS ( S-LATS for short) for the cosine annealing learning rate scheduler
[Equation] ( x [Equation] The final threshold is D = µη max T/ 2 , which satisfies D ∝ µ .
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 thus be expressed in a unified form of d ( t ) =
g ( t/T ) · D .
We evaluate LATS and S-LATS under identical [Table] final thresholds D = 0 .
5 , 1 . 0 , 5 . 0 , which can be gleaned from Tab.
The results show LATS and S-LATS are indistinguishable from accuracy and sparsity.
Thus, we turn to a simplified threshold scheduler in the following discussion.
5. 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 analogue 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.3. 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 L 1 coefficient µ ( t ) =
β tT , where 0 < β < 1 is a constant.
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]
[Equation]
where the analytic form of g PGH is shown below
[Equation]
5.3. L INK TO EARLY PRUNING
For the PGH scheduler, we interpolate β between 0 and 1 and get a series of different PGH schedulers.
As shown in Fig.
1, the increasing in threshold slows over time, which is caused by decaying L 1 penalty.
For 0 < β < 1 , if pruning is ignored when penalty is below a preset threshold, we get a family of early pruning algorithms.
They will stop pruning at different stages.
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.
[Figure] There are two limit cases when β approaches 0 or 1.
It is obvious that β → 1 leads to S-LATS for no decay is applied.
When β → 0 , the penalty is always zero except for the beginning.
In this case, PGH scheduler degenerates to a magnitude-based pruning after weight initialization followed by a normal training stage.
This is also referred to as sparse-to-sparse training or pruning at initialization method.
E XPERIMENTS
In this section, we test proposed threshold scheduler S-LATS on both deep ANNs and SNNs.
The favorable performances against the previous studies are confirmed.
In all experiments we switch sparisty levels by changing D , which is equivalent to tuning L 1 penalty coefficient.
We also tune hyperparameter β in PGH scheduler to maintain different phases of early pruning algorithm.
Compared to the dense baseline, no tuning on other training hyperparameters is needed, which minimize the effort when applying to other networks.
6. S-LATS
S-LATS achieves state-of-the-art performances on both ANNs (ResNet-50, MobileNet-V1 (Howard et al., 2017)) and SNNs (SEW ResNet-18 (Fang et al., 2021)).
The results on ResNet-like networks and MobileNet-V1 are illustrated in Fig.
We add Gradual Magnitude Pruning (GMP) (Zhu &
Gupta, 2017) into comparison for few recent studies are conducted on MobileNet-V1. Notably, our method surges ahead of all the other baselines under < 98% sparsity for pruning on ResNet-50, which is shown in Table 3.
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 [Table]
[Figure]
extreme sparsity ( ≥ 99%), which suggests the theory might be imperfect under such conditions. An ablation study shows S-LATS outperforms the default threshold scheduler of STDS, which is shown in Appendix F
[Figure]
Pruning using large batch size.
Powerprop (Schwarz et al., 2021), adopts a batch size of 4096 for pruning of ResNet-50 and achieves fascinating performances under high sparsity.
Hence, we explore the large batch size setting.
Due to limited resources, we only increase it to 1024 and enlarge the learning rate correspondingly.
The rightmost of Fig.
2 shows it indeed leads to higher performance, which outperforms all other SOTA studies.
Astonishingly, even though the performance of the dense network slightly degrades, the accuracy of the sparse ones is improved overall.
On the basis of the above finding, we believe applying an even larger batch size, like 4096 used in Powerprop, to our method may lead to top-notch performance tradeoff.
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]
i t y ( % )
D =0.1 D =0.11 D =0.
D =0.
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.
Note that our method is agnostic about the structure of network. Hence, some layers are completely pruned as shown in the right of Fig.
[Table] wherein three consecutive layers within a residual block tend to be pruned simultaneously.
However, owing to the skip connection in ResNet, the feature can still pass through shortcuts to the final FC layer, and thus the whole networks are still normally trained.
The aforementioned results are collected in Tab. 4.
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, and yet becomes 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 “Leave at least one path from input through output”.
We believe our method can for sure further benefit from knowledge in the prosperous field of network pruning. A CKNOWLEDGMENTS
This work is supported by grants from the National Key R&D Program of China under
Grant 2020AAA01035, the Key-Area Research Development Program of Guangdong Province (2021B0101400002) and the National Natural Science Foundation of China under contract
No. 62088102, No. 62176003, No. 62006132, No. 62027804 and No. 61825101.
The computing resources of Pengcheng Cloudbrain are used in this research.
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]
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]
=
[Equation]
t
[Equation]
+ d [Equation] t ) sign( w [Equation]
t
[Equation]
− [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]
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. 40 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
[Equation]
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.
We will show 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 wants to investigate 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.
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 [Figure] 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]
[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) [Equation]
, | θ | < 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
Sine: d ( t ) =
12 (1 + sin( π ( tT − 12 )))
D • 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.
In fact, simply setting threshold trainable indicates one only cares about whether the final threshold is optimal, which turns out to be a tangential issue in soft threshold pruning.
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 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]
For cosine annealing learning rate scheduler, we have
[Equation]
To evaluate the sum of cos( iπ/N ) , we give the following results
[Equation]
Recall Eq. 30, 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) [Equation] ) where κ > 0 is a constant ( κ = 0 . 9 in aforementioned studies), the form of LATS cannot be reduced like cosine annealing scheduler in most cases.
To see this, we write the corresponding LATS as [Equation] Note that 1 − iT makes up an arithmetic progression, and the threshold is the sum of their powers. Simplifying the sums of powers of arithmetic progression requires the so-called Bernoulli number
Jacobi (1834); Knuth (1993) when κ is integer. For a general κ > 0 , the expression of Eq.includes the generalized harmonic numbers H ( − p ) n := (cid:80) n [-− 1 i =
-]k =1 k p , which is further based on the Hurwitz zeta function Coffey (2008). In such case, we cannot analytically compute LATS, which forces us to do the summation in Eq. 34.
It thus brings about the accumulative error.
In brief, we cannot expect each learning rate scheduler corresponds to an analytical and simple form of LATS.
To handle this, we turn to an approximation rather than precisely evaluating d ( t ) . Returning to Eq. 29, with d (0 , 0) ignored, we split the right hand side into two terms B (cid:80) n − {+1 i =
 h ( i/N ) and b · h ( n/N ) .+}
The first term could be viewed as left Riemann sum in two steps 1) interpolate points 0 , 1 /N, 2 /N, . . .
, ( n − 1) /N, n/N with constant spacing 1 /N , the width of the rectangles 2) evaluate the sum of rectangle areas with height h ( i/N ) . It leads to the approximation of the Riemann integral [Equation]
The second term match b/B part of a residual tiny rectangle [Equation] To sum up, Eq. 35 and Eq. 36 together make a numerical approximation of integral
BN (cid:82) n/N + b/BNh ( x )d x , which is shown schematically in Fig. Note that we could write training progress as t/T now, where t = Bn + b is the current iteration id and T = BN is the total training iterations.
So far, we successfully replace the summation with integration. However, the current final threshold is BN (cid:82) 1 0 h ( x )d x , which is different from the real one d ( N − 1 ,B ) . To keep the final threshold d ( T ) =
d ( N − 1 ,B ) , we normalize the integral as follows
[Equation]
[Equation]
For most learning rate scheduler functions, integration is much easier than summation.
S-LATS enables us to apply our pruning method on wider varieties of deep learning applications.
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 1 N 2 N · · · · · · n − 1 N nN n +1 N n N + b B N Training Progress 0. Lea r n i ng R a t e S c hedu l e r h ( · ) 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.
Apparently, our proposed method towers over previous sparsifyduring-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]
S-LATS 68.25 ±0. GMP 61. STR 64. STR 62.
STR 61. WoodFisher 63. S-LATS 66.73 ±0. S-LATS 65.63 ±0.
S-LATS 64.93 ±0.
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 [Equation]
In the presence of weight decay [Equation] the update rule described by Eq.
17 has an additional penalty term as follows
[Equation]
Following derivation in Eq.
[Equation] we have
[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.
39, 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]
(43) where u th , t f are the firing threshold and firing time respectively.
The ODE in Eq.42 can be discretized via the Euler method and transformed into an RNN-like iterative computing manner as follows
[Equation]
(44) 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 between the number of spikes and weight sparsity is manifested in the slightly decreased average Sparsity (%) 0. A v g .
F i r i ng
R a t e 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. 9.
[Table]
[Table]
[Table]
I S PARSITY BUDGETS
[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.