Dense-to-Sparse Gate for Mixture-of-Experts

1 INTRODUCTION . Transformer models have achieved a huge success in improving model quality via scaling data and model sizes ( Brown et al. , 2020 ; Dosovitskiy et al. , 2021 ; Ying et al. , 2021 ) . Specifically , Kaplan et al . ( 2020 ) explored the scaling law of transformer models which shows that the model quality scales as a power-law with data sizes , model sizes and the computation . However , with the rapid increasing of the model sizes , it is hard to further scale the model to extremely large size due to the limited computation power of available hardware devices . To address these challenges , sparselygated Mixture-of-Experts ( MoE ) , a popular form of conditional computation , is proposed to further increase the model size while without increasing the computational cost ( e.g. , FLOPS ) proportionally ( Bengio et al. , 2013 ; Shazeer et al. , 2017 ; Fedus et al. , 2021 ; Lepikhin et al. , 2021 ; Lewis et al. , 2021b ; Roller et al. , 2021 ) . Specifically , tokens are routed by a sparse gate to corresponding experts that only contain part of the full model for processing , leading to lower computational costs compared to a dense model with the same model size . The routing strategy defined by the gate is key to MoE models , as it not only determines the model sparse structure ( i.e. , which experts tokens are routed to ) , but also affects the specialization of experts ( i.e. , how each expert is trained ) ( Roller et al. , 2021 ) . State-of-the-art approaches ( Fedus et al. , 2021 ; Lepikhin et al. , 2021 ) usually adopt a pre-defined sparse gate ( e.g. , Top-1 or Top-2 ) , and then initialize and train both the gate functions and experts jointly . However , such a join-training over pre-defined gates could significantly limit both the training efficiency and model quality . Particularly , at the beginning of training a MoE model , both the gate and the experts are randomly initialized . In this case , the gate does not have evidence to decide which expert to process a token , and the experts also do not have experiences to process a randomly-assigned input token . Joint-updating their models requires a vast amount of trial and errors to converge . Furthermore , the pre-defined gates significantly limit the opportunity of exploring all experts to only 1 or 2 at a time , and have to rely on some random routing noises to get feedback from other experts and finally reinforce their future routing choices . Our observation shows that such random routing in the initial stage and longdistance reinforce-based model updating in existing approaches could affect both the training time and final model quality . 1https : //anonymous.4open.science/r/MoE-3D0D/README.moe.md In this work , to overcome the limitations in existing approaches , we advocate a new training mechanism to decide the sparse gate gradually for MoE models , which is named as Dense-to-Sparse gate ( DTS-gate ) . Instead of pre-defining the sparse gate in the previous works , DTS-gate starts with a dense gate that routes tokens to all the experts so that each expert gets trained sufficiently at the beginning . Then , after a short warming-up training for all experts , the gate adaptively learns the weights of routing to each expert and gradually routes tokens to fewer experts , making the training structure sparser and continuously reducing the computation cost , while keeping the model quality improving as usual . In particular , to implement the DTS-gate , our idea is to carefully control the temperature of a softmax-based routing function , so that to adjust the weights distribution among experts and control the sparsity of the MoE layer during training . In short , DTS-gate advances in two aspects . First , compared to the joint-training of gate and experts from scratch , DTS-gate provides a mechanism of training gate after experts , which could reduce a lot of random error-trails at the beginning . Second , compared to the reinforce-based model updating , starting with a dense gate allows us to get training feedbacks from all experts and adjust the routing weights directly to the right direction , which not only speedups the convergence of the gate , but also benefits for the experts specialization . Experiments show that compared with the state-of-the-art Switch-Gate in GPT-MoE ( 1.5B ) model with OpenWebText dataset ( 40GB ) , DTS-Gate can obtain 2.0x speed-up to reach the same validation perplexity ( Figure 3 ( a ) ) , as well as higher FLOPs-efficiency of a 1.42x speed-up ( Figure 3 ( b ) ) . Experiments also verify the ability of DST-Gate for scaling models with more experts or more MoE layers . 2 RELATED WORK . Static Sparse Neural Networks Exploiting the sparsity in deep neural networks can reduce both storage and computation requirements for model training and inference . One of the most widely used sparsification methods is weight pruning ( LeCun et al. , 1990 ; Han et al. , 2015 ) . Previous studies propose to prune away redundant or less useful weights based on various pruning criteria ( e.g. , the importance of individual weights ( Han et al. , 2015 ) or groups of weights ( Wen et al. , 2016 ; Luo et al. , 2017 ; He et al. , 2017 ) ) and then fine-tune remained weights to regain the lost accuracy . After pruning and fine-tuning , parts of weights are permanently removed , inducing a static sparsity pattern in DNNs . The sparsity pattern/structure is a trade-off between model effectiveness and hardware efficiency ( Mao et al. , 2017 ) . Early works attempt to increase the sparsity ratio or model accuracy by employing unstructured sparsification methods , while recent works focus more on structured sparsity for practical speedup on hardware . Interestingly , Frankle & Carbin ( 2018 ) points out training a sparse network from scratch is superior or comparable to pruning-based methods . While Liu et al . ( 2018 ) draw a different conclusion under different settings ( e.g. , training hyper-parameters , pruning structures ) . Therefore , it still needs further investigation to answer which one is critical to sparse neural networks , the sparsity structure or remained weights . Our DTS-gate is analogous to pruning-based methods that training all experts first and then learning the sparse gate routing . Conditional Computation with MoE Different from static sparse neural networks that permanently remove some weights , conditional computation ( Bengio et al. , 2013 ) activates only some relevant parts of the model on a per-example basis , which can be regarded as a dynamic sparsity structure that remains all weights but brings sparsity into the computation . The MoE model ( Shazeer et al. , 2017 ) , as a specific form of conditional computation , contains many experts and a trainable gating network which selects a sparse combination of experts to process each input sample . Conditional computation is capable of reducing inference cost ( without reducing model capacity ) or increasing model capacity ( without increasing inference cost ) from a model acceleration or scaling perspective . On the other hand , the activated parts ( experts in MoE ) can be regarded as structured sparse blocks , which does not introduce additional computational overhead . However , conditional computation models are often difficult to train , since they require learning discrete routing decisions from individual examples to experts and the gating network tends to converge to a state that only selects the same few experts ( Eigen et al. , 2013 ) . Shazeer et al . ( 2017 ) , Lepikhin et al . ( 2021 ) and Fedus et al . ( 2021 ) add auxiliary load-balancing losses to mitigate this self-reinforcement phenomenon and improve training efficiency . In such MoE models , the gating network and experts , as two critical components , are jointly trained which may interfere with each other . In DTS-gate , we consider K = 𝒏 K = 2 K = 1 decoupling the training of experts and the gating network by involving all experts starting with a high temperature in Gumbel-Softmax and then training the gating network to be sparser and select the best expert through decaying this temperature . Lewis et al . ( 2021b ) formulates token-expert allocation as a linear assignment problem and guarantees balanced compute loads . Roller et al . ( 2021 ) replaces the gating network with a hash-based routing strategy , e.g . random hash , clustered hash dispersed hash . 3 METHODS . 3.1 PRELIMINARY . The main components of a sparse MoE layer are a network G as sparse gate and an expert network E. Sparse Gate Network The sparse gate network G takes a token xs ( xs ∈ RD ) as input and produces the probability of it with respective to N experts . Equation 1 formulates the gate function , where Wg ( Wg ∈ RD×N ) is the embedding of N experts and the top-K outputs are normalized via softmax distribution . The sparse gate network is widely adopted by Shazeer et al . ( 2017 ) ; Fedus et al . ( 2021 ) ; Lepikhin et al . ( 2021 ) and achieves promising results with constant computation cost . g ( xs ) = softmax ( TopK ( xs ·Wg ) ) ( 1 ) Expert Network The expert network E is a set of N experts { e1 , ... , eN } , where each is a FFN layer and contains its own parameters in our models . For each expert ei ( ei : RD → RD ) , it takes the token xs as input to produce its own output ei ( xs ) . The final output of the expert network ys is the linearly weighted combination of each expert ’ s output on the token by the gate ’ s output , formulated as Equation 2. ys = N∑ i=1 g ( xs ) iei ( xs ) ( 2 ) 3.2 DENSE-TO-SPARSE GATE . Sparse gating has demonstrated its superior model efficiency in both training and inference , but prior work tends to convergence to a sub-optimal model on FLOPs-efficiency or sample-efficiency because of the jointly training of the randomly initialized gate network and expert network . In this paper , we propose a new mechanism for the gate function , named Dense-to-Sparse gate ( DTS-Gate ) , which starts as a dense gate that routes tokens to all experts and then gradually becomes sparser . DTS-Gate benefits from the sufficient training of experts in the early stage and then make the experts selection becomes sparser on the basis of specialized experts . This dense-to-sparse process only occupies a small fraction compared with the total training time , which usually takes days to weeks . Sparsity in Mixture of Experts Sparsity naturally exists in MoE models . Here we train a 24- layer transformer decoder model , with every other FFN layer replaced by 16-expert MoE layer . Each MoE layer is densely connected , where tokens are distributed to all experts but with weighted differently . As shown in Figure 2 , the expert weight matrix is extremely sparse and the loads are naturally imbalanced among experts . Dynamic Sparsity Gate In order to control the sparsity during training , we adopt the softmax temperature to adjust the weights distribution among experts . Formulated as Equation 3 , Wg is the experts embedding , ζ is the extra noise and sampled from Gumbel ( 0 , 1 ) distribution ( Jang et al. , 2017 ) , and τ is the softmax temperature which controls the distribution . When the τ increases , the distribution becomes more uniform , which evolves more experts into the computation of each token . As the τ approaching 0 , the distribution becomes one-hot , which is more confident for the gate network . g ( xs ) = e ( xs·Wg+ζ ) /τ∑N s′=1 e ( xs′ ·Wg+ζ ) /τ ( 3 ) Sparsity Scheduler By scheduling the temperature of Equation 3 , we can control the sparsity of the MoE layer at different training stages . There is a trade-off between small temperature , where experts distributions are nearly one-hot and lead to large variance of gradients in learning but less FLOPs in computation . Large temperature is just the opposite . DTS-Gate starts at a large τ and anneal to a small τ with linearly decaying . Adaptive Capacity With τ decreasing , the weights among experts approach one-hot distribution , where most of them are close to zero . To optimize such sparsity into actucal training time , we use threshold drop the experts whose weights fall below the threshold as Figure 1 , and no extra communication or computation will be wasted . Furthermore , we don ’ t enforce a balanced assignment algorithm for DTS-Gate because we observe naturally imbalanced distribution of experts from Dense-MoE models , shown as Figure 2 ( b ) . To meet the demand of above two designs , we enable each expert with adaptive capacity to make them well specialized .

Summary: This paper studied the problem of training the gating network for the mixture-of-experts based model architectures. To make the gating network training more stable and robust, the authors proposed a dense-to-sparse gating training algorithm that uses Gumbel noise and temperature tuning. The authors conducted experiments on large NLP datasets.