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 is conceived in the true spirit of sparse training (i.e., start with a sparse neural network; train at any moment in time  a sparse neural network and don’t use information from non-existing connections; do inference with the trained sparse neural network) to enable neural networks training at the edge on low-resource devices. In order to do so, the paper proposes a new sparse training algorithm with dynamic sparsity, named MEST, with a couple of variations for various settings (e.g., hard or soft memory constraints). Also, MEST, in contrast with the typical sparse training algorithms,  introduces the magnitude combined with the gradient as a metric to detect unimportant connections. Nevertheless, the paper is also considering data efficiency for the first time in sparse training (up to my knowledge). The extensive and rigorous empirical validation shows that MEST and its variants reach state-of-the-art performance in sparse training. Notably, even if it is not one of the paper claims, it seems that MEST can outperform dense training on convolutional neural networks using just half of the computational resources. Convolutional layers being perhaps the last type of neural network layer on which it is still arguable in the literature if truly sparse training can outperform dense training (see for details Table 2, left column, ResNet-50 on ImageNet).