Selfish Sparse RNN Training

1 INTRODUCTION . Recurrent neural networks ( RNNs ) ( Elman , 1990 ) , with a variant of long short-term memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) , have been highly successful in various fields , including language modeling ( Mikolov et al. , 2010 ) , machine translation ( Kalchbrenner & Blunsom , 2013 ) , question answering ( Hirschman et al. , 1999 ; Wang & Jiang , 2017 ) , etc . As a standard task to evaluate models ’ ability to capture long-range context , language modeling has witnessed great progress in RNNs . Mikolov et al . ( 2010 ) demonstrated that RNNs perform much better than backoff models for language modeling . After that , various novel RNN architectures such as Recurrent Highway Networks ( RHNs ) ( Zilly et al. , 2017 ) , Pointer Sentinel Mixture Models ( Merity et al. , 2017 ) , Neural Cache Model ( Grave et al. , 2017 ) , Mixture of Softmaxes ( AWD-LSTM-MoS ) ( Yang et al. , 2018 ) , ordered neurons LSTM ( ON-LSTM ) ( Shen et al. , 2019 ) , and effective regularization like variational dropout ( Gal & Ghahramani , 2016 ) , weight tying ( Inan et al. , 2017 ) , DropConnect ( Merity et al. , 2018 ) have been proposed to significantly improve the performance of RNNs . At the same time , as the performance of deep neural networks ( DNNs ) improves , the resources required to train and deploy deep models are becoming prohibitively large . To tackle this problem , various dense-to-sparse methods have been developed , including but not limited to pruning ( LeCun et al. , 1990 ; Han et al. , 2015 ) , Bayesian methods ( Louizos et al. , 2017a ; Molchanov et al. , 2017 ) , distillation ( Hinton et al. , 2015 ) , L1 Regularization ( Wen et al. , 2018 ) , and low-rank decomposition ( Jaderberg et al. , 2014 ) . Given a pre-trained model , these methods work effectively to accelerate the inference . Recently , some dynamic sparse training ( DST ) approaches ( Mocanu et al. , 2018 ; Mostafa & Wang , 2019 ; Dettmers & Zettlemoyer , 2019 ; Evci et al. , 2020 ) have been proposed to bring efficiency for both , the training phase and the inference phase by dynamically changing the sparse connectivity during training . However , previous approaches are mainly for CNNs . For RNNs , the long-term dependencies and repetitive usage of recurrent cells make them more difficult to be sparsified ( Kalchbrenner et al. , 2018 ; Evci et al. , 2020 ) . More importantly , the state-of-the-art performance achieved by RNNs on language modeling is mainly associated with the optimizer , averaged stochastic gradient descent ( ASGD ) ( Polyak & Juditsky , 1992 ) , which is not compatible with the existing DST approaches . The above-mentioned problems heavily limit the performance of the off-the-shelf sparse training methods in the RNN field . For instance , while “ The Rigged Lottery ” ( RigL ) achieves state-of-the-art sparse training results with various CNNs , it fails to match the performance of the iterative pruning method in the RNN setting ( Evci et al. , 2020 ) . In this paper , we introduce an algorithm to train sparse RNNs with a fixed number of computational costs throughout training . We abbreviate our sparse RNN training method as Selfish-RNN because our method encourages cell weights to obtain their parameters selfishly . The main contributions of this work are five-fold : • We propose an algorithm to train sparse RNNs from scratch with a fixed number of parameters . This advantage constrains the training costs to a fraction of the costs needed for training a dense model , allowing us to choose suitable sparsity levels for different types of training platforms . • We introduce SNT-ASGD , a sparse variant of the non-monotonically triggered averaged stochastic gradient descent optimizer , which overcomes the over-sparsified problem of the original NT-ASGD ( Merity et al. , 2018 ) caused by dynamic sparse training . • We demonstrate state-of-the-art sparse training performance with various RNN models , including stacked LSTMs ( Zaremba et al. , 2014 ) , RHNs , ordered neurons LSTM ( ONLSTM ) on Penn TreeBank ( PTB ) dataset ( Marcus et al. , 1993 ) and AWD-LSTM-MoS on WikiText-2 dataset ( Melis et al. , 2018 ) . • We present an approach to analyze the evolutionary trajectory of the sparse connectivity optimized by dynamic sparse training from the perspective of graph . With this approach , we show that there exist many good structural local optima ( sparse sub-networks having equally good performance ) in RNNs , which can be found in an efficient and robust manner . • Our analysis shows two surprising phenomena in the setting of RNNs contrary to CNNs : ( 1 ) random-based weight growth performs better than gradient-based weight growth , ( 2 ) uniform sparse distribution performs better than Erdős-Rényi ( ER ) sparse initialization . These results highlight the need to choose different sparse training methods for different architectures . 2 RELATED WORK . Dense-to-Sparse . There are a large amount of works operating on a dense network to yield a sparse network . We divide them into three categories based on the training cost in terms of memory and computation . ( 1 ) Iterative Pruning and Retraining . To the best of our knowledge , pruning was first proposed by Janowsky ( 1989 ) and Mozer & Smolensky ( 1989 ) to yield a sparse network from a pre-trained network . Recently , Han et al . ( 2015 ) brought it back to people ’ s attention based on the idea of iterative pruning and retraining with modern architectures . Some recent works were proposed to further reduce the number of iterative retraining e.g. , Narang et al . ( 2017 ) ; Zhu & Gupta ( 2017 ) . Frankle & Carbin ( 2019 ) proposed the Lottery Ticket Hypothesis showing that the sub-networks ( “ winning tickets ” ) obtained via iterative pruning combined with their “ lucky ” initialization can outperform the dense networks . Zhou et al . ( 2019 ) discovered that the sign of their initialization is the crucial factor that makes the “ winning tickets ” work . Our work shows that there exists a much more efficient and robust way to find those “ winning ticketts ” without any special initialization . The aforementioned methods require at least the same training cost as training a dense model , sometimes even more , as a pre-trained dense model is involved . We compare our method with state-of-the-art pruning method proposed by Zhu & Gupta ( 2017 ) in Appendix I . With fewer training costs , our method is able to discover sparse networks that can achieve lower test perplexity than iterative pruning . ( 2 ) Learning Sparsity During Training . There are also some works attempting to learn the sparse networks during training . Louizos et al . ( 2017b ) and Wen et al . ( 2018 ) are examples that gradually enforce the network weights to zero via L0 and L1 regularization , respectively . Dai et al . ( 2018 ) proposed a singular value decomposition ( SVD ) based method to accelerate the training process for LSTMs . Liu et al . ( 2020a ) proposed Dynamic Sparse Training to discover sparse structure by learning binary masks associated with network weights . However , these methods start with a fully dense network , and hence are not memory efficient . ( 3 ) One-Shot Pruning . Some works aim to find sparse neural networks by pruning once prior to the main training phase based on some salience criteria , such as connection sensitivity ( Lee et al. , 2019 ) , signal propagation , ( Lee et al. , 2020 ) , and gradient signal preservation ( Wang et al. , 2020 ) . These techniques can find sparse networks before the standard training , but at least one iteration of dense model needs to be trained to identify the sparse sub-networks , and therefore the pruning process is not applicable to memory-limited scenarios . Additionally , one-shot pruning generally can not match the performance of dynamic sparse training , especially at extreme sparsity levels ( Wang et al. , 2020 ) . Sparse-to-Sparse . Recently , many works have emerged to train intrinsically sparse neural networks from scratch to obtain efficiency both for training and inference . ( 1 ) Static Sparse Training . Mocanu et al . ( 2016 ) introduced intrinsically sparse networks by exploring the scale-free and small-world topological properties in Restricted Boltzmann Machines . Later , some works expand static sparse training into CNNs based on expander graphs and show comparable performance ( Prabhu et al. , 2018 ; Kepner & Robinett , 2019 ) . ( 2 ) Dynamic Sparse Training . Mocanu et al . ( 2018 ) introduced Sparse Evolutionary Training ( SET ) which initializes a sparse network and dynamically changes the sparse connectivity by a simple remove-and-regrow strategy . At the same time , DeepR ( Bellec et al. , 2018 ) trained very sparse networks by sampling the sparse connectivity based on a Bayesian posterior . The iterative configuration updates have been proved to converge to a stationary distribution . Mostafa & Wang ( 2019 ) introduced Dynamic Sparse Reparameterization ( DSR ) to train sparse neural networks while dynamically adjusting the sparsity levels of different layers . Sparse Networks from Scratch ( SNFS ) ( Dettmers & Zettlemoyer , 2019 ) improved the sparse training performance by growing free weights according to their momentum . It requires extra computation and memory to update the dense momentum tensor for each iteration . Further , Evci et al . ( 2020 ) introduced RigL which activates weights with the highest magnitude gradients . This approach grows weights expected to receive gradients with high magnitudes , while amortizing a large number of memory requirements and computational cost caused by momentum . Due to the inherent limitations of deep learning software and hardware libraries , all of the above works simulate sparsity using a binary mask over weights . More recently , Liu et al . ( 2020b ) proved the potentials of DST by developing for the first time an independent software framework to train very large truly sparse MLPs trained with SET . However , all these works mainly focus on CNNs and MLPs , and they are not designed to match state-of-the-art performance for RNNs . We summarize the properties of all approaches compared in this paper in Table 1 . Same with SET , our method can guarantee Backward Sparse , which does not require any extra information from the removed weights . Additionally , we discuss the differences among SET , pruning techniques , and our method in Appendix H . 3 SPARSE RNN TRAINING . Our sparse RNN training method is illustrated in Figure 1 with LSTM as a specific case of RNNs . Note that our method can be easily applied to any other RNN variants . The only difference is the number of cell weights . Before training , we randomly initialize each layer at the same sparsity ( the fraction of zero-valued weights ) , so that the training costs are proportional to the dense model at the beginning . To explore more sparse structures , while to maintain a fixed sparsity level , we need to optimize the sparse connectivity together with the corresponding weights ( a combinatorial optimization problem ) . We apply dynamic sparse connectivity and SNT-ASGD to handle this combinatorial optimization problem . The pseudocode of the full training procedure of our algorithm is shown in Algorithm 1 .

Summary: In this paper, the authors studied the possibility of sparsity exploration in Recurrent Neural Networks (RNNs) training. The main contributions include two parts: (1) Selfish-RNN training algorithm in Section 3.1 (2) SNT-ASGD optimizer in Section 3.2. The key idea of the Selfish-RNN training algorithm is a non-uniform redistribution across cell weights for better regularization. The authors mentioned previous sparse training techniques mainly focus on Multilayer Perceptron Networks (MLPs) and Convolutional Neural Networks (CNNs) rather than RNNs. This claim seems to be doubtful because one-time SVD + fine-tuning usually works very well for most RNN training applications in the industry.