XMoE: Sparse Models with Fine-grained and Adaptive Expert Selection (2024)

Abstract: Sparse models, including sparse Mixture-ofExperts (MoE) models, have emerged as an effective approach for scaling Transformer models. However, they often suffer from computational inefficiency since a significant number of parameters are unnecessarily involved in computations via multiplying values by zero or low activation values. To address this issue, we present XMoE, a novel MoE designed to enhance both the efficacy and efficiency of sparse MoE models. XMoE leverages small experts and a threshold-based router to enable tokens to selectively engage only essential parameters. Our extensive experiments on language modeling and machine translation tasks demonstrate that XMoE can enhance model performance while decreasing the computation load at MoE layers by over 50% without sacrificing performance. Furthermore, we present the versatility of XMoE by applying it to dense models, enabling sparse computation during inference. We provide a comprehensive analysis and make our code available at https://github.com/ysngki/XMoE.

heading: 1 Introduction
text: Recently, remarkable advancements in large language models have been achieved through scaling up their sizes (Brown et al., 2020; Chung et al., 2022; Touvron et al., 2023). However, this progress has come with a significant increase in training costs, posing a challenge to further scaling. To address this issue, sparse models, such as sparse Mixture-of-Experts (MoE) models, have emerged as an alternative approach. These models allow for scaling the model size without a corresponding increase in computational cost (Lepikhin et al., 2021; Jaszczur et al., 2021; Fedus et al., 2022; Rajbhandari et al., 2022; Xue et al., 2024; Jiang et al., 2024). The key to this ability lies in sparsely activated MoE layers. These layers comprise multiple sub-networks, or “experts”, typically implemented
as Feed-Forward Networks (FFNs) (Vaswani et al., 2017). Unlike traditional models that utilize all parameters for each input token, MoE models selectively activate a subset of experts. This approach effectively decouples computational costs from model size, paving the way for more efficient scaling.
While MoE models are effective for scaling, this paper argues that MoE models exacerbate the issue of computational inefficiencies. Initially identified in the dense model T5 (Raffel et al., 2020), this issue is characterized by a significant portion of computations in the FFN layer being wasted on multiplying values by zero (Zhang et al., 2022; Li
11664
et al., 2023b). As illustrated in Figure 1, the computational inefficiency issue is also prevalent in sparse models and even worsens as the number of experts increases. This observation suggests that only a small portion of the parameters in an expert is useful for the input, while others are unnecessarily involved in the computation. Consequently, selecting one expert for each input can already lead to a significant waste of computation. In order to alleviate this problem, a more fine-grained and adaptive strategy for parameter selection is necessary.
Figure 2 shows the overview of a novel MoE design, named XMoE, which allows tokens to select fewer parameters to improve the efficiency without hindering model performance. To achieve this, XMoE proposes to exploit small experts and a threshold-based router. First, considering that expert is the smallest unit of parameter selection in MoE models, utilizing small experts is the prerequisite for the more fine-grained selection. It allows models to choose the useful parameters precisely without activating the redundant parameters.
In order to ensure the effectiveness, a novel adaptive router is further exploited by XMoE. Different from the widely-used top-k router that dispatches each input to a fixed number of experts, this adaptive router allows tokens to self-determine the number of required experts based on a pre-defined threshold. Intuitively, an easily processed token can be routed to a single small expert, while a critical token may require multiple experts (Zhou et al., 2022; Li et al., 2023a). An adaptive router allows models to leverage the difference of the input complexity to dynamically allocate computational resources. This not only enhances model efficiency but also yields potential quality improvements when computational resources are constrained.
In conjunction with the aforementioned design, XMoE aims to enhance the efficiency of MoE mod-
els, with a focus on improving quality with a fixed computational budget or reducing computational costs without compromising performance. Extensive experiments in language modeling and machine translation demonstrate performance gains over existing MoE methods. XMoE also enables a reduction in Floating Point Operations (FLOPs) by over 50% with minimal impact on performance. Additionally, our investigation extends to training dense models using XMoE to facilitate sparse computation during inference. This approach not only matches the performance of dense counterparts but also facilitates a substantial reduction in FLOPs.
Our contributions can be summarized as: (1) We identify a computational inefficiency issue in current sparse MoE models. (2) We propose XMoE to improve the efficiency of MoE models with small experts and a novel routing strategy. (3) Extensive experiments in language modeling and machine translation demonstrate XMoE as a promising alternative to existing sparse and dense models.

