🤖 AI Summary
To address the high memory overhead of Mixture-of-Experts (MoE) models due to their large parameter count and the significant accuracy degradation caused by existing expert-level pruning methods, this paper proposes a fine-grained atomic expert pruning framework. Its core innovation lies in the first-ever mapping of Hessian information from the expert parameter space to the atomic expert output space—reducing computational complexity to *O(d²)*—and leveraging the Optimal Brain Surgeon principle to estimate atomic expert importance via forward/backward passes on a small calibration dataset. Evaluated on DeepSeek-MoE and Qwen-MoE, the method achieves 20–25% model compression with negligible accuracy loss while reducing inference FLOPs by approximately 20%, substantially improving MoE model deployment efficiency.
📝 Abstract
Mixture-of-Experts (MoE) architectures in large language models (LLMs) deliver exceptional performance and reduced inference costs compared to dense LLMs. However, their large parameter counts result in prohibitive memory requirements, limiting practical deployment. While existing pruning methods primarily focus on expert-level pruning, this coarse granularity often leads to substantial accuracy degradation. In this work, we introduce HEAPr, a novel pruning algorithm that decomposes experts into smaller, indivisible atomic experts, enabling more precise and flexible atomic expert pruning. To measure the importance of each atomic expert, we leverage second-order information based on principles similar to Optimal Brain Surgeon (OBS) theory. To address the computational and storage challenges posed by second-order information, HEAPr exploits the inherent properties of atomic experts to transform the second-order information from expert parameters into that of atomic expert parameters, and further simplifies it to the second-order information of atomic expert outputs. This approach reduces the space complexity from $O(d^4)$, where d is the model's dimensionality, to $O(d^2)$. HEAPr requires only two forward passes and one backward pass on a small calibration set to compute the importance of atomic experts. Extensive experiments on MoE models, including DeepSeek MoE and Qwen MoE family, demonstrate that HEAPr outperforms existing expert-level pruning methods across a wide range of compression ratios and benchmarks. Specifically, HEAPr achieves nearly lossless compression at compression ratios of 20% ~ 25% in most models, while also reducing FLOPs nearly by 20%. The code can be found at href{https://github.com/LLIKKE/HEAPr}{https://github.com/LLIKKE/HEAPr}.