ModernMOE:将高效专家设计引入扩散Transformer的架构革新

针对当前扩散Transformer扩展时盲目堆砌参数的问题,研究提出ModernMOE(MMOE)架构,系统性地将大语言模型中的高效专家机制适配至AIGC基础模型。该架构引入路由专家、共享与轻量级专家、门控残差路由及注意力残差信息复用等创新。在单节点八卡H100、批量大小256、训练40万步的受限预算下,MMOE实现了比密集模型和中间稀疏基线更快的收敛速度与更低的FID分数,达成最佳质量-成本平衡。路由分析证实了专家在深度上的稳定特化,证明了AIGC模型可借鉴LLM的平衡扩展路径,而非单纯增加稀疏度。

当前,大型语言模型通过容量增长与效率控制的结合实现了成功扩展,但在AIGC基础模型领域,尤其是以扩散Transformer为骨干的网络中,尽管已引入稀疏专家机制,但现有工作大多仅侧重于扩大总参数量和稀疏比例,未能引入使LLM扩展实用的效率机制,导致生成质量与训练部署成本之间缺乏平衡。本研究旨在回答一个核心问题:能否以更平衡的方式将LLM高效扩展的架构原则适配到AIGC基础模型中?为此,作者提出了ModernMOE(MMOE),这是一种对SiT风格扩散Transformer的现代化改造。其核心贡献在于系统性地整合了路由专家、共享与轻量级专家、门控残差路由以及注意力残差信息复用等现代专家组件。与简单地将MoE作为插件替换不同,MMOE深入研究了这些组件在扩散Transformer内部组合时对收敛性、效率和生成质量的具体影响,试图在保持生成质量的同时严格控制计算开销,为AIGC模型的扩展提供了一条新的技术路径。

在技术方法层面,MMOE并未采用单一的专家替换策略,而是构建了一个复杂的专家交互系统。首先,模型引入了路由专家机制,允许不同输入动态选择特定的专家路径,从而提升计算效率。其次,为了进一步优化资源利用,设计中包含了共享专家与轻量级专家,前者用于处理通用特征,后者则专注于特定细节的生成,这种分层设计有助于在减少参数总量的同时保持模型的表达能力。在训练策略上,MMOE采用了门控残差路由技术,通过门控网络动态调整专家权重,并结合注意力残差信息复用机制,确保在去噪过程中关键信息的传递不丢失。这种架构设计不仅提高了模型的收敛速度,还通过减少冗余计算显著降低了每步训练的成本。

此外,研究还详细分析了路由机制在深度网络中的表现,发现专家在深层网络中表现出稳定的特化趋势,这有助于模型在不同生成阶段更精准地分配计算资源,从而实现效率与质量的协同优化。实验设置方面,研究在极具挑战性的单节点八卡H100 GPU集群上进行,批量大小设定为256,总训练步数为40万步,这一预算在大型模型训练中相对紧凑,极具实际参考价值。在匹配的培训和采样协议下,MMOE在每一个记录的检查点均达到了比密集模型和中间稀疏专家基线更低的FID分数,表明其收敛速度显著更快。特别是在稀疏变体中,MMOE实现了最佳的质量-成本平衡,证明了其架构设计的有效性。进一步的路由分析显示,专家在深度上表现出稳定的特化现象,轻量级路由被大量使用,且在去噪过程中路由变化较为温和,这表明模型在生成过程中保持了高度的稳定性和一致性。

这些关键结果不仅验证了MMOE在生成质量上的优势,也揭示了其在计算效率上的巨大潜力,为后续研究提供了重要的基准数据。从行业意义来看,MMOE的研究表明AIGC基础模型可以借鉴LLM的平衡扩展路径,通过引入经过验证的高效设计而非单纯增加参数和稀疏度来实现性能提升。这一发现对开源社区具有重要启示,它鼓励研究者关注架构层面的创新而非仅仅依赖算力堆砌。对于工业落地而言,MMOE在有限预算下实现的高质量生成能力,显著降低了AIGC模型的部署门槛,使得在资源受限环境下运行高性能生成模型成为可能。此外,该研究为后续关于专家混合模型在扩散模型中的应用提供了新的思路,特别是在路由机制和专家特化方面的深入分析,为优化生成过程的稳定性和效率提供了理论支持。总体而言,MMOE不仅是一项技术突破,更为AIGC领域的可持续发展指明了方向,推动了模型效率与质量平衡的新标准。

Sources