확산 기반 텍스트-이미지 생성의 확장성에 관하여

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 설계의 성능을 구분짓는 요소임을 발견했습니다. 또한, 채널 수를 늘리는 것보다 트랜스포머 블록을 늘리는 것이 텍스트-이미지 정렬 성능을 개선하는 데 더 파라미터 효율적임을 확인했습니다. 이를 통해 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.