Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models (2024)

Abstract: A pivotal advancement in the progress of large language models (LLMs) is the emergence of the Mixture-of-Experts (MoE) LLMs. Compared to traditional LLMs, MoE LLMs can achieve higher performance with fewer active parameters, but it is still hard to deploy them due to their immense parameter sizes. Different from previous weight pruning methods that rely on specifically designed hardware, this paper mainly aims to enhance the deployment efficiency of MoE LLMs by introducing plug-and-play expert-level sparsification techniques. Specifically, we propose, for the first time to our best knowledge, posttraining approaches for task-agnostic and taskspecific expert pruning and skipping of MoE LLMs, tailored to improve deployment efficiency while maintaining model performance across a wide range of tasks. Extensive experiments show that our proposed methods can simultaneously reduce model sizes and increase the inference speed, while maintaining satisfactory performance. Code will be made available at https://github.com/Lucky-Lance/ Expert_Sparsity.

heading: 1 Introduction
text: Large language models (LLMs) have shown remarkable abilities across various domains (OpenAI, 2023; Zhou et al., 2024), as evidenced by the widespread use of ChatGPT and Gemini (Team et al., 2023). Recent notable advancement in this area is the introduction of the opensourced Mixture-of-Experts (MoE) LLM, Mixtral 8x7B (Jiang et al., 2024), which sparsely activates only a portion of its parameters during the training and inference process. This model surpasses the performance of dense Transformer-based LLMs, such as LLaMA-2 70B (Touvron et al., 2023a,b), with fewer active parameters (13B) during inference.
∗Equal contribution †Corresponding author ‡Project lead
MoE LLMs achieve a reduction in on-the-fly (active) parameters by choosing only top-k experts for the inference of each token, thereby enhancing inference speed (Sanseviero et al., 2023). However, the static parameters, particularly those required for constructing the MoE architecture, still demand considerable memory and storage for deployment. For example, loading the Mixtral 8x7B model in bf16 format requires at least two A100-80G GPUs. Notably, in this MoE model, the eight experts constitute around 96% (45B out of 47B) of the total number of parameters.
On the other hand, not all experts are equal in the MoE model. Recent studies, such as (Chi et al., 2022), have demonstrated this discrepancy in expert training outcomes. The differing levels of training among each expert highlight the importance and practicality of identifying and pruning less significant experts, thereby improving the deployment efficiency of MoE models.
Unlike existing post-training weight pruning schemes for LLMs, which primarily target unstructured sparsity and N:M semi-structured sparsity (Sun et al., 2023; Frantar and Alistarh, 2023), our approach focuses on expert-level sparsity for model sparsification. The aforementioned finegrained weight pruning techniques are effective in reducing the total number of parameters. However, they face challenges in plug-and-play deployment due to the necessity for specific hardware designs (e.g., FPGA) (Zhou et al., 2021), which demands an extensive co-design of hardware and systems.
In this paper, we systematically explore expertlevel sparsity in MoE LLMs and, for the first time to our best knowledge, introduce hardware-friendly post-training methods for either permanently removing unimportant experts (expert pruning) or dynamically skipping experts during inference (dynamic expert skipping). Our proposed method significantly reduces memory usage for deploying MoE LLMs and enhances their inference speed.
6159
Initially, we investigate how to prune less important experts while maintaining satisfactory performance, utilizing an efficient post-training approach. We aim to minimize the token reconstruction loss in a layer-by-layer manner. Given the limited number of experts in a single MoE layer of the LLM, we meticulously enumerate and choose combinations of experts that yield the lowest token reconstruction loss, subsequently, concatenating them to obtain the final pruned MoE model. This strategy significantly lowers the memory demands for deploying MoE LLMs. We examine expert-level pruning for both task-agnostic and task-specific (first in literature) models, tailoring our strategies to optimize performance across a wide range of applications.
Building on this foundation, we further dive into strategies for accelerating the inference speed of MoE LLMs without compromising their robustness. Specifically, based on the model’s fixed expert count, we introduce an online method for dynamically skipping certain experts. This approach, which is complementary to our expert pruning strategy, allows for on-the-fly adjustment of the number of active experts during inference, thus enhancing the inference speed. By integrating the dynamic (expert) skipping approach with expert pruning, we achieve a more streamlined and efficient deployment for MoE LLMs.
Experiments on Mixtral 8x7B (Instruct) models (Jiang et al., 2024) demonstrate that our methods significantly reduce memory usage and increase inference speed. Take the example of posttraining pruning two experts. As shown in Fig. 1, we halve the number of GPUs needed, allowing deployment on a single 80G GPU and achieving a 1.2× inference speedup. The pruning also results in mild performance loss, specifically, around 2.9 points for task-agnostic and 6.2 points (reducible to
1.6 with task-specific fine-tuning) for task-specific models. Further combination of dynamic skipping with expert pruning can lead to the same inference speedup with dropping 4 experts while achieving much higher model performances. To the best of our knowledge, this study is the first to discuss expert-level sparsity and propose efficient schemes for expert pruning and skipping for MoE LLMs.

