Sur la compression architecturale des modèles de diffusion texte-image

On Architectural Compression of Text-to-Image Diffusion Models

May 25, 2023

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

cs.AI

Résumé

Les résultats exceptionnels de génération texte-image (T2I) des modèles de diffusion stable (SDMs) s'accompagnent de demandes computationnelles substantielles. Pour résoudre ce problème, les recherches récentes sur les SDMs efficaces ont priorisé la réduction du nombre d'étapes d'échantillonnage et l'utilisation de la quantification des réseaux. En complément de ces approches, cette étude met en lumière la puissance de la compression architecturale classique pour la synthèse T2I à usage général, en introduisant des SDMs avec distillation de connaissances et suppression de blocs (BK-SDMs). Nous éliminons plusieurs blocs résiduels et d'attention de l'U-Net des SDMs, obtenant ainsi une réduction de plus de 30 % du nombre de paramètres, des MACs par étape d'échantillonnage et de la latence. Nous effectuons un pré-entraînement basé sur la distillation avec seulement 0,22 million de paires LAION (moins de 0,1 % des paires d'entraînement complètes) sur un seul GPU A100. Malgré un entraînement avec des ressources limitées, nos modèles compacts peuvent imiter le SDM original en bénéficiant des connaissances transférées et obtiennent des résultats compétitifs par rapport à des modèles plus volumineux comportant plusieurs milliards de paramètres sur le benchmark zéro-shot MS-COCO. De plus, nous démontrons l'applicabilité de nos modèles pré-entraînés légers dans la génération personnalisée avec un ajustement fin via 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.