Lite Transformer with Long-Short Range Attention

By Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, Song Han. 2020

392 citations recorded by the source corpus.

Transformer has become ubiquitous in natural language processing (e.g., machine translation, question answering); however, it requires enormous amount of computations to achieve high performance, which makes it not suitable for mobile applications since mobile phones are tightly constrained by the hardware resources and battery. In this paper, we investigate the mobile setting (under 500M Mult-Adds) for NLP tasks to facilitate the deployment on the edge devices. We present Long-Short Range Attention (LSRA), where one group of heads specializes in the local context modeling (by convolution) while another group captures the long-distance relationship (by attention). Based on this primitive, we design Lite Transformer that is tailored for the mobile NLP application. Our Lite Transformer demonstrates consistent improvement over the transformer on three well-established language tasks: machine translation, abstractive summarization, and language modeling. It outperforms the transformer on WMT’14 English-French by 1.2 BLEU under 500M Mult-Adds and 1.7 BLEU under 100M Mult-Adds, and reduces the computation of transformer base model by 2.5x. Further, with general techniques, our Lite Transformer achieves 18.2x model size compression. For language modeling, our Lite Transformer also achieves 3.8 lower perplexity than the transformer around 500M Mult-Adds. Without the costly architecture search that requires more than 250 GPU years, our Lite Transformer outperforms the AutoML-based Evolved Transformer by 0.5 higher BLEU under the mobile setting.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

In Efficiency & Compression

Citations in the corpus

Conformer: Convolution-augmented Transformer for Speech Recognition (2020, methodology): “Another possible way of the architecture is to split the input into parallel branches of multi-headed self attention module and a convolution module with their output concatenated as suggested in [17].”
Of Non-Linearity and Commutativity in BERT (2021, background): “Furthermore, based on works on understanding transformers, in [38] the authors introduce an improved self-attention mechanism, resulting in state of the art performance.”
CvT: Introducing Convolutions to Vision Transformers (2021, methodology): “In NLP and speech recognition, convolutions have been used to modify the Transformer block, either by replacing multi-head attentions with convolution layers [38], or adding additional convolution layers in parallel [39] or sequentially [13], to capture local relationships.”
Random Feature Attention (2021, background): “Answering to this, many approaches have been proposed, e.g., learning the attention spans (Sukhbaatar et al., 2019; Wu et al., 2020), and enhancing the attention computation with recurrent (Hao et al., 2019; Chen et al., 2019) or convolutional (Wu et al., 2019; Mohamed et al., 2019) components.”
HAT: Hardware-Aware Transformers for Efficient Natural Language Processing (2020, methodology): “…adopt the average iteration time of 2.88 for decoding (Gu et al., 2019 Latency BLEU Transformer (Vaswani et al., 2017) 4.3s 25.85 Levenshtein (Gu et al., 2019) 6.5s 25.20 Evolved Transformer (So et al., 2019) 3.7s 25.40 Lite Transformer (Wu et al., 2020) 3.4s 25.79 HAT (Ours) 3.4s 25.92 decoding).”
EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty (2024, background): “…et al., 2015), quantization (Hubara et al., 2018; Shen et al., 2020; Kim et al., 2021; Zadeh et al., 2020; Zafrir et al., 2019), pruning (Gale et al., 2019; Sanh et al., 2020; Kur-tic et al., 2022; Voita et al., 2019), and innovative network architecture designs (Gu & Dao, 2023; Wu et al., 2020).”
Representing Long-Range Context for Graph Neural Networks with Global Attention (2022, background): “Recently, modiﬁcations of the transformer architecture emerge to further improve the efﬁciency [34, 19, 6].”
A Survey on Vision Transformer (2020, background): “Beyond compressing pre-defined transformer models to small ones, some works try to design compact models directly [132, 61].”
TransReID: Transformer-based Object Re-Identification (2021, background): “Two recent works [40, 13] show that a token embedding is mainly determined by its nearby tokens.”
A Survey of Transformers (2021, background): “Similar to low-rank self-attention [45] that decomposes attention into a locality-constrained attention and a low-rank global attention, Lite Transformer [147] proposes to replace each attention module in Transformer with a two-branch structure, where one branch uses attention to capture long-range ”
Long-Short Transformer: Efficient Transformers for Language and Vision (2021, methodology): “To compress the model footprint on edge devices, Lite Transformer [34] combines convolution and self-attention, but it still has quadratic complexity for long sequences.”
Big Little Transformer Decoder (2023, methodology): “Well-known approaches include efficient architecture design [25, 31, 34, 57, 64, 72], quantization [8, 29, 51, 71, 76, 78, 79], pruning [11, 13, 32, 33, 39, 46, 61], and neural architecture search [5, 53, 54, 63, 73, 77].”
Post-Training Quantization for Vision Transformer (2021, background): “[31] proposed Short Range Attention (LSRA) to conduct transformer on edge devices, where one group of heads specializes in the local context modeling (by convolution) while another group specializes in the long-distance relationship modeling.”
CDDFuse: Correlation-Driven Dual-Branch Feature Decomposition for Multi-Modality Image Fusion (2022, methodology): “Considering to balance the performance and computational efﬁciency, we use the LT block [67] as the basic unit of BTE.”
MambaOut: Do We Really Need Mamba for Vision?* (2024, neutral): “To address this issue, a variety of token mixers with linear complexity to token length have been introduced [78], such as dynamic convolution [44, 90, 92], Linformer [84], Longformer [7], Big Bird [105], and Per-former [14].”
Conformer-Kernel with Query Term Independence for Document Retrieval (2020, positive): “…r and hence finding alternatives to explicitly computing the QKᵀ matrix to achieve a complexity of O(n × r)—e.g., [Kitaev et al., 2019, Roy et al., 2020, Tay et al., 2020, Wang et al., 2020], or (iii) A hybrid of both approaches—e.g., [Child et al., 2019, Beltagy et al., 2020, Wu et al., 2020].”
Mesh Graphormer (2021, methodology): “Recent studies [40, 8] show that the combination of convolution and self-attention in a transformer encoder is helpful to improve representation learning.”
Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models (2021, methodology): “, 2020], whereas neural network (NN) models often compose different modules (attention, MLP), and in many applications the MLP layers are the main training bottleneck [Wu et al., 2020].”
Sub-Linear Memory: How to Make Performers SLiM (2020, background): “The FLOPs count has a direct effect on energy consumption (Wu* et al., 2020), a crucial factor for on-device applications.”