[论文] ToMoE：通过动态结构剪枝将稠密大语言模型转换为混合专家模型

# 通过动态结构剪枝将稠密大语言模型转换为混合专家模型  
来源：https://arxiv.org/html/2501.15316  
高尚千 [email protected]  
致谢：通讯作者  
所属机构：佛罗里达州立大学计算机科学系  
丁华 [email protected]  
雷扎·希尔卡万德 [email protected]  
所属机构：马里兰大学帕克分校计算机科学系  
林志恒 [email protected]  
所属机构：三星美国研究院  
唐政 [email protected]  
所属机构：三星美国研究院  
李政浩 [email protected]  
所属机构：佛罗里达州立大学计算机科学系  
袁龙革 [email protected]  
所属机构：佛罗里达州立大学计算机科学系  
李方毅 [email protected]  
所属机构：宾夕法尼亚大学工程与应用科学学院  
张泽宇 [email protected]  
所属机构：亚马逊通用人工智能部门  
阿里雷扎·甘杰丹尼沙利加尼 [email protected]  
所属机构：马里兰大学帕克分校计算机科学系  
钱露 [email protected]  
所属机构：中佛罗里达大学计算机科学系  
许洁 [email protected]  
所属机构：佛罗里达大学健康结果与生物医学信息学系  
许彦昌 [email protected]  
所属机构：三星美国研究院  

###### 摘要  
大型语言模型（LLMs）展现出卓越的能力，但因高计算需求而面临部署挑战。传统剪枝方法通过永久移除参数来降低这些成本，这不可避免地导致性能下降。为缓解此问题，我们提出 ToMoE 方法，该方法通过揭示稠密模型内部固有的专家结构，将稠密 LLMs 转换为混合专家（MoE）模型，且无需任何权重更新。ToMoE 利用动态结构剪枝将专家构建与路由器训练统一在单一阶段，从而实现一致的高性能。值得注意的是，即使不对模型权重进行微调，ToMoE 在 Phi-2、LLaMA-2、LLaMA-3 和 Qwen-2.5 等模型上也持续优于最先进的剪枝和 MoE 技术。本文代码已开源：https://github.com/gaosh/ToMoE。

