Modèles de diffusion évolutifs basés sur des transformateurs

Nous explorons une nouvelle classe de modèles de diffusion fondés sur l’architecture Transformer. Nous entraînons des modèles de diffusion latente d’images en remplaçant le squelette U-Net couramment utilisé par un Transformer agissant sur des patches latents. Nous analysons la capacité de mise à l’échelle de nos Transformateurs de diffusion (DiTs) à travers la complexité du passage avant, mesurée en Gflops. Nous constatons que les DiTs présentant une charge computationnelle plus élevée en Gflops — due à une augmentation de la profondeur ou de la largeur du Transformer, ou à un nombre accru de tokens d’entrée — affichent systématiquement des scores FID plus faibles. En plus de présenter de bonnes propriétés de mise à l’échelle, nos modèles DiT-XL/2 les plus volumineux surpassent tous les modèles de diffusion antérieurs sur les benchmarks ImageNet conditionnés par classe à 512x512 et 256x256, atteignant un score FID état-de-l’art de 2,27 sur ce dernier.