Year: 2021

Abstract:
Mixture-of-experts (MoE) is becoming popular due to its success in improving the model quality, especially in Transformers. By routing tokens with a sparse gate to a few experts that each only contains part of the full model, MoE keeps the model size unchanged and significantly reduces per-token computation, which effectively scales neural networks. However, we found that the current approach of jointly training experts and the sparse gate introduces a negative impact on model accuracy, diminishing the efficiency of expensive large-scale model training. In this work, we proposed Dense-To-Sparse gate (DTS-Gate) for MoE training. Specifically, instead of using a permanent sparse gate, DTS-Gate begins as a dense gate that routes tokens to all experts, then gradually and adaptively becomes sparser while routes to fewer experts. MoE with DTS-Gate naturally decouples the training of experts and the sparse gate by training all experts at first and then learning the sparse gate. Experiments show that compared with the state-of-the-art Switch-Gate in GPT-MoE(1.5B) model with OpenWebText dataset(40GB), DTSGate can obtain 2.0x speed-up to reach the same validation perplexity, as well as higher FLOPs-efficiency of a 1.42x speed-up. Our code is available 1.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHODS

heading: 3.1 PRELIMINARY
text: 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)

heading: 3.2 DENSE-TO-SPARSE GATE
text: 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.

heading: 4 EXPERIMENTS

heading: 4.1 EXPERIMENTAL SETUP
text: Dataset Models are trained on two datasets separately: English Wikipedia (17.5GB) (Devlin et al., 2019) and WebText (40GB) (Radford et al.). Wikipedia is a single domain of text, which contains about 2,500 words, and WebText is a large and diverse dataset in various domains, which is collected using 45 million links from Common Crawl. We apply GPT2-BPE (vocabulary size of 50257) on both datasets for tokenization and binarize the dataset for loading efficiently.
Evaluation Our experiments are conducted on language modeling task, and perplexity (PPL) is reported as the evaluation metric of model performance. We focus exclusively on comparing FLOPsefficiency, defined as the best model performance (PPL) can be achieved given the fixed number of floating-point operations. Meanwhile, we don’t add the FLOPs of balance loss throughout experiments in this paper, as we focus on gate distribution with or without balance loss. We don’t report actual training time because the actual training time mixes the computation complexity with system implementation and leads to unfair comparisons as new algorithms without well optimized system implementation. In the appendix, we will show the routing pattern of DTS is relative stable, with heavy-loaded experts almost keep fixed. And we can utilize system techniques, e.g., partitioning hot experts among multiple GPUs and merging multiple cold experts onto less GPUs, for load balance between GPUs.
Model Architectures We test on three models, including GPT-Small, GPT-Medium and GPTMoE, and the detailed architectures are shown in Table 1, where GPT-Small and GPT-Medium are two standard architectures defined by Radford et al.. We scale the feed-forward layer (FFN) of GPT-
Medium(345M) from 1 to 16 as the MoE Layer, and replace every other FFN as MoE to match the total parameters of GPT-Large(1.5B). We also extend GPT-Small(117M) with various MoE-layer settings to evaluate our DTS-Gate method.
Baselines As for MoE-based models, we compare 4 methods totally, including Switch(Top-1 gate proposed by Fedus et al. (2021)), DTS-Dynamic(decay the temperature from 2 to 0.3 linearly in the first 15000 steps, and switch to Top1 at the 20000 step), and Dense-MoE(select all the experts for each token). The rank of theoretical FLOPs per token is GPT-MoE-Dense > DTS > Switch = GPT-Small(excluding the cost of gate network and all to all communication in MoE layer).
Training Hyperparameters All models use the same hyperparameters of 4,000 warmup steps, polynomial learning rate scheduler, Adam optimizer, clip norm of 5.0, weight decay of 0.1 and dropout of 0.1. We use GeLU activation functions (Hendrycks & Gimpel, 2016) and set sequence length as 1024. We use CrossEntropy as the criterion and 0.01 as the coefficient of balance loss in Switch. To save the training cost, we use batch size of 256 and sequence length of 768 on the GPTSmall model, other hyper-parameters keep constant. We set the threshold of gates as 0.001 across the training, which is a hyper-parameter to determine how large the expert’s weight is important and is a trade-off between computation and accuracy from our point of view.

