DiEP: Adaptive Mixture-of-Experts Compression through Differentiable Expert Pruning

We thank reviewer sDD7 for providing positive feedback and helpful suggestions. Below please find a response.

W1: Concerns on the Pruning Cost

We appreciate the reviewer's feedback and the opportunity to provide a detailed analysis of the computational cost of our method. We totally agree that analyzing pruning cost is crucial. Our experiments show DiEP is highly efficient, providing a superior cost-accuracy balance. We use Memory-hour Cost (GB·h)—a metric capturing both peak GPU memory and time—to holistically measure the total computational effort on two distinct MoE models.

Table 1: Pruning Cost and Accuracy on Mixtral 8x7B

Analysis:

• Superior Efficiency: DiEP achieves a memory-hour cost of just 31.97 GB·h. This is a 74% reduction compared to NAEE and a 70% reduction compared to MC-MoE.
• Best-in-Class Performance-per-Cost: While DiEP's cost is on par with S-MoE, it achieves a significantly higher accuracy (59.9 vs. 55.9). This demonstrates that DiEP offers the best trade-off, delivering the highest accuracy for a given, minimal computational budget.

Table 2: Pruning Cost and Accuracy on Deepseek-MoE16B

Analysis:

• Scalablility: The efficiency gains of DiEP are even more pronounced on models with a larger number of experts. DiEP's memory-hour cost is only 15.42 GB·h, which is 24% lower than the next most efficient method, S-MoE (20.42 GB·h).
• Highest Accuracy at Lowest Cost: Crucially, DiEP not only has the lowest pruning cost but also achieves the highest accuracy (54.1), improving upon S-MoE's accuracy by +1.4 points while being more efficient.

Addressing the Hidden Costs of Other Methods:

It is important to note that methods like S-MoE, while appearing fast, have hidden computational costs that scale poorly. S-MoE's multi-step expert merging process involves:

1. Expert similarity analysis (CKA): This requires $O(S \cdot N^2)$ computations, where S is the number of samples and N is the number of experts. This quadratic scaling with the number of experts makes it increasingly expensive for modern MoE architectures.
2. Learning merging coefficients: This introduces an additional optimization overhead.

In contrast, DiEP's gradient-based approach avoids these expensive, quadratically scaling steps. Its computational cost is primarily determined by a few forward/backward passes on a small calibration set, making it highly scalable and efficient.

Conclusion: The experimental data clearly demonstrates that the "finetuning" or pruning cost of DiEP is not a concern; rather, it is a significant strength. It delivers state-of-the-art accuracy at a minimal and scalable computational cost, making it a practical solution for task-agnostic deployment.

W2: Justification for Design Choices

We thank the reviewer for this insightful question regarding our core design principles. These choices were motivated by the multi-scale nature of expert redundancy in MoE models.

1. Rationale for α (Intra-layer) and β (Inter-layer) Scores

We decompose expert importance into two components to model redundancy at different granularities:

• α (Intra-layer scores): Learns the local, relative importance of experts competing within the same layer.
• β (Inter-layer scores): Learns the global, depth-aware importance of entire layers, addressing the fact that shallower layers are often more critical than deeper ones (as shown in our Figure 4).

By combining them ($s_i^{(l)} = \bar{\alpha}_i^{(l)} \cdot \beta^{(l)}$), we create a unified score that enables a principled, non-uniform pruning across the entire model in a single, global sort. This factorized approach is designed to be more expressive and effective than using a single, monolithic importance score.

2. Rationale and Empirical Validation for the α:β Update Ratio

The update ratio is an empirical choice aimed at achieving optimization stability. The high-dimensional α scores learn fine-grained expert rankings, while the low-dimensional β scores make coarse-grained, systemic adjustments. Our hypothesis is that allowing α to update more frequently helps stabilize the local expert rankings before the more impactful β scores are adjusted.

To rigorously justify our choice, we conducted a new ablation study on the α:β update ratio. The results, shown below, are for 25% and 50% pruning ratios (formatted as 25%/50%).

Key Findings:

• The experimental results clearly validate our design choice. The 3:1 update ratio consistently achieves the best or near-best performance across almost all tasks and pruning ratios, culminating in the highest average scores for both 25% (63.8) and 50% (59.9) pruning.
• The data shows that giving more updates to the intra-layer scores (α) generally leads to better performance (e.g., 2:1 and 4:1 outperform 1:1 and 1:2). This supports our hypothesis that stabilizing the local expert rankings is crucial.
• The 3:1 ratio strikes an optimal balance, outperforming both less frequent (e.g., 2:1) and more frequent (e.g., 4:1) α update schedules in terms of average performance.

