SnapGen: Domando modelos de texto a imagen de alta resolución para dispositivos móviles con arquitecturas y entrenamiento eficientes

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

December 12, 2024

Autores: 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

Resumen

Los modelos de difusión texto-a-imagen (T2I) existentes enfrentan varias limitaciones, incluyendo grandes tamaños de modelo, ejecución lenta y generación de baja calidad en dispositivos móviles. Este artículo tiene como objetivo abordar todos estos desafíos mediante el desarrollo de un modelo T2I extremadamente pequeño y rápido que genere imágenes de alta resolución y alta calidad en plataformas móviles. Proponemos varias técnicas para lograr este objetivo. En primer lugar, examinamos sistemáticamente las elecciones de diseño de la arquitectura de red para reducir los parámetros del modelo y la latencia, asegurando al mismo tiempo una generación de alta calidad. En segundo lugar, para mejorar aún más la calidad de generación, empleamos la destilación de conocimiento entre arquitecturas de un modelo mucho más grande, utilizando un enfoque multinivel para guiar el entrenamiento de nuestro modelo desde cero. En tercer lugar, permitimos una generación en pocas etapas mediante la integración de orientación adversarial con destilación de conocimiento. Por primera vez, nuestro modelo SnapGen demuestra la generación de imágenes de 1024x1024 px en un dispositivo móvil en alrededor de 1.4 segundos. En ImageNet-1K, nuestro modelo, con solo 372M parámetros, logra un FID de 2.06 para la generación de 256x256 px. En los benchmarks de T2I (es decir, GenEval y DPG-Bench), nuestro modelo con apenas 379M parámetros supera a modelos a gran escala con miles de millones de parámetros en un tamaño significativamente menor (por ejemplo, 7 veces más pequeño que SDXL, 14 veces más pequeño que 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).