Towards Structured Dynamic Sparse Pre-Training of BERT

1 INTRODUCTION . The increasing task performance gains of large , pre-trained language models have fueled interest in computationally efficient unsupervised training ( Kaplan et al. , 2020 ) . In recent years , sparsity has regained popularity as a technique for improving the computational efficiency of deep learning models ( Hoefler et al. , 2021 ) . Current sparsity methods can be distinguished into approaches that impose sparsity on the weights of neural networks via weight sparsity ( Frankle & Carbin , 2019 ; Gale et al. , 2019 ; Bellec et al. , 2017 ; Mostafa & Wang , 2019 ; Evci et al. , 2019 ; Dettmers & Zettlemoyer , 2019 ; Mocanu et al. , 2018 ; Jayakumar et al. , 2020 ) , or techniques that dynamically route activations to only interact with a subset of the network weights via conditional sparsity ( Shazeer et al. , 2017 ; Lepikhin et al. , 2020 ; Fedus et al. , 2021 ; Lewis et al. , 2021 ) . In weight sparse training ( Frankle & Carbin , 2019 ; Gale et al. , 2019 ) , the number of network parameters is reduced by imposing sparsity patterns on the network weights . As a result , weight sparse training can lead to significant savings in FLOPs , making it promising for scaling to larger network architectures for a given compute budget . One of the most promising candidates for weight sparse training is dynamic sparsity ( DynSparse ) , which reduces FLOPs while only requiring training of sparse subsets of the over-parameterized network ( Bellec et al. , 2017 ; Mostafa & Wang , 2019 ; Evci et al. , 2019 ; Dettmers & Zettlemoyer , 2019 ; Mocanu et al. , 2018 ; Jayakumar et al. , 2020 ; Liu et al. , 2021a ) . In DynSparse approaches , the sparsity pattern imposed on the weights is continuously modified during training using pruning and re-allocation strategies . This evolution leads to a joint exploration of both network topology and parameters , which has been shown to outperform static sparsity baselines ( Bellec et al. , 2017 ; Mostafa & Wang , 2019 ; Evci et al. , 2019 ; Dettmers & Zettlemoyer , 2019 ) . However , so far , the limited performance on language modeling task ( Evci et al. , 2019 ) has resulted in DynSparse training not seeing wide adoption for large-scale language modeling tasks despite recent advances ( Jayakumar et al. , 2020 ) . Given the high cost and energy consumption of unsupervised training of large-scale language models ( Strubell et al. , 2019 ; Patterson et al. , 2021 ) , dynamic sparsity bears the potential to make pre-training more efficient and affordable . To this end , we adopt and investigate DynSparse training techniques ( Dettmers & Zettlemoyer , 2019 ; Evci et al. , 2019 ) for pre-training of BERT bidirectional language encoder ( Devlin et al. , 2018 ) based on the highly scalable Transformer architecture ( Vaswani et al. , 2017 ) . Our work achieves Pareto improvements versus the dense baseline using both structured and unstructured DynSparse training of BERT . 1.1 CONTRIBUTIONS . Investigating dynamic always-sparse training for BERT pre-training . We adapt the DynSparse training algorithm to BERT pre-training ( Section 2.1 ) . In particular , we find that gradient-based re-allocation ( Evci et al. , 2019 ) results in a collapse of the explored network parameters ( Figure 11 ) , which we mitigate through the use of random parameter re-allocation . Achieving scalable FLOPs efficient dynamic sparse training . We compare dense and sparse methods for a given FLOPs budget and demonstrate both algorithmic scalability and Pareto improvement on the FLOPs scale , as shown in Figure 3 . Adapting dynamic always-sparse training to block structures . We extend the unstructured DynSparse training towards block-sparse structure ( Section 3.2 ) . In particular , we find that the choice of metric during block pruning has a strong influence on the task performance , as shown in Figure 7 . Pareto improvements for structured DynSparse training We show that the resulting structured DynSparse training of BERT without structured regularization gives Pareto improvements compared to the dense BERT baseline , as shown in Figure 1 . In the following section , we report the results of explorative experiments conducted to motivate our study of DynSparse training of BERT . The rest of the paper then concentrates on DynSparse training , with methodology discussed in Section 2 and results presented in Section 3 . 1.2 IDENTIFYING SUITABLE ANGLE OF ATTACKS FOR SPARSE PRE-TRAINING . Sparse training of unsupervised language models is relatively under-explored , compared to sparse training of the supervised fine-tuning objective ( Radiya-Dixit & Wang , 2020 ; Chen et al. , 2020 ; Sanh et al. , 2020 ) . Consequently , we design two explorative experiments to assess whether DynSparse training is a suitable sparse training algorithm for pre-training of BERT . Firstly , we analyze the importance of trainable parameters by keeping a random pattern of weights non-trainable ( constant non-zero ) or zero-valued throughout training . This experiment allows us to disentangle the role of ’ zero ’ vs. ’ untrainable ’ weights in the connectivity patterns , to shed light on the parameter dependence of BERT pre-training . Like zeroed weights , the constant weights are unable to encode new information about the task . Still , they might promote the propagation of signals or gradient-flow through the network , which has been considered a core aspect of some sparse training algorithms in vision ( Evci et al. , 2020 ; Lubana & Dick , 2021 ; Tessera et al. , 2021 ) . Non-zero parameters might also lead to better utilization of the remaining trainable parameters . However , as shown in Figure 1 ( a ) , we find that none of these effects plays a relevant role in the training dynamics , since the task performance of the network with sparsified weights ( dashed orange line ) matches the one with the same fraction of untrained weights ( solid blue line ) . Different from vision models that are often based on convolutions , the transformer architecture contains large dense matrices and multiplicative interaction ( Jayakumar et al. , 2019 ) . While zeroing parameters is not expected to affect the training dynamics , the performance remains bounded by the number of trainable parameters . Secondly , we would like to analyze the effect of sparsification at different stages of the training process . For this purpose , we keep a random subset of the network parameters untrainable in the first half of the pre-training , before making them trainable in the second half ( and vice versa ) . Unlike magnitude pruning , which eliminates learned information , freezing and unfreezing parameters ensures symmetry between the different phases ( ignoring the linearly decaying learning rate schedule ) . The agreement in the task performance towards the end of training in Figure 1 ( b ) indicates that representation is continuously built up during training , with no particular effect of when the sparsification is applied . This lack of preference is interesting , given that pre-training has been found to lead to a reduction of the intrinsic dimension ( Li et al. , 2018 ) with respect to downstream tasks ( Aghajanyan et al. , 2020 ) . Our results suggest that sparse pre-training does not necessarily profit from an initial dense training phase . Therefore , we can distribute computation in a way that is both algorithmic and computationally beneficial . In DynSparse training , the network representation is `` always-sparse '' , i.e . it does not rely on the representation of the underlying dense network , making the approach suited for sparse training of large language modeling architectures . Consequently , we believe that BERT pre-training is well suited for DynSparse training . 2 METHODOLOGY . Throughout this work , we study the self-supervised pre-training objective from the original BERT model ( Devlin et al. , 2018 ) , which consists of the Masked Language Model ( MLM ) loss , corresponding to the task performance in predicting a random subset of masked tokens , and the noisier Next Sentence Prediction ( NSP ) loss for binarized next sentence prediction . We focus on a single phase of pre-training with a sequence length of 128 , using the Adam optimizer . All hyperparameters are given in Appendix A for a training length of 10 epochs . 2.1 ADAPTING UNSTRUCTURED DYNSPARSE ALGORITHM TO BERT . In the present work , we first study and adapt the unstructured DynSparse training schematically shown in Figure 2 to pre-training of the BERT language models . Specifically , we initialize the sparsity pattern randomly with the same fixed sparsity ratio on all fully connected encoder weights ( non-embedding weights ) . The weights are initialized using a truncated normal distribution ( see also Figure 9 ) . During an update step of DynSparse training ( see Algorithm 1 ) we use magnitude pruning to remove a time t dependent fraction pr ( t ) of the network parameters . The same fraction of parameters is re-allocated elsewhere in the weight tensor . To complete the sparsity update step , all newly allocated parameters and their corresponding first and second-order moments of the Adam optimizer are initialized to zero . Given that DynSparse training has been primarily developed for vision architectures ( Dettmers & Zettlemoyer , 2019 ; Evci et al. , 2019 ) and did not show competitive performance on the language tasks , we find it necessary to reassess some of the algorithm choices for BERT . In particular , during the re-allocation step of DynSparse training , we use random re-allocation of pruned weights instead of gradient-based techniques as in RigL ( Evci et al. , 2019 ) . For one , this avoids potential issues from a collapse of the explored parameter space ( compare Figure 11 ) . More importantly , the absence of dense gradient computation makes our approach always-sparse , such that the full dense model is never actually instantiated . We found that the cosine decay of the pruning ratio introduced in Evci et al . ( 2019 ) outperforms constant pruning schedules and leads to a reduction of the changes in network topology during training . We refer to the maximum pruning ratio pr simply as `` pruning ratio '' throughout the paper . All DynSparse hyperparameters are optimized for a sparsity ratio of 0.9 ( for more details , refer to Appendix A.1 ) . 2.2 BLOCK-SPARSE DYNSPARSE ALGORITHM . Extending unstructured DynSparse training towards using structured sparse computation requires modification to both the prune and update steps in Figure 2 . Magnitude pruning can be justified as a simple compression algorithm resulting in unstructured sparsity . However , there is no unique way to extend the magnitude pruning metric to blocks of parameters . Choosing a good metric for block pruning is essential , as magnitude pruning has been surprisingly successful in preserving the task performance of sparse networks ( Gale et al. , 2019 ) . In the following , we evaluate the selection of norms consisting of L∞-norm , L2-norm and L1-norm as different criteria for estimating the importance of blocks of parameters . For a block of weights B = W { r , s| ( r , s ) ∈B } taken from a weight tensor W indexed by B , the Lp-norm is given by Lp ( B ) = ∑ i , j |Bi , j |p 1/p , ( 1 ) where the exponent p , e.g . p =∞ , 2 , 1 , controls the relative importance of individual parameters of a block according to their magnitude . In the limit of block size 1×1 , the block pruning according to Eq . ( 1 ) reduces to magnitude pruning , allowing us to investigate the task performance with increasing block sizes . For small values of p→ 0 , each parameter in the block contributes equally towards the importance of the block element as |Bi , j |p → 1 , while for large values of p→∞ the importance of the block collapses towards the parameter with the largest magnitude , with Lp→∞ ( B ) → max ( |B| ) . Therefore , the pruning metric for blocks controls the extent to which the magnitude of each of the parameters in a block contributes to the importance of the block itself . Algorithm 1 : Structured DynSparse training Input : total number of training step T , total number of sparsity updates n− 1 , pruning ratio pr ( t ) at time t , sparsity ratio s , block size B ; Initialize : Impose random block-sparse pattern on non-embedding weights with uniform constant sparsity ratio across all layers . Initialize weights sampled from a truncated normal distribution ; for k = 1 to n do Train network with static sparsity pattern for T/n steps ; For each weight tensor : prune fraction pr ( t ) of blocks with smallest Lp-norm defined in Eq . ( 1 ) ; re-allocate the same fraction pr ( t ) of blocks ; re-allocated parameters , first and second-order moments of Adam optimizer are all initialized to zero end Structured regularization Group Lasso regularization is commonly used as a structured sparsityinducing regularizer ( Wen et al. , 2017 ; Narang et al. , 2017 ) . We introduce the Group Lasso regularization in the update ∆W of weight tensor W following the decoupling of weight decay and Adam optimizer from Loshchilov & Hutter ( 2017 ) . More specifically , the entry ( i , j ) of the parameter update ∆Wij is adjusted to ∆W reg ij as ∆W regij = ∆Wij − lr ( t ) · λgroup · wstd · √ B Wij√ ∑ ( r , s ) ∈B ( i , j ) W 2r , s + , ( 2 ) where B ( i , j ) denotes the set of weight indices that belong to the same block as the ( i , j ) th weight element and η ( t ) is the linearly decaying learning rate ( see Appendix A ) . The remaining coefficients are the Group Lasso coefficient λgroup , and the small constant = 10−6 for numerical stability . The extra pre-factors wstd = 0.02 and √ B corresponding to the standard deviation of the weights at initialization ( see Appendix A ) and the square-root of the block size respectively are chosen such as to ensure that the regularization coefficients of weight decay and group lasso are comparable in magnitude .

Summary: This paper did an extensive investigation of applying dynamic sparse training for BERT.  The main contributions of this work are several empirical observations: 1. The authors empirically show that random re-allocation of weights is better than gradient-based approach (RigL) for unstructured DynSparse training of Bert.  2. The authors investigated structured DynSparse training and found L1 norm is the best metric to prune the weights. 3. For a specific range of size of BERT, the DynSparse training shows Pareto improvement (MLM validation loss vs. Non-embedding FLOPS) over the dense BERT.