heading: 2 Background

heading: 2.1 Mixture-of-Experts (MoE)
text: MoE is a family of neural network architecture that enables conditional computation by sparsely activating small sub-networks, so-called experts, based on a pre-defined router. The core component of MoE models is the MoE layer, which consists of a set of experts {E1, · · · , EN} and a routing function (Shazeer et al., 2017; Zoph et al., 2022; Dai et al., 2024). Each expert E is a parameterized function. The routing function is utilized to route individual tokens to their assigned experts. In this work, we consider MoE for Transformer, where FFN layers within Transformer are substituted with large MoE layers, in which each expert is an independent FFN.
Trainable Router. A commonly used router is based on a gating network consisting of a trainable weight matrix followed by a softmax function. For the input token x with its intermediate representation denoted as h ∈ Rd, the router computes the probability distribution over experts as:
p = Softmax(W · h), (1)
where W ∈ RN×d, N denotes the total number of experts and d denotes the hidden size of the model. The set of top-k experts E is decided based on p, where |E| = k. Each expert processes tokens
independently. Then, the final output of the token is the weighted sum of the output of its k experts:
y = ∑
i∈E piEi(h). (2)
Load Imbalance. Learnable routers can easily cause the load imbalance issue, as shown in Shazeer et al. (2017). If there is no constraint employed during training, most tokens would be dispatched to a small number of experts, leading to a large portion of experts being insufficiently trained (Lepikhin et al., 2021). In addition, if experts are distributed across different nodes, most nodes must wait for others to finish the computation, thus hindering the training efficiency.
Capacity. Expert capacity is introduced to MoE models to avoid the severe load imbalance issue by limiting the maximum number of tokens routed to each expert (Lepikhin et al., 2021; Fedus et al., 2022; Rajbhandari et al., 2022; Puigcerver et al., 2024). Suppose a given batch’s token number is T , and the expert number is N . The expert capacity C is:
C = T
N · γ (3)
where γ refers to the preset capacity factor. If an expert is underutilized, the unused capacity buffers are filled with padding tokens. Once an expert is at capacity, additional tokens are dropped, which means being passed to the subsequent Transformer block directly.
MoE for Dense Models. Although MoE models are proposed for training large-scale sparse models, recent studies bring the concept of MoE for either pruning or training dense models. Zhang et al. (2022) propose to divide the feed-forward network layers in a pretrained dense Transformer into several small experts based on heuristic strategies. Each input token selectively activates top-k experts instead of utilizing the whole FFNs. Chen et al. (2023) propose to view MoE as a regularization method for training dense Transformers. Although our work proposes XMoE for sparse model training, it can also be utilized for training dense models to enable sparse computation during inference.

heading: 2.2 FFN as Memory
text: Recent studies suggest that an FFN layer can be conceptualized as a memory layer composed of
numerous key-value memory pairs (Lample et al., 2019; Geva et al., 2021). In this view, each column in the first matrix of an FFN layer is a key vector, and the corresponding row in the second matrix is the value vector. It is observed that only some memory values benefit an input token (Dai et al., 2022), leading to most memory pairs contributing to redundant computations. We hypothesize that in MoE models, which incorporate multiple FFNs, the beneficial memories tend to be distributed across different FFNs. This dispersion of useful memories diminish overall efficiency.

