The Mixture-of-Experts (MoE) technique can scale up the model size of Transformers with an affordable computational overhead. We point out that existing learning-to-route MoE methods suffer from the routing fluctuation issue, i.e., the target expert of the same input may change along with training, but only one expert will be activated for the input during inference. The routing fluctuation tends to harm sample efficiency because the same input updates different experts but only one is finally used. In this paper, we propose STABLEMOE with two training stages to address the routing fluctuation problem. In the first training stage, we learn a balanced and cohesive routing strategy and distill it into a lightweight router decoupled from the backbone model. In the second training stage, we utilize the distilled router to determine the token-to-expert assignment and freeze it for a stable routing strategy. We validate our method on language modeling and multilingual machine translation. The results show that STABLEMOE outperforms existing MoE methods in terms of both convergence speed and performance. The code is available at https://github. com/Hunter-DDM/stablemoe . In recent years, large-scale Transformers (Devlin et al., 2019; Dong et al., 2019; Raffel et al., 2020; Clark et al., 2020; Bao et al., 2020; Brown et al., 2020) have shown a striking ability to model languages. However, with the model scale growing, the training speed will go slower, and the extremely large memory requirement also introduces a heavy burden of engineering. Mixture of Experts (MoE) (Jacobs et al., 1991; Jordan and Jacobs, 1994; Shazeer et al., 2017), in a much easier way, enables Transformers to scale up the number of parameters meanwhile introducing an affordable  Contribution during internship at Microsoft Research. computational overhead. MoE-based Transformers have a set of expert modules, and only a few experts will be activated for each input token. In this way, we can expand the model scale by adding expert modules, which will keep the computational and memory overhead within a tolerable range. Most existing MoE methods (Lepikhin et al., 2021; Fedus et al., 2021; Lewis et al., 2021) decide the token-to-expert routing according to the dynamically changing token representations. However, we point out that they face the routing fluctuation problem. As shown in Figure 1, the same input may be assigned to different experts along with training. However, during inference, only one expert will be activated for the input. The routing fluctuation problem tends to harm sample efficiency because the same input updates different experts while only one is finally used. Taking BASE Layer (Lewis et al., 2021) as an example, during the whole training process, we examine the token-to-expert assignment for tokens in the validation set. For an input token, we define the last fluctuation step as the last step where its target expert is different from the final step. We plot the cumulative token percentage with regard to the last fluctuation step (annotated as its percentage accounting for all training steps) in Figure 2. We find that the last fluctuation step of 40.9% tokens exceeds 20%, which means 40.9% tokens do not have a stable target expert when 20% of all training steps have been done. Furthermore, 29.1% tokens still change their target experts after half of the whole training process, and 15.4% tokens even change the target expert after 80% of all training steps, which is nearing the training ending. These statistics prove that the routing fluctuation problem indeed exists in previous MoE methods. In this paper, we propose STABLEMOE with two training stages to address the routing fluctuation problem. In the first training stage, we follow the learning-to-route paradigm and aim to learn a 7085 Training Epochs Expert 3 Inference  1  2  3   Transformer Blocks Transformer Blocks The car is Expert 2 Expert 1 Expert 3 Transformer Blocks Transformer Blocks Expert 2 Expert 1 Expert 3 Transformer Blocks Transformer Blocks Expert 2 Expert 1 Transformer Blocks Transformer Blocks Expert 2 Expert 1 Expert 3 The car is The car is The car is Figure 1: Illustration of the routing fluctuation problem. balanced and cohesive routing strategy. We design a balance loss to guarantee the assignment is balanced. In addition, inspired by Lewis et al. (2021), we adopt a sigmoid gating mechanism, which enables the task objective to propagate supervised signal back to the routing strategy, to facilitate learning a more cohesive assignment. As the routing strategy is being learned, we synchronously distill it into a lightweight router decoupled from the backbone model. In the second training stage, we utilize the distilled router to determine the token-to-expert assignment. The distilled router is frozen in this stage to provide a stable routing strategy, which addresses the routing fluctuation problem in the remaining training. We conduct experiments on language modeling and multilingual machine translation. The results show that STABLEMOE outperforms existing MoE methods in terms of both convergence speed and performance. (1) We point out the routing fluctuation problem in existing learning-to-route MoE methods. (2) We propose STABLEMOE to address the routing fluctuation problem. (3) We conduct substantial experiments under various settings to show the advantages of STABLEMOE over existing MoE methods. We first introduce the MoE mechanism designed for Transformers (Vaswani et al., 2017). Given a standard L -layer Transformer model and an input sequence X containing T tokens, the Transformer output HL is calculated by HL = [ h L 1 ; h L 2 ; ... ; h LT ] , (1) h lt = FFN (cid:16) u lt (cid:17) + u lt , (2) u l 1: T = self-att (cid:16) h l  1 1: T (cid:17) + h l  1 1: T , (3) where h lt is the hidden state of t -th token after the l -th layer, Self-Att(  ) is the self-attention module, and FFN(  ) is short for the feed-forward network. For simplicity, we omit the layer normalization. We implement MoE for Transformers by inserting MoE layers, that are composed of a set of FFNs, into two neighboring Transformer blocks. At an MoE layer, for each input token, only a few or one expert will be activated, controlled by a gating function g (  ) : h lt = N (cid:88) i =1 g i (cid:16) h l  1 t (cid:17) FFN i (cid:16) h l  1 t (cid:17) + h l  1 t , (4) where N is the total number of experts, and FFN i is the i -th expert. Here, the gating function g i (  ) is sparse for computational efficiency. For simplicity, we omit the layer normalization. 3. In the first training stage, we follow the learning-to-route paradigm and aim to learn a balanced and cohesive routing strategy. As the routing strategy is being learned, we synchronously distill it into a lightweight router decoupled from the backbone model. In the second training stage, we utilize the distilled router to determine the token-to-expert assignment. The distilled router is frozen in this stage to provide a stable routing strategy. During inference, we also use the frozen distilled router for consistent routing. 3.1 Training Stage 1: Learn Routing Strategy Let h l  1 t  R d be the input representation of token t and E  RN  d be the centroids of N experts. For each MoE layer, we assign each token to one expert FFN (Fedus et al., 2021; Lewis et al., 2021; Roller et al., 2021). The assignment score is: s t,i = E (cid:62) i h l  1 t , (5) where s t,i is the assignment score between token t and expert i , indicating their affinity. We use a greedy assignment algorithm, i.e., sending each token to the expert with the highest affinity. Then, we calculate the expert FFN output as: a t = arg max i ( s t,i ) , (6) h lt =  ( s t,a t ) FFN a t (cid:16) h l  1 t (cid:17) + h l  1 t , (7) where a t is the expert index that token t is sent to, and  is the sigmoid gate (Lewis et al., 2021). Considering the sigmoid gate  ( s t,a t ) , if FFN a t is beneficial for token t , optimizing the training objective (e.g., minimizing the cross-entropy loss for language modeling) will urge the gate to be greater; otherwise, the gate will tend to be smaller. The gate signal urges similar tokens to be assigned to the same expert that is beneficial to them, thus producing cohesive token-to-expert assignments. Balance Loss We design a balance loss L bal to avoid imbalanced assignments that will result in a high computational bottleneck in the MoE layer and thus limit the computational efficiency: L bal =  N (cid:88) i =1   ( |A i |  n ) n (cid:88) t A i  ( s t,i )   , (8) where  is a hyper-parameter, A i denotes the set of tokens assigned to expert i , and n denotes the average number of tokens per expert. Intuitively, if an expert is overloaded, the balance loss will urge its assignment scores to be smaller. Otherwise, if an expert is unoccupied, the balance loss will increase its assignment scores to capture more tokens. Distilled Router As the routing strategy is being learned, we synchronously distill it into a lightweight router decoupled from the backbone model to mimic the original routing strategy. Let X be the input sequence and  E be the distilled expert centroids, we use word embeddings D(  ) to extract the routing features. We use the cross-entropy loss as the distillation loss L dis :  h l  1 t = D( X t ) ,  s t,i =  E (cid:62) i  h l  1 t , (9) L dis =  T (cid:88) t =1 log exp ( s t,a t ) (cid:80) Ni =1 exp ( s t,i ) , (10) where  h l  1 t is the distilled routing feature of token t ,  s t,i is the distilled assignment score between token t and expert i , and a t is the expert index that token t is actually sent to. In practice, D(  ) 7087 Methods Assignment Algorithm Gating Function Balance Loss Switch Transformer Greedy softmax Yes BASE Layer Auction (Bertsekas, 1992) sigmoid No Hash Layer Fixed Hashing { 0 , 1 } No STABLEMOE Training Stage 1 Greedy sigmoid Yes Training Stage 2 Fixed Routing sigmoid No Table 1: Comparison of three core elements among STABLEMOE and existing MoE-based Transformers. can also be other feature extractors such as CNNs or Transformers (we investigate other variants of distilled routers in Section 4.4.3), but the word embedding is the fastest one and achieves the best performance. At the end of training stage 1, we freeze all parameters for the distilled router (i.e., D(  ) and  E ) to prepare a stable routing strategy for training stage 2 and the inference stage. Given frozen D(  ) and  E , in training stage 2, we directly use them for a stable routing strategy. Keeping other processes the same as in training stage 1, we calculate the output of the MoE layer as follows:  h l  1 t = D( X t ) ,  s t,i =  E (cid:62) i  h l  1 t , (12)  a t = arg max i ( s t,i ) , (13) h lt =  ( s t,  a t ) FFN  a t (cid:16) h l  1 t (cid:17) + h l  1 t . (14) Notice that the sigmoid gate  (  ) still uses original assignment score s t,  a t as input, so the gate signal can also be learned in training stage 2. Since the routing strategy has been fixed in training stage 2, we no longer need the balance loss and distillation loss. Therefore, the training loss for training stage 2 contains only the task loss: LS 2 = L task . During inference, we also use the frozen distilled router for routing. The fixed routing strategy, which is consistent with training stage 2, makes information learned in MoE layers be utilized more thoroughly and thus leads to better performance. We compare three core elements, including the assignment algorithm, the gating function, and the balance loss, among STABLEMOE and existing MoE-based Transformers. In Table 1, we summarize their differences. Assignment Algorithm Switch Transformer and the training stage 1 in STABLEMOE simply assign each token to the expert with the highest affinity. BASE Layer adopts the auction algorithm (Bertsekas, 1992) to find a global balanced assignment with the maximum affinity sum. Hash layer and the training stage 2 in STABLEMOE have token-level fixed routing strategies, which have good stability. Gating Function Hash Layer uses a hard gating function, which means an expert is either fully activated or not activated, no any intermediate state. Switch Layer, BASE Layer, and STABLEMOE have soft gating functions, which can judge the affinity between a token and its target expert and determine a proper ratio to use the expert. Soft gating mechanisms also urge models to learn a more cohesive token-to-expert assignment. Balance Loss BASE Layer and Hash Layer do not apply any balance losses. By contrast, Switch Transformer and the training stage 1 in STABLEMOE design balance losses to control the balance of the token-to-expert assignment. In summary, combing two training stages, STABLEMOE has a stable, cohesive, and balanced routing strategy, while the other three MoE methods cannot meet them all simultaneously. (Lewis et al., 2021) and Roller et al. (2021), we use the combination of the corpora in RoBERTa (Liu et al., 2019) and the English subset of the CC100 (Con-neau et al., 2020) corpus. The corpus contains about 100B tokens, and we randomly sample 5M tokens for validation and 20M tokens for test. Multilingual Machine Translation We follow Wang et al. (2020) and Ma et al. (2020) to use a collection of parallel data in different languages from the WMT datasets. 1 The dataset contains 32.5 million parallel data for language pairs between English and other 9 languages, including French (Fr), Czech (Cs), German (De), Finnish (Fi), Latvian (Lv), Estonian (Et), Romanian (Ro), Hindi (Hi), and Turkish (Tr). In our experiments, we combine the original parallel data with 180 million back-translation data as described in (Ma et al., 2020) and call the augmented dataset WMT for short. We conduct experiments based on fairseq 2 . All experiments are conducted on NVIDIA V100 GPUs with 32 GB memory. Language Modeling We adopt the tokenizer of GPT-2 (Radford et al., 2019), which uses byte-pair encoding (Sennrich et al., 2016) with a vocabulary size of 50,257. We set up two settings for STABLEMOE, a base one and a large one. For both settings, we insert one MoE layer after the middle Transformer block. We train the model for 60K steps in total (6K for training stage 1 and 54K for training stage 2). The dimension of the distilled routing features is 50, which brings 2.51M extra parameters for routing. The balance factor  is set to 0.3. We 1 http://www.statmt.org 2 https://github.com/facebookresearch/fairseq use Adam (Kingma and Ba, 2015) with  1 = 0 . 9 and  2 = 0 . 98 as the optimizer. The rest of the hyper-parameters are summarized in Appendix A. Multilingual Machine Translation Following (Ma et al., 2020), we use the Sentence-Piece (Kudo and Richardson, 2018) model to tokenize sentences. The vocabulary is learned from the training set and consists of 64,000 tokens. We insert two MoE layers, one after the third encoder block and one after the third decoder block. We train the model for 352K steps in total (30K for training stage 1 and 322K for training stage 2). The dimension of the distilled routing features is also set to 50. The balance factor  is set to 0.3. We use Adam with  1 = 0 . 9 and  2 = 0 . 98 as the optimizer. The rest of the hyper-parameters are summarized in Appendix B. 4.3 Results 4.3.1 Language Modeling We compare STABLEMOE with Switch Transformer, BASE Layer, Hash Layer, and the standard Transformer. All MoE models have the same number of shared parameters as the standard Transformer. Under the base setting, in addition, we compare two larger dense Transformers that add FFNs in a dense manner to achieve the same number of total parameters as MoE models. The deeper model stacks more FFNs, while the wider model uses FFNs with a larger hidden size. The floating point operations (FLOPs) per sequence are profiled by the torchprofile toolkit. Models # Params FLOPs De Ro Fr Cs Et Hi Tr Fi Lv Avg Standard Transformer 77M 290B 39.8 36.0 32.5 29.1 27.2 24.5 23.6 21.8 20.3 28.31 Larger Transformer 90M 317B 40.6 36.9 33.7 29.8 27.8 25.4 24.6 22.2 20.9 29.10 Switch Transformer 480M 317B 42.3 37.1 33.8 31.0 28.6 26.0 24.3 23.0 21.2 29.70 BASE Layer 480M 317B 42.6 37.8 34.2 31.0 29.0 26.9 25.1 23.2 21.6 30.16 Hash Layer 480M 317B 42.7 37.0 34.6 31.3 28.7 26.5 23.9 23.1 21.7 29.94 STABLEMOE 480M 317B 43.0 37.4 34.7 31.5 29.3 26.8 24.7 23.6 21.9 30.32 Under the base setting, STABLEMOE outperforms existing MoE methods on both the validation and the test sets by 0.3-0.8 perplexity. Compared with dense models, STABLEMOE achieves about 3.7 lower perplexity than the standard Transformer, and about 1.3 higher perplexity than the deeper larger model. Under the large setting, consistently, STABLEMOE outperforms the other MoE methods, and achieves about 2.6 lower perplexity than the standard Transformer. We also compare the convergence speed of different models under the base setting. The results are plotted in Figure 4, which takes the validation perplexity as y-axis and the training wall time as x-axis. Although larger dense models achieve better validation perplexity at last, their training speed is quite slow. With regard to the convergence speed, MoE-based Transformers usually exceed dense models. Further, among the MoE methods, STABLEMOE has the fastest convergence speed. We compare STABLEMOE with Switch Transformer, BASE Layer, Hash Layer, the standard Transformer, and a larger Transformer. All MoE-16 32 64 Number of Experts 18.5 19.0 19.5 20.0 20.5 21.0 V a li d PPLBASE Layer Hash Layer StableMoE Figure 5: Comparison of MoE-based Transformers with different numbers of experts. based models have the same number of shared parameters as the standard Transformer. Except the standard Transformer, the other models have the same FLOPs. We translate other languages to English (X  En) and report the test BLEU on WMT in Table 3. STABLEMOE achieves the best average test BLEU among the compared MoE methods. Keeping the same FLOPs, STABLEMOE outperforms the dense model by 1.22 test BLEU. With the MoE technique, we expand the number of parameters by 523% and the FLOPs just increase by 9.3%. Number of Experts Figure 5 shows the results of BASE Layer, Hash Layer, and STABLEMOE with different numbers of experts. As the number of experts goes larger, the validation perplexity of each model tends to further descend. Consistently, STABLEMOE performs the best with dif-7090 3 Sublayers (454M) 10 Sublayers (1.51B) Number of Expert Sublayers (Parameters) 18.0 18.5 19.0 19.5 20.0 V a li d PPLBASE Layer Hash Layer StableMoE Figure 6: Comparison of MoE models with different numbers of expert sublayers (i.e., number of parame-ters). ferent numbers of experts. In addition, it is worth noting that STABLEMOE with 16 experts outperforms BASE Layer with 32 experts, and STABLEMOE with 32 experts achieves a similar perplexity to BASE Layer with 64 experts. Number of Expert Parameters We compare MoE models with different numbers of expert parameters by setting different expert sublayers. Models with 3 and 10 expert sublayers have 454M and 1.51B expert parameters, respectively. From Figure 6, we observe that more expert parameters bring better performance, and STABLEMOE consistently performs the best under both settings. Position of MoE Layers We investigate the effect of the inserting position of the MoE layer. By default, the MoE layer stacks 3 MoE sublayers and is inserted after the L 2 -th Transformer block (mid-dle). We also attempt to insert the MoE layer before the first Transformer block (bottom), and after the last Transformer block (top). In addition, we also investigate the effect if we scatter 3 MoE sublayers uniformly into the standard Transformer, i.e., after the L 4 -th, 2 L 4 -th, and 3 L 4 -th blocks, respectively. As shown in Table 4, among the above four settings, Models Valid PPL BASE Layer 20.04 + Fixed Routing Strategy (Stage 2) 19.41 (0.63  ) STABLEMOE with Only Stage 1 19.48 + Fixed Routing Strategy (Stage 2) 19.28 (0.20  ) Table 5: Effects of the fixed routing strategy. inserting stacked MoE sublayers into the middle position allows STABLEMOE to achieve the best performance. Ratio Between Two Training Stages We investigate the balance point of the ratio between two training stages in STABLEMOE. Given a fixed number of total steps, allocating more steps to training stage 1 can help to learn and distill a better routing strategy. On the other hand, a larger ratio of training stage 2 means longer stable training. Under the base setting of language modeling, we attempt to allocate 6K, 15K, and 30K steps to training stage 1 and show the results in Table 6. We find that if we use word embeddings as the distilled router, allocating 6K steps (10% of the total steps) to training stage 1 is a good balance point. We speculate that the word embedding is simple enough to be learned fast, so longer stable training is more important to achieve better performance. Based on the base setting of language modeling, we design two experiments to investigate how much performance improvement the fixed routing strategy can bring. On the one hand, we equip BASE Layer with a stable routing strategy to address its routing fluctuation problem. Specifically, as in STABLEMOE, we use word embeddings to distill the routing strategy of BASE Layer in the first 6K training steps, and freeze the distilled router for stable routing in the remaining training. As shown in Table 5, the fixed routing strategy decreases the validation perplexity of BASE Layer by 0.63. On the other hand, we attempt to disable the training stage 2 in STABLEMOE and always train the model as in training stage 1. As a result, the validation perplexity of STABLEMOE becomes 0.20 higher than the full version that has a fixed routing strategy. These two cases support that the fixed routing strategy, which addresses the routing fluctuation problem, can bring better performance for MoE-based Transformers. In Table 6, in addition to word embedding, we also investigate four variants of the distilled router including CNN and three Transformers with different numbers of layers. We allocate 15K steps to training stage 1 for all of them. From the table, we find that using word embedding achieves the best performance, while the 3-layer Transformer does not perform well. For the routing strategy distillation, the distilling signal from a 32-category classifica-tion objective may not be informative enough to learn a complex router. By contrast, it is more suitable for simpler routers. Therefore, we recommend using word embedding, which is simple and effective, as the distilled router in STABLEMOE. We compare the degree of routing fluctuations between STABLEMOE and BASE Layer to show our advantage with regard to the routing stability. During the 60K training steps, we examine the token-to-expert assignment for tokens in the validation set every 500 steps. For each token, we define the last fluctuation step as the last step where its target expert is different from the final step. We plot the cumulative token percentage about the last fluctuation step in Figure 7. For ease of reading, we annotate the x-axis as the percentage it accounts for all training steps. From the figure, we find that the routing fluctuation problem is notable for BASE Layer. By contrast, for STABLEMOE, there is no routing fluctuation in training stage 2 since we apply a fixed routing strategy. Jacobs et al. (1991); Jordan and Jacobs (1994) propose Mixture of Experts (MoE) to compute different examples with independent expert modules. Shazeer et al. (2017) introduce MoE to build large-scale language models based on LSTMs (Hochre-iter and Schmidhuber, 1997). Recently, as Transformers become popular, many pieces of work design MoE-version FFNs to build MoE-based Transformers. GShard (Lepikhin et al., 2021), Switch Transformer (Fedus et al., 2021), and BASE Layer (Lewis et al., 2021) follow the learning-to-route paradigm and dynamically learn how to route each input token to experts. However, we point out that these learning-to-route methods face the routing fluctuation problem. Hash Layer (Roller et al., 2021) propose a non-parametric routing strategy, which uses a pre-designed token-level hash table to determine the token-to-expert assignment. The static routing strategy will not fluctuate, but the randomly determined hash table limits the upper bound of its performance. Our work includes the advantages of learning-to-route methods to learn a balanced and cohesive routing strategy, and further addresses the routing fluctuation problem through applying a frozen lightweight router that mimics the original routing strategy. In this paper, we point out the routing fluctuation problem that exists in previous learning-to-route MoE methods. In order to address this problem, we propose STABLEMOE with two training stages. We first learn a balanced and cohesive routing strategy and synchronously distill it into a lightweight router decoupled from the backbone model. Then, we freeze the distilled router for a stable routing strategy in the remaining training. We validate STABLEMOE on language modeling and multilingual machine translation. The results show that STA 7092 BLEMOE outperforms existing MoE methods in terms of both convergence speed and performance. William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. CoRR , abs/2101.03961. Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. 1991. Adaptive mixtures of local experts. Neural Computing , 3(1):7987. Michael I. Jordan and Robert A. Jacobs. 1994. Hierarchical mixtures of experts and the EM algorithm. Neural Computing , 6(2):181214. Diederik P. Kingma and Jimmy Ba. 2015. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015 . Damai Dai, Zhifang Sui, and Baobao Chang are supported by the National Key Research and Development Program of China 2020AAA0106701 and NSFC project U19A2065.
abstain abstain abstain objective objective method abstain abstain other abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain method method abstain result abstain abstain objective other abstain method abstain abstain method abstain method abstain abstain objective result abstain other method method other other method objective abstain method other method other abstain other method other other other other other other abstain method method other objective abstain method method other other abstain other method abstain method method abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain method method abstain method method abstain method method method method abstain abstain other abstain abstain method abstain method method abstain abstain method abstain abstain method abstain method abstain abstain other abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain result abstain abstain abstain abstain abstain abstain abstain abstain abstain result result abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain result abstain abstain abstain result abstain abstain abstain abstain result abstain other other other other abstain other other method abstain objective objective abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other