HAT : Hardware-Aware Transformers for Efficient Natural Language Processing (2020)

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 SuperTransformer 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 at https://github.com/mit-han-lab/hardware-aware-transformers .

## introduction transformer @xcite has been widely used in natural language processing tasks. by stacking multiple identical encoder/decoder layers with attention modules, it provides a significant performance improvement over previous convolutional or recurrent neural network models @xcite . nevertheless, it is challenging to deploy transformers on mobile devices due to the high computation cost. for instance, in order to translate a sentence with only 30 words, a transformer-big model needs to execute 13g flops and takes 20 seconds on a raspberry pi. such long latency will hurt the user experience on edge devices. thus we need hardware-efficient transformers (figure 1 ). there are two common pitfalls when evaluating the efficiency of a transformer. (1) flops does not reflect the measured latency. although flops is used as an metric for efficiency in prior arts @xcite , it is not a good latency proxy. as in figure 2 (right), models with the same flops can result in very different measured latencies; (2) different hardware prefers different transformer architecture. as in inspired by the success of neural architecture search (nas) @xcite @xcite , we propose to search for hardware-aware transformers (hat) by directly involving the latency feedback into the design loop. in this way, we do not need flops as the latency proxy and can search specialized models for various hardware. we first construct a large search space with arbitrary encoder-decoder attention and heterogeneous transformer layers. traditional transformer has an information bottleneck between the encoder and decoder. arbitrary encoder-decoder attention breaks the bottleneck, allowing all decoder layers to attend to multiple and different encoder layers instead of only the last one. thus low-level information from the encoder can also be used by the decoder. motivated by figure 2 , we introduce heterogeneous transformer layers to allow different layers to have different architecture adapting various hardware. to perform a low-cost search in such a large design space, we first train a transformer supernet -supertransformer, which contains many sub-transformers sharing the weights. we train all subtransformers simultaneously by optimizing the uniformly sampled subtransformers from the su-pertransformer. the performance of a subtransformer with inherited weights from the supertransformer can provide a good relative performance approximation for different architectures trained from-scratch. unlike conventional nas, we only need to pay the supertransformer training cost for once and can evaluate all the models in the design space with it. finally, we conduct an evolutionary search to find the best subtransformer under the hardware latency constraint. experiments show that hat can be naturally incorporated with model compression techniques such as quantization and knowledge distillation. we evaluate hat with wmt'14 en-de, wmt'14 en-fr, wmt'19 en-de, and iwslt'14 de-en tasks on raspberry pi arm cpu, intel xeon cpu, and nvidia titan xp gpu. compared with previous work @xcite @xcite , hat achieves up to 3× speedup, 3.7× smaller size over transformer-big without loss of accuracy. with 12,041× less search cost, hat outperforms the evolved transformer with 2.7× speedup and 3.6× smaller size. it also achieves up to 1.9× speedup over levenshtein and lite transformers with no bleu score loss. with 4-bit quantization, hat can further reach 25× model size reduction. hat has three contributions: (1) hardware-aware and specialization. to our best knowledge, we are the first to directly involve the hardware feedback in the model design, to reduce nlp model latency for target hardware, instead of relying on proxy signals (flops). for different hardware platforms, specialized models for low-latency inference are explored. (2) low-cost neural architecture search with a large design space. we propose arbitrary encoder-decoder attention to break the information bottleneck; and heterogeneous layer to let different layers alter its capacity. a weight-shared supertransformer is trained to search for efficient models at a low cost. (3) design insights. based on the search results, we reveal some design insights: attending to multiple encoder layers is beneficial for the decoder; gpu prefers shallow and wide models while arm cpu prefers deep and thin ones. elastic layer num in encoder elastic head num (self attention) elastic hidden dim in ffn encoder layer 2 encoder layer m elastic embedding dim elastic head num (self attention) elastic hidden dim in ffn elastic embedding dim elastic head num (en-decoder attention) decoder layer n elastic layer num in decoder arbitrary encoder-decoder attention concat ❶ train a supertransformer by uniformly sampling subtransformers with weight sharing 12/3/2019 _ e _1895988. e:///u e / a a /d ad / _ e _1895988. 1/1 c ea ed b ca lina cani f m he n n p jec iot 12/3/2019 _g _1132940. g e:///u e / a a g/d ad / _g _1132940. g 1/1 c ea ed b mi ha pe i hche f m he n n p jec gpu 12/3/2019 electronic de ice le:///users/hanrui ang/do nloads/noun_cpu_2880768.s g 1/1 c ea ed b m ha ad kh i a f he n p jec cpu latency predictor ❷ collect hardware latency datasets 12/3/2019 noun_database_2997685.svg file:///users/hanruiwang/downloads/noun_database_2997685.svg 1/1 layer num embed dim heads num latency ❸ train a latency predictor for each hardware evolutionary search engine 12/3/2019 _ c a _2500412. :///u / a a /d ad / _ c a _2500412. 1/1 c ea ed b a ice de ig f he n p jec ❹ evolutionary search latency subtransformer architecture val loss encoder layer 1 decoder layer 1 e.g. 4 heads e.g. 2 heads subtransformer architecture supertransformer figure 3: hat overview. a large design space is constructed with arbitrary encoder-decoder attention and heterogeneous layers. (1) train a weight-shared supertransformer by iteratively optimizing randomly sampled subtransformers. it can provide a performance proxy for subtransformers. (2) collect (subtransformer architecture, latency) data pairs on the target hardware. (3) train a latency predictor for each hardware to provide fast and accurate latency feedback. (4) perform an evolutionary search with hardware latency constraint to find the model with the lowest validation loss. (5) finally, the searched model is trained from scratch to get the final performance. elastic #layers in encoder elastic #heads (self attention) elastic hidden dim in ffn encoder layer 2 encoder layer m elastic embedding dim elastic #heads (self attention) elastic hidden dim in ffn elastic embedding dim elastic #heads (en-decoder attention) decoder layer n elastic #layers in decoder arbitrary encoder-decoder attention concat ❶ training supertransformer by uniformly sampling and training all subtransformers 12/3/2019 _ e _1895988. e:///u e / a a /d ad / _ e _1895988. 1/1 c ea ed b ca lina cani f m he n n p jec iot 12/3/2019 _g _1132940. g e:///u e / a a g/d ad / _g _1132940. g 1/1 c ea ed b mi ha pe i hche f m he n n p jec gpu 12/3/2019 electronic de ice le:///users/hanrui ang/do nloads/noun_cpu_2880768.s g 1/1 c ea ed b m ha ad kh i a f he n p jec cpu hardware latency predictor ❷ collecting hardware latency dataset 12/3/2019 noun_database_2997685.svg file:///users/hanruiwang/downloads/noun_database_2997685.svg 1/1 #layers embed dim #heads latency ❸ training hardware latency predictor evolutionary search engine 12/3/2019 _ c a _2500412. :///u / a a /d ad / _ c a _2500412. 1/1 c ea ed b a ice de ig f he n p jec ❹ evolutionary search latency subtransformer architecture val loss encoder layer 1 decoder layer 1 e.g. 4 heads e.g. 2 heads subtransformer architecture supertransformer encoder layer j self attention decoder layer m encoder layer i kj vj ki vi concat en-decoder attention q v k ## proposed approaches an overview of the hat framework is shown in figure 3 . we firstly train a supertransformer with a large design space. then, for a given hardware platform, we collect a dataset of (subtransformer architecture, measured latency) pairs for different models, and train a latency predictor. finally, we conduct an evolutionary search with a latency constraint to find an efficient model specialized for the target hardware.