Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts

Switch Diffusion Transformer: Synergizing Denoising Tasks with Sparse Mixture-of-Experts

10 Jul 2024 | Byeongjun Park, Hyojun Go, Jin-Young Kim, Sangmin Woo, Seokil Ham, Changick Kim
The paper introduces the Switch Diffusion Transformer (Switch-DiT), a novel architecture that synergizes denoising tasks in diffusion models by establishing inter-task relationships without compromising semantic information. Switch-DiT employs sparse mixture-of-experts (SMoE) layers within each transformer block to facilitate parameter isolation and handle conflicts between tasks. The diffusion prior loss regularizes the gating network, encouraging similar tasks to share denoising paths while isolating conflicting ones. Extensive experiments on unconditional and class-conditional image generation datasets validate the effectiveness of Switch-DiT in improving image quality and convergence rate, demonstrating its ability to construct tailored denoising paths across various generation scenarios. The project page is available at https://byeongjun-park.github.io/Switch-DiT/.The paper introduces the Switch Diffusion Transformer (Switch-DiT), a novel architecture that synergizes denoising tasks in diffusion models by establishing inter-task relationships without compromising semantic information. Switch-DiT employs sparse mixture-of-experts (SMoE) layers within each transformer block to facilitate parameter isolation and handle conflicts between tasks. The diffusion prior loss regularizes the gating network, encouraging similar tasks to share denoising paths while isolating conflicting ones. Extensive experiments on unconditional and class-conditional image generation datasets validate the effectiveness of Switch-DiT in improving image quality and convergence rate, demonstrating its ability to construct tailored denoising paths across various generation scenarios. The project page is available at https://byeongjun-park.github.io/Switch-DiT/.
Reach us at info@study.space