MMOE: Modernizing Diffusion Transformers with Efficient Expert Design
Abstract Overview
This paper presents MMOE, a modernization of SiT-style diffusion transformers that imports several efficiency-oriented mixture-of-experts designs from large language models. The architecture combines routed experts, shared and lightweight experts such as copy/zero/constant routes, gate-residual routing, and attention-residual aggregation, while keeping the denoising objective, conditioning interface, and sampler unchanged. Experiments are conducted mainly on class-conditional ImageNet generation under a controlled single-node budget of eight H100 GPUs, batch size 256, and 400k training steps. The study emphasizes matched internal comparisons to analyze how each component affects convergence, memory, training cost, and routing behavior.
Novelty
The distinctive contribution is not simply adding MoE layers to a diffusion transformer, but systematically adapting multiple modern LLM-MoE efficiency mechanisms into a unified diffusion-transformer backbone. The paper also studies these components in a controlled modernization path, including lightweight zero-computation experts and attention-residual reuse, to assess their combined effect on quality-cost tradeoffs.
Results
Under matched training and sampling protocols on ImageNet-256, MMOE achieves the lowest FID among the paper's dense and intermediate sparse baselines at each recorded checkpoint, reaching 3.75 at 400k steps versus 5.20 for dense SiT in the main 50-step ODE, 50k-sample setting. It also improves over dense SiT at smaller model scales and on ImageNet-512, and the lightweight expert design reduces per-block activation memory relative to AMOE while lowering AMOE's measured training time from 120h to 67h. Routing analysis indicates stable block-wise specialization, substantial use of lightweight routes in earlier blocks, and low adjacent-step routing switch rates during denoising.
Key Points
- MMOE integrates routed experts, lightweight copy/zero/constant experts, gate-residual routing, and attention-residual aggregation into a SiT-style diffusion transformer without changing the external training or sampling interface.
- In the paper's controlled internal comparisons, attention-residual aggregation provides the largest single quality gain, and the full MMOE model shows the best reported quality-cost balance among the tested sparse variants.
- The authors explicitly limit their claims to matched internal baselines under a single-machine budget, noting that comparisons with previously published diffusion-MoE systems are not direct because training budgets and protocols differ.