In conclusion, our α/β decomposition provides a principled way to model multi-scale redundancy, and our chosen 3:1 update ratio is not arbitrary but is empirically validated to be the most effective schedule for stable and high-performing optimization. We will incorporate these new results into our revised manuscript.

Q1: Elaboration on $γ_{2}$

We sincerely thank the reviewer for this insightful question. It correctly identifies a crucial aspect of our adaptive skipping mechanism and allows us to elaborate on the design rationale and provide new empirical evidence for its effectiveness.

In fact, traditional routing weights effectively measure the importance of an expert for a given token. However, they do not capture the redundancy between the selected experts. A scenario can easily arise where the router assigns high scores to two experts (e.g., top-2) that are functionally very similar. In this case, activating both is computationally wasteful. The core motivation for our design is precisely this: to prune not just unimportant experts, but also redundant ones. This is where $γ_{2}$ comes in. It acts as a redundancy penalty.

Generalizability of$γ_{2}$ (CKA Similarity): A Task-Agnostic Structural Prior

You are correct that $γ_{2}$ is computed offline on a calibration set and remains fixed. The key to its generalizability lies in what it represents: $γ_{2}$ is not a data-specific feature but rather a measure of the inherent, task-agnostic structural similarity between expert weight matrices.

The expert weights are the result of pre-training on trillions of tokens. The functional similarities and differences captured by CKA are therefore a deeply embedded, stable property of the model itself. While the activation of experts is dynamic, their fundamental functional relationship is static. Therefore, $γ_{2}$ acts as a robust and transferable structural prior that remains valid across unseen prompts and tasks.

Here we conducted an ablation study as empirical evidence based on the NAEE common framework[1] . We started with the full model and then compared the performance of a skipping mechanism based solely on routing weights ($γ_{1}$) against our proposed method which combines both routing weights and similarity ($γ_{1} \times γ_{2}$).

Analysis of Results:

• $γ_{2}$ Improves Accuracy: The results clearly show that incorporating the similarity-based $γ_{2}$ leads to better performance. Our method achieves an average accuracy of 64.1, recovering 0.4 points of the accuracy lost by the $γ_{1}$-only approach, while still providing a slightly better speedup.

Conclusion: In summary, our two-factor $γ$ is a principled design. $γ_{1}$ (from routing weights) provides a dynamic, per-token importance signal, while $γ_{2}$ (from CKA) provides a static, generalizable structural redundancy penalty. The empirical evidence confirms that this combination leads to a more robust and accurate adaptive skipping mechanism that successfully transfers to unseen inputs at inference time. We will add this detailed justification and ablation study to our revised paper.

[1]. Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experts large language models. ACL 2024.

We thank Reviewer 3XFq for the thoughtful feedback. Below, we address each concern.

W1: Training Overhead & Calibration Data

We agree that practical compression methods should incur minimal cost. DiEP is designed as a lightweight, one-shot pruning approach.

1. Efficient Training

DiEP’s calibration is fast and cost-effective:

Cost Comparison (Appendix A.3): On Mixtral 8x7B, DiEP takes just 0.23 hours / 31.97 GB·h, a 74% reduction over NAEE (124.58 GB·h).

Quality vs. Cost: For similar cost to S-MoE, DiEP achieves +4.0 accuracy. Compared to MC-MoE, it's 70% cheaper with +28 points accuracy gain.

Thus, DiEP offers better performance at lower overhead than many training-free methods.

2. Minimal Calibration Data

As shown in Table 9 (Appendix A.6), DiEP achieves optimal performance on Mixtral 8x7B with just 128 samples. This is a tiny, almost negligible amount of data that can be easily sourced. The table also shows that even with as few as 32 samples, DiEP avoids performance collapse and still functions effectively.

Task-Agnostic Data: Our domain generalization test (see W2) shows DiEP pruned on C4 performs well on MathQA/CommonseQA, indicating no need for in-domain data.

Conclusion: The overhead and data requirements of DiEP are intentionally designed to be minimal to ensure broad applicability. It uses minimal task-agnostic data, and delivers superior pruned models.

W2&Q1: Robustness & Theory

We thank the reviewer for raising the importance of robustness under domain shift. DiEP’s design encourages generalization both theoretically and empirically.

1. Theory: Stable & Knowledge-Preserving

