SnapGen : Dompter les modèles texte-image haute résolution pour appareils mobiles avec des architectures et un entraînement efficaces

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

December 12, 2024

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

Résumé

Les modèles de diffusion texte-vers-image (T2I) existants rencontrent plusieurs limitations, notamment des tailles de modèle importantes, des temps d'exécution lents et une génération de faible qualité sur les appareils mobiles. Cet article vise à relever tous ces défis en développant un modèle T2I extrêmement petit et rapide qui génère des images haute résolution et de haute qualité sur les plateformes mobiles. Nous proposons plusieurs techniques pour atteindre cet objectif. Tout d'abord, nous examinons systématiquement les choix de conception de l'architecture du réseau pour réduire les paramètres du modèle et la latence, tout en garantissant une génération de haute qualité. Ensuite, pour améliorer encore la qualité de la génération, nous utilisons la distillation des connaissances entre architectures à partir d'un modèle beaucoup plus grand, en utilisant une approche multi-niveaux pour guider l'entraînement de notre modèle à partir de zéro. Troisièmement, nous permettons une génération en quelques étapes en intégrant un guidage adversarial avec la distillation des connaissances. Pour la première fois, notre modèle SnapGen démontre la génération d'images de 1024x1024 px sur un appareil mobile en environ 1,4 seconde. Sur ImageNet-1K, notre modèle, avec seulement 372M de paramètres, atteint un FID de 2,06 pour une génération de 256x256 px. Sur les bancs d'essai T2I (c'est-à-dire GenEval et DPG-Bench), notre modèle, avec seulement 379M de paramètres, surpasse des modèles à grande échelle avec des milliards de paramètres, tout en étant significativement plus petit (par exemple, 7 fois plus petit que SDXL, 14 fois plus petit 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).