Source-linked AI summary
Lite Transformer with Long-Short Range Attention
Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, Song Han
TL;DR
High-performing Transformer NLP models are computationally expensive for mobile and edge devices. The paper introduces Lite Transformer with LSRA, specializing different heads in local and long-distance context modeling, and reports consistent improvements over Transformer across three language tasks under mobile resource constraints.
Problem
Transformer-based NLP models require substantial computation, making high-performing models unsuitable for many mobile and edge-device settings.
Method
Lite Transformer uses LSRA, combining specialized attention and convolutional branches for long-distance and local context modeling, respectively.
Results
Lite Transformer consistently improves over Transformer across machine translation, abstractive summarization, and language modeling under mobile resource constraints.
Takeaways & Limitations
Manual design guided by bottleneck analysis and specialized heads can outperform AutoML-based Evolved Transformer while avoiding its costly architecture search.
Takeaways & Limitations
The evaluation targets mobile settings defined by computation and parameter constraints, including a 500M Mult-Adds budget.
Abstract
from arXiv · showhide
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 that are tightly constrained by the hardware resources and battery. In this paper, we present an efficient mobile NLP architecture, Lite Transformer to facilitate deploying mobile NLP applications on edge devices. The key primitive is the Long-Short Range Attention (LSRA), where one group of heads specializes in the local context modeling (by convolution) while another group specializes in the long-distance relationship modeling (by attention). Such specialization brings consistent improvement over the vanilla transformer on three well-established language tasks: machine translation, abstractive summarization, and language modeling. Under constrained resources (500M/100M MACs), Lite Transformer outperforms transformer on WMT'14 English-French by 1.2/1.7 BLEU, respectively. Lite Transformer reduces the computation of transformer base model by 2.5x with 0.3 BLEU score degradation. Combining with pruning and quantization, we further compressed the model size of Lite Transformer by 18.2x. For language modeling, Lite Transformer achieves 1.8 lower perplexity than the transformer at around 500M MACs. Notably, Lite Transformer outperforms the AutoML-based Evolved Transformer by 0.5 higher BLEU for the mobile NLP setting without the costly architecture search that requires more than 250 GPU years. Code has been made available at https://github.com/mit-han-lab/lite-transformer.
1 INTRODUCTION
Lite Transformer addresses the high computational cost of Transformer-based NLP on mobile devices through LSRA, which separates local and long-distance context modeling. Experiments report improvements over Transformer across translation, summarization, and language modeling, while manual design also avoids costly AutoML search.
- Transformer achieves strong NLP performance but requires more than 10G Mult-Adds for a 30-word sentence, exceeding many edge-device capabilities.
- LSRA trades FFN computation for wider attention and assigns specialized heads to local convolutional and global attention-based context modeling.Lite Transformer is built by stacking the LSRA primitive.
- 1.2 BLEU under 500M Mult-Adds and 1.7 BLEU under 100M Mult-Adds are reported on WMT 2014 English-French versus Transformer.
- 18.2× model size compression is achieved when Lite Transformer is combined with pruning and quantization.
- 0.5 higher BLEU than AutoML-based Evolved Transformer is achieved under mobile settings, without its more than 250 GPU years of architecture search.
2 RELATED WORK
Prior work addresses sequence modeling, transformer design, automated architecture search, and model compression, but efficient mobile NLP remains constrained by computation and search cost.
- RNNs and CNNs: RNNs have long supported sequence modeling but are difficult to parallelize across sequences because of temporal dependency.
- Transformers: Attention-based transformers capture global context through pairwise correlations and achieve state-of-the-art performance by stacking self-attention layers.
- Automated Model Design: AutoML-based NLP architecture search can impose prohibitive computational and environmental costs, with search emissions nearly 5× a car’s average lifetime emissions.
- Transformers: Transformer research has explored scaling batch size, relative position representations, weighted multi-head attention, and other architectural variants.
- Automated Model Design: Neural architecture search automates model design, increasingly incorporating hardware resource constraints into optimization.
- Model Acceleration: Model acceleration also compresses existing networks through neuron or channel pruning and network quantization.
3 IS BOTTLENECK EFFECTIVE FOR 1-D ATTENTION?
The section argues that the conventional channel bottleneck may be inefficient for one-dimensional attention, especially when sequence lengths are small. It motivates a flattened transformer block that preserves channel dimensions and leaves more computation available for attention optimization.
- Attention Computation: Attention models short- and long-term relationships through pairwise dot products, but this operation introduces massive computation.
- Bottleneck Design: The conventional bottleneck reduces attention channels before computation and expands them afterward to control attention cost.
- Bottleneck Design: For translation sequences of roughly 20–30 tokens, attention costs O(4Nd^2 + N^2d), while the FFN costs O(2 × 4Nd^2).
- Flattened Transformer: The authors argue that reducing and restoring channels is not optimal for one-dimensional attention and replace it with a flattened transformer block.
- Flattened Transformer: The flattened design makes attention the major computational component, creating more scope for further attention optimization.
- Flattened Transformer: On WMT’14 English–French, the flattened modification achieves comparable performance at slightly larger computation before further optimization.
4 LONG-SHORT RANGE ATTENTION (LSRA)
LSRA separates global and local context modeling into parallel specialized branches. Attention handles global relationships, while convolution efficiently captures local context, reducing redundancy and computation.
- Attention maps contain sparse global relationships and diagonal local relationships, indicating that conventional attention models both context types.
- LSRA specializes context modeling by assigning global relationships to an attention branch and local relationships to a convolutional branch.The two branches operate in parallel, with channel-wise input splitting before the outputs are mixed by the FFN.
- Splitting the input along the channel dimension reduces overall computation by 2× while preserving separate global and local processing.
- Lite Transformer attention focuses on global contexts without the diagonal pattern, leaving local-context capture to the convolution branch.
5 EXPERIMENTAL SETUP
The experiments define mobile NLP settings through computation and parameter constraints and evaluate Lite Transformer across translation, summarization, and language modeling. The setup compares its efficiency and performance with transformer-based baselines under these resource limits.
- 5.1 MOBILE SETTINGS: 500M Mult-Adds is the mobile computation limit for machine translation with 30-token sequences, corresponding to 1G FLOPs.The limit is motivated by the approximately 48G FLOPS performance of an ARM Cortex-A72 mobile CPU.
- 5.1 MOBILE SETTINGS: 10M parameters is the defined mobile model-size constraint, reflecting application download, runtime-memory, and disk-space limitations.
- 5.3 ARCHITECTURE: The experiments use transformer-base-derived sequence-to-sequence models and fairseq’s transformer-base reimplementation as the backbone.
- 5.3 ARCHITECTURE: The architecture flattens the transformer bottleneck and replaces self-attention with parallel attention and convolutional branches.The implementation uses four heads per branch and task-specific convolution kernel sizes.
6 RESULTS
Lite Transformer consistently improves mobile translation and language-task efficiency, while preserving summarization quality and supporting substantial compression. Its gains extend across resource constraints and remain compatible with pruning and quantization.
- Machine Translation: 1.6 BLEU improvement was achieved over the transformer on IWSLT’14 De-En around 100M Mult-Adds.
- Machine Translation: 1.2 and 1.7 BLEU improvements were obtained on WMT’14 En-De and En-Fr, respectively, under 100M Mult-Adds.
- Abstractive Summarization: Similar F1-Rouge scores to the transformer were obtained on CNN-DailyMail with 2.4× less computation and 2.5× less model size.
- Language Modeling: 1.8 lower test perplexity was achieved on WIKITEXT-103 under similar resource constraints.
- Model Compression: 18.2× model size compression was achieved on WMT En-Fr using pruning and quantization, with negligible BLEU score degradation.
- Comparison with Automated Design: 0.5 and 0.2 BLEU advantages over Evolved Transformer were reported under 100M and 300M Mult-Adds, respectively.
- Comparison with Automated Design: More than 250 GPU years and 626155 lbs CO2 emissions are associated with the neural architecture search used by comparison systems.
7 CONCLUSION
The paper concludes that LSRA enables Lite Transformer to specialize local and long-distance modeling for mobile NLP. Across multiple language applications, the model improves over Transformer and surpasses Evolved Transformer under mobile constraints.
- LSRA assigns some heads to local context modeling and others to long-distance relationship modeling.
- Lite Transformer is designed for mobile settings below 500M Mult-Adds to facilitate deployment on edge devices.
- Consistent improvements over Transformer are demonstrated across multiple language applications.
- Lite Transformer surpasses Evolved Transformer under mobile settings despite Evolved Transformer requiring costly architecture search.
A.1 ADDITIONAL VISUALIZATION OF ATTENTION WEIGHTS
The visualization compares attention patterns from a base Transformer and LSRA, illustrating how the proposed design separates local and global feature extraction.
- Attention weights are visualized for the base Transformer and LSRA using fully trained models on WMT En-Fr translation.
- Conventional attention emphasizes local relationships, while LSRA uses convolution for locality and attention for global feature extraction.