DiffMoE: Seleção Dinâmica de Tokens para Transformers de Difusão Escaláveis

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

Resumo

Os modelos de difusão têm demonstrado um sucesso notável em várias tarefas de geração de imagens, mas seu desempenho é frequentemente limitado pelo processamento uniforme das entradas em diferentes condições e níveis de ruído. Para abordar essa limitação, propomos uma abordagem inovadora que aproveita a heterogeneidade inerente do processo de difusão. Nosso método, DiffMoE, introduz um pool global de tokens em nível de lote que permite que especialistas acessem distribuições globais de tokens durante o treinamento, promovendo comportamentos especializados. Para liberar todo o potencial do processo de difusão, o DiffMoE incorpora um preditor de capacidade que aloca dinamicamente recursos computacionais com base nos níveis de ruído e na complexidade das amostras. Por meio de uma avaliação abrangente, o DiffMoE alcança desempenho de ponta entre os modelos de difusão no benchmark do ImageNet, superando substancialmente tanto arquiteturas densas com 3x parâmetros ativados quanto abordagens MoE existentes, mantendo 1x parâmetros ativados. A eficácia de nossa abordagem vai além da geração condicionada por classe para tarefas mais desafiadoras, como a geração de texto para imagem, demonstrando sua ampla aplicabilidade em diferentes aplicações de modelos de difusão. Página do Projeto: 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/