Convergence Proof (Appendix B.2): Our alternating update converges to a stationary point under standard conditions.

Knowledge-Preserving Regularization: The key to DiEP's robustness lies in our objective function: $\min_{\alpha, \beta} \mathcal{L}_{ce}(y, F'(x; \alpha, \beta)) + \lambda \cdot \Phi(\alpha, \beta)$The term $\Phi(\alpha, \beta) = ||F'(x; \alpha, \beta) - F(x)||_F$ acts as a powerful knowledge-preserving regularizer. which transfers generalization to the pruned model $F′$, mitigating overfitting to limited calibration data.

Future plans include extending to stochastic settings, analyzing non-convex landscapes, and formalizing the role of λ.

2. Empirical: Robustness to Domain Shift

We simulate domain shift by pruning on general (C4), in-domain (Math), or mixed (C4+Math) data, and testing on CommonseQA & MathQA.

Table 1: Pruning on General (C4)

DiEP on general data retains performance, beating NAEE by +9.51 points under aggressive pruning.

Table 2: Pruning on In-Domain (Math)

Even when biased toward Math, DiEP generalizes well to CommonseQA.

Table 3: Pruning on Mixed (C4+Math)

DiEP improves further with in-domain data; in fact, outperforms full model on MathQA.

We will highlight these as "robustness to domain shift" results in the paper.

W3: Deeper Analysis on the Importance of Layers and Experts

We thank the reviewer for highlighting the importance of interpretability. Understanding why certain experts and layers are retained by DiEP provides valuable insight into the method’s effectiveness. We offer the following analysis based on the learned α and β scores.

1. Why DiEP Prioritizes Shallow Layers (High β Scores)

As shown in Figure 4b, DiEP assigns consistently higher β scores to earlier layers (roughly 1–15). This pattern aligns with well-established findings on the hierarchical structure of Transformers:

Shallow layers are known to extract foundational syntactic and semantic features [1–3], forming the base for all subsequent reasoning. These layers handle diverse, low-level signals shared across tasks and domains. Preserving experts in these layers is essential, as early-stage degradation can propagate and amplify in deeper stages.

Deeper layers, in contrast, operate on refined features and are responsible for more abstract contextualization. Their computations tend to be more task-specific and less universally critical [4], making them more suitable targets for pruning.

Thus, DiEP’s β scores reflect a principled, data-driven retention of layers based on their functional necessity, not just position.

2. Why DiEP Prefers Certain Experts (High α Scores)

Figure 4a reveals notable variance in intra-layer α scores, with only a subset of experts per layer being deemed essential. This supports the MoE hypothesis of functional specialization [5, 6]:

Generalist experts exhibit high α scores across data points, likely specializing in high-frequency, broadly useful linguistic patterns. Specialist experts may focus on niche phenomena and are invoked less frequently.

Crucially, DiEP’s joint modeling of α (expert importance) and β (layer importance) allows fine-grained pruning decisions:

• A layer may be important (high β), but only a few of its experts (high α) are truly impactful.
• DiEP thus preserves these "hub" experts while pruning redundant or less frequently activated ones. This coupled, multi-scale selection is a key strength of DiEP: it balances global layer utility with local expert relevance, leading to sparse yet highly functional subnetworks.

Conclusion

The α–β scores learned by DiEP reveal how MoE models encode hierarchical and distributed knowledge:

• β scores prioritize depth-wise importance, aligning with representational hierarchy.
• α scores capture intra-layer specialization, aligning with expert sparsity and task coverage.

By jointly modeling both, DiEP achieves interpretable and structured pruning.

References
[1] BERT Rediscovers the Classical NLP Pipeline. ACL 2019.
[2] Analyzing Multi-Head Self-Attention. EMNLP 2020.
[3] A Primer in BERTology. TACL 2020.
[4] Sparsely-Gated MoE Layer. ICLR 2017.
[5] Ultimate Expert Specialization in MoE LMs. arXiv:2401.06066.
[6] Efficient Expert Pruning for MoE LLMs. ACL 2024.

Q2: Adaptive Pruning Ratio

We thank the reviewer for this insightful question. Below we elaborate on both current DIEP capabilities and potential extensions.

1. Dynamic Ratios via Threshold-Based Pruning (Already Supported)

DiEP's core output is a continuous importance score ($s_i^{(l)}$) for each expert. Instead of pruning a fixed percentage (top-K), one can simply set an importance threshold, T, and prune all experts with scores below it. This naturally leads to a dynamic pruning ratio that adapts to the model's specific redundancy profile, a flexibility already supported by our method.

