MEST: Accurate and Fast Memory-Economic Sparse Training Framework on the Edge

1 Introduction . To promote the broader applications of deep learning on the edge , a surge of research efforts have been devoted to removing the over-parameterization in neural networks for accelerated inference . Specifically , existing works have explored various strategies such as heuristics-based pruning [ 1 , 2 ] , regularization-based pruning [ 3 , 4 ] , and recently prevailing network architecture search [ 5 , 6 , 7 , 8 ] . Recently , a new trend of exploring sparsity for training acceleration of neural networks has emerged to embrace the promising training-on-the-edge paradigm . The first works in this direction use the pruning-at-initialization approach such as SNIP [ 9 ] and GraSP [ 10 ] that first obtains a fixed sparse model structure and then follows with a traditional training process . However , the whole process is still computation- and memory-intensive , and therefore not compatible with the end-to-end edge training paradigm . Such a sparse training methodology with the pre-fixed structure also faces the problem of compromised accuracy . † These authors contributed equally . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Furthermore , sparse training with dynamic sparsity mask such as SET [ 14 ] , DeepR [ 15 ] , DSR [ 13 ] , RigL [ 11 ] , and SNFS [ 12 ] have been proposed , showing great potential towards end-toend edge training . Specifically , they start with a sparse model structure picked intuitively from the initialized ( but not trained ) dense model , and then heuristically explore various sparse topologies at the given sparsity ratio , together with the sparse model training . The underlying principle of sparse training is that the total epoch number is the same as dense training , but the speed of each training iteration ( batch ) is significantly improved , thanks to the reduced computation amount due to sparsity . This paper proposes a novel Memory-Economic Sparse Training ( MEST ) framework targeting for accurate and fast execution on edge devices . Specifically , we boost the accuracy with the MEST+EM ( Elastic Mutation ) to effectively explore various sparse topologies . And we propose another enhancement through applying a soft memory bound , namely , MEST+EM & S , which relaxes on the memory footprint during sparse training with the target sparsity ratio met by the end of training . In Figure 1 , the accuracy of the proposed MEST is compared against state-of-the-art ( SOTA ) sparse training algorithms under different sparsity ratios , using the same unstructured sparsity scheme . Furthermore , as with inference acceleration , we find that sparse training closely relates to the adopted sparsity scheme such as unstructured [ 16 ] , structured [ 17 , 18 ] , or fine-grained structured [ 19 ] scheme , which can result in varying accuracy , training speed , and memory footprint performance for sparse training . With our effective MEST framework , this paper systematically investigates the sparse training problem with respect to the sparsity schemes . Specifically , besides the directly observable model accuracy , we conduct thorough analysis on the memory footprint by different sparsity schemes during sparse training , and in addition , we measure the training speed performance under various schemes on the mobile edge device . On top of that , the paper proposes to employ data efficiency for further acceleration of sparse training . The prior works [ 20 , 21 , 22 , 23 ] show that the amount of information provided by each training example is different , and the hardness of having an example correctly learned is also different . Some training examples can be learned correctly early in the training stage and will never be “ forgotten ” ( i.e. , misclassified ) again . And removing those easy and less informative examples from the training dataset will not cause accuracy degradation on the final model . However , the research of connecting training data efficiency to a sparse training scene is still missing , due to the dynamic sparsity mask . In this work , we explore the impact of model sparsity , sparsity schemes , and sparse training algorithm on the amount of removable training examples . And we also propose a data-efficient two-phase sparse training approach to effectively accelerate sparse training by removing less informative examples during the training process without harming the final accuracy . The contributions of this work are summarized as follows : • A novel Memory-Economic Sparse Training ( MEST ) framework with enhancements by Elastic Mutation ( EM ) and the soft memory bound targeting for accurate and fast execution on the edge . • A systematic investigation of the impact of sparsity scheme on the accuracy , memory footprint , as well as training speed with real edge devices , providing guidance for future edge training paradigm . • Exploring the training data efficiency in the sparse training scenario for further training acceleration , by proposing a two-phase sparse training approach for in-situ identification and removal of less informative examples during the training without hurting the final accuracy . • On CIFAR-100 with ResNet-32 , comparing with representative SOTA sparse training works , i.e. , SNIP , GraSP , SET , and DSR , our MEST increases accuracy by 1.91 % , 1.54 % , 1.14 % , and 1.17 % ; achieves 1.76× , 1.65× , 1.87× , and 1.98× training acceleration rate ; and reduces the memory footprint by 8.4× , 8.4× , 1.2× , and 1.2× , respectively . 2 Background and Related Work . This section introduces representative neural network sparsity schemes and their impacts on memory footprint in sparse training , as well as the neural network sparse training strategies . 2.1 Sparsity Scheme . Neural network pruning has been well investigated for inference acceleration . The majority of works in this direction apply a pretraining-pruning-retraining flow , which is not compatible with the trainingon-the-edge paradigm . According to the adopted sparsity scheme , those works can be categorized as unstructured [ 16 , 1 ] , structured [ 24 , 2 , 25 , 26 , 17 , 3 , 27 , 28 , 29 , 30 , 31 , 18 , 32 , 33 ] , and fine-grained structured [ 19 , 34 , 35 , 36 , 37 , 38 , 39 , 40 , 41 ] including the pattern-based and block-based ones . Detailed discussion about these sparsity schemes is provided in Appendix A . Although these sparsity schemes are mainly proposed for accelerating inference , we find that they also play an important role in sparse training in terms of accuracy , memory footprint , and training speed . For memory footprint , we focus on the two major components that vary with the sparsity scheme : model representation together with gradients produced during training . The detailed analysis is provided in Appendix B . 2.2 Sparse Training . Majority of the sparse training works can be categorized into two groups : fixed-mask sparse training and dynamic-mask sparse training . Additionally , there exist works [ 42 , 43 , 44 ] that prune dense networks in the early training stage , but they are out of scope for sparse training on the edge . 2.2.1 Sparse Training with Fixed Sparsity Mask . The fixed-mask approach [ 9 , 10 , 45 , 46 , 47 ] has been proposed to decouple pruning and training such that after pruning , the sparse model training can be executed on edge devices . SNIP [ 9 ] preserves the loss after pruning based on connection sensitivity . GraSP [ 10 ] prunes connections in a way that accounts for their role in the network ’ s gradient flow . SynFlow [ 45 ] proposes iterative synaptic flow pruning , which avoids layer collapse and preserves the total flow of synaptic strengths throughout the network . Since the proposed pruning algorithm does not incorporate back propagation , it achieves global pruning at initialization without data . Based on the unstructured SNIP objective , 3SP [ 47 ] further introduces a computation-aware weighting of the pruning score . This actively biases pruning by removing more computation-intensive channels which either have small effect on the loss or have significant computation cost . However , the pruning-at-initialization process is still computation and memory-intensive and therefore not compatible with the end-to-end sparse training on the edge . And these works ( except 3SP ) employ the unstructured sparsity scheme . 2.2.2 Sparse Training with Dynamic Sparsity Mask . To reduce the computation as well as memory footprint during the whole training phase , sparse training is exploited in many works [ 15 , 14 , 13 , 12 , 11 ] , which can adjust the sparsity topology during training as well as maintain a low memory footprint . Sparse Evolutionary Training ( SET ) [ 14 ] uses magnitude-based pruning and random growth at the end of each training epoch . DeepR [ 15 ] combines dynamic sparse parameterization with stochastic parameter updates for training . This method is primarily demonstrated with sparsification of fully-connected layers and applied to relatively small and shallow networks . DSR [ 13 ] develops a dynamic reparameterization method to achieve high parameter efficiency in training sparse deep residual networks . SNFS [ 12 ] develops sparse momentum , an algorithm which uses exponentially smoothed gradients ( momentum ) to identify layers and weights which reduce the error efficiently . RigL [ 11 ] proposes to iteratively update sparse model topology during training by calculating dense gradients only at the update step . Note that , though SNFS and RigL are sparse training , they actually involve computation of all the gradients corresponding to both pruned and non-zero weights , and therefore their memory footprint is equivalent to that of dense training . 3 Sparse Training on the Edge . 3.1 Sparsity Scheme in Sparse Training on the Edge . It is common to see that sparse training works [ 9 , 10 , 15 , 14 , 13 , 12 , 11 , 45 ] represent the training speed performance using the training FLOP count . Actually , such FLOPS can not account for the actual execution overheads caused by the sparse data operations . For example , the unstructured sparsity exhibits an irregular memory access pattern , leading to significant execution overhead . Moreover , the dense model can take advantage of Winograd [ 48 ] to significantly accelerate the computation speed , which may apply in sparse computation . Therefore , this paper directly measures the training speed performance on a mobile device . We investigate the training acceleration performance by different sparsity schemes including unstructured , structured , and two state-of-the-art fine-grained structured ( i.e. , block and pattern ) , through a prototype implementation on a mobile edge device . To do so , we conduct compiler-level optimizations , leveraging a computation-graph based compilation approach , including optimizations on computation graph itself , tensor optimizations , etc . More details are provided in Appendix C. Figure 2 shows representative training speed performance under various sparsity schemes , using an example CONV layer adopted from ResNet32 . We evaluate the acceleration rate by measuring the total forward and backward execution time of the sparse CONV layer on a Samsung Galaxy S20 smartphone , and then normalizing with respect to that of the corresponding dense layer . Surprisingly , even under the same sparsity ratio , we found that the acceleration rates of different sparsity schemes are varied significantly . When the sparsity ratio is below 70 % and 80 % for block-based sparsity and unstructured sparsity schemes , respectively , they can not achieve any acceleration and even slow down the computation speed , compared to the corresponding dense model . Thus , choosing an appropriate sparsity scheme is an essential factor for sparse training for acceleration purposes .

Summary: This paper proposes the memory-economic sparse training (MEST) framework that targets efficient sparse training on edge devices. The author proposes an elastic mutation (EM) and a soft memory bound algorithm for sparse training. Their training methods can keep the entire training process sparse and avoiding involve dense forward/backward computation. This paper also investigates the impact of different sparsity schemes on sparse training performance in accuracy and speed acceleration. Besides, the paper also introduces data-efficient training in sparse training scenes. Compared to previous works, this work is more hardware friendly and can achieve similar even higher accuracy. These proposed methods are more hardware-friendly and achieve higher accuracy. The proposed framework is promising to be implemented for real-world sparse training acceleration. This paper can be impactful for both the community of academia and industry.