密集训练、稀疏推理：混合专家语言模型的训练机制再思考
Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models
Pan, Bowen, Yikang Shen, Haokun Liu, Mayank Mishra, Gaoyuan Zhang, Aude Oliva, Colin Raffel, Rameswar Panda
arXiv (Cornell University) · 2024-04-08
摘要
混合专家（MoE）语言模型在不牺牲性能的前提下，可将计算成本降低至密集模型的1/2至1/4，因此在计算受限场景下更具效率。不过，为达到与密集模型相当的性能，MoE模型通常需要多2至4倍的参数，这导致其GPU显存需求更高，在自回归生成这类I/O受限场景中效率反而更低。本研究针对MoE模型提出了一种密集训练与稀疏推理相结合的框架（DS-MoE），该框架在训练阶段对所有专家模块采用密集计算方式，推理阶段则采用稀疏计算模式，从而兼顾计算效率与参数效率。我们在大语言模型训练上的实验表明，DS-MoE模型比传统稀疏MoE模型更具参数效率，其总参数规模与性能虽与密集模型相当，但计算开销更低（仅需激活30%至40%的模型参数）。借助vLLM进行的性能测试显示，DS-MoE-6B模型的运行速度比Mistral-7B等同类密集模型快1.86倍，较DeepSeekMoE-16B、Qwen1.5-MoE-A2.7B等同类MoE模型也快1.50倍至1.71倍。
本摘要由 AI（Hy-MT2-Pro）翻译生成，仅供参考，请以原文为准。
AI 综述
基于全文生成- 主要观点
- 论文针对混合专家（MoE）语言模型在 I/O 受限场景（如自回归逐 token 生成）中效率偏低的问题给出诊断：MoE 相比稠密模型可用 2–4× 更少的计算量达到相当性能，但通常需要 2–4× 的参数量才能追平稠密模型，导致显存需求大、生成吞吐差。作者将这种参数低效主要归因于稀疏训练——每个 token 只激活并优化部分专家，路由器与专家均只获得稀疏梯度，并在专家并行下因专家使用不均衡造成 GPU 利用率下降。据此提出 DS-MoE 框架：训练时对每层全部专家做稠密计算与稠密梯度回传，推理时按路由分数做 top-K 或阈值筛选的稀疏激活。为使稀疏推理在训练后可行，引入互信息损失，在最大化专家分布熵（负载均衡）的同时最小化条件熵（单 token 上的专家集中）。实验结论为：DS-MoE 在总参数量与性能上与同规模稠密模型相当，而推理只激活 30%–40% 的参数；参数效率优于标准稀疏 MoE，DS-MoE-3B 与 SMoE-5B 性能及计算量相当但 MLP 专家参数量仅为其一半；模型规模越大对稀疏性的容忍度越高。
- 研究方法
- 在 Pile 数据集子集上预训练（共 300B token；1B 模型使用 30B 子集，3B 与 6B 使用 100B 子集），采用 CodeGen 分词器，序列长度 2048。构建同配方对照实验：Dense-1B/3B/6B、SMoE-1B/1.5B/5B（训练与推理均激活 top-2 专家，switch loss 负载均衡，dMoE 实现）与 DS-MoE-1B/3B/6B，各模型超参以表格逐项列出。评测使用 LM evaluation harness，在 PiQA、HellaSwag、WinoGrande、SciQ、Arc-e、Arc-c 零样本任务与 Wikitext 困惑度上报告结果，并统计跨任务的平均激活参数量与激活隐藏层比例。推理吞吐用 vLLM 实测，与 Mistral-7B、DeepSeekMoE-16B、Qwen1.5-MoE-A2.7B 比较。消融包括 MI 损失权重 α 的两档取值（2e-4 与 4e-4，在 24% 与 29% 两个稀疏度下评测）以及 threshold、TopK、Threshold-TopK 三种专家采样策略的 Wikitext 困惑度—激活参数量对比。架构上以 Mixture-of-Attention（MoA）头替换稠密自注意力层，采用 GQA 在专家间共享键值缓存、按专家分配查询向量。
- 学术贡献
- 方法层面提出“稠密训练—稀疏推理”的混合训练范式，把 MoE 的参数低效问题从架构设计转向训练阶段的梯度传播，并给出路由器完整梯度的形式化改写（∇S = [e1(X),…,eN(X)]ᵀ∇O、∇ej(X)=Sj∇O），与传统掩码式稀疏梯度形成对照。将互信息损失具体化为 L_MI = −H(e) + (1/|X|)Σ H(e|X)，论证其同时满足专家负载均衡与单 token 专家集中两个目标，从而为训练后的稀疏推理提供可部署性约束。架构上把自注意力层也专家化（MoA），在专家间共享 K/V 缓存，使激活参数比例的统计覆盖注意力层与 MLP 层两类模块。实证层面在同训练配方下给出参数效率对照数据：DS-MoE-3B（934M 激活、34%）对标 SMoE-5B（1212M 激活、25%）与 Dense-3B，以及 DS-MoE-6B 在 vLLM 下相对 Mistral-7B 的 1.86×、相对 DeepSeekMoE-16B 与 Qwen1.5-MoE-A2.7B 的 1.50–1.71× 吞吐优势。论文完整列出各模型超参数、α 取值、优化器设置、硬件与训练时长，为 MoE 的高效训练与推理研究提供了可复现的同规模基线设置。
- 存在局限
- 验证规模限于 1B–6B 的中等规模语言模型，论文对 70B 以上模型稀疏度更高的判断仅是基于趋势外推，材料中未提供相应实验。阈值采样虽在 Wikitext 困惑度与激活参数量间取得最佳折中，但同一批次内不同 token 激活的专家数不一致，论文自承其不利于实际部署，须退化为 TopK 或 Threshold-TopK，而这两种策略又不及阈值策略的折中效果。互信息损失权重 α 的作用不稳定：α=4e-4 训练的模型在 24% 与 29% 两个稀疏度上均保持性能，α=2e-4 的模型在 24% 稀疏度下性能下降而在 29% 下最好，说明该超参需按目标稀疏度人工选取，材料未给出选取规则。评测集中于常识推理与问答的零样本准确率及 Wikitext 困惑度，未涉及指令微调、长文本生成与多语言场景。稠密训练需前向计算并保留全部专家输出，其额外训练算力与显存开销未与稀疏训练做量化对比；与 Mistral-7B、DeepSeekMoE-16B、Qwen1.5-MoE-A2.7B 的吞吐比较属跨来源模型对比，训练配方与数据并不一致。层利用率分析部分在提供的材料中仅给出装置说明（在 Wikitext 上统计自注意力层与 MLP 层激活专家比例），未给出结论。
- 进一步研究方向
- 在 30B、70B 及以上规模复现 Dense/SMoE/DS-MoE 的同配方对照训练，检验“模型越大稀疏容忍度越高”的外推是否成立，并拟合激活比例随参数规模变化的函数关系。将固定 K 的 TopK 或 Threshold-TopK 直接作为训练目标而非事后采样策略，测量批推理中 token 间激活专家数的方差及其对吞吐与延迟的影响，验证训练目标与推理采样一致性是否带来增益。对 MI 损失权重 α 开展多档网格扫描，绘制稀疏度—下游性能的帕累托前沿，并尝试按层深或按训练进程自适应的 α 调度，检验其能否免除按目标稀疏度人工选参。把训练 FLOPs、激活显存与推理吞吐纳入统一成本模型，测算稠密训练相对稀疏训练的总成本盈亏平衡点。将评测扩展到指令微调后的生成质量、长上下文与多语言任务，考察仅激活 30%–40% 参数对生成类指标而非零样本准确率与困惑度的影响。补齐并扩展层利用率分析：在同一 DS-MoE-3B 模型上比较自注意力层与 MLP 层的激活专家比例差异，据此检验是否应对两类层设置不同的稀疏度或不同的 K 值。
本综述由 AI（deepseek-flash）阅读开放获取全文自动生成，仅供参考，不代表原文作者观点，请以原文为准。
参考文献
相关论文
话题：Advanced Neural Network Applications