heading: 2 Related Works

heading: 2.1 Mixture-of-Experts Models
text: First introduced in (Jacobs et al., 1991), a Mixtureof-Experts (MoE) model contains multiple separate networks, and each network processes a subset of the entire dataset. This separation can be viewed as a modular transformation of a multi-layer network. MoE structure is used for designing Recurrent Neural Networks (RNNs) in (Shazeer et al., 2017) and further extended to encoder-decoder Transformer-based models (Lepikhin et al., 2020). With the recent development of decoder-only GPT family of models (Brown et al., 2020; Touvron et al., 2023a,b), MoE models based on this structure gain popularity (Jiang et al., 2024). In this paper, we focus on post-training expert pruning/skipping methodologies for MoE LLMs.

heading: 2.2 Expert Pruning for MoE Models
text: Expert pruning within MoE models has garnered attention in the realm of Natural Language Processing, particularly in machine translation tasks. In these contexts, the translation of specific languages often renders the expertise of other language specialists superfluous. The most activated experts are reserved in (Kim et al., 2021) to prune a machine translation MoE model, and (Koishekenov et al., 2022) proposes expert pruning metrics based on gate statistics collected during decoding. Al-
though these methods actively deal with expert pruning for MoE models, they are still limited to the machine translation domain with linguistic models. Researchers in (Chen et al., 2022) provide a dropping-while-training method that progressively drops the non-professional experts for target downstream tasks, and experiments are carried out on Switch Transformers models (Fedus et al., 2022). However, in the LLM era, it is usually difficult to afford such a training paradigm.

heading: 2.3 Post-training Pruning for LLMs
text: Post-training pruning (Kwon et al., 2022; Hubara et al., 2021) has become a popular topic for neural network sparsification in recent years. Given a trained model, post-training pruning aims at achieving the optimal model sparsification outcome by utilizing model parameters together with some calibration data. Recent works extend pruning methods to LLMs (Sun et al., 2023; Frantar and Alistarh, 2023). However, these pruning methods primarily focus on sparsifying the weight matrices of linear layers in the LLMs and require dedicated hardware. To the best of our knowledge, efficient post-training expert pruning methods have not been discussed for decoder-only LLMs with MoE structures.

heading: 3 Method
text: To enhance the deployment efficiency of MoE LLMs, we concentrate on expert-level model sparsity and innovatively propose post-training techniques designed to reduce memory usage and increase inference speed. In this section, we offer a comprehensive explanation of our proposed methods for expert pruning (Sec. 3.2) and dynamic (expert) skipping (Sec. 3.3), considering both memory consumption and inference speed.

heading: 3.1 Preliminary
text: In the decoder-only sparse MoE Transformer models, as discussed in (Gale et al., 2023), the FeedForward Network (FFN) sub-layers of the traditional dense model are replaced with MoE layers, each containing n experts. Specifically, within the MoE layer of the Mixtral 8x7B model, featuring 8 experts as detailed in (Jiang et al., 2024), each token x in the input sequence is routed to the top-2 experts based on the routing weights w.
The inference process for each input token x within the MoE layer of the Mixtral 8x7B decoder layer is depicted in Fig. 2. Initially, the router computes routing logits l = {l0, . . . , ln−1} and routing
weights w = Softmax(l) for the experts. Then the top-k experts, where k = 2 for the Mixtral 8x7B model, are selected based on their routing weights to process the token. Each of these k selected experts, applying a SwiGLU transformation Ei(·) (i ∈ {0, 1, . . . , n− 1}), contributes to the final output. This output is a weighted sum of the individual expert outputs, with weights w̃i being the normalized values of the corresponding routing weights for the selected experts. The normalized weight for expert ej (j ∈ {0, 1, . . . , k − 1})1 is calculated as follows:
w̃ej = wej∑k−1
m=0wem , (1)
yielding the MoE layer’s output for the token x as:
z =
k−1∑
j=0
w̃ej · Eej (x). (2)
Aside from these specified mechanisms, the remaining aspects of the MoE network mirror those of a standard decoder-only Transformer model.

