DiffMoE: Dynamische Token-Auswahl für skalierbare Diffusions-Transformatoren

DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers

March 18, 2025

Autoren: 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

Zusammenfassung

Diffusionsmodelle haben bemerkenswerte Erfolge bei verschiedenen Bildgenerierungsaufgaben gezeigt, doch ihre Leistung wird oft durch die gleichmäßige Verarbeitung von Eingaben unter verschiedenen Bedingungen und Rauschpegeln eingeschränkt. Um diese Einschränkung zu überwinden, schlagen wir einen neuartigen Ansatz vor, der die inhärente Heterogenität des Diffusionsprozesses nutzt. Unsere Methode, DiffMoE, führt einen batch-weiten globalen Token-Pool ein, der es Experten ermöglicht, während des Trainings auf globale Token-Verteilungen zuzugreifen und so spezialisiertes Expertenverhalten zu fördern. Um das volle Potenzial des Diffusionsprozesses auszuschöpfen, integriert DiffMoE einen Kapazitätsprädiktor, der Rechenressourcen dynamisch basierend auf Rauschpegeln und Probenkomplexität zuweist. Durch umfassende Evaluierungen erreicht DiffMoE state-of-the-art Leistungen unter Diffusionsmodellen auf dem ImageNet-Benchmark und übertrifft dabei sowohl dichte Architekturen mit 3x aktivierten Parametern als auch bestehende MoE-Ansätze deutlich, während nur 1x aktivierte Parameter beibehalten werden. Die Wirksamkeit unseres Ansatzes erstreckt sich über die klassenbedingte Generierung hinaus auf anspruchsvollere Aufgaben wie Text-zu-Bild-Generierung, was seine breite Anwendbarkeit in verschiedenen Diffusionsmodellanwendungen unterstreicht. Projektseite: 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/