Mise à l'échelle des Transformers de Diffusion jusqu'à 16 milliards de paramètres

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

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

cs.AI

Résumé

Dans cet article, nous présentons DiT-MoE, une version éparse du Transformer de diffusion, qui est scalable et compétitive par rapport aux réseaux denses tout en offrant une inférence hautement optimisée. Le DiT-MoE intègre deux conceptions simples : le routage d'experts partagés et la perte d'équilibre au niveau des experts, permettant ainsi de capturer des connaissances communes et de réduire la redondance entre les différents experts routés. Lorsqu'il est appliqué à la génération d'images conditionnelles, une analyse approfondie de la spécialisation des experts révèle quelques observations intéressantes : (i) La sélection des experts montre une préférence pour la position spatiale et l'étape de débruitage, tout en étant insensible aux différentes informations conditionnelles de classe ; (ii) À mesure que les couches MoE deviennent plus profondes, la sélection des experts passe progressivement d'une position spatiale spécifique à une dispersion et un équilibre. (iii) La spécialisation des experts tend à être plus concentrée au début de l'étape de débruitage, puis devient progressivement uniforme après la moitié. Nous attribuons cela au processus de diffusion qui modélise d'abord les informations spatiales de basse fréquence, puis les informations complexes de haute fréquence. Sur la base de ces observations, une série de DiT-MoE atteint expérimentalement des performances comparables à celles des réseaux denses tout en nécessitant une charge de calcul bien moindre lors de l'inférence. Plus encourageant encore, nous démontrons le potentiel de DiT-MoE avec des données d'images synthétisées, en mettant à l'échelle un modèle de diffusion à 16,5 milliards de paramètres qui obtient un nouveau score SoTA FID-50K de 1,80 dans des résolutions de 512x512. La page du projet : 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.