15 KiB
DeepSeekMoE:迈向混合专家语言模型的极致专家专精
原文标题:DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models arXiv:2401.06066 | 发布:11 Jan 2024 作者/机构:Damai Dai、Chengqi Deng、Chenggang Zhao 等 + DeepSeek-AI(深度求索),联合北京大学、清华大学、南京大学 原文 PDF:02-DeepSeekMoE-2401.06066.pdf
一、摘要(全文翻译)
在大语言模型时代,混合专家(Mixture-of-Experts, MoE)是在扩大模型参数规模时管理计算成本的一种颇具前景的架构。然而,像 GShard 这类从 N 个专家中激活 top-K 个专家的传统 MoE 架构,在确保专家专精(expert specialization,即每个专家获取互不重叠且聚焦的知识)方面面临挑战。为此,我们提出了面向极致专家专精的 DeepSeekMoE 架构。它包含两项主要策略:(1)将专家细粒度地切分为 mN 个专家,并从中激活 mK 个,从而实现更灵活的激活专家组合;(2)隔离出 K_{s} 个专家作为共享专家,旨在捕捉通用知识并缓解路由专家中的冗余。从 2B 参数的适度规模出发,我们证明 DeepSeekMoE 2B 达到了与 GShard 2.9B 相当的性能,而后者拥有 1.5 倍的专家参数与计算量。此外,DeepSeekMoE 2B 还几乎逼近了与其总参数量相同的稠密对照模型的性能,而后者构成了 MoE 模型的性能上界。随后,我们将 DeepSeekMoE 扩展到 16B 参数,并表明它仅用约 40% 的计算量就达到了与 LLaMA2 7B 相当的性能。进一步,我们将 DeepSeekMoE 扩展到 145B 参数的初步尝试持续验证了该架构相较 GShard 架构的显著优势,并显示它仅用 28.5%(甚至可能 18.2%)的计算量就取得了与 DeepSeek 67B 相当的性能。
二、核心贡献
- 架构创新:提出 DeepSeekMoE —— 一个以「极致专家专精」为目标的 MoE 架构,通过细粒度专家切分(fine-grained expert segmentation)与共享专家隔离(shared expert isolation)两项策略,在不增加专家总参数量与计算量的前提下提升专家的专业化程度。
- 实证验证:在 12 个 zero-shot / few-shot 基准上做验证实验,DeepSeekMoE 2B 大幅超越同规模 GShard 2B,并追平拥有 1.5 倍专家参数与计算量的 GShard 2.9B;同时几乎逼近同总参数量稠密模型的性能,即 MoE 模型的理论上界。
- 可扩展性:将模型扩展到 16B 总参数并在 2T tokens 上训练,DeepSeekMoE 16B 仅用约 40% 计算量即达到 DeepSeek 7B 与 LLaMA2 7B 的相当水平;进一步将 145B 版本的初步尝试显示其持续优于 GShard 架构,并以 28.5%(甚至 18.2%)的计算量对标 DeepSeek 67B。
- MoE 对齐:对 DeepSeekMoE 16B 做监督微调(SFT)构建聊天模型,DeepSeekMoE Chat 16B 在聊天场景下与 DeepSeek Chat 7B、LLaMA2 SFT 7B 相当,说明 MoE 模型同样能从指令微调中获益,且保持了良好的通用性与适应性。
- 公开发布:开源 DeepSeekMoE 16B 的模型 checkpoint;该模型无需量化即可部署在单张 40GB 显存的 GPU 上。
三、方法与技术要点
3.1 通用 MoE 架构(预备知识)
论文先给出 Transformer 语言模型中通用的 MoE 范式作为基线。标准 Transformer 语言模型由 L 层 Transformer 块堆叠而成,第 l 层的自注意力与 FFN 计算为:
\mathbf{u}_{1:T}^{l} = \operatorname{Self-Att}\left(\mathbf{h}_{1:T}^{l-1}\right) + \mathbf{h}_{1:T}^{l-1}, \qquad \mathbf{h}_{t}^{l} = \operatorname{FFN}\left(\mathbf{u}_{t}^{l}\right) + \mathbf{u}_{t}^{l}
构造 MoE 语言模型的典型做法,是按一定间隔把 Transformer 中的 FFN 替换为 MoE 层:每个 MoE 层包含多个结构上与标准 FFN 完全相同的专家,每个 token 被分配给 1 个(Switch Transformer)或 2 个(GShard)专家。第 l 个 FFN 被替换为 MoE 层后,输出隐藏状态为:
\mathbf{h}_{t}^{l} = \sum_{i=1}^{N}\left(g_{i,t}\operatorname{FFN}_{i}\left(\mathbf{u}_{t}^{l}\right)\right) + \mathbf{u}_{t}^{l}
g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t}\in\operatorname{Topk}\left(\{s_{j,t}\mid 1\leqslant j\leqslant N\},K\right),\\ 0, & \text{otherwise},\end{cases} \qquad s_{i,t} = \operatorname{Softmax}_{i}\left({\mathbf{u}_{t}^{l}}^{T}\mathbf{e}_{i}^{l}\right)
其中 N 为专家总数,\mathbf{e}_{i}^{l} 是第 l 层第 i 个专家的质心(centroid),s_{i,t} 为 token 与专家的亲和度,g_{i,t} 为门控值。g_{i,t} 的稀疏性(N 个门控值中只有 K 个非零)保证了 MoE 层的计算效率——每个 token 只在 K 个专家中被计算。
为什么需要新架构:论文指出这类传统架构存在两个痛点,共同限制了专家专精。(1)知识混杂(knowledge hybridity):常用 MoE 只用 8 或 16 个专家,单个专家被分配到的 token 覆盖面过广,导致它不得不在同一份参数里塞进彼此差异极大的知识,难以同时用好;(2)知识冗余(knowledge redundancy):不同专家接收到的 token 都需要某些通用知识,于是多个专家会各自重复学习这些共享知识,造成专家参数浪费。二者叠加使现有 MoE 无法逼近其理论性能上界。
3.2 细粒度专家切分
针对知识混杂,论文的做法是:在保持专家参数量与计算成本不变的前提下,把每个专家 FFN 的中间隐藏维度压缩到原来的 $1/m$,从而把一个专家切成 m 个更小的专家;由于每个专家变小了,相应地把激活专家数也提高到 m 倍,以维持相同的计算量。公式上只需把 top-K 的求和范围从 N 扩到 $mN$、把激活数从 K 扩到 $mK$:
\mathbf{h}_{t}^{l} = \sum_{i=1}^{mN}\left(g_{i,t}\operatorname{FFN}_{i}\left(\mathbf{u}_{t}^{l}\right)\right) + \mathbf{u}_{t}^{l}, \qquad g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t}\in\operatorname{Topk}\left(\{s_{j,t}\mid 1\leqslant j\leqslant mN\},mK\right),\\ 0, & \text{otherwise},\end{cases}
其收益可以从组合数上直观理解:当 N=16 时,典型 top-2 路由只有 \binom{16}{2}=120 种激活组合;若把每个专家切成 4 个小专家,细粒度路由的可能组合数变成 $\binom{64}{8}=4{,}426{,}165{,}368$。组合灵活性的暴涨意味着模型能更精确、更有针对性地组装所需知识,而落在每个小专家上的知识类型更单一,专家因此能保持更高的专业化程度。
3.3 共享专家隔离
针对知识冗余,论文隔离出 K_{s} 个专家作为共享专家:不管路由器怎么算,每个 token 都会被确定性地分配给这些共享专家,由它们专门捕捉并沉淀跨上下文通用的知识。为了保证计算量不变,路由专家中被激活的专家数相应减少 K_{s} 个。完整的 DeepSeekMoE 层可写为:
\mathbf{h}_{t}^{l} = \sum_{i=1}^{K_{s}}\operatorname{FFN}_{i}\left(\mathbf{u}_{t}^{l}\right) + \sum_{i=K_{s}+1}^{mN}\left(g_{i,t}\operatorname{FFN}_{i}\left(\mathbf{u}_{t}^{l}\right)\right) + \mathbf{u}_{t}^{l}
g_{i,t} = \begin{cases} s_{i,t}, & s_{i,t}\in\operatorname{Topk}\left(\{s_{j,t}\mid K_{s}+1\leqslant j\leqslant mN\},mK-K_{s}\right),\\ 0, & \text{otherwise},\end{cases} \qquad s_{i,t} = \operatorname{Softmax}_{i}\left({\mathbf{u}_{t}^{l}}^{T}\mathbf{e}_{i}^{l}\right)
于是在 DeepSeekMoE 中:共享专家数为 $K_{s}$,路由专家总数为 $mN-K_{s}$,非零门控数为 $mK-K_{s}$。作者说明,共享专家隔离的原型可追溯到 Rajbhandari et al. (2022),关键区别在于后者是从工程视角出发,而本文是从算法视角出发,把它作为提升专家专精的手段。
3.4 负载均衡设计
自动学习的路由策略会出现负载不均衡,带来两个缺陷:一是路由崩塌(routing collapse),模型总只选少数几个专家,其余专家得不到充分训练;二是在专家分布到多设备时,负载不均会加重计算瓶颈。
专家级均衡损失用于缓解路由崩塌,定义如下($N'=mN-K_{s}$,$K'=mK-K_{s}$,\mathbb{1}(\cdot) 为指示函数,\alpha_{1} 为专家级均衡因子):
\mathcal{L}_{\mathrm{ExpBal}} = \alpha_{1}\sum_{i=1}^{N'} f_{i}P_{i}, \quad f_{i} = \frac{N'}{K'T}\sum_{t=1}^{T}\mathbb{1}(\text{Token } t \text{ selects Expert } i), \quad P_{i} = \frac{1}{T}\sum_{t=1}^{T}s_{i,t}
设备级均衡损失则针对跨设备计算均衡。作者强调:过度强制专家级均衡会损害模型性能,因此不必在专家层面做严格约束,真正要保证的是设备之间计算量均衡。把全部路由专家划分为 D 组 \{\mathcal{E}_{1},\dots,\mathcal{E}_{D}\} 并每组部署在一个设备上后:
\mathcal{L}_{\mathrm{DevBal}} = \alpha_{2}\sum_{i=1}^{D} f_{i}'P_{i}', \quad f_{i}' = \frac{1}{|\mathcal{E}_{i}|}\sum_{j\in\mathcal{E}_{i}} f_{j}, \quad P_{i}' = \sum_{j\in\mathcal{E}_{i}} P_{j}
实践中的取舍是:专家级均衡因子设得较小(防止路由崩塌即可),设备级均衡因子设得较大(促进设备间计算均衡)。
四、主要实验结果
| 基准 / 任务 | 本模型 | 对比对象 | 说明 |
|---|---|---|---|
| Pile(Loss,验证实验) | DeepSeekMoE 2B 1.808 | GShard 2B 1.867;GShard ×1.5 1.808;Dense ×16 1.806 | 100B tokens;DeepSeekMoE 逼近稠密上界 |
| HellaSwag(0-shot Acc.) | 54.8 | GShard 2B 50.5;GShard ×1.5 54.4 | 同总参数、同激活参数下领先 |
| HumanEval(Pass@1) | 4.9 | GShard 2B 3.7;GShard ×1.5 3.0;Dense ×16 4.3 | 验证实验中代码生成最强 |
| TriviaQA(5-shot EM) | 16.6 | GShard 2B 10.2;GShard ×1.5 15.7 | 知识密集型任务优势明显 |
| Pile(BPB,16B) | DeepSeekMoE 16B 0.74 | DeepSeek 7B 0.75;LLaMA2 7B 0.76 | 2T tokens;算力为 DeepSeek 7B 的 40.5%、LLaMA2 7B 的 39.6% |
| HellaSwag(0-shot) | 77.1 | DeepSeek 7B 75.4;LLaMA2 7B 75.6 | 超过两个 7B 稠密模型 |
| GSM8K(8-shot EM) | 18.8 | DeepSeek 7B 17.4;LLaMA2 7B 15.5 | 数学推理领先 |
| MATH(4-shot EM) | 4.3 | DeepSeek 7B 3.3;LLaMA2 7B 2.6 | 数学推理领先 |
| HumanEval(Pass@1) | 26.8 | DeepSeek 7B 26.2;LLaMA2 7B 14.6 | 代码生成大幅超 LLaMA2 7B |
| MBPP(3-shot Pass@1) | 39.2 | DeepSeek 7B 39.0;LLaMA2 7B 21.8 | 同上 |
| MMLU(5-shot Acc.) | 45.0 | DeepSeek 7B 48.2;LLaMA2 7B 45.8 | 落后:注意力参数仅约 0.5B(DeepSeek 7B 为 2.5B) |
| CMMLU(5-shot Acc.) | 42.5 | LLaMA2 7B 32.6;DeepSeek 7B 47.2 | 中文能力远超 LLaMA2,但低于 DeepSeek 7B |
| CHID(0-shot Acc.) | 89.4 | LLaMA2 7B 37.9;DeepSeek 7B 89.3 | 中文成语理解近乎碾压 LLaMA2 |
| MBPP(Chat,3-shot Pass@1) | DeepSeekMoE Chat 16B 46.2 | LLaMA2 SFT 7B 27.8;DeepSeek Chat 7B 39.0 | 同一份 SFT 数据,聊天模型仍保持代码优势 |
| DROP(Chat,1-shot EM) | 33.8 | LLaMA2 SFT 7B 40.0;DeepSeek Chat 7B 41.7 | 落后:长文阅读理解是短板 |
| Pile(Loss,145B) | DeepSeekMoE 145B 1.876 | DeepSeek 67B 1.905;GShard 137B 1.961 | 245B tokens 初步训练;算力为 DeepSeek 67B 的 28.5% |
| GSM8K(8-shot EM,Half Activated) | DeepSeekMoE 142B 13.8 | DeepSeek 67B 11.8;GShard 137B 6.4 | 仅激活一半专家参数、18.2% 算力仍超 DeepSeek 67B |
总体看,DeepSeekMoE 在语言建模与知识密集型任务(Pile、HellaSwag、TriviaQA、NaturalQuestions)上优势最突出——作者归因于 MoE 模型中 FFN 参数远多于注意力参数,而 Transformer 的 FFN 承担知识记忆功能。在代码生成与数学推理上,16B 模型全面超越 LLaMA2 7B(HumanEval 26.8 vs 14.6,MBPP 39.2 vs 21.8),部分归功于预训练语料中数学与代码文本更丰富。多选题任务(MMLU、CEval、CMMLU)是持续短板:16B 与 145B 版本均落后 DeepSeek 同代稠密模型(16B 对 DeepSeek 7B:MMLU 45.0 vs 48.2、CEval 40.6 vs 45.0;145B 对 DeepSeek 67B:MMLU 39.4 vs 45.1),论文明确将其归因于注意力参数容量不足。此外,SFT 之后 DeepSeekMoE 与 DeepSeek 7B 在多选题上的差距有所收窄。
五、结论、局限与影响
- 结论:DeepSeekMoE 通过细粒度专家切分与共享专家隔离,显著提升了专家专精程度。2B 规模的验证实验证明该架构能逼近 MoE 模型的理论性能上界;消融实验显示两种策略各自都带来提升,且专家切分越细性能越好;利用率分析显示禁用 top 路由专家时 DeepSeekMoE 比 GShard ×1.5 更敏感(说明专家冗余更低、更不可替代),直接禁用共享专家会让 Pile loss 从 1.808 恶化到 2.414(说明共享专家捕捉的是路由专家无法覆盖的基础知识)。16B 版本在 2T tokens 上训练后以约 40% 计算量对标 7B 稠密模型,经 SFT 对齐后仍保持该优势;145B 的初步结果持续优于 GShard,并以 28.5%(半激活版本 18.2%)计算量对标 DeepSeek 67B。
- 局限性:论文自述的局限包括:(1)在多选题任务(MMLU、CEval、CMMLU 类)上表现不足,根源是 DeepSeekMoE 16B 的注意力参数仅约 0.5B(DeepSeek 7B 为 2.5B),作者此前在 DeepSeek 7B 上的研究发现注意力容量与多选题表现正相关;(2)DeepSeekMoE 145B 只是「初步探索」,仅训练 245B tokens,尚未完成最终版本与完整训练,因此其结论属阶段性验证;(3)在 16B 规模未采用更细的切分粒度(专家为 0.25 倍标准 FFN),原因是专家过小会降低计算效率,论文称在 16B 以上的更大规模仍可采用更细粒度。此外论文正文未对 145B 的最终目标训练量给出明确数字。
- 影响:该工作给出了 MoE 走向「极致专家专精」的两条正交设计原则,并以极小代价(同参数、同计算量)获得显著收益,其中「细粒度 + 共享专家」的组合后来成为 DeepSeek-V2/V3 系列 MoE 的基础架构;它同时验证了 MoE 模型可从 SFT 对齐中获益(与早期「MoE 难从微调获益」的观察相反),并开源了可单卡 40GB 部署、推理速度约为 7B 稠密模型 2.5 倍的 16B checkpoint,对工业界的低成本大模型部署有直接参考价值。
六、关键术语对照表
| English | 中文 |
|---|---|
| Mixture-of-Experts (MoE) | 混合专家 |
| Expert Specialization | 专家专精 / 专家专业化 |
| Fine-Grained Expert Segmentation | 细粒度专家切分 |
| Shared Expert Isolation | 共享专家隔离 |
| Shared Expert | 共享专家 |
| Routed Expert | 路由专家 |
| Knowledge Hybridity | 知识混杂 |
| Knowledge Redundancy | 知识冗余 |
| Routing Collapse | 路由崩塌 |
| Expert-Level Balance Loss | 专家级均衡损失 |
| Device-Level Balance Loss | 设备级均衡损失 |
| Expert Parallelism | 专家并行 |
| Activated Parameters | 激活参数 |
| Bits per Byte (BPB) | 每字节比特数 |
| Top-K Routing | Top-K 路由 |
| Feed-Forward Network (FFN) | 前馈网络 |
| Dense Counterpart | 稠密对照模型 |