Hierarchical pruning improves efficiency and accuracy in ai language models
Fine-Grained Units, Flat Decisions: Hierarchical Micro-Expert Pruning for Mixture-of-Experts Models
Artificial Intelligence
Summary
Large AI language models sometimes waste computing power on parts of their networks that aren't very useful. The authors studied a way to carefully remove small pieces inside these models without hurting their ability to understand and generate language. They introduced UC-RAMP, a two-step method that decides how much to keep from each expert part and then fine-tunes which channels to keep. Their approach led to better performance on various tasks and faster response times when generating text.
What this means in practice
- •For machine learning engineers: Optimize large MoE language models for faster inference with minimal accuracy loss by selectively pruning expert channels post-training.
- •For cloud ai service providers: Reduce computational costs and latency of hosted language models by deploying pruned MoE models that maintain task accuracy.$Commercial implications: Enables development of cost-efficient AI language services that are faster and less expensive to operate at scale.
Authors
Zukang Xu, Xinpei Gao, Xing Hu, Zhe Jiang, Dawei Yang
Abstract
Micro-expert pruning reduces the capacity of mixture-of-experts (MoE) models by removing coupled channels within routed experts. However, finer pruning granularity alone does not guarantee the preservation of downstream capabilities, as similar language-modeling perplexity can accompany substantially different task performance. We revisit micro-expert mask construction as two distinct decisions: allocating retained capacity across experts and selecting channels within that capacity. We propose UC-RAMP, a two-stage post-training pruning framework that explicitly separates these decisions. CAPE uses cross-domain calibration evidence to redistribute expert widths under an exact layer-wise budget, while RAMP refines channel composition at fixed widths through one-out, one-in exchanges guided by a channel-interaction surrogate. Across three MoE models at 25% and 50% routed-expert channel pruning, UC-RAMP achieves the highest ten-task mean accuracy among evaluated pruning methods in five of six settings. At 50% pruning, it outperforms the strongest evaluated baseline by 8.19, 3.96, and 6.14 percentage points across the three models, respectively. In a single-A100 deployment at 25% pruning, UC-RAMP reduces measured time to first token by 19.7--23.9% and time per output token by 20.2--22.3% relative to the unpruned model. The code is available at https://github.com/xzktx003/UC_RAMP.