Taming Sparsely Activated Transformer with Stochastic Experts

1 INTRODUCTION . Large neural network models have shown to be effective in many natural language processing tasks such as machine translation ( Lewis et al. , 2020 ; Conneau & Lample , 2019 ) , natural language understanding ( Devlin et al. , 2019 ; Liu et al. , 2019 ; He et al. , 2020 ) , and natural language generation ( Radford et al. , 2019 ; Brown et al. , 2020 ) . These models are usually densely activated . That is , a model uses all its parameters to process all inputs . One drawback of these models is the prohibitive training cost . Moreover , the extreme size drastically reduces inference speed , further limiting the models ’ practicality . To address these issues , sparsely activated models ( SAMs , Shazeer et al . 2017 ) have been proposed . A SAM adaptively selects a subset of its parameters for different inputs during model training and inference . This makes it possible to train SAMs that are an order of magnitude larger than densely activated models without significant increase in computational cost . For example , the sparsely activated GShard ( Lepikhin et al. , 2020 ) consists of over 600 billion parameters and the Switch Transformer ( Fedus et al. , 2021 ) 1.5 trillion parameters , while GPT-3 ( Brown et al. , 2020 ) , which is arguably the largest densely activated model , consists of only 175 billion parameters . The building block of SAMs is the expert layer , which contains an attention mechanism and multiple feed-forward neural networks ( FFNs ) in parallel . Each FFN is referred to as an expert . During ∗Work was done during an internship at Microsoft . training , an input is routed to a fixed number of experts , such that the number of floating point operations ( FLOPs ) of one forward pass remains constant , regardless of the total number of experts . Thus , training SAMs is much more cost-efficient than training densely activated models . For example , training of Switch-large ( Fedus et al. , 2021 ) and that of T5-large ( Raffel et al. , 2019 ) require the same forward FLOPs , despite that the former is 35 times larger ( 26.3 vs. 0.74 billion parameters ) . However , SAMs have been reported to be parameter inefficient . For example , although the Switchlarge model is 35 times larger than T5-large , its performance on the GLUE benchmark ( Wang et al. , 2019a ) is only slightly better ( 88.5 vs. 87.8 ) . There are also cases where the performance of SAMs is even worse than smaller densely activated models . For example , the performance of Switchlarge is worse than T5-large on the ARC Reasoning Challenge ( 66.0 vs. 68.8 ) ( Clark et al. , 2018 ) . In another example , although GShard ( Lepikhin et al. , 2020 ) shows substantial gains over densely activated models , a diminishing return with larger number of parameters has been observed . Most on-going research has focused on improving SAMs by developing effective routing methods . Since only a subset of model parameters ( i.e. , experts ) are updated for each input during training , we need to decide which experts to be activated given an input . Existing works ( Shazeer et al. , 2017 ; Lepikhin et al. , 2020 ; Fedus et al. , 2021 ; Yang et al. , 2021 ) use a gating network for input routing . However , the gating mechanism suffers from the notorious load imbalance issue : the gate ’ s weight could collapse such that nearly all the inputs are routed to the same expert . Therefore , many methods are proposed to mitigate this issue , such as noisy gating ( Shazeer et al. , 2017 ) , expert capacity ( Lepikhin et al. , 2020 ) , load balancing loss ( Lepikhin et al. , 2020 ; Fedus et al. , 2021 ) , and k Top-1 gating ( Yang et al. , 2021 ) . However , these routing methods have not been proved effective to make SAMs more parameter efficient . To understand why SAMs are not parameter efficient , we analyze the performance of several classic MoE models . Our analysis reveals that a SAM does not always outperform a densely activated model of a similar size , confirming the results reported in Yang et al . ( 2021 ) . Moreover , we also observe that the widely-used routing method based on the gating mechanism does not work better than randomly routing inputs to experts , Inspired by our findings , we propose a new SAM , THOR ( Transformer witH StOchastic ExpeRts ) . Unlike classic SAMs , such as the Switch Transformer , experts in THOR are randomly activated ( with no need of any gating mechanism ) for each input during training and inference . THOR models are trained by minimizing both the cross-entropy loss and a consistency regularization term , such that experts can learn not only from training data but also from other experts as teachers so that all the experts make consistent predictions . To validate the effectiveness of THOR , we have conducted extensive experiments on machine translation using three settings : low-resource , rich-resource , and multilingual . Results show that THOR models outperform state-of-the-art MoE models by an average of 2 BLEU score on twelve low-resource translation tasks . In the rich-resource setting , THOR achieves new state-of-the-art results on the two widely-used translation benchmarks , WMT ’ 16 En-De and WMT ’ 14 En-Fr . On multilingual translation tasks , the THOR model with 300 million parameters achieves 2 BLEU score improvement over a state-of-the-art MoE model of the same size . Moreover , our model achieves state-of-the-art results on these tasks — the same BLEU score that is achieved by the Z-code MoE model ( Kim et al. , 2021 ) with 5.5 billion parameters ( 18 times larger ) . 2 BACKGROUND . Transformer . The Transformer ( Vaswani et al. , 2017 ) model has demonstrated its superior performance in many sequence-to-sequence natural language processing tasks , such as neural machine translation . The model contains an encoder and a decoder . The encoder consists of multiple encoder layers , each having an identical structure . An encoder layer employs a self-attention mechanism and a feed-forward neural network ( FFN ) . The decoder is similarly constructed , except for an additional cross-attention mechanism in each decoder layer . Sparsely Activated Models . The building block of SAMs is the expert layer , which is similar to the Transformer layer . Each of these expert layers contain an attention mechanism and multiple FFNs in parallel , where each FFN is referred to as an expert . Let { Ei } Ni=1 denote the experts , and N denotes the total number of experts . A gating mechanism decides to which expert ( s ) an input should be routed . At each expert layer , given an input vector x ∈ Rd , where d is the embedding dimension , the gate value of routing x to expert Ei is pi ( x ) = [ Softmax ( Wgx ) ] i , ( 1 ) where Wg ∈ RN×d is the trainable weight matrix of the gating mechanism . Given the gate values { pi ( x ) } Ni=1 , we select the top-K experts to form an activated set of experts T ⊂ { 1 · · ·N } , where |T | = K. Then the output xout of the expert layer is xout = ∑ i∈T pi ( x ) Ei ( x ) . ( 2 ) Notice that in Eq . 2 , input x only activates K instead of N experts , where K N , e.g. , K = 2 and N = 2048 in GShard ( Lepikhin et al. , 2020 ) . This implies that the number of FLOPs required for one forward pass does not increase with the number of experts N . Therefore , SAMs can scale to an enormous size without any significant increase in training time and inference time . The gate weight matrixWg ( Eq . 1 ) is trained together with the rest of the model parameters . Because there is no constraint on the learned weights , it is possible that Wg collapses such that one row dominates , i.e. , all the inputs are routed to one expert . This problem is referred to as load imbalance . Existing works adopt various ad-hoc heuristics to mitigate this issue , e.g. , adding Gaussian noise to Eq . 1 ( noisy gating , Shazeer et al . 2017 ) , limiting the maximum number of inputs that can be routed to an expert ( expert capacity , Lepikhin et al . 2020 ) , imposing a load balancing loss ( Lepikhin et al. , 2020 ; Fedus et al. , 2021 ) , and using linear assignment ( Lewis et al. , 2021 ) . There are other works that remove the gating mechanism such that load imbalance is no longer an issue , e.g. , by incorporating hash functions ( Roller et al. , 2021 ) . Besides the load imbalance issue , there are also heated discussions on how to construct T in Eq . 2 . For example , Shazeer et al . ( 2017 ) ; Lepikhin et al . ( 2020 ) ; Yang et al . ( 2021 ) conjecture that routing inputs to K > 1 experts is necessary , while Fedus et al . ( 2021 ) argue that using K = 1 is sufficient and more computationally efficient . 3 ANALYSIS OF SPARSELY ACTIVATED MODELS . We investigate behavior of the gating mechanism of several classic MoE models . We conduct experiments on a multilingual translation task , { De , Vi } → En . More details are presented in Appendix A . We consider two MoE models proposed in Shen et al . ( 2019 ) , referred to as MoE ( dec ) and MoE ( tok ) , respectively , and three variants of the Switch Transformer proposed in Fedus et al . ( 2021 ) . The number of experts is set to two for all the MoE models . We compare them with the Transformer ( Vaswani et al. , 2017 ) model of the same model size . Figure 1 shows the validation losses and BLEU scores of three models : Transformer , MoE ( dec ) , and MoE ( tok ) . We see that the two MoE models perform very similarly , and neither outperforms the Transformer by a significant margin . To interpret the results of Figure 1 , we examine the load of each expert and the confidence scores of routing inputs to different experts . An expert ’ s load is defined as the proportion of inputs that are assigned to it . For an input that is routed to an expert , its routing confidence score ( output of the gating mechanism ) determines the level of preference , e.g. , if the routing confidence score is 0.5 , then the gate has no preference for either expert . For each expert , we compute the average routing confidence score over all the inputs assigned to it . Figure 2 shows that after the early stage of training ( i.e. , the first 200 iterations ) , the gate weight collapses and nearly all the inputs are routed to expert 2 . Also , the average routing confidence score of expert 2 is close to 1.0 , which means that the gate strongly prefers expert 2 to expert 1 . In this case , only one of the experts is sufficiently trained . Figure 3 depicts a different scenario , where the inputs are randomly dispatched to the experts . Notice that after approximately 4000 iterations , the two experts are equally loaded , and the probabilities of assigning any input to expert 1 and expert 2 are almost identical , indicating that the gating mechanism has no preference for either expert . We have identified two behaviors of the gating mechanism : load imbalance and random routing . The former is also reported in recent papers ( Shazeer et al. , 2017 ; Lepikhin et al. , 2020 ; Fedus et al. , 2021 ) . We further investigate the Switch Transformer ( Fedus et al. , 2021 ) , which is a state-of-theart MoE variant that incorporates various methods to resolve the load imbalance issue . In addition , because behavior of the gating mechanism in the Switch Transformer mimics random routing ( see Appendix A ) , we examine the effect of discarding the gate and randomly assigning inputs to experts . Figure 4 demonstrates the validation losses and BLEU scores of the Transformer and three variants of the Switch Transformer , where inputs are routed according to tokens ( referred to as Switch ( t ) ) , sentences ( Switch ( s ) ) , or are routed randomly ( Switch ( r ) ) . Similar to the results in Figure 1 , we see that the four models perform similarly . This shows that even after we alleviate load imbalance , model performance is not improved ( i.e. , the Switch Transformers do not outperform the vanilla Transformer ) , and the performance of the Switch Transformer does not vary much among different routing methods , including random routing . We remark that in this paper , we focus on natural language processing tasks , in particular neural machine translation . There are other works in different research fields ( e.g. , computer vision ) that draw different conclusions than ours ( Riquelme et al. , 2021 ) . We attribute this to the intrinsic differences between image classification and language generation , e.g. , each input in the former belongs to a clearly-defined category , while no such knowledge exists in the latter . In summary , the experiments reveal • A sparsely activated model does not always outperform a densely activated model of the same model size . • The widely-used routing method based on the gating mechanism does not work better than randomly routing inputs to experts .

Summary: The paper proposes a new routing mechanism for sparse models in the context of language tasks. Rather than learning a parametric router that learns how to assign tokens to experts, the proposed algorithm (THOR), randomly selects two experts per mini batch, and applies those experts independently to every input. A consistency loss is used to force experts to provide similar predictions. A number of experiments are provided suggesting the algorithm outperforms previous works.