🤖 machine learning

MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

The paper introduces ModernMOE (MMOE), a modernized diffusion transformer architecture that adapts efficient LLM scaling principles—such as routed experts, shared lightweight experts, and residual routing—to achieve faster convergence and superior quality-cost balance in AIGC generation compared to dense and traditional sparse baselines, all within an accessible single-machine training budget.

Original authors: Yanhao Jia, Jiepeng Wang, Haibin Huang, Chi Zhang, Erik Cambria, Xuelong Li

Published 2026-07-28

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yanhao Jia, Jiepeng Wang, Haibin Huang, Chi Zhang, Erik Cambria, Xuelong Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where computers are learning to paint, dream, and create art from scratch. This is the realm of Artificial Intelligence, specifically a branch called "Generative AI," where machines generate new images, videos, and stories. To do this, they use massive digital brains called "Foundation Models." Think of these models as giant libraries of patterns; the more books (data) they read and the bigger their shelves (parameters) are, the better they get at creating. However, there's a catch: making these libraries bigger usually means they take forever to read and require a supercomputer to run, which is expensive and slow.

Recently, scientists discovered a clever trick used by language models (the ones that write text) to solve this. Instead of reading every single book in the library for every question, they built a system with many different "experts." When a question comes in, a smart gatekeeper decides which few experts are needed to answer it, while the rest can take a coffee break. This is called a "Mixture of Experts" (MoE). It's like having a team of specialists where you only call the plumber when there's a leak, rather than waking up the whole team of doctors, chefs, and mechanics. The big question for image creators was: Can we use this same "call only the experts you need" trick to make image generators faster and cheaper without making them less creative?

This paper introduces a new system called MMOE (ModernMOE) to answer that question. The researchers took a standard image generator, which usually forces every single part of its brain to work on every single pixel, and gave it a modern upgrade. Instead of just adding more experts and hoping for the best, they redesigned the team's workflow. They added a few special "lazy" experts who can do simple things like copying an image or doing nothing at all, saving energy. They also added a "gate-residual" system, which lets the gatekeeper remember what it decided in the previous step, so it doesn't have to rethink everything from scratch. Finally, they let the experts share notes with each other across different layers of the brain, so information flows more smoothly.

The team tested this new MMOE system on a single powerful computer with eight graphics cards, training it for 400,000 steps. They compared it against older, "dense" models where everything works all the time, and other "sparse" models that just have more experts but no smart shortcuts. The results suggest that MMOE is the winner in a race for efficiency. It learned to create high-quality images faster than the others, reaching a lower error score (called FID) at every checkpoint. It didn't just get better; it got better cheaper. The analysis showed that the system learned to use its "lazy" experts often, especially in the early stages of creating an image, and that the experts specialized in different tasks without getting confused.

The paper argues that simply making models bigger and more complex isn't the only way to improve them. Instead, by borrowing smart efficiency tricks from language models—like using lightweight experts and sharing information—the team suggests we can build image generators that are both high-quality and practical to run on smaller budgets. While the study was limited to specific image types and didn't test every possible scenario, the results strongly suggest that this "modernized" approach is a promising path forward, offering a better balance between how good the images look and how much it costs to make them.

Technical Summary: MMOE: Modernizing Diffusion Transformers with Efficient Expert Design

Problem Statement
While Large Language Models (LLMs) have successfully scaled by decoupling total parameter capacity from activated computation using sparse Mixture-of-Experts (MoE) and efficiency mechanisms (e.g., zero-computation experts, gate-residual routing), AIGC Foundation Models (AFMs), specifically diffusion-transformer backbones, have largely followed a different trajectory. Current efforts to apply MoE to diffusion transformers (e.g., DiT-MoE, Race-DiT) primarily focus on enlarging the total expert pool and sparsity ratios without importing the efficiency mechanisms that make LLM scaling practical. This approach often results in models that are difficult to train and deploy on resource-constrained devices, failing to balance generation quality against training and deployment costs. The paper poses the question: Can the architectural principles behind efficient LLM scaling be adapted to AFMs to improve generation quality without proportionally increasing effective computation or training burden?

Methodology: ModernMOE (MMOE)
The authors propose ModernMOE (MMOE), a modernization of SiT-style diffusion transformers that systematically integrates four efficiency mechanisms proven in LLMs into the diffusion backbone. Crucially, MMOE keeps the external interface (denoising objective, conditioning, latent representation, and sampler) unchanged, modifying only the internal block-level computation.