## 1 引言  
尽管 LLMs 展现出了执行多样化任务的卓越能力（8 (https://arxiv.org/html/2501.15316#bib.bib29)；34 (https://arxiv.org/html/2501.15316#bib.bib30)；56 (https://arxiv.org/html/2501.15316#bib.bib31)；50 (https://arxiv.org/html/2501.15316#bib.bib32)；3 (https://arxiv.org/html/2501.15316#bib.bib33)；16 (https://arxiv.org/html/2501.15316#bib.bib36)；55 (https://arxiv.org/html/2501.15316#bib.bib34)；33 (https://arxiv.org/html/2501.15316#bib.bib35)），但其庞大的模型规模常常限制了在资源有限设备上的可用性。因此，大量研究（45 (https://arxiv.org/html/2501.15316#bib.bib12)；5 (https://arxiv.org/html/2501.15316#bib.bib11)；21 (https://arxiv.org/html/2501.15316#bib.bib37)）聚焦于最小化这些模型的计算和内存开销。结构剪枝（45 (https://arxiv.org/html/2501.15316#bib.bib12)）已成为解决此挑战的一个有前景的方案，因为与非结构化剪枝不同，它在无需专用实现的情况下即可实现压缩。然而，结构剪枝方法的问题在于，它会大幅降低模型容量，导致与稠密模型相比出现显著的性能差距。即使部分恢复这一差距，其微调成本也极其高昂。为了在参数数量与性能之间取得更好的权衡，稀疏混合专家（MoE）模型（58 (https://arxiv.org/html/2501.15316#bib.bib38)；38 (https://arxiv.org/html/2501.15316#bib.bib39)）被设计为仅激活模型参数的一个子集（对应于所选专家）。近期提出的 MoE 模型，如 DeepseekMoE（13 (https://arxiv.org/html/2501.15316#bib.bib40)），证明了其能在使用少量激活参数的同时，匹配具有相似总参数量的稠密模型的性能。沿着这一思路，将稠密模型转换为 MoE 模型，可能为弥补结构剪枝方法留下的性能差距提供一种有前景的途径。  
与之前从稠密模型构建 MoE 模型的方法（71 (https://arxiv.org/html/2501.15316#bib.bib61)；37 (https://arxiv.org/html/2501.15316#bib.bib60)；73 (https://arxiv.org/html/2501.15316#bib.bib41)）不同，我们的研究发现，MoE 本质上存在于稠密模型中，并且无需更新模型权重（持续预训练）即可被揭示。具体而言，我们表明这些专家可以通过动态结构剪枝来识别。这些结果代表了一项前所未有的新颖贡献。  
MoE 模型的核心思想是条件计算，即根据输入 token 动态选择专家。这一概念与动态剪枝方法（25 (https://arxiv.org/html/2501.15316#bib.bib15)）密切相关，后者根据输入特征做出剪枝决策。利用这一关联，我们提出使用动态结构剪枝从稠密模型构建 MoE 模型。具体来说，对于多头自注意力（MHA）层，我们应用 top-K 路由和静态剪枝进行压缩；对于 MLP 层，我们使用 top-1 专家路由将其转换为 MoE 层。为动态结构剪枝学习的路由机制可直接用作 MoE 层的路由模块。通过可微离散操作，MoE 转换过程可表述为一个可微动态剪枝问题。基于此表述，我们可以高效地将稠密模型转换为 MoE 模型，其成本与常规结构剪枝方法相当或更低。图 1 (https://arxiv.org/html/2501.15316#S2.F1) 展示了我们的方法、静态剪枝以及原始 LLM 之间的比较。  
基于上述发现和技术，我们提出了 ToMoE，以通过动态剪枝有效地将稠密 LLMs 转换为 MoE 模型。本工作的贡献可总结如下：
- • 通过动态剪枝实现稠密到 MoE 的转换：我们引入了一种通过动态剪枝将稠密模型转换为 MoE 模型的新方法。具体而言，我们在 MHA 层中沿头维度实现 top-K 路由和静态剪枝，并在 MLP 层中对学习到的专家采用 top-1 路由。此公式确保了稀疏高效的计算，同时保留了模型容量。
- • 路由与专家的联合优化：所提出的方法通过求解一个正则化优化问题来联合优化路由模块和专家配置。我们的方法利用可微操作实现高效灵活的 MoE 构建。
- • 一致的性能提升：我们的方法在仅训练路由器（不微调模型权重）的情况下，在多种任务上持续优于最先进的结构剪枝和 MoE 构建技术。这一性能提升在 Phi-2、LLaMA-2、LLaMA-3 和 Qwen-2.5 等广泛使用的公开模型上得到了证明。
- • 详细分析：我们从多个角度广泛分析了 ToMoE 生成的模型，包括参数分配、路由器行为以及不同设计组件的消融实验。我们希望这些分析能为该领域的未来研究提供有价值的见解和指导。

## 2 相关工作  
(a) (b) (c)  
图 1： (a)：原始 LLM 使用所有参数处理所有输入文本。(b)：LLMs 的静态剪枝永久移除模型参数，剩余参数用于处理所有输入文本。我们的方法 (c)：使用动态剪枝的 LLMs 采用不同的子网络（以不同颜色表示）处理不同的 token。我们引入 MoE 以在推理时实现固定的预期计算预算。  
**剪枝**：结构剪枝（40 (https://arxiv.org/html/2501.15316#bib.bib59)；35 (https://arxiv.org/html/2501.15316#bib.bib1)；45 (https://arxiv.org/html/2501.15316#bib.bib12)）是一种有吸引力的现实部署技术，因为它移除了冗余参数以减小模型大小，且无需专用实现。结构剪枝方法主要分为两类：静态剪枝（4 (https://arxiv.org/html/2501.15316#bib.bib2)；48 (https://arxiv.org/html/2501.15316#bib.bib6)；18 (https://arxiv.org/html/2501.15316#bib.bib20)）和动态剪枝（25 (https://arxiv.org/html/2501.15316#bib.bib15)；10 (https://arxiv.org/html/2501.15316#bib.bib19)；2 (https://arxiv.org/html/2501.15316#bib.bib3)；14 (https://arxiv.org/html/2501.15316#bib.bib4)）。静态剪枝根据与输入无关的重要性度量移除参数。例如，LLM-Pruner（45 (https://arxiv.org/html/2501.15316#bib.bib12)）使用基于梯度的标准消除非必要的耦合结构。结构剪枝的问题在于它常常在稠密模型与之相比时产生显著的性能差距（45 (https://arxiv.org/html/2501.15316#bib.bib12)；5 (https://arxiv.org/html/2501.15316#bib.bib11)）。相反，动态剪枝根据依赖于输入的度量移除权重。动态剪枝的早期尝试（25 (https://arxiv.org/html/2501.15316#bib.bib15)；10 (https://arxiv.org/html/2501.15316#bib.bib19)）聚焦于卷积神经网络，其中针对输入样本选择性地激活通道。最近的工作，如 D-LLM（67 (https://arxiv.org/html/2501.15316#bib.bib5)），通过根据输入 token 选择性跳过层，将条件计算的概念引入 LLMs。动态剪枝方法的问题在于，对于不同的输入，它们没有固定的计算预算，这在小批量设置或预填充阶段服务 LLMs 时会产生问题。另一方面，我们的方法将稠密 LLM 转换为具有固定每 token 预算的稀疏 MoE 模型。另一条研究路线将上下文稀疏性应用于 LLMs（43 (https://arxiv.org/html/2501.15316#bib.bib62)；72 (https://arxiv.org/html/2501.15316#bib.bib63)；36 (https://arxiv.org/html/2501.15316#bib.bib64)），即根据某些条件选择性地激活神经元。尽管取得了一些有希望的结果，但由于其不规则的稀疏模式，通常更难实现更好的推理效率。相比之下，MoE 模型在系统层面有更全面的支持，使其成为扩展模型的热门选择。因此，我们的方法主要关注将稠密模型转换为 MoE 模型。  
**MoE**：稀疏混合专家（MoE）模型通过保持甚至增强模型容量，同时不按比例增加计算成本，改进了纯结构剪枝。例如，Sparsely-Gated MoE（58 (https://arxiv.org/html/2501.15316#bib.bib38)）采用可训练的门控网络为每个输入选择一小部分专家，使模型能够高效地扩展到数千个专家（38 (https://arxiv.org/html/2501.15316#bib.bib39)）。更近期的方法，如 DeepSeekMoE（13 (https://arxiv.org/html/2501.15316#bib.bib40)）进一步解决了专家专业化问题，以相似数量的激活参数匹配了稠密模型的性能。以前从稠密模型构建 MoE 的方法（71 (https://arxiv.org/html/2501.15316#bib.bib61)；37 (https://arxiv.org/html/2501.15316#bib.bib60)；73 (https://arxiv.org/html/2501.15316#bib.bib41)）将专家构建和路由器训练分成两个独立阶段，常常导致次优性能。相比之下，我们的方法将专家构建直接整合到剪枝过程中，将其视为与路由器学习统一的步骤，从而在无需微调的情况下显著提高了性能。

## 3 ToMoE  
图 2：ToMoE 对 MLP 层使用 top-1 路由，对 MHA 层沿头维度使用静态和动态剪枝。  
大多数最新的 LLMs，如 GPT（54 (https://arxiv.org/html/2501.15316#bib.bib42)）、LLaMA（64 (https://arxiv.org/html/2501.15316#bib.bib9)）等，采用解码器唯一架构，因此我们的方法专注于解码器唯一架构。一个典型的解码器块由多头注意力（MHA）层和多层感知机（MLP）层组成。为清晰起见，我们用 T 表示序列长度，d 表示隐藏维度，d_{mid} 表示 MLP 中间维度，H 表示注意力头数。为降低解码器唯一架构的计算成本，我们提出将原始模型转换为 MoE 模型。对于 MHA 层，我们沿头维度 d/H 使用 top-K 路由和静态剪枝。MHA 层的 top-K 路由和静态剪枝确保在预填充或模型服务期间，所有 token 保持相同的头维度，从而支持并行处理。对于 MLP 层，我们的方法沿 MLP 中间维度 d_{mid} 将其转换为 MoE 层，并采用 top-1 路由。我们的方法与之前动态剪枝方法的一个关键区别是，转换后的模型对所有输入保持一致的计算成本。这一特性对于高效处理至关重要。  

### 3.1 专家嵌入  
受近期使用超网络（26 (https://arxiv.org/html/2501.15316#bib.bib17)；22 (https://arxiv.org/html/2501.15316#bib.bib56)；24 (https://arxiv.org/html/2501.15316#bib.bib43)）生成剪枝决策的成功启发，我们采用超网络生成专家嵌入：  
**E_all = HN(z)**，  (1)  
其中 z 是超网络的输入，从随机分布中采样，**E_all = [E_1, ..., E_l, ..., E_L]** 包含所有层的嵌入，且 **E_l ∈ R^{N × d_e}**，其中 N 是专家数量，d_e 是专家嵌入维度。每个嵌入 **E_{l,i}** 将用于生成专家的配置。超网络生成 **E_all** 的目的是引入跨不同层和操作的层间依赖关系。这种设计在实践中已被证明可以加速学习过程（24 (https://arxiv.org/html/2501.15316#bib.bib43)）。更多细节见附录 A (https://arxiv.org/html/2501.15316#A1)。  

### 3.2 专家构建  
本节将讨论如何从 MLP 层构建专家。在解码器层中，MLP 的公式为：  
**f_MLP(X) = σ(X W_G) ⊙ (X W_U) W_D**，  
其中矩阵 **W_U ∈ R^{d × d_{mid}}**, **W_G ∈ R^{d × d_{mid}}** 和 **W_D ∈ R^{d_{mid} × d}** 分别表示上投影、门控投影和下投影矩阵。此外，σ 表示非线性激活函数，⊙ 表示哈达玛积（逐元素乘积）。  
假设目标是使用 N 个专家，在结构剪枝的设定下，每个专家可表示为：  
**f_{MLP}^i(X_t) = σ(X_t W_G S_i) ⊙ (X_t W_U S_i) S_i^⊤ W_D**,  (2)  
其中 i=1, ..., N, 且 **S_i = Diag(s_i)** (**s_i ∈ R^{d_{mid}}**)。