One-sentence Summary

Researchers from Xi'an Jiaotong University and China Telecom's TeleAI introduce Mosaic Pruning (MoP), a hierarchical cluster-then-select framework that leverages an Activation Variability Score and cross-domain expert similarity to retain functionally complementary experts and mitigate the generalization gap inherent in single-corpus pruning, achieving $7.24\%$7.24% higher performance on general tasks and $8.92\%$8.92% on specialized tasks such as math reasoning and code generation.

Key Contributions

• The paper introduces Mosaic Pruning (MoP), a post-training pruning framework for Sparse Mixture-of-Experts (SMoE) models that employs a hierarchical "cluster-then-select" process to retain a diverse, functionally complementary set of experts rather than optimizing for a single calibration corpus.

• The method groups experts by a similarity metric that captures performance across different task domains, then selects the most representative expert from each cluster using a proposed Activation Variability Score, which preserves specialist capabilities across diverse downstream tasks.

• Extensive experiments across multiple MoE model families demonstrate that MoP outperforms existing pruning baselines, achieving a 7.24% average improvement on general tasks and an 8.92% gain on specialized tasks such as mathematical reasoning and code generation.

Introduction

Mixture-of-Experts (MoE) models such as Mixtral 8x7B activate only a subset of their parameters during inference, delivering stronger performance than larger dense models like Llama 2 70B while keeping computational cost per token low. Yet their practical deployment is limited by massive static memory footprints (over 80GB for Mixtral) and by redundancy among experts, since many experts are functionally similar or contribute little to most tasks. Prior post-training pruning, exemplified by Enumeration Pruning, removes experts that contribute least to token reconstruction loss on generic corpora, but this fails to generalize to specialized domains, causing a "functional collapse" on tasks such as mathematical reasoning and code generation because generalist experts are retained while domain-specialist experts are discarded. To address this, the authors introduce Global Variability-aware Pruning (GVP), which scores each expert with an Activation Variability Score derived from KL divergence, and then propose Mosaic Pruning (MoP), a hierarchical cluster-then-select framework that groups experts by a Task Performance Similarity matrix and keeps the most critical representative within each group based on the diversity score. This preserves a functionally comprehensive set of experts and provides a strong inductive bias toward cross-task generalization; experiments on Mixtral-8x7B and Qwen1.5-MoE-A2.7B demonstrate average improvements of 7.24% on general tasks and 8.92% on specialized tasks.

Dataset

Dataset Composition and Sources

The authors evaluate their method using two calibration dataset sources and a structured set of evaluation benchmarks across multiple capability domains.

Calibration Datasets

• The main baseline, Enumeration Pruning, uses a calibration dataset sourced from the general-purpose C4 corpus.
• The proposed GVP and MoP methods rely on a mixed-diversity calibration dataset, designed to improve functional clustering. The exact composition of this dataset is described in a referenced section of the paper.

Evaluation Benchmark Subsets

• General capability assessment: seven standard benchmarks, including ARC-c/e, BoolQ, HellaSwag, MMLU, Open-BookQA, and WinoGrande.
• Mathematical reasoning: GSM8K and MATH.
• Code generation: HumanEval and MBPP.
• Expert diversity validation: six distinct domain-specific test sets, namely Math-GSM8K (math), Code-HumanEval (code), Commonsense-PIQA (commonsense reasoning), Summarization-XSum (summarization), Science-ARC-C (science), and Logic-MNLI (logic).

How the Data Is Used

• The calibration data is used exclusively for the pruning process, with different sources assigned to different methods: C4 for Enumeration Pruning, and the mixed-diversity set for GVP and MoP.
• The evaluation benchmarks are used to assess model performance after pruning, covering general capability, specialized skills, and expert diversity.
• Models are tested on two open-source MoE architectures: Mixtral-8x7B-Instruct (47B parameters, 8 experts per layer) and Qwen1.5-MoE-A2.7B-Chat (14.3B parameters, 60 experts per layer), allowing the authors to evaluate robustness and scalability across very different expert counts.
• All experiments run on a server equipped with four NVIDIA A100 80G GPUs.

Processing and Construction Details

• The mixed-diversity calibration dataset is constructed to support functional clustering, though the paper does not detail its exact filtering rules or schema in the provided text.
• The benchmark samples for expert diversity validation are drawn from six distinct domains, each represented by a single task type per domain, ensuring broad coverage for diversity checks.
• No cropping strategy or additional metadata construction is described in the supplied material; the pipeline centers on calibration data selection and benchmark evaluation only.

