拡散モデルに基づくテキストから画像への生成のスケーラビリティについて

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

April 3, 2024

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

cs.AI

要旨

大規模なモデルとデータサイズのスケーリングは、LLM（大規模言語モデル）の進化において非常に成功を収めてきました。しかし、拡散モデルに基づくテキストから画像（T2I）生成モデルのスケーリング則は十分に探究されていません。また、コストを削減しながら性能を向上させるためにモデルを効率的にスケーリングする方法も明らかではありません。異なるトレーニング設定と高額なトレーニングコストにより、公平なモデル比較は極めて困難です。本研究では、拡散モデルに基づくT2Iモデルのスケーリング特性を実証的に調査するため、デノイジングバックボーンとトレーニングセットのスケーリングに関する広範かつ厳密なアブレーション実験を行いました。これには、0.4Bから4BパラメータのスケーリングされたUNetおよびTransformerバリアントを、最大600M枚の画像データセットでトレーニングすることが含まれます。モデルスケーリングに関しては、クロスアテンションの位置と量が既存のUNet設計の性能を左右することがわかりました。また、チャネル数を増やすよりも、Transformerブロックを増やすことがテキストと画像のアラインメントを改善するためのパラメータ効率が高いことが明らかになりました。その後、SDXLのUNetよりも45%小さく、28%高速な効率的なUNetバリアントを特定しました。データスケーリングの側面では、トレーニングセットの品質と多様性が、単純なデータセットサイズよりも重要であることを示しました。キャプションの密度と多様性を増やすことで、テキストと画像のアラインメント性能と学習効率が向上します。最後に、モデルサイズ、計算量、データセットサイズのスケールに応じてテキストと画像のアラインメント性能を予測するスケーリング関数を提供します。

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.