2. Learning the Optimal Ratio via Differentiable Regularization

Furthermore, our differentiable framework can be extended to learn the pruning ratio directly. We propose a simple yet powerful extension:

The added term is $\mathcal{L_s}(\alpha, \beta) = \sum_{l=1}^{L} \sum_{i=1}^{N} |s_i^{(l)}| = \sum_{l=1}^{L} \sum_{i=1}^{N} |\bar{\alpha}_i^{(l)} \cdot \beta^{(l)}|$, which encourages the model to shrink unimportant scores. The final pruning ratio thus becomes an emergent property of the learned solution, shaped by the calibration data and task constraints.

This formulation allows the model to "justify" the retention of each expert, resulting in pruning strategies that are both principled and task-adaptive.

Conclusion This is a valuable direction for future work that our framework is well-positioned to explore. We will add this discussion to our revised manuscript.

Q3: Continual Learning Evaluation

We thank the reviewer for this critical question. To assess DiEP's robustness to shifting expert utility, we conducted a continual learning experiment on Hellaswag, BoolQ, and MathQA as the initial calibration dataset. We then evaluate the performance of the pruned models to measure how well they generalize and retain knowledge. A lower "Forgetting" score indicates better knowledge preservation.

Experimental Results: The tables below show the performance (25%/50% ratio).

NAEE

DiEP

Summary

The results clearly show that DiEP is more robust in a continual learning setting. By regularizing towards the original model's function, DiEP learns to preserve experts that are fundamentally important for general capabilities, rather than just for the specific calibration task. This makes our pruning method inherently more robust to the challenges of continual learning.

Dear Reviewer 3XFq,

Thank you very much for your timely reply! We sincerely appreciate the time and thoughtful feedback you have provided throughout the review process, as well as your recognition that our responses have addressed your concerns.

Best regards,

Authors

We thank reviewer xJgk for providing positive feedback and helpful suggestions. Below please find a response.

W1: Novelty Concerns

We thank the reviewer for their insightful feedback. We agree that our DIEP builds upon the established concept of differentiable search. However, we respectfully argue that DiEP is not a straightforward adaptation. Instead, it introduces significant and novel algorithmic and architectural innovations to address the unique challenges of MoE models, which differ fundamentally from the dense CNN/RNN architectures targeted by early NAS methods.

We highlight two key areas of novelty:

1. A Fundamentally Different Problem Formulation and a Novel Optimization Framework

There is a core difference in both the goal and the optimization mechanics between DARTS and DiEP.

• DARTS targets operation search within a DAG-based cell. It solves a bilevel optimization problem to find the best set of operations (e.g., convolution, pooling) and their connections:

  $$\min_\alpha \mathcal{L}_{val}(w^*(\alpha), \alpha) \quad$$ $$\text{s.t. } w^*(\alpha) = \arg\min_w \mathcal{L}_{train}(w,\alpha)$$

This formulation is memory-intensive, relies on a validation set, and is designed to search a discrete set of operations, not prune entire expert sub-networks.

