Training Flexible Depth Model by Multi-Task Learning for Neural Machine Translation

Findings of the Association for Computational Linguistics: EMNLP 2020, pages 4307–4312 November 16 - 20, 2020. c⃝2020 Association for Computational Linguistics 4307 Training Flexible Depth Model by Multi-Task Learning for Neural Machine Translation Qiang Wang1∗, Tong Xiao2,3, Jingbo Zhu2,3 1Machine Intelligence Technology Lab, Alibaba DAMO Academy 2Northeastern University, Shenyang, China 3NiuTrans Co., Ltd., Shenyang, China zhiniao.wq@alibaba-inc.com {xiaotong,zhujingbo}@mail.neu.edu.com Abstract The standard neural machine translation model can only decode with the same depth configuration as training. Restricted by this feature, we have to deploy models of various sizes to maintain the same translation latency, because the hardware conditions on different terminal devices (e.g., mobile phones) may vary greatly. Such individual training leads to increased model maintenance costs and slower model iterations, especially for the industry. In this work, we propose to use multi-task learn- ing to train a flexible depth model that can adapt to different depth configurations during inference. Experimental results show that our approach can simultaneously support decod- ing in 24 depth configurations and is superior to the individual training and another flexible depth model training method——LayerDrop. 1 Introduction As neural machine translation models become heav- ier and heavier (Vaswani et al., 2017), we have to re- sort to model compress techniques (e.g., knowledge distillation (Hinton et al., 2015; Kim and Rush, 2016)) to deploy smaller models in devices with limited resources, such as mobile phones. How- ever, a practical challenge is that the hardware conditions of different devices vary greatly. To ensure the same calculation latency, customizing distinct model sizes (e.g., depth, width) for dif- ferent devices is necessary, which leads to huge model training and maintenance costs (Yu et al., 2019). For example, we need to distill the pre- trained large model into N individual small models. The situation becomes worse for the industry when considering more translation directions and more frequent model iterations. ∗Work done during Ph.D. study at Northeastern Univer- sity. An ideal solution is to train a single model that can run in different model sizes. Such attempts have been explored in SlimNet (Yu et al., 2019) and LayerDrop (Fan et al., 2020). SlimNet allows running in four width configurations by joint train- ing of these width networks, while LayerDrop can decode with any depth configuration by applying Dropout (Srivastava et al., 2014) on layers during training. In this work, we take a further step along the line of flexible depth network like LayerDrop. As shown in Figure 1, we first demonstrate that when there is a large gap between the predefined layer dropout during training and the actual pruning ratio during inference, LayerDrop’s performance is poor. To solve this problem, we propose to use multi- task learning to train a flexible depth model by treating each supported depth configuration as a task. We reduce the supported depth space for the aggressive model compression rate and propose an effective deterministic sub-network assignment method to eliminate the mismatch between training and inference in LayerDrop. Experimental results on deep Transformer (Wang et al., 2019) show that our …

Content truncated to 500 words