拡散トランスフォーマーを160億パラメータまでスケーリング

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

著者: Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Junshi Huang

cs.AI

要旨

本論文では、拡散Transformerのスパース版であるDiT-MoEを提案する。DiT-MoEは、高密度ネットワークと競合するスケーラビリティを持ちながら、高度に最適化された推論を実現する。DiT-MoEには、共有エキスパートルーティングとエキスパートレベルのバランス損失という2つのシンプルな設計が含まれており、これにより異なるルーティングされたエキスパート間の共通知識の捕捉と冗長性の削減が可能となる。条件付き画像生成に適用した際、エキスパートの専門化に関する深い分析からいくつかの興味深い観察が得られた：(i) エキスパートの選択は空間位置とノイズ除去のタイムステップに偏りを示すが、異なるクラス条件情報には鈍感である。(ii) MoE層が深くなるにつれて、エキスパートの選択は特定の空間位置から分散とバランスへと徐々にシフトする。(iii) エキスパートの専門化は初期のタイムステップでより集中し、その後半で徐々に均一化する傾向がある。これは、拡散プロセスが最初に低周波の空間情報をモデル化し、その後高周波の複雑な情報をモデル化するためであると考えられる。上記のガイダンスに基づき、一連のDiT-MoEは実験的に高密度ネットワークと同等の性能を達成しつつ、推論時の計算負荷を大幅に削減する。さらに、合成画像データを用いてDiT-MoEの可能性を示し、16.5Bパラメータの拡散モデルをスケーリングし、512×512解像度設定で1.80の新たなSoTA FID-50Kスコアを達成した。プロジェクトページ: https://github.com/feizc/DiT-MoE。

English

In this paper, we present DiT-MoE, a sparse version of the diffusion Transformer, that is scalable and competitive with dense networks while exhibiting highly optimized inference. The DiT-MoE includes two simple designs: shared expert routing and expert-level balance loss, thereby capturing common knowledge and reducing redundancy among the different routed experts. When applied to conditional image generation, a deep analysis of experts specialization gains some interesting observations: (i) Expert selection shows preference with spatial position and denoising time step, while insensitive with different class-conditional information; (ii) As the MoE layers go deeper, the selection of experts gradually shifts from specific spacial position to dispersion and balance. (iii) Expert specialization tends to be more concentrated at the early time step and then gradually uniform after half. We attribute it to the diffusion process that first models the low-frequency spatial information and then high-frequency complex information. Based on the above guidance, a series of DiT-MoE experimentally achieves performance on par with dense networks yet requires much less computational load during inference. More encouragingly, we demonstrate the potential of DiT-MoE with synthesized image data, scaling diffusion model at a 16.5B parameter that attains a new SoTA FID-50K score of 1.80 in 512times512 resolution settings. The project page: https://github.com/feizc/DiT-MoE.