L ego- MT : Learning Detachable Models for Massively Multilingual Machine Translation (2023)

Abstract: Multilingual neural machine translation (MNMT) aims to build a unified model for many language directions. Existing monolithic models for MNMT encounter two challenges: parameter interference among languages and inefficient inference for large models. In this paper, we revisit the classic multi-way structures and develop a detachable model by assigning each language (or group of languages) to an individual branch that supports plug-and-play training and inference. To address the needs of learning representations for all languages in a unified space, we propose a novel efficient training recipe, upon which we build an effective detachable model, Lego-MT.For a fair comparison, we collect data from OPUS and build a translation benchmark covering 433 languages and 1.3B parallel data. Experiments show that Lego-MT with 1.2B parameters brings an average gain of 3.2 spBLEU. It even outperforms M2M-100 with 12B parameters. The proposed training recipe brings a 28.2 \times speedup over the conventional multi-way training method.code and data repo: https://github.com/CONE-MT/Lego-MT.git .

## introduction multilingual neural machine translation (mnmt) translates languages by mapping a source sentence to a unified representation space and decoding a target sentence from this space @xcite @xcite @xcite . traditional mnmt models use a shared network to align representations in different languages. recently, scaling up the size of mnmt models brings significant quantitative improvements and new qualitative capabilities (m2m-100, @xcite @xcite . beyond mnmt, recent large-scale language models (e.g., chat-gpt) also show promising results on zero-shot (or 1 https://github.com/cone-mt/lego-mt . (2) lego-mt is a multi-way structure that includes both multilingual (denoted as m) and language-specific encoders and decoders for english (denoted as e), chinese (denoted as c) and nepali (denotes as n). the architecture is detachable at inference time where only a specific encoder and decoder are needed. u (unified space) represents hidden representations generated by encoders. to find out the best recipe for massive multilingual translation, we revisit the classic multi-way (or multi-branch) architecture @xcite , whose philosophy is to allocate an individual encoder and decoder for each language (or group of languages), as shown in figure 1 . the immediate benefit of this structure is: despite appealing, there remain two big challenges when training multi-way structures: representation alignment between different languages due to the lack of shared parameters; and low gpu efficiency during training because unused parameters occupy gpu memory but do not have any computations. furthermore, the feature of random language mixture in a batch makes it infeasible to use an online-loading method (i.e., loading during usage) to accelerate training since it will cause impractical @xcite . to address these challenges, we propose a novel training recipe, which results in our new detachable model, lego-mt. we classify the training data into different language-centric groups such that we only need to load specific branches into gpu memory, eliminating the need to load different modules constantly. the language-centric group is trained in sequential order. second, during each languagecentric training, we introduce a multilingual branch and propose a new triple-flow method to help a model learn to map to and translate from a unified space. specifically, a unified space is a type of representation space rather than a module. it creates a common representation of language that can be used across multiple language tasks. to evaluate our training recipe for massive mnmt, we construct a many-to-many translation dataset foot_0 covering 7 language-centric groups, 433 languages, based on the open-source website opus foot_1 @xcite . lego-mt-1.2b yields average gains of 3.2 spbleu, and even outperforms m2m-100-12b which has 10× inference parameters. furthermore, the proposed training recipe brings a 28.2× speedup compared with the conventional multiway training method. we also conduct comprehensive experiments on branch combinations, thanks to the detachable nature of the model. we find that low-resource languages prefer multilingual branches and high-resource languages prefer language-specific branches. in addition, we also observe that the unseen combination of a highresource language encoder and a high-resource language decoder can achieve better performance, showing that lego-mt can align different branches into a unified space effectively. the main contributions can be summarized as follows: ## related work in this part, we review recent related multilingual machine translation models. we classify them into three categories: fully / group-shared @xcite , and mixture-of-expert (moe). the fully-shared model is the most prevalent model in multilingual neural machine translation (mnmt). this model employs a single architecture to translate in all directions @xcite @xcite @xcite @xcite and has demonstrated efficacy in aiding low-resource directions. however, fullyshared models are often subject to capacity bottlenecks and trade-offs between translation quality and the number of languages @xcite @xcite . group-shared models incorporate individual parameters for each group and represent a popular solution for sharing language-specific encoders or decoders @xcite . @xcite ; @xcite ; @xcite ; @xcite ; @xcite proposed an mnmt model only for shared languagespecific modules. lass @xcite learns language-specific sub-networks for each language direction for multilingual translation. adpater methods @xcite during training, we introduce an efficient training method by classifying multilingual data into language-centric groups. the language-centric groups are trained in a sequential way. during each training phase, only language-specific parameters are loaded into gpu memory. the training maintains three flows: enc-flow (language-specific encoder + multilingual decoder) for training specific encoder, dec-flow (multilingual encoder + language-specific decoder) to train language-specific decoder, and mix-flow (multilingual encoder + multilingual decoder) to avoid the overfitting of multilingual encoder and decoder to each language-centric training data. u means the unified space, hidden representations generated by encoders. mixture-of-expert (moe) has recently emerged as a prominent research direction @xcite @xcite @xcite @xcite , which are sparsely activated, with each inference only activating a subset of parameters. researchers have applied moe to massively multilingual translation and introduced various regularization strategies to enhance performance @xcite . despite promising results, moe's objective differs from ours, as it still requires the entire structure to be stored in gpu memory during inference. the encoder-decoder structure has demonstrated considerable flexibility through the utilization of the lego-nn @xcite . the lego-nn can be applied to various tasks with decoder modules being detachable, in contrast, the lego-mt model design allows for the performance of massively mnmt with all modules being detachable. ## experiments while lego-mt is generic, we focus the experiments on m2m-100-1.2b as backbone models since m2m-100 is a leading mt model. ## analysis on lego-mt ablation studies on triple-flow training we design three flows in lego-mt: mix-flow, enc-flow, and dec-flow. mix-flow contains a multilingual encoder and a multilingual decoder, which is essential in regularizing language-specific training. we start from m-flow and see how enc-flow and dec-flow affect the final performance, which gives more insights into the design of our framework. for simplification, we use chinese-centric data in top-10 shards and select 7 @xmath0, which includes high-resource languages (be, de, fa, jv) and low-resource languages (ne, pa, sw). we train lego-mt on the selected set and observe results in it is mainly because that language-specific decoder may cause a large distribution shift on multilingual encoders, resulting in catastrophic forgetting. that's why we split the training into two stages and keeps dec-flow in the second stage. analysis on inference path section due to the plug-and-play features, there are several possible inference paths for a single translation direction. at the inference stage, there are three alternative solutions for language-centric translation: mix-flow, enc-flow, and dec-flow. figure 2 shows the comparison between these inference paths. for lowresource languages (eg., ceb, az, ne), mix-flow (m-encoder + m-decoder) works better than either enc-flow (e-encoder + m-decoder) or dec-flow (m-encoder + d-decoder). high-resource languages (eg., en,de,zh, ar) prefer languagespecific branches. dec-flow (a multilingual encoder and a language-specific decoder) achieves better performance among these paths. this demonstrates that specific parameters are more important when the amount of data in a language is huge. in summary, the mix-flow (m-encoder + m-decoder) is recommended for inference tasks with low-resource languages, and the dec-flow (m-encoder + d-decoder) is more appropriate for high-resource languages. ## conclusion with the increasing scale of languages, using a single model to translate all directions brings new challenges in practice. this paper proposes an efficient training recipe, which results in a detachable multilingual translation model, lego-mt. to validate the effectiveness of our algorithm, we develop a massive mnmt translation dataset, covering 433 languages. results on flores-101 show that lego-mt-1.2b achieves large performance improvements over strong baselines under a fair comparison. it even outperforms the result of m2m-12b with a gain of 4 bleu on many-to-one.