FlexMoE: 面向MoE语言模型的一体通用嵌套式专家内剪枝

# FlexMoE：面向MoE语言模型的一体化嵌套专家内剪枝

来源：https://arxiv.org/html/2606.27866

范墨 新加坡国立大学 [email protected]  
韩玉轩 新加坡国立大学 [email protected]  
张耿 新加坡国立大学 [email protected]  
赵王波 新加坡国立大学 [email protected]  
游洋 新加坡国立大学 [email protected]  

###### 摘要

混合专家（MoE）语言模型通过稀疏激活的专家来扩展模型能力，使这种架构成为现代大型模型的标准方案。然而，稀疏激活并未消除存储和部署所有专家所带来的负担，而且不同的设备、用户和工作负载之间的可用部署预算可能存在显著差异。现有的MoE压缩方法仍然大多是固定预算的，通常针对每个选定的目标预算优化一个压缩后的端点。我们研究了一种不同的设置：将大型预训练MoE大语言模型转化为一个跨预算的、可部署的子网络嵌套族。我们的方法首先根据重要性对专家FFN通道进行排序，然后让每个专家学习一个离散动作来剪枝其通道。通过逐步增加成本压力，一次动作训练就能导出一系列从高预算到低预算的动作掩码，每个掩码标识出一个可靠的、嵌套在排序后基模型中的较小子网络。此外，我们还在中等剪枝预算（40%）下使用一次恢复性微调来恢复退化后的模型质量，并将恢复后的模型迁移到其他未见过的预算上。总体而言，我们的框架超越了近期MoE压缩基线。具体来说，在Qwen2-57B-A14B上，我们的方法在无需微调的情况下，剪枝50%的路由专家参数，仍能保留约99.8%的基模型性能。在部署方面，我们的剪枝子网络实现了真实的内存减少和吞吐量提升，并进一步通过内核级协同设计支持实时在线预算切换。

## 1 引言

参考图说明（请见图1）

图1：现有固定预算MoE压缩方法与FlexMoE的比较。现有方法通常优化或实例化一个针对特定目标预算的压缩端点，包括专家剪枝/跳过方法（如NAEE (Lu et al., 2024)）、专家权重分解方法（如MoE-SVD (Li et al., 2025) 和 TD-MoE (Xu et al., 2026)）。相比之下，FlexMoE 实现了“一次剪枝，多处部署”的结构性剪枝和多预算部署，无需为每个部署目标维护单独的压缩检查点。

混合专家（MoE）已成为扩展大型模型最有效的方案之一：通过增加总参数容量，同时每个token只激活一小部分专家，MoE模型可以在更低的活跃计算量下实现高质量 (Lepikhin et al., 2021; Fedus et al., 2022)。这一设计已经驱动了大量成功的系统，从顶级大规模研究系统，到一系列紧凑而强大的开源模型，表明稀疏专家架构是扩展现代基础模型日益标准化的方式。

然而，稀疏激活的专家并未消除存储和部署所有专家的负担。在部署中，我们并不总是需要、或者有预算（我们在这里使用“预算”一词来表示给定服务设置下模型可用的部署侧资源包，包括设备硬件容量、内存占用、延迟/吞吐量目标、上下文长度需求或服务级别目标（SLO））来使用其最大容量运行的最大MoE模型。在实践中，这类预算因用户、设备、平台和工作负载而异：一个服务可能需要在充裕的预算下进行长上下文或多智能体推理，而另一个服务可能在更严格的延迟和成本约束下处理轻量级聊天机器人请求 (Fu et al., 2024; Gao et al., 2026; Ma et al., 2026)。这使得MoE压缩不仅仅是一个模型简化问题，也是一个部署适配问题。

然而，现有的MoE压缩方法仍然大多是固定预算的。专家剪枝方法在结构上移除、跳过或合并专家，而低秩分解或专家内剪枝等压缩方法通常针对选定的目标预算优化一个压缩端点 (Lu et al., 2024; Bai et al., 2025; Yang et al., 2024b; Gu et al., 2025; Li et al., 2025; Xu et al., 2026)。当部署目标事先已知时，这很有用，但当预算因场景而异或随时间变化时，就很不方便：切换到另一个操作点通常需要重新运行整个压缩流程、重新加载模型或跨不同预算维护多个模型。此外，近期的系统研究强调，真实的大语言模型服务必须在非平稳流量和混合请求需求下运行，在线重新配置通常很有价值，但难以清晰实现 (Gao et al., 2026; Qian et al., 2026)。

