HAT: Hardware-Aware Transformers for Efficient Natural Language Processing

Abstract

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 arbitrary encoder-decoder attention and heterogeneous layers. Then we train a Super-Transformer that covers all candidates in the design space, and efficiently produces many SubTransformers with weight sharing. Finally, we perform an evolutionary search with a hardware latency constraint to find a specialized 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 3× speedup, 3.7× smaller size over baseline Transformer; 2.7× speedup, 3.6× smaller size over Evolved Transformer with 12,041× less search cost and no performance loss. HAT is open-sourced.

AutoTinyBERT: Automatic Hyper-parameter Optimization for Efficient Pre-trained Language Models

“…We use Lat(•) to predict the latency of the candidates to filter out the candidates that do not meet the latency constraint. Lat(•) is built with the method by Wang et al (2020a), which first samples about 10k architectures from A and collects their inference time on target devices, and then uses a feed-forward network to fit the data. For more details of evolutionary algorithm, please refer to Appendix C. Note that we can use different methods in search process, such as random search and more advanced search, which is left as future work.…”

“…More- 6 The first 16 models https://github.com/ google-research/bert from 2L128D to 8L768D. 4-192-768-12-192 4-256-480-12-192 17.0/17.0× over, we introduce HAT (Wang et al, 2020a), as a baseline of one-shot learning. HAT focuses on the search space of non-identical layer structures.…”

“…In this work, we only consider the case of identical structure for each Transformer layer, instead of the non-identical Transformer (Wang et al, 2020a) or other heterogeneous modules (Xu et al, 2021) (such as convolution units). It has two advantages: (1) it reduces an exponential search space of O( * for m = 1 → M do 7:…”

“…To make SuperPLM more effective, we propose practical techniques including the head sub-matrix extraction and efficient batch-wise training, and particularly limit the search space to the models with identical layer structure. Furthermore, by using SuperPLM, we leverage search algorithm (Xie and Yuille, 2017;Wang et al, 2020a) to find hyperparameters for various latency constraints.…”

LV-BERT: Exploiting Layer Variety for BERT

“…Traditional NAS methods Zhu et al, 2020) use downstream task performance as the objective to search for task-specific models. Instead, similar to the work by Khetan and Karnin (2020) et al, 2015).…”