★ FEATURED ARTICLE 为什么知识混合、知识重复会阻碍专家形成“专业领域”直观理解MoE 的理想状态是每个专家成为一个“专科医生”比如专家 A 专攻数学、专家 B 专攻历史、专家 C 专攻美术。但实际训练中会出现两个问题知识混合路由器把语义差异很大的 token 分给同一个专家。这个专家被迫同时学数学、历史、美术结果每门都学不精无法形成专业特长。知识重复多个专家都在处理相似语义的 token大家重复学习相同的知识专家之间没有分工浪费了模型容量。两者叠加MoE 退化成“多个平庸的稠密模型”理论上“参数多、容量大”的优势无法发挥。例子设 4 个 token语义类别为t0: 数学 t1: 历史 t2: 美术 t3: 数学3 个专家TC 路由 Top-1。路由得分tokenE0E1E2选中t0 (数学)0.80.10.1E0t1 (历史)0.70.20.1E0t2 (美术)0.10.60.3E1t3 (数学)0.30.20.5E2问题分析知识混合E0 同时处理 t0数学和 t1历史E0 被迫成为“数学历史”混合专家难以专精。知识重复t0数学分给 E0t3数学分给 E2E0 和 E2 都在学数学重复劳动没有分工。这就是 MoE 训练中需要解决的问题。下面看 DeepSpeed-MoE 和 Switch Transformer 分别怎么处理。DeepSpeed-MoE 怎么处理DeepSpeed-MoE 从三个层面解决参数效率、蒸馏压缩、系统优化。参数效率提升PR-MoE 与 MoSPR-MoE 的核心设计PR-MoEPyramid Residual MoE有两大改动固定 MLP 专家残差纠错每个 token 都先经过一个共享的固定 MLP专家只学习“残差修正量”而不是从零学习完整变换。这样专家不需要重复学习基础特征只需专注于自己的“专业修正”减少知识重复。金字塔式专家数量浅层专家少深层专家多。因为浅层 token 表示还不稳定语义结构不清晰放太多专家容易造成知识混合深层语义已经分化可以容纳更多专家做精细分工。数值例子残差纠错设固定 MLP 对 token 的输出为h_fixed [1.0, 2.0]两个专家对同一 token 的残差输出Expert A: r_A [0.5, -0.3] Expert B: r_B [0.1, 0.8]路由权重Top-1 选 Expert A权重 0.7最终输出 h_fixed 0.7 * r_A [1.0, 2.0] 0.7 * [0.5, -0.3] [1.0 0.35, 2.0 - 0.21] [1.35, 1.79]对比传统 MoE专家要从零学习完整输出容易重复学习基础特征PR-MoE 中专家只学残差分工更清晰减少知识重复。数值例子金字塔式专家数量设 6 层模型浅层第1-2层每层 2 个专家 深层第3-6层每层 8 个专家浅层总专家数2×242×24深层总专家数4×8324×832。参数分布向深层倾斜。浅层 token 语义未分化少专家避免知识混合深层语义丰富多专家精细分工。MoS 蒸馏压缩MoSMixture-of-Students用分阶段知识蒸馏把 PR-MoE 压缩成更浅的学生模型前期用教师信号蒸馏稳定学习教师分布后期关闭蒸馏只优化语言模型损失让学生自主泛化。效果模型尺寸再缩小3.7 倍推理速度比同质量稠密模型更快。数值例子简化教师模型 12 层学生模型 6 层。前期损失L0.5LCE0.5LKD\mathcal{L} 0.5\mathcal{L}_{\text{CE}} 0.5\mathcal{L}_{\text{KD}}L0.5LCE0.5LKD学生快速逼近教师分布后期损失LLCE\mathcal{L} \mathcal{L}_{\text{CE}}LLCE学生自主优化避免欠拟合。系统优化升级DeepSpeed-MoE 在系统层做了三件事自适应并行不同层专家数量不同灵活组合专家并行、专家切片、数据并行、张量切片保证每层获得最合适的并行方式。张量切片降低 All-to-All 复杂度从O(N2)O(N^2)O(N2)降到O(N)O(N)O(N)。分层 All-to-All先节点内高速通信再跨节点交换减少跨机器通信量。显式数据布局转换把 MoE 稀疏重排改写为显式布局转换关键内核延迟下降 6 倍以上。数值例子All-to-All 通信量设 64 块 GPU每块 GPU 有MMM个 token 需要发送给其他 GPU。原始 All-to-All每块 GPU 向其他 63 块各发一份通信量约O(N2)64×644096O(N^2) 64 \times 64 4096O(N2)64×644096单位。张量切片后每块 GPU 只发切片通信量约O(N)64O(N) 64O(N)64单位。分层 All-to-All节点内 8 块 GPU 先做高速 All-to-All通信量8×8648 \times 8 648×864跨 8 个节点再做一次通信量8×8648 \times 8 648×864总通信量远小于全局 4096。延迟例子原内核延迟 6 ms显式布局转换后 1 ms下降 6 倍。Switch Transformer 怎么处理Switch Transformer 的核心是Top-1 路由 超容量处理 稳定性技巧用极简设计降低知识混合和重复。路由器计算对每个 token 表示xxx路由器计算logitsxWr \text{logits} xW_rlogitsxWrpsoftmax(logits) p \text{softmax}(\text{logits})psoftmax(logits)ppp是每个专家的概率分布。数值例子设 2 个专家token t0 的表示x[1,2]x [1, 2]x[1,2]路由器权重W_r [[1, 0], [0, 1]]则logits [1*1 2*0, 1*0 2*1] [1, 2]Softmax:e12.718,e27.389 e^1 2.718, \quad e^2 7.389e12.718,e27.389p[2.718/10.107,7.389/10.107][0.269,0.731] p [2.718/10.107, 7.389/10.107] [0.269, 0.731]p[2.718/10.107,7.389/10.107][0.269,0.731]实际路由决策Top-1Switch Transformer 采用Top-1每个 token 只被分配给得分最高的单个专家。上例中t0 被分配给专家 1概率 0.731。为什么 Top-1 能缓解知识混合/重复每个 token 只进一个专家专家收到的 token 语义更集中减少“一个专家同时学多种知识”的混合问题实现简单跨设备通信少专家不会同时被大量调用softmax 概率主要用于辅助损失统计前向计算只用选中的专家。对比 Top-2Top-2 会让一个 token 进两个专家虽然表达能力强但专家收到的 token 语义更杂知识混合风险更高通信量也更大。数值例子4 个 token2 个专家Top-1 路由tokenlogitssoftmax选中专家t0[2.0, 0.5][0.817, 0.183]E0t1[1.5, 0.8][0.669, 0.331]E0t2[0.3, 2.1][0.139, 0.861]E1t3[2.2, 0.4][0.859, 0.141]E0E0 收到 t0、t1、t3E1 收到 t2。超容量处理Expert Capacity Overflow给每个专家设容量上限expert_capacitytokens_per_batchnum_experts×capacity_factor \text{expert\_capacity} \frac{\text{tokens\_per\_batch}}{\text{num\_experts}} \times \text{capacity\_factor}expert_capacitynum_expertstokens_per_batch×capacity_factor超出容量的 token不执行该专家 FFN只通过残差传到下一层也不产生该专家的梯度。数值例子4 个 token2 个专家capacity_factor 1.0capacity42×1.02 \text{capacity} \frac{4}{2} \times 1.0 2capacity24×1.02每个专家最多处理 2 个 token。上例中E0 被分配 t0、t1、t3共 3 个超过容量 2按路由概率排序t3 (0.859) t0 (0.817) t1 (0.669)E0 保留 t3、t0t1被丢弃t1 只通过残差传递E1 收到 t2容量 2未超。作用防止单个专家过载避免“热门专家处理一切”被丢弃的 token 不产生该专家梯度减少知识重复残差保证信息不完全丢失。Router Z-loss为避免路由 logits 在低精度下出现极端值Switch Transformer 引入 Router Z-lossLz1B∑i1B(log∑j1Eelogitsi,j)2 \mathcal{L}_z \frac{1}{B} \sum_{i1}^{B} \left( \log \sum_{j1}^{E} e^{\text{logits}_{i,j}} \right)^2LzB1i1∑B(logj1∑Eelogitsi,j)2它惩罚 logits 幅度过大降低 softmax 对极端输入的敏感性。数值例子token Alogits [2.0, 0.5]Ze2.0e0.57.3891.6499.038 Z e^{2.0} e^{0.5} 7.389 1.649 9.038Ze2.0e0.57.3891.6499.038logZ2.202 \log Z 2.202logZ2.202Lz(2.202)24.849 \mathcal{L}_z (2.202)^2 4.849Lz(2.202)24.849token Blogits [0.5, 0.3]Ze0.5e0.31.6491.3502.999 Z e^{0.5} e^{0.3} 1.649 1.350 2.999Ze0.5e0.31.6491.3502.999logZ1.098 \log Z 1.098logZ1.098极端 logits 的 z-loss 更大训练时会被惩罚促使 logits 保持温和softmax 分布不会过度尖锐路由更稳定。较小的初始化路由器随机初始化困难。Switch Transformer 从截断正态分布中抽取权重均值为 0适当减小初始化尺度。数值直觉设权重初始标准差为σ\sigmaσ输入维度ddd。logits 的方差约为Var(logits)≈d⋅σ2⋅Var(x) \text{Var}(\text{logits}) \approx d \cdot \sigma^2 \cdot \text{Var}(x)Var(logits)≈d⋅σ2⋅Var(x)σ\sigmaσ太大→\rightarrow→logits 方差大→\rightarrow→softmax 极端→\rightarrow→早期路由剧烈波动σ\sigmaσ太小→\rightarrow→logits 接近 0→\rightarrow→softmax 均匀→\rightarrow→路由无区分度截断正态 小尺度→\rightarrow→logits 方差适中早期梯度稳定。数值例子设d512d 512d512Var(x)1\text{Var}(x) 1Var(x)1。σ0.1\sigma 0.1σ0.1Var(logits)≈512×0.015.12\text{Var}(\text{logits}) \approx 512 \times 0.01 5.12Var(logits)≈512×0.015.12标准差约 2.26softmax 较温和σ1.0\sigma 1.0σ1.0Var(logits)≈512\text{Var}(\text{logits}) \approx 512Var(logits)≈512标准差约 22.6softmax 极端尖锐训练不稳定。所以小初始化能降低早期梯度方差提高稳定性。先 Self-Attention后 FFN/MoESwitch Transformer 子层顺序是Self-Attention → FFN/MoE为什么这个顺序重要先 Self-Attention 建立全局语义每个 token 与上下文所有 token 交互生成包含丰富上下文的情境化表示。后 FFN/MoE 做专家特征增强路由器基于已经融合上下文的信息分配专家token 语义更清晰专家分配更准确减少知识混合。如果颠倒FFN/MoE 先处理缺乏上下文的原始 embedding非线性变换可能扰乱 token 之间的几何关系Self-Attention 的相似度计算准确度下降最终削弱模型表达能力。数值例子设两个 tokenx [1.0, 1.0] y [1.0, 1.1]它们很相似余弦相似度接近 1。如果先 FFN假设 FFN 是f(v)[v12,v22]f(v) [v_1^2, v_2^2]f(v)[v12,v22]这种非线性f(x) [1.0, 1.0] f(y) [1.0, 1.21]相似度仍然较高但数值被放大。如果 FFN 更极端比如f(v)[max(0,v1−1.05),max(0,v2−1.05)]f(v) [\max(0, v_1 - 1.05), \max(0, v_2 - 1.05)]f(v)[max(0,v1−1.05),max(0,v2−1.05)]f(x) [0, 0] f(y) [0, 0.05]x 变成全零y 只剩一点原本相似的 token 几何关系被破坏Self-Attention 再计算相似度就失真了。先 Self-Attention 后 FFNtoken 先在原始空间计算相似度关系准确FFN 再对融合后的表示做非线性增强不会破坏注意力已经建立的依赖关系。总结对比问题DeepSpeed-MoE 的处理Switch Transformer 的处理知识混合PR-MoE 残差纠错专家只学修正量金字塔式浅层少专家Top-1 路由每个 token 只进一个专家专家语义更集中知识重复固定 MLP 共享基础特征专家不重复学基础超容量处理丢弃超额 token减少重复梯度负载均衡自适应并行、分层 All-to-All辅助负载均衡损失 超容量训练稳定系统优化、显式布局转换Router Z-loss、较小初始化、先 Attention 后 FFN推理加速MoS 蒸馏压缩 3.7 倍内核延迟降 6 倍Top-1 简化通信稀疏激活一句话DeepSpeed-MoE 通过“残差纠错 金字塔专家 蒸馏 系统优化”让专家分工更清晰、训练更便宜Switch Transformer 通过“Top-1 路由 超容量 z-loss 小初始化 先 Attention 后 FFN”用极简设计减少知识混合和重复同时保持训练稳定。两者分别代表了现代 MoE 在“结构效率”和“路由简化”
阅读完成 · 觉得有帮助？
