Mosaic Pruning: A Hierarchical Framework for Generalizable Pruning of Mixture-of-Experts Models

TL;DR

Mosaic Pruning is a hierarchical framework that constructs a diverse, representative set of experts in Mixture-of-Experts models, enhancing their generalization across multiple domains and tasks while reducing memory overhead.

Contribution

It introduces a structured clustering and selection process for pruning experts, improving domain generalization and performance in Mixture-of-Experts models.

Findings

Achieves 7.24% improvement on general tasks

Achieves 8.92% improvement on specialized tasks

Outperforms prior pruning methods significantly

Abstract

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.