heading: 3 Method

heading: 3.1 FFN Decomposition
text: Considering that widely-used activation functions such as ReLU and GELU operate on an elementwise basis, it is reasonable to conceptualize an FFN layer as a composition of several smaller FFN layers
y = W2 σ(W1 · x) = ∑
i
W i2 σ(W i 1 · x), (4)
where W ij is the parameters of i-th small FFN layer and σ denotes the activation function. Instead of maintaining a single large FFN, XMoE trains multiple small FFNs and compose them to produce an output. In subsequent discussions, we assume that the dimensions are identical for each experts unless stated otherwise. Our preliminary experiments showed that the GELU activation function performed slightly better. Therefore, we will use GELU by default.

heading: 3.2 Threshold-based Router
text: XMoE’s router consists of a trainable weight matrix W ∈ Rd×N , where each column of W serves as a centroid representing an expert. Given an input token x and its intermediate representation h ∈ Rd, the probability of choosing the i-th experts is pi, according to the Eq (1).
Considering the distribution p varies with layers, tokens and model configurations, manual selection of an optimal value for the top-k parameter k can be challenging (Yang et al., 2021; Li et al., 2023a). A higher value of k can ensure effectiveness but at the cost of increased computational overhead per token, potentially impacting efficiency. Conversely, a lower value of k may reduce computational load but restrict the model’s capacity to handle complex tokens. To navigate this trade-off, XMoE employs
a threshold-based selection mechanism, enabling tokens to self-determine the number of experts they should be routed to based on a predefined threshold parameter t, which ranges from 0 to 1.
The overview of the selection procedure is illustrated in Figure 3. Initially, the probabilities p are sorted in descending order: p = [pi1 , pi2 , ..., piN ], where pi1 ≥ pi2 ≥ ... ≥ piN and ij indicates the index of top-j expert. The router then identifies the smallest index m such that the cumulative sum of probabilities up to index m is greater than or equal to the threshold t:
argmin m
m∑
j=1
pij ≥ t. (5)
The input tokens are dispatched to their first m experts. According to the Eq (2), an expert’s contribution to the output is directly proportional to its assigned probability p. Thus a high cumulative sum of probabilities indicates that tokens are routed to the most relevant experts, while other experts with lower probabilities have minimal influence and can be ignored.
Priority. The threshold-based router in MoE models permits a token to select multiple experts, potentially overwhelming certain experts and leading to dropped tokens due to limited expert capacities. To mitigate this, XMoE assigns a priority r to a token dispatched to a specific expert. The experts process tokens with higher priorities first. Specifically, the prioritization is based on the likelihood of a token considering the expert as its preferred choice, determined by a heuristic rule. If a token views an expert as its top-i choice with probability p, the priority r is set heuristically as (p− i). Appendix A.1 presents the algorithmic implementation.

heading: 3.3 Auxiliary Loss
text: Following Fedus et al. (2022), we utilize an auxiliary loss as a part of the training objective. It encourages input tokens to be uniformly dispatched to the experts. We only impose this constraint on the top-1 assignment. The details are provided in Appendix A.2.

heading: 3.4 Complexity
text: Within an MoE layer, a total of C · N tokens are processed by the experts, where C denotes the capacity and N is the number of experts. Let O(E) represents the computational complexity associated
with an expert handling a single token. The computational complexity of an MoE layer, as indicated by Eq. 3, can be expressed as:
O(MoE) ∝ γ · T ·O(E). (6)
A small γ can reduce the computational cost but aggravate token dropping since experts may not have enough capacity to process received tokens. On the contrary, increasing γ could lead to additional computational overhead and tend to waste computing resources on padding tokens. XMoE allows to trade off the efficacy and efficiency by concurrently increasing γ and diminishing O(E).