Method

The authors propose a hierarchical expert pruning framework designed to balance model performance and functional diversity. The methodology progresses through three distinct stages: an initial baseline approach, an intermediate variability-aware strategy, and a final hierarchical clustering-based solution. Refer to the framework diagram for a conceptual overview of these strategies.

The baseline strategy, Enumeration Pruning (Enu), focuses on minimizing the perturbation to the model's output upon expert removal. It searches a calibration dataset to find an optimal subset of experts that minimizes the output reconstruction loss. The optimization problem is formally expressed as:

where $\mathcal{L}_{\text{recon}}$Lrecon​ quantifies the cumulative difference between the outputs of the original unpruned Mixture of Experts (MoE) layer and the pruned layer. While effective for models with a small number of experts, exhaustive enumeration becomes computationally infeasible for larger configurations. In such cases, a greedy strategy is employed to iteratively remove the expert contributing least to the reconstruction loss increase. However, this approach tends to preserve functionally generalist experts at the expense of specialists crucial for downstream tasks.

To address the over-reliance on reconstruction loss in Enu, the authors introduce Global Variability-aware Pruning (GVP). This method prioritizes functional specialists using a unified variability metric. The core metric is the Activation Variability Score ($S_{\text{var}}$Svar​), which quantifies an expert's functional specialization by calculating the Kullback-Leibler divergence between its normalized activation distribution and a uniform distribution over all tokens:

A high $S_{\text{var}}$Svar​ score indicates a concentrated activation pattern, suggesting the expert is a functional specialist. The GVP strategy operates in two stages. First, it retains a core set of general experts ($\mathcal{E}_{\text{general}}$Egeneral​) by minimizing reconstruction loss to ensure baseline performance. Second, from the remaining candidate pool, it selects the most specialized experts based on the global ranking of $S_{\text{var}}$Svar​ scores. Despite its improvements, GVP's global ranking can retain functionally similar yet highly specialized experts, leading to redundancy.

To fundamentally address the redundancy issue in GVP, the authors propose Mosaic Pruning (MoP), a hierarchical framework that systematically preserves expert functional diversity. The first stage of MoP mirrors GVP by retaining a shared expert set to secure baseline performance. The core innovation lies in the second stage, which replaces GVP’s simple global ranking with a structured cluster-then-select process based on functional similarity. Refer to the workflow diagram for a detailed visualization of this process.

The MoP pipeline begins with data-driven domain discovery. A mixed-diversity calibration dataset is utilized to evaluate expert performance across different functional domains. The authors apply K-Means clustering to the input hidden states of all tokens to partition the token embeddings into semantically coherent domains, minimizing the Within-Cluster Sum of Squares (WCSS):

where $\boldsymbol{\mu}_k$μk​ is the centroid of the $k$k-th cluster. The number of clusters $K$K is set to match the number of experts to be selected, creating a one-to-one mapping between discovered functional clusters and experts.

Next, a similarity matrix is constructed to evaluate expert functional similarity. For each pair of experts in the candidate pool, the authors compute a similarity score based on their performance across the discovered domains. Spearman’s rank correlation coefficient is chosen as the core metric because it is nonparametric, robust to outliers, and captures monotonic relationships in expert performance. The performance vector $\mathbf{v}_i^{\text{perf}}$viperf​ for an expert $i$i is constructed using its average reconstruction error in each domain. The Spearman correlation $\rho$ρ between the rank-transformed performance vectors of two experts $i$i and $j$j is calculated, and the final performance similarity score is normalized to the interval $[0, 1]$[0,1]:

Agglomerative hierarchical clustering is then applied to group experts. The similarity matrix is converted into a distance matrix, and Ward’s linkage method is used to iteratively merge the closest pairs of clusters. The merge cost is defined as the increase in the total Error Sum of Squares (ESS) resulting from the merger:

This process continues until the candidate pool is partitioned into the target number of clusters.

Finally, intra-cluster representative selection is performed. The Activation Variability Score ($S_{\text{var}}$Svar​) is used to identify the most specialized expert within each functional cluster. For each cluster $G_k$Gk​, the representative expert $e_k^*$ek∗​ is chosen as:

The resulting set of representatives forms a group of experts that are both individually specialized and functionally complementary, effectively mitigating redundancy while maintaining diverse capabilities.

Experiment

