dMoE: 具有可学习块专家的 dLLMs

TL;DR dMoE 通过将扩散语言模型中的 token 级专家路由聚合为块级路由，解决并行解码与 MoE 不匹配导致的激活专家过多问题，在保持几乎全部性能的同时将唯一激活专家数从 69.5 降至 14.6，显著降低内存并加速推理。

问题

问题背景

扩散大语言模型 (dLLMs) 通过迭代式去掩码-重掩码过程，天然支持块级并行解码，绕过了自回归模型的顺序生成瓶颈。为在可控激活参数下扩展容量，近期工作将 dLLM 与 Mixture-of-Experts (MoE) 架构结合，但两者的设计假设存在根本冲突，严重拖累推理效率。

现有方法局限

传统 MoE 层以令牌级路由为核心：每个 token 独立计算专家分布并激活 top-k 专家。然而在 dLLM 中，一次前向传播需同时处理一个 block 内的多个 token（这些 token 有双向依赖），导致：

• 唯一激活专家数激增：block 内不同 token 可能选择不同的专家，一个 block 可能触发数十个唯一专家，远超 MoE 设计的稀疏激活本意。
• 推理变为 memory-bound：每个专家的参数加载与中间激活需要大量显存带宽，批处理无法有效掩盖延迟，端到端延迟大幅增加。

为什么这个问题难且重要

dLLM 的并行解码势能依赖于 block 处理的效率，而 MoE 的稀疏激活则依赖 token 级细粒度路由，两者在路由粒度上天然不匹配。简单降低 top-k 或扩大 block 都会损害生成质量。dMoE 指出，这种不匹配造成 平均唯一激活专家数从 69.5 降至 14.6 的优化空间，直接关联 76.64%–79.84% 的显存节省与 1.14×–1.66× 的延迟加速。对于要在推理服务中部署大规模扩散语言模型的团队而言，该瓶颈若不解决，成本与吞吐将成为落地阻碍。

行业类比

类似在推荐系统中，对一批用户请求逐个路由到不同精排模型，会导致每个模型重复加载、GPU 利用率碎片化；dMoE 的块级路由相当于对请求进行聚类后批量调用同一模型集合，从而大幅减少上下文切换与内存带宽压力。

核心洞察

• dMoE 首次明确揭示并解决了扩散语言模型 (dLLM) 与 MoE 架构集成的根本性矛盾：块并行解码需要一次处理多个具有双向依赖的 token，而传统 MoE 对每个 token 独立路由，导致推理时唯一激活专家数暴增，使推理卡在显存带宽瓶颈。此前的工作要么忽略该不匹配，要么仅在 token 级优化，dMoE 则从块级视角重新设计路由，为 dLLM MoE 的高效部署提供了架构层面的新思路。
• dMoE 的方法极其简洁且实用：它不是增加复杂的门控网络或训练蒸馏，而是将同一块内所有 token 的专家分布直接聚合为一个块级分布，再依此统一路由，大幅减少唯一激活的专家数（69.5→14.6）并保留 99.11% 的性能。这种无训练的“即插即用”策略对工程落地极为友好，可直接替换现有 MoE 层，在保持模型质量的同时将内存占用降低约 78%，端到端延迟加速最多 1.66 倍，显著提升了 dLLM MoE 的推理效率与硬件利用率。

方法

问题背景

Diffusion Large Language Models (dLLMs) 采用块并行解码，一次前向处理多个 token（称为一个 block），且 token 间存在双向依赖，这与自回归模型左到右的顺序生成截然不同。当 dLLM 集成 Mixture-of-Experts (MoE) 架构以扩展容量时，传统 token 级专家路由 与块并行处理之间出现根本性不匹配：每个 token 独立选择专家，导致同一个 block 内激活的专家数量随 block 大小线性增长，唯一激活专家数可能高达 69.5，推理严重受限于内存带宽。

dMoE 方法核心

dMoE 提出一种块级 MoE 路由框架，将专家选择粒度从 token 提升到 block，解决上述不匹配。其工作流程如下：

• 输入：一个 block 内所有 token 的隐藏表示。
• Token 级分布计算：对每个 token 独立计算路由 logits，通过 softmax 得到其对各专家的概率分布（standard token-level routing）。
• 块级聚合（核心模块）：将 block 内所有 token 的专家分布进行聚合（例如加权求和或平均），得到一个统一的 block-level expert distribution。聚合权重可通过可学习参数实现（“Learnable Block Experts”），从而自适应地融合不同 token 的路由偏好。
• 块级专家选择：基于块级分布进行 top-k 选择，确定该 block 最终激活的 k 个专家。
• 共享前向：block 内所有 token 共享同一组激活专家进行后续计算，消除了重复专家实例化。

通过将独立路由变为一致路由，dMoE 将唯一激活专家数从 69.5 降至 14.6，极大降低了内存访问压力。

工程收益与差异

实验证实，dMoE 可保留原始性能的 99.11%，同时减少 76.64%–79.84% 内存占用，端到端延迟加速 1.14×–1.66×。与同类工作的差异：现有 dLLM-MoE 方法均遵循 token 级独立路由，dMoE 首次通过块内分布聚合实现了路由粒度的跃升，无需修改模型结构或训练目标，仅靠路由策略调整即化解了块并行与 token 路由的根本矛盾，是一种轻量、即插即用的方案。

实验

实验设计概述

实验在一系列扩散大语言模型（dLLM）基准上进行，对比 dMoE 与传统 token 级 MoE 的表现。评估重点围绕 推理效率 与 生成质量 两个维度：统计每个前向块中 唯一激活的专家数量，测量 显存峰值 与 端到端延迟，并同步监控 困惑度/下游任务指标 以验证性能损失。实验覆盖多种模型规模与解码配置，确保结论的通用性。

