🤖 machine learning

MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

本文介绍了 ModernMOE (MMOE)，这是一种现代化的扩散 Transformer 架构，它借鉴了高效的 LLM 缩放原则——例如路由专家、共享轻量级专家和残差路由——旨在实现比稠密模型及传统稀疏基准模型更快的收敛速度和更优的质量-成本平衡，且所有这些均能在易于实现的单机训练预算内完成 AIGC 生成任务。

原作者： Yanhao Jia, Jiepeng Wang, Haibin Huang, Chi Zhang, Erik Cambria, Xuelong Li

发布于 2026-07-28

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Yanhao Jia, Jiepeng Wang, Haibin Huang, Chi Zhang, Erik Cambria, Xuelong Li

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一个计算机正在学习从无到有地绘画、做梦并进行艺术创作的世界。这就是人工智能的领域，特别是被称为“生成式人工智能”（Generative AI）的一个分支，在这里，机器可以生成新的图像、视频和故事。为了实现这一点，它们使用被称为“基础模型”（Foundation Models）的巨大数字大脑。可以将这些模型想象成巨大的模式图书馆；它们阅读的书籍（数据）越多，书架（参数）越大，它们就变得越擅长创作。然而，这里有一个难点：让这些图书馆变得更大通常意味着它们需要花费极长的时间来阅读，并且需要超级计算机来运行，这既昂贵又缓慢。

最近，科学家们发现了一个用于语言模型（那些编写文本的模型）解决此问题的巧妙技巧。他们并没有针对每个问题都去阅读图书馆里的每一本书，而是构建了一个拥有许多不同“专家”的系统。当一个问题进来时，一个聪明的门卫会决定需要哪几位专家来回答，而其他专家则可以去喝杯咖啡。这被称为“混合专家模型”（Mixture of Experts，简称 MoE）。这就像拥有一支专家团队，你只在发生漏水时才请水管工，而不是把医生、厨师和机械师整个团队都叫醒。大的疑问在于：图像创作者能否使用这种同样的“只调用所需的专家”的技巧，在不降低创造力的前提下，让图像生成器变得更快、更便宜？

这篇论文介绍了一个名为 MMOE（ModernMOE）的新系统来回答这个问题。研究人员采用了一个标准的图像生成器——这类生成器通常会强迫其大脑的每一个部分都在处理每一个像素——并为其进行了现代化的升级。他们并没有只是简单地增加更多专家并听天由命，而是重新设计了团队的工作流程。他们添加了一些特殊的“懒惰”专家，这些专家可以做一些简单的任务，比如复制图像或什么都不做，从而节省能量。他们还添加了一个“门控残差”（gate-residual）系统，让门卫能够记住它在上一步所做的决定，这样它就不必从头开始重新思考一切。最后，他们让专家们在脑部的不同层级之间共享笔记，以便信息能够更顺畅地流动。

团队在一台配备了八块显卡的单台高性能计算机上测试了这个新的 MMOE 系统，对其进行了 4 亿步的训练。他们将其与旧有的、“稠密型”（dense）模型（即所有部分始终都在工作的模型）以及其他“稀疏型”（sparse）模型（即只是拥有更多专家但没有智能捷径的模型）进行了对比。结果表明，MMOE 是效率竞赛中的获胜者。它学会了比其他模型更快地创建高质量图像，在每一个检查点都达到了更低的误差得分（称为 FID）。它不仅变得更好，而且变得更“廉价”了。分析显示，该系统学会了经常使用其“懒惰”专家，尤其是在创建图像的早期阶段，并且专家们在承担不同任务时各司其职，不会产生混淆。

论文认为，仅仅通过让模型变得更大、更复杂并不是提高它们的唯一途径。相反，通过借鉴语言模型的智能效率技巧——例如使用轻量级专家和共享信息——该团队建议我们可以构建出既高质量又具备实用运行能力的图像生成器。虽然这项研究局限于特定的图像类型，且并未测试所有可能的场景，但结果有力地表明，这种“现代化”的方法是一个充满前景的方向，它在图像视觉效果与制作成本之间提供了更好的平衡。

技术摘要：MMOE：通过高效专家设计实现扩散 Transformer 的现代化

问题陈述
虽然大语言模型（LLMs）通过将总参数容量与激活计算解耦（利用稀疏混合专家模型 MoE 以及效率机制，如零计算专家、门控残差路由），成功实现了规模化扩展；但生成式人工智能基础模型（AFMs），特别是扩散 Transformer（diffusion-transformer）骨干网络，其发展轨迹却大不相同。目前将 MoE 应用于扩散 Transformer 的尝试（如 DiT-MoE、Race-DiT）主要集中在扩大总专家池和稀疏比例，而未能引入使 LLM 规模化变得实用的效率机制。这种方法往往导致模型难以在资源受限的设备上进行训练和部署，无法在生成质量与训练及部署成本之间取得平衡。本文提出了这样一个问题：能否借鉴 LLM 高效规模化的架构原则，来改进 AFM，从而在不按比例增加有效计算量或训练负担的情况下提升生成质量？

