MoE-Pruner: Pruning Mixture-of-Experts Large Language Model using the Hints from Its Router

(2) The distilled step shows only minor improvements, as seen in Table 4. What could be the underlying reason for this limited gain?

As shown in Section 4.2, our distillation step improves the pruned student model performance with minimal computational budget. We only use 1000 samples from the C4 dataset and the training can be done in less than one hour. We followed Wanda’s [4] idea which enforces a limited computational budget during the fine-tuning stage to be computationally efficient while restoring pruned model performance effectively. Considering the limited computational budget, we are able to restore the performance of pruned MoE models by a non-trivial amount.

(3) This method builds incrementally on the baseline, Wanda, by incorporating gate weights within the MoE pruning framework.

Our method builds upon Wanda and extends it to MoE LLMs while maintaining the simplicity of Wanda's pruning metric. A key insight is the consideration of router weights ($Gate$) in MoE architectures. Consider a simple Mixture-of-Experts with two experts and each with only one weight: $y=Gate_1 \cdot w_1\cdot x + Gate_2 \cdot w_2 \cdot x$. If $\vert w_1 \vert \leq \vert w_2 \vert$, traditional pruning methods would remove $w_1$. However, in MoE architectures:

• If $Gate_1 \approx 1$ and $Gate_2 = 0$, $\vert Gate_1 \cdot w_1 \cdot x \vert \gg \vert Gate_2 \cdot w_2 \cdot x \vert$, making $w_2$ the correct choice for removal.

This motivating example shows that previous pruning methods for LLMs do not consider the importance of router weights which only exist in MoE architecture and may result in lower performance after pruning MoE. Therefore, we propose MoE-Pruner designed explicitly for MoE LLMs to handle such a limitation while maintaining the simplicity of Wanda's pruning metric.

Thank you for your suggestion, we have added a detailed motivation and analysis of our method in the revised version, section 3.2, line 212.

Some previous studies have focused on pruning entire experts, as seen in [1, 2, 3]. Could this method potentially be integrated with these approaches to enhance overall pruning effectiveness?

Yes, MoE-Pruner is compatible with expert-pruning methods and can complement them effectively:

1. Integration with NAEE [1]:

• NAEE minimizes reconstruction loss through an enumerative expert pruning approach. MoE-Pruner could enhance NAEE by using its router weight-aware metric to guide pruning.
• NAEE is limited to architectures with few experts (e.g., Mixtral-8x7B and MiniCPM-8x2B). It cannot process DeepSeek-V2-Lite and Qwen1.5-MoE with 64 or more experts.

2. Expert Merging (e.g., [2]):

• Expert merging approaches like [2] focus on merging similar experts using a similarity matrix. MoE-Pruner’s router weights could guide the merging process by identifying underutilized experts.

3. Integration with STUN [3]:

• STUN reduces pruning complexity for models with many experts. MoE-Pruner could be combined with STUN’s structured-then-unstructured pruning for better scalability.

These integrations show that MoE-Pruner could enhance pruning efficiency and scalability.

If you have other questions, we are happy to answer.

References:

[1]. Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, Hongsheng Li. Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models. ACL 2024.

[2]. Zeliang Zhang, Xiaodong Liu, Hao Cheng, Chenliang Xu, Jianfeng Gao. Diversifying the Expert Knowledge for Task-Agnostic Pruning in Sparse Mixture-of-Experts. arXiv:2407.09590.

[3]. Jaeseong Lee, Seung-won Hwang, Aurick Qiao, Daniel F Campos, Zhewei Yao, Yuxiong He. STUN: Structured-Then-Unstructured Pruning for Scalable MoE Pruning. arXiv:2409.06211.

[4]. Mingjie Sun, Zhuang Liu, Anna Bair, J. Zico Kolter. A Simple and Effective Pruning Approach for Large Language Models. ICLR 2024.

Complexity analysis - The paper does not provide a complexity analysis showing the pruning and inference speedup.

We conduct an efficiency evaluation experiment on memory reduction and inference speedup using the A100 GPU. To provide a comprehensive comparison, we also evaluated the NAEE [1] expert pruning method. Comparison results are shown in the table below. Our proposed MoE-Pruner method at the structured 2:4 sparsity pattern outperforms NAEE in terms of both average performance and inference speedup, while incurring only a small memory overhead for storing sparse tensor indices. The original average performance and memory of Mixtral-8x7B are 69.16 and 87.49GB, respectively. Speedup is measured according to the original Mixtral 8x7B baseline. The following table has been added to the Table 3 efficiency evaluation experiments in our revised version:

Poor presentation - The paper is poorly organized. Section 2 (Preliminaries) is too short while Sections 3.1 and 3.2 are essentially a literature review, serving the same purpose. Figure 1 does not clearly explain the method as the score S does not interact with other components in the figure.

We have restructured the manuscript as follows:

• Section 2 (Preliminaries): Expanded to include details on MoE initialization methods, including sparse upcycling initialization [2], and its relevance to MoE pruning.
• Section 3 (Methodology): Streamlined to focus on MoE-Pruner’s formulation and motivation. We included a detailed motivating example that demonstrates how router weights influence pruning decisions in MoE architectures.
• Figure 1: Removed to avoid redundancy. Instead, we emphasize how the pruning metric ($S_{ij}$) directly interacts with other components.

What is pruning and inference speedup of MoE-Pruner compared to other baselines?

Please refer to our response to weakness 3 for the efficiency evaluation experiment about memory reduction and inference speedup.

How did the authors extend Wanda to MoE models, is it equivalent to setting Gate_j=1 in MoE-Pruner?

Yes, if we set Gate value for each expert to be 1 in MoE-Pruner, then our method becomes the same as Wanda.

The method is poorly motivated (only from L232-237). A deeper analysis (for example, Section 3, Wanda) is preferred.

Our method builds upon Wanda and extends it to MoE LLMs while maintaining the simplicity of Wanda's pruning metric. A key insight is the consideration of router weights ($Gate$) in MoE architectures. Consider a simple Mixture-of-Experts with two experts and each with only one weight: $y=Gate_1 \cdot w_1 \cdot x + Gate_2 \cdot w_2 \cdot x$. If $\vert w_1 \vert \leq \vert w_2 \vert$, traditional pruning methods would remove $w_1$. However, in MoE architectures:

• If $Gate_1 \approx 1$ and $Gate_2 = 0$, $\vert Gate_1 \cdot w_1 \cdot x \vert \gg \vert Gate_2 \cdot w_2 \cdot x \vert$, making $w_2$ the correct choice for removal.

This motivating example shows that previous pruning methods for LLMs do not consider the importance of router weights which only exist in MoE architecture and may result in lower performance after pruning MoE. Therefore, we propose MoE-Pruner designed explicitly for MoE LLMs to handle such a limitation while maintaining the simplicity of Wanda's pruning metric.

Thank you for your suggestion, we have added a detailed motivation and analysis of our method in the revised version, section 3.2, line 212.

If you have other questions, we are happy to answer.

References

[1]. Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, Hongsheng Li. Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models. ACL 2024.

[2] Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, Neil Houlsby. Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints. ICLR 2023.

Response 3: I would still recommend the author to explicitly clarify how the observed "differences between two MoE initialization methods" influence the pruning method's design. Otherwise, the two main contents of the paper seem a little disconnected in their current form.

Thank you for this insightful observation. In our experiments, we observed that router weights (Gate values) differ significantly based on the MoE initialization method:

• The router weights (Gate) in an upcycled model (e.g. Mixtral-8x7B) are not fully optimized in the MoE because they do not exist in the dense checkpoint and are only added after sparse upcycling.
• The router weights in MoE model trained from scratch (e.g. DeepSeek-V2-Lite) are fully optimized during training, providing more reliable indicators of expert utilization and importance.

By explicitly considering how the MoE initialization method affects the distribution and reliability of router and expert weights, we design a universal pruning metric that is adaptive to these differences:

• For sparse upcycled models, the pruning method compensates for less reliable router weights by focusing on the inherent importance of expert weights and prevents the potential loss of valuable weights that might be mistakenly pruned due to under-optimized router weights.
• For models trained from scratch, the pruning method takes full advantage of reliable router weights to make more informed pruning decisions and results in a more effective pruning process that maintains model performance by preserving critical experts and weights.

We acknowledge that our paper did not sufficiently clarify this connection, leading to a disconnection between the MoE initialization methods and the pruning metric design. We will revise section 3.2 to explicitly discuss how MoE initialization methods impact router weight reliability and influence pruning design and add experimental results demonstrating how our pruning method performs on models initialized via both methods, highlighting the adaptations made for each case.

Thank you for looking into our rebuttal and providing insightful suggestions to improve our paper quality. Here are our response to your followup questions:

Response 1: The reported speedup (1.14x) seems minor given the significant performance drop (5%). Additionally, the benchmark used in the table is unclear. Is such a performance drop acceptable in this context?