The experiments were conducted on Mixtral-8x7B and Qwen1.5-MoE across general, mathematical, coding, and multi-domain benchmarks. MoP consistently outperformed all baselines across pruning rates and models, showing that preserving functional diversity enables more robust model compression while also delivering memory savings and faster inference. Ablation studies confirmed that each component contributes progressively, and MoP remains superior even when the calibration dataset is held equal. Activation analysis revealed that MoP retains clear domain specialization, whereas enumeration-based pruning causes functional homogenization, and on specialized math and code tasks the pruned model preserves complex reasoning abilities far better than existing methods.

Mosaic Pruning (MoP) consistently outperforms baseline pruning methods such as Enumeration, Random, and Frequency pruning across all expert retention rates and benchmarks, with an average improvement of 7.24% over Enumeration Pruning. The method preserves functional diversity among experts, leading to robust compression. It also achieves significant deployment efficiency gains, including memory savings and an average inference speedup of 1.20×. MoP achieves the highest average scores at both 6 and 4 experts per layer compared to all baselines. Performance gaps widen at higher pruning rates, with MoP retaining much stronger average accuracy than Frequency or Random pruning. The method's advantage is attributed to preserving expert functional diversity, preventing the functional homogenization seen in baseline methods.

An ablation study across two models shows progressive performance gains when moving from Enumeration Pruning to GVP to MoP, with MoP achieving the best average scores in both cases. Even when Enumeration Pruning is recalibrated on the same mixed-diversity dataset as MoP, MoP still outperforms it on all benchmarks. MoP achieves the highest average accuracy on both Mixtral and Qwen among all pruning methods. Using the same mixed calibration dataset, MoP surpasses Enumeration Pruning (Mixed) across every evaluated benchmark. Performance improves progressively from Enumeration Pruning to GVP to MoP, confirming the contribution of each component.

Mosaic Pruning (MoP) reduces the number of expert modules in MoE models, leading to substantial memory savings and faster inference. Across both Mixtral and Qwen architectures, pruning fewer experts yields meaningful speedups, with an average inference speedup of 1.20x. The method achieves these efficiency gains while preserving model performance. Removing expert modules significantly lowers peak memory usage across both models. Inference speed increases consistently as experts are pruned, with speedups ranging from 1.12x to 1.52x. MoP achieves an average 1.20x speedup due to fewer experts and reduced inter-device communication overhead.

MoP consistently outperforms Enumeration Pruning on GSM8K across both models and expert counts, with particularly large gains at higher pruning rates. On the MATH benchmark, improvements are less uniform: MoP achieves substantial gains for Qwen at high pruning, but results for Mixtral are mixed. The overall pattern supports that MoP preserves functional diversity and avoids the functional collapse seen in existing pruning methods. MoP yields large accuracy improvements over Enumeration Pruning on GSM8K, especially when pruning Mixtral down to fewer experts. On MATH, MoP shows a strong advantage for Qwen at higher pruning rates, while Mixtral results are more mixed. Across all specialized benchmarks, MoP averages an 8.92% performance gain over Enumeration Pruning, reflecting better preservation of expert specialization.

On code generation benchmarks with zero-shot evaluation, the MoP pruning method consistently outperforms Enumeration Pruning across both Mixtral and Qwen models at various expert retention levels. The performance gap is especially pronounced at higher pruning rates, where Enumeration Pruning collapses to near-zero scores while MoP retains meaningful capability. MoP preserves substantially higher HumanEval and MBPP scores than Enumeration Pruning at the same expert counts, particularly for Mixtral with 4 experts retained. Enumeration Pruning shows severe functional collapse at aggressive pruning rates, with HumanEval scores dropping below 2 points on both models. At a moderate pruning level for Qwen (50 experts), MoP achieves a HumanEval score over 10 times higher than Enumeration Pruning, while MBPP gains are smaller but still positive.

Mosaic Pruning (MoP) consistently outperforms baseline pruning methods across benchmarks, models, and expert retention rates by preserving functional diversity among experts, which prevents the functional collapse seen in baselines. Ablation studies confirm that each component contributes progressively, with MoP surpassing even recalibrated Enumeration Pruning on mixed-diversity datasets. Beyond accuracy, MoP yields meaningful deployment gains, including memory savings and an average 1.20x inference speedup. On specialized tasks like GSM8K, MATH, and code generation, MoP retains substantially higher performance at aggressive pruning rates where baselines often collapse.