Mosaic Pruning for MoE Models

• MoP is a hierarchical pruning technique that preserves functional diversity in MoE models through domain-aware clustering and selection.
• It combines generalist selection with Spearman rank-based clustering and activation variability scoring to effectively retain specialized experts.
• Empirical results show MoP outperforms conventional methods in both general and specialized tasks, enabling a 'prune once, deploy everywhere' approach.

Mosaic Pruning (MoP) denotes a hierarchical, domain-aware pruning technique for Mixture-of-Experts (MoE) models, devised to mitigate generalization collapse observed in conventional post-training pruning. Rather than relying solely on global reconstruction loss minimization with respect to a single corpus, MoP partitions experts into functional clusters spanning both generalist and highly specialized behaviors. By structurally combining data-driven domain discovery, Spearman rank-based clustering of expert performance profiles, and an intra-cluster selection criterion grounded in activation concentration (Activation Variability Score), MoP produces pruned MoE models that retain a representative tile from each key functional specialization. Empirical results demonstrate substantial gains in both general-purpose and high-specialization domains (e.g., numerical reasoning, code generation), positioning MoP as a robust, generalizable, and performant pruning framework—enabling "prune once, deploy everywhere" capability for modern LLMs (Hu et al., 25 Nov 2025).

1. Motivation and Challenges in Pruning MoE Models

The increased adoption of Sparse Mixture-of-Experts (SMoE) architectures—where only a fraction of the model's experts are consulted per input—has led to challenges in post-training pruning, particularly for LLMs. Conventional pruning approaches such as Enumeration Pruning retain experts based solely on reconstruction loss minimization across a general-purpose calibration corpus. These approaches are agnostic to the functional specializations of individual experts, often resulting in a set of generalists while discarding domain-specific experts that are indispensable for specialized tasks.

This leads to two critical deployment bottlenecks:

• Catastrophic domain generalization failure: Pruned models typically collapse on domain-specific tasks (e.g., mathematics, program synthesis) not well-represented in the single calibration set.
• Inefficient re-pruning and calibration: Each new application domain requires retraining or costly pruning, reducing the utility of a single pruned checkpoint for heterogeneous downstream use.

Mosaic Pruning addresses these limitations by enforcing functional diversity through a structured, hierarchical process (Hu et al., 25 Nov 2025).

2. Hierarchical "Cluster-Then-Select" Framework

MoP operationalizes its core goal—to distill both generalists and domain-specific specialists into the pruned model—in two principal stages:

Stage 1: Generalist Selection (Enumeration Pruning Analog)

Given a trained MoE layer with $n$ experts and calibration data $\mathcal{D}_{\text{cache}}$, select $m$ generalist experts via:

$$E_{\text{general}} = \arg\min_{S \subset E,\, |S|=m} L_{\text{recon}}(\mathcal{D}_{\text{cache}}, S)$$

where $L_{\text{recon}}$ measures output mean squared reconstruction error. The set $E_{\text{cand}} = E \setminus E_{\text{general}}$ constitutes candidate specialists.

Stage 2: Domain Discovery, Clustering, and Specialist Selection

2.1 Domain Discovery

• Apply K-Means clustering to hidden state activations $\{x_t\}_{t=1}^{N_{\text{total}}}$ from calibration data, partitioning tokens into $K = r-m$ domains.

2.2 Expert Performance Profiling

• For each candidate expert $i$, compute domain-wise reconstruction errors:

$$v^{\text{perf}}_{i,k} = \frac{1}{|T_k|} \sum_{t \in T_k} \| E_i(x_t) - z_{\text{real},t}\|_2^2$$

where $E_i(\cdot)$ means only expert $i$ is active on input, and $T_k$ is the set of tokens in domain $k$.

2.3 Functional Similarity and Clustering

• Rank-transform $v^{\text{perf}}_i$ per expert, and compute pairwise Spearman correlation $\rho(R_i, R_j)$. Normalized similarity score:

$$S_{\text{perf}}(i, j) = \frac{1 + \rho(R_i, R_j)}{2}$$

