Condense, Don't Just Prune: Enhancing Efficiency and Performance in MoE Layer Pruning

TL;DR

MoE architectures offer scalable capacity but incur substantial memory costs in real-world deployment. CD-MoE condenses sparse MoE layers into dense ones by retaining shared experts and a small, carefully selected subset of routing experts, using greedy layer- and expert-level condensation guided by Jensen-Shannon divergence. On DeepSeekMoE-16B, CD-MoE achieves about $27.5\%$ memory reduction and $1.26\times$ speed while preserving roughly $90\%$ of zero-shot accuracy, with lightweight fine-tuning recovering up to $98\%$ of the original performance on a single $80$G A100. These results demonstrate practical efficiency gains for fine-grained MoEs and suggest promising avenues for combining condensation with quantization or distillation to further reduce resource demands.

Abstract

Mixture-of-Experts (MoE) has garnered significant attention for its ability to scale up neural networks while utilizing the same or even fewer active parameters. However, MoE does not alleviate the massive memory requirements of networks, which limits their practicality in real-world applications, especially in the era of large language models (LLMs). While recent work explores the possibility of removing entire layers of MoE to reduce memory, the performance degradation is still notable. In this paper, we propose ConDense-MoE (CD-MoE), which, instead of dropping the entire MoE layer, condenses the large, sparse MoE layer into a smaller, denser layer with only a few experts activated for all tokens, while maintaining hardware friendliness. Our approach is specifically designed for fine-grained MoE with shared experts, where Feed-Forward Networks are split into many small experts, with certain experts isolated to serve as shared experts that are always activated, such as DeepSeekMoE and QwenMoE. We demonstrate the effectiveness of our method. Specifically, for the DeepSeekMoE-16B model, our approach maintains 90% of the average accuracy while reducing memory usage by 27.5% and increasing inference speed by 1.26 times. Moreover, we show that by applying lightweight expert fine-tuning -- only to the condensed layers -- and using 5 hours on a single 80G A100 GPU, we can successfully recover 98% of the original performance. Our code is available at: https://github.com/duterscmy/CD-MoE/tree/main.

Figures (6)

• Figure 1: Left: The structure of the Deepseek MoE layer. $w_i'$ represents the weights after normalization. Right: The structure of the ConDense-MoE layer, where only the most important top-k experts are retained. $\bar{w_i}$ represents the fixed weights that are pre-computed during condensing using the average weight of all calibration tokens.
• Figure 2: Fluctuations in the $JS$ divergence between the the outputs of the condensed model and the original dense model across different layers.
• Figure 3: CD-MoE against baselines on zero-shot tasks w/o fine-tuning. Left: Average accuracy with varying Memory Ratio against the original model. Right: Average accuracy with varying SpeedUp against the original model. The Gray dotted line is the original model result. CD-MoE-S represents the shared experts and no routing experts, and CD-MoE-SR represents shared with routing experts. Baseline indicated performance of the dense model.
• Figure 4: CD-MoE with lightweight fine-tuning. Left: SFT results on CD-MoE-S with increasing number of condensed layers. Right: SFT results on CD-MoE-SR with increasing number of condensed layers. Baseline indicated performance of the dense model.
• Figure 5: Left: fluctuations in the $KL$ divergence. Right: fluctuations in the perplexity.