Sobre a Compressão Arquitetural de Modelos de Difusão Texto-para-Imagem

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

Resumo

Os resultados excepcionais de geração de texto para imagem (T2I) dos modelos Stable Diffusion (SDMs) vêm acompanhados de demandas computacionais substanciais. Para resolver esse problema, pesquisas recentes sobre SDMs eficientes priorizaram a redução do número de etapas de amostragem e a utilização de quantização de rede. De forma ortogonal a essas direções, este estudo destaca o poder da compressão arquitetônica clássica para síntese T2I de propósito geral, introduzindo SDMs com conhecimento destilado e blocos removidos (BK-SDMs). Eliminamos vários blocos residuais e de atenção da U-Net dos SDMs, obtendo uma redução de mais de 30% no número de parâmetros, MACs por etapa de amostragem e latência. Realizamos pré-treinamento baseado em destilação com apenas 0,22 milhões de pares LAION (menos de 0,1% dos pares de treinamento completos) em uma única GPU A100. Apesar de serem treinados com recursos limitados, nossos modelos compactos podem imitar o SDM original ao se beneficiar do conhecimento transferido e alcançar resultados competitivos em comparação com modelos maiores com bilhões de parâmetros no benchmark zero-shot MS-COCO. Além disso, demonstramos a aplicabilidade de nossos modelos pré-treinados leves na geração personalizada com ajuste fino do 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.