方法论：ModernMOE (MMOE)
作者提出了 ModernMOE (MMOE)，这是一种对 SiT 式扩散 Transformer 的现代化改造，它系统地将 LLM 中被证明有效的四种效率机制集成到了扩散骨干网络中。至关重要的是，MMOE 保持了外部接口（去噪目标、条件控制、潜变量表示和采样器）不变，仅修改了内部的块级（block-level）计算。

该架构引入了以下组件：

MoE++ 专家层： 不同于标准的路由 MLP，专家池包括：
- 重型专家（Heavy Experts）： 标准 MLP。
- 轻量级零计算专家（Lightweight Zero-Computation Experts）： “复制”（Identity）、“零”（输出为零）和“常数”（学习到的向量混合）专家。这些专家允许模型在不需要重度更新的 Token 上跳过完整的 MLP 执行。
门控残差路由（Gate-Residual Routing）： 路由器整合了来自前一层专家层门控决策的残数，使得路由倾向能够在不同层之间继承，从而稳定决策。
共享专家（Shared Experts）： 一个共享的密集专家始终处于激活状态，提供公共计算路径，同时配合路由专家实现针对 Token 的专业化处理。
注意力残差聚合（Attention-Residual Aggregation）： 模型重用来自先前完成的块状态的信息。在注意力子层和 MLP 子层之前，模型使用学习到的伪查询（pseudo-query）聚合来自已完成块状态列表的表示，从而实现跨深度的信息重用，而无需引入新的生成目标。

前向传播过程维持了循环状态（已完成的块状态、当前部分状态和门控残数），以促进这种聚合与路由。

核心贡献

提出 AFM 现代化公式： 本文将 AFM 的规模化定义为一个现代化问题，系统地探讨了 LLM MoE 设计是否能在不单纯增加参数量或稀疏比例的情况下，提升扩散 Transformer 的性能。
统一架构： 提出了 MMOE，它将路由专家、共享专家、轻量级零计算专家、门控残差路由以及注意力残差聚合集成到一个统一的 SiT 式骨干网络中。
受控实证研究： 作者提供了一个严格且受控的比较研究，展示了从 Dense SiT $\to$ Standard MoE $\to$ Shared MoE $\to$ MoE++ $\to$ AttnRes-MoE $\to$ MMOE 的现代化演进路径，并使用了匹配的可获取训练预算（单个 8-GPU H100 节点，Batch Size 为 256，400k 步）。

结果
在 Class-conditional ImageNet-256 生成任务上的评估结果如下：

收敛性与质量： 在所有记录的检查点中，MMOE 的 Fréchet Inception Distance (FID) 均低于 Dense SiT 和中间阶段的稀疏专家基准模型。它在每个训练步长内收敛更快，在 400k 步时达到了 3.75 的 FID，而 Dense 基准模型为 5.20。
效率： 虽然标准的 MoE 变体由于路由开销增加了实际运行时间（wall-clock time），但 MMOE 在所有稀疏变体中实现了最佳的质量-成本平衡。它将最昂贵变体（AMOE）的训练时间从 120 小时缩短至 67 小时，同时保持了更优的 FID。
可扩展性： 这种改进在不同模型规模（S/2, B/2, L/2）和分辨率（ImageNet-512）下均成立，表明其益处并不局限于 XL/2 设置。
路由分析： 对收敛模型的分析显示：
- 稳定的专业化： 表现出强烈的块级专家专业化特征（高负载 Gini 系数），而非均匀使用。
- 深度相关的轻量级使用： 轻量级专家（复制/零/常数）在早期块中被大量使用（44.6% 的路由），并在后期块中减少（23.4%），这支持了“早期去噪步骤受益于廉价变换”的假设。
- 低切换率： 在相邻的去噪步骤之间，前两个专家的集合变化仅为 ~1.4% 至 2.7%，表明路由决策非常稳定。
显存节省： 与不含轻量级专家的 Attention-residual MoE 相比，使用轻量级专家使每块的前向激活显存减少了约 20%，反向传播显存减少了约 32%。

意义与主张
本文声称，通过引入经过验证的效率设计而非仅仅增加总参数量和稀疏比例，AFM 可以遵循 LLM 的“平衡规模化路径”。MMOE 的意义在于证明了：

质量-成本平衡： 通过结合稀疏路由、廉价专家路径和表示重用，可以实现高质量生成，并获得比单纯依赖更大专家池更好的质量-成本权衡。
在可获取预算下的可行性： 这些改进可以在单机预算（8x H100）内实现，这使得先进的扩散 Transformer 研究比以往需要大规模多节点集群和数百万次迭代的研究更具可及性。
机制迁移： 特定的 LLM 效率机制（零计算专家、门控残差、注意力残差）能有效地迁移到扩散 Transformer 中，从而提高收敛速度和专业化稳定性。

作者对其主张保持谦逊，指出他们并非声称在与更大规模训练（如 7M 次迭代）的方法相比具有 SOTA 级别的 FID，而是强调其贡献在于受控的现代化研究，以及在受限的单节点预算下展现出的卓越性能。他们也承认了局限性，包括研究侧重于类条件生成，以及在分布式训练效率方面仍需进一步研究（目前受通信限制）。

类似论文