Mosaic Pruning: A Hierarchical Framework for Generalizable Pruning of Mixture-of-Experts Models
This addresses the practical deployment issue of high memory overhead in SMoE models for large language model applications, offering a more generalizable pruning method that is incremental over prior domain-specific approaches.
The paper tackled the problem of poor generalization in post-training pruning of Sparse Mixture-of-Experts models, which suffer from catastrophic performance degradation when applied to new domains, by introducing Mosaic Pruning, a hierarchical framework that selects a functionally complementary set of experts to retain diverse capabilities, resulting in gains of 7.24% on general tasks and 8.92% on specialized tasks.
Sparse Mixture-of-Experts (SMoE) architectures have enabled a new frontier in scaling Large Language Models (LLMs), offering superior performance by activating only a fraction of their total parameters during inference. However, their practical deployment is severely hampered by substantial static memory overhead, as all experts must be loaded into memory. Existing post-training pruning methods, while reducing model size, often derive their pruning criteria from a single, general-purpose corpus. This leads to a critical limitation: a catastrophic performance degradation when the pruned model is applied to other domains, necessitating a costly re-pruning for each new domain. To address this generalization gap, we introduce Mosaic Pruning (MoP). The core idea of MoP is to construct a functionally comprehensive set of experts through a structured ``cluster-then-select" process. This process leverages a similarity metric that captures expert performance across different task domains to functionally cluster the experts, and subsequently selects the most representative expert from each cluster based on our proposed Activation Variability Score. Unlike methods that optimize for a single corpus, our proposed Mosaic Pruning ensures that the pruned model retains a functionally complementary set of experts, much like the tiles of a mosaic that together form a complete picture of the original model's capabilities, enabling it to handle diverse downstream tasks.Extensive experiments on various MoE models demonstrate the superiority of our approach. MoP significantly outperforms prior work, achieving a 7.24\% gain on general tasks and 8.92\% on specialized tasks like math reasoning and code generation.