DiffMoE: スケーラブルなDiffusion Transformersのための動的トークン選択

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ベンチマークにおいて拡散モデルの中で最先端の性能を達成し、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/