MMNMT : Modularizing Multilingual Neural Machine Translation with Flexibly Assembled M o E and Dense Blocks (2023)

Abstract: Mixture-of-Experts (MoE) based sparse architectures can significantly increase model capacity with sublinear computational overhead, which are hence widely used in massively multilingual neural machine translation (MNMT). However, they are prone to overfitting on low-resource language translation. In this paper, we propose a modularized MNMT framework that is able to flexibly assemble dense and MoE-based sparse modules to achieve the best of both worlds. The training strategy of the modularized MNMT framework consists of three stages: (1) Pre-training basic MNMT models with different training objectives or model structures, (2) Initializing modules of the framework with pre-trained couterparts (e.g., encoder, decoder and embedding layers) from the basic models and (3) Fine-tuning the modularized MNMT framework to fit modules from different models together. We pre-train three basic MNMT models from scratch: a dense model, an MoE-based sparse model and a new MoE model, termed as MoE-LGR that explores multiple Language-Group-specifc Routers to incorporate language group knowledge into MNMT. The strengths of these pre-trained models are either on low-resource language translation, high-resource language translation or zero-shot translation. Our modularized MNMT framework attempts to incorporate these advantages into a single model with reasonable initialization and fine-tuning. Experiments on widely-used benchmark datasets demonstrate that the proposed modularized MNMT framwork substantially outperforms both MoE and dense models on high- and low-resource language translation as well as zero-shot translation. Our framework facilitates the combination of different methods with their own strengths and recycling off-the-shelf models for multilingual neural machine translation. Codes are available at https://github.com/lishangjie1/MMNMT.

## introduction multilingual neural machine translation translates multiple languages within a single model via multitask learning, facilitating the deployment of machine translation service in practice and improving low/zero-resource language translation @xcite @xcite @xcite . however, as the number of languages translated by an mnmt model increase, the capacity of the model has to be increased accordingly, otherwise translation quality will degrade for all language pairs, especially for high-resource languages. this phenomenon is referred to as capacity bottleneck @xcite . moe models provide an effective way to increase model capacity while the computation cost is sublinear to the number of parameters. due to this advantage of moe models over dense models, sparse architectures built on moe models are widely explored for massively multilingual neural machine translation that requires large model capacity. however, moe models suffer from overfitting on low-resource languages @xcite , which is not observed in mnmt built on dense models. this inspires us to ask a question: can we achieve the best of both worlds of dense and moebased sparse architectures for multilingual nmt? to answer this question, we propose mmnmt that modularizes multilingual nmt with flexibly assembled dense and moe blocks. specifically, the training strategy of mmnmt consists of three stages: ularized mnmt model with blocks from the pre-trained basic multilingual nmt models, e.g., using the pre-trained dense encoder to initialize the encoder of mmnmt. such initialization can be done in a flexible module assembling way. using dense modules to initialize moe models in our training strategy is in line with our preliminary experiments and recent studies on moe-based language models. these studies find that training moe-based language models from an off-the-shelf dense model (e.g., t5 @xcite ) is more efficient than training from scratch @xcite . our preliminary experiments also demonstrate that dense mnmt models are superior to moe-based sparse models on lowresource language translation. in order to diversify the options of pre-trained basic models and improve zero-shot translation in mmnmt, we further propose moe-lgr that incorporates a language group router into moe models. routing mechanism plays an important role in token dispatch and resource allocation for moe models. previous works in this line mainly focus on load balancing across experts in moe @xcite @xcite @xcite to prevent experts from being not specialized or overly specialized. however, the language information of each token is not fully explored in routers, which could enhance the generation ability of the decoder and mitigate the off-target problem in zero-shot translation. to address this issue, we introduce moe-lgr. particularly, we categorize languages into multiple groups according to linguistic typology and language embedding clustering, and learn a router per language group to enhance the difference of routing for different language groups. the main contributions of this work are summarized as follows: multilingual neural machine translation multilingual neural machine translation has been gaining increasing interest in recent years @xcite @xcite @xcite . however, mnmt models tend to be inferior to bilingual nmt counterparts on translating high-resource languages due to the capacity bottleneck @xcite . a wide variety of approaches have been proposed to alleviate this issue, e.g., deepening models to increase model capacity @xcite , exploring lightweight language-specific modules @xcite @xcite , word alignments (lin et al., 2020), contrastive learning @xcite and mixture-of-experts @xcite @xcite . mixture-of-experts based mnmt models usually suffer from overfitting problem @xcite . alleviating the overfitting problem is a key focus of our proposed modularized mnmt framework. mixture-of-experts sparsely-gated mixture-of-experts (moe) @xcite selects the top-k experts in moe layers through a routing mechanism. a strand of research focuses on load balancing, such as base layer @xcite , hash layer @xcite , thor @xcite , expert choice routing @xcite . in this work, we propose a language group routing mechanism to provide language information and enhance the difference of routing across language groups, which significantly improves the performance of zero-shot translation. recently, ## preliminary experiments and findings to have a deep understanding on the strengths and weaknesses of dense and moe models in multilingual nmt, we conducted preliminary experiments to compare them. details of these experiments can be found in section 5.1. results are shown in ## methodology to mitigate the overfitting issue and achieve the best of both worlds, we propose a general framework mmnmt to modularize multilingual nmt so as to assemble desirable modules from both moe and dense models. the training strategy of mmnmt consists of basic model pre-training, module initializing and fine-tuning. ## experiments we used the moe branch of fairseq foot_1 to implement our moe models. we conducted extensive experiments to examine the effectiveness of the proposed mmnmt against moe and dense model baselines in multilingual machine translation. ## conclusion in this paper, we have presented mmnmt, a modularized multilingual neural machine translation framework that is capable of flexibly assembling both dense and sparse blocks to achieve the best of both worlds for high/low-resource language translation and zero-shot translation. experiments and in-depth analyses demonstrate that our framework combined with different modules significantly outperforms the moe and dense baselines on highand low-resource language translation as well as zero-shot translation.