Get More at Once: Alternating Sparse Training with Gradient Correction
Abstract
Recently, a new trend of exploring training sparsity has emerged, which removeparameters during training, leading to both training and inference efficiency improvement.
This line of works primarily aims to obtain a single sparse modelunder a pre-defined large sparsity ratio.
It leads to a static/fixed sparse inferencemodel that is not capable of adjusting or re-configuring its computation complexity(i.e.
, inference structure, latency) after training for real-world varying and dynamichardware resource availability.
To enable such run-time or post-training networkmorphing, the concept of ‘dynamic inference’ or ‘training-once-for-all’ has beenproposed to train a single network consisting of multiple sub-nets once, but eachsub-net could perform the same inference function with different computing complexity.
However, the traditional dynamic inference training method requires ajoint training scheme with multi-objective optimization, which suffers from verylarge training overhead.
In this work, for the first time, we propose a novel alternating sparse training (AST) scheme to train multiple sparse sub-nets for dynamicinference without extra training cost compared to the case of training a singlesparse model from scratch.
Furthermore, to mitigate the interference of weightupdate among sub-nets without losing the generalization of optimization, we propose gradient correction within the inner-group iterations to reduce their weightupdate interference.
We validate the proposed AST on multiple datasets againststate-of-the-art sparse training method, which shows that AST achieves similar orbetter accuracy, but only needs to train once to get multiple sparse sub-nets withdifferent sparsity ratios.
More importantly, comparing with the traditional jointtraining based dynamic inference training methodology, the large training overheadis completely eliminated without affecting the accuracy of each sub-net.
Introduction
For Deep Neural Networks (DNNs), sparsification (i.e., pruning) has been widely explored in the lastdecade aiming to reduce the computational and memory cost by removing unimportant parameters.
Early works of DNN pruning focused on exploring the post-training sparsity [10, 7, 11] to improvethe inference efficiency, which prunes a well-trained dense model followed by additional fine-tuningto recover accuracy.
Recently, exploring in-training sparsity has emerged as a promising techniqueto improve the training efficiency by pruning parameters during training [3, 4, 9, 19, 14].
Thesuccess of these in-training sparsification methods is mainly rewarded from the prune-and-regrowtechnique, which periodically and aggressively eliminates the unimportant non-zero weights from themodel, and then regrows certain portion of the pruned candidates back.
Until now, such in-training
Loss
Loss 𝐗𝐗 𝐭𝐭 𝐘𝐘 𝐭𝐭 𝐘𝐘 𝐭𝐭
Loss
Loss
Sumloss 𝐘𝐘 𝐭𝐭
Loss
Loss 𝐘𝐘 𝐭𝐭+𝟏𝟏 𝐘𝐘 𝐭𝐭
Mini-Batches (b) 𝐗𝐗 𝐧𝐧 𝐗𝐗 𝐭𝐭 𝐗𝐗 𝐧𝐧 𝐗𝐗 𝐭𝐭 𝐗𝐗 𝐧𝐧 𝐗𝐗 𝐭𝐭 𝐗𝐗 𝐭𝐭+
Switch
Mini-Batches
Training Direction Mini-Batches (a) (c)
Grad Correction 𝐗𝐗 𝐭𝐭+
Loss 𝐘𝐘 𝐭𝐭+
Switch
Disabled connection
Active connection
Inner group
Pruned Channel sparsity works primarily focused on training a single sparse model under a pre-defined sparsity ratio,generating a static and fixed sparse network.
In other words, the sparsity ratio has to be pre-calculatedbefore training depending on the target hardware resources and application requirement.
However,obtaining multiple sparse models (towards storage/energy vs. accuracy trade-off) from the samemodel architecture requires extra fine-tuning or training overhead.
To enable run-time or post-training network morphing, the concept of ‘dynamic inference’ or ‘trainingonce-for-all’ has emerged recently [24, 2].
Slimmable Neural Network (S-Net) [24] and its optimizedcounterpart (US-Net) [22] allow the model to switch between different architectures among thepre-defined sub-net candidates with different channel widths.
Such run-time dynamics allows users totune the tradeoffs between model complexity and inference accuracy.
To mitigate the weight updateinterference, S-Net/US-Net forces the smaller sub-net to be completely nested by the larger sub-net.
Based on that, OFA [2] and BigNAS [23] are extended to support a much larger number of sub-netsacross more dimensions (e.g., depth, width, kernel size, and resolution).
However, all prior methodsrely on a joint training scheme with multi-objective optimization, which requires minimizing the lossof all sub-nets collectively.
Specifically, all sub-nets have to perform forward and backward passesin each mini-batch iteration, resulting in much larger training overhead compared to the individualnetwork training.
To address these concerns, we propose a novel alternating sparse training (AST) method for dynamicinference which aims to train multiple sparse sub-nets at one training procedure without extra trainingcost, compared to training a single sparse model independently [3, 4, 9, 19, 14].
To generate andsparsely train each sub-net, we adapt the prune-and-regrow scheme from the in-training sparsitymethod [14], as it achieves high training sparsity among other schemes.
AST alternately trainssub-nets with different sparsity values for different mini-batch iterations, as illustrated in Figure 1.
Inspired by the implicit regularization of SGD that was revealed by [18, 16], we further demonstratesuch alternating training scheme for dynamic inference helps to maximize the inner product ofgradients between the consecutive mini-batch iterations (i.e., sub-nets in our case).
However, evenwith the help of the implicit regularization, we find that negative inner products still remain, whichrepresents the conflicting gradient direction between each two sub-nets.
Such phenomenon causesthe interference of weight update, while it is helpful to escape from local minima and improvegeneralization of optimization for individual network training by mini-batch SGD [6, 1].
To mitigatethe interference of weight updates between sub-nets but without losing the generalization, we furtherpropose gradient correction to remove the conflicting gradient direction between sub-nets only withininner-group iterations, which is defined as N consecutive mini-batch iterations, where N is thenumber of sub-nets (e.g., N = 2 as shown in Figure 1).
Overall, compared to the sparse training and dynamic inference works, AST achieves both trainingefficiency with sparsity (against dynamic inference) and inference dynamics with multiple sparsesub-nets (against individual sparse training).
We conduct comprehensive experiments on CIFAR-10,
CIFAR-100 and ImageNet datasets with various DNN models and sparsity granularities.
Compared to the SoTA sparse training on CIFAR, the proposed AST algorithm achieves similar or better accuracywhile training multiple sparse sub-nets all at once (reduced training cost).
Furthermore, compared tothe dynamic inference method, AST consistently achieves better accuracy on all the experiments.
Related Works and Background2.1 Sparse Training
As summarized in [8], sparse training works mainly can be categorized into three groups according tothe time point the sparsity is applied: 1) Post-training sparsification [3, 4, 9, 19, 14], which removesthe weights by fine-tuning based on a pre-trained model; 2) Before-training [11, 20] sparsification,which obtains a sparse model before the main training procedure; and 3) In-training sparsification,which removes weights during the training process from scratch [14, 26, 4, 3].
In-training sparsification Different from post-training and before-training sparsification, exploiting sparsity during training only requires a single training process from scratch.
The pruning topologyis gradually perfectized along with the weights optimization.
Compared to the before-training sparsification, pruning the model during training makes the gradient visible to the algorithms, so the reflectionof the gradient can be used to correct the final pruning decision, leading to better accuracy comparedto post-training.
Motivated by this, the prune-and-regrow technique [3] periodically removes theunimportant non-zero weights from the sparse model and regrows certain pruned weights back duringeach mini-batch iteration of training process.
As a representative work, RigL [4] first prunes a certainratio r based on weight magnitude as:
[Equation]
where TopK ( v, k ) returns the weight/activation tensor retaining the top k -proportion of elementsfrom v , and s is the targeted sparsity ratio.
After that, the r proportion of new connections arere-generated based on the gradient magnitude in the same mini-batch iteration:
[Equation]
By doing so, such prune-and-regrow scheme optimizes the sparse connection with a fixed sparsityratio r during the entire training process.
Based on this, recent works develop different weightimportance criterion to perform prune-and-regrow.
SNFS [3] uses the momentum magnitude ofweight to do pruning.
MEST [26] uses the sum of weight and gradient magnitude to indicate theimportance of weights for pruning, and further randomly select weights to grow back.
In addition,
GraNet [14] employs the same rule as Rigl but adopts a smaller sparsity as initialization (e.g., 0.5)and gradually prunes the model to the target sparsity (e.g., 0.8, 0.9) by following a cosine decay.
Inaddition, more similar to our work, [19, 16] propose to alternately train a dense model and one ofits sparse variant (i.e., sub-net) returning both an accurate dense model and a sub-net.
However, therationality behind it is not clearly described.
Overall, the objective of our work is obtaining multiplesparse sub-networks with the same amount of training effort as the individual network sparsification.
More recently, based on the in-training sparsification mechanism, [13] proposes dynamic sparsetraining ensemble method to independently generate multiple sparse sub-nets for ensemble.
Suchmethod is orthogonal to our work which could be further combined to improve accuracy.
2.2 Joint Training for Dynamic Inference
Dynamic inference works [24, 22, 23, 2, 21] primarily aim to train a single network which consists ofmultiple sub-nets to perform inference independently.
Since the weights of the sub-nets are partiallyshared between each other, they can be switched at run-time permitting dynamic inference accuracycomplexity trade-off.
Such concept is firstly proposed by Slimmable Neural Network (S-Net) [24] andits optimized counterpart (US-Net) [22], which train a single neural network executable at differentchannel-widths.
Inspired by S-Net, OFA [2] and BigNAS [23] are proposed to construct a dynamic
DNN that includes a much larger number of sub-nets across many more dimensions (e.g., depth,width, kernel size, and resolution) .
Note that, as revealed by [2], to enable sub-nets switching withun-compromised individually inference accuracy, the structures of sub-nets are defined by a subsetrule : the smaller sub-net must be the completely subset of the larger sub-net.
Specifically, this line of
Sub-Net Sub-Net
Sub-Net
Sub-Net 1 𝒕 𝒕 +𝒕 +𝒕 +
Iterations
Switch Switch
Switch
Initialization
Disabled connection
Active connection
Loss t
Loss t+ Loss t+
Loss t+ works train a network in a joint training fashion, which can be expressed as:
[Equation]
(3) where X is the mini-batch of inputs with corresponding targets Y , and N is the number of sub-nets.
L ( · ; · ) calculates the cross-entropy loss of DNN output and target.
f ( X, { W i } ) computes the outputof sub-net parameterized by { W i } .
However, such method consumes much larger training timecompared to training a single individual model, since all sub-nets have to perform forward andbackward once in each mini-batch iteration as shown in Eq.
Sparse Train Once Get More
3.1 Rationale of Alternating Sparse Training (AST)
Different from joint training for dynamic inference such as S-Net [24] and US-Net [22], whichupdates all sub-nets in each mini-batch iteration, we propose the alternating sparse training (AST)scheme which trains multiple sub-nets in an alternating fashion over time, where a single sub-net istrained in one iteration.
Note that, the sub-net is defined a part weight of the network, which can runinference independently.
All the sub-nets are partial shared with each other within a single network.
As shown in Figure 2, assuming we have three sub-nets with different sparsity ratio (e.g., sub-net 1< sub-net 2 < sub-net 3), such training scheme will start to train the sub-net 1 in the first iterationfollowed by the training of the sub-net 2, and sub-net 3 in the second and third iteration, respectively.
It only updates the active connections while ignoring the disabled weights (i.e., pruned for currentsparse sub-net), then repeatedly switches sub-nets every three consecutive iterations until the trainingcompletes.
The rationale behind AST is inspired by the findings in Reptile [18], which is a meta-learningalgorithm originally aiming to learn an initialization of a model from multiple tasks for fast adaptionon new tasks.
It shows that mini-batch SGD imposes an implicit regularization to maximize thedot-product of consecutive mini-batches.
In the following, we will show how this proof adapts to ourcase.
Assuming that we alternately train two sub-nets, sub-net 1 and sub-net 2 with two consecutivemini-batches B 1 and B 2 respectively, according to the Taylor Series, the gradient of sub-net 2 g 2calculated by SGD can be expressed by:
[Equation]
Where H 1 is Hessian of the sub-net 1 and α is the current learning rate.
Similar to Reptile, the termαH 1 g 2 serves to maximize the dot-product of the consecutive sub-nets, where the expectation can beexpressed as:
[Equation]
[Figure] [Figure]
Thus, − αH 1 g 2 is the direction that serves to maximize the inner product of two consecutive minibatches, which means to improve the gradient alignment of two sub-nets towards better learninggeneralization.
Such proof can be easily adapted to more than two sub-nets.
The detailed explanationof Eq. 4 is relegated into the Appendix.3.2 Proposed AST with Gradient Correction on Consecutive Inner-group Sub-nets (AST-GC)
[Figure]
Although the implicit regularization of mini-batch SGDhelps to maximize the inner product gradient of consecutive sub-nets as discussed in Section 3.1, we indeedobserve that there still remain partial number of negative inner product during training, causing conflictingdirection of weight updates among sub-nets as shownin Figure 4.
However, in terms of optimizing a network by mini-batch SGD, the noise caused by such conflicting gradient direction helps to escape from saddlepoints or local minima and improve the generalizationas well [6, 1].
Due to that, inspired by the gradient projection method, which is used in multi-task learning [25]and continual learning [12], we propose a gradient correction technique to remove the conflictinggradients within the inner-group iterations during training, meanwhile allow the gradient to havenegative direction for the inter-group sub-nets.
As shown in Figure 3, the inner-group iterationsis defined as the consecutive sub-nets within N mini-bath iterations, where N is the total numberof sub-nets (e.g., N =3 as shown in Figure 3).
In contrast, inter-group represents the relationshipbetween inner-group iterations.
Specifically, the gradient correction adapts a simple procedure withinthe inner-group sub-nets: if the gradients between two consecutive sub-nets are in conflict (i.e., theircosine similarity is negative), we project the gradient of current sub-net onto the normal plane of thegradient of its prior sub-net.
Otherwise, the original gradient is unchanged.
Considering two sub-netsin inner-group iterations with the gradient g i and g j respectively, the technique can be mathematicallyillustrated as:
[Equation]
By doing so, the proposed gradient correction within inner-group iterations have two benefits: 1) itcould guarantee that all sub-nets within the consecutive inner-group are updated in the non-conflictingdirection, reducing the weight interference between each other; 2) the normal update using themini-batch SGD between inter-group sub-nets will help to improve the generalization of the model.
3.3 Sparse Sub-net Training
Another important aspect of our AST scheme is to generate and train “ sparse ” sub-nets.
Followingmost sparse training works, we adopt the prune-and-regrow scheme from GraNet [14] as the backbonetechnique.
Specifically, AST starts from a random initialized sparse model, and then applies pruneand-regrow mechanism as described in Eq.
1 and Eq.
2 for each training iteration with the sub-netspecific sparsity ratio.
The pseudocode of the proposed AST algorithm is relegated into the Appendix.
Given to the fact that alternating training scheme switches the sub-net per iteration, the following
[Figure]
[Table]
questions arises: 1) What is the optimal architecture relationship between sub-nets?
How thepruning should be scheduled for each sub-net?
We address these two questions with the followingobservations:
Observation 1: Enabling the freedom of exploring unique architectures of each sub-net elevates thelearning capacity of AST over the S-Net.
To validate this observation, we generalize the sub-net relations into the following two categories:
• Completely subset (CS): As proposed by S-Net [24], the high-sparsity models are fully
contained in the low-sparsity models (Figure
(a)).
Under the context of prune-and-regrow,
the regrowing process is only performed in the lowest sparsity model, while the rest of
the sub-nets rigorously extend the sparsity from the previous low-sparsity model by using
magnitude-based pruning only.
The one-time regrow guarantees the fully-subset relation-
ships among different sub-nets.
• Non disjoint (ND): Each sparse sub-net performs prune-and-regrow individually to optimize
the overall pruning decision during training.
As depicted in Figure 5(b), sub-nets can freely
exploit
their own architectures while the intersections remain non-empty.
Compared to the
CS scheme, the non-disjoint AST empowers the subset networks with more architecturefreedom.
Figure 6 shows the layer-wise sparsity gap and non-overlap between two NDtrained sub-nets that target different final sparsity values ( s f ).
The non-overlap is computedby XORing the binary sparse masks between sub-nets, the percentages of “1” in the resultanttensor represents the level of non-overlap.
Apparently, the distinction of connections islarger than the sparsity difference, which implies the existence of the unique connectionsgenerated by the ND prune-and-regrow in each sub-net.
With the ND scheme, we observe that the small amount of architecture freedom shown in Figure 6 can lead to the improved overall performance:
Table 1 summarizes the accuracy of the wide
ResNet-32 [26] trained by different AST schemes on CIFAR-10 dataset.
Assisted by the comprehensive architecture exploration, the non-disjoint AST scheme achieves higher accuracy comparedto the conservative completely-subset (CS) training.
Thus, in this work, we use ND-scheme for allexperiments.
Observation 2: Intermittent sparsity increment among sub-nets stabilizes AST process.
As introduced in Section 2.1, the sparsity of the in-training sparsification method is periodically up-
[Table]
dated (e.g., 1,000 iterations) based on a pre-defined sparsity schedule [14].
Regarding the ASTscheme, the sub-net model architectures are consecutively switched and trained, but the successivearchitecture switching does not imply the necessity of consecutive sparsity update of each sub-net.
rewarded from the extended adjustment period ∆
On the contrary, the intensive sparsity incrementof all sub-nets could destabilize the training.
Inthis work, the sparsity of each sub-net increasesperiodically, but the sparsity increment of eachsub-net is intermittently performed with the adjustment period ∆ τ .
In the meantime, sub-netsare still consecutively switched during the ∆ τ .
Figure 7 shows the CIFAR-10 validation accuracy of an AST-trained ResNet-32 sub-net withfinal sparsity of 98%.
During the early stage oftraining with high learning rate, the stabilizedvalidation accuracy indicates the effectivenessof the extended adjustment.
Experiments4.1 Settings
The proposed AST method is thoroughly verified with multiple datasets, including CIFAR-10, CIFAR100, and ImageNet.
We use similar training scheme as prior works [14, 26], which employed 160epochs for CIFAR-10/100 and 100 epochs for ImageNet training.
The multiple sparse sub-nets aretrained from scratch and pruned by the proposed AST algorithm.
The cosine annealing learning rateschedule is employed with 0.1 initial learning rate for the CIFAR datasets training.
For ImageNetdataset, we introduce extra 5 epochs for the initial warmup training.
The pruning candidates areglobally selected while the first layer remains dense.
Regarding the regrowing process, the percentageof regrow candidates gradually decrease from 0.5 to 0.0 with cosine annealing schedule.
In themeantime, the extended adjustment period between sub-nets is set to 100 and 400 for CIFAR and
ImageNet experiments, respectively.
While we believe a more fine-grained hyperparameter tuningcould lead to better accuracy, we choose the above scheme for simplicity and reproducibility.
Thereported sub-net accuracy values are obtained from a single background model checkpoint.
In allexperiments, we report the average accuracy with its variation in 3 runs.
4.2 Main Results
CIFAR-10 and CIFAR-100.
Table 2 shows the CIFAR-10/100 accuracy of the proposed ASTalgorithm, where we used the wide ResNet-32 model that was employed in [26].
Following thetypical high sparsity results reported previously, we train three sub-nets with 90%, 95%, and 98%sparsity from scratch all at once with identical epochs as prior works.
Given the fact that the highinitial sparsity is beneficial to the overall memory saving of the training process, we report the resultswith both dense ( s i = 0% ) and highly sparse ( s i = 90% ) initial models.
Compared to the prior
[Table]
individually-trained SoTA works [14, 26], the proposed AST+GC algorithm obtains three highlysparse models through one-time training with negligible accuracy degradation.
The resultant 3 ×total training cost reduction frees the model from energy-urged extra training, leading to power andlatency benefits.
The joint training scheme [24] requires multiple forward pass in each iteration, theincreased computational cost leading to the averaged-yet-suboptimal performance.
In addition to thereduced training cost, the proposed AST method outperforms the joint-training scheme [24] by up to1.3% and 2.9% inference accuracy, while achieving up to 2.63 × training cost reduction on CIFAR-10and CIFAR-100 datasets, respectively.
We also verify the AST scheme with more sub-nets and largesparsity gaps in Section 4.4.
ImageNet-2012.
We further evaluate the proposed method with ResNet-50 on ImageNet in Table 3.
It is noticeable that the 0 .
5 × and 0 .
25 × model of the jointly-trained S-Net [24] corresponding to72 .
98% and 92 .
23% weight sparsity, the averaged overall sparsity (82.6%) is less than our targetedsparsity (85%).
The proposed AST training scheme outperforms the joint-training scheme by 7.5%inference accuracy with 80% highly sparse initial models.
Compared to the individual trainingscheme, the proposed AST scheme achieves a comparable or even better performance to SNIP [11]and SET [15], while reduces the total training cost up to 2.38 × .
The experimental results withdifferent settings and models are shown in the Appendix.
4.3 Extend to Structured Fine-grained Sparsity
Demonstrating the sparsity-induced speedup on GPU has now emerged as a feasible solution to depictthe effectiveness of pruning algorithms.
The recent Nvidia Ampere architecture is equipped with the
Sparse Tensor Cores to accelerate the neural network computation on GPU with N : M structuredfine-grained sparsity [27].
Varying the group size M and the number of sparse elements N leads todifferent overall sparsity values.
Motivated by this, we extend the proposed AST algorithm to trainmultiple sub-nets with different N : M sparsity patterns all at once.
Specifically, the prune-and-regrowis performed based on the group-wise summed weight or gradient magnitude.
Powered by theopen-sourced Nvidia-ASP library, the convolution computation can be accelerated when the sparseweight group size M is divisible by 4 (e.g., 4, 8, 16).
The proposed AST algorithm collectivelytrains four ResNet-18 sub-nets with 2:4, 3:4, 7:8, and 15:16 sparse patterns, corresponding to 50%,75%, 87.5%, and 93.75% overall sparsity, respectively.
Starting from scratch, the percentage of the
N : M sparse groups is gradually increased from 0% to 100% within each sub-net.
Compared to theindividually trained dense model, the proposed AST scheme achieving up to 2.3 × inference speedup on GPU with 4 × less training efforts and negligible accuracy loss, as shown in Table 13.
Theinference time is measured on a Nvidia 3090 GPU with FP32 data precision.
4.4 Ablation Study
AST with more than 3 sub-nets.
We also verify the effectiveness of the proposed AST schemefor training different number of subset networks.
Table 5 shows the CIFAR-10 performance of
AST by training 2, 3, 4, 5 sparse sub-networks collectively based on ResNet-18.
Compared to theindividually-trained baseline sparse model, training more sub-nets with AST reduces the total training
[Table] [Table]
effort with the cost of the marginally degraded accuracy, especially in the high sparsity models (e.g.,98% sparsity).
Given the tradeoffs between accuracy and total training cost, training three sub-netshave the best overall performance.
The flexibility of training a large number of sub-nets all at onceenables the proposed AST algorithm to select suitable architecture for different power budgets.
AST sub-nets with large sparsity difference.
Besides the collective training of highly sparse models,learning the largely-varied subset architectures is also essential.
As shown in Table 5, the proposed
AST algorithm is still able to optimize the model performance with the large sparsity gaps (e.g., 70%vs.
Extended AST training efforts.
As shown in Figure 2, AST iteratively selects different sub-net foreach mini-batch iteration.
Even though the batch shuffling is employed to the training, each sub-netcannot be fully trained by the whole training set inside each epoch.
A straight-forward solution isextending the total training efforts (epochs).
Assume the unit training cost is 160 epochs (1 × ), weevaluate the AST performance with three wide ResNet-32 [26] sub-nets with different training effort,as shown in Figure 8.
Compared to the individually trained GraNet [14] baseline (total=3 × ), ASTachieves the same accuracy in all three sparse models with only ∼ 2 × averaged total training effort.
Conclusion
In this work, we first propose Alternating Sparse Training (AST) scheme to train multiple sparseneural networks all at once.
Then, we demonstrate the benefits of gradient correction (GC) withtheoretical analysis and experimental results.
As one of the earliest research works in this domain, the
AST algorithm achieves high accuracy and high training efficiency without any repeated or ensembledtraining steps [24].
Inspired by the prune-and-regrow scheme [14], the proposed AST-GC schemeexploits multiple sparse sub-nets simultaneously while achieving comparable or even higher accuracyagainst the individually-trained SoTA methods.
Obtaining multiple models with AST brings practicalbenefits to the energy-driven hardware computation along with superior performance on accuracy.
Impact and limitation.
In the current work, AST can work well on a certain number of sub-nets.
Extending it to support a larger number of sub-nets is an interesting direction.
In addition, to enlargethe practical efficiency for both training and dynamic inference, we will further apply quantizationand investigate more structured sparsity pattern upon AST.
Currently, we have demonstrated thepractical value for dynamic inference on GPU with Nividia Ampere architecture by using N:Mstructured sparsity pattern.
We would like to further demonstrate the actual efficiency of AST onmore hardware platforms, such as CPU, mobile/IoT devices, specific hardware accelerators, etc.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justification to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section 4. • Did you include the license to the code and datasets?
[Yes] The code and the data areproprietary.• Did you include the license to the code and datasets?
[N/A]Please
do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[N/A](d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No](b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See section 4.(c)
Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See appendix4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
[Table]
(b) Did you mention the license of the assets?
[No]
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[No]
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[No]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]5.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
Require: Model weight W , number of sub-nets N , initial sparsity s i , target sparsity [ s 1 f , ..., s Nf ] , gradual pruning frequency ∆ T , extended adjustment period ∆ τ .
1: random initialize W with initial sparsity s i 2: for each training iteration t do 3: get current number of training iterations t c 4: switching to current sub-net W i where i ← t c mod ( N ∗ ∆ τ ) 5: training W i ←
SGD ( W i ) 6: if ( t mod ∆ T == 0) then 7: gradual pruning with the pruning rate produced by Eq.
6 with target sparsity s if 8: prune-and-regrow by Eq.1 and Eq.
9: if i !
=
then 10: gradient correction within the inner-group by Eq.
11: end if 12: end if 13: end for
Appendix6.1 Algorithm
The pesudocode of AST is presented in Algorithm 1.
To generate and sparsely train each sub-net, ASTadapts the gradual pruning scheme combining with prune-and-regrow from GraNet [14].
Specifically,given initial sparsity s i , target sparsity s f , gradual pruning frequency ∆ T , starting and end epoch ofgradual pruning t 0 and t f , pruning iterations n , the pruning rate of each pruning iteration is definedas:
[Equation]
6.2 Detailed explanation of rational of AST
Following the proof of Reptile [18], the expectation of αH 1 g 2 can be further expressed as the gradientinner product of two consecutive sub-nets:
[Equation]
As shown in Eq.6.2, it is clear to see that the term − αH 1 g 2 is the direction hat serves to maximizethe inner product of two consecutive mini-batches.
Thus, it proves that the proposed AST has implicitregularization to alignment the weight update between sub-nets.
6.3 Detailed experimental setup of AST
The training hyper-parameters of the compared individual sparse training works are same for CIFAR10 and CIFAR-100.
But this line of works adapt different hyper-parameters to achieve good accuracyon ImageNet.
The report accuracy of Rigl [4] uses 4096 batchsize and trains the model on 100 epochswith initial learning rate 1.6.
GraNet [14] uses 64 bathsize, 100 training epochs, and set the initiallearning rate to 0.1.
Mest [26] uses a larger 2048 batchsize and trains the model on 150 epochs withthe initial learning rate 2.048.
In this case, to evaluate our method and fair compare with them, weconduct the basic training settings like the original ResNet by using 256 batchsize and training themodel on 100 epochs with 0.1 initial learning rate.
We believe a more fine-grained hyper-parametersetting could lead to better accuracy.
We only run AST on ImageNet once due to the limited resourcesby using four Nvidia RTX A4000 GPUs.
[Table] [Table]
[Table]
6.4 Additional experimental results
To evaluate the effect of initial sparsity
[Table]
we conduct the experiment by using 50% initial sparsity of
ResNet-50 on ImageNet as shown in Table.7.
Compared to 80% initial sparsity, 50% initial sparsitycould achieve 0.6% accuracy gain on both 80% and 90% prune ratio.
Furthermore, we also conductthe experiment on the smaller ResNet-18, which achieve 62.3% and 62.1% accuracy on 80% and90% prune ratio respectively.
6.5 Training cost comparison
In addition to the training rounds comparison, we analyze the training cost of the proposed ASTmethod in terms of the detailed metrics.
Table 8 summarizes the number of total FLOPs of trainingmultiple sparse deep neural networks.
Compared to the joint-training scheme [24] or seperatelytrained GraNet [14], the proposed AST algorithm achieves up to 2.63 × training cost reduction, whilemaintaining the similar inference accuracy as the individual training baseline.
Furthermore, thedetailed computation cost of the ImageNet experiments are reported in Table 3.
6.6 The impact of different extended adjustment (EA) periods
Regarding the Observation 1 , we have summarized the comparison results between the investigated Completely-subset (CS) scheme and the Non-disjoint scheme (ND) in Table 1 of the originalmanuscript.
As a result, on CIFAR-10 dataset, the ND scheme has the outperformed performancecompared to the constrained completely-subset scheme (CS).
For the Observation 2 , to validate the effectiveness of the proposed extended adjustment method (EA),we add an ablation study on various adjustment periods ∆ τ , which is used to determine the frequencyof sub-nets switching.
As shown in Table R1, compared to the smaller adjustment period (i.e.,∆ τ = 0 , 90 ), ∆ τ = 300 achieves the best accuracy on all three sparsity levels.
The reason is thatperforming sub-net switching frequently elevates the instability of model optimization.
[Table]
41 ± 0.
72.57 ± 0.
68.42 ± 0.
6.7 AST vs. Naive fine-tuning
We analyze the impact of fine-tuning based on the following three perspectives:
A short time of fine-tuning from the dense pre-trained model.
Start with sparse pre-training, fine-tune the high sparsity models with a short epochs, whilekeep the overall training cost (time) as same as AST.3.
To further clarify the advantages of AST, we also investigate another perspective: a short
time of fine-tuning from the sparse pre-trained model.
Same as the experimental setup in the main paper, we conduct the experiments based on the wide
ResNet-32 model on CIFAR-100 dataset.
Given the dense pre-trained model, we separately fine-tunethe dense model to achieve 90%, 95%, and 98% sparsity with minimum efforts.
As shown in Table 10,fine-tuning from a dense model in a short period cannot achieve comparable accuracy as the proposed
AST algorithm.
Furthermore, the 160 epochs of pre-training and additional fine-tuning elevate theoverall training costs.
In addition to the dense model fine-tuning, we address the second concern of the reviewer byperforming the sparse progressive training while keeping the overall training cost to be the same as asingle AST training.
With the wide ResNet-32 model, we first sparsify the model to 90% sparsityfrom scratch with 60 epochs.
Subsequently, we prune the 90% sparse model to 95% and then to 98%sparsity with 50 epochs of fine-tuning.
Compared to the single AST training, the total training effort
[Table]
[Table]
is the same (60+50+50=160 epochs).
As shown in Table 11, such an individual pruning method failedto achieve the performance as the proposed AST training method.
The large accuracy gap suggeststhe necessity of the proposed alternative sparsification training.
Furthermore, we investigate the impact of fine-tuning based on a pre-trained sparse model.
We firstfully train a sparse subnet with 90% sparsity (with 160 epochs) and prune the resultant model to 95%and 98% with a minimum amount of fine-tuning.
As shown in Table 12, fine-tuning the 90% sparsemodel to 95% or 98% sparsity with up to 30 epochs cannot achieve comparable accuracy as AST,with even higher total training effort.
The experimental results in Table 10, Table 11, and Table 12 suggest that it is difficult for individualfine-tuning to achieve the level of high sparsity and high accuracy as the proposed AST, regardless ofthe initial sparsity of the inherited model checkpoint.
[Table] [Table]