DiffMoE: Dynamische Tokenselectie voor Schaalbare Diffusie Transformers

DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers

March 18, 2025

Auteurs: Minglei Shi, Ziyang Yuan, Haotian Yang, Xintao Wang, Mingwu Zheng, Xin Tao, Wenliang Zhao, Wenzhao Zheng, Jie Zhou, Jiwen Lu, Pengfei Wan, Di Zhang, Kun Gai

cs.AI

Samenvatting

Diffusiemodellen hebben opmerkelijke successen geboekt in diverse beeldgeneratietaken, maar hun prestaties worden vaak beperkt door de uniforme verwerking van invoer onder verschillende omstandigheden en ruisniveaus. Om deze beperking aan te pakken, stellen we een nieuwe aanpak voor die gebruikmaakt van de inherente heterogeniteit van het diffusieproces. Onze methode, DiffMoE, introduceert een batchgewijze globale tokenpool die experts in staat stelt om tijdens de training toegang te krijgen tot globale tokendistributies, wat gespecialiseerd expertgedrag bevordert. Om het volledige potentieel van het diffusieproces te benutten, integreert DiffMoE een capaciteitsvoorspeller die rekenbronnen dynamisch toewijst op basis van ruisniveaus en samplecomplexiteit. Door middel van uitgebreide evaluatie behaalt DiffMoE state-of-the-art prestaties onder diffusiemodellen op de ImageNet-benchmark, waarbij het zowel dichte architecturen met 3x geactiveerde parameters als bestaande MoE-benaderingen aanzienlijk overtreft, terwijl het slechts 1x geactiveerde parameters behoudt. De effectiviteit van onze aanpak strekt zich uit verder dan klasse-conditionele generatie naar uitdagendere taken zoals tekst-naar-beeldgeneratie, wat de brede toepasbaarheid ervan aantoont in verschillende diffusiemodeltoepassingen. Projectpagina: https://shiml20.github.io/DiffMoE/

English

Diffusion models have demonstrated remarkable success in various image generation tasks, but their performance is often limited by the uniform processing of inputs across varying conditions and noise levels. To address this limitation, we propose a novel approach that leverages the inherent heterogeneity of the diffusion process. Our method, DiffMoE, introduces a batch-level global token pool that enables experts to access global token distributions during training, promoting specialized expert behavior. To unleash the full potential of the diffusion process, DiffMoE incorporates a capacity predictor that dynamically allocates computational resources based on noise levels and sample complexity. Through comprehensive evaluation, DiffMoE achieves state-of-the-art performance among diffusion models on ImageNet benchmark, substantially outperforming both dense architectures with 3x activated parameters and existing MoE approaches while maintaining 1x activated parameters. The effectiveness of our approach extends beyond class-conditional generation to more challenging tasks such as text-to-image generation, demonstrating its broad applicability across different diffusion model applications. Project Page: https://shiml20.github.io/DiffMoE/