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

TL;DR

This paper explores architectural configurations of Diffusion Mixture-of-Experts models, demonstrating that careful tuning of design factors significantly improves performance and efficiency beyond routing innovations.

Contribution

It systematically studies architectural factors in Diffusion MoE models and provides a practical training recipe that enhances performance with fewer parameters.

Findings

Careful architecture tuning yields significant performance gains.

Proposed architectures outperform strong baselines.

Efficient training recipes enable effective Diffusion MoE models.

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.