DiffMoE: Selección Dinámica de Tokens para Transformadores de Difusión Escalables

DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers

March 18, 2025

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

Resumen

Los modelos de difusión han demostrado un éxito notable en diversas tareas de generación de imágenes, pero su rendimiento a menudo se ve limitado por el procesamiento uniforme de las entradas en diferentes condiciones y niveles de ruido. Para abordar esta limitación, proponemos un enfoque novedoso que aprovecha la heterogeneidad inherente del proceso de difusión. Nuestro método, DiffMoE, introduce un grupo global de tokens a nivel de lote que permite a los expertos acceder a distribuciones globales de tokens durante el entrenamiento, fomentando un comportamiento especializado de los expertos. Para liberar todo el potencial del proceso de difusión, DiffMoE incorpora un predictor de capacidad que asigna dinámicamente recursos computacionales en función de los niveles de ruido y la complejidad de las muestras. A través de una evaluación exhaustiva, DiffMoE logra un rendimiento de vanguardia entre los modelos de difusión en el benchmark de ImageNet, superando sustancialmente tanto a las arquitecturas densas con 3x parámetros activados como a los enfoques MoE existentes, mientras mantiene 1x parámetros activados. La efectividad de nuestro enfoque se extiende más allá de la generación condicionada por clases a tareas más desafiantes, como la generación de texto a imagen, demostrando su amplia aplicabilidad en diferentes aplicaciones de modelos de difusión. Página del proyecto: 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/