A Lightweight Mixture-of-Experts Neural Machine Translation Model with Stage-wise Training Strategy (2024)

Abstract: Dealing with language heterogeneity has always been one of the challenges in neural machine translation (NMT).The idea of using mixture-of-experts (MoE) naturally excels in addressing this issue by employing different experts to take responsibility for different problems.However, the parameter-inefficiency problem in MoE results in less performance improvement when boosting the number of parameters.Moreover, most of the MoE models are suffering from the training instability problem.This paper proposes MoA (Mixture-of-Adapters), a lightweight MoE-based NMT model that is trained via an elaborately designed stage-wise training strategy.With the standard Transformer as the backbone model, we introduce lightweight adapters as experts for easy expansion.To improve the parameter efficiency, we explicitly model and distill the language heterogeneity into the gating network with clustering.After freezing the gating network, we adopt the Gumbel-Max sampling as the routing scheme when training experts to balance the knowledge of generalization and specialization while preventing expert over-fitting.Empirical results show that MoA achieves stable improvements in different translation tasks by introducing much fewer extra parameters compared to other MoE baselines.Additionally, the performance evaluations on a multi-domain translation task illustrate the effectiveness of our training strategy.

## introduction in recent years, neural machine translation (nmt), a key component of natural language processing (nlp), has been studied extensively with significant progress @xcite . texts from various domains often exhibit unique expression styles. domain diversity leads to heterogeneous data distribution of a large multisource dataset. when training an nmt model with the global optimization strategy, data from diverse domains tend to adjust model parameters to fitting their respective distributions, which harms the convergence of the model. in literature, some works @xcite @xcite @xcite regarded this problem as domain shift and tried to address it by transfer learning. however, domain knowledge is required in these works, which introduces a new data collection problem. how to deal with the heterogeneity of language in nmt tasks remains challenging. the core concept of moe is using multiple experts to divide a problem space into homogeneous regions @xcite , which has a natural advantage in solving the problem of language heterogeneity. recently, previous works @xcite @xcite explored the mixture-of-experts (moe) structure in nmt tasks. these studies demonstrate the impressive capacity of moe in handling various data distributions. they boost the number of parameters from million to billion while maintaining low computational requirements. however, moe is reported to be parameter-inefficient @xcite @xcite i.e., a huge number of parameters only brings a small performance improvement. as an illustration, compared with a dense model, an moe model only offers an average improvement of 0.3 bleu with 20 times more parameters @xcite . meanwhile, training the gating network implicitly by an overall optimization makes most of the moe models suffer from the training instability problem. it is crucial to meticulously design a training strategy to prevent instability. for instance, expert load imbalance may occur during training of an moe model: the gating network may route most data to a small number of experts, meanwhile many other experts do not get sufficiently trained at all @xcite . moreover, the routing fluctuation @xcite issue, i.e. the gating network may route the same data to different experts along with training, is also one of the factors leading to training instability. in this paper, we propose moa (mixture-of-adapters), a lightweight moe-based nmt model that is trained using a stage-wise training strategy. our model is composed of three components: (i) an encoder-decoder based backbone model; (ii) a gating network responsible for routing data to suitable experts by their encoded features; (iii) a set of lightweight adapters @xcite as the experts transplanted in every decoder layer of the backbone model. with the stage-wise training strategy, these three components are trained sequentially. specifically, the backbone model is trained using a standard machine translation task. meanwhile, we pre-inject an adapter in every decoder layer in this training stage, and use these adapters for parameter initialization of the other adapters in the adapter training stage. in the training stage of the gating network, the language heterogeneity is modeled explicitly using clustering and distilled into the gating network through a multi-classification task. such an explicit learning strategy improves the parameter efficiency and ensures the training stability of our model. moreover, to balance the knowledge of generalization and specialization and prevent the over-fitting problem, we employ the gumbel-max sampling as the routing scheme when training the adapters. empirical results show that moa achieves stable improvement in different translation tasks by introducing much fewer extra parameters compared to the other moe baselines. additionally, the performance evaluations and the ablation studies on the multi-domain translation task illustrate the effectiveness of our training strategy. ## related works the moe structure @xcite has been widely studied in the machine translation area @xcite @xcite . with the same core concept, different moe models draw attention to different design strategies. one difference lies in what to use as experts. most of the moe models @xcite @xcite adopt feed-forward networks (ffn) as experts. based on transformer @xcite , many works @xcite @xcite inject extra moe layers or substitute the ffn layers with moe layers. instead of using ffn layers, @xcite use the attention heads as experts to achieve stronger performance than the standard multi-head attention layer. another difference is the training strategy. @xcite activate two or more experts to obtain nonzero derivatives for the gating networks in back-propagation. @xcite only activate one expert per time, they train the gating network by auxiliary losses. @xcite use a two-stage training strategy to address the routing fluctuation problem. different from the above works that use load balancing loss to prevent expert load imbalance, @xcite formulate token-to-expert allocation as a linear assignment problem that requires no auxiliary load balancing loss. @xcite propose gating dropout to reduce cross-machine communication and speed up the training process. moreover, according to the granularity of different routing schemes, moe models can be divided into three levels: token-level, sentence-level, and task-level. most of the above works adopt tokenlevel schemes, where different experts will be activated for different tokens. the sentence-level routing scheme refers to tokens from a sentence that share the same gating result. when selecting experts by task boundaries as opposed to making input-level decisions, e.g., for multilingual machine translation tasks, the routing scheme is regarded as task-level @xcite . ## model architecture moa consists of three components: (i) an encoderdecoder based backbone model; (ii) a gating network that responsibility is to route data to suitable experts by their encoded features; and (iii) a set of lightweight adapters as experts transplanted in the end of every decoder layer of the backbone model. the backbone model is based on the encoderdecoder structure, where the encoder/decoder block is composed of a stack of several identical layers. it theoretically can be any encoder-decoder based model, we use the powerful transformer @xcite in our experiment. given a source sentence @xmath0, ..., x n ), the encoder block maps it to a sequence of hidden states @xmath1, ..., h n ). then, h is fed to the decoder block to generate an output sequence @xmath2, ..., y m ) with an autoregressive process. the gating network makes use of the hidden states h to discriminate different data distributions. first, @xmath3, then two linear transformations are introduced with a tanh activation in between to compute adapter scores s, the adapters transfer the decoded hidden states from generic to specific. different from previous moe models using original feed-forward network (ffn) @xcite with large inner dimensions as experts, introducing extra lightweight adapters makes the size of experts can be more flexibly controlled. in each adapter, the output z i of the i-th decoder layer is first normalized with layer normalization, then zi is fed to an ffn with a small inner dimension, followed by a residual connection, to obtain the adapter output, in inference, only the adapter with the biggest score in each decoder layer is activated. unlike inference, the gumbel-max sampling is adopted as the routing scheme in the adapter training stage, which will be discussed in the next section. ## stage-wise training most of the moe models train their gating network along with an overall optimization of the final task. although some auxiliary losses are introduced to avoid potential risks such as expert load imbalance, this implicit learning approach introduces another discrete latent variable learning problem and increases the training difficulty of the gating networks on how to distinguish different data distributions, which leads to the parameter-inefficiency problem in moe. in this paper, we train moa with a stage-wise training strategy. each training stage is elaborately designed to improve model performance with as few extra parameters as possible. next, we will discuss our training process in detail. ## experimental settings to evaluate the effectiveness of our method, we conduct a set of experiments on both several standard machine translation tasks and a multi-domain machine translation task. the translation quality is measured by the bleu-4 @xcite score. next, we will provide a comprehensive description of our experimental settings. ## conclusion this paper proposes moa, a lightweight moebased nmt model that is trained via an elaborately designed stage-wise training strategy. the lightweight adapters are introduced as experts for easy expansion. by modeling the language heterogeneity with clustering and distilling the knowledge into the gating network explicitly, moa improves the parameter efficiency and avoids training instability. the gumbel-max sampling is adopted as the routing scheme when training the adapters to balance the knowledge of generalization and specialization and avoid over-fitting. empirical results show the effectiveness of the proposed method.