heading: 4.2 RESULTS AND ANALYSIS

heading: 4.2.1 END TO END COMPARISON
text: To improve the computation efficiency, only part parameters are used for each token in sparse models with the cost of model performance. DTS-Gate aims to shift the model training from dense to sparse, and keep the inference cost same as sparse models, and results are shown in Figure 3. 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.
Comparison with Dense Models Dense models, where all parameters are participated in computation for each token, tends to achieve high model performance due to its large parameter capacity but with high computation cost. As shown by Figure 3(a), MoE-Dense performs better than other models at the same iteration. By utilizing the dense-to-sparse training mechanism, MoEDTS shrinks the performance gap between MoE-Dense and MoE-Switch. As for FLOPs-efficiency, MoE-DTS outperforms both dense models because of its sparse computation.
Comparison with Sparse Models MoE-Switch pre-defines its static Top-1 gating network and jointly training the gate and experts networks. Different from MoE-Switch, MoE-DTS utilizes temperature to adjust the distribution of the token-to-experts (one-hot or uniform) and threshold to remove computation of experts with low weights. MoE-DTS performs better than MoE-Switch in sample-efficiency because of more experts involved in training and updates at beginning, shown as Figure 3(a). As for FLOPs-efficiency, DTS-Gate first involve more experts into warm-up training, which is poor FLOPs-efficency. But with the training going on, DTS-Gate can obtain greater than 5% improvements in FLOPs-efficiency compared with the state-of-the-art Switch-Gate in GPTMoE(1.5B) model with OpenWebText dataset(40GB).
It is worth noting that several hyper-parameters of MoE-DTS, such as max/min temperature and decay iterations, needs to be carefully determined. In this paper, we just propose the training mechanism of Dense-To-Sparse and prove its effectiveness through a simple policy.

heading: 4.2.2 SCALABILITY
text: In this section, we conduct experiments about the scalability and FLOPs-efficiency of DTS-Gate. Although it doesn’t adopt a balanced routing algorithm, experiment results show that DTS-Gate can be utilized in scaling models, because it wouldn’t converge to only a few experts .
Increase the Expert Number Based on the 12-layer GPT-Small model(117M), we replace the 7th FFN layer by one MoE layer and vary its experts number within {1, 4, 8, 16}. As shown by Figure 4(a), MoE-DTS keeps consistent improvements when increasing expert numbers.
Increase the MoE layer number Similarly, we also vary the number of MoE layer within {0, 1, 4, 6} and each MoE layer contains 8-experts. Figure 4(b) shows that by increasing MoE layers, MoE-DTS can achieve better model performence with same FLOPs.

heading: 4.2.3 EFFECT OF SPARSITY SCHEDULER
text: Our work proposes the mechanism of Dense-To-Sparse training for MoE models, and employ Gumbel-Softmax to sparsify the MoE layers by decreasing the temperature τ . In this subsection, several experiments are conducted to show the effect of different temperature schedulers(τ−scheduler). The used models are GPT-MoE, 24-layer decoder with 12 MoE-layer (16 experts) and dataset is OpenWebText(40GB).
We decay τ from max value to min value in the first 15000 iterations and switch to Top1 at the 20000 step. Experiments with different max value and min value are evaluated and the results are shown in Figure 5.
Max/Min Temperature Under small temperatures, the weight distribution of experts is close to one-hot, which leads to the one-token-one-expert distribution and high computation efficiency, but the variance of gradients is large. In the contrast, large temperatures result in nearly uniform distribution gate weights, which evolves more experts into training but the variance of gradients is small. We schedule the temperature τ to start from a large value and anneal it to a small τ gradually. Figure 5 shows the sample-efficiency and FLOPs-efficiency of different τ− schedulers. By comparing {1 to 0.1} and {2 to 0.1}, we find larger begin temperatures will lead to high sampleefficiency but low FLOPs-efficiency, because more experts are participated in computation. By comparing {2 to 0.1} and {2 to 0.3}, we find smaller end temperatures cause poor model performance because of the large variance of gradients. After analysis about these experiments, we choose {2 to 0.3} as our training configs.
Decay Iterations After max/min temperatures are determined, we need to care about the decay iterations. {2 to 0.1}, {2 to 0.3} and {2 to 0.5} spend about 11800, 13200 and 15000 iterations respectively for decaying temperature from 2 to 0.5. Higher decay rate performs good at sampleefficiency at the expense of FLOPs-efficiency. Choosing suitable decay rate is a balance between sample-efficiency and FLOPs-efficiency. Actually, considering about the training cost, we usually set 5%− 15% as the decay iterations.
Dynamic V.S. Static Both large temperatures and small temperatures have their own advantages and disadvantages, it would be better to design a τ− scheduler according the models and hardware for the balance between model-performance and FLOPs-efficiency.