heading: 3.2 Post-training Expert Pruning
text: As demonstrated in the preliminary subsection, the parameters of experts occupy the major proportion of the whole MoE LLM model. However, for a single token, only a small subset of these experts are activated, leading to considerable inefficiencies in parameter utilization. Existing posttraining weight pruning methods for LLMs (e.g., Wanda (Sun et al., 2023), SparseGPT (Frantar and Alistarh, 2023)), while effective in reducing model
1ej is the index of the j-th largest element of w, i.e. the index of the j-th selected expert
parameters, do not support efficient deployment of MoE LLMs without specialized hardware implementations (Mishra et al., 2021). Therefore, we introduce a heuristic search method to prune the number of experts in a post-training manner.
Task-agnostic Expert Pruning for General Tasks. Different from existing pruning schemes leveraging unstructured or semi-structured weight sparsity (Sun et al., 2023; Frantar and Alistarh, 2023) for LLMs, our proposed post-training expert pruning method aims at reducing the parameter numbers of MoE LLMs by permanently discarding less important experts, thereby improve the inference efficiency. It is a post-training pruning method and does not require any parameter update. Fig. 3 (a) illustrates our proposed post-training expert pruning method. We conduct expert pruning in a layer-wise manner. Specifically, the pruning method contains two steps as follows.
Firstly, we set up a small calibration dataset, then perform inference on the original MoE model with all experts over the dataset, and cache the inputoutput token pairs of each MoE layer. We use samples from the pre-training dataset C4 (Raffel
et al., 2019) as calibration data, since pre-training datasets are often more comprehensive and not dominated by knowledge specific to any particular domain (Sun et al., 2023).
Secondly, after caching input-output pairs for each MoE layer, we enumerate expert combinations based on the preset parameter r, denoting the number of preserved experts. Assume the function of the MoE layer at layer l is F(·), with the cached input represented by x. During each enumeration, we maintain r experts and eliminate the remaining experts along with their associated routing weights. Subsequently, we employ the pruned MoE layer, denoted as F ′(·,C), to recalculate the corresponding output, where C represents a subset containing r experts selected from the original n experts. Inspired by channel pruning in CNNs (He et al., 2017), the Frobenius norm of the difference between cached output F(x) and the output of pruned layer F ′(x,C) is used to quantify the discrepancy between the model before and after expert pruning, and we denote it as reconstruction loss. The expert subset corresponding to the minimum reconstruction loss is retained. The n− r experts
left are considered to contribute least to the original MoE model and are thus discarded. The expert pruning process in layer l can be formulated as:
min C
∥F ′(x,C)−F(x)∥F (3)
s.t. C ⊆ {expert0, . . . , expertn−1}, |C| = r.
We heuristically search for expert subset C with the lowest reconstruction loss in each layer separately and obtain an MoE model with r experts by the concatenation of each pruned layer. After removing the insignificant experts, the pruned model can be easily loaded using existing packages (e.g., Huggingface Transformers (Wolf et al., 2020)) with just a change of the model configuration. Especially, with 2 experts pruned, the deployment budget is reduced to a single 80G GPU for loading the Mixtral 8x7B (Instruct) model with bf16 data type.
Task-specific Expert Pruning for Domainspecific Tasks. Previous research on post-training pruning for LLMs (Sun et al., 2023; Frantar and Alistarh, 2023) typically considers the performance over general tasks. For the first time, in our work, we investigate the task-specific post-training pruning for MoE LLMs. Our above-proposed expert pruning strategy is adept at conserving the knowledge encapsulated by the experts of an MoE model for general tasks. However, as a pre-training dataset, C4 spans a wide array of domains, posing challenges when pruning experts for domain-specific tasks (e.g., models tailored for mathematics (Yu et al., 2023; Wang et al., 2024)). We evaluate the 5-shot performance of the C4 pruned MoE LLM model with r = 6 on math tasks (GSM8K) (Cobbe et al., 2021), resulting in a performance degradation plummeting from 58.61 to 41.02 for the Mixtral 8x7B model. To address this challenge, we propose to shift the calibration dataset from C4 to the training set of MATH (Hendrycks et al., 2021), to concentrate the pruning process on the mathematics domain.
Remark. For a better comparison between general tasks and domain-specific tasks, we visualize the distribution of pruned experts selected by C4 and MATH, as shown in Fig. 4. For both the Mixtral 8x7B and Mixtral 8x7B Instruct model, identical expert combinations are observed in only four specific layers. This suggests that there are distinct differences in the distributions of pre-training datasets and domain-specific datasets. Further details and discussions are provided in Sec. 4.2.

