Title: DiEP: Adaptive Mixture-of-Experts Compression through Differentiable Expert Pruning
Venue: NeurIPS 2025
Accepted: yes
Keywords: Expert Pruning

Differentiable Optimization

Language Model

Mixture-of-Experts
Abstract: Despite the significant breakthrough of Mixture-of-Experts (MoE), the increasing scale of these MoE models presents huge memory and storage challenges. Existing MoE pruning methods, which involve reducing parameter size with a uniform sparsity across all layers, often lead to suboptimal outcomes and performance degradation due to varying expert redundancy in different MoE layers. To address this, we propose a non-uniform pruning strategy, dubbed Differentiable Expert Pruning (DiEP), which adaptively adjusts pruning rates at the layer level while jointly learning inter-layer importance, effectively capturing the varying redundancy across different MoE layers. By transforming the global discrete search space into a continuous one, our method handles exponentially growing non-uniform expert combinations, enabling adaptive gradient-based pruning. Extensive experiments on five advanced MoE models demonstrate the efficacy of our method across various NLP tasks. Notably, \textbf{DiEP} retains around 92\% of original performance on Mixtral 8$\times$7B with only half the experts, outperforming other pruning methods by up to 7.1% on the challenging MMLU dataset.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes DiEP, a differentiable expert pruning framework for compressing sparse Mixture-of-Experts (MoE) models. Unlike existing methods that impose uniform pruning across layers or rely on costly discrete search, DiEP introduces a two-level optimization strategy: intra-layer expert importance scores (α) and inter-layer importance modulation (β). These are optimized on a small calibration dataset using a lightweight reconstruction-based loss. A novel expert skipping mechanism is also proposed to further reduce inference latency by dynamically skipping experts based on routing weight ratios and precomputed similarity statistics. DiEP is shown to achieve competitive performance (e.g., 92% retention with 50% experts) across multiple large-scale MoE models such as Mixtral, DeepSeek-MoE, and Qwen2.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- The paper proposes a novel method. Reformulating expert selection as a differentiable global optimization problem is both novel and practically effective.\n- The method introduces comparatively minimal overhead while delivering notable acceleration across various models.\n- The paper is well-structured and supports the proposed method with comprehensive experiments.\n\n### Weaknesses\n\n- The finetuning cost of the proposed technique remains unclear. \n- Some design choices lack rigorous justification—for example, the use of α and β to quantify expert and layer importance, as well as the rationale behind the choice of the α:β update ratio.\n\nQUESTIONS:\n- Could you elaborate on the γ computation, particularly the role of γ₂ derived from CKA similarity? My understanding is that γ₂ is computed offline using the calibration dataset and remains fixed during inference. If so, how can this static similarity-based threshold generalize to unseen prompts or tasks at inference time? Is there theoretical or empirical evidence supporting the transferability of γ₂ across different inputs?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces DiEP, a differentiable pruning framework for Mixture-of-Experts (MoE) models that addresses expert redundancy by learning to prune both intra-layer and inter-layer experts adaptively. Rather than relying on heuristic-based or exhaustive search methods that use uniform pruning across all MoE layers, DiEP formulates pruning as a continuous optimization problem. It introduces two sets of learnable parameters: intra-layer importance scores and inter-layer importance scores, which together guide a globally optimized pruning decision. Additionally, the paper presents an adaptive expert skipping mechanism during inference based on routing weights and expert similarity.\nExperiments across five MoE-based LLMs, including Mixtral 8x7B, DeepSeek-MoE, and Qwen2, show that DiEP achieves better performance than prior pruning methods while reducing up to 50% of experts and improving inference efficiency.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n\nThis paper addressed an important scalability challenge in MoE-based LLMs: making these models deployable by reducing computational and memory overhead without significant accuracy loss.\n\nThe proposed method is technically sound and well-motivated, grounded in clear observations of expert redundancy variability across MoE layers.\n\nThe continuous relaxation and differentiable optimization approach is well-integrated and executed with practical efficiency.\n\nEmpirical evaluation is conducted across diverse model families and tasks, including ablation and visualization studies that substantiate the claims.\n\n**Weakness**\n\nThe method introduces additional training overhead (although minor) and requires calibration data, which might not be readily available in all deployment settings.\n\nThe paper could provide more theoretical guarantees or robustness analysis, especially under domain shift or data drift conditions.\n\nWhile the visualization analyses are compelling, deeper analysis on why certain layers or experts are deemed more important would strengthen the interpretability of the method.\n\nQUESTIONS:\nHow sensitive is the method to the choice of calibration dataset? You mention results on C4 and Math datasets; however, can you clarify how performance varies when the calibration data is misaligned with downstream tasks?\n\nCan the method handle dynamic pruning ratios? Would it be possible to adaptively determine the pruning ratio during optimization rather than treating it as a fixed hyperparameter?\n\nLimitations in Multi-task or Continual Learning. How would DiEP behave in a continual learning or multi-task setting where expert utility might shift?\n\nLIMITATIONS:\nYes, the authors have adequately discussed the limitations and trade-offs involved, such as reliance on calibration data and the inability of prior uniform pruning approaches to handle inter-layer heterogeneity. The discussion on societal impact is minimal, but the work primarily targets efficiency improvements without introducing major ethical concerns.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes DiEP, a differentiable expert pruning framework that formulates expert selection as a continuous optimization problem. By leveraging gradient-based optimization and incorporating an adaptive expert skipping mechanism, DiEP aims to reduce memory consumption and accelerate inference while preserving the performance of sparse Mixture-of-Experts (MoE) models. The method is evaluated across multiple language tasks, showing superior performance over existing MoE pruning baselines and establishing a new benchmark for efficient MoE deployment.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- The paper is well-motivated, addressing the important and practical challenge of reducing the computational cost of MoE models through expert path pruning.\n\n- The proposed use of differentiable NAS for expert path selection is intuitive and well-formulated. While not novel, it aligns with recent trends in neural architecture optimization and is a reasonable application to the MoE setting.\n\n- The writing is clear and the paper is easy to follow, with appropriate background and methodological descriptions.\n\n- Experiments are conducted on multiple strong open-source MoE language models, including Mistral and DeepSeek, demonstrating the practical utility of the method.\n\nWeaknesses\n\n- The core technical contribution is limited in novelty. The idea of applying differentiable NAS to select or prune expert paths has strong conceptual overlap with earlier work in CNN model search and compression. This paper primarily adapts these techniques to MoE models without introducing new architectural or algorithmic innovations tailored to the unique characteristics of MoEs.\n\n- A major concern lies in the lack of domain generalization. The routing paths are statically searched and fixed on a single dataset, which introduces data-specific biases. MoE models are designed to leverage dynamic expert selection for better generalization across diverse inputs (e.g., mathematical reasoning vs commonsense QA). Static pruning may significantly impair this flexibility, reducing the model’s robustness and transferability.\n\n- The paper does not sufficiently engage with recent MoE pruning or compression papers, e.g., [1,2].\n\n[1] MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts\n[2] Mixture Compressor for Mixture-of-Experts LLMs Gains More\n\nQUESTIONS:\nSee Weakness.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}