Source-linked AI summary
HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
Hanrui Wang, Zhanghao Wu, Zhijian Liu, Han Cai, Ligeng Zhu, Chuang Gan, Song Han
TL;DR
Transformers are difficult to deploy on mobile and resource-constrained hardware because of high computation cost, and FLOPs may not reflect measured latency across platforms. HAT uses hardware-aware neural architecture search with a large, weight-shared SuperTransformer design space to discover specialized low-latency models, achieving speedups and smaller models across machine-translation evaluations.
Problem
Deploying Transformers on mobile devices is challenging because of high computation cost, while FLOPs does not reliably reflect measured latency and hardware platforms prefer different architectures.
Method
HAT searches a large design space with arbitrary encoder-decoder attention and heterogeneous layers using a weight-shared SuperTransformer, then incorporates hardware latency feedback to find specialized models.
Results
HAT achieves up to 1.6×, 3×, and 3.4× speedup with up to 1.4×, 3.7×, and 4× smaller size than baselines across the reported evaluations, with over 12000× smaller total GPU-hour than Evolved Transformer.
Takeaways & Limitations
HAT discovers high-performance, low-latency models for efficient Transformer deployment across various hardware platforms at substantially lower search cost.
Abstract
from arXiv · showhide
Transformers are ubiquitous in Natural Language Processing (NLP) tasks, but they are difficult to be deployed on hardware due to the intensive computation. To enable low-latency inference on resource-constrained hardware platforms, we propose to design Hardware-Aware Transformers (HAT) with neural architecture search. We first construct a large design space with $\textit{arbitrary encoder-decoder attention}$ and $\textit{heterogeneous layers}$. Then we train a $\textit{SuperTransformer}$ that covers all candidates in the design space, and efficiently produces many $\textit{SubTransformers}$ with weight sharing. Finally, we perform an evolutionary search with a hardware latency constraint to find a specialized $\textit{SubTransformer}$ dedicated to run fast on the target hardware. Extensive experiments on four machine translation tasks demonstrate that HAT can discover efficient models for different hardware (CPU, GPU, IoT device). When running WMT'14 translation task on Raspberry Pi-4, HAT can achieve $\textbf{3}\times$ speedup, $\textbf{3.7}\times$ smaller size over baseline Transformer; $\textbf{2.7}\times$ speedup, $\textbf{3.6}\times$ smaller size over Evolved Transformer with $\textbf{12,041}\times$ less search cost and no performance loss. HAT code is https://github.com/mit-han-lab/hardware-aware-transformers.git
1 Introduction
HAT targets Transformer deployment latency by using hardware feedback rather than FLOPs alone, while searching a flexible architecture space with shared weights. It is designed to find specialized models for different hardware platforms at lower search cost.
- Motivation: 13G FLOPs and 20 seconds are required to translate a 30-word sentence with Transformer-Big on a Raspberry Pi.Such latency hurts user experience on edge devices.
- Motivation: FLOPs does not reliably predict measured latency, and different hardware platforms favor different Transformer architectures.Embedding dimension strongly affects Raspberry Pi latency but has little GPU impact.
- Proposed direction: HAT directly incorporates hardware latency feedback into neural architecture search to avoid using FLOPs as a latency proxy.This enables specialized models for various hardware platforms.
- Design space: Arbitrary encoder-decoder attention lets decoder layers attend to multiple encoder layers, while heterogeneous layers vary architecture across layers.These choices address the encoder-decoder information bottleneck and adapt model capacity.
- Search framework: A weight-shared SuperTransformer provides a performance proxy for many SubTransformers, reducing the cost of evaluating architectures in the design space.The search then uses evolutionary optimization under hardware latency constraints.
- Evaluation: HAT is evaluated on WMT’14 En-De, WMT’14 En-Fr, WMT’19 En-De, and IWSLT’14 translation tasks.The supplied passage identifies the four evaluation tasks but does not provide the full task results.
2 Proposed Approaches
HAT combines a flexible Transformer design space with a weight-shared SuperTransformer and hardware-specific latency prediction. Evolutionary search then selects latency-constrained SubTransformers for final training.
- Design Space: HAT breaks standard Transformer design conventions by allowing arbitrary encoder-decoder attention and heterogeneous layers.Decoder layers can attend to multiple encoder layers, and layer capacities can differ.
- Design Space: Different layers can use different numbers of heads, hidden dimensions, and embedding dimensions, while encoder and decoder dimensions remain internally consistent.Encoder and decoder layer counts are also elastic.
- SuperTransformer: A SuperTransformer contains every SubTransformer as a weight-shared subnetwork, sharing common embedding, fully connected, attention, and early-layer weights.This allows architectures to be evaluated without fully training each one.
- SuperTransformer: Uniformly sampled SubTransformers are optimized during SuperTransformer training, whose validation evaluations provide relative performance proxies after training.The SuperTransformer uses the same training steps as a baseline Transformer model.
- Latency Modeling: An offline latency predictor is trained from architecture–measured-latency pairs because online measurement would require hundreds of inferences per sampled SubTransformer.The latency predictor has an average prediction error (RMSE) of 0.1s.
- Evolutionary Search: Evolutionary search queries predicted latency and validation loss, retains models below the hardware constraint, and trains the selected models from scratch.The search therefore enforces the target hardware latency requirement during architecture selection.
3 Experiments
The experiments evaluate HAT across four translation tasks, multiple baselines, and three hardware platforms using measured latency, translation quality, model size, and search-cost metrics. The search space and evolutionary-search procedure are configured to produce hardware-specialized models.
- Experiment Setup: Experiments cover WMT’14 En-De, WMT’14 En-Fr, WMT’19 En-De, and IWSLT’14 De-En translation tasks.The baselines include Transformer, Levenshtein Transformer, Evolved Transformer, and Lite Transformer.
- Experiment Setup: Latency is measured on Raspberry Pi-4 ARM CPU, Intel Xeon CPU, and Nvidia TITAN Xp GPU using repeated translations of fixed-length sentences.Measurements average the middle 80% of 300 runs; sentence lengths are 30 tokens for WMT and 23 for IWSLT.
- Implementation Details: The WMT SuperTransformer varies embedding dimension, hidden dimension, attention heads, decoder depth, and encoder-decoder attention choices across approximately 10^15 SubTransformers.The encoder depth is fixed at six because encoder computation accounts for less than 5% of measured latency.
- Implementation Details: For each hardware platform, HAT collects 2000 architecture-latency samples, trains a three-layer MLP latency predictor, and runs evolutionary search for 30 iterations.The evolutionary search uses a population of 125 with mutation and crossover populations of 50 each.
- Training: Baseline and searched SubTransformer training uses task-specific optimization schedules, including Adam with cosine learning-rate scheduling for WMT.WMT models train for 40K steps, while IWSLT models train for 50K steps.
- Evaluation: Figure 7 compares inference latency and BLEU trade-offs for WMT’14 En-De and En-Fr across the three hardware platforms.The figure reports up to 3× faster inference and 3.7× smaller size than Transformer-Big.
4 Results
HAT consistently finds hardware-efficient translation models that preserve or improve quality while reducing latency, model size, and search cost. The results also show hardware-specific architectural patterns and compatibility with inherited-weight training, quantization, and knowledge distillation.
- Performance Comparisons: 3× faster and 3.7× smaller than Transformer-Big on ARM CPU with the same BLEU; HAT also achieves over 2× speedup on Intel CPU.On Nvidia GPU, dimension scaling alone provides little latency reduction, whereas HAT finds low-latency, high-performance models.
- Performance Comparisons: Up to 1.6×, 3×, and 3.4× speedups accompany up to 1.4×, 3.7×, and 4× smaller sizes versus Transformer and Evolved Transformer baselines.Table 2 compares latency, model size, FLOPs, BLEU, and training cost, including CO2 emissions and cloud cost.
- Performance Comparisons: 1.3× faster than Transformer with higher BLEU and 1.9× faster than Levenshtein with 0.7 higher BLEU, HAT also outperforms Lite Transformer at similar latency.These comparisons use Raspberry Pi ARM CPU latency on WMT’14 En-De.
- Design Insights: 10% of decoder layers attend to three encoder layers and 40% attend to two, supporting arbitrary encoder-decoder attention choices.The observed attention patterns are presented as evidence for the necessity of arbitrary encoder-decoder attention.
- Design Insights: GPU-specialized models are wide and shallow, whereas Raspberry Pi models are deep and thin, reflecting different latency sensitivities to dimensions and depth.These patterns also guide manual hardware-specific Transformer design.
- Ablation Study: HAT achieves higher BLEU with 1.5× lower latency and 1.5× smaller size than the largest SubTransformer, showing that larger models do not necessarily perform better.Evolutionary search also finds lower-loss models than random search.
- Search Efficiency: 12,041× lower search cost than Evolved Transformer results from training candidate models together in a weight-sharing SuperTransformer rather than individually.The inherited-weight performance proxy preserves the relative ordering of SubTransformers closely enough to support search.
- Compression Compatibility: 4-bit quantization reduces HAT model size by 25× with only 0.1 BLEU loss, while 8-bit quantization achieves 0.1 higher BLEU than its full-precision version.HAT can also be combined with knowledge distillation, which raises student BLEU from 25.8 to 26.1.
5 Related Work
Related work spans Transformer modifications, decoding strategies, and neural architecture search for efficient model design. HAT is positioned as complementary to these approaches rather than as a replacement for them.
- Transformer: Prior Transformer research modifies attention, adds convolutional modules, propagates information across encoder layers, or replaces attention with alternative mechanisms.HAT is described as orthogonal to these Transformer modifications and combinable with them.
- Transformer: Non- or partially-autoregressive models reduce decoding iterations but can suffer from low performance, motivating alternative efficiency strategies.The related work also includes learned combinations of encoder outputs as decoder inputs.
- Neural Architecture Search: Neural architecture search research includes black-box evolutionary or reinforcement-learning optimization and differentiable architecture methods, especially in computer vision.HAT applies hardware-feedback-driven evolutionary search to Transformer design.
6 Conclusion
HAT addresses efficient Transformer deployment across hardware platforms by combining hardware-aware architecture search with a large design space and an efficient weight-shared SuperTransformer.
- HAT is proposed to solve efficient Transformer deployment challenges across various hardware platforms.
- The framework searches a large design space using an efficient weight-shared SuperTransformer.
- HAT consumes four orders of magnitude less search cost than the prior Evolved Transformer while discovering high-performance, low-latency models.
A.1 SubTransformer Performance Proxy
Validation loss of SubTransformers inherited from the SuperTransformer tracks the BLEU scores of independently trained SubTransformers, supporting validation loss as a performance proxy.
- Larger validation loss corresponds to lower BLEU scores for the evaluated SubTransformers.
- Validation loss can serve as a performance proxy for SubTransformers.
A.2 Visualizations of Searched Models on WMT’14 En-De Task
HAT finds different WMT’14 En-De SubTransformer architectures for Raspberry Pi CPU and Nvidia GPU while maintaining similar BLEU scores.
- The Raspberry Pi model is deep and thin, whereas the GPU model is shallow and wide.
- 28.10 BLEU for Raspberry Pi CPU and 28.15 BLEU for Nvidia GPU indicate similar translation quality.
- The searched SubTransformers for Raspberry Pi ARM CPU and Nvidia GPU are different.
A.3 Latency, BLEU and SacreBLEU of searched HAT models.
The paper reports specific latency, BLEU, and SacreBLEU measurements for searched HAT models presented in Figures 7 and 8.
- Table 8 reports latency, BLEU, and SacreBLEU scores for the searched HAT models.