DiffMoE : Sélection dynamique de tokens pour des transformeurs de diffusion évolutifs

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

Résumé

Les modèles de diffusion ont démontré un succès remarquable dans diverses tâches de génération d'images, mais leurs performances sont souvent limitées par le traitement uniforme des entrées dans des conditions et niveaux de bruit variables. Pour remédier à cette limitation, nous proposons une approche novatrice qui exploite l'hétérogénéité inhérente au processus de diffusion. Notre méthode, DiffMoE, introduit un pool global de tokens au niveau du lot, permettant aux experts d'accéder aux distributions globales de tokens pendant l'entraînement, favorisant ainsi un comportement spécialisé des experts. Pour libérer tout le potentiel du processus de diffusion, DiffMoE intègre un prédicteur de capacité qui alloue dynamiquement les ressources de calcul en fonction des niveaux de bruit et de la complexité des échantillons. À travers une évaluation approfondie, DiffMoE atteint des performances de pointe parmi les modèles de diffusion sur le benchmark ImageNet, surpassant largement à la fois les architectures denses avec 3x paramètres activés et les approches MoE existantes tout en maintenant 1x paramètres activés. L'efficacité de notre approche s'étend au-delà de la génération conditionnelle par classe à des tâches plus complexes telles que la génération texte-à-image, démontrant son applicabilité étendue à différentes applications des modèles de diffusion. Page du projet : 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/