Get More at Once: Alternating Sparse Training with Gradient Correction
Abstract
Recently, a new trend of exploring training sparsity has emerged, which removes parameters during training, leading to both training and inference efficiency improvement.
This line of works primarily aims to obtain a single sparse model under a pre-defined large sparsity ratio.
It leads to a static/fixed sparse inference model that is not capable of adjusting or re-configuring its computation complexity (i.e., inference structure, latency) after training for real-world varying and dynamic hardware resource availability.
To enable such run-time or post-training network morphing, the concept of ‘dynamic inference’ or ‘training-once-for-all’ has been proposed to train a single network consisting of multiple sub-nets once, but each sub-net could perform the same inference function with different computing complexity.
However, the traditional dynamic inference training method requires a joint training scheme with multi-objective optimization, which suffers from very large 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 dynamic inference without extra training cost compared to the case of training a single sparse model from scratch.
Furthermore, to mitigate the interference of weight update among sub-nets without losing the generalization of optimization, we propose gradient correction within the inner-group iterations to reduce their weight update interference.
We validate the proposed AST on multiple datasets against state-of-the-art sparse training methods, which shows that AST achieves similar or better accuracy, but only needs to train once to get multiple sparse sub-nets with different sparsity ratios.
More importantly, comparing with the traditional joint training based dynamic inference training methodology, the large training overhead is completely eliminated without affecting the accuracy of each sub-net.
Code is available at https://github.com/mengjian0502/AST .
Introduction
For Deep Neural Networks (DNNs), sparsification (i.e., pruning) has been widely explored in the last decade 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 improve the inference efficiency, which prunes a well-trained dense model followed by additional fine-tuning to recover accuracy.
Recently, exploring in-training sparsity has emerged as a promising technique to improve the training efficiency by pruning parameters during training [3, 4, 9, 19, 14].
The success of these in-training sparsification methods is mainly rewarded from the prune-and-regrow technique, which periodically and aggressively eliminates the unimportant non-zero weights from the
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 model, and then regrows certain portion of the pruned candidates back.
Until now, such in-training 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-calculated before training depending on the target hardware resources and application requirement.
However, obtaining multiple sparse models (towards storage/energy vs. accuracy trade-off) from the same model 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 optimized counterpart (US-Net) [22] allow the model to switch between different architectures among the pre-defined sub-net candidates with different channel widths.
Such run-time dynamics allows users to tune the tradeoffs between model complexity and inference accuracy.
To mitigate the weight update interference, 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-nets across more dimensions (e.g., depth, width, kernel size, and resolution).
However, all prior methods rely on a joint training scheme with multi-objective optimization, which requires minimizing the loss of all sub-nets collectively.
Specifically, all sub-nets have to perform forward and backward passes in each mini-batch iteration, resulting in much larger training overhead compared to the individual network training.
To address these concerns, we propose a novel alternating sparse training (AST) method for dynamic inference which aims to train multiple sparse sub-nets at one training procedure without extra training cost, compared to training a single sparse model independently [3, 4, 9, 19, 14].
To generate and sparsely train each sub-net, we adapt the prune-and-regrow scheme from the in-training sparsity method [14], as it achieves high training sparsity among other schemes.
AST alternately trains sub-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 demonstrate such alternating training scheme for dynamic inference helps to maximize the inner product of gradients between the consecutive mini-batch iterations (i.e., sub-nets in our case).
However, even with the help of the implicit regularization, we find that negative inner products still remain, which represents the conflicting gradient direction between each two sub-nets.
Such phenomenon causes the interference of weight update, while it is helpful to escape from local minima and improve generalization of optimization for individual network training by mini-batch SGD [6, 1].
To mitigate the interference of weight updates between sub-nets but without losing the generalization, we further propose gradient correction to remove the conflicting gradient direction between sub-nets only within inner-group iterations, which is defined as N consecutive mini-batch iterations, where N is the number 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 training efficiency with sparsity (against dynamic inference) and inference dynamics with multiple sparse sub-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 accuracy while training multiple sparse sub-nets all at once (reduced training cost).
Furthermore, compared to the dynamic inference method, AST consistently achieves better accuracy on all the experiments.
Related Works and Background 2.1 Sparse Training
As summarized in [8], sparse training works mainly can be categorized into three groups according to the time point the sparsity is applied: 1) Post-training sparsification [3, 4, 9, 19, 14], which removes the 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 topology is 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 reflection of the gradient can be used to correct the final pruning decision, leading to better accuracy compared to post-training.
Motivated by this, the prune-and-regrow technique [3] periodically removes the unimportant non-zero weights from the sparse model and regrows certain pruned weights back during each mini-batch iteration of training process.
As a representative work, RigL [4] first prunes a certain ratio r based on weight magnitude as:
[Equation]
) where TopK ( v, k ) returns the weight/activation tensor retaining the top k -proportion of elements from v , and s is the targeted sparsity ratio.
After that, the r proportion of new connections are re-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 sparsity ratio r during the entire training process.
Based on this, recent works develop different weight importance criterion to perform prune-and-regrow.
SNFS [3] uses the momentum magnitude of weight to do pruning.
MEST [26] uses the sum of weight and gradient magnitude to indicate the importance 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.
In addition, more similar to our work, [19, 16] propose to alternately train a dense model and one of its sparse variant (i.e., sub-net) returning both an accurate dense model and a sub-net.
However, the rationality behind it is not clearly described.
Overall, the objective of our work is obtaining multiple sparse 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 sparse training ensemble method to independently generate multiple sparse sub-nets for ensemble.
Such method 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 of multiple sub-nets to perform inference independently.
Since the weights of the sub-nets are partially shared 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] and its optimized counterpart (US-Net) [22], which train a single neural network executable at different channel-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 with un-compromised individually inference accuracy, the structures of sub-nets are defined by a subset rule : the smaller sub-net must be the completely subset of the larger sub-net.
Specifically, this line of works train a network in a joint training fashion, which can be expressed as:
[Equation]
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 output of sub-net parameterized by { W i } .
However, such method consumes much larger training time compared to training a single individual model, since all sub-nets have to perform forward and backward once in each mini-batch iteration as shown in Eq.
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+
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], which updates 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 is trained in one iteration.
Note that, the sub-net is defined a part weight of the network, which can run inference 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< sub-net 2 < sub-net 3), such training scheme will start to train the sub-net 1 in the first iteration followed 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 current sparse sub-net), then repeatedly switches sub-nets every three consecutive iterations until the training completes.
The rationale behind AST is inspired by the findings in Reptile [18], which is a meta-learning algorithm originally aiming to learn an initialization of a model from multiple tasks for fast adaption on new tasks.
It shows that mini-batch SGD imposes an implicit regularization to maximize the dot-product of consecutive mini-batches.
In the following, we will show how this proof adapts to our case.
Assuming that we alternately train two sub-nets, sub-net 1 and sub-net 2 with two consecutive mini-batches B 1 and B 2 respectively, according to the Taylor Series, the gradient of sub-net 2 gcalculated 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 be expressed as:
[Equation]
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 learning generalization.
Such proof can be easily adapted to more than two sub-nets.
The detailed explanation of Eq. 4 is relegated into the Appendix.
3.2 Proposed AST with Gradient Correction on Consecutive Inner-group Sub-nets (AST-GC)
Although the implicit regularization of mini-batch SGD helps to maximize the inner product gradient of consecutive sub-nets as discussed in Section 3.1, we indeed observe that there still remain partial number of negative inner product during training, causing conflicting direction of weight updates among sub-nets as shown in 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 saddle
[Figure]
points or local minima and improve the generalization as well [6, 1].
Due to that, inspired by the gradient projection method, which is used in multi-task learning [25] and continual learning [12],
[Figure]
we propose a gradient correction technique to remove the conflicting gradients within the inner-group iterations during training, meanwhile allow the gradient to have negative direction for the inter-group sub-nets.
As shown in Figure 3, the inner-group iterations is defined as the consecutive sub-nets within N mini-bath iterations, where N is the total number of sub-nets (e.g.,
N =3 as shown in Figure 3).
In contrast, inter-group represents the relationship between inner-group iterations.
Specifically, the gradient correction adapts a simple procedure within the inner-group sub-nets: if the gradients between two consecutive sub-nets are in conflict (i
[Figure]
e., their cosine similarity is negative), we project the gradient of current sub-net onto the normal plane of the gradient of its prior sub-net.
Otherwise, the original gradient is unchanged.
Considering two sub-nets in inner-group iterations with the gradient g i and g j respectively, the technique can be mathematically illustrated as:
[Equation]
By doing so, the proposed gradient correction within inner-group iterations have two benefits: 1) it could guarantee that all sub-nets within the consecutive inner-group are updated in the non-conflicting direction, reducing the weight interference between each other; 2) the normal update using the mini-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.
Following most sparse training works, we adopt the prune-and-regrow scheme from GraNet [14] as the backbone technique.
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-net specific 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 questions arises: 1) What is the optimal architecture relationship between sub-nets?
How the pruning should be scheduled for each sub-net?
We address these two questions with the following observations:
Observation 1: Enabling the freedom of exploring unique architectures of each sub-net elevates the learning 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
[Figure]
[Table]
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 relationships 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 architecture freedom.
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 computed by XORing the binary sparse masks between sub-nets, the percentages of “1” in the resultant tensor represents the level of non-overlap.
Apparently, the distinction of connections is larger than the sparsity difference, which implies the existence of the unique connections generated by the ND prune-and-regrow in each sub-net.
the non-disjoint (ND) AST scheme.
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 compared to the conservative completely-subset (CS) training.
Thus, in this work, we use ND-scheme for all experiments.
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 updated (e.g., 1,000 iterations) based on a predefined sparsity schedule [14].
Regarding the
AST scheme, the sub-net model architectures are consecutively switched and trained, but the successive architecture switching does not imply the necessity of consecutive sparsity update of each sub-net.
On the contrary, the intensive sparsity increment of all sub-nets could destabilize the training.
In this work, the sparsity of each sub-net increases periodically, but the spar-
[Table]
sity increment of each sub-net is intermittently performed with the adjustment period ∆ τ .
In the meantime, sub-nets are still consecutively switched during the ∆ τ .
Figure 7 shows the CIFARvalidation accuracy of an AST-trained ResNet-32 sub-net with final sparsity of 98%.
During the early stage of training with high learning rate, the stabilized validation accuracy indicates the effectiveness of the extended adjustment.
Experiments 4.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 employedepochs for CIFAR-10/100 and 100 epochs for ImageNet training.
The multiple sparse sub-nets are trained from scratch and pruned by the proposed AST algorithm.
The cosine annealing learning rate schedule is employed with 0.1 initial learning rate for the CIFAR datasets training.
For ImageNet dataset, we introduce extra 5 epochs for the initial warmup training.
The pruning candidates are globally selected while the first layer remains dense.
Regarding the regrowing process, the percentage of regrow candidates gradually decrease from 0.5 to 0.0 with cosine annealing schedule.
In the meantime, 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 tuning could lead to better accuracy, we choose the above scheme for simplicity and reproducibility.
The reported sub-net accuracy values are obtained from a single background model checkpoint.
In all experiments, 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 AST algorithm, where we used the wide ResNet-32 model that was employed in [26].
Following the typical 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 high initial sparsity is beneficial to the overall memory saving of the training process, we report the results with both dense
( s i = 0% ) and highly sparse ( s i = 90% ) initial models.
Compared to the prior individually-trained SoTA works [14, 26], the proposed AST+GC algorithm obtains three highly sparse 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 and latency benefits.
The joint training scheme [24] requires multiple forward pass in each iteration, the increased computational cost leading to the averaged-yet-suboptimal performance.
In addition to the reduced training cost, the proposed AST method outperforms the joint-training scheme [24] by up to 1.3% and 2.9% inference accuracy, while achieving up to 2.63 × training cost reduction on CIFARand CIFAR-100 datasets, respectively.
We also verify the AST scheme with more sub-nets and large sparsity 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 to
[Table] [Table]
98% and 92 .
23% weight sparsity, the averaged overall sparsity (82.6%) is less than our targeted sparsity (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 training scheme, 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 × .
In addition, compared to AST without gradient correction (GC), AST+GC improves accuracy by 0.6% and 0.8% respectively.
The experimental results with different 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 depict the 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 structured fine-grained sparsity [27].
Varying the group size M and the number of sparse elements N leads to different overall sparsity values.
Motivated by this, we extend the proposed AST algorithm to train multiple sub-nets with different N : M sparsity patterns all at once.
Specifically, the prune-and-regrow is performed based on the group-wise summed weight or gradient magnitude.
Powered by the open-sourced Nvidia-ASP library, the convolution computation can be accelerated when the sparse weight group size M is divisible by 4 (e.g., 4, 8, 16).
The proposed AST algorithm collectively trains 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 the individually trained dense model, the proposed AST scheme achieving up to 2.3 × inference speed up on GPU with 4 × less training efforts and negligible accuracy loss, as shown in Table 14.
The inference 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 scheme for 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 the individually-trained baseline sparse model, training more sub-nets with AST reduces the total training 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-nets have the best overall performance.
The flexibility of training a large number of sub-nets all at once enables the proposed AST algorithm to select suitable architecture for different power budgets.
[Table] [Table]
AST sub-nets with different sparsity difference.
Besides the collective training of highly sparse models, learning the largely-varied subset architectures is also essential.
As shown in Table 5 and
[Table]
the proposed AST algorithm is still able to optimize the model performance with the large sparsity gaps (e.g., 50% vs. 95%).
Extended AST training efforts.
As shown in Figure 2, AST iteratively selects different sub-net for each mini-batch iteration.
Even though the batch shuffling is employed to the training, each sub-net cannot be fully trained by the whole training set inside each epoch.
A straight-forward solution is extending the total training efforts (epochs).
Assume the unit training cost is 160 epochs (1 × ), we evaluate 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 × ), AST achieves 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 sparse neural networks all at once.
Then, we demonstrate the benefits of gradient correction (GC) with theoretical 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 ensembled training steps [24].
Inspired by the prune-and-regrow scheme [14], the proposed AST-GC scheme exploits multiple sparse sub-nets simultaneously while achieving comparable or even higher accuracy against the individually-trained SoTA methods.
Obtaining multiple models with AST brings practical benefits 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 enlarge the practical efficiency for both training and dynamic inference, we will further apply quantization and investigate more structured sparsity pattern upon AST.
Currently, we have demonstrated the practical value for dynamic inference on GPU with Nividia Ampere architecture by using N:M structured sparsity pattern.
We would like to further demonstrate the actual efficiency of AST on more hardware platforms, such as CPU, mobile/IoT devices, specific hardware accelerators, etc.
Acknowledgments.
We thank the anonymous reviewers for their constructive comments.
This work is supported in part by the National Science Foundation under Grant No.1931871, No.2144751,
No.1652866, No.1931871, No.2144751, and the Center for Brain-inspired Computing (C-BRIC) in
JUMP
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how 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 referencing the 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 are proprietary.
• 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 instructions block and only keep the Checklist section heading above along with the questions/answers below.
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] (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 to them?
[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 they were chosen)?
[Yes] See section 4.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See appendix 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[No] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/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, if applicable?
[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 amount spent 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
Appendix 6.1 Algorithm
The pesudocode of AST is presented in Algorithm 1.
To generate and sparsely train each sub-net, AST adapts 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 of gradual pruning t 0 and t f , pruning iterations n , the pruning rate of each pruning iteration is defined as:
[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 gradient inner 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 maximize the inner product of two consecutive mini-batches.
Thus, it proves that the proposed AST has implicit regularization 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 accuracy on ImageNet.
The report accuracy of Rigl [4] uses 4096 batchsize and trains the model on 100 epochs with initial learning rate 1.6.
GraNet [14] uses 64 bathsize, 100 training epochs, and set the initial learning rate to 0.1.
Mest [26] uses a larger 2048 batchsize and trains the model on 150 epochs with the initial learning rate 2.048.
In this case, to evaluate our method and fair compare with them, we conduct the basic training settings like the original ResNet by using 256 batchsize and training the model on 100 epochs with 0.1 initial learning rate.
We believe a more fine-grained hyper-parameter setting could lead to better accuracy.
We only run AST on ImageNet once due to the limited resources by 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.8.
Compared to 80% initial sparsity, 50% initial sparsity could achieve 0.6% accuracy gain on both 80% and 90% prune ratio.
Furthermore, we also conduct the experiment on the smaller ResNet-18, which achieve 62.3% and 62.1% accuracy on 80% and 90% prune ratio respectively.
6.5 Training cost comparison
In addition to the training rounds comparison, we analyze the training cost of the proposed AST method in terms of the detailed metrics.
Table 9 summarizes the number of total FLOPs of training multiple sparse deep neural networks.
Compared to the joint-training scheme [24] or seperately trained GraNet [14], the proposed AST algorithm achieves up to 2.63 × training cost reduction, while maintaining the similar inference accuracy as the individual training baseline.
Furthermore, the detailed 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 original manuscript.
As a result, on CIFAR-10 dataset, the ND scheme has the outperformed performance compared 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 frequency of sub-nets switching.
As shown in Table R1, compared to the smaller adjustment period (i
[Table]
e., ∆ τ = 0 , 90 ), ∆ τ = 300 achieves the best accuracy on all three sparsity levels.
The reason is that performing sub-net switching frequently elevates the instability of model optimization.
[Table]
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, while keep the overall training cost (time) as same as AST.
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-tune the dense model to achieve 90%, 95%, and 98% sparsity with minimum efforts.
As shown in Table 11, 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 the overall training costs.
In addition to the dense model fine-tuning, we address the second concern of the reviewer by performing the sparse progressive training while keeping the overall training cost to be the same as a single AST training.
With the wide ResNet-32 model, we first sparsify the model to 90% sparsity from 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 is the same (60+50+50=160 epochs).
As shown in Table 12, such an individual pruning method failed to achieve the performance as the proposed AST training method.
The large accuracy gap suggests the necessity of the proposed alternative sparsification training.
[Table] [Table]
Furthermore
[Table]
we investigate the impact of fine-tuning based on a pre-trained sparse model.
We first fully 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 13, fine-tuning the 90% sparse model 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 11, Table 12, and Table 13 suggest that it is difficult for individual fine-tuning to achieve the level of high sparsity and high accuracy as the proposed AST, regardless of the initial sparsity of the inherited model checkpoint
[Table]
[Table] [Table]