DiffMoE: Динамический выбор токенов для масштабируемых диффузионных трансформеров

DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformers

March 18, 2025

Авторы: 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

Аннотация

Диффузионные модели продемонстрировали впечатляющие результаты в различных задачах генерации изображений, однако их производительность часто ограничивается равномерной обработкой входных данных при различных условиях и уровнях шума. Чтобы устранить это ограничение, мы предлагаем новый подход, который использует присущую диффузионному процессу неоднородность. Наш метод, DiffMoE, вводит пул глобальных токенов на уровне батча, что позволяет экспертам получать доступ к глобальным распределениям токенов во время обучения, способствуя специализированному поведению экспертов. Для полного раскрытия потенциала диффузионного процесса DiffMoE включает предсказатель емкости, который динамически распределяет вычислительные ресурсы в зависимости от уровня шума и сложности образца. В ходе всесторонней оценки DiffMoE достигает наилучших результатов среди диффузионных моделей на бенчмарке ImageNet, значительно превосходя как плотные архитектуры с 3x активированными параметрами, так и существующие подходы MoE, сохраняя при этом 1x активированных параметров. Эффективность нашего подхода выходит за рамки условной генерации по классам и распространяется на более сложные задачи, такие как генерация изображений по тексту, демонстрируя его широкую применимость в различных приложениях диффузионных моделей. Страница проекта: 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/