FlexMoE: One-for-All Nested Intra-Expert Pruning for MoE Language Models

2026-06-26

AI Summary

This work addresses the challenge of repeatedly compressing Mixture-of-Experts (MoE) models under varying deployment budgets by proposing the first nested pruning framework that supports arbitrary budget switching. The approach constructs a family of nested subnetworks through expert FFN channel importance ranking, discrete pruning action learning, and progressive cost-constrained optimization, followed by a single unified fine-tuning step to recover performance across all budgets. This enables one-time training, multi-level deployment, and runtime online budget adaptation. Evaluated on Qwen2-57B-A14B, the method retains 99.8% of the original performance while using only 50% of expert parameters, substantially reducing memory consumption and improving throughput, outperforming existing MoE compression techniques.

Abstract

Mixture-of-Experts (MoE) language models scale model ability with sparsely activated experts, making this architecture a standard recipe for modern large models. However, sparse activation does not remove the deployment burden of storing and serving all experts, and the available deployment budget can vary substantially across devices, users, and workloads. Existing MoE compression methods are still largely fixed-budget, typically optimizing one compressed endpoint at each chosen target budget. We study a different setting: converting a large pretrained MoE LLM into a nested family of deployable subnetworks across budgets. Our method first ranks expert FFN channels by their importance, then lets each expert learn a discrete action to prune its channels. By gradually increasing cost pressure, a single action-training run exports a series of action masks from high to low budgets, each of which identifies a reliable smaller subnetwork nested in the ranked base model. Moreover, we use a single recovery fine-tune at a mid pruning budget (40%) to recover degraded model quality and transfer the recovered model to other unseen budgets. Overall, our framework surpasses recent MoE compression baselines. Specifically, on Qwen2-57B-A14B, our method retains ~99.8% of base performance while pruning 50% of routed expert parameters even without fine-tuning. For deployment, our pruned subnetworks deliver real memory reduction and throughput gains, and further support realtime online budget switching with kernel-level co-design.

Problem

• Mixture-of-Experts
• model compression
• deployment budget
• nested subnetworks
• sparse activation

Innovation

• nested pruning
• Mixture-of-Experts
• budget-adaptive compression
• channel importance ranking
• online budget switching