The N iu T rans System for WNGT 2020 Efficiency Task (2020)

Abstract: This paper describes the submissions of the NiuTrans Team to the WNGT 2020 Efficiency Shared Task. We focus on the efficient implementation of deep Transformer models (Wang et al., 2019; Li et al., 2019) using NiuTensor, a flexible toolkit for NLP tasks. We explored the combination of deep encoder and shallow decoder in Transformer models via model compression and knowledge distillation. The neural machine translation decoding also benefits from FP16 inference, attention caching, dynamic batching, and batch pruning. Our systems achieve promising results in both translation quality and efficiency, e.g., our fastest system can translate more than 40,000 tokens per second with an RTX 2080 Ti while maintaining 42.9 BLEU on newstest2018.

## introduction in recent years, the transformer model and its variants @xcite @xcite @xcite have established state-of-the-art results on machine translation (mt) tasks. however, achieving high performance requires an enormous amount of computations @xcite , limiting the deployment of these models on devices with constrained hardware resources. the efficiency task aims at developing mt systems to achieve not only translation accuracy but also memory efficiency or translation speed across different devices. this competition constraints systems to translate 1 million english sentences within 2 hours. our goal is to improve the quality of translations while maintaining enough speed. we participated in both cpus and gpus tracks in the shared task. our system was built with niutensor, an opensource tensor toolkit written in c++ and cuda based on dynamic computational graphs. niutensor is developed for facilitating nlp research and industrial deployment. the system is lightweight, high-quality, production-ready, and incorporated with the latest research ideas. we investigated with a different number of encoder/decoder layers to make trade-offs between translation performance and speed. we first trained several strong teacher models and then compressed teachers to compact student models via knowledge distillation @xcite . we find that using a deep encoder (up to 35 layers) and a shallow decoder (1 layer) gives reasonable improvements in speed while maintaining high translation quality. we also optimized the transformer model decoding in engineering, such as caching the decoder's attention results and using low precision data type. we present teacher models and training details in section 2, then in section 3 we describe how to obtain lightweight student models for efficient decoding. optimizations for the decoding across different devices are discussed in section 4. we show the details of our submissions and the results in section 5. section 6 summarizes this paper and describes future work. ## lightweight student models after the training of deep transformer teachers, we compressed the knowledge in an ensemble into a single model through knowledge distillation @xcite . then we analyzed the decoding time of each part in the deep transformer. we further pruned the encoder and decoder layers to improve the decoding efficiency. ## submissions and results we submitted five systems to this shared task, one for the cpu track and four for the gpu track, summarized as for the gpu tracks, our systems were measured on a tesla t4 gpu. gpu versions were compiled with cuda 10.1, and the executable file is about 96 mib. our models differ in encoder and decoder layers. the base model (35-6) has 35 encoder layers and six decoder layers and achieves 44.6 bleu on the newstest2018. then we see a speedup of more than one-third and a slight decrease of only 0.2 bleu by reducing the decoder layer to 1 (35-1). we continue to reduce the number of encoder layers for more accelerations. the 18-1 system reduces the translation time by one-third with only half of the encoder layers compared to the 35-1 model. our fastest system consists of 9 encoder layers and one decoder layer, which has one-third parameters of the 35-6 model, achieves 40 bleu on the wmt 2019 test set, and speeds up the baseline by 3x. for the cpu track, we used the entire machine, which has 96 virtual cores. our cpu version is compiled with mkl static library, and the executable file is 22 mib. we used a tiny model for the cpu with 256 hidden dimensions and kept other hyper-parameters as the 9-1 model in the gpu version. interestingly, using half of the hidden size significantly reduces the translation quality. the main reason is that the parameters of large models cannot be reused when using smaller dimensions. this also proves that reducing the number of encoder and decoder layers is a more effective compression method. the cpu system achieves 37.2 bleu on the newstest2018 and is 1.2x faster than the fastest gpu system. we made fewer efforts to reduce the model size and memory footprint. our systems use a global memory pool, and we sort the input sentences in descending order of length. thus the memory consumption will reach a peak in the early stage of decoding and then decrease. our base model contains 152 million parameters, and the file size is 291 mib when stored in 16-bit floats. the docker image size ranges from 724 mib to 930 mib for our gpu systems, while the cpu version is 452 mib. all systems running in docker are slightly slow down, and we plan to improve this in subsequent versions. ## conclusion to maximize the decoding efficiency while ensuring sufficiently high translation quality, we explored different techniques, including knowledge distillation, model compression, and decoding algorithms. the deep encoder and shallow decoder networks achieve impressive performance in both translation quality and speed. we speed up the decoding by 3x with lightweight models and efficient implementations. for the gpu system, we plan to optimize the fp16 inference by reducing the type conversion and applying kernel fusion @xcite for transformer models. for the cpu system, we will further speed up the inference by restricting the output vocabulary to a subset of likely candidates given the source @xcite and using low precision data type @xcite @xcite .