The performance drop reported (from 69.16 to 64.58) reflects the results after one-shot structured pruning without fine-tuning. To address this, our paper proposes expert-wise knowledge distillation to recover pruned model performance using the pretrained MoE model with a limited compute budget. After applying this distillation (using 1,000 samples and training in less than one hour), the pruned model achieves a significantly improved average performance of 67.07, narrowing the gap to the original model.

The 1.14× speedup was based on PyTorch 2.4.1, which only supports the CUTLASS backend for sparse tensor computations. With PyTorch 2.5.0 and the newly supported cuSPARSELt backend, the speedup improves to 1.31×, offering a more substantial inference efficiency gain. We provide the updated table reflecting these results:

The benchmark is the same as the average zero-shot performance on 9 evaluation tasks we used in the paper: ARC-challenge, ARC-easy, Boolq, HellaSwag, MMLU, OpenBookQA (OBQA), PIQA, RTE, and WinoGrande. We also list the zero-shot performance below:

To summarize, the performance drop after pruning (69.16 to 64.58) can be mitigated by our expert-wise knowledge distillation, achieving 67.07 performance. With additional data and compute resources, we anticipate further narrowing the gap to the pretrained model’s performance.

(Continued response to weakness 3)

Compared with the expert pruning method NAEE [2], our method shows much higher average performance across all four MoE models at 50% structured sparsity. Especially if you look at the MMLU performance, which is the most comprehensive benchmark we used, expert pruning results in a very large performance degradation (Mixtral-8x7B 67.03 to 47.68, MiniCPM-8x2B 56.23 to 23.39), while our method shows a moderate loss (Mixtral-8x7B 67.03 to 62.17, MiniCPM-8x2B 56.23 to 42.42).

1. In line 154, the authors mentioned that upcycling initialization will lead to higher expert similarity in MoE models. Since the similarity between experts is relatively high, pruning experts should theoretically not result in a significant performance drop. Why do the authors reach the opposite conclusion?

This is a very good and interesting question. Sparse upcycled MoE models start from the same expert weight and have a higher similarity between experts. Pruning experts seem to have a lower performance drop. However, as demonstrated in our experiments (Table above), expert pruning results in severe degradation (e.g., Mixtral-8x7B MMLU: 67.03 → 47.68), while our 2:4 structured pruning exhibits a much smaller loss (67.03 → 62.17).

The reason is that sparse upcycled MoE models usually use a smaller number of experts (e.g. 8 experts) and have a more uniform expert activation frequency. If we consider a top-2 activation method with 8 experts, then each expert will process 25% of the input tokens. Pruning half of the expert will result in severe performance drop. Our 2:4 structured pruning tries to preserve each expert performance while maintaining a 50% structured sparsity. Experimental results show that our methods achieve better performance on sparse upcycled MoE models compared with expert pruning method.

2. In line 158, the authors mentioned that train from scratch will show imbalanced expert activation frequency, indicating that least-used expert pruning could help compress model size and not bring performance degradation. However, I believe this may be task-dependent; these seemingly least-used experts could be very useful for specific tasks such as math or coding. Did the authors conduct any related validation?

Thank you for your insightful suggestion. For a model trained from scratch, we use DeepSeek-V2-Lite and test its performance on a representative math datasets, GSM8K, with NAEE [2] and MoE-Pruner:

Results demonstrate that expert pruning (NAEE) causes a drastic performance drop (GSM8K: 38.51 → 0.99), while MoE-Pruner retains better performance (38.51 → 9.17). We use random expert pruning as NAEE [2] uses an enumerate method, so it only applies to the MoE model with a smaller number of experts, e.g. 8, but not 64. This large performance drop is also proven in the NAEE [2] paper Table 3, suggesting that expert pruning may be better for task-dependent setting or followup fine-tuning scenarios. We will incorporate this finding into our draft.

If you have other questions, we are happy to answer.

References:

[1]. Enshu Liu, Junyi Zhu, Zinan Lin, Xuefei Ning, Matthew B. Blaschko, Shengen Yan, Guohao Dai, Huazhong Yang, Yu Wang. Efficient Expert Pruning for Sparse Mixture-of-Experts Language Models: Enhancing Performance and Reducing Inference Costs. arXiv:2407.00945.

[2]. Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, Hongsheng Li. Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models. ACL 2024.

[3]. Alexandre Muzio, Alex Sun, Churan He. SEER-MoE: Sparse Expert Efficiency through Regularization for Mixture-of-Experts. arXiv:2404.05089.