Spending Your Winning Lottery Better After Drawing It

Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2019 ) suggests that a dense neural network contains a sparse sub-network that can match the performance of the original dense network when trained in isolation from scratch . Most works retrain the sparse sub-network with the same training protocols as its dense network , such as initialization , architecture blocks , and training recipes . However , till now it is unclear that whether these training protocols are optimal for sparse networks . In this paper , we demonstrate that it is unnecessary for spare retraining to strictly inherit those properties from the dense network . Instead , by plugging in purposeful “ tweaks ” of the sparse subnetwork architecture or its training recipe , its retraining can be significantly improved than the default , especially at high sparsity levels . Combining all our proposed “ tweaks ” can yield the new stateof-the-art performance of LTH , and these modifications can be easily adapted to other sparse training algorithms in general . Specifically , we have achieved a significant and consistent performance gain of 1.05 % − 4.93 % for ResNet18 on CIFAR-100 over vanilla-LTH . Moreover , our methods are shown to generalize across datasets ( CIFAR-10 , CIFAR-100 , TinyImageNet ) and architectures ( Vgg16 , ResNet-18/ResNet-34 , MobileNet ) . All codes will be publicly available . 1 INTRODUCTION . Deep neural networks ( NN ) have achieved significant progress in many tasks such as classification , detection , and segmentation . However , most existing models are computationally extensive and overparameterized , thus it is difficult to deploy these models in real-world devices . To address this issue , many efforts have been devoted to compressing the heavy model into a lightweight counterpart . Among them , network pruning ( LeCun et al. , 1990 ; Han et al. , 2015a ; b ; Li et al. , 2016 ; Liu et al. , 2019 ) , which identifies sparse sub-networks from the dense networks by removing unnecessary weights , stands as one of the most effective methods . Previous methods ( Han et al. , 2015b ) usually prune the dense network after the standard training process to obtain the sparse sub-networks . The performance of the pruned network , however , decreases heavily as parts of the weights are removed , and retraining is thus required to recover the original performance ( Han et al. , 2015b ) . The recent Lottery Ticket Hypothesis ( LTH ) ( Frankle & Carbin , 2019 ) represents a major paradigm shift from conventional after-training pruning . LTH suggests that a dense NN contains sparse subnetworks , named “ winning tickets ” , which can match the performance of the original dense NN when trained in isolation from scratch . Such a winning ticket can be “ drawn ” by finding the sparse weight mask , from dense training accompanied with iterative magnitude pruning ( IMP ) . The found sparse mask is then applied to the original dense NN , and the masked sparse subnetwork is subsequently re-trained from scratch . Using a similar metaphor , We call the sparse re-training step as “ spending ” the lottery , after it is drawn . In most ( if not all ) LTH literature ( Frankle & Carbin , 2019 ; Frankle et al. , 2019 ; Renda et al. , 2020 ) , the re-training step takes care of the masked subnetwork , which is re-trained with the same initialization ( or rewinding ) and same training recipe as its dense network . In plain language , “ You spend the same lottery ticket in the same way you draw it ” . Recent evidence seems to support this convention by attributing LTH ’ s success in recovering the original pruned solution Evci et al . ( 2020a ) . However , till now it is still unclear that whether the architecture blocks , initialization 1 regimes , or training recipes are necessarily optimal for the sparse network . Our question of curiosity is hence : “ Can you spend the same lottery ticket in a different yet better way than how you draw it ” ? Contrary to the common beliefs , this paper demonstrates that it is unnecessary for sparse network retraining ( “ spending the lottery ” ) to stick to the protocol of dense network training or sparse mask finding ( “ drawing the lottery ” ) . Instead , having sparse re-training purposely misaligned in some way from dense training can make the found sparse subnetwork work even better . Specifically , we investigate two possible aspects of modified sparse re-training : • Architecture tweaking : after the sparse subnetwork is found , we modify the network architecture by : ( a ) injecting more residual skip-connections that are non-existent in dense networks ; and ( b ) changing the ReLU neurons into smoother activation functions such as Swish ( Ramachandran et al. , 2017 ) and Mish ( Misra , 2019 ) . • Training recipe tweaking : when training the ( found or modified ) sparse subnetwork architecture , we modify the training approach by : ( c ) changing the “ lottery ticket initialization ” by learned layer-wise scaling ; and ( d ) replacing the one-hot labels with either naive or knowledge-distilled soft labels . Each idea could be viewed as certain type of learned smoothening ( we will explain later ) , and is plug-and-play in the sparse re-training stage of any LTH algorithm . Those techniques can be applied either alone or altogether , and can significantly boost the sparse re-training performance in large models and at high sparsities . Note that all above “ tweaks ” only affect the sparse re-training stage ( we never alter the found sparse mask ) , but not the dense training/masking finding stage . In fact , our experiments will show that they boost sparse re-training much more than dense counterparts . Our contributions can be summarized as : • In contrast to the common wisdom that LTH sparse re-training needs to inherit ( masked ) network architecture , initialization , and training protocol from dense training , we for the first time demonstrate that purposely re-tweaking them will actually improve the sparse re-training step . That urges our rethinking of the LTH ’ s true value . • We investigate two groups of techniques to tweak the sparse subnetwork architecture and training recipe respectively . For the former , we inject new skip connections and replace new activation neurons . For the latter , we re-scale the initialization and soften the labels . Each of the techniques improves sparse re-training ( much more than they can help dense counterparts ) , and altogether they lead to further boosts . • Our extensive experimental results demonstrate that by plugging these techniques in LTH sparse retraining , we can significantly improve the performance of “ winning tickets ” at high sparsity levels and large models , setting the state-of-the-art performance bar of LTH . Furthermore , we show that they can benefit other sparse training algorithms in general , and provide visualizations to analyze their successes . 2 BACKGROUND WORK . 2.1 THE LOTTERY TICKET HYPOTHESIS . The LTH ( Frankle & Carbin , 2019 ) implies that initialization is the key for sparse network retraining . Beyond image classification ( Liu et al. , 2019 ; Savarese et al. , 2020 ; Wang et al. , 2020 ; You et al. , 2020 ; Ma et al. , 2021 ; Chen et al. , 2020a ) , LTH has also been widely explored in numerous contexts , such as natural language processing ( Gale et al. , 2019 ; Yu et al. , 2019 ; Prasanna et al. , 2020 ; Chen et al. , 2020b ; c ) , reinforcement learning ( Yu et al. , 2019 ) , self-supervised learning ( Chen et al. , 2020a ) , lifelong learning ( Chen et al. , 2021c ) , generative adversarial networks ( Chen et al. , 2021d ; Kalibhat et al. , 2020 ; Chen et al. , 2021a ) , and graph neural networks ( Chen et al. , 2021b ) . However , when retraining the sparse network , these works still strictly follow the same training recipe from dense networks . The most recent work ( Tessera et al. , 2021 ) reveals that focusing on initialization appears insufficient . Optimization strategies , regularization , and architecture choices also play significant roles in sparse network training . However , ( Tessera et al. , 2021 ) only compares sparse networks to their equivalent capacity dense networks , and most of their experiments are con- 2 ducted on multi-layer perceptron ( MLP ) . Thus , it is unclear whether their conclusion can generalize to the ticket finding from dense CNNs . 2.2 SMOOTHNESS . Introducing smoothness into NNs , including on the weights , logits , or training trajectory , is a common techniques to improve the generalization and optimization ( Jean & Wang , 1994 ) . For labels , smoothness is usually introduced by replacing the hard target with soft labels ( Szegedy et al. , 2016 ) or soft logits ( Hinton et al. , 2015a ) . This uncertainty of labels helps to alleviate the overconfidence ( Hein et al. , 2019 ) and improves the generalization . Smoothness can also implemented by replacing the activation functions ( Misra , 2019 ; Ramachandran et al. , 2017 ) , adding skip-connections in NNs ( He et al. , 2016 ) , or averaging along the trajectory of gradient descent ( Izmailov et al. , 2018 ) . These methods contribute to more stable gradient flows ( Tessera et al. , 2021 ) and smoother loss landscapes , but most of them have not been considered nor validated on sparse NNs . 3 METHODOLOGY The LTH ( Frankle & Carbin , 2019 ) suggests that the “ winning tickets ” can be found via the following three steps : ( 1 ) training a dense network from scratch ; ( 2 ) pruning unnecessary weights to form the mask of a sparse sub-network ; and ( 3 ) re-training the subnetwork from the same initialization used in the dense model . For the third step , the retraining of sparse network usually inherits all properties , such as architecture blocks and training recipes , from dense networks . However , our experiments validate that those are not necessarily optimal for training sparse networks . To verify , we investigate two aspects of “ tweaks ” dedicated to the sparse re-training step : model architecture , and training recipe . 3.1 MODEL ARCHITECTURE TWEAKING . Replacing Smoother Activations Most deep NNs apply the Rectified Linear Units ( ReLU ) ( Nair & Hinton , 2010 ) as the activation function . However , the gradient of ReLU changes suddenly around zero , and this non-smooth nature of ReLU is an obstacle to sparse retraining because it leads to high activation sparsity into the subnetwork , likely blocking the gradient flow ( Appendix A1.1 ) . To mitigate this issue and encourage healthier gradient flow , we replace the ReLU to Swish ( Ramachandran et al. , 2017 ) and Mish ( Misra , 2019 ) . Different from ReLU , Swish and Mish are both smooth non-monotonic activation functions . The non-monotonic property allows for the gradient of small negative inputs , which leads to a more stable gradient flow ( Tessera et al. , 2021 ) . Meanwhile , the loss landscape of Swish and Mish are 3 proved to have smoother transition ( Misra , 2019 ) , which makes the loss function easier to optimize and hence makes the sparse network generalize better . Injecting New Skip Connections High sparsity networks easily suffer from the layer-collapse ( Tanaka et al. , 2020 ) , i.e. , the premature pruning of an entire layer . This could make the sparse network untrainable , as the gradient can not be backpropagated through that layer . The skip connection ( or named ” residual-addition ” ) ( He et al. , 2016 ) was initially proposed to avoid gradient vanishing problem , and enables the training of a very deep model . It is later proven that skip connections can smooth the loss surfaces ( Li et al. , 2017b ) . That naturally motivates us to consider using more skip connections on the sparse networks to smoothen the landscape and preserve gradients , besides its possible mitigation effect when encountering collapsed layers . Inspired by ( Ding et al. , 2021 ) , we propose to “ artificially ” add new skip-connections during our sparse re-training . Figure 3 illustrates this architectural modifications to the traditional Resnet-18 block . Similar to existing residual connection in traditional ResNet-18 block , our newly introduced skip-connections add input of each ( 3 × 3 ) convolution block , to their output before the activation . The original motivation behind residual connections comes from their ability to allow gradient information to flow to earlier layers of the NN , thereby reducing the vanishing gradient problem during training ( He et al. , 2016 ) . With high activation sparsity present in sparse subnetworks , additional skip-connections can facilitate healthy gradient flow and improve their trainability . Furthermore , ( Li et al. , 2017b ) observed that with an increase in depth of networks , neural loss landscape becomes highly chaotic and leads to drop in generalization and trainability . They further observed that skip connections promote flatness and prevent transition to chaotic behaviour . Inspired by them , we added skip-connections to stabilize the initial chaotic training regime of sparse NN at high sparsity level and manage to prevent the transition to a sub-optimal behaviour .

Summary: The authors propose several "tweaks" to sparse re-training of lottery tickets that have been identified by iterative magnitude pruning (IMP). These tweaks include  1) the replacement of ReLU activation functions by smooth activation functions, 2) the use of soft-labels instead of one-hot labels,  3) learned layer-wise rescaling of the initial parameters to enable effective gradient steps initially,  4) the injection of new skip connections to avoid layer collapse.  The tweaks are compared and combined with weight rewinding. Experiments on CIFAR10, CIRAR100, any TinyImageNet suggest a superior performance of the tweaks over vanilla IMP.