[Paper] MoE-DiffuSeq：使用 Sparse Attention 和 Mixture of Experts 提升长文档 Diffusion Models

Source: arXiv - 2512.20604v1

概述

本文介绍了 MoE‑DiffuSeq，一种新框架，将 mixture‑of‑experts（MoE）路由与自定义的 sparse‑attention 机制相结合，使基于扩散的文本生成能够适用于超长文档。通过解决现有扩散模型（例如 DiffuSeq）中臭名昭著的内存和计算瓶颈，作者将该技术推向更接近实际应用的场景，如科学论文撰写、代码库合成以及多轮对话机器人。

关键贡献

• 稀疏注意力扩散骨干：一种针对性的注意力方案，随着序列长度近似线性扩展，显著降低 GPU 内存使用。
• 混合专家路由：对每个 token 动态激活仅一小部分专家子网络，进一步降低 FLOPs，同时保持模型容量。
• 软吸收状态：嵌入扩散去噪步骤中，加速收敛并提升 token 级别的重建保真度。
• 全面基准测试：在长文本数据集（科学摘要、代码库、对话日志）上的实证结果显示，训练/采样速度提升 2–3 倍，并在 BLEU、ROUGE 以及人工评估的一致性上取得可衡量的提升。
• 开源实现：作者发布代码和预训练检查点，降低开发者尝试基于扩散的生成模型的门槛。

方法论

1. Base diffusion model – 基于 DiffuSeq，后者将文本生成视为逆扩散过程：噪声的 token 序列逐步去噪，恢复为可读文本。
2. Sparse attention layer – 与传统的全自注意力（O(N²) 计算成本）不同，模型仅在滑动窗口以及一组学习得到的“全局” token 上计算注意力。这样将每层的复杂度降低到 O(N·k)，其中 k ≪ N。
3. Mixture‑of‑Experts (MoE) routing – 每个 transformer 块包含多个专家前馈网络。轻量级门控网络为每个 token 选择 top‑k 专家，仅在前向/反向传播时激活这些专家。这样在不成比例增加计算量的情况下实现高容量模型。
4. Soft absorbing state – 在扩散步骤中，允许一小部分概率质量“吸收”到稳定状态，从而有效缩短收敛所需的扩散时间步数。
5. Training & sampling – 模型使用标准的变分扩散损失进行训练，但额外加入 MoE 正则化（负载均衡损失）和稀疏注意力掩码。采样遵循常规的逆扩散调度，现因吸收状态而加速。

Results & Findings

• 效率：训练速度提升约 2.5×，在 2k‑token 序列上采样延迟降低超过 50 %。
• 质量：在自动指标和人工评估上均有一致提升，尤其在长跨度保持全局连贯性方面表现突出。
• 可扩展性：得益于 MoE 稀疏性，从 1 B 到 4 B 参数的扩展仅导致内存占用略有增长。

实际影响

• 开发者工具：IDE 插件可以自动生成大量文档或代码片段，并且现在可以在不产生高延迟的情况下依赖扩散模型。
• 内容平台：新闻编辑部和科学出版机构可以使用 MoE‑DiffuSeq 起草长篇文章，快速获得保留结构的初稿。
• 对话式 AI：处理多轮、上下文丰富对话的客服机器人能够在数百轮对话中保持连贯性，而不会导致 GPU 成本激增。
• 边缘友好部署：由于每个 token 只激活一小部分专家，推理可以在多块 GPU 或专用加速器集群上进行分片，使大规模生成更具成本效益。
• 开源生态：发布的代码与 Hugging Face Transformers 集成，开发者可以将 MoE‑DiffuSeq 以最小的摩擦插入现有流水线。

限制与未来工作

• 专家不平衡：尽管使用了负载均衡损失，一些专家仍可能被低利用，尤其在高度同质的语料库上。
• 稀疏注意力超参数调优：选择合适的窗口大小和全局标记数量仍需针对特定数据集进行实验。
• 扩散步数：虽然软吸收状态减少了步数，但模型仍需要数十次逆扩散迭代，这可能成为超低延迟应用的障碍。
• 未来方向 作者提出的包括：
  1. 自适应门控，学习在每个标记上动态变化活跃专家的数量。
  2. 融入检索增强生成，以进一步提升事实准确性。
  3. 探索混合自回归‑扩散调度，以进一步减少推理步数。

作者

• Alexandros Christoforos
• Chadbourne Davis

论文信息

• arXiv ID: 2512.20604v1
• Categories: cs.CL
• Published: 2025年12月23日
• PDF: 下载 PDF