SnapGen: Domando Modelos de Texto-para-Imagem de Alta Resolução para Dispositivos Móveis com Arquiteturas e Treinamento 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

Resumo

Os modelos de difusão texto-para-imagem (T2I) existentes enfrentam várias limitações, incluindo tamanhos de modelo grandes, tempo de execução lento e geração de baixa qualidade em dispositivos móveis. Este artigo tem como objetivo abordar todos esses desafios desenvolvendo um modelo T2I extremamente pequeno e rápido que gera imagens de alta resolução e alta qualidade em plataformas móveis. Propomos várias técnicas para alcançar esse objetivo. Primeiramente, examinamos sistematicamente as escolhas de design da arquitetura da rede para reduzir os parâmetros do modelo e a latência, garantindo ao mesmo tempo uma geração de alta qualidade. Em segundo lugar, para melhorar ainda mais a qualidade da geração, empregamos a destilação de conhecimento entre arquiteturas de forma cruzada de um modelo muito maior, utilizando uma abordagem multinível para orientar o treinamento do nosso modelo do zero. Em terceiro lugar, possibilitamos uma geração em poucas etapas integrando orientação adversarial com destilação de conhecimento. Pela primeira vez, nosso modelo SnapGen demonstra a geração de imagens de 1024x1024 px em um dispositivo móvel em cerca de 1,4 segundos. No ImageNet-1K, nosso modelo, com apenas 372M de parâmetros, alcança um FID de 2,06 para geração de 256x256 px. Nos benchmarks T2I (ou seja, GenEval e DPG-Bench), nosso modelo com meros 379M de parâmetros supera modelos em grande escala com bilhões de parâmetros em um tamanho significativamente menor (por exemplo, 7x menor que SDXL, 14x menor 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).