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 work aims to combine ideas presented in recent dynamic sparse training (DST) methods with gradual magnitude pruning. Additionally the paper has some novel experiments that confirm previous work on when-is-best to prune during training. Although an interesting and potentially impactful idea, the story and results need a bit more work.