SlowGuess/ACL_2020_OCR huggingface dataset:
Hanrui Wang $^{1}$ , Zhanghao Wu $^{1}$ , Zhijian Liu $^{1}$ , Han Cai $^{1}$ , Ligeng Zhu $^{1}$ , Chuang Gan $^{2}$ , Song Han $^{1}$
$^{1}$ Massachusetts Institute of Technology,  $^{2}$ MIT-IBM Watson AI Lab
{hanrui, zhwu, zhijian, hancai, ligeng, chuangg, songhan}@mit.edu
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 Sub Transformers 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 \times$  speedup,  $3.7 \times$  smaller size over baseline Transformer;  $2.7 \times$  speedup,  $3.6 \times$  smaller size over Evolved Transformer with  $12,041 \times$  less search cost and no performance loss. HAT is open-sourced.
Transformer (Vaswani et al., 2017) 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 (Kim, 2014).
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
Figure 1: Framework for searching Hardware-Aware Transformers. We first train a SuperTransformer that contains numerous sub-networks, then conduct an evolutionary search with hardware latency feedback to find one specialized SubTransformer for each hardware.
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 (Howard et al., 2017; Wu et al., 2020), 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 Table 1, the Transformer model optimized on one hardware is sub-optimal for another because latency is influenced by different factors on different hardware platforms. For example, the embedding size has significant impact on the Raspberry Pi latency but hardly influences the GPU latency (Figure 2).
Inspired by the success of Neural Architecture Search (NAS) (Bender et al., 2018; Guo et al., 2019; Pham et al., 2018; Cai et al., 2019a), we propose to search for Hardware-Aware Transformers (HAT)
7675
Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 7675-7688 July 5 - 10, 2020. ©2020 Association for Computational Linguistics
Figure 2: Latency of different Transformer models on different hardware. We find (1) FLOPs does not reflect the real measured latency; (2) Latency influencing factors of different hardware are contrasting. Thus we need to consider hardware latency feedback to design specialized models for different hardware.
Measured On →	GPU Latency	ARM CPU Latency
Specialized For ↓	BLEU
HAT (GPU)	28.10	147 ms	6491 ms
HAT (ARM CPU)	28.15	184 ms	6042 ms
Table 1: BLEU score and measured inference latency of HAT on WMT'14 En-De task. The efficient model for GPU is not efficient for ARM CPU and vice versa.
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 SubTransformers sharing the weights. We train all SubTransformers simultaneously by optimizing the uniformly sampled SubTransformers from the SuperTransformer. 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 (Vaswani et al., 2017; So et al., 2019; Gu et al., 2019; Wu et al., 2020), HAT achieves up to  $3 \times$  speedup,  $3.7 \times$  smaller size over Transformer-Big without loss of accuracy. With  $12,041 \times$  less search cost, HAT outperforms the Evolved Transformer with  $2.7 \times$  speedup and  $3.6 \times$  smaller size. It also achieves up to  $1.9 \times$  speedup over Levenshtein and Lite Transformers with no BLEU score loss. With 4-bit quantization, HAT can further reach  $25 \times$  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.
$\bullet$  Train a SuperTransformer by uniformly sampling SubTransformers with weight sharing
$\Theta$  Collect Hardware Latency Datasets
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.
Figure 4: Arbitrary Encoder-Decoder Attention. Each encoder-decoder attention in one decoder layer can attend to the outputs from multiple encoder layers, fully leveraging the features extracted by the encoder.
2 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.
2.1 Design Space
We construct a large design space by breaking two conventions in the Transformer design: (1) All decoder layers only attend to the last encoder layer; (2) All the layers are identical.
Arbitrary Encoder-Decoder Attention. Different encoder layers extract features on different abstraction levels. Conventionally, all the decoder lay
ers only attend to the last encoder layer. It forms an information bottleneck that forces all the decoder layers to learn solely from the high abstraction level and ignore the low-level information. To break the bottleneck, we propose Arbitrary Encoder-Decoder Attention to learn the most suitable connections between the encoder and the decoder. Each decoder layer can choose multiple encoder layers to attend. The key and value vectors from encoder layers are concatenated in the sentence length dimension (Figure 4) and fed to the encoder-decoder cross attention module. The mechanism is efficient because it introduces no additional parameters. The latency overhead is also negligible. For example, with each decoder layer attending to two encoder layers, the latency of Transformer-Base on Nvidia TITAN Xp GPU barely increases by  $0.4\%$ . It improves the model capacity by allowing attention to different abstraction levels.
Heterogeneous Transformer Layers. Previous Transformers repeat one architecture for all layers. In HAT, instead, different layers are heterogeneous, with different numbers of heads, hidden dim, and embedding dim. In attention layers, different heads are used to capture various dependencies. However, Voita et al. (2019) shows that many heads are redundant. We thereby make attention head number elastic so that each attention module can decide its necessary number of heads.
In the FFN layer, the input features are cast to a higher dimension (hidden dim), followed by an
7677
Figure 5: Weight Sharing of the SuperTransformer. All SubTransformers share the front portion of word embeddings, and weights in the fully-connected layers.
activation layer. Traditionally, the hidden dim is set as  $2 \times 4 \times 4$  of the embedding dim, but this is sub-optimal since different layers need different capacities depending on the feature extraction difficulty. We hence make the hidden dim elastic.
Moreover, we also support elastic embedding dim of encoder and decoder, but it is consistent inside encoder/decoder. The number of encoder & decoder layers are also elastic to learn the proper level of feature encoding and decoding. Other design choices such as the length of  $Q, K, V$  vectors in attention modules can be naturally incorporated in our framework, which we leave for future work.
2.2 SuperTransformer
It is critical to have a large design space in order to find high-performance models. However, training all the models and comparing their BLEU scores is infeasible. We thus propose SuperTransformer, a supernet for performance approximation, which can judge the performance of a model without fully training it. The SuperTransformer is the largest model in the search space with weight sharing (Pham et al., 2018; Liu et al., 2019; Cai et al., 2019a). Every model in the search space (a SubTransformer) is a part of the SuperTransformer. All SubTransformers share the weights of their common parts. For elastic embedding dim, all SubTransformers share the front portion of the longest word embedding and corresponding FC layer weights. As in Figure 5, for elastic FFN hidden dim, the front part of the FC weights is shared. For elastic head number in attention modules, the whole  $Q, K, V$  vectors (the lengths are fixed in our design space) are shared by dividing into head_number parts. Elastic layer numbers let all SubTransformers share the first several layers.
Figure 6: The latency predictor is very accurate, with an average prediction error (RMSE) of 0.1s.
In the SuperTransformer training, all possible SubTransformers are uniformly sampled, and the corresponding weights are updated. In practice, the SuperTransformer only needs to be trained for the same steps as a baseline Transformer model, which is fast and low-cost. After training, we can get the performance proxy of sampled models in the design space by evaluating the corresponding SubTransformers on the validation set without training.
2.3 Evolutionary Search for SubTransformer
Given a latency requirement, we perform an evolutionary search to find a satisfactory SubTransformer. There are two ways to evaluate the hardware latency of a SubTransformer: (1) Online measurement in which we measure the models during the search process. (2) Offline, where we train a latency predictor to provide the latency. We apply the offline method here because it is fast and accurate. For the online method, a single sampled SubTransformer requires hundreds of inferences to get an accurate latency, which lasts for minutes and slows down the searching. For the offline method, we encode the architecture of a SubTransformer into a feature vector, and predict its latency instantly with a multi-layer perceptron (MLP). Trained with thousands of real latency data points, the predictor yields high accuracy (Figure 6). Note that the predicted latency is only used in the search process, and we report real measured latency in the experiment section. Compared with deducing a closed-form latency model for each hardware, the latency predictor method is more general and faster.
We use an evolutionary algorithm to conduct the search process. As in Figure 3, the search engine queries the latency predictor for SubTransformer latency, and validates the loss on the validation set. The engine only adds SubTransformers with
latency smaller than the hardware constraint to the population. We then train the searched models from scratch to obtain the final performance.
We conduct experiments on four machine translation tasks: WMT'14 En-De, WMT'14 En-Fr, WMT'19 En-De, and IWSLT'14 De-En, consisting of 4.5M, 36.3M, 43.0M, and 160K pairs of training sentences, respectively. For WMT'14 En-De, we apply 32K source-target BPE vocabulary, train on WMT'16, validate on newstest2013 and test on newstest2014, replicating Wu et al. (2019b); For WMT'14 En-Fr, we use 40K source-target BPE vocabulary, validate on newstest2012&2013, and test on newstest2014, replicating Gehring et al. (2017). WMT'19 En-De adopts 49.6K source-target BPE vocabulary, validates on newstest2017, and tests on newstest2018, the same as Junczys-Dowmunt (2019). We use 10K joint BPE vocabulary in lower case for IWSLT'14 De-En (Grave et al., 2017).
3.2 Experiment Setups
Baselines. Our baseline models are Transformer (Vaswani et al., 2017), Levenshtein Transformer (Gu et al., 2019), both with the Ott et al. (2019) implementation, Evolved Transformer (So et al., 2019) and Lite Transformer (Wu et al., 2020).
Evaluation Metrics. For evaluation, we use beam four and length penalty 0.6 for WMT, and beam five for IWSLT (Vaswani et al., 2017). All BLEUs are calculated with case-sensitive tokenization<sup>1</sup>, but we also apply the compound splitting  $\mathrm{BLEU}^2$  for WMT, the same as Vaswani et al. (2017). We test the model with the lowest validation set loss for WMT and the last ten checkpoints averaged for IWSLT.
We test the latency of the models by measuring translation from a source sentence to a target sentence with the same length. The length is the average output length on the test set - 30 for WMT and 23 for IWSLT. For each model, we measure the latency for 300 times, remove the fastest and slowest  $10\%$  and then take the average of the rest  $80\%$ . We conduct experiments on three representative hardware platforms: Raspberry Pi-4 with an ARM Cortex-A72 CPU, Intel Xeon E5-2640 CPU, and Nvidia Titan Xp GPU.
SuperTransformer Setups. The SuperTransformer for WMT has the following design space: [512, 640] for embedding dim, [1024, 2048, 3072] for hidden dim, [4, 8] for the head number in all attention modules, [1, 2, 3, 4, 5, 6] for decoder layer number. Due to decoder auto-regression, encoder only accounts for less than  $5\%$  of the measured latency; thereby, we set the encoder layer number fixed as 6. For arbitrary encoder-decoder attention, each decoder can choose to attend to the last one, two, or three encoder layers. The SuperTransformer design space for IWSLT is the same as WMT except for [2048, 1024, 512] for hidden dim and [4, 2] for head number. We set the  $Q, K, V$  vector dim fixed as 512. The design space contains around  $10^{15}$  possible SubTransformers and covers a wide range of model size and latency (largest =  $6 \times$  smallest). We train the SuperTransformers of WMT for 40K steps and 50K steps for IWSLT.
Hardware-Aware Evolutionary Search Setups. The input of the latency predictor is a feature vector of SubTransformer architecture with ten elements: layer number, embed dim, average hidden dim, average self-attention heads, of both encoder and decoder; plus average encoder-decoder attention heads, and the average number of encoder layers each decoder layer attends. A dataset of 2000 (SubTransformer architecture, measured latency) samples for each hardware is collected, and split into train-valid:test=8:1:1. We normalize the features and latency, and train a three-layer MLP with 400 hidden dim and ReLU activation. We choose three-layer because it is more accurate than the one-layer model, and over three layers do not improve accuracy anymore. With the predictor, we conduct an evolutionary search for 30 iterations in the SuperTransformer, with population 125, parents population 25, mutation population 50 with 0.3 probability and crossover population 50.
Training Settings. Our training settings are in line with Wu et al. (2019b) and Wu et al. (2020). For WMT, we train for 40K steps with Adam optimizer and a cosine learning rate (LR) scheduler (Kingma and Ba, 2015; Loshchilov and Hutter, 2017), where the LR is linearly warmed up from  $10^{-7}$  to  $10^{-3}$ , and then cosine annealed. For IWSLT, we train for 50K steps with inverse square root LR scheduler. The baseline Transformers are trained with the same settings as the searched SubTransformers for fair comparisons.
$^{1}$ https://github.com/moses-smt/mosesdecoder
$^{2}$ https://github.com/tensorflow/tensor2tensor
7679
Figure 7: Inference latency and BLEU trade-offs of WMT'14 En-De and En-Fr on three hardware platforms. HAT consistently outperforms the baseline Transformers and achieves up to  $3 \times$  faster inference speed and  $3.7 \times$  smaller size over Transformer-Big. Specific latency, BLEU and SacreBLEU (Post, 2018) are in Appendix Table 8.
Figure 8: Inference latency and BLEU trade-offs of WMT'19 and IWSLT'14 tasks on Nvidia GPU.
4.1 HAT Performance Comparisons
In Figure 7, 8 and Appendix Table 8, we compare HAT with Transformer baselines on four tasks. The embedding dims are 512 and 1024 for the Transformer-Base and Big, respectively. The hidden dims are  $4 \times$  and  $2 \times$  of the embedding dim for WMT and IWSLT. The IWSLT models are smaller to prevent overfitting (Wu et al., 2019b). We obtain a series of baseline models with layer number scaling (yellow) and dimension scaling (blue). We set different latency constraints on three hardware to get a series of HAT models. HAT consistently outperforms baselines with a large gap under different latency constraints. On ARM CPU, HAT is
$3 \times$  faster and  $3.7 \times$  smaller than Transformer-Big with the same BLEU. On Intel CPU, HAT achieves over  $2 \times$  speedup. On Nvidia GPU, the blue dash line is nearly vertical, indicating that dimension scaling can hardly reduce the latency. In this case, HAT can still find models with low latency and high performance.
We further compare various aspects of HAT with Transformer (Vaswani et al., 2017) and Evolved Transformer (So et al., 2019) in Table 2. HAT achieves up to  $1.6 \times, 3 \times$ , and  $3.4 \times$  speedup with up to  $1.4 \times, 3.7 \times$ , and  $4 \times$  smaller size than baselines. We report FLOPs for translating a 23-token sentence for IWSLT and 30 for WMT. We show the overall GPU hours for training the SuperTransformer and the searched SubTransformer. We also calculate the cloud computing costs with different modes: "preemptable" is cheaper ( $0.74/h)) than "on-demand" ($ 2.48/h) (Strubell et al., 2019). HAT is highly affordable since the total GPU-hour is over  $12000 \times$  smaller than the Evolved Transformer, and is even smaller than Transformer-Big by virtue of the compact model size.
In Table 3, we compare HAT with other latest models. We scale down all models to have similar BLEU scores with Levenshtein for fair comparisons. We adopt the average iteration time of 2.88 for decoding (Gu et al., 2019), without limiting the length of the output sentence (12 tokens after
7680
Hardware-Aware	Hetero.Layers	Latency	#Params	FLOPs(G)	BLEU	GPUHours	CO2e(lbs)	CloudComp. Cost
IWSLT'14De-En	Transformer	X	X	3.3s	32M	1.5	34.5	2	5	$12 - $40
HAT (Ours)	✓	✓	2.1s	23M	1.1	34.5	4	9	$24 - $80
WMT'14En-Fr	Transformer	X	X	23.2s	176M	10.6	41.2	240	68	$178 - $595
Evolved Trans.	X	X	20.9s	175M	10.8	41.3	2,192,000	626,000	$1.6M - $5.5M
HAT (Ours)	✓	✓	7.8s	48M	3.4	41.4	216	61	$159 - $534
HAT (Ours)	✓	✓	9.1s	57M	3.9	41.8	224	64	$166 - $555
WMT'14En-De	Transformer	X	X	20.5s	176M	10.6	28.4	184	52	$136 - $456
Evolved Trans.	X	X	7.6s	47M	2.9	28.2	2,192,000	626,000	$1.6M - $5.5M
HAT (Ours)	✓	✓	6.0s	44M	2.7	28.2	184	52	$136 - $456
HAT (Ours)	✓	✓	6.9s	48M	3.0	28.4	200	57	$147 - $495
Table 2: Comparisons of latency, model size, FLOPs, BLEU and training cost in terms of  $\mathrm{CO}_{2}$  emissions (lbs) and cloud computing cost (USD) for Transformer, the Evolved Transformer and HAT. The training cost estimation is adapted from Strubell et al. (2019). The training time is for one Nvidia V100 GPU, and the latency is measured on the Raspberry Pi ARM CPU. The cloud computing cost is based on AWS.
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
Figure 9: Evolutionary search can find better SubTransformers in the SuperTransformer than random search.
decoding). HAT runs  $1.3 \times$  faster than Transformer with higher BLEU;  $1.9 \times$  faster than Levenshtein with 0.7 higher BLEU. Under similar latency, HAT also outperforms Lite Transformer. These results demonstrate HAT's effectiveness in lower latency scenarios. Our framework can also be adopted to speedup those models.
Design Insights. For all HAT WMT models in Figure 7,  $10\%$  of all decoder layers attend to
Table 3: Raspberry Pi ARM CPU latency and BLEU comparisons with different models on WMT'14 En-De. HAT has the lowest latency with the highest BLEU.
SubTransformer	Latency	#Params	BLEU
WMT'14	Largest	10.1s	71M	28.1
En-De	Searched HAT	6.9s	48M	28.4
WMT'14	Largest	10.1s	71M	41.4
En-Fr	Searched HAT	9.1s	57M	41.8
Table 4: The searched HAT compared with the largest SubTransformer in the design space. Larger models do not necessarily have better performance. HAT models have lower latency, smaller size, and higher BLEU.
three encoder layers,  $40\%$  attend to two encoder layers. That demonstrates the necessity of arbitrary encoder-decoder attentions.