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: The paper proposes an algorithm, DTS-Gate, to warm-up experts in an MoE learning setup. Rather than applying standard top-k selection per input, the algorithm initially applies dense routing (i.e. k = total number of experts), to avoid quick collapse, and gradually relaxes it to become sparser (and cheaper). Experiments and ablations in language tasks are provided to support the advantages of the algorithm.