heading: 3.5 Comparison with Top-k Router
text: Top-1 Router. The top-1 router exclusively assigns each token to a single expert. It potentially leads to inefficiencies when experts receive fewer tokens than their capacity, resulting in wasted computations processing padding tokens. In contrast, the threshold-based router dispatch each token to at least one expert. Consequently, the token assignment of the top-1 router is a subset of that produced by the threshold-based router. This enables models with the threshold-based router to implicitly reduce computational inefficiencies. Top-k Router. Top-k router allows each token to choose multiple experts, thus unlikely leading to computation waste, especially when the capacity is limited. However, this router processes tokens with the same amount of computation, regardless of the difference in complexity between tokens. In contrast, an adaptive router allows the model to dynamically allocate computational resources to important tokens while reducing unnecessary computations, leading to better utilization of training computation.
The threshold-based router and the top-k router share parameters of the same shape, with their only distinction lying in their respective selection strategies. This structural similarity indicates that the threshold-based routing strategy and the top-k strategy can be seamlessly interchanged without necessitating any modifications.

heading: 3.6 Dense Model Applications
text: A dense model can be treated as an MoE model with only one expert per layer. XMoE can be applied by decomposing FFN layers in the dense model into multiple smaller ones. In contrast to training a sparse model, we densely train this model by setting the threshold t = 1.0 and γ as N . This setup ensures that each token is processed by all experts without token dropping. It is expected that the router can learn to measure the importance of different experts. During inference, both t and γ can be adjusted to enable sparse computation.

heading: 4 Experiments

heading: 4.1 Tasks and Datasets
text: Language Modeling. We pretrain models on two English-language datasets, OpenWebText (Gokaslan and Cohen, 2019) and Wikitext-103 (Merity et al., 2017). The former is a public reproduction of WebText used for GPT-2 training (Radford et al., 2019). The latter is a smaller language modeling corpus containing Wikipedia articles.
Machine Translation. We have collected 5 language pairs from WMT23 datasets. Each language pair is trained independently. The detailed statistics for these language pairs can be found in Table 4.

heading: 4.2 Model Configuration
text: Language Modeling. We adopt the Transformer decoder with 12 layers. Our implementation is based on Megatron-LM (Shoeybi et al., 2019). We use the tokenizer of GPT-2 of which the vocabulary size is 50, 257. The model is trained for 60K steps in total on OpenWebText and 4k steps on Wikitext103. The learning rate is 4.5× 10−4 and the cosine learning rate decay is exploited. During training, we use float16 for acceleration. The threshold t is set to 0.90 based on the results of pilot experiments.
Machine Translation. We adopt the Transformerbased architecture with 12 encoder layers and 6 decoder layers. Our implementation is based on
fairseq (Ott et al., 2019). The vocabulary is learned from the training data for each language pair using byte-pair encoding. The threshold t is set to 0.90. Automatic mixed precision is enabled to accelerate training. We report the detokenized BLEU and chrF2++ using sacreBLEU1.
For sparse models, an MoE layer is utilized to replace the dense FFN layer in every alternate Transformer block, following previous practice (Zhou et al., 2022) . The capacity factor γ is set to 1.0 for models with top-1 routers. For models with smaller experts, we increase the capacity factor according to Eq (6) to ensure the consumed training compute identical. We run our experiments on one node with 4 NVIDIA A100 GPUs.

heading: 4.3 Baselines
text: Top-1 routing, which assigns each token to one expert, is widely used in MoE models, such as Switch Transformers (Fedus et al., 2022), BASE Layers (Lewis et al., 2021) and Hash Layers (Roller et al., 2021). Switch Transformer utilizes a learnable top1 router with an auxiliary loss to alleviate the load imbalance issue. BASE Layer proposes to view the token assignment as a linear assignment problem. They force each expert to process an equal number of tokens during training while using greedy assignments at test time. Hash Layers replace the learnable router with simple hash functions. We implement Hash Layers with a random hash function, which is suggested to have strong performance (Roller et al., 2021). We extend the top-1 router of Switch Transformer to the top-k (k > 1) router.

