Efficient Expert Pruning for Sparse Mixture-of-Experts Language Models: Enhancing Performance and Reducing Inference Costs

Authors: Enshu Liu, Junyi Zhu, Zinan Lin, Xuefei Ning, Matthew B. Blaschko, Shengen Yan, Guohao Dai, Huazhong Yang, Yu Wang
Primary area: foundation or frontier models, including LLMs
Keywords: Large language model, mixture of experts, pruning

Abstract:
The rapid advancement of large language models (LLMs) has led to architectures with billions to trillions of parameters, posing significant deployment challenges due to their substantial demands on memory, processing power, and energy consumption. Sparse Mixture-of-Experts (SMoE) architectures have emerged as a solution, activating only a subset of parameters per token, thereby achieving faster inference while maintaining performance. However, SMoE models still face limitations in broader deployment due to their large parameter counts and significant GPU memory requirements. 
In this work, we introduce a gradient-free evolutionary strategy named Efficient Expert Pruning (EEP) to enhance the pruning of experts in SMoE models. Specifically, EEP searches the pruning pattern and use expert merging as an memory-efficient way of fine-tuning the pruned model. EEP relies solely on model inference (i.e., no gradient computation) and achieves greater sparsity while maintaining or even improving performance on downstream tasks. EEP can be used to reduce both the total number of experts (thus saving GPU memory) and the number of active experts (thus accelerating inference).
For example, in the task-specific setting, we demonstrate that pruning up to 75\% of experts in Mixtral $8\times7$B-Instruct results in a substantial reduction in parameters with minimal performance loss, or pruning 50\% of experts and activating one fewer expert to achieve 1.41$\times$ speedup. Our experiments include four different model sizes from Mixtral, Qwen1.5 and Qwen2, and utilize more than 10 datasets as well as various settings. Results show that our method outperforms the related baselines by a large margin, demonstrating a significant advancement in this direction. Results of our method can be reproduced using the code provided in the supplementary material.

Decision: reject
Reviewer scores: 6, 6, 3, 6

Reviews and discussion:
<thread>
creation date:
1730313363617

thread_type:
official_review

<thread_object>
creator: reviewer_n5tl

<content>
summary:
This paper proposes Efficient Expert Pruning (EEP), a gradient-free evolutionary strategy to optimize the pruning and merging of experts within Sparse Mixture-of-Experts (SMoE) models. EEP aims to enhance performance while saving GPU memory and reducing inference costs by minimizing both the total number of experts and the number of active experts during inference. This is achieved through an evolutionary search process that optimizes two matrices: a router mapping matrix for expert selection and an expert merging matrix to retain pruned experts' knowledge. EEP is evaluated across various downstream tasks and shows considerable improvement over baselines in terms of memory efficiency, inference speed, and overall performance.

strengths:
1. This paper is well-written, well-organized including problem formulation, methodological exposition, and experimental setup, making it accessible to readers with a background in the field. Besides, the implementation details and code are provided for reproducing.
2. The EEP method employs a gradient-free evolutionary search strategy for expert pruning and merging in the SMoE model, which is a relatively novel method.
3. This paper conducts extensive experiments on multiple datasets and models to validate the effectiveness of the method, demonstrating significant improvements over existing baselines, which greatly enhances the practicality of SMoE pruning.

weaknesses:
1. In the NAEE [1] method, the authors use the C4 dataset as the calibration dataset and evaluate the zero-shot accuracy of the pruned model, which is a common experimental setup in the field of LLM compression. The authors should conduct similar experiments and provide data in this regard to allow for a more comprehensive assessment of the effectiveness of the EEP method.
2. I am concerned about the search cost of EEP. The total time spent on EEP for each dataset is several hours, and this is just for the smallest Mixtral 8x7B Instruct model. For larger datasets and larger models, EEP will take even longer, which is often unacceptable.
3. The out-of-distribution accuracy of EEP is much lower than that of NAEE [1] on many datasets (Tables 5, 12 and 13), so I am concerned about the generalization ability of EEP.
4. Some notations are confusing.
    - The authors use bold uppercase symbols to represent matrices. In line 91, $W_{RM}$ and $W_{EM}$ should be bolded to ensure consistency with the notation used in Section 4.2.
    - The symbol $\boldsymbol{W}_G$ in Equation 2 is not explained. I suspect it refers to the router network $\boldsymbol{W}_R$ in line 179.

