Successfully Applying the Stabilized Lottery Ticket Hypothesis to the Transformer Architecture (2020)

Abstract: Sparse models require less memory for storage and enable a faster inference by reducing the necessary number of FLOPs. This is relevant both for time-critical and on-device computations using neural networks. The stabilized lottery ticket hypothesis states that networks can be pruned after none or few training iterations, using a mask computed based on the unpruned converged model. On the transformer architecture and the WMT 2014 English-to-German and English-to-French tasks, we show that stabilized lottery ticket pruning performs similar to magnitude pruning for sparsity levels of up to 85%, and propose a new combination of pruning techniques that outperforms all other techniques for even higher levels of sparsity. Furthermore, we confirm that the parameter’s initial sign and not its specific value is the primary factor for successful training, and show that magnitude pruning cannot be used to find winning lottery tickets.

## introduction current neural networks are heavily growing in depth, with many fully connected layers. as every fully connected layer includes large matrices, models often contain millions of parameters. this is commonly seen as an over-parameterization @xcite . different techniques have been proposed to decide which weights can be pruned. in structured pruning techniques @xcite , whole neurons or even complete layers are removed from the network. unstructured pruning only removes individual connections between neurons of succeeding layers, keeping the global network architecture intact. the first technique directly results in smaller model sizes and faster inference, while the second offers more flexibility in the selection of which parameters to prune. although the reduction in necessary storage space can be realized using sparse matrix representations @xcite , most popular frameworks currently do not have sufficient support for sparse operations. however, there is active development for possible solutions @xcite @xcite . this paper compares and improves several unstructured pruning techniques. the main contributions of this paper are to: in concurrent work, @xcite test the stabilized lottery ticket on the transformer architecture and the wmt 2014 @xmath0, as well as other architectures and fields. this paper extends the related works by demonstrating and comparing the applicability of different pruning techniques on a deep architecture for two translation tasks, as well as proposing a new combination of pruning techniques for improved performance. ## pruning techniques in this section, we give a brief formal definition of each pruning technique. for a more detailed description, refer to the respective original papers. in the given formulas, a network is assumed to be specified by its parameters θ. when training the network for t iterations, θ t for @xmath1, t ] represents the parameters at timestep t. magnitude pruning (mp) relies on the magnitude of parameters to decide which weights can be pruned from the network. different techniques to select which parameters are selected for pruning have been proposed @xcite @xcite . in this work, we rely on the implementation from @xcite where the parameters of each layer are sorted by magnitude, and during training, an increasing percentage of the weights are pruned. it is important to highlight that mp is the only pruning technique not requiring multiple training runs. lottery ticket (lt) pruning assumes that for a given mask m, the initial network θ 0 already contains a sparse sub-network θ 0 m that can be trained to the same accuracy as θ 0 . to determine m, the parameters of each layer in the converged model θ t are sorted by magnitude, and m is chosen to mask the smallest ones such that the target sparsity s t is reached. we highlight that even though m is determined using θ t , it is then applied to θ 0 before the sparse network is trained. to reach high sparsity without a big loss on accuracy, @xcite recommend to prune iteratively, by training and resetting multiple times. stabilized lottery ticket (slt) pruning is an adaptation of lt pruning for larger models. @xcite propose to apply the computed mask m not to the initial model θ 0 , but to an intermediate checkpoint θ t where 0 < t t is chosen to be early during the training. they recommend to use 0.001t ≤ @xmath2.07t and refer to it as iterative magnitude pruning with rewinding. we highlight that @xcite always choose θ t from the first, dense model, while this work choses θ t from the last pruning iteration. constant lottery ticket (clt) pruning assumes that the specific random initialization is not important. instead, only the corresponding choice of signs affects successful training. to show this, @xcite propose to compute θ t m as in slt pruning, but then to train f (θ t m) as the sparse model. here, f sets all remaining parameters p in each layer l to sign(p) • α l , i.e., all param-eters in each layer have the same absolute value, but their original sign. in all of our experiments, α l is chosen to be α @xmath3. slt-mp is a new pruning technique, proposed in this work. it combines both slt pruning and mp in the following way: first, slt pruning is used to find a mask m with intermediate sparsity s i . this might be done iteratively. θ t m with sparsity s i is then used as the initial model for mp (i.e., θ 0 = θ t m). here, in the formula for mp, s 0 = s i . we argue that this combination is beneficial, because in the first phase, slt pruning removes the most unneeded parameters, and in the second phase, mp can then slowly adapt the model to a higher sparsity. mp-slt is analogue to slt-mp: first, mp is applied to compute a trained sparse network θ t with sparsity s i . this trained network directly provides the corresponding mask m. θ t m is then used for slt pruning until the target sparsity is reached. this pruning technique tests whether mp can be used to find winning lottery tickets. ## experiments we train the models on the wmt 2014 @xmath4, consisting of about 4.5m and 36m sentence pairs, respectively. newstest2013 and 2014 are chosen to be the development and test sets. all experiments have been performed using the base transformer architecture as described in @xcite . 1 the models are trained for 500k iterations on a single v3-8 tpu, saving checkpoints every 25k iterations. for all experiments, we select the best model based on the bleu score on the development set. for mp, we only evaluate the last 4 checkpoints, as earlier checkpoints do not have the targeted sparsity. intermediate mp sparsity levels s t are computed as s @xmath5, (s 0s t )(1 -t 400000 ) 3 } @xcite . for efficiency reasons, weights are only pruned every 10k iterations. unless stated otherwise, we start with initial sparsity s 0 = 0. the final sparsity s t is individually given for each experiment. we prune only the matrices, not biases. we report the approximate memory consumption of all trained models using the compressed sparse column (csc) format @xcite , which is the default for sparse data storage in the scipy toolkit @xcite . our initial experiments have shown that adafactor leads to an improvement of 0.5 bleu compared to adam. hence, we select it as our optimizer with a learning rate of lr(t) = 1 max(t,w) for @xmath6. we note that this differs from the implementation by @xcite , in which adam has been used. we highlight that for all experiments that require a reset of parameter values (i.e., lt, slt, clt, slt-mp, and mp-slt), we reset t to 0, to include the warmup phase in every training run. a shared vocabulary of 33k tokens based on word-pieces @xcite is used. the reported case-sensitive, tokenized bleu scores are computed using sacrebleu @xcite , ter scores are computed using multeval @xcite . all results are averaged over two separate training runs. for all experiments that require models to be reset to an early point during training, we select a checkpoint after 25k iterations. all iterative pruning techniques except slt-mp are pruned in increments of 10 percentage points up to 80%, then switching to 5 points increments, and finally pruning to 98% sparsity. slt-mp is directly trained using slt pruning to 50% and further reduced by slt to 60%, before switching to mp. ## experimental results in this section, we evaluate the experimental results for @xmath7. importantly, we note that the magnitude of the parameters in both the initial and the final models increases with every pruning step. this causes the model with 98% sparsity to have weights greater than 100, making it unsuitable for checkpoint averaging, as the weights become too sensitive to minor changes. @xcite report that they do successfully apply checkpoint averaging. this might be because they choose θ t from the dense training run for resetting, while we choose θ t from the most recent sparse training. clt the underlying idea of the lt hypothesis is, that the untrained network already contains a sparse sub-network which can be trained individually. @xcite show that only the signs of the remaining parameters are important, not their specific random value. while @xcite perform their experiments on mnist and cifar-10, we test this hypothesis on the wmt 2014 @xmath8. surprisingly, clt pruning outperforms slt pruning on most sparsity levels (see slt-mp across all sparsity levels, the combination of slt pruning and mp outperforms all other pruning techniques. for high sparsity values, slt-mp models are also superior to the slt models by @xcite , even though they start of from a better performing baseline. we hypothesize that by first discarding 60% of all parameters using slt pruning, mp is able to fine-tune the model more easily, because the least useful parameters are already removed. we note that the high weight magnitude for sparse slt models prevents successful mp training. therefore, we have to reduce the number of slt pruning steps by directly pruning to 50% in the first pruning iteration. however, as seen by comparing the scores for 50% and 60% sparsity on slt and slt-mp, this does not hurt the slt performance. for future work, we suggest trying different sparsity values s i for the switch between slt and mp. mp-slt switching from mp to slt pruning causes the models to perform worse than for pure mp or slt pruning. this indicates that mp cannot be used to find winning lottery tickets. ## conclusion in conclusion, we have shown that the stabilized lottery ticket (slt) hypothesis performs similar to magnitude pruning (mp) on the complex transformer architecture up to a sparsity of about 85%. especially for very high sparsities of 90% or more, mp has proven to perform reasonably well while being easy to implement and having no additional training overhead. we also have successfully verified that even for the transformer architecture, only the signs of the parameters are important when applying the slt pruning technique. the specific initial parameter values do not significantly influence the training. by combining both slt pruning and mp, we can improve the sparsity-accuracy tradeoff. in slt-mp, slt pruning first discards 60% of all parameters, so mp can focus on fine-tuning the model for maximum accuracy. finally, we show that mp cannot be used to determine winning lottery tickets. in future work, we suggest performing a hyperparameter search over possible values for t in slt pruning (i.e., the number of training steps that are not discarded during model reset), and over s i for the switch from slt to mp in slt-mp. we also recommend looking into why clt pruning works in our setup, while frankle et al. (2020) present opposing results.