受这些观察启发，我们提出了 **FlexMoE**，一个嵌套的专家内剪枝框架，将单个预训练MoE模型转化为一个跨预算、可物质化的子网络族（我们这里用“预算”表示剪枝参数占路由专家总参数的比例）。如图2所示，我们的流程首先根据估计的重要性对专家FFN通道进行重新排序，以实现保留顶部通道的切片。然后，我们让每个专家在一个逐步增加剪枝压力的训练循环中，从预定义的离散比率集合中学习一个保留比率，从而单次训练就能得到跨多个预算的一系列动作掩码，每个掩码标识出一个嵌套在排序后预训练MoE模型中的可靠剪枝子网络。可选地，我们进一步执行一次中等预算的微调阶段，从参数剪枝中恢复一组共享权重，从而实现整个预算族中模型的复用。这产生了一个“一次训练，多处部署”的剪枝流程。除了在实时在线预算调整场景中实现真正的吞吐量优势外，我们还进一步通过内核级优化，将预算族与面向部署的系统协同设计相结合。

我们在 Mixtral-8x7B、Phi-3.5-MoE 和 Qwen2-57B-A14B 上进行了全面实验。在这些基座上，FlexMoE 超越了强大的 MoE 压缩基线，同时提供了更灵活的部署。特别是在 Qwen2-57B-A14B 上，剪枝后的子网络在 50% 剪枝预算下保留了大约 99.8% 的基模型性能，在 80% 预算下仍保留了约 92.9%。我们进一步表明，从单次中等预算微调策略中恢复的共享权重能够很好地迁移到未见过的更高和更低预算，并且由此产生的子网络族在部署中提供了真正的吞吐量提升和更灵活的接口，尤其是在与我们探索性的算法-系统协同设计结合用于在线预算切换时。

## 2 相关工作

### 2.1 MoE 压缩

如图1所示，MoE压缩方法主要沿两个方向进行。大量MoE参数剪枝工作是在专家层面减少、跳过或合并专家。NAEE表明，预训练MoE大语言模型中的专家并非同等重要，并研究了专家剪枝和跳过以提高部署效率 (Lu et al., 2024)。DiEP通过可微分专家剪枝进一步推进了这一方向，在压缩目标下学习保留哪些专家，而不仅仅依赖启发式专家排序 (Bai et al., 2025)。在专家合并方面，HC-SMoE聚类并合并功能相似的专家，以减少模型大小而无需重新训练 (Chen et al., 2024)。对于跳过，MoNE用轻量级新手替换冗余的专家输出，而不是保留完整的专家计算，提供了另一种减少参数和部署成本的方法 (Zhang et al., 2025)。

另一条线是在内部压缩专家权重，同时很大程度上保留原始MoE结构。MoE-SVD 和 D2-MoE 是预训练MoE大语言模型的代表性方法，而 TD-MoE 进一步将这一方向扩展到每个层内的跨专家联合张量分解 (Li et al., 2025; Gu et al., 2025; Xu et al., 2026)。MoE-I2 也相关，它结合了专家间剪枝和专家内分解 (Yang et al., 2024b)。我们的工作与这个专家内压缩家族最为接近，但大多数是基于低秩分解的。然而，我们的剪枝方法是基于训练的、细粒度的。我们还对来自同一预训练MoE基模型的每个专家内部施加了一个可切分的嵌套结构，该模型共享相同的路由器分配和专家拓扑，从而实现了更灵活的部署接口和预算切换。

### 2.2 嵌套子网络、掩码学习与排序

更广泛的“一次训练，多处部署”思想源于弹性嵌套子网络训练。US-Nets 引入了通用可瘦子网络，后来 MatFormer、Flextron 和 AmoebaLLM 将嵌套或多合一参数共享扩展到 Transformer 和大语言模型 (Yu and Huang, 2019; Devvrit et al., 2024; Cai et al., 2024; Fu et al., 2024)。更近期的MoE工作已经开始通过学习从粗到细的专家排序或可瘦专家宽度，将弹性直接引入MoE语言模型 (Wang et al., 2025; Tastan et al., 2026)。我们的工作不同：我们不是在进行MoE预训练期间引入弹性，而是从预训练的MoE大语言模型出发，在后训练阶段导出一系列跨预算的嵌套子网络。我们的方法还借鉴了可微分结构学习和基于梯度的重要性排序。MaskLLM 和 Gumbel-Softmax 提供了使用梯度优化学习离散结构决策的标准工具 (Fang et al., 2024; Jang et al., 2017)，而 RECAP 和更早的泰勒准则表明，分组一阶显著性对于结构化剪枝决策是有效的 (Ilhan et al., 2024; Molchanov et al., 2019)。我们在一个新设置中结合了这些要素：在按重要性排序的专家通道上学习每个专家的通道前缀切片动作。

