Sulla Compressione Architetturale dei Modelli di Diffusione da Testo a Immagine

On Architectural Compression of Text-to-Image Diffusion Models

May 25, 2023

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

cs.AI

Abstract

I risultati eccezionali nella generazione da testo a immagine (T2I) dei modelli Stable Diffusion (SDM) comportano notevoli esigenze computazionali. Per risolvere questo problema, le ricerche recenti sugli SDM efficienti si sono concentrate sulla riduzione del numero di passaggi di campionamento e sull'utilizzo della quantizzazione della rete. In direzione ortogonale a questi approcci, questo studio evidenzia il potenziale della compressione architetturale classica per la sintesi T2I generica, introducendo SDM con distillazione della conoscenza e rimozione di blocchi (BK-SDM). Eliminiamo diversi blocchi residui e di attenzione dalla U-Net degli SDM, ottenendo una riduzione di oltre il 30% nel numero di parametri, nelle operazioni MAC per passaggio di campionamento e nella latenza. Eseguiamo un pre-addestramento basato su distillazione con soli 0,22 milioni di coppie LAION (meno dello 0,1% delle coppie di addestramento complete) su una singola GPU A100. Nonostante siano stati addestrati con risorse limitate, i nostri modelli compatti sono in grado di imitare l'SDM originale, beneficiando della conoscenza trasferita, e ottengono risultati competitivi rispetto a modelli più grandi con miliardi di parametri nel benchmark zero-shot MS-COCO. Inoltre, dimostriamo l'applicabilità dei nostri modelli pre-addestrati leggeri nella generazione personalizzata con il fine-tuning di 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.