Sobre la Escalabilidad de la Generación de Imágenes a partir de Texto Basada en Difusión

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

April 3, 2024

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

cs.AI

Resumen

El aumento en el tamaño de los modelos y los datos ha sido bastante exitoso para la evolución de los LLM (Modelos de Lenguaje de Gran Escala). Sin embargo, la ley de escalamiento para los modelos de texto a imagen (T2I) basados en difusión no ha sido completamente explorada. Además, no está claro cómo escalar eficientemente el modelo para obtener un mejor rendimiento a un costo reducido. Los diferentes ajustes de entrenamiento y el elevado costo del mismo hacen que una comparación justa entre modelos sea extremadamente difícil. En este trabajo, estudiamos empíricamente las propiedades de escalamiento de los modelos T2I basados en difusión mediante la realización de extensas y rigurosas ablaciones sobre el escalamiento tanto de los backbones de eliminación de ruido como del conjunto de entrenamiento, incluyendo el entrenamiento de variantes escaladas de UNet y Transformer que van desde 0.4B hasta 4B de parámetros en conjuntos de datos de hasta 600M de imágenes. Para el escalamiento del modelo, encontramos que la ubicación y la cantidad de atención cruzada distinguen el rendimiento de los diseños existentes de UNet. Además, aumentar los bloques de transformadores es más eficiente en términos de parámetros para mejorar la alineación texto-imagen que aumentar el número de canales. Luego, identificamos una variante eficiente de UNet, que es un 45% más pequeña y un 28% más rápida que la UNet de SDXL. En cuanto al escalamiento de datos, demostramos que la calidad y la diversidad del conjunto de entrenamiento importan más que simplemente el tamaño del conjunto de datos. Aumentar la densidad y diversidad de los textos descriptivos mejora el rendimiento de la alineación texto-imagen y la eficiencia del aprendizaje. Finalmente, proporcionamos funciones de escalamiento para predecir el rendimiento de la alineación texto-imagen como funciones de la escala del tamaño del modelo, el cómputo y el tamaño del conjunto de datos.

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.