参考图说明（请见图2）

图2：FlexMoE流程的整体可视化。首先，通过估计的重要性对每个专家FFN投影通道进行重新排序，以实现保留顶部通道的切片。然后，让每个专家在一个逐步增加剪枝压力的优化循环中持续学习一个保留比率，从而单次训练就能得到跨多个预算的一系列动作掩码，每个保存的掩码标示出一个嵌套在重新排序后的基MoE模型中的剪枝子网络。随后，在中等预算（40%）下进行可选的单点微调阶段，从参数剪枝中恢复一组共享权重，从而实现整个预算族中恢复后权重的复用。

## 3 方法

我们考虑一个预训练的MoE大语言模型，包含 \(L\) 个MoE层和每层 \(E\) 个专家。对于路由到第 \(l\) 层第 \(e\) 个专家的隐藏状态 \(h\)，我们将专家FFN写为：

\[
\mathrm{FFN}_{l,e}(h)=W_{l,e}^{\mathrm{down}}\Big(\phi(W_{l,e}^{\mathrm{gate}}h) \odot W_{l,e}^{\mathrm{up}}h\Big),
\tag{1}
\]

其中 \(\phi(\cdot)\) 是激活函数，中间FFN宽度为 \(d_{\mathrm{ff}}\)。

### 3.1 专家内通道排序

直接将通道前缀切片应用于预训练专家是脆弱的。因此，我们执行一次性的隐藏通道重排序步骤，将每个专家FFN转换为按重要性排序的布局，使得更小的前缀在前缀切片剪枝下保留更重要的参数。对于专家 \((l,e)\)，我们定义一个结构化参数组 \(\Theta_{l,e,j}\)，对应第 \(j\) 个FFN隐藏通道，由 \(W_{l,e}^{\mathrm{gate}}\) 和 \(W_{l,e}^{\mathrm{up}}\) 的相应行以及 \(W_{l,e}^{\mathrm{down}}\) 的匹配列组成。换句话说，\(\Theta_{l,e,j}\) 收集了FFN中连接到隐藏通道 \(j\) 的所有参数。只要对 \(W_{l,e}^{\mathrm{gate}}\)、\(W_{l,e}^{\mathrm{up}}\) 和 \(W_{l,e}^{\mathrm{down}}\) 中的相应行/列一致地应用相同的排列，FFN隐藏通道就是排列不变的 (Navon et al., 2023)。这允许我们按重要性重新排序隐藏通道而不改变完整的专家函数，同时使前缀通道切片动作更有意义。

我们在一小组校准批次 \(\mathcal{B}_{\mathrm{rank}}\) 上为每个隐藏通道估计一阶泰勒显著性。对于批次 \(b \in \mathcal{B}_{\mathrm{rank}}\)，批次级和最终的组显著性计算如下：

\[
g_{\theta}^{(b)}=\frac{\partial\mathcal{L}^{(b)}}{\partial\theta},\qquad s_{l,e,j}^{(b)}=\sum_{\theta \in \Theta_{l,e,j}}(\theta\,g_{\theta}^{(b)})^{2},\qquad s_{l,e,j}=\frac{1}{|\mathcal{B}_{\mathrm{rank}}|}\sum_{b\in \mathcal{B}_{\mathrm{rank}}}s_{l,e,j}^{(b)}.
\tag{2}
\]

直观地，\(s_{l,e,j}\) 估计了参数组 \(\Theta_{l,e,j}\)（通道 \(j\)）对损失 \(\mathcal{L}\) 的影响程度。这遵循了结构化显著性的分组一阶泰勒视角，与 RECAP (Molchanov et al., 2019; Ilhan et al., 2024) 中使用的排序信号密切相关。然而，在我们的设置中，我们仅将其作为预处理步骤在动作学习之前对专家FFN隐藏通道进行排序，而不是作为迭代剪枝-恢复过程的一部分。

然后，我们按降序排序 \(s_{l,e,j}\)，得到一个排列 \(\pi_{l,e}\)，或等价地一个排列矩阵 \(P_{l,e}\)，它被一致地应用于专家FFN隐藏维度：

\[
\tilde{W}_{l,e}^{\mathrm{gate}}=P_{l,e}W_{l,e}^{\mathrm{gate}},\quad \tilde{W}_{l,e}^{\mathrm{up}}=P_{l,e}W_{l,e}^{\mathrm{up}}
\]