Lite Transformer with Long-Short Range Attention

1 INTRODUCTION . Transformer ( Vaswani et al. , 2017 ) is widely used in natural language processing due to its high training efficiency and superior capability in capturing long-distance dependencies . Building on top of them , modern state-of-the-art models , such as BERT ( Devlin et al. , 2019 ) , are able to learn powerful language representations from unlabeled text and even surpass the human performance on the challenging question answering task . However , the good performance comes at a high computational cost . For example , a single transformer model requires more than 10G Mult-Adds in order to translate a sentence of only 30 words . Such extremely high computational resources requirement is beyond the capabilities of many edge devices such as smartphones and IoTs . Therefore , it is of great importance to design efficient and fast transformer architecture specialized for real-time NLP applications on the edge . Automatic neural architecture search ( Zoph & Le , 2017 ; So et al. , 2019 ) is a choice for high accuracy model design , but the massive search cost ( GPU hours and CO2 emission ) raises severe environmental concerns ( Strubell et al. , 2019 ) , shown in Figure 1b . In this paper , we focus on the efficient inference for mobile devices , where the total number of MultAdds is constrained below 500M . A straightforward way to reduce the computation of the transformer is to shrink the embedding size directly . Although it can effectively reduce both model size and computation , it also weakens the model capacity capturing the long and short distance relationship at the same time . To this end , we systematically studied the computation breakdown of the transformer ∗ indicates equal contributions . and observed that the computation ( Mult-Adds ) is dominated by the feed-forward network ( FFN ) . We discovered that the prevailing bottleneck-structured transformer block is not efficient . We then present a novel Long-Short Range Attention ( LSRA ) primitive . LSRA trades off the computation in FFN for wider attention layers . It stretches the bottleneck to introduce more dependency capturing capability for the attention layer , and then shrink the embedding size to reduce the total computation amount while maintaining the same performance . Instead of having one module for “ general ” information , LSRA dedicates specialized heads to model long and short distance contexts . Inspired by Wu et al . ( 2019b ) , LSRA introduces convolution in a parallel branch to capture local dependencies so that the attention branch can focus on global context capture . By stacking this primitive , we build Lite Transformer for mobile NLP applications . Extensive experiments demonstrate that our Lite Transformer model offers significant improvements over the transformer on three language tasks : machine translation , abstractive summarization , and language modeling . For machine translation , on IWSLT 2014 German-English , it outperforms the transformer by 3.1 BLEU under 100M Mult-Adds ; on WMT 2014 English-German , it surpasses the transformer by 0.4 BLEU under 500M Mult-Adds and 1.2 BLEU under 100M Mult-Adds ; on WMT 2014 English-French , it also achieves consistent improvements over the transformer : 1.2 BLEU under 500M Mult-Adds and 1.7 BLEU under 100M Mult-Adds . Further , combined with general model compression techniques ( pruning and quantization ) , our Lite Transformer can achieve 18.2× model size compression . For the summarization task , on CNN-DailyMail , it reduces the computation of the transformer base model by 2.4× . For language modeling , it achieves 1.8 lower perplexity than the transformer around 500M Mult-Adds . Guided by our design insights , our manually-designed Lite Transformer achieves 0.5 higher BLEU than the AutoML-based Evolved Transformer ( So et al. , 2019 ) , which requires more than 250 GPU years to search , emitting as much carbon as five cars in their lifetimes ( see Figure 1b ) . It indicates that AutoML is not a panacea : careful analysis and design insights ( i.e. , removing the bottleneck , specialized heads ) can effectively prune the search space and improve the sample efficiency . The contribution of this paper has four aspects : 1 . We systematically analyze the commonly used computation bottleneck structure in modern neural networks and find that the bottleneck design is not optimal for 1-D attention if using FLOPs as figure of merit . 2 . We propose a specialized multi-branch feature extractor , Long-Short Range Attention ( LSRA ) , as the basic building block of our transformer , where convolution helps capture the local context and attention concentrates on global context . 3 . We build Lite Transformer based on our LSRA . Under mobile computation resource constraints ( 500M Mult-Adds ) , our Lite Transformer demonstrates coherent improvement on three widely used machine translation datasets . With extra experiments on other tasks , Lite Transformer is efficient for multiple language applications . 4 . Even compared to AutoML-searched Evolved Transformer , our Lite Transformer offers 0.5 higher BLEU score on WMT En-De dataset under the mobile setting , saving the design cost by 20000× in CO2 emissions . It alerts us to rethink the practicality of AutoML in terms of design cost and “ green AI ” . 2 RELATED WORK . RNNs and CNNs . Recurrent neural networks ( RNNs ) have prevailed various sequence modeling tasks for a long time ( Sutskever et al. , 2014 ; Luong et al. , 2015 ; Bahdanau et al. , 2015 ; Wu et al. , 2016 ) . However , RNNs are not easy to parallelize across the sequence due to its temporal dependency . Recently , some work has demonstrated that RNN is not an essential component to achieve stateof-the-art performance . For instance , researchers have proposed highly-efficient convolution-based models ( Kalchbrenner et al. , 2016 ; Gehring et al. , 2017 ; Kaiser et al. , 2018 ; Wu et al. , 2019b ) . Convolution is an ideal primitive to model the local context information ; however , it lacks the ability to capture the long-distance relationship , which is critical in many sequence modeling tasks . Transformers . As an alternative , attention is able to capture global-context information by pairwise correlation . Transformer ( Vaswani et al. , 2017 ) has demonstrated that it is possible to stack the self-attentions to achieve state-of-the-art performance . Recently , there have been a lot of variants to the transformer ( Ahmed et al. , 2017 ; Ott et al. , 2018 ; Chen et al. , 2018 ; Paulus et al. , 2018 ; Shaw et al. , 2018 ; Sukhbaatar et al. , 2019a ; b ; Child et al. , 2019 ) . Among them , Ott et al . ( 2018 ) proposed to scale up the batch size ; Shaw et al . ( 2018 ) leverages the relative position representations ; Ahmed et al . ( 2017 ) introduces the weighted multi-head attention ; Sukhbaatar et al . ( 2019a ) applies adaptive masks for long-range information on character-level language modeling with very long sequences . All these attempts are orthogonal to our work , as their methods can also be applied in our architecture . Automated Model Design . Due to the vast architecture design space , automating the design with neural architecture search ( NAS ) becomes popular ( Zoph & Le , 2017 ; Zoph et al. , 2018 ; Pham et al. , 2018 ; Cai et al. , 2019a ) . To make the design efficient , integrating the hardware resource constraints into the optimization loop begins to emerge , such as MnasNet ( Tan et al. , 2019 ) , ProxylessNAS ( Cai et al. , 2019b ) and FBNet ( Wu et al. , 2019a ) . In the NLP community , the evolved transformer ( So et al. , 2019 ) adopts the neural architecture search ( Zoph & Le , 2017 ) to design basic blocks and finds a better # parameter-BLEU trade-off for the transformer . However , AutoML-based model designs require significant amount of GPU hours to find the ‘ best ’ model , which is not affordable for most researchers . Model Acceleration . Apart from designing efficient models directly ( Liu et al. , 2019b ; Li et al. , 2020 ) , another approach to achieve efficient inference is to compress and accelerate the existing large models . For instance , some have proposed to prune the separate neurons ( Han et al. , 2015b ; 2016 ) or the entire channels ( He et al. , 2017 ; Liu et al. , 2017 ; He et al. , 2018 ) ; others have proposed to quantize the network ( Courbariaux et al. , 2016 ; Zhu et al. , 2017 ; Krishnamoorthi , 2018 ; Wang et al. , 2019 ) to accelerate the model inference . Recently , AutoML has also been used to automate the model compression and acceleration ( He et al. , 2018 ; Yang et al. , 2018 ; Wang et al. , 2019 ; Liu et al. , 2019a ) . All these techniques are compressing existing models and are therefore orthogonal to our approach . We aim to explore how to make use of the domain knowledge to design an efficient architecture from the beginning , rather than compressing an existing model . 3 IS BOTTLENECK EFFECTIVE FOR 1-D ATTENTION ? . The attention mechanism has been widely used in various applications , including 1-D ( language processing ( Vaswani et al. , 2017 ) ) , 2-D ( image recognition ) , and 3-D ( video recognition ( Wang et al. , 2018 ) ) . It computes pairwise dot-product between all the input elements to model both short-term and long-term relationships . Despite its effectiveness , the operation introduces massive computation . Assume the number of elements ( e.g. , length of tokens in language processing , number of pixels in image , etc . ) fed to attention layer is N , and the dimension of features ( i.e. , channels ) is d , the computation needed for the dot-product is N2d . For images and videos , N is usually very large . For example , the intermediate feature map in a video network ( Wang et al. , 2018 ) has 16 frames , each with 112×112 resolution , leading to N = 2× 105 . The computation of convolution and fully-connected layers grows linearly w.r.t . N , while the computation of attention layers grows quadratically w.r.t . N . The computation of attention module will soon overwhelm with a large N . Lite Transformer with Long-Short Range Attention , ICLR ’ 20 1 Lite Transformer with Long-Short Range Attention To address the dilemma , a common practice is first to reduce the number of channels d using a linear projection layer before applying attention and increase the dimension afterwards ( as shown in Figure 2 ) . In the original design of transformer ( Vaswani et al. , 2017 ) , the channel dimension in the attention module is 4× smaller than that in the FFN layer . Similarly , in the non-local video network ( Wang et al. , 2018 ) , the channel number is first reduced by half before applying the non-local attention module . This practice saves the computation by 16× or 4× . Nevertheless , it also decreases the contexts capture ability of attention layers with a smaller feature dimension . The situation could be even worse for language processing , as attention is the major module for contexts capture ( unlike images and videos where convolutions conduct the major information capture ) . For tasks like translation , the length of the input sequence N tends to be small , which is around 20-30 in common cases . A transformer block consists of an attention ( or two for decoder ) , followed by a feed-forward network ( FFN ) . For the attention layer , the Mult-Adds would be O ( 4Nd2 +N2d ) ; for FFN , the Mult-Adds is O ( 2× 4Nd2 ) . Given a small N , it is doubtful if the bottleneck design is a good trade-off between computation and accuracy on 1D attention . To verify the idea , we first profile the computation breakdown in the transformer in Figure 2 . Surprisingly , for the original transformer ( denoted as ‘ Base ’ in the figure ) , the FFN layer actually consumes much of the computation . This is not desirable since FFN itself can not perform any contexts captures . In conclusion , due to the small N , the bottleneck design can not significantly reduce the computation in 1D attention , while the limited benefit for computation reduction is further compromised by the large FFN layer . It also harms the capacity of attention layer due to the smaller dimension , which is the major contexts capture unit in the transformer . Therefore , we argue that the bottleneck design is not optimal for 1-D attention . We instead design a ‘ flattened ’ version of the transform block that does not reduce and increase the channel dimension . With the new design , the attention part now takes up the major computation in the flattened transformer model in Figure 2 , leaving a larger space for further optimization . We also test the performance change of such modification on WMT ’ 14 En-Fr dataset . We can achieve comparable performance at a slightly larger computation , which can be easily reduced with further optimization that is discussed in the next section .

Summary: This paper presents a new technique (LSRA) improving Transformer for constrained scenarios (e.g., mobile settings). It combines two attention modules to provide both global and local information separately for a translation task. In this manner, the authors place the attention and the convolutional module side by side, thus having different perspectives (globally and locally) of the sentence. They test their approach to 2 common translation benchmarks.