Sparse Training via Boosting Pruning Plasticity  with Neuroregeneration

1 Introduction . Neural network pruning is the most common technique to reduce the parameter count , storage requirements , and computational costs of modern neural network architectures . Recently , posttraining pruning [ 49 , 29 , 18 , 47 , 10 , 54 , 74 , 5 , 57 , 75 ] and before-training pruning [ 31 , 30 , 67 , 63 , 6 , 11 ] have been two fast-rising fields , boosted by lottery tickets hypothesis ( LTH ) [ 10 ] and singleshot network pruning ( SNIP ) [ 31 ] . The process of post-training pruning typically involves fully pre-training a dense network as well as many cycles of retraining ( either fine-tuning [ 18 , 17 , 39 ] or rewinding [ 12 , 54 ] ) . As the training costs of the state-of-the-art models , e.g. , GPT-3 [ 4 ] and FixEfficientNet-L2 [ 64 ] have exploded , this process can lead to a large amount of overhead cost . Recently emerged methods for pruning at initialization significantly reduce the training cost by identifying a trainable sub-network before the main training process . While promising , the existing methods fail to match the performance achieved by the magnitude pruning after training [ 11 ] . ∗Partial of this work have been done when Shiwei Liu worked as an intern at JD Explore Academy . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Compared with the above-mentioned two classes of pruning , during-training pruning is a class of methods that reap the acceleration benefits of sparsity early on the training and meanwhile achieve promising performance by consulting the information obtained during training . There are some works [ 77 , 13 , 33 ] attempting to gradually prune the network to the desired sparsity during training , while they mainly focus on the performance improvement . Up to now , the understanding of duringtraining pruning has been less explored due to its more complicated dynamical process , and the performance gap still exists between pruning during training and full dense training . To better understand the effect of pruning during the optimization process ( not at inference ) , we study the ability of the pruned models to recover the original performance after a short continued training with the current learning rate , which we call pruning plasticity ( see Section 3.1 for a more formal definition ) . Inspired by the neuroregeneration mechanism in the nervous system where new neurons and connections are synthesized to recover the damage in the nervous system [ 26 , 41 , 73 ] , we examine if allowing the pruned network to regenerate new connections can improve pruning plasticity , and hence contribute to pruning during training . We consequently propose a parameter-efficient method to regenerate new connections during the gradual pruning process . Different from the existing works for pruning understanding which mainly focus on dense-to-sparse training [ 42 ] ( training a dense model and prune it to the target sparsity ) , we also consider sparse-to-sparse training ( training a sparse model yet adaptively re-creating the sparsity pattern ) which recently has received an upsurge of interest in machine learning [ 44 , 3 , 9 , 48 , 8 , 37 , 36 ] . In short , we have the following main findings during the course of the study : # 1 . Both pruning rate and learning rate matter for pruning plasticity . When pruned with low pruning rates ( e.g. , 0.2 ) , both dense-to-sparse training and sparse-to-sparse training can easily recover from pruning . On the contrary , if too many parameters are removed at one time , almost all models suffer from accuracy drops . This finding makes a connection to the success of the iterative magnitude pruning [ 10 , 54 , 5 , 6 , 65 ] , where usually a pruning process with a small pruning rate ( e.g. , 0.2 ) needs to be iteratively repeated for good performance . Pruning plasticity also gradually decreases as the learning rate drops . When pruning happens during the training phase with large learning rates , models can easily recover from pruning ( up to a certain level ) . However , pruning plasticity drops significantly after the second learning rate decay , leading to a situation where the pruned networks can not recover with continued training . This finding helps to explain several observations ( 1 ) for gradual magnitude pruning ( GMP ) , it is always optimal to end pruning before the second learning rate drop [ 77 , 13 ] ; ( 2 ) dynamic sparse training ( DST ) benefits from a monotonically decreasing pruning rate with cosine or linear update schedule [ 8 , 9 ] ; ( 3 ) rewinding techniques [ 12 , 54 ] outperform fine-tuning as rewinding retrains subnetworks with the original learning rate schedule whereas fine-tuning often retrains with the smallest learning rate . # 2 . Neuroregeneration improves pruning plasticity . Neuroregeneration [ 41 , 73 ] refers to the regrowth or repair of nervous tissues , cells , or cell products . Conceptually , it involves synthesizing new neurons , glia , axons , myelin , or synapses , providing extra resources in the long term to replace those damaged by the injury , and achieving a lasting functional recovery . Such mechanism is closely related to the brain plasticity [ 51 ] , and we borrow this concept to developing a computational regime . We show that , while regenerating the same number of connections as pruned , the pruning plasticity is observed to improve remarkably , indicating a more neuroplastic model being developed . However , it increases memory and computational overheads and seems to contradict the benefits of pruningduring-training . This however raises the question : can we achieve efficient neuroregeneration during training with no extra costs ? We provide an affirmative answer to this question . # 3 . Pruning plasticity with neuroregeneration can be leveraged to substantially boost sparse training performance . The above-mentioned findings of pruning plasticity can generalize to the final performance level under a full continued training to the end . Imitating the neuroregeneration behavior [ 41 , 73 ] , we propose a new sparse training method – gradual pruning with zero-cost neuroregeneration ( GraNet ) , which is capable of performing regeneration without increasing the parameter count . In experiments , GraNet establishes the new state-of-the-art performance bar for dense-to-sparse training and sparse-to-sparse training , respectively . Particularly , the latter for the first time boosts the sparse-to-sparse training performance over various dense-to-sparse methods by a large margin without extending the training time , with ResNet-50 on ImageNet . Besides the consistent performance improvement , we find the subnetworks that GraNet learns are more accurate than the ones learned by the existing gradual pruning method , providing explanations for the success of GraNet . 2 Related Work . Post-Training Pruning . Methods that yield a sparse neural network from a pre-trained network by pruning the unimportant weights or neurons , to the best of our knowledge , were proposed in [ 24 ] and [ 50 ] . After that , various pruning methods have emerged to provide increasingly efficient methods to identify sparse neural networks for inference . The pruning criterion includes weight magnitude [ 18 , 10 ] , gradient [ 61 ] Hessian [ 29 , 19 , 59 ] , Taylor expansion [ 47 , 46 ] , etc . Low-rank decomposition [ 7 , 23 , 17 , 71 ] are also used to induce structured sparsity in terms of channels or filters . Most of the above-mentioned pruning methods require many pruning and re-training cycles to achieve the desired performance . During-Training Pruning . Instead of inheriting weights from a pre-trained model , some works attempt to discover well-performing sparse neural networks with one single training process . Gradual Magnitude Pruning ( GMP ) , introduced in [ 77 ] and studied further in [ 13 ] , gradually sparsifies the neural network during the training process until the desired sparsity is reached . Besides , [ 40 ] and [ 68 ] are prior works that enforce the network to sparse during training via L0 and L1 regularization , respectively . [ 60 , 34 , 55 , 70 , 28 ] moved further by introducing trainable sparsity heuristics to learn the sparse masks and weights simultaneously . These methods are all classified as dense-to-sparse training as they start from a dense network . Dynamic Sparse Training ( DST ) [ 44 , 3 , 48 , 8 , 9 , 36 , 35 , 25 ] is another class of methods that prune models during training . The key factor of DST is that it starts from a random initialized sparse network and optimizes the sparse topology as well as the weights simultaneously during training ( sparse-to-sparse training ) . Without an extended training time [ 37 ] , sparse-to-sparse training usually falls short of dense-to-sparse training in terms of the prediction accuracy . For further details , see the survey of [ 43 , 21 ] . Before-Training Pruning . Motivated by SNIP [ 31 ] , many works [ 67 , 63 , 6 ] have emerged recently to explore the possibility of obtaining a trainable sparse neural network before the main training process . [ 11 ] demonstrates that the existing methods for pruning at initialization perform equally well when the unpruned weights are randomly shuffled , which reveals that what these methods discover is the layer-wise sparsity ratio , rather than the indispensable weight values and positions . Our analysis shows that both the mask positions and weight values are crucial for GraNet . 3 Methodology for Pruning Plasticity . The primary goal of this paper is to study the effect of pruning as well as neuroregeneration on neural networks during the standard training process . Therefore , we do not consider post-training pruning and before-training pruning . Below , we introduce in detail the definition of pruning plasticity and the experimental design that we used to study pruning plasticity . 3.1 Metrics . Let us denote Wt ∈ Rd as the weights of the network and mt ∈ { 0 , 1 } d as the binary mask yielded from the pruning method at epoch t. Thus , the pruned network can be denoted as Wt mt . Let T be the total number of epochs the model should be trained . Let CONTRAINk ( Wt mt , a ) refers to the function that continues to train the pruned model for k epochs with the learning rate schedule a . Definition of Pruning plasticity . We define pruning plasticity as tCONTRAINk ( Wt mt , at ) − tPRE , where tPRE is the test accuracy measured before pruning and tCONTRAINk ( Wt mt , at ) is the test accuracy measured after k epoch of continued training CONTRAINk ( Wt mt , at ) . Specifically , to better understand the effect of pruning on the current model status and to avoid the effect of learning rate decay , we fix the learning rate as the one when the model is pruned , i.e , at . This setting is also appealing to GMP [ 77 , 13 ] and DST [ 44 , 9 , 48 , 37 ] in which most of the pruned models are continually trained with the current learning rate for some time . Final performance gap . Nevertheless , we also investigate the effect of pruning on the final performance , that is , continually training the pruned networks to the end with the remaining learning rate schedule CONTRAINT−t ( Wt mt , a [ t+1 : T ] ) . In this case , we report tCONTRAINT−t ( Wt mt , a [ t+1 : T ] ) − tFINAL , where tFINAL is the final test accuracy of the unpruned models .

Summary: This paper understands the during-training pruning from the perspective of training plasticity which is the ability of the pruned networks to recover the original performance. A technique named neuroregeneration is further shown to improve the training plasticity. Based on these findings, a gradual magnitude pruning (GMP) method with neuroregeneration (GraNet), and its dynamic sparse training variant are proposed. Experiments show the effectiveness of GradNet and GradNet-ST.