GroupBERT: Enhanced Transformer Architecture with Efficient Grouped Structures

1 INTRODUCTION . Deep neural networks have emerged as the leading solution to enabling end-to-end language processing ( Hochreiter & Schmidhuber , 1997 ; Sutskever et al. , 2014 ; Chung et al. , 2014 ) . Recently , the Transformer model based on the self-attention mechanism ( Vaswani et al. , 2017 ) has become the most promising architecture for language applications , especially when used as a pre-trained foundation model ( Devlin et al. , 2018 ; Radford et al. , 2019 ; Brown et al. , 2020 ; Bommasani et al. , 2021 ) . Attention based models are also increasingly showing promising results for established applications in domains different from natural language processing ( Dosovitskiy et al. , 2020 ) . Complementary to the Transformer ’ s improved ability to model long-range dependencies in sequences is its superior potential to scale to larger sizes ( Kaplan et al. , 2020 ) and its suitability for execution on existing accelerators . This makes these models favoured over traditional recurrent language models . Given the increased computational demand of these models , there is a growing and pressing interest to develop more efficient architectures ( Strubell et al. , 2019 ) . Some previous proposals were able to reduce the computational burden of the Transformer with improved task performance , but often with a corresponding slower execution , as will be discussed further in Section 4 . We demonstrate a set of modifications to the the structure of the Transformer layer that improve FLOP utilization by the encoder stack . The proposed GroupBERT model relies on grouped matrix multiplications and convolutions , and delivers a more efficient version of BERT , superior in both task performance and computation . GroupBERT utilizes grouped operations in a novel way which makes the model more FLOP efficient . However grouped operations are characterised by a reduced computational load for a given memory access ( Masters et al. , 2021 ) , i.e arithmetic intensity . This property would make them undesirable for traditional accelerators , which rely on large dense computations and reduced memory access . For this investigation we use the IPU ( Jia et al. , 2019 ) hardware , since its architecture uses on-chip SRAM for model execution , making it suitable to explore the potential of these efficient building blocks . We achieve a further task performance boost by increasing the depth of the model as GroupBERT extends each Transformer layer to contain four modules : one multi-head attention ( MHA ) , one grouped convolution module , and two grouped feed-forward modules ( GFFN ) . The MHA and grouped convolution modules process token information along the sequence dimension , and each is followed by the general computation GFFN module . While there are twice as many modules in the proposed GroupBERT layer , the overall increase in computation is modest as we utilize sparse grouped operations , for a total FLOP increase of about 60 % . Not only does GroupBERT deliver better performance per FLOP , but it is also executed faster as measured in total time-to-train . By employing both attention and convolution , the model has components dedicated to both short and long-range interactions , making a more efficient use of the more expensive attention mechanism . We also utilize the parameters of GroupBERT more efficiently during training , by discarding dropout for pre-training on a large corpus of text and by improving stability to use higher learning rates . With all these innovations , GroupBERT Base is only slightly larger than BERT Base , yet it achieves better validation MLM loss than BERT Large using less than half of its FLOPs . 2 ARCHITECTURE . In this work , we propose an efficient modification of the Transformer encoder layer called GroupBERT . The original Transformer layer consists of two modules : multi-head attention ( MHA ) and feedforward network ( FFN ) . Each of these modules also includes dropout , a shortcut connection , and layer normalization ( Srivastava et al. , 2014 ; He et al. , 2016 ; Ba et al. , 2016 ) . GroupBERT includes four modules in every layer , as illustrated in Figure 1 . We add a convolution module in sequence with the MHA to efficiently model local interactions between tokens and to allow the attention mechanism to focus on long-range interactions . We then complement every sequence processing block with a dedicated fully-connected module . For better efficiency , we introduce grouped projections to the FLOPs intensive FFN module , making the layer structure more FLOP efficient . 2.1 GROUPED FEED-FORWARD MODULES . The FFN module plays a crucial part in the unparalleled task performance of Transformers ( Dong et al. , 2021 ; Lee-Thorp et al. , 2021 ) . Although it is an essential complement to sequence processing modules it introduces a significant computational burden , as two thirds of the FLOPs are concentrated in the FFN module . To make this integral part of the transformer more lightweight we utilize structured sparsity in a form of sparsely grouped matrix multiplication . Consider a dense matrix multiplication of matrices H ∈ Ra×b and W ∈ Rb×c : ( HW ) i , j = ∆ b∑ n=1 hi , n · wn , j ( 1 ) A sparsely grouped version of W corresponds to a block diagonal matrix W ( G ) with G groups , a matrix of similar dimension to W and a sparsity ratio of 1/G . An equivalent alternative formulation of a block-diagonal matrix is a grouped convolution for a 1-dimensional 1× 1 convolution ( Iandola et al. , 2020 ) . This reduces the number of stored parameters , and can be implemented efficiently without zero-multiplication as : ( HW ( G ) ) i , j = ∆ b∑ n=1 hi , n · wn , j = b/G∑ n=1 h i , n + b G · bj−1 / c G c · w n + b G · bj−1 / c G c , j ( 2 ) We propose a novel scheme to utilize grouped transformations in a transformer FFN module . Our first finding is that parameters in the expanding FFN matrix contribute more to task performance , and sparsity is particularly damaging for these fan-out matrices . The second matrix is less sensitive to parameter reduction due to the sparse input and the reduction of projection dimension . Therefore , introducing sparsity in the second matrix results in a Pareto efficient balance between compute and task-performance . Our second finding is that the locality constraint of grouped projections on the hidden dimension is detrimental to the model . However we remove this constraint by using an output linear projection . This is similar to the output projection matrix used in the MHA block , where each attention head acts only on a slice of the hidden dimension . We find the optimal value for the number of groups to be G = 4 , bringing the parameter count of GFFN to be 75 % of its dense counterpart , while also delivering a more Pareto efficient task performance ( Fig 6 ) . Alternative schemes of applying grouped transformations ( Iandola et al. , 2020 ; Mehta et al. , 2021 ; Jiang et al. , 2020 ) ( Figure 2 ) in a transformer fail to outperform the baseline in a setting of BERT pre-training ( see Section 3.3 ) . 2.2 CONVOLUTION BLOCK . Sequential locality plays an important role for contextualizing tokens in language models . At the same time , long-range interactions have proven to be vital for state-of-the-art performance . Transformers inherently support long-range content-based interactions via self-attention and usually incorporate a form of positional encoding , allowing attention to also capture position-based interactions ( Dai et al. , 2019 ) . Although this gives self-attention strong representational power , a convolution is a more efficient implementation of strictly local , position-based fusion . For this reason we adopt a dedicated convolutional module to improve overall efficiency . The design of our convolution module is similar to Gulati et al . ( 2020 ) , in which convolutions were introduced into a speech recognition Transformer . We apply a gate consisting of a pointwise convolution followed by a Gated Linear Unit ( GLU ) that has been beneficial in language applications ( Dauphin et al. , 2017 ; Wu et al. , 2019a ; 2020 ) . Unlike Gulati et al . ( 2020 ) , we use grouped convolutions in place of depthwise convolutions to add representational capacity . We find that the best trade-off between task performance and computational cost is achieved by using a grouped convolution with group size 16 and kernel size 7 , computed over the sequence dimension . The module also includes an additional layer normalization and a Swish activation ( Ramachandran et al. , 2017 ) . With this module included , fewer attention heads show a strong locality preference since such interactions are readily captured by convolutions . This effect is visible in the attention maps of Figure 3 , showing weaker locality in the model that includes convolutions . To measure this effect quantitatively , we calculate the entropy across target positions for each head and source position . We then average , and normalize by the maximum possible value ( see Appendix C ) . For this measure , zero means that every head attends to a single position exclusively , while one means that every head is position agnostic , although there could still be a joint position and content term . BERT Base has an average entropy ratio of 0.75 and BERT Base + Conv has 0.92 , indicating a shift of positional fusion work from attention to convolution . 2.3 EFFICIENT PARAMETER UTILIZATION . In line with earlier research on the Transformer architecture ( Wang et al. , 2019 ; Liu et al. , 2020 ; Xiong et al. , 2020 ) , we move layer normalization ( Ba et al. , 2016 ) from its position after the module ’ s residual ( `` postnorm '' ) to the first position within each residual block ( `` prenorm '' ) . While this modification does not directly improve task performance , it stabilizes training and allows the use of a larger learning rate that would otherwise trigger the model with postnorm to diverge . We increase the learning rate by a factor of 4× compared to the postnorm baseline . Similarly to Lan et al . ( 2020 ) , we find the use of dropout to be detrimental to the pre-training stage . Due to the substantial size of the dataset , this kind of regularization is not required . While removing dropout yields improvements to the pre-training loss , this does not apply to downstream tasks that rely on smaller datasets . Consequently , we include dropout only when fine-tuning on supervised tasks , that have smaller datasets than the pre-training corpus . 3 RESULTS . To evaluate the architecture modifications , we chose BERT ( Devlin et al. , 2018 ) pre-training and fine-tuning . The large dataset and challenging training objective mean that task performance improves consistently with model size ( Lan et al. , 2020 ) and the risk of over-fitting is reduced . This makes it possible to clearly distinguish architecture modifications that benefit efficiency . Our evaluation of GroupBERT for language representation learning shows that the architecture is : 1 . Training FLOP-efficient across a range of model sizes ( Sections 3.3 , 3.4 ) . 2 . Training time-efficient across a range of compute budgets ( Sections 3.3 , 3.4 ) . 3 . Improved by each constituent part ( Section 3.5 ) . 3.1 EXPERIMENTS . Each experiment consists of two pre-training phases and a fine-tuning phase consisting of multiple training runs , started from the pre-trained model . All phases use the AdamW optimiser ( Loshchilov & Hutter , 2019 ) , with β1 = 0.9 , β2 = 0.999 , = 10−6 . The learning rate follows a linear warm-up decay schedule , whereby the warmup phase lasts for min ( 104 , 0.1· total steps ) steps , and the peak learning rate depends on the training phase and model size . The model is defined over a vocabulary of 30522 WordPiece tokens ( Wu et al. , 2016 ) . Weights are initialized using a truncated normal distribution of standard deviation 0.02 . For all experiments we use 2 Graphcore M2000 IPU systems . Pre-training phase one optimises the Masked Language Model ( MLM ) and Next-Sentence Prediction ( NSP ) loss for corrupted sentence pairs . Masked and padded sequences of length 128 are grouped into batches of approximately 512 sequences , with slight variations depending on the model size ( see Appendix A ) . The model is trained for 10 epochs of Wikipedia + BookCorpus ( Zhu et al. , 2015 ) , corresponding to approximately 8 ·105 optimisation steps . For all experiments with GroupBERT , baseline BERT and all other models tested , the learning rate is set to the value that produces the best validation loss . Pre-training phase two uses sequence length 384 , 5 epochs , and approximately 2·105 optimisation steps . SQuAD 1.1 fine-tuning ( Rajpurkar et al. , 2016 ) adds a token span prediction layer and the whole model is fine-tuned to perform extractive question answering . Training uses target batch size 32 and we train for 2-3 epochs with various learning rates ( Appendix B ) and report results for the best hyperparameters setting . We report F1 and Exact match scores , which show higher variance than MLM loss values . On the grounds of larger variance , we fine-tune each pre-training checkpoint five times using different seeds for every hyperparameter setting . Fine-tuning has been shown to be quite a brittle process in recent studies ( Dodge et al. , 2020 ; Zhang et al. , 2021 ; Mosbach et al. , 2021 ) . In particular , many instabilities are caused by fine-tuning without using bias correction , an implementation that was adopted following the original experimental setup of BERT . This omission in the optimizer was observed to cause a collapse of the training process . For this reason , we included a bias-correction term to the AdamW implementation for fine-tuning .

Summary: Pre-trained language models have become a critical component in NLP applications. However, due to its large-scale pre-training computation, pre-trained models are usually too expensive to train. Therefore, in this paper, authors design GroupBERT, which introduces a set of modifications to the structure of the Transformer network. More specifically, they first apply grouped transformation to reduce the computations cost of dense feed-forward layers, and then use grouped convolution to enhance the learning of local and global interactions. Experimental results demonstrate the effectiveness of the proposed method.