DiffMoE: Selezione Dinamica dei Token per Trasformatori Diffusivi Scalabili

DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers

March 18, 2025

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

Abstract

I modelli di diffusione hanno dimostrato un notevole successo in vari compiti di generazione di immagini, ma le loro prestazioni sono spesso limitate dall'elaborazione uniforme degli input in condizioni e livelli di rumore variabili. Per affrontare questa limitazione, proponiamo un approccio innovativo che sfrutta l'eterogeneità intrinseca del processo di diffusione. Il nostro metodo, DiffMoE, introduce un pool globale di token a livello di batch che consente agli esperti di accedere alle distribuzioni globali di token durante l'addestramento, promuovendo un comportamento specializzato degli esperti. Per sfruttare appieno il potenziale del processo di diffusione, DiffMoE incorpora un predittore di capacità che alloca dinamicamente le risorse computazionali in base ai livelli di rumore e alla complessità del campione. Attraverso una valutazione completa, DiffMoE raggiunge prestazioni all'avanguardia tra i modelli di diffusione sul benchmark ImageNet, superando sostanzialmente sia le architetture dense con 3x parametri attivati che gli approcci MoE esistenti, mantenendo 1x parametri attivati. L'efficacia del nostro approccio si estende oltre la generazione condizionata alla classe a compiti più impegnativi come la generazione di immagini da testo, dimostrando la sua ampia applicabilità in diverse applicazioni dei modelli di diffusione. Pagina del progetto: 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/