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: The authors propose a new Transformer architecture. Compared the the traditional Transformer, the proposed architecture has two new features.:1) grouped FFN and 2) a convolution module. The grouped FFN is an efficient replacement of the original FFN. The convolution module is responsible for capturing local dependencies. Experiments demonstrate that the proposed variant achieves up to 2x efficiency gain in terms of both FLOPs and time-to-train.