PreMoe: Lightening MoEs on Constrained Memory by Expert Pruning and Retrieval

TL;DR

MoE models scale model capacity but suffer from prohibitive memory footprints due to storing all experts and backbone components. PreMoe tackles this with Probabilistic Expert Pruning guided by the Task-Conditioned Expected Selection Score ($TCESS$) and Task-Adaptive Expert Retrieval (TAER) that pre-computes and loads only task-relevant experts, supplemented by 4-bit quantization. Empirical results on DeepSeek-R1 and Pangu-MoE variants show near-full reasoning performance under substantial expert pruning (up to ~87.5% reduction) and meaningful throughput gains, enabling deployment on memory-constrained platforms. The approach demonstrates that task-specific expert specialization is a powerful lever for compressing MoE models while preserving core capabilities, broadening access to state-of-the-art sparse architectures. Overall, PreMoe offers a principled, scalable path to dynamic, task-aware loading of large MoE models for cloud and edge environments.

Abstract

Mixture-of-experts (MoE) architectures enable scaling large language models (LLMs) to vast parameter counts without a proportional rise in computational costs. However, the significant memory demands of large MoE models hinder their deployment across various computational environments, from cloud servers to consumer devices. This study first demonstrates pronounced task-specific specialization in expert activation patterns within MoE layers. Building on this, we introduce PreMoe, a novel framework that enables efficient deployment of massive MoE models in memory-constrained environments. PreMoe features two main components: probabilistic expert pruning (PEP) and task-adaptive expert retrieval (TAER). PEP employs a new metric, the task-conditioned expected selection score (TCESS), derived from router logits to quantify expert importance for specific tasks, thereby identifying a minimal set of critical experts. TAER leverages these task-specific expert importance profiles for efficient inference. It pre-computes and stores compact expert patterns for diverse tasks. When a user query is received, TAER rapidly identifies the most relevant stored task pattern and reconstructs the model by loading only the small subset of experts crucial for that task. This approach dramatically reduces the memory footprint across all deployment scenarios. DeepSeek-R1 671B maintains 97.2\% accuracy on MATH500 when pruned to 8/128 configuration (50\% expert reduction), and still achieves 72.0\% with aggressive 8/32 pruning (87.5\% expert reduction). Pangu-Ultra-MoE 718B achieves 97.15\% on MATH500 and 81.3\% on AIME24 with 8/128 pruning, while even more aggressive pruning to 4/64 (390GB memory) preserves 96.95\% accuracy on MATH500. We make our code publicly available at https://github.com/JarvisPei/PreMoe.

Figures (6)

• Figure 1: Overview of our proposed PreMoe framework.
• Figure 2: Heatmap of average activated router logits across BigBench tasks (rows) and experts (columns) in DeepSeek-R1. Red indicates high positive logits, blue negative ones. The sparse patterns of red highlights task-specific expert specialization.
• Figure 3: Example of Reasoning generation of DeepSeek-R1 with different pruning ratios.
• Figure A.1: Comparison Reasoning generation of DeepSeek-R1 when collecting TCESS patterns with or without considering the model's reasoning output. The top example uses our default collection strategy (query + reasoning output), while the bottom example derives patterns only from the input query, leading to repetitive output for DeepSeek-R1 (8/32).
• Figure A.2: Example 2 of Reasoning generation of DeepSeek-R1 with different pruning ratios (8/256, 8/128, 8/32) using PreMoe.