[1] Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models. ACL 2024 Main.

questions:
1. In Section 5.1, the authors state that the dataset is randomly divided into training and test set, with the training set used for evolutionary search and the test set used for evaluation. How exactly is this division performed? Table 14 seems to contain some data on the trainset size, but it is not comprehensive enough. 
2. The authors used 50 MMLU datasets as the training set and 7 as the test set, utilizing a large number of datasets as the training set to test accuracy on out-of-distribution datasets. I would like to know the total search time EEP took in this case, and whether using fewer training sets would impact the out-of-distribution accuracy.
3. How does the accuracy of the EEP pruned model obtained from one dataset perform on other datasets? For example, if BoolQ is used as the training set, what is the accuracy of the pruned model on other validation sets?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730471167546

thread_type:
official_review

<thread_object>
creator: reviewer_osj6

<content>
summary:
This paper addresses the challenge of pruning in sparse mixture of experts (SMoE) models. The authors propose a gradient-free evolutionary strategy called Efficient Expert Pruning (EEP) to improve the pruning process for experts within these models. Experimental results on diverse models and datasets demonstrate the effectiveness of this approach.

strengths:
(1) The writing is clear and easy to follow.  

(2) The proposed method is interesting, suggesting that evolutionary strategies can effectively optimize parameter combinations for model pruning.   

(3) The insights presented are valuable, demonstrating that MoE models contain redundancy and hold significant potential for pruning while maintaining performance and improving efficiency.

weaknesses:
(1) Lack of Reliable Evaluations: I find it puzzling that Mixtral-8x22B's average performance improves despite being pruned from 8 experts down to 2. This raises questions about the robustness of the evaluation methods, especially given the relatively small datasets used in this study. In contrast, NAEE—a key baseline—evaluates on the comprehensive MMLU dataset, which could serve as a more reliable benchmark.

(2) In Table 6, even after pruning from 8 experts down to 4 and then 2, the speedup is small. What factors might explain this limited increase in speed?

(3) The method’s efficiency is largely limited by its computational demands and the dataset used. The pruning dataset is relatively small, yet the process still requires significant time—up to several hours.

(4) Are the training and evaluation samples drawn from the same distribution? This could make sense as a task-agnostic pruning approach. However, given the long pruning time even with limited samples, why not fine-tune a roughly pruned model instead?

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730474314331

thread_type:
official_review

<thread_object>
creator: reviewer_mrrf

<content>
summary:
This paper introduces Efficient Expert Pruning (EEP), a gradient-free evolutionary strategy for optimizing Sparse Mixture-of-Experts (SMoE) language models. The work addresses the challenge of reducing model size and computation costs while maintaining performance. The authors propose a two-phase approach combining expert pruning and expert merging, which can be executed using only model inference capabilities without requiring gradient computation or extensive GPU resources. The method is implemented through router mapping and expert merging matrices, enabling evolutionary search for optimal pruning configurations.

strengths:
1. The paper is exceptionally well-written and organized, with comprehensive experimental validation across multiple model scales and diverse datasets. The methodology is clearly presented with thorough ablation studies and reproducible code.
2. The proposed two-phase EEP framework presents a novel and effective paradigm for model compression. The authors provide insightful analysis of router network behavior that justifies their approach.
3. The method achieves remarkable results while only requiring inference capabilities, with pruned models sometimes even outperforming the original uncompressed models.

weaknesses:
1. The analysis of why performance sometimes improves after pruning in Sec. 5.6. remains somewhat speculative, lacking rigorous theoretical foundations and comprehensive empirical validation.
2. The performance gap between global setting and task-specific setting in OOD experiments suggests that EEP's impressive results might be more dependent on task-specific optimization than claimed, potentially limiting its general applicability across different tasks.
3. The search process of EEP is computationally expensive and time-consuming, particularly for large datasets. The paper does not address this efficiency bottleneck.

