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

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 learning to train a flexible depth model that can adapt to different depth configurations during inference. Experimental results show that our approach can simultaneously support decoding in 24 depth configurations and is superior to the individual training and another flexible depth model training method——LayerDrop.

## introduction as neural machine translation models become heavier and heavier @xcite , we have to resort to model compress techniques (e.g., knowledge distillation @xcite ) to deploy smaller models in devices with limited resources, such as mobile phones. however, 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 different devices is necessary, which leads to huge model training and maintenance costs @xcite . for example, we need to distill the pretrained large model into n individual small models. the situation becomes worse for the industry when considering more translation directions and more frequent model iterations. an ideal solution is to train a single model that can run in different model sizes. such attempts have been explored in slimnet @xcite and layerdrop @xcite . slimnet allows running in four width configurations by joint training of these width networks, while layerdrop can decode with any depth configuration by applying dropout @xcite 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 multitask 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 @xcite show that our approach can simultaneously support decoding in 24 depth configurations and is superior to the individual training and layerdrop. ## flexible depth by multi-task learning we propose to use multi-task learning to solve the above problems. all tasks are trained jointly and share the same parameters. concretely, unlike lay-erdrop, which allows up to m × n possible depth configurations, our approach sets a smaller depth configuration space (m i , n i ) k @xmath0, n i ) as a task. another major difference from layerdrop is that each task's sub-network is unique and deterministic in our method, resulting in consistent sub-network used between training and inference. algorithm 1: training flexible depth model by multi-task learning. reduce depth space. for depth d, in principle, layerdrop can be pruned to any depth of φ(d) = {0, 1, 2, . . . , d}. however, consider the actual situation of model compression for resourcelimited devices, it is unnecessary if the compressing rate is too low, e.g., @xmath1. therefore, for an aggressive compress rate, we replace the entire space φ(d) with the set of all positive divisors of d 2 : the physical meaning of φ(d) is to compress every d/d layers into one layer, where @xmath2. guideline for deterministic sub-network assignment. the use of deterministic sub-networks is critical to maintaining the consistency between training and inference. however, for each @xmath3, it is not trivial to decide which d layers should be selected to construct the d-layer subnetwork. here we propose two metrics to guide the procedure. the first is task balance (tb), whose motivation is to make every layer have as uniform tasks as possible. we use the standard deviation of the number of tasks per layer to measure it quantitatively: proposed method. guided by these two metrics, we design an effective sub-network assignment method optimal. we record the usage state s i of each layer to ensure not to put too many tasks on the same layer. at initialization, we set s i as alive. for @xmath4, optimal prioritizes to process large depth. optimal uniformly assigns one layer for every @xmath5. in each chunk, we pick the middle layer of ceil(c/2) -1 (called middleleft). note that, layerdrop uses the leftmost layer in each chunk (called left), as shown in eq. 3. although left and middleleft have the same ald, we found that there is a large gap in tb. for example, when @xmath6, left's tb is 1.5, which is much higher than middleleft's 0.78 (lower is better). then, optimal records which layers are used and picks the less used layers as much as possible. each used layer is marked as dead. if current alive layers cannot accommodate the picked depth d, we pass it and choose a smaller d until the alive layers are sufficient, or reset all layers as alive. algorithm 1 describes the training process of our method. during training, compared with individual training and layerdrop from scratch, our fdm finetunes on the individually pretrained m m -n and uses sequence-level knowledge distillation (seq-kd) @xcite to help shallower networks training. we note that in conventional seq-kd, the student model cannot finetune on the teacher model directly because the ## conclusion we demonstrated layerdrop is not suitable for fdm training because of (1) the huge sub-network space in training and (2) the mismatch between training and inference. then we proposed to use multi-task learning to mitigate it. experimental results show that our approach can decode with up to 24 depth configurations and obtain comparable or better performance than individual training and layerdrop. in the future, we plan to explore more effective fdm training methods, and combining flexible depth and width is also one of the attractive directions.