Sur l'évolutivité de la génération d'images à partir de texte basée sur la diffusion

On the Scalability of Diffusion-based Text-to-Image Generation

April 3, 2024

Auteurs: Hao Li, Yang Zou, Ying Wang, Orchid Majumder, Yusheng Xie, R. Manmatha, Ashwin Swaminathan, Zhuowen Tu, Stefano Ermon, Stefano Soatto

cs.AI

Résumé

L'augmentation de la taille des modèles et des données a été un succès majeur dans l'évolution des LLMs (Large Language Models). Cependant, les lois de mise à l'échelle pour les modèles de génération d'images à partir de texte (T2I) basés sur la diffusion ne sont pas encore pleinement explorées. Il est également difficile de déterminer comment mettre à l'échelle efficacement ces modèles pour améliorer leurs performances tout en réduisant les coûts. Les différents paramètres d'entraînement et les coûts élevés de formation rendent une comparaison équitable des modèles extrêmement complexe. Dans ce travail, nous étudions empiriquement les propriétés de mise à l'échelle des modèles T2I basés sur la diffusion en effectuant des ablations extensives et rigoureuses sur la mise à l'échelle des architectures de débruitage et des ensembles d'entraînement, incluant l'entraînement de variantes de UNet et de Transformers allant de 0,4 à 4 milliards de paramètres sur des jeux de données allant jusqu'à 600 millions d'images. Pour la mise à l'échelle des modèles, nous constatons que l'emplacement et la quantité d'attention croisée distinguent les performances des architectures UNet existantes. De plus, augmenter le nombre de blocs de transformers est plus efficace en termes de paramètres pour améliorer l'alignement texte-image que d'augmenter le nombre de canaux. Nous identifions ensuite une variante efficace de UNet, qui est 45% plus petite et 28% plus rapide que le UNet de SDXL. Concernant la mise à l'échelle des données, nous montrons que la qualité et la diversité de l'ensemble d'entraînement sont plus importantes que la simple taille du jeu de données. Augmenter la densité et la diversité des légendes améliore les performances d'alignement texte-image et l'efficacité de l'apprentissage. Enfin, nous fournissons des fonctions de mise à l'échelle pour prédire les performances d'alignement texte-image en fonction de la taille du modèle, de la puissance de calcul et de la taille du jeu de données.

English

Scaling up model and data size has been quite successful for the evolution of LLMs. However, the scaling law for the diffusion based text-to-image (T2I) models is not fully explored. It is also unclear how to efficiently scale the model for better performance at reduced cost. The different training settings and expensive training cost make a fair model comparison extremely difficult. In this work, we empirically study the scaling properties of diffusion based T2I models by performing extensive and rigours ablations on scaling both denoising backbones and training set, including training scaled UNet and Transformer variants ranging from 0.4B to 4B parameters on datasets upto 600M images. For model scaling, we find the location and amount of cross attention distinguishes the performance of existing UNet designs. And increasing the transformer blocks is more parameter-efficient for improving text-image alignment than increasing channel numbers. We then identify an efficient UNet variant, which is 45% smaller and 28% faster than SDXL's UNet. On the data scaling side, we show the quality and diversity of the training set matters more than simply dataset size. Increasing caption density and diversity improves text-image alignment performance and the learning efficiency. Finally, we provide scaling functions to predict the text-image alignment performance as functions of the scale of model size, compute and dataset size.