SnapGen: Укрощение моделей текст-в-изображение высокого разрешения для мобильных устройств с эффективными архитектурами и обучением

SnapGen: Taming High-Resolution Text-to-Image Models for Mobile Devices with Efficient Architectures and Training

December 12, 2024

Авторы: Dongting Hu, Jierun Chen, Xijie Huang, Huseyin Coskun, Arpit Sahni, Aarush Gupta, Anujraaj Goyal, Dishani Lahiri, Rajesh Singh, Yerlan Idelbayev, Junli Cao, Yanyu Li, Kwang-Ting Cheng, S. -H. Gary Chan, Mingming Gong, Sergey Tulyakov, Anil Kag, Yanwu Xu, Jian Ren

cs.AI

Аннотация

Существующие модели диффузии текста в изображение (T2I) сталкиваются с несколькими ограничениями, включая большие размеры моделей, медленное время выполнения и низкое качество генерации на мобильных устройствах. Цель данной статьи заключается в решении всех этих проблем путем разработки крайне маленькой и быстрой модели T2I, способной генерировать изображения высокого разрешения и качества на мобильных платформах. Мы предлагаем несколько техник для достижения этой цели. Во-первых, мы систематически изучаем выбор дизайна архитектуры сети для сокращения параметров модели и задержки, обеспечивая при этом высокое качество генерации. Во-вторых, для дальнейшего улучшения качества генерации мы используем дистилляцию знаний между архитектурами от гораздо более крупной модели, применяя многоуровневый подход для направления обучения нашей модели с нуля. В-третьих, мы обеспечиваем генерацию на несколько шагов, интегрируя направление адверсариального обучения с дистилляцией знаний. Впервые наша модель SnapGen демонстрирует генерацию изображений размером 1024x1024 пикселя на мобильном устройстве примерно за 1.4 секунды. На наборе данных ImageNet-1K наша модель с всего лишь 372 миллионами параметров достигает значения FID 2.06 для генерации изображений размером 256x256 пикселей. На бенчмарках T2I (т.е. GenEval и DPG-Bench) наша модель с всего лишь 379 миллионами параметров превосходит крупномасштабные модели с миллиардами параметров при значительно меньшем размере (например, в 7 раз меньше, чем у SDXL, и в 14 раз меньше, чем у IF-XL).

English

Existing text-to-image (T2I) diffusion models face several limitations, including large model sizes, slow runtime, and low-quality generation on mobile devices. This paper aims to address all of these challenges by developing an extremely small and fast T2I model that generates high-resolution and high-quality images on mobile platforms. We propose several techniques to achieve this goal. First, we systematically examine the design choices of the network architecture to reduce model parameters and latency, while ensuring high-quality generation. Second, to further improve generation quality, we employ cross-architecture knowledge distillation from a much larger model, using a multi-level approach to guide the training of our model from scratch. Third, we enable a few-step generation by integrating adversarial guidance with knowledge distillation. For the first time, our model SnapGen, demonstrates the generation of 1024x1024 px images on a mobile device around 1.4 seconds. On ImageNet-1K, our model, with only 372M parameters, achieves an FID of 2.06 for 256x256 px generation. On T2I benchmarks (i.e., GenEval and DPG-Bench), our model with merely 379M parameters, surpasses large-scale models with billions of parameters at a significantly smaller size (e.g., 7x smaller than SDXL, 14x smaller than IF-XL).