Distilling Multiple Domains for Neural Machine Translation (2020)

Abstract: Neural machine translation achieves impressive results in high-resource conditions, but performance often suffers when the input domain is low-resource. The standard practice of adapting a separate model for each domain of interest does not scale well in practice from both a quality perspective (brittleness under domain shift) as well as a cost perspective (added maintenance and inference complexity). In this paper, we propose a framework for training a single multi-domain neural machine translation model that is able to translate several domains without increasing inference time or memory usage. We show that this model can improve translation on both high- and low-resource domains over strong multi-domain baselines. In addition, our proposed model is effective when domain labels are unknown during training, as well as robust under noisy data conditions.

## introduction neural machine translation (nmt) @xcite @xcite can achieve high quality when trained using deep architectures on large amounts of relevant data @xcite . however, training data for generic translation models is typically not balanced or diverse with respect to domain. as a result, translation quality can be inconsistent across domains, with lowerquality outputs for low-resource domains such as chat compared to high-resource domains such as news @xcite . one way to address this is by building domainadapted models @xcite ) that target a specific domain. in this case, in-domain data is used to specialize the machine translation (mt) model for the target domain, for example by fine-tuning a generic model on this data @xcite @xcite . this yields "expert" models that are better than models trained on the in-domain data alone and more specialized than a generic translation system. however, for an mt application that needs to translate multiple domains, this approach would require maintaining and running separate expert systems for each domain, which becomes prohibitively expensive as the number of domains increases. additionally, in a real-world scenario, the domain of the input text might be unknown at inference time, adding the complexity of detecting which system should be used for a given input. one alternative to expert models are multidomain mt systems @xcite @xcite @xcite . specifically, the goal of a multidomain method is to obtain a single nmt model that approaches the performance obtained through multiple expert models. under this framework, access to several domain-specific corpora is assumed at training time, but the domain information is not known at inference time. in this paper, we address the problem of multidomain mt. our goal is to exploit knowledge about the heterogeneous nature of the data to train a single fixed-capacity model that approaches the quality of the experts across all domains. achieving high quality on each domain without damaging quality on other domains and without increasing the model complexity is an ambitious goal that matches the setup of many user-facing mt systems @xcite . the main contributions of the paper are: ## multi-domain distillation for mt we assume our training data d tr is composed of n disjoint labeled domains d i (@xmath0, . . . , n}) containing parallel sentences (s, t): the goal is to build a fixed-capacity model that performs well across all n domains. in this work, we assume that all domains are equally important and we measure performance as the unweighted average across all domain-specific test sets d tst i . however, it is desirable for a multi-domain model to not trade improved performance on one domain with degradation on another. for this reason, we also evaluate performance drop across all domains w.r.t. a baseline model trained on d tr . our approach builds on several observations made in previous work. @xcite showed that knowledge distillation using an increased capacity teacher model is an effective method for reducing the complexity of training data. although the exact mechanisms are still not well understood @xcite @xcite , smallersized student models trained on the output of the teacher perform better than the same models trained on the original data. their performance is on par with that of the very large teacher model, which is impractical to use in practice. in multi-domain mt, increased depth alone does not generally provide the best performance, and increasing capacity through specialization of networks, either as completely separate neural models or stacked models, is the optimal strategy in practice @xcite . we exploit this intuition and generalize the sequence-level knowledge distillation approach of @xcite to the multi-domain case by distilling the output of multiple domain-specific teachers. this technique is referred to as multi-domain knowledge distillation (mdkd; figure 1 ), and it consists of three steps: ## mdkd using unsupervised clusters the mdkd framework works well when the domain-labeled training corpus contains domains that are well-defined, since high-quality deep domain experts can be trained. however, as noted in the previous section, domain labels may not always correspond to consistent, separable domains. in addition, in many cases, domain labels might not be available at all. in this section, we investigate whether the mdkd technique can be used to improve the multidomain performance of a single model without knowing domain labels at test or training time. instead of relying on corpus-level domain labels, we cluster the heterogeneous training data at the sentence level and treat the clusters obtained as regular domains. clustering the training data to cluster the training data, we first compute sentence embeddings of all the source training sentences using the multilingual variant of bert (mbert; @xcite , which has 768 dimensions and is trained on wikipedia data from 104 languages. we then apply k-means clustering @xcite to compute the clusters on the inferred sentence embeddings. different numbers of clusters are computed: for @xmath1, we have four domains, bleu avg parl law med it gold labels 49.8 39.3 59.5 57.1 43.2 3 clusters 49.2 39.4 58.6 55.8 42.8 4 clusters 49.1 39.0 58.7 55.8 42.9 5 clusters 49.3 39.1 58.9 56.1 42.9 similarly, for @xmath3, we compute models with 4, 5, and 6 clusters. we leave finding the optimal number of clusters as future work; in preliminary experiments, the method we employed to automatically compute this number led to a prohibitively large amount of clusters (more in appendix b.3). cluster classifications for each domain in the test set are shown as heat maps in figure 2 . interestingly, the unsupervised clusters do not overlap strongly with the gold domain labels, even for the @xmath4. ## ablations 6.1 improvements due to distillation our proposed multi-domain knowledge distillation models train deep in-domain teachers and distill these teachers into the shallower students. in this section, we aim to understand how much of the gains seen from the mdkd models can be attributed to the knowledge distillation component. to this end, we train a knowledge distillation baseline model that builds a single deep teacher for the entire data. the training data is distilled using this teacher model and a student is trained on the concatenation of the distilled data and the original data, similarly to the mdkd models. we train unbalanced and upsampled teachers, from which we distill unbalanced and upsampled students, respectively. ## related work knowledge distillation was first introduced for classification tasks as a way to compress large networks or ensembles of networks into smaller models that achieve similar performance @xcite . @xcite extended this to neural machine translation, and since then many researchers have proposed further applications of sequence-level knowledge distillation for nmt, for example for non-autoregressive translation models @xcite . most prior approaches to multi-domain neural machine translation (see @xcite for a survey) require knowledge of the input domain at test time. @xcite used word-level and sentence-level domain tags on the source sentence. similarly, @xcite performed multi-domain nmt by breaking the word embeddings into generic and domain-specific components. @xcite trained speakerspecific nmt models by treating each speaker as a domain and adapting the softmax bias term for each domain. these models work well when the domain is known at training and inference time, but requiring labeled data at inference time is a major limitation in a real-world setting. @xcite introduced the setup that we follow in this paper, where domains are known at training but not at inference. in addition to the target token approach evaluated in this paper, they trained a second model that adds a domain classifier on top of the nmt encoder; this achieved similar bleu scores but is less parameter-efficient than their target token model. @xcite considered a case where no domain labels are used at all; instead, a generic model is adapted on the fly using similar training sentences to the input. our multi-domain knowledge distillation technique is architecture-agnostic and thus complementary to these approaches, since it can be combined with any multi-domain nmt model. knowledge distillation has been applied to nmt domain adaptation by @xcite , who used a domain-specific teacher and a generic teacher to improve domain-adapted expert models. most similar to our mdkd approach is the application of knowledge distillation to multilingual nmt by @xcite , who trained single-language teacher models and then distilled these models to a multilingual student model. knowledge distillation has also been previously applied to domain-aware nmt by @xcite . however, unlike our work, they did not train domain-specific teachers; instead, they used source factors like domain and casing information to inform a single teacher model. concurrently to our work, mghabbar and ratnamogan (2020) also proposed multi-domain knowledge distillation using domain-specific teachers and a domain-agnostic student. their method differs from ours in several aspects, including our use of sequence-level knowledge distillation, domain-specific distillation data, and single-best distillation outputs for each training sentence. our work and mghabbar and ratnamogan (2020) are complementary, as both propose different effective approaches for leveraging knowledge distillation to train multi-domain nmt. ## conclusions we have introduced multi-domain knowledge distillation, a new method for multi-domain nmt that distills multiple expert models into a single student that shows high quality across all domains. we have kept both model architecture and capacity fixed and shown that mdkd leads to significantly better multi-domain models without any increase in translation time or memory usage. since the approach is architecture-independent, it is easy to combine with other multi-domain nmt models. in this paper, we have fixed the depth and the architecture of the teachers; however, improving the teachers using different architectures may also lead to better empirical results. our experiments have covered two data quality conditions: when the domains are well-defined and separable, individually trained deep domain experts outperform all the multi-domain baselines and mdkd bridges a large portion of the gap between these baselines and the deep experts. a second set of experiments has revealed a dataset for which the domains were not clearly separable and some domains exhibited train/test mismatches. in this setting, training domain-specific expert models is not a robust strategy, as the expert models performed significantly worse that the baselines. despite using distillation from these experts, mdkd was very robust to this noisy setting: not only was there no quality degradation, but we even observed modest improvements over the baselines. finally, we performed experiments in which we assumed that the domain labels are unknown and are obtained through clustering of the train data. the resulting mdkd models outperformed the baselines again, showing that gold domain labels are not strictly needed. for future work, we plan to expand the automatic domain induction methods and test the mdkd framework on generic mt with data exhibiting varying degrees of heterogeneity: as mdkd distills domain-specific models to create multiple simpler data distributions, we want to investigate if inducing train-time specializations and using them for distillation through mdkd can lead to better quality.