heading: 5 Results

heading: 5.1 Language Modeling
text: Table 1 reports the perplexity results of models with equivalent computational complexity in MoE layers. A consistent performance gain is observed across both the OpenWebText and WikiText-103 datasets when reducing the expert size. Specifically, when the expert size decreases to 384, XMoE with 323M parameters achieves the best performance. It surpasses top-k routing and Switch Transformer by 0.33 and 0.65 perplexity points on OpenWebText, respectively. The difference between XMoE and top-k routing lies in the routing strategy, while the difference between top-k routing and Switch
1BLEU:nrefs:1 | case:mixed | eff:no | tok:13a | smooth:exp | version:2.3.2. chrF2:nrefs:1 | case:mixed | eff:yes | nc:6 | nw:2 | space:no | version:2.3.2
Transformer lies in the expert size. This suggests that a threshold-based router can assist models in leveraging expert capacities, and smaller experts can enhance model quality. There are intriguing outcomes when we increases the parameter count to 556M by expanding the intermediate dimension of experts. It is shown that top-k routing with a size of 1024 outperforms that with a size of 512 on OpenWebText, as does XMoE. We attribute this to the optimization of sparse models and leave this for future investigation. Efficiency. Figure 4 illustrates the impact of the number of floating point operations (FLOPs) within MoE layers on perplexity score. We see that models exhibit poorer performance as FLOPs decrease. This trend is expected, as limited expert capacity can lead to more dropped tokens, thereby increasing the number of inadequately processed tokens.
It is observed that XMoE with small expert size is robust to the reductions in FLOPs. Notably, XMoE with an expert size of 384 outperforms Switch Transformer (referred to as Exp8-Size3072) on WikiText-103, while consuming only 25% of the FLOPs of the latter. This finding suggests that
a significant portion of parameters in models with large expert sizes are unnecessarily engaged in computations. By replacing these large experts with smaller ones, XMoE not only improves performance but also saves computational resources.
Dense Model. Dense models can be trained as XMoE by partitioning their FFN layers into smaller ones. As shown in Table 2, these modified models exhibit a remarkable reduction of over 50% in computational complexity with only a marginal decrease in performance across both datasets. On the
WikiText-103 dataset, XMoE can outperform the dense model while achieving a 75% reduction in Floating Point Operations (FLOPs). These findings underscore the potential of sparse models as promising alternatives to their dense counterparts.

heading: 5.2 Machine Translation
text: We compare XMoE with Switch Transformer and Top-k routing. Models are trained to translate other languages to English. All models have the same number of shared parameters and have the same FLOPs. The results are detailed in Table 3, revealing that XMoE outperforms its counterparts. The observations generally align with the finding on language modeling that a reduction in expert size can always lead to better performance. However, the extent of the improvement depends upon factors such as the specific language being translated, and the evaluation metrics employed.

heading: 5.3 Analysis
text: Efficiency. Figure 5 shows the number of experts that tokens require to satisfy the threshold vs. training step. Initially, a substantial portion of experts is required to satisfy the threshold, but as training progresses, a significant reduction in the required number of experts is observed. This
phenomenon is consistent with existing research indicating that dense models exhibit sparse activation once trained (Li et al., 2023b). Our proposed XMoE leverages this emergent sparsity to to enhance computational efficiency by employing small experts and a threshold-based router. The use of small-scale experts facilitates fine-grained expert selection, thereby mitigating the activation of redundant parameters. Concurrently, the thresholdbased router encourages tokens to select the fewest experts. While the threshold-based routing may initially lead to an excessive selection of experts, the computations is supposed to exhibit sparsity after training.
Effectiveness. The visualization in Figure 6 illustrates the average percentage of positive values after the activation function across different configurations. It can be seen that a reduction in expert size consistently correlates with an increase in this percentage. This trend suggests that models with small experts can more effectively leverage the parameters within selected experts, thereby yielding performance improvements.

