Multilingual Neural Machine Translation with Knowledge Distillation

1 INTRODUCTION . Neural Machine Translation ( NMT ) has witnessed rapid development in recent years ( Bahdanau et al. , 2015 ; Luong et al. , 2015b ; Wu et al. , 2016 ; Gehring et al. , 2017 ; Vaswani et al. , 2017 ; Wu et al. , 2018 ; Song et al. , 2018 ; Shen et al. , 2018 ; Guo et al. , 2018 ; He et al. , 2018 ; Gong et al. , 2018 ) , including advanced model structures ( Gehring et al. , 2017 ; Vaswani et al. , 2017 ) and human parity achievements ( Hassan et al. , 2018 ) . While conventional NMT can well handle single pair translation , training a separate model for each language pair is resource consuming , considering there are thousands of languages in the world1 . Therefore , multilingual NMT ( Johnson et al. , 2017 ; Firat et al. , 2016 ; Ha et al. , 2016 ; Lu et al. , 2018 ) is developed which handles multiple language pairs in one model , greatly reducing the offline training and online serving cost . Previous works on multilingual NMT mainly focus on model architecture design through parameter sharing , e.g. , sharing encoder , decoder or attention module ( Firat et al. , 2016 ; Lu et al. , 2018 ) or sharing the entire models ( Johnson et al. , 2017 ; Ha et al. , 2016 ) . They achieve comparable accuracy with individual models ( each language pair with a separate model ) when the languages are similar to each other and the number of language pairs is small ( e.g. , two or three ) . However , when handling more language pairs ( dozens or even hundreds ) , the translation accuracy of multilingual model is usually inferior to individual models , due to language diversity . It is challenging to train a multilingual translation model supporting dozens of language pairs while achieving comparable accuracy as individual models . Observing that individual models are usually ∗Authors contribute equally to this work . 1https : //www.ethnologue.com/browse of higher accuracy than the multilingual model in conventional model training , we propose to transfer the knowledge from individual models to the multilingual model with knowledge distillation , which has been studied for model compression and knowledge transfer and well matches our setting of multilingual translation . It usually starts by training a big/deep teacher model ( or ensemble of multiple models ) , and then train a small/shallow student model to mimic the behaviors of the teacher model , such as its hidden representation ( Yim et al. , 2017 ; Romero et al. , 2014 ) , its output probabilities ( Hinton et al. , 2015 ; Freitag et al. , 2017 ) or directly training on the sentences generated by the teacher model in neural machine translation ( Kim & Rush , 2016a ) . The student model can ( nearly ) match the accuracy of the cumbersome teacher model ( or the ensemble of multiple models ) with knowledge distillation . In this paper , we propose a new method based on knowledge distillation for multilingual translation to eliminate the accuracy gap between the multilingual model and individual models . In our method , multiple individual models serve as teachers , each handling a separate language pair , while the student handles all the language pairs in a single model , which is different from the conventional knowledge distillation where the teacher and student models usually handle the same task . We first train the individual models for each translation pair and then we train the multilingual model by matching with the outputs of all the individual models and the ground-truth translation simultaneously . After some iterations of training , the multilingual model may get higher translation accuracy than the individual models on some language pairs . Then we remove the distillation loss and keep training the multilingual model on these languages pairs with the original log-likelihood loss of the ground-truth translation . We conduct experiments on three translation datasets : IWSLT with 12 language pairs , WMT with 6 language pairs and Ted talk with 44 language pairs . Our proposed method boosts the translation accuracy of the baseline multilingual model and achieve similar ( or even better ) accuracy as individual models for most language pairs . Specifically , the multilingual model with only 1/44 parameters can match or surpass the accuracy of individual models on the Ted talk datasets . 2 BACKGROUND . 2.1 NEURAL MACHINE TRANSLATION . Given a set of bilingual sentence pairsD = { ( x , y ) ∈ X×Y } , an NMT model learns the parameter θ by minimizing the negative log-likelihood − ∑ ( x , y ) ∈D logP ( y|x ; θ ) . P ( y|x ; θ ) is calculated based on the chain rule ∏Ty t=1 P ( yt|y < t , x ; θ ) , where y < t represents the tokens preceding position t , and Ty is the length of sentence y . The encoder-decoder framework ( Bahdanau et al. , 2015 ; Luong et al. , 2015b ; Sutskever et al. , 2014 ; Wu et al. , 2016 ; Gehring et al. , 2017 ; Vaswani et al. , 2017 ) is usually adopted to model the conditional probability P ( y|x ; θ ) , where the encoder maps the input to a set of hidden representations h and the decoder generates each target token yt using the previous generated tokens y < t as well as the representations h . 2.2 MULTILINGUAL NMT . NMT has been extended from the translation of a single language pair to multilingual translation ( Dong et al. , 2015 ; Luong et al. , 2015a ; Firat et al. , 2016 ; Lu et al. , 2018 ; Johnson et al. , 2017 ; Ha et al. , 2016 ) , considering the large amount of languages pairs in the world . Some of these works focus on how to share the components of the NMT model among multiple language pairs . Dong et al . ( 2015 ) use a shared encoder but different decoders to translate the same source language to multiple target languages . Luong et al . ( 2015a ) use the combination of multiple encoders and decoders , with one encoder for each source language and one decoder for each target language respectively , to translate multiple source languages to multiple target languages . Firat et al . ( 2016 ) share the attention mechanism but use different encoders and decoders for multilingual translation . Similarly , Lu et al . ( 2018 ) design the neural interlingua , which is an attentional LSTM encoder to bridge multiple encoders and decoders for different language pairs . In Johnson et al . ( 2017 ) and Ha et al . ( 2016 ) , multiple source and target languages are handled with a universal model ( one encoder and decoder ) , with a special tag in the encoder to determine which target language to translate . In Gu et al . ( 2018a ; b ) and Neubig & Hu ( 2018 ) , multilingual translation is leveraged to boost the accuracy of low-resource language pairs with better model structure or training mechanism . It is observed that when there are dozens of language pairs , multilingual NMT usually achieves inferior accuracy compared with its counterpart which trains an individual model for each language pair . In this work we propose the multilingual distillation framework to boost the accuracy of multilingual NMT , so as to match or even surpass the accuracy of individual models . 2.3 KNOWLEDGE DISTILLATION . The early adoption of knowledge distillation is for model compression ( Bucilu et al. , 2006 ) , where the goal is to deliver a compact student model that matches the accuracy of a large teacher model or the ensemble of multiple models . Knowledge distillation has soon been applied to a variety of tasks , including image classification ( Hinton et al. , 2015 ; Furlanello et al. , 2018 ; Yang et al. , 2018 ; Anil et al. , 2018 ; Li et al. , 2017 ) , speech recognition ( Hinton et al. , 2015 ) and natural language processing ( Kim & Rush , 2016a ; Freitag et al. , 2017 ) . Recent works ( Furlanello et al. , 2018 ; Yang et al. , 2018 ) even demonstrate that student model can surpass the accuracy of the teacher model , even if the teacher model is of the same capacity as the student model . Zhang et al . ( 2017 ) propose the mutual learning to enable multiple student models to learn collaboratively and teach each other by knowledge distillation , which can improve the accuracy of those individual models . Anil et al . ( 2018 ) propose online distillation to improve the scalability of distributed model training and the training accuracy . In this paper , we develop the multilingual distillation framework for multilingual NMT . Our work differs from Zhang et al . ( 2017 ) and Anil et al . ( 2018 ) in that they collaboratively train multiple student models with codistillation , while we use multiple teacher models to train a single student model , the multilingual NMT model . 3 METHOD . As mentioned , when there are many language pairs and each pair has enough training data , the accuracy of individual models for those language pairs is usually higher than that of the multilingual model , given that the multilingual model has limited capacity comparing with the sum of all the individual models . Therefore , we propose to teach the multilingual model using the individual models as teachers . Here we first describe the idea of knowledge distillation in neural machine translation for the case of one teacher and one student , and then introduce our method in the multilingual setting with multiple teachers ( the individual models ) and one student ( the multilingual model ) . 3.1 ONE TEACHER AND ONE STUDENT . Denote D = { ( x , y ) ∈ X × Y } as the bilingual corpus of a language pair . The log-likelihood loss ( cross-entropy with one-hot label ) on corpus D with regard to an NMT model θ can be formulated as follows : LNLL ( D ; θ ) =− ∑ ( x , y ) ∈D logP ( y|x ; θ ) , logP ( y|x ; θ ) = Ty∑ t=1 |V |∑ k=1 1 { yt = k } logP ( yt = k|y < t , x ; θ ) , ( 1 ) where Ty is the length of the target sentence , |V | is the vocabulary size of the target language , yt is the t-th target token , 1 { · } is the indicator function that represents the one-hot label , and P ( ·|· ) is the conditional probability with model θ . In knowledge distillation , the student ( with model parameter θ ) not only matches the outputs of the ground-truth one-hot label , but also to the probability outputs of the teacher model ( with parameter θT ) . Denote the output distribution of the teacher model for token yt as Q ( yt|y < t , x ; θT ) . The cross entropy between two distributions serves as the distillation loss : LKD ( D ; θ , θT ) = − ∑ ( x , y ) ∈D Ty∑ t=1 |V |∑ k=1 Q { yt = k|y < t , x ; θT } logP ( yt = k|y < t , x ; θ ) . ( 2 ) The difference between LNLL ( D ; θ ) and LKD ( D ; θ , θT ) is that the target distribution of LKD ( D ; θ , θT ) is no longer the original one-hot label , but teacher ’ s output distribution which is more smooth by assigning non-zero probabilities to more than one word and yields smaller variance in gradients ( Hinton et al. , 2015 ) . Then the total loss function becomes LALL ( D ; θ , θT ) = ( 1− λ ) LNLL ( D ; θ ) + λLKD ( D ; θ , θT ) , ( 3 ) where λ is the coefficient to trade off the two loss terms .

Summary: Train a multilingual NMT system using the technique of Johnson et al (2017), but augment the standard cross-entropy loss with a distillation component based on individual (single-language-pair) teacher models. Periodically compare the validation BLEU score of the multilingual model with that of each individual model, and turn off distillation for language pairs where the multilingual model is better. On three different corpora (IWSLT, WMT, TED) with into-English translation from numbers of source languages ranging from 6 (WMT) to 44 (TED), this technique outperforms standard distillation for every language pair, and outperforms the individual models for most language pairs. Supplementary experiments justify the strategy of selectively turning off distillation, and quantify the effect using only the top 8 vocabulary items in distillation.