questions:
Could the authors elaborate on how the choice of group numbers and search iterations affects the final performance? A more systematic analysis of these hyperparameters would be valuable for implementation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730882929079

thread_type:
official_review

<thread_object>
creator: reviewer_qte1

<content>
summary:
This paper works on the topic of pruning experts in SMoE models. They propose a gradient-free evolutionary strategy named EEP. In EEP, there are two matrices named expert merging matrix and experts mapping matrix, used to search for pruning configurations. The goal is to achieve sparsity while reaching good performance on downstream tasks. The paper shows the performance of the proposed method on four different model sizes and uses more than 10 datasets.

strengths:
SMoE can achieve faster inference while maintaining performance. To search for the subset of experts, this paper proposes a gradient-free approach based on evolutionary strategy. By using two phases expert pruning and expert merging, the method achieves high sparsity.

weaknesses:
1. While the method is gradient-free, it does not necessarily imply a low search cost. The evolutionary strategy used requires multiple iterations and epochs, and the search cost is highly dependent on the settings for iterations, epochs, and the data size used for calculating the validation metric. To better illustrate the efficiency of the method, I suggest the authors provide concrete runtime numbers for different dataset sizes, comparing them to gradient-based alternatives. This would help clarify whether the method can be effectively extended to more complex datasets.

2. The method heavily relies on evaluation results from selected data to conduct the search process. While Table 15 shows in-distribution results for two dataset sizes, it is unclear whether similar trends hold for out-of-distribution cases. A comprehensive study on the influence of dataset size is recommended. Specifically, it would be beneficial to include a plot showing the performance vs. dataset size across a range of sizes for both in-distribution and out-of-distribution settings. Additionally, the authors should explain the counterintuitive result where a smaller dataset size provides better performance than a larger dataset size.

3. For out-of-distribution evaluation, 7 datasets from MMLU are used for validation. It would be helpful to know whether the number of validation datasets affects performance. To address this, I recommend conducting an ablation study that varies the number of validation datasets (e.g., from 1 to 7) and reporting how this impacts performance. This would offer insight into the method’s robustness in extreme cases where only a single dataset is available for validation.

4. The writing of the paper could be improved, as the key ideas are not easy to follow. A significant portion of the text is devoted to background and preliminaries, such as the description of the self-attention mechanism in Section 3, which is well-known and could be condensed or moved to the appendix. This would create more space to emphasize the main algorithm and core contributions in the main paper. I suggest reorganizing the structure to focus on the core algorithm earlier in the text for better clarity and accessibility.

questions:
1. Why is the full model performance significantly higher in this paper compared to the NAEE paper? For example, in Table 2, the full model (Mixtral 8×7B-Instruct) achieves 85.8 on ARC-c with 8 experts, while in the NAEE paper's Table 1, the performance is only 62.20.

2. Can the proposed method still work for out-of-distribution cases when there is only one dataset available for validation?

3. What does the performance vs. dataset size curve look like during the search process? What is the rationale for choosing specific dataset sizes for validation?

4. How is the performance of the proposed method compared to dynamic MoE methods?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732700448629

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to all Reviewers

comment:
We sincerely appreciate the reviewers' detailed and helpful suggestions on our work. We are excited to see that the reviewers think the paper is exceptionally well-written and organized (mRrF, oSj6, N5TL), the experiments are comprehensive (mRrF, N5TL), the approach is novel, well-justified, and interesting (mRrF, oSj6, N5TL), the analysis on router behavior is insightful (mRrF, oSj6), and the results are remarkable with significant improvements (mRrF, N5TL).

Following the reviewers' suggestions, we conducted additional 23 task-specific (in-distribution) search and evaluation experiments, 2 task-agnostic (out-of-distribution) search and more than 100 OOD evaluation experiments, as well as more computational cost profiling results, detailed in our rebuttals to each reviewer. We have also updated the revised manuscript (changes highlighted in red) based on reviewers' feedback.

