SnapGen: Domare modelli di testo-immagine ad alta risoluzione per dispositivi mobili con architetture e addestramento efficienti

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

December 12, 2024

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

Abstract

I modelli di diffusione testo-immagine (T2I) esistenti affrontano diverse limitazioni, tra cui dimensioni del modello ingenti, lentezza di esecuzione e generazione di bassa qualità su dispositivi mobili. Questo articolo si propone di affrontare tutte queste sfide sviluppando un modello T2I estremamente piccolo e veloce che genera immagini ad alta risoluzione e di alta qualità su piattaforme mobili. Proponiamo diverse tecniche per raggiungere questo obiettivo. In primo luogo, esaminiamo sistematicamente le scelte progettuali dell'architettura di rete per ridurre i parametri del modello e la latenza, garantendo nel contempo una generazione di alta qualità. In secondo luogo, per migliorare ulteriormente la qualità della generazione, utilizziamo la distillazione della conoscenza tra architetture da un modello molto più grande, utilizzando un approccio multi livello per guidare l'addestramento del nostro modello da zero. In terzo luogo, abilitiamo una generazione a pochi passaggi integrando la guida avversaria con la distillazione della conoscenza. Per la prima volta, il nostro modello SnapGen dimostra la generazione di immagini da 1024x1024 px su un dispositivo mobile in circa 1,4 secondi. Su ImageNet-1K, il nostro modello, con soli 372M parametri, raggiunge un FID di 2,06 per la generazione da 256x256 px. Sui benchmark T2I (ad esempio, GenEval e DPG-Bench), il nostro modello con soli 379M parametri supera modelli su larga scala con miliardi di parametri in dimensioni significativamente minori (ad esempio, 7 volte più piccolo di SDXL, 14 volte più piccolo di 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).