heading: 4.2.4 EFFECT OF BALANCED LOADS
text: Current works utilize either extra balance loss (Lepikhin et al., 2021; Fedus et al., 2021) or balanced assignment algorithm (Lewis et al., 2021a) to avoid the training of sparse MoE-layer converge at the status of only a few experts are frequently used. Balanced loads are naturally easy for efficient system implementation. However, it is still unknown about which is better for model performance. We conduct several experiments to explore this mystery. Meanwhile, DTS-Gate proposed by this paper is orthogonal to the balance loss, because DTS-Gate cares about the weights’ polarization at the token level while balance loss cares about the loads of each expert at model level/layer level.
With and Without Balance Loss Balance loss can improve the performance of switch transformer Fedus et al. (2021). Therefore, we conduct experiment with and without balance loss on DTS, to study whether it can also improve the model performance of DTS. To analyze the influence of balance loss on DTS-Gate, we replace the 7th layer of GPT-Small(117M) models with one MoE layer(16experts), and training this model with and without balance loss on OpenWebText and Wikipedia separately. As shown by Figure 6(a), this model performs almost the same on Wikipedia with and without balance loss, while it performs better without balance loss on OpenWebText at the
begining and then gradually coincide with the one with balance loss. The difference may introduced by that OpenWebText is a large and diverse dataset in various domains but Wikipedia only comes from one domain. Therefore, imbalanced distribution is suitable for training on OpenWebText. But Why OpenWebText with and without balance loss are similar at the end of training. We check the experts’ loads in next paragraph.
Expert Loads Distribution We directly load the checkpoints from the end to end experiments, where MoE-DTS is the 1.5B MoE models described as above. Figure 6(b) shows that the experts’s loads of the 3rd to 11th MoE layer is almost balanced, which is learned without balance loss.

heading: 5 CONCLUSION AND FUTURE WORK
text: MoE models suffer from the model quality challenge due to the difficulty of training the pre-defined sparse gate and the experts jointly. To address this challenge in MoE training, we presented DTSGate that (1) designs the dense-to-sparse training mechanism and (2) provides the scheduling policy for sparsity adjusting. Our evaluations show that DTS-Gate can not only achieve better model quality in Transformers with given computation budget but also achieve better FLOPs-efficiency when comparing with previous works in MoE training. DTS-Gate is a stepping stone in exploring MoE training. On the other hand, DTS-Gate opens challenges for system execution due to the computation in the early stage and the adaptive capacity of experts. In the future, we would like to design and implement system-level optimizations to achieve efficient training in both model quality and system execution.

heading: A APPENDIX
text: A.1 EXPERT LOADS IN DTS

Review summaries:
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.

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.

This work proposes the Dense-to-Sparse Gate or DTS-Gate. It’s a simple idea of starting training with soft (continuous) routing to all experts and then gradually reducing to hard (discrete) routing. The authors claim that this improves the quality of traditional approaches. The work proposes the novel idea and details how to schedule the sparsity adjustment. The authors suggest this improves traditional approaches by 5% FLOPs-efficiency.

This paper introduces the idea of using dense-to-sparse (DTS) training gates in mixture-of-experts (MoE) based on a gradual sparsification process. The novel produced method is evaluated on a Transformer model. Overall, it looks to me that the paper has rather limited novelty. While the paper is well-written, in my opinion, the empirical evaluation support is not perfectly aligned with the paper claims.