• Construct distance $D(i, j) = 1 - S_{\text{perf}}(i, j)$.
• Perform agglomerative clustering (Ward's linkage), merging clusters to minimize cluster mean error increase, halting at $K$ clusters.

2.4 Intra-Cluster Specialist Identification

• For each expert $i$, compute the Activation Variability Score:

$$S_{\text{var}}(i) = \sum_{t=1}^{N_{\text{total}}} (p_{t,i} / Z_i) \log_2 \left((p_{t,i}/Z_i) N_{\text{total}}\right)$$

where $p_{t,i}$ is the gating logit of expert $i$ on token $t$, and $Z_i = \sum_t p_{t,i}$. High $S_{\text{var}}$ denotes sharp specialization.

• Within each cluster $G_k$, select $e_k^* = \arg\max_{i \in G_k} S_{\text{var}}(i)$.

The final retained expert set is $$E_{\text{final}} = E_{\text{general}} \cup \{e_1^*, \dotsc, e_K^*\}$$.

3. Algorithmic Pseudocode and Workflow

The MoP process can be summarized as follows:

1. Compute $E_{\text{general}}$ by minimizing $L_{\text{recon}}$ with $m$ experts.
2. Set $E_{\text{cand}} = E \setminus E_{\text{general}}$, $K = r - m$.
3. Cluster calibration activations with K-Means $(K)$, assign domains.
4. For each $i \in E_{\text{cand}}$, compute $v^{\text{perf}}_i \in \mathbb{R}^K$.
5. Compute pairwise normalized Spearman similarity matrix $S_{\text{perf}}(i, j)$ for all $(i, j)$.
6. Run Ward's agglomerative clustering over $E_{\text{cand}}$ using $D(i, j)$.
7. Within each cluster, compute $S_{\text{var}}(i)$, select maximal $e_k^*$.
8. Output the union $E_{\text{final}}$ as the pruned expert set.

This hierarchical, cluster-driven mechanism is designed to guarantee both coverage (domain diversity) and efficiency (no domain is left without a representative specialist).

4. Theoretical Rationale and Comparison to Prior Approaches

MoP is motivated by the observation that minimizing global reconstruction loss (Enumeration Pruning) generally yields a "generic" pool of experts with mediocre specialization, resulting in severe degradation on domains outside the calibration set—a phenomenon described as functional collapse. This is especially acute in MoE models harboring latent specialists crucial for non-mainstream tasks such as arithmetic reasoning or code synthesis.

By introducing a dedicated clustering step of experts' domain error profiles, MoP injects an explicit inductive bias for functional diversity. This ensures that each discovered specialization is represented by the most focused expert (as measured by activation variability), preventing the dilution of domain-specific capacity. The framework thus aims to bridge the generalization gap plaguing prior pruning pipelines, promoting robust transfer of pruned models across diverse tasks without repeated per-domain reconfiguration (Hu et al., 25 Nov 2025).

5. Empirical Evaluation and Results

Extensive experiments substantiate MoP over several major MoE LLMs in both general and specialist-driven evaluations:

Model	Experts Post-Pruning	Enumeration Pruning (Avg)	MoP (Avg)	Absolute Gain
Mixtral-8×7B	6	72.73	74.01	+1.28
Qwen1.5-MoE7B	50	55.11	59.10	+4.00

• General Capability (Benchmarks: ARC-c/e, BoolQ, HellaSwag, MMLU, OBQA, WinoGrande): MoP achieves approximately 7.24% average relative gain over Enumeration Pruning.
• Specialized Tasks (e.g., GSM8K, MATH, HumanEval, MBPP): MoP improves domain-specific expert retention, with an overall average improvement of approximately 8.92%.
  • For instance, in code generation, Qwen1.5-MoE7B saw gains from 1.2, 8.2 to 14.0, 9.3 in HumanEval and MBPP, respectively.
• Diversity Validation: Expert activation heatmaps reveal that Enumeration Pruning leads to a few generalists dominating all domains, whereas MoP selects domain-specialized experts.

6. Domain Coverage and Generalizability

A salient property of MoP is the guarantee that no latent specialization, as discovered through unsupervised clustering of the token-embedding space, is omitted in the pruned configuration. Diversity validation analyses confirm that, unlike prior art which yields experts with overlapping functional profiles, MoP's approach preserves both broad and narrow competencies in the surviving experts, a key factor in facilitating stable out-of-domain and cross-task transfer.

This procedure validates the "prune once, deploy everywhere" paradigm. A plausible implication is that hardware efficiency and memory savings do not come at the cost of eliminating niche model capacity. Moreover, by leveraging small, mixed-diversity calibration sets for both clustering and selection, MoP's pipeline remains efficient and easily adaptable to evolving deployment scenarios without repeated re-training.

7. Conclusion

Mosaic Pruning introduces a systematic, inductively-biased approach to expert selection in SMoE LLMs. By combining unsupervised domain discovery, Spearman-rank functional clustering, and KL-based intra-cluster activation variability, MoP distills a panel of both generalists and domain-specialized experts. This principled design translates into substantial empirical gains (7–9%) across both generalist benchmarks and high-specialization tasks, substantiating MoP as an effective universal pruning method for the next generation of deployable, robust MoE LLMs (Hu et al., 25 Nov 2025).