Sobre la Compresión Arquitectónica de Modelos de Difusión de Texto a Imagen

On Architectural Compression of Text-to-Image Diffusion Models

May 25, 2023

Autores: Bo-Kyeong Kim, Hyoung-Kyu Song, Thibault Castells, Shinkook Choi

cs.AI

Resumen

Los excepcionales resultados de generación de texto a imagen (T2I) de los modelos Stable Diffusion (SDMs) vienen acompañados de demandas computacionales sustanciales. Para resolver este problema, investigaciones recientes sobre SDMs eficientes han priorizado la reducción del número de pasos de muestreo y el uso de cuantización de redes. De manera ortogonal a estas direcciones, este estudio destaca el poder de la compresión arquitectónica clásica para la síntesis T2I de propósito general mediante la introducción de SDMs con eliminación de bloques y destilación de conocimiento (BK-SDMs). Eliminamos varios bloques residuales y de atención de la U-Net de los SDMs, logrando una reducción de más del 30% en el número de parámetros, operaciones MAC por paso de muestreo y latencia. Realizamos un preentrenamiento basado en destilación con solo 0.22 millones de pares de LAION (menos del 0.1% de los pares de entrenamiento completos) en una sola GPU A100. A pesar de ser entrenados con recursos limitados, nuestros modelos compactos pueden imitar al SDM original al beneficiarse del conocimiento transferido y logran resultados competitivos frente a modelos más grandes con miles de millones de parámetros en el benchmark zero-shot MS-COCO. Además, demostramos la aplicabilidad de nuestros modelos preentrenados livianos en la generación personalizada mediante el ajuste fino con DreamBooth.

English

Exceptional text-to-image (T2I) generation results of Stable Diffusion models (SDMs) come with substantial computational demands. To resolve this issue, recent research on efficient SDMs has prioritized reducing the number of sampling steps and utilizing network quantization. Orthogonal to these directions, this study highlights the power of classical architectural compression for general-purpose T2I synthesis by introducing block-removed knowledge-distilled SDMs (BK-SDMs). We eliminate several residual and attention blocks from the U-Net of SDMs, obtaining over a 30% reduction in the number of parameters, MACs per sampling step, and latency. We conduct distillation-based pretraining with only 0.22M LAION pairs (fewer than 0.1% of the full training pairs) on a single A100 GPU. Despite being trained with limited resources, our compact models can imitate the original SDM by benefiting from transferred knowledge and achieve competitive results against larger multi-billion parameter models on the zero-shot MS-COCO benchmark. Moreover, we demonstrate the applicability of our lightweight pretrained models in personalized generation with DreamBooth finetuning.