Task-Specific Expert Pruning for Sparse Mixture-of-Experts

TL;DR

The paper tackles the practicality of large sparse MoE models by turning them into efficient single-expert dense models tailored to downstream tasks. It introduces a one-pass fine-tuning framework that progressively drops non-professional experts based on a task-specific proficiency criterion, culminating in a single expert per MoE layer by mid-training. Empirical results show the single-expert eager pruning preserves most MoE benefits on GLUE and SQuAD, often surpassing dense baselines, while delivering substantial inference speedups due to eliminated intra-MoE communication. The approach highlights the importance of the proficiency metric and timely dropping timing, offering a scalable path to deploying MoE capabilities in resource-constrained environments.

Abstract

The sparse Mixture-of-Experts (MoE) model is powerful for large-scale pre-training and has achieved promising results due to its model capacity. However, with trillions of parameters, MoE is hard to be deployed on cloud or mobile environment. The inference of MoE requires expert parallelism, which is not hardware-friendly and communication expensive. Especially for resource-limited downstream tasks, such sparse structure has to sacrifice a lot of computing efficiency for limited performance gains. In this work, we observe most experts contribute scarcely little to the MoE fine-tuning and inference. We further propose a general method to progressively drop the non-professional experts for the target downstream task, which preserves the benefits of MoE while reducing the MoE model into one single-expert dense model. Our experiments reveal that the fine-tuned single-expert model could preserve 99.3% benefits from MoE across six different types of tasks while enjoying 2x inference speed with free communication cost.

Figures (7)

• Figure 1: The contribution of experts varies in different downstream tasks The distribution of the contribution is quite long-tailed, indicating some experts are more "professional" on such tasks.
• Figure 2: The pipeline of our method. We split the fine-tuning process into several training windows evenly and prune the non-professional experts at the end of each training window.
• Figure 3: Different pruning criterion for professional experts. We compare the most professional expert chosen by hit rate $R$ and alpha score $\mathcal{A}$ on five different tasks. The most professional expert chosen by $\mathcal{A}$ better narrow preserve the performance of standard MoE-ft.
• Figure 4: The effect of training window length on GLUE tasks.$\gamma$ is a parameter to control the length of training window $\mathcal{L}$, where $\mathcal{L} = \gamma * \frac{N}{E}$. $N$ and $E$ denote the total training steps and the number of experts respectively.
• Figure 5: Comparison of different drop threshold. On the left figure, we illustrate the fine-tuned MoE accuracy on MRPC task with a different $\mathcal{A}$ threshold. On the right figure, we depict the number of survival experts at half of the schedule, which denoted as $K_{half}$.