• DiEP targets global, non-uniform expert pruning across all MoE layers. We formulate this as a single-level optimization problem with a novel reconstruction objectivet: $$\min_{\alpha, \beta} \mathcal{L}_{ce}(y, F'(x; \alpha, \beta)) + \lambda \cdot \Phi(\alpha, \beta)$$

Where:$F'$ is the pruned model, and $F$ is the original full model. $\Phi(\alpha, \beta) = ||F'(x; \alpha, \beta) - F(x)||_F$ is a reconstruction regularization term.

This approach avoids the validation set. The architecture parameters ($\alpha, \beta$) are learned directly on the training data alongside the model weights (via an alternating strategy), making the process more efficient and robust.

This algorithmic shift from bilevel, validation-based search to single-level, reconstruction-guided optimization is a primary contribution of our work.

2. Architecture-Specific Innovations

MoE models present a unique challenge: managing redundancy across dozens large expert sub-networks, where redundancy varies significantly from layer to layer. Our method is explicitly designed for this context.

• Targeting the Right Granularity: Instead of searching for micro-level operations like DARTS, DiEP introduces learnable, continuous proxies designed for the macro-level task of expert selection: Intra-layer scores ($\alpha$) learn the relative significance of experts within a single MoE layer. Inter-layer scores ($\beta$) learn to modulate the pruning pressure across different layers, enabling a non-uniform pruning budget. This mechanism directly addresses the observed phenomenon that expert utility is not uniform across model depth in MoEs.

• Efficient Design: The combination of $(\alpha, \beta)$ allows us to compute a global importance $s_i^{(l)} = \bar{\alpha}_i^{(l)} \cdot \beta^{(l)}$ for every expert in the model. This enables a global sort to find the least important experts, regardless of their layer. This entire mechanism has a lightweight memory complexity of $O(L \times N)$ (layers $\times$ experts), making it highly scalable for modern MoEs with many layers and experts. A naive application of DARTS-style search to this problem would be computationally infeasible.

References DARTS: Differentiable Architecture Search. ICLR 2019.

W2: Domain Generalization and Static Pruning

We sincerely thank the reviewer for raising this critical point. Here our response is structured around two key points: a clarification of our methodology and direct empirical evidence from our new experiments.

1. Clarification: Static Pruning of Experts vs. Dynamic Routing of Tokens

A crucial distinction must be made: DiEP performs static pruning of the expert pool, but the resulting pruned model retains its fully dynamic routing mechanism. For any given input, the router still dynamically selects the top-k experts from the remaining, optimized set of experts. We are not fixing the routing paths for tokens; we are identifying and removing a subset of experts that are fundamentally redundant, thereby creating a more efficient expert library. The core strength of MoE — adapting to diverse inputs by selecting different experts — is preserved.

2. Empirical Evidence: Robust Generalization from General-Domain to Specific-Domain Tasks

To test for domain generalization, we pruned our model using general-domain (C4), in-domain (Math), and mixed (C4+Math) calibration sets. We then evaluated performance on two specialized tasks: CommonseQA (commonsense reasoning) and MathQA (mathematical reasoning).

Table 1: Pruning on General-Domain Data (C4)

When calibrating on the general C4 dataset, DiEP demonstrates strong robustness.

Table 2: Pruning on In-Domain data (Math)

Even when calibrated on Math data, DiEP excels on the out-of-distribution (OOD) CommonseQA task.

Table 3: Pruning on Mixed-Domain Data (C4+Math)

We can see adding a small amount of in-domain data to the calibration set improves performance.

Conclusion: These experiments directly counter the concern about domain generalization. DiEP's strength lies in its ability to leverage a small, task-agnostic dataset to learn a globally optimized, non-uniform pruning mask that preserves the most critical experts for a wide range of downstream tasks.

W3: Engagement with Recent MoE Compression Literature

We sincerely thank the reviewer for highlighting these important recent papers. Here we make the detailed comparison with such approaches.

Comparison with Mixture Compressor (MC-MoE) and Delta Decompression (D²-MoE)

These two methods represent the state-of-the-art in compression of MoE models. They differ significantly from DiEP in both their paradigm and methodology.

• Paradigm Difference:

  • MC-MoE & D²-MoE: They apply techniques like mixed-precision quantization (MC-MoE) or weight decomposition into shared components (D²-MoE) directly to a trained model using heuristics.
  • DiEP: DiEP is a lightweight, gradient-based fine-tuning method. It leverages a small, task-agnostic dataset to learn expert importance scores via optimization, allowing for a more data-aware and principled pruning decision.

• Methodology Difference:

  • MC-MoE & D²-MoE: These are hybrid compressors. They modify and transform the weights of all experts through quantization, SVD, or decomposition. They do not entirely remove experts.
  • DiEP: DiEP is a pure pruning method. It identifies and completely removes the least important expert sub-networks, resulting in a model with fewer total experts.

Experimental Comparison

We compare DiEP against MC-MoE[1] and D²-MoE [2]. The results on several key benchmarks are presented below.

The results clearly show that DiEP significantly outperforms both D²-MoE and MC-MoE, especially at r=25%. On the challenging MMLU benchmark, DiEP achieves a score of 64.9, which is 4.3 points higher than D²-MoE and 4.6 points higher than MC-MoE, despite MC-MoE operating at a higher bit-rate.

Comparison with MoE++

MoE++ and our DiEP method are orthogonal, as MoE++ focuses on accelerating inference by routing simple tokens to "zero-computation experts" to reduce FLOPs, without shrinking the model's static size. In contrast, DiEP is designed specifically to compress the static model footprint, reducing storage and memory requirements. Given their fundamentally different objectives, a direct compression performance comparison is not applicable.

[1] Mixture Compressor for Mixture-of-Experts LLMs Gains More, ICLR2025. [2] Delta Decompression for MoE-based LLMs Compression, ICML2025.