Playing the lottery with rewards and multiple languages: lottery tickets in RL and NLP

1 INTRODUCTION . The lottery ticket phenomenon ( Frankle & Carbin , 2019 ; Frankle et al. , 2019 ; Zhou et al. , 2019 ) occurs when small , sparse sub-networks can be found in over-parameterized deep neural networks ( DNNs ) which , when trained in isolation , can achieve similar or even greater performance than the original , highly over-parameterized network . This phenomenon suggests that over-parameterization in DNN training is beneficial primarily due to proper initialization rather than regularization during the training process itself ( Allen-Zhu et al. , 2019 ; 2018 ; Du & Lee , 2018 ; Du et al. , 2019 ; Neyshabur et al. , 2014 ; 2019 ) . However , despite extensive experiments in Frankle & Carbin ( 2019 ) and Frankle et al . ( 2019 ) , it remains unclear whether the lottery ticket phenomenon is an intrinsic feature of DNN behavior or whether it is dependent on other factors such as supervised learning , network architecture , specific tasks ( e.g. , image classification ) , the bias in the dataset , or artifacts from the optimization algorithm itself . As discussed in Frankle & Carbin ( 2019 ) and Liu et al . ( 2019 ) , large learning rates severely damage the lottery ticket effect , and for larger models ( such as VGG and ResNets ) and datasets ( e.g. , ImageNet ) , heuristics like learning rate warmup ( Frankle & Carbin , 2019 ) and late rewinding ( Frankle et al. , 2019 ) are needed to induce high performance and reliable winning tickets . Recent work has also questioned the effectiveness of the lottery ticket hypothesis , raising concerns about the generality of this phenomenon ( Liu et al. , 2019 ; Gale et al. , 2019 ) . In this work , we address the question of whether the lottery ticket phenomenon is merely an artifact of supervised image classification with feed-forward convolutional neural networks , or whether ∗Work done while at Facebook AI Research ; currently at Horizon Robotics †To whom correspondence should be addressed . this phenomenon generalizes to other domains , architectural paradigms , and learning regimes ( e.g. , environments with reward signals ) . Many natural language processing ( NLP ) models feature complex gating mechanics paired with recurrent dynamics , either of which may significantly impact the optimization process and , consequently , the lottery ticket phenomenon . Furthermore , prior work has suggested that this phenomenon is not present in Transformer models ( Gale et al. , 2019 ) , calling the broad applicability of lottery tickets into question . In reinforcement learning ( RL ) , the data distribution shifts as the agent learns from often reward signals , significantly modifying the optimization process and the resultant networks . Pre-trained feature extractors have proven successful in computer vision ( Kornblith et al. , 2019 ; Razavian et al. , 2014 ; Yosinski et al. , 2014 ) , but in RL , agents often fail to generalize even to extremely similar situations ( Raghu et al. , 2018 ; Lanctot et al. , 2017 ; Cobbe et al. , 2018 ; Ruderman et al. , 2019 ) . To answer this question , we evaluate whether the lottery ticket hypothesis holds for NLP and RL , both of which are drastically different from traditional supervised image classification . To demonstrate the lottery ticket phenomenon , we ask whether sparsified subnetworks initialized as winning tickets outperform randomly initialized subnetworks at convergence . We note that , though desirable , we do not require that subnetworks match the performance of the full network , as originally stated in Frankle et al . ( 2019 ) . We exclude this requirement because we are primarily interested in whether appropriate initialization impacts the performance of subnetwork training , consistent with the revised definition of the lottery ticket hypothesis in Frankle & Carbin ( 2019 ) . For NLP , we evaluate language modelling on Wikitext-2 with LSTMs ( Merity et al. , 2017 ) and machine translation on the WMT ’ 14 English-German translation task with Transformers ( Vaswani et al. , 2017 ) . For RL , we evaluate both classic control problems and Atari games ( Bellemare et al. , 2013 ) . Perhaps surprisingly , we found that lottery tickets are present in both NLP and RL tasks . In NLP , winning tickets were present both in recurrent LSTMs trained on language modeling and in Transformers ( Vaswani et al. , 2017 ) trained on a machine translation task while in RL , we observed winning tickets in both classic control problems and Atari games ( though with high variance ) . Notably , we are able to find masks and initializations which enable a Transformer Big model to train from scratch to achieve 99 % the BLEU score of the unpruned model on the Newstest ’ 14 machine translation task while using only a third of the parameters . Together , these results demonstrate that the lottery ticket phenomenon is a general property of deep neural networks , and highlight their potential for practical applications . 2 RELATED WORK . Our work is primarily inspired by the lottery ticket hypothesis , first introduced in Frankle & Carbin ( 2019 ) , which argues that over-parameterized neural networks contain small , sparse sub-networks ( with as few as 0.1 % of the original network ’ s parameters ) which can achieve high performance when trained in isolation . Frankle et al . ( 2019 ) revised the lottery ticket hypothesis to include the notion of late rewinding , which was found to significantly improve performance for large-scale models and large-scale image classification datasets . In addition , the revised lottery ticket hypothesis relaxed the need for subnetworks to match the performance of the full network to simply exceeding the performance a randomly initialized subnetwork . For brevity , we will refer exclusively to the revised definition throughout the paper . However , both of these works solely focused on supervised image classification , leaving it unclear whether the lottery ticket phenomenon is present in other domains and learning paradigms . Recent work ( Liu et al. , 2019 ) challenged the lottery ticket hypothesis , demonstrated that for structured pruning settings , random sub-networks were able to match winning ticket performance . Gale et al . ( 2019 ) also explored the lottery ticket hypothesis in the context of ResNets and Transformers . Notably , they found that random sub-networks could achieve similar performance to that of winning ticket networks for both model classes . However , they did not use iterative pruning or late rewinding , both of which have been found to significantly improve winning ticket performance Frankle et al . ( 2019 ) ; Frankle & Carbin ( 2019 ) . More broadly , pruning methods for deep neural networks have been explored extensively ( Han et al. , 2015 ) . Following Frankle et al . ( 2019 ) and Frankle & Carbin ( 2019 ) , we use magnitude pruning in this work , in which the smallest magnitude weights are pruned first ( Han et al. , 2015 ) . To determine optimal pruning performance , Molchanov et al . ( 2017b ) greedily prune weights to determine an oracle ranking . Also , Ayinde et al . ( 2019 ) and Qin et al . ( 2019 ) have attempted to rank channels by redundancy and preferentially prune redundant filters , and Molchanov et al . ( 2017a ) used variational methods to prune models . However , all of these works only evaluated these approaches for supervised image classification . 3 APPROACH . 3.1 GENERATING LOTTERY TICKETS . Pruning methods In our experiments , we use both one-shot and iterative pruning to find winning ticket initializations . In one-shot pruning , the full network is trained to convergence , and then a given fraction of parameters are pruned , with lower magnitude parameters pruned first . To evaluate winning ticket performance , the remaining weights are reset to their initial values , and the sparsified model is retrained to convergence . However , one-shot pruning is very susceptible to noise in the pruning process , and as a result , it has widely been observed that one-shot pruning under-performs iterative pruning methods ( Frankle & Carbin , 2019 ; Han et al. , 2015 ; Liu et al. , 2019 ) . In iterative pruning ( Frankle & Carbin , 2019 ; Han et al. , 2015 ) , alternating cycles of training models from scratch and pruning are performed . At each pruning iteration , a fixed , small fraction of the remaining weights are pruned , followed by re-initialization to a winning ticket and another cycle of training and pruning . More formally , the pruning at iteration k + 1 is performed on the trained weights of the winning ticket found at iteration k. At iteration k with an iterative pruning rate p , the fraction of weights pruned is : rk = 1− ( 1− p ) k , 1 ≤ k ≤ 20 We therefore perform iterative pruning for all our experiments unless otherwise noted , with an iterative pruning rate p = 0.2 . For our RL experiments , we perform 20 pruning iterations . Pruning was always performed globally , such that all weights ( including biases ) of different layers were pooled and their magnitudes ordered for pruning . As a result , the fraction of parameters pruned across layers may vary given a total pruning fraction . Late rewinding In the original incarnation of the lottery ticket hypothesis ( Frankle & Carbin , 2019 ) , winning tickets were reset to their values at initialization . However , Frankle et al . ( 2019 ) found that resetting winning tickets to their values to an iteration early in training resulted in dramatic improvements in winning ticket performance on large-scale image datasets , even when weights were reset to their values only a few iterations into training . late rewinding can therefore be defined as resetting winning tickets to their weights at iteration j in training , with the original lottery ticket paradigm taking j = 0 . Unless otherwise noted , we use late rewinding throughout , with a late rewinding value of 1 epoch used for all RL experiments . We also compare late rewinding with normal resetting for NLP in section 4.1 and on several representative RL games in section 4.2 . 3.2 NATURAL LANGUAGE PROCESSING . To test the lottery ticket hypothesis in NLP , we use two broad model and task paradigms : twolayer LSTMs for language modeling on Wikitext-2 ( Merity et al. , 2017 ) and Transformer Base and Transformer Big models ( Vaswani et al. , 2017 ) on the WMT ’ 14 En2De News Translation task . Language modeling using LSTMs For the language modeling task , we trained an LSTM model with a hidden state size of 650 . It contained a dropout layer between the two RNN layers with a dropout probability of 0.5 . The LSTM received word embeddings of size 650 . For training , we used truncated Backpropagation Through Time ( truncated BPTT ) with a sequence length of 50 . The training batch size was set to 30 , and models were optimized using Adam with a learning rate of 10−3 and β1 = 0.9 , β2 = 0.999 , = 10−3 . As in the RL experiments , we use global iterative pruning with an iterative pruning rate of 0.2 and 20 total pruning iterations . We also employ late rewinding where the initial weights of a winning ticket were set to the weights after first epoch of training the full network . For ticket evaluation , we trained the model for 10 epochs on the training set , after which we computed the log perplexity on the test set . We also perform two ablation studies without late rewinding and using one-shot pruning , respectively . Machine translation using transformers For the machine translation task , we use the FAIRSEQ framework1 ( Ott et al. , 2019 ) , following the setup described in ( Ott et al. , 2018 ) to train Transformerbase model on the pre-processed dataset from ( Vaswani et al. , 2017 ) . We train models for 50,000 updates and apply checkpoint averaging . We report case-sensitive tokenized BLEU with multi-bleu.pl2 on the Newstest 2014 3 .

Summary: This paper studies the existence of “winning ticket” initialization in natural language processing (NLP) and reinforcement learning (RL). The lottery ticket hypothesis has been found effective in over-parameterized deep neural networks, which provides better sub-network initialization if not outperforming the original full network. Experiments show that winning ticket initializations generally outperform parameter-matched random initializations on recurrent LSTM models, large-scale Transformer, and discrete-action space RL tasks, including both classic control and pixel control.  A substantial number of parameters can be saved.