MMOE:用高效专家设计现代化扩散Transformer

MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

精选理由

这篇论文把LLM里高效扩展的思路用到了扩散Transformer上,MMOE在单机8卡H100就跑出了更低的FID,比加参数更实用。

AI 摘要

MMOE (ModernMOE) 是一种对SiT风格扩散Transformer的现代化改造,系统性地集成了路由专家、共享与轻量专家、门控残差路由及注意力残差信息复用。所有实验在单个8卡H100节点上以batch size 256训练40万步。在相同预算下,MMOE在每个记录检查点都达到更低的FID,即每训练步收敛更快,且优于密集和中等稀疏专家基线。路由分析显示专家在深度上稳定专业化,轻量路由被大量使用,去噪过程中步骤间的路由变化适度。这表明AIGC基础模型可以跟随LLM的平衡扩展路径,而不是简单地增加总参数量和稀疏比。

原文 · arXiv cs.LG

MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

Modern large language models scale successfully by pairing capacity growth with efficiency, keeping per-token and deployment costs under control as capacity grows. AIGC Foundation Models (AFMs), especially diffusion-transformer backbones, have begun to adopt sparse experts, but recent efforts mostly enlarge total parameter counts and sparsity ratios without importing the efficiency mechanisms that made LLM scaling practical, so generation quality is seldom balanced against training and deployment cost. This raises a natural question: can the architectural principles behind efficient LLM scaling be adapted to AFMs in a more balanced way? We introduce ModernMOE (MMOE), a modernization of SiT-style diffusion transformers that systematically adapts routed experts, shared and lightweight experts, gate-residual routing, and attention-residual information reuse to AIGC generation. Rather than treating MoE as a single plug-in replacement, MMOE studies how different modern expert components affect convergence, efficiency, and generation quality when composed inside a diffusion transformer. Every experiment in this paper is trained on a single eight-GPU H100 node with batch size 256 for 400k steps, an accessible single-machine budget. Under matched training and sampling protocols and at this budget, MMOE reaches lower FID at every recorded checkpoint, that is, it converges faster per training step, than dense and intermediate sparse-expert baselines, and among the sparse variants it attains the best quality-cost balance. Routing analysis further shows stable expert specialization across depth, substantial use of lightweight routes, and modest step-to-step routing changes during denoising. These results suggest that AFMs can follow the balanced scaling path of LLMs by importing proven efficiency designs, rather than by simply increasing total parameters and sparsity ratios.