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: The paper studies the feasibility to use the dynamic sparsity (DynSparse) method of continuously pruning and re-allocating network weights during training for large unsupervised (self-supervised) language models, namely BERT. Towards this end, the authors present several parameter studies (w.r.t. hyperparameters) of the original BERT training procedure (Mask Language Modeling and Next Sentence Prediction) to assess the Pareto-curve of training BERT with sparse vs dense.