The architecture introduces the following components:

MoE++ Expert Layer: Instead of standard routed MLPs, the expert pool includes:
- Heavy Experts: Standard MLPs.
- Lightweight Zero-Computation Experts: "Copy" (identity), "Zero" (output zero), and "Constant" (learned vector mixture) experts. These allow the model to skip full MLP execution for tokens where a heavy update is unnecessary.
Gate-Residual Routing: The router incorporates a residual from the previous expert layer's gate decision, allowing routing tendencies to be inherited across layers, stabilizing decisions.
Shared Experts: A shared dense expert is always active to provide a common computation path, alongside routed experts for token-dependent specialization.
Attention-Residual Aggregation: The model reuses information from previously completed block states. Before both the attention and MLP sub-layers, the model aggregates representations from a list of completed block states using a learned pseudo-query, enabling cross-depth information reuse without new generative objectives.

The forward pass maintains recurrent states (completed block states, current partial state, and gate residuals) to facilitate this aggregation and routing.

Key Contributions

Formulation of AFM Modernization: The paper frames AFM scaling as a modernization problem, systematically examining whether LLM MoE designs can improve diffusion transformers without simply inflating parameter counts or sparsity ratios.
Unified Architecture: The proposal of MMOE, which integrates routed experts, shared experts, lightweight zero-computation experts, gate-residual routing, and attention-residual aggregation into a single SiT-style backbone.
Controlled Empirical Study: The authors provide a rigorous, controlled comparison of the modernization path (Dense SiT $\to$ Standard MoE $\to$ Shared MoE $\to$ MoE++ $\to$ AttnRes-MoE $\to$ MMOE) under a matched, accessible training budget (single eight-GPU H100 node, batch size 256, 400k steps).

Results
Evaluated on class-conditional ImageNet-256 generation:

Convergence and Quality: MMOE achieves lower Fréchet Inception Distance (FID) at every recorded checkpoint compared to dense SiT and intermediate sparse-expert baselines. It converges faster per training step, reaching an FID of 3.75 at 400k steps compared to 5.20 for the dense baseline.
Efficiency: While standard MoE variants increased wall-clock training time due to routing overhead, MMOE achieved the best quality-cost balance among sparse variants. It reduced the training time of the most expensive variant (AMOE) from 120 hours to 67 hours while maintaining superior FID.
Scalability: The improvements hold across different model scales (S/2, B/2, L/2) and resolutions (ImageNet-512), indicating the benefits are not limited to the XL/2 setting.
Routing Analysis: Analysis of the converged model reveals:
- Stable Specialization: Strong block-wise expert specialization (high load Gini coefficient) rather than uniform usage.
- Depth-Dependent Lightweight Use: Lightweight experts (copy/zero/constant) are heavily utilized in early blocks (44.6% of routes) and decrease in later blocks (23.4%), supporting the hypothesis that early denoising steps benefit from cheap transformations.
- Low Switch Rates: The top-2 expert set changes for only ~1.4% to 2.7% of tokens between adjacent denoising steps, indicating stable routing decisions.
Memory Savings: The use of lightweight experts reduces forward activation memory by ~20% and backward memory by ~32% per block compared to attention-residual MoE without lightweight experts.

Significance and Claims
The paper claims that AFMs can follow the "balanced scaling path" of LLMs by importing proven efficiency designs rather than merely increasing total parameters and sparsity ratios. The significance of MMOE lies in demonstrating that:

Quality-Cost Balance: High-quality generation can be achieved with better quality-cost tradeoffs by combining sparse routing with cheap expert paths and representation reuse, rather than relying solely on larger expert pools.
Feasibility on Accessible Budgets: These improvements are achievable within a single-machine budget (8x H100), making advanced diffusion transformer research more accessible compared to prior works requiring massive multi-node clusters and millions of iterations.
Mechanism Transfer: Specific LLM efficiency mechanisms (zero-computation experts, gate-residuals, attention-residuals) transfer effectively to diffusion transformers, improving convergence speed and specialization stability.

The authors remain modest regarding their claims, noting that they do not claim state-of-the-art FID compared to methods trained at much larger scales (e.g., 7M iterations), but rather that their contribution is the controlled modernization study and the demonstration of superior performance under a constrained, single-node budget. They acknowledge limitations, including the focus on class-conditional generation and the need for further study on distributed training efficiency, which is currently communication-bound.

More like this