Efficient Training of Diffusion Mixture-of-Experts Models: A Practical Recipe

TL;DR

This work argues that the under-explored architectural design space of diffusion MoE models largely explains the gap with LLM MoEs, and proposes architecture-centric improvements inspired by DeepSeekMoE, including narrowed FFN intermediates, shared experts, and 2D RoPE positional encoding. The authors introduce DSMoE (latent diffusion) and JiTMoE (pixel-space diffusion) and demonstrate consistent, parameter-efficient gains across model sizes, achieving state-of-the-art-like results with fewer activated parameters and faster convergence. Quantitative gains on ImageNet 256×256 with Rectified Flow, plus ablation analyses and visual results, establish a practical training recipe that extends to both latent and pixel diffusion frameworks. Overall, architecture choices, rather than routing complexity alone, drive substantial improvements and efficiency in diffusion MoE models with broad applicability.

Abstract

Recent efforts on Diffusion Mixture-of-Experts (MoE) models have primarily focused on developing more sophisticated routing mechanisms. However, we observe that the underlying architectural configuration space remains markedly under-explored. Inspired by the MoE design paradigms established in large language models (LLMs), we identify a set of crucial architectural factors for building effective Diffusion MoE models--including DeepSeek-style expert modules, alternative intermediate widths, varying expert counts, and enhanced attention positional encodings. Our systematic study reveals that carefully tuning these configurations is essential for unlocking the full potential of Diffusion MoE models, often yielding gains that exceed those achieved by routing innovations alone. Through extensive experiments, we present novel architectures that can be efficiently applied to both latent and pixel-space diffusion frameworks, which provide a practical and efficient training recipe that enables Diffusion MoE models to surpass strong baselines while using equal or fewer activated parameters. All code and models are publicly available at: https://github.com/yhlleo/EfficientMoE.

Figures (7)

• Figure 1: Training MSE loss curves of DiffMoE and DSMoE models. DSMoE, equipped with improved MoE designs, consistently achieves lower diffusion MSE losses than the baselines. Moreover, this performance gap further widens as the model scales up.
• Figure 2: Top: Average number of experts activated per class across MoE layers. Bottom: Per-layer expert activation frequency averaged over all classes.
• Figure 3: Training MSE loss curves of DSMoE with "E16" and "E48". Across all model sizes, E48 consistently achieves lower MSE losses than E16, indicating faster and more stable convergence.
• Figure 4: Training MSE loss curves of DSMoE with different position encoding methods.
• Figure 5: Ablations on replacing all FFN layers with MoE layer (i.e., w/o Interleave), applying GQA ainslie2023gqa (i.e., w/ GQA), and removing shared experts (i.e., S0A3).