heading: 3.3 Dynamic Skipping During Inference
text: Our expert pruning strategy effectively reduces memory consumption during model deployment. However, each token is still processed by k selected experts, without a reduction in runtime FLOPs. Intuitively, not all tokens require the selection of all top-k experts during the token generation process. Consequently, we introduce a scheme that dynamically skips certain experts for individual token inference, to further enhance inference efficiency.
As shown in Fig. 2, during the inference process, top-k experts are chosen with routing weights w = {we0 , we1 , . . . , wek−1} for each token x in an MoE layer. For simplicity, we assume k = 2 (as in Mixtral 8x7B)2. Our proposed dynamic expert skipping method is illustrated in Fig. 3 (b). Without loss of generality, assume experts with indices e0 and e1 are chosen, and we1 < we0 . To accelerate inference speed, if we1 < βwe0 , we do not assign x to expert e1, where β is a hyper-parameter separately set for each MoE layer. In our implementation, we forward the model over the sampled calibration data and set β as the median value of we1 we0
for each MoE layer. The dynamic skipping scheme can be carried out on the fly to speed up inference, and can be used simultaneously with expert pruning. In our experiments, we observe a 1.2× to 1.3× inference speedup with r = 6.

heading: 4 Experiment
text: In this section, a series of experiments are carried out to evaluate our proposed methods. We introduce experiments of expert pruning for general tasks in Sec. 4.1, domain-specific tasks in Sec. 4.2, and dynamic expert skipping results in Sec. 4.3.

heading: 4.1 Expert Pruning for General Tasks
text: In this subsection, we evaluate the proposed expert pruning method on some general tasks, which can comprehensively reflect the knowledge retention of the model after expert pruning.
Experiment Setup. Similar to Wanda (Sun et al., 2023), we choose calibration data from the C4 (Raffel et al., 2019) dataset and combine them into 128 token sequences, each with a length of 20483. We perform expert pruning on both Mixtral 8x7B and Mixtral 8x7B Instruct models, resulting in MoE
2Deeper theoretical insight and a broader application to the top-k scenario of dynamic skipping can be found in Sec. A.2.
3For more experiments about the influence of the calibration dataset size, please refer to Sec. A.3.
models with two experts discarded (r = 6) and four experts discarded (r = 4) in each layer. Pruning a Mixtral 8x7B model takes about 30 minutes for r = 6 and 90 minutes for r = 4.
After expert pruning, we evaluate the performance of the pruned MoE models following Wanda (Sun et al., 2023). Specifically, we report zero-shot accuracies of 8 tasks from EleutherAI LM Harness (Gao et al., 2023). We also test the token generation speed4, together with the peak GPU memory usage during model inference.
Comparison with Weight Pruning Methods. We compare our proposed expert pruning method with the representative weight pruning algorithm Wanda. For a fair comparison, we set r = 4 in our method and test Wanda with the commonly used structured 2:4 sparsity pattern. This will lead to around 50% parameter reduction for both methods. The results are shown in Tab. 1. The inference speedup of 2:4 structured model relies on specially designed hardware (Mishra et al., 2021) and scripts. In our experiments, we even observe a lower inference speed compared with the dense weight model5.
4We revise the script provided in https://github.com/ AutoGPTQ/AutoGPTQ/ to test token generating speed.
5Our implementation is based on https://pytorch.org/ tutorials/prototype/semi_structured_sparse.html.
Besides, our expert pruning method excels Wanda with the 2:4 sparsity pattern in both memory usage and benchmark performance.
Comparison with Other Expert Pruning Baselines. We also set up two baseline methods. One baseline is randomly dropping experts in each layer. The other method calculates the activation frequency of each expert during forward passes on the calibration data and discards those with the lowest activation frequencies in each layer. Comparison results are listed in Tab. 2. The method based on activation frequency gets the worst performance. This phenomenon implies that although the MoE model might show a tendency for expert selection during the inference process, simply carrying out expert pruning based on activation frequency might not always lead to desirable results. In contrast, our proposed method achieves better results. Compared to the origin model with 8 experts, our model achieves a 2.9-point performance drop with 24% parameter reduction and a 7.1-point performance drop with 48% parameter reduction on average without any extra training.
Memory Usage and Generation Speed. The memory usage statistics are shown in Fig. 1. It takes 2 A100-80G GPUs to load and forward the
original 8-expert model with bf16 data type. After pruning 2 and 4 experts, only one 80G GPU is needed for the inference process. For token generation speed analysis, during model inference, we still need to route each token to two experts. However, reducing the number of GPUs required to load the model can decrease the time consumed by GPU intercommunication, resulting in a much higher token generation speed. We observe a 1.20× token generation speedup for the model with 2 experts pruned and a 1.27× speedup with 4 experts pruned.

heading: 4.2 Expert Pruning for Domain-Specific Tasks
text: Experiment Setup. We investigate on mathematical reasoning tasks. We randomly sample sentences from the train set of MATH (Hendrycks et al., 2021) and combine them into 128 token sequences, each with a length of 2048. We carry out expert pruning with r = 6 and r = 4, then test 5-shot GSM8K (Cobbe et al., 2021) results.
Baselines to Compare. We compare the 5-shot GSM8K result with randomly pruned models, models pruned using samples from C4 as calibration data, the 2:4 structured model obtained by Wanda, as well as the original MoE model with 8 experts. Vanilla Wanda uses C4 for data calibration. To test the influence of the calibration dataset on model
performance, we also leverage the MATH dataset for Wanda pruning.
Evaluation Results. Tab. 3 illustrates the 5-shot evaluation results on the GSM8K dataset. The performance witnesses a significant drop after random pruning, or pruning with calibration data obtained from C4. However, this degradation dramatically reduces after leveraging the MATH dataset for calibration data construction. A similar phenomenon is also observed with Wanda. This indicates that when facing domain-specific tasks, using datasets corresponding to these specific tasks can yield better expert pruning results than using pre-training datasets. It also implies that our proposed method for changing the calibration dataset for domainspecific tasks can also be applied to other pruning algorithms. Besides, our expert pruning strategy (r = 4) significantly outperforms Wanda with 2:4 structured sparsity.
More Discussion. Using samples from the MATH dataset can greatly improve the domainspecific performance on mathematics tasks. However, as our method follows a post-training manner without any training, the expert pruning scheme still leads to a significant performance drop. To reduce the performance gap between pruned models and original models, we fully fine-tune the MoE models with different expert numbers on the MetaMathQA (Yu et al., 2023) dataset and compare their performances. We fine-tune and compare models pruned by C4 (r = 6), MATH (r = 6, r = 7), and the original 8-expert model. The zero-shot GSM8K@1 and MATH@1 accuracies after finetuning different MoE models are shown in Tab. 4.
As can be seen, the fine-tuning process significantly reduces the performance drop incurred by expert pruning and leads to comparable results with tuning full-expert models. Specifically, for the Mixtral 8x7B Instruct model, the accuracy of the pruned 7-expert model on the GSM8K test set exceeds that of the 8-expert model. This suggests that for certain practical downstream tasks, a large number of experts might not be a necessity for achieving good performance. Also, the pruned models using samples from MATH as the calibration dataset outperform those using C4 after tuning, further highlighting the effectiveness of adopting domain-specific calibration datasets for task-specific models.

heading: 4.3 Dynamic Expert Skipping Results
text: This subsection assesses the effectiveness of our proposed dynamic expert skipping approach. Additionally, we explore the combination of both methods to further improve inference efficiency.
Experiment Setup. We perform tests on taskagnostic models for better representativeness. The setup is similar to Sec. 4.1. We first prune the Mixtral 8x7B and Mixtral 8x7B Instruct model using calibration data sampled from the C4 dataset and get the pruned models with r = 6 and r = 4. During the testing of different benchmarks, we dynamically skip certain experts. For evaluation, we report the zero-shot accuracies of 8 tasks from EleutherAI LM Harness (Gao et al., 2023). Our proposed dynamic expert skipping method does not influence the memory usage for model inference, so we just report the inference speed in this subsection.
Baselines to Compare. We suggest that our proposed dynamic expert skipping can be seamlessly
integrated with the expert pruning approach. For setting up baselines, we evaluate zero-shot accuracies of the original 8-expert model without dynamic expert skipping, as well as the accuracies of models pruned to r = 6 and r = 4 without dynamic expert skipping. Subsequently, we incorporate dynamic expert skipping into the inference process of these models, evaluate accuracies over benchmarks, and measure token generation speedup.
Evaluation Results. The evaluation results are illustrated in Tab. 5. We show the average accuracies of the 8 zero-shot tasks at the “LM-eval” column, together with the inference speedup ratio compared to the original 8-expert models. As can be seen, based on expert pruning, the dynamic expert skipping method can further enhance the inference speed with just negligible performance drops. We can achieve nearly 90% performance of the Mixtral 8x7B Instruct model with half parameters and a 1.33× token generation speedup. Another notable observation is that the Mixtral 8x7B Instruct model using both expert pruning and dynamic skipping with r = 6 achieves the same inference speedup as the model using only expert pruning with r = 4 while getting a much higher accuracy over the LMeval benchmark. This phenomenon also proves the efficiency of our dynamic expert skipping approach. For a more comprehensive evaluation, we perform dynamic skipping on task-specific models and observe similar results. Experiment details are shown in the Appendix (Sec. A.4).

heading: 4.4 More Analysis
text: Discussion about Inference Speed. In Fig. 1, we also observe a notable inference acceleration when pruning experts from r = 6 to r = 4, a phenomenon not attributable to decreased interGPU communication overhead. We think that this enhancement stems primarily from improved temporal and spatial locality. This includes reduced cache misses, optimized memory prefetching, and faster block loading. Readers may refer to papers related to memory-intensive LLM inference (such
as (Alizadeh et al., 2023)) to find more explanation. More Ablations. We carry out efficient layerwise expert pruning in this work. Readers may be curious about the effectiveness compared with a layer-by-layer progressive searching paradigm, where the pruning of subsequent layers is aware of the pruning result of previous layers. To this end, we compare these two pruning paradigms in Tab. 6, leveraging the Mixtral 8x7B model and calibration samples from the C4 dataset. Although the layerby-layer progressive manner can get slightly better results with fewer experts pruned, when it comes to a high expert pruning rate (e.g., 50% pruning rate with r = 4), more performance drop is observed. We attribute this to the possible overfitting of the small calibration dataset.

heading: 5 Conclusion and Discussion
text: In this paper, based on the structural characteristics of MoE LLMs and the shortcomings of current weight pruning schemes, we focus on expert-level model sparsification and, for the first time, provide post-training expert pruning together with dynamic (expert) skipping methods to enhance the deployment efficiency of MoE LLMs. Our methods can significantly reduce memory usage and enhance inference speed while maintaining high model performance. Looking ahead, we aim to further refine our pruning/skipping techniques and incorporate them with weight pruning or parameter quantization strategies, achieving more effective deploying approaches for MoE LLMs.

heading: Limitations
text: Our method can reduce memory usage and improve inference speed for more efficient deployment of MoE LLMs. Despite its advancements, there are still some limitations. Firstly, our method for expert pruning is based on the enumeration of expert combinations. This is feasible for pruning currently popular MoE LLMs with 4 or 8 experts. However, with the number of experts in each MoE layer increasing (e.g., 32 experts in one MoE layer), it will be cumbersome to perform our pruning algorithm. Secondly, we conduct experiments on the opensourced Mixtral 8x7B and Mixtral 8x7B Instruct models as they are by far the most popular MoE LLMs. With the development of MoE LLMs, we will carry out experiments on other MoE LLMs in the future to give a more comprehensive analysis of the generalizability and scalability of our method.

heading: Ethics Statement
text: Our research focuses on improving the deployment efficiency of Mixture-of-Experts (MoE) large language models (LLMs) through expert-level sparsification techniques, aiming to reduce model sizes and enhance inference speed without compromising performance. While our methods offer potential benefits for deploying advanced LLMs more broadly and efficiently, we acknowledge the importance of considering the ethical implications of deploying such models. These include ensuring the responsible use of LLMs, mitigating biases in model outputs, and addressing privacy concerns. We commit to making our code available for transparency and encourage the community to use our findings responsibly, considering the societal impacts of deploying LLMs.

heading: Acknowledgement
text: This project is funded in part by National Key R&D Program of China Project 2022ZD0161100, by the Centre for Perceptual and Interactive Intelligence (CPII) Ltd under the Innovation and Technology Commission (ITC)’s InnoHK, by General Research Fund of Hong Kong RGC Project 14204021. Hongsheng Li is a PI of CPII under the InnoHK.

heading: A Appendix

heading: A.1 Expert Selection Tendency in MoE Models
text: To investigate the tendency of expert selection during model inference, we analyze the Mixtral 8x7B model using samples from the C4 dataset (Raffel et al., 2019) (Fig. 5 (a)) and the MATH dataset (Hendrycks et al., 2021) (Fig. 5 (b)), respectively. C4 is a pre-training dataset representing the general relationships between experts, MATH is a dataset designed for a specific downstream task. We visualize the probability of different experts being selected in layer 0, layer 15, and layer 31 of the model during inference. As top-2 experts are chosen by default, each grid in the plot represents the frequency of two experts being selected simultaneously during the forward pass over the sample set (the x-axis and the y-axis of the plot each represent one expert). As can be seen from Fig. 5, the model exhibits a certain tendency in the selection of experts, particularly when tailored to specific downstream tasks.
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7
Layer 0
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7
Layer 15
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7
Layer 31
0.000
0.025
0.050
0.075
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7
Layer 0
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7
Layer 15
0 1 2 3 4 5 6 7
0 1 2 3 4 5 6 7
Layer 31
0.00
0.03
0.06
0.09
(a) (b)
Figure 5: Frequency visualization of expert selection in layer 0, layer 15, and layer 31 for the Mixtral 8x7B model on samples of (a) C4 (Raffel et al., 2019) and (b) MATH (Hendrycks et al., 2021) dataset respectively. The model exhibits certain preferences in the selection of experts.

heading: A.2 Theoretical Insight and Broader Application of Dynamic Skipping
text: Consider a single MoE layer. Suppose in a top-k setting with a total of n experts, for every single token x, the routing weights of the n experts are w1, . . . , wn, ∑n i=1wi = 1. Without loss of generality, we assume w1 ≥ w2 ≥ · · · ≥ wk ≥ · · · ≥ wn, and the output feature vector of each expert are f1, . . . , fk, . . . , fn, where fi = Ei(x).
Without dynamic skipping, the original output
of this layer should be:
z = 1
∑k m=1wm
k∑
m=1
wmfm.
Suppose after dynamic skipping, only top-i experts remain (1 ≤ i ≤ k). Then the output is:
ẑ = 1∑i
m=1wm
i∑
m=1
wmfm.
As defined above in the main paper, the reconstruction loss can be calculated by:
L = ||ẑ − z||2. Through experiment, we observe that the distribution of ||f1− f2||2 is fairly concentrated for Mixtral 8x7b on the C4 calibration set. Thus for simplification, we assume ||fm − fn||2(m ̸= n) to be a fixed value, say D. Therefore,
L = ||ẑ − z||2
= || 1∑i m=1wm
i∑
m=1
wmfm − 1
∑k m=1wm
k∑
m=1
wmfm||2
= ||∑kn=1 ∑i m=1wnwmfm − ∑i n=1 ∑k m=1wnwmfm||2
( ∑i m=1wm)( ∑k m=1wm)
= ||∑kn=1 ∑i m=1wnwmfm − ∑k n=1 ∑i m=1wnwmfn||2
( ∑i m=1wm)( ∑k m=1wm)
= ||∑kn=i+1 ∑i m=1wnwm(fm − fn)||2
( ∑i m=1wm)( ∑k m=1wm)
≤ || ∑k n=i+1 ∑i m=1wnwm||2 ·D
( ∑i m=1wm)( ∑k m=1wm)
(= holds if all (fm − fn) are of the same direction)
= ∑k m=i+1wm∑k m=1wm D.
We set an upper bound H (H ≤ D) on L (make L ≤ H) to trade-off between accuracy and inference speed, then have
∑k m=i+1wm∑k m=1wm D ≤ H.
Let β = HD , we have:
k∑
m=i+1
wm ≤ β · k∑
m=1
wm.
Therefore, in this generalized setting, dynamic skipping should reserve the top-i∗ experts where
i∗ = min i
s.t. k∑
m=i+1
wm ≤ β · k∑
m=1
wm.
Specifically, in the top-2 setting, for computational simplicity (reduce the use of additions), let β =
H D−H (a little abuse of notation), then the dynamic skipping criteria for top-2 is
w2 ≤ βw1.

heading: A.3 Experiments on the Sizes of Calibration Datasets
text: We prune the Mixtral 8x7b model with different sizes of calibration datasets. To be specific, we randomly sample 1, 2, 4, 16, 64, and 256 sequences (each composed of 2048 tokens) from the C4 dataset to form calibration datasets. Then the model is pruned to r = 6 and evaluated on various benchmarks. The average LM-eval results are reported in Tab. 7.
As can be seen, using 64 and 128 sequences can result in the highest overall results (≥ 64). Using a small set of sequences will possibly lead to performance degradation (especially for using just 1 sequence), but our method is somewhat robust to the size of datasets, as seen from the table.

heading: A.4 Dynamic Skipping for Domain-specific Tasks
text: We also perform dynamic expert skipping on domain-specific tasks (mathematical reasoning tasks). We calibrate β for each layer using samples from the training set of MATH and evaluate 5-shot accuracy on the GSM8K dataset. We also test and report the token generation speed of each MoE model. The results are shown in Tab. 8. In this case, dynamic expert skipping leads to more performance drops. But for the Mixtral 8x7B Instruct model, expert pruning with 2 experts and combining dynamic skipping also leads to the same inference speedup with pruning 4 experts, while achieving higher evaluation accuracy.

heading: A.5 Actual Memory Reduction
text: A more detailed statistical comparison between our expert pruning method with baseline methods on
the Mixtral 8x7B model is shown in Tab. 9.

heading: Method Sparsity Memory (MB)

heading: A.6 Relationships with Other Network Pruning and Parameter Quantization Methods
text: As plug-and-play techniques, both our proposed expert pruning and dynamic skipping methods are orthogonal to other model light-weighting schemes (e.g., weight pruning (Frantar and Alistarh, 2023; Sun et al., 2023), token pruning (Kim et al., 2022; Ding et al., 2023)) and are compatible with weight quantization approaches (Frantar et al., 2022; Lin et al., 2023).

heading: A.7 More Experiment Details
text: In this part, we give more experimental details for a better understanding of our proposed methods.
Calibration Set Construction for Expert Pruning. For task-agnostic models, we use the samples from C4 (Raffel et al., 2019) as the calibration dataset. Following the setting of Wanda, we sample from the first part of the training data6. For task-specific (mathematics) models, we use samples from the training set of MATH (Hendrycks et al., 2021). The structure of the MATH dataset is different from C4, so we reconstruct the dataset in the format of C4 and randomly sample from it.
6https://huggingface.co/datasets/allenai/c4/ blob/main/en/c4-train.00000-of-01024.json.gz
Calibration Set Construction for Dynamic (Expert) Skipping. To calculate β for dynamic expert skipping in each MoE layer, we forward the MoE model over the calibration dataset and set β as the median value of we1we0 separately for each layer. We choose to use the median value over the calibration dataset as in this case, the skipping will happen with around 50% possibility. Here we provide the value of β for the Mixtral 8x7B model with calibration data sampled from C4 and MATH respectively. As can be seen, the parameter in each layer differs significantly.
C4: 0.402,0.494,0.463,0.484,0.478,0.491,0.523, 0.521,0.544,0.570,0.574,0.489,0.503,0.618,0.568, 0.535,0.559,0.519,0.537,0.487,0.469,0.461,0.461, 0.469,0.458,0.418,0.433,0.418,0.406,0.433,0.447, 0.535
MATH: 0.503,0.586,0.505,0.531,0.509,0.422, 0.511,0.461,0.447,0.478,0.529,0.454,0.472,0.531, 0.499,0.486,0.503,0.491,0.430,0.440,0.402,0.423, 0.386,0.407,0.395,0.354,0.340,0.351,0.334,0.368, 0.365,0.346
Model Fine-tuning. In the part of task-specific expert pruning for domain-specific tasks, we finetune the Mixtral 8x7B and Mixtral 8x7B Instruct models with 8 experts, 7 experts, and 6 experts on the MetaMathQA (Yu et al., 2023) dataset. The training is conducted on 16 A100-80G GPUs. We train the model for 900 steps, using a learning rate of 2e-5 with the cosine learning rate scheduler.

heading: Future Work
text: With the development of MoE LLMs, we will carry out experiments on other MoE LLMs in the future to give a more comprehensive analysis of the generalizability and scalability of our method