heading: 5.4 Hyperparameters
text: Threshold. We investigate the effect of threshold t on the perplexity score during both training and inference stages. From Figure 7, we see that increasing the threshold generally improves perfor-
0.0 0.2 0.4 0.6 0.8 1.0 0.0
0.0 0.2 0.4 0.6 0.8 1.0 0.0
1.0
mance. However, XMoE with a threshold of 0.9 slightly outperforms the model with a threshold of 1.0. It is noteworthy that at a threshold of 1.0, each token is sent to all experts by the thresholdbased router. In contrast, at a threshold of 0.0, the router behaves like a top-1 router. Hence, replacing the top-1 router with the top-N router, where N represents the number of experts, is also a simple approach to enhance performance.
Wall Time. Figure 8 illustrates the per-batch wall time during inference. We observe that a reduc-
0.2
tion in FLOPs at the Mixture of Experts (MoE) layers correlates with a decrease in overall wall time. However, XMoE with smaller experts exhibit significantly higher latency compared to those with larger experts. The reason is that smaller experts require more sparse computation, which is not well supported on computation hardware such as TPUs and GPUs (Li et al., 2023b). In addition, the increased number of experts introduces additional computational overhead to the routing process due to operations like sorting and ranking across the expert pool. These observations highlight the limitations associated with further reducing expert dimensions. The size of experts should be properly chosen in order to fully utilize the advantages offered by XMoE.

heading: 6 Conclusion
text: This paper proposes a novel MoE design, XMoE, with the primary objective of improving the efficacy and efficiency of sparse MoE models. By employing small experts and a threshold-based router, XMoE demonstrates performance enhancements while significantly reducing FLOPs, leveraging the inherent sparsity of the model. Our research sheds light on the utilization of sparsity to improve model quality. As for future directions, we aim to further harness the advantages of sparse computation, fo-
cusing on enhancements from both hardware and algorithmic perspectives.

heading: 7 Limitations
text: Our experiments were conducted on language modeling and machine translation tasks. To ascertain the effectiveness of XMoE across a broader spectrum of NLP tasks, additional experiments are necessary. Additionally, due to the limited computational resources in our experiments, the largest model explored in this paper comprises 556 million parameters, notably smaller than the parameter counts in prevalent large-scale language models, which often exceed billions. To substantiate the claims made in this paper, further investigations in larger-scale configurations are needed. The expert size is also an important factor to XMoE, and setting it to 1 could yield valuable insights. Regrettably, our current implementation makes this model unfeasible. We will leave these further investigation as our future work.

heading: A Appendix
text: A.1 Expert Selection
The process for expert selection concerning a given token is presented in Algorithm 1. Tokens are dispatched to experts according to I . Subsequent to the routing phase, each expert independently handles the tokens assigned to them. Given that the capacity of each expert is restricted to C, only the top C tokens, as per the priority R, are processed by each expert. Any tokens beyond capacity are disregarded.
Algorithm 1 Expert Selection Procedure Require: Probability distribution p =
[p1, p2, ..., pn], Threshold t Ensure: Indices of selected experts I , correspond-
ing priorities R 1: Sort p in descending order 2: Initialize I = [], R = [] and sum = 0 3: for i = 1 to N do 4: sum = sum+ pi 5: Append i to I 6: Append pi − i to R 7: if sum ≥ t then 8: Break 9: end if
10: end for 11: return I , R
A.2 Load Balancing Loss
Let N represent the total number of experts involved in the evaluation process. The auxiliary loss function is formulated as follows:
loss = N · N∑
i=1
fi · pi, (7)
where fi denotes the fraction of tokens that rank the i-th expert as their top choice, and pi represents the sum of probabilities assigned to the top-ranked selection by these tokens.

heading: Future Work
text: We will leave these further investigation as our future work