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

초록

확산 모델(Diffusion models)은 다양한 이미지 생성 작업에서 뛰어난 성과를 보여왔지만, 다양한 조건과 노이즈 수준에 걸쳐 입력을 균일하게 처리함으로써 성능이 제한되는 경우가 많습니다. 이러한 한계를 해결하기 위해, 우리는 확산 과정의 본질적 이질성을 활용한 새로운 접근 방식을 제안합니다. 우리의 방법인 DiffMoE는 배치 수준의 글로벌 토큰 풀을 도입하여 전문가들이 훈련 중에 글로벌 토큰 분포에 접근할 수 있게 함으로써 전문가의 특수화된 행동을 촉진합니다. 확산 과정의 전체 잠재력을 발휘하기 위해, DiffMoE는 노이즈 수준과 샘플 복잡성에 따라 계산 자원을 동적으로 할당하는 용량 예측기(capacity predictor)를 통합합니다. 포괄적인 평가를 통해, DiffMoE는 ImageNet 벤치마크에서 확산 모델 중 최첨단 성능을 달성하며, 활성화된 매개변수가 3배인 조밀한 아키텍처와 기존의 MoE 접근 방식을 크게 능가하면서도 1배의 활성화된 매개변수를 유지합니다. 우리의 접근 방식의 효과는 클래스 조건부 생성을 넘어 텍스트-이미지 생성과 같은 더 도전적인 작업까지 확장되어, 다양한 확산 모델 애플리케이션에서의 광범위한 적용 가능성을 입증합니다. 프로젝트 페이지: 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/