核心发现

• 专家激活大幅收缩：块级路由将平均唯一激活专家从 69.5 降到 14.6，减少近 80%，从根本上缓解了专家切换带来的显存碎片与 I/O 压力。
• 性能几乎无损：在如此高的专家压缩率下，dMoE 仍保留了 99.11% 的原始模型性能，证明块内 token 间存在高度一致的专家偏好，聚合策略是合理的。
• 显存与延迟显著改善：显存占用降低 76.64%–79.84%，端到端延迟获得 1.14×–1.66× 加速。由于每个块只激活少数专家，KV 缓存与权重加载的开销被成倍削减，这对 batch 部署场景 尤为关键。

与基线的深度对比

传统 token 级 MoE 针对 dLLM 的块并行解码存在结构性错配：每个 token 独立路由导致一个块内激活专家集合庞大，推理沦陷于 memory-bound 瓶颈。dMoE 将路由粒度提升至块级别，通过聚合块内 token 的专家分布生成统一的 块级专家分布，然后据此执行 Top-K 路由。这种设计并非简单取平均，而是保留了可学习的聚合权重，允许模型自适应地平衡不同 token 的路由贡献。对比结果显示，块级路由是弥合 dLLM 与 MoE 架构鸿沟的轻量级方案，无需修改模型核心结构，仅在前向流程中引入一个无害的聚合算子即可实现显著的效率提升，同时避免了传统方案中因专家激活过多导致的吞吐塌陷。

行业影响

落地场景

dMoE 直接面向已采用或规划采用 扩散大语言模型 (dLLMs) 与 混合专家 (MoE) 架构的团队。任何需要大容量但低延迟的文本生成业务均可受益，典型场景包括：

• 实时对话系统：客服机器人、语音助手后端，要求首 token 延迟极低。dMoE 显著减少唯一激活专家数，将推理从显存带宽瓶颈中解放，使大容量 dLLM+MoE 模型可部署于更低成本的硬件。
• 批量内容生成：营销文案、技术报告自动生成。扩散模型天然并行解码，但逐 token 专家路由导致显存碎片化，dMoE 以块级路由大幅降低显存占用，提升单卡吞吐量。
• 代码补全与重构：IDE 插件场景要求毫秒级响应，dMoE 加速推理同时保持模型质量，使大型代码模型可集成进轻量级编辑器环境。

商业价值

• 降本：实验显示 显存占用减少 76.64%–79.84%，相同硬件下可服务更大模型或更高并发，直接降低 TCO。对云服务商，可将 GPU 实例规格下探一档，仍满足 SLA。
• 体验提升：端到端延迟 加速 1.14×–1.66×，且性能保留 99.11%。对延迟敏感业务（如实时语音交互、证券数据分析），响应时间的改善可直接转化为用户留存与转化率提升。
• 扩展性：dMoE 作为插件式路由策略，不改动模型权重，可与现有训练好的 dLLM+MoE 检查点无缝结合，降低迁移成本。

与现有产品/工作流的接口

dMoE 本质是对 MoE 门控层的一次块级替换，接口非常轻量：

1. 训练侧：在已有 dLLM 训练流程中，将标准 token-level gating 替换为 block-level expert distribution aggregation。论文提供的代码可直接嵌入主流训练框架（如 PyTorch + HuggingFace），需要的改动仅限于前向传播中增加一个块内概率聚合操作。
2. 推理侧：dMoE 减少唯一激活专家后，降低了显存碎片与操作启动开销。可直接接入现有的 vLLM 或 TensorRT-LLM 等推理引擎，通过自定义 MoE 层或融合算子实现，无需改变服务化流程。
3. 监控与调试：块级路由更集中，专家负载均衡更可预测，便于设计 expert parallelism 策略与容量因子，简化分布式部署的调参。

具体落地 Use Case

• 电商智能导购：大型电商平台使用 dLLM+MoE 提供个性化商品推荐与问答。dMoE 将推理延迟从 500ms 降至 300ms 以内（加速 1.66×），同时允许在单张 A10 卡上运行千亿参数模型，使中小型商户也能负担实时 AI 导购服务。
• 在线教育自适应题库生成：教育科技公司用扩散模型生成不同难度的习题及解析，要求快速响应学生做题行为。dMoE 使模型在生成复杂推理步骤时仍保持低延迟，提升学习体验，同时降低云服务支出。

代码开源 fscdc/dMoE，可与现有 dLLM 工作流快速集成。

局限

• **方法通用性受限**：dMoE 专为 dLLMs 的块并行解码设计，依赖双向依赖的 block 结构进行专家分布聚合。对于标准自回归模型或非块式扩散模型，聚合策略可能不适用，且论文未在其他生成架构上验证。此设计深度耦合扩散掩码修复过程，可能限制其在更广泛 LLM 生态中的推广。
• **性能-效率权衡的边界**：尽管 dMoE 在多数基准上保留了 99.11% 的性能，但实验未详尽报告极端压缩比（如块大小极大或极小）下的表现退化点。实际部署中，块大小选择、专家数量配置对延迟和内存的影响可能需要进一步自动调优，而论文未提供动态块尺寸的自适应策略。
• **训练与推理的额外开销**：论文主要关注推理时的内存与延迟，但引入块级路由需要计算 token 级分布并聚合，可能增加训练时的计算图复杂度和全局 barrier 同步成本。在大规模集群训练中，这种聚合操作的通信开销及其对并行策略的影响尚待评估。