PreMoe: 専門家の刈り込みと検索による制約されたメモリ上のMoEの軽量化

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

May 23, 2025

著者: Zehua Pei, Ying Zhang, Hui-Ling Zhen, Xianzhi Yu, Wulong Liu, Sinno Jialin Pan, Mingxuan Yuan, Bei Yu

cs.AI

要旨

Mixture-of-experts（MoE）アーキテクチャは、大規模言語モデル（LLMs）を膨大なパラメータ数にスケーリングすることを可能にし、計算コストの比例的な増加を伴わない。しかし、大規模なMoEモデルのメモリ需要は、クラウドサーバーから消費者向けデバイスまで、さまざまな計算環境での展開を妨げている。本研究では、まずMoE層内の専門家活性化パターンに顕著なタスク固有の特化が存在することを示す。これに基づいて、メモリ制約のある環境での大規模MoEモデルの効率的な展開を可能にする新しいフレームワークであるPreMoeを提案する。PreMoeは、確率的専門家プルーニング（PEP）とタスク適応型専門家検索（TAER）という2つの主要なコンポーネントを特徴とする。PEPは、ルーターのロジットから導出されたタスク条件付き期待選択スコア（TCESS）という新しいメトリックを使用し、特定のタスクに対する専門家の重要性を定量化し、最小限の重要な専門家セットを特定する。TAERは、これらのタスク固有の専門家重要度プロファイルを活用して効率的な推論を実現する。TAERは、多様なタスクに対するコンパクトな専門家パターンを事前に計算し保存する。ユーザークエリを受信すると、TAERは最も関連性の高い保存済みタスクパターンを迅速に特定し、そのタスクに不可欠な専門家の小さなサブセットのみをロードしてモデルを再構築する。このアプローチにより、すべての展開シナリオでのメモリ使用量が大幅に削減される。DeepSeek-R1 671Bは、8/128構成（50％の専門家削減）にプルーニングされた場合、MATH500で97.2％の精度を維持し、さらに積極的な8/32プルーニング（87.5％の専門家削減）でも72.0％の精度を達成する。Pangu-Ultra-MoE 718Bは、8/128プルーニングでMATH500で97.15％、AIME24で81.3％の精度を達成し、さらに積極的な4/64プルーニング（390GBメモリ）でもMATH500で96.95％の精度を維持する。我々は、コードをhttps://github.com/JarvisPei/PreMoeで公開している。

English

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.