Here, we want to highlight several points about our work.

- Firstly, we want to clarify a common misunderstanding of our work that *the search cost of EEP is high*.

    -  Although we did not focus on optimizing the search cost when designing EEP, **we have shown in our original submission that EEP's search cost can be easily reduced by using smaller training sets** (Table 15). To further support this, we added more experiments showing EEP's performance and search cost under 200, 100, and 50 data points (see the rebuttal and revision for details). The results show that **EEP can complete the pruning process in less than 20 minutes while achieving better performance than the full model and the baselines**. This computation time is comparable to the baseline method NAEE. 
    
    -  Additionally, our computation time for EEP includes both the GPU time (model inference) and CPU time (model merging, which is the main overhead for small dataset size like 50). This makes the comparison somewhat unfair to EEP since GPU time is typically more expensive than CPU time. With further optimizations, such as performing model merging on GPUs, we expect EEP's computation time to be even faster.

- Secondly, we want to re-emphasize the key features of EEP that set it apart from baseline methods:

    - **Improving inference efficiency**: Through extensive experiments, we demonstrate that **EEP significantly improves the inference efficiency of MoE models by reducing both GPU memory consumption (by pruning the total number of experts) and inference latency (by pruning the number of active experts)**. Since deployment computational costs are often much higher than training costs (e.g., see [this report](https://www.datacenterdynamics.com/en/news/openai-training-and-inference-costs-could-reach-7bn-for-2024-ai-startup-set-to-lose-5bn-report/)), these benefits provided by EEP have substantial practical implications.
    
    - **Gradient-free:** Since EEP is a gradient-free method, **EEP requires much less GPU memory than gradient-based methods** (e.g., fine-tuning), and **it can run on devices available only for inference, expanding its applicability to more use cases.**
    
    - **Robustness in both in-distribution and out-of-distribution settings**: Extensive experiments demonstrate that EEP performs well regardless of whether the search and test datasets are the same or different. We notice that the reviewers are more interested in the out-of-distribution settings. But we want to mention that the **in-distribution (i.e., task- or domain-specific) setting where EEP outperforms the baselines significantly also demonstrates important applications**. Task-specific LLMs fine-tuned from general capability LLMs, such as Code LLama, BioMistral, and Llama Guard, are widely used in many applications. Recent papers [1] also argue that composing different task-specific models/modules rather than using one general-purpose model and expecting it to perform well on all tasks might be better in terms of scalability, efficiency, and so on. **Using EEP to adapt a general-purpose MoE model towards one specific downstream task while making it more efficient aligns with these trends.**

We hope that our rebuttal and revision effectively address the reviewers' concerns. If there are any follow-up questions, we would be glad to discuss them further.

[1] Xiao, Chaojun, et al. "Configurable foundation models: Building llms from a modular perspective." arXiv preprint arXiv:2409.02877 (2024).
</content>
</thread_object>

</thread>


Metareview:
metareview:
The deployment of Sparse Mixture-of-Experts models remains challenging due to their large parameter counts and GPU memory demands. To address this, this paper introduces Efficient Expert Pruning, a gradient-free evolutionary strategy that prunes and merges experts in SMoE models based solely on model inference, reducing parameters and memory usage while maintaining or improving performance. Experiments on multiple datasets and models, including Mixtral and Qwen, show that EEP achieves sparsity and inference speedups, outperforming baselines and enabling more efficient SMoE deployment.

Three reviewers have slightly positive feedback and most reviewers acknowledge the paper presentation & methods' efficacy, especially impressed by the high sparsity of compressing Mixtral MoE. However, it is well-known that Mixtral MoE is highly redundant, also for Qwen MoE. My major concern is the generalization ability of the proposed approach to other types of MoEs which have less redundancy due to improved optimization and designs. For example, deepseek and modularformer. Also, the discussions with related MoE merging papers are encouraged to be included. Overall, we recommend a rejection.