Tous les prompts ne se valent pas : élagage basé sur les prompts pour les modèles de diffusion texte-image

Not All Prompts Are Made Equal: Prompt-based Pruning of Text-to-Image Diffusion Models

June 17, 2024

Auteurs: Alireza Ganjdanesh, Reza Shirkavand, Shangqian Gao, Heng Huang

cs.AI

Résumé

Les modèles de diffusion texte-image (T2I) ont démontré des capacités impressionnantes en génération d'images. Cependant, leur intensité computationnelle empêche les organisations aux ressources limitées de déployer des modèles T2I après les avoir affinés sur leurs données internes cibles. Bien que les techniques d'élagage offrent une solution potentielle pour réduire la charge computationnelle des modèles T2I, les méthodes d'élagage statique utilisent le même modèle élagué pour toutes les invites d'entrée, négligeant les besoins variables en capacité des différentes invites. L'élagage dynamique aborde ce problème en utilisant un sous-réseau distinct pour chaque invite, mais il empêche le parallélisme par lots sur les GPU. Pour surmonter ces limitations, nous introduisons l'Élagage Adaptatif sur Mesure pour les Invites (Adaptive Prompt-Tailored Pruning, APTP), une nouvelle méthode d'élagage basée sur les invites conçue pour les modèles de diffusion T2I. Au cœur de notre approche se trouve un modèle de routage des invites, qui apprend à déterminer la capacité requise pour une invite texte donnée et la dirige vers un code d'architecture, en fonction d'un budget de calcul total souhaité pour les invites. Chaque code d'architecture représente un modèle spécialisé adapté aux invites qui lui sont assignées, et le nombre de codes est un hyperparamètre. Nous entraînons le routeur d'invites et les codes d'architecture en utilisant l'apprentissage contrastif, garantissant que des invites similaires sont mappées à des codes proches. De plus, nous utilisons le transport optimal pour empêcher les codes de converger vers un seul. Nous démontrons l'efficacité d'APTP en élaguant Stable Diffusion (SD) V2.1 en utilisant CC3M et COCO comme ensembles de données cibles. APTP surpasse les bases de référence d'élagage à modèle unique en termes de scores FID, CLIP et CMMD. Notre analyse des clusters appris par APTP révèle qu'ils sont sémantiquement significatifs. Nous montrons également qu'APTP peut découvrir automatiquement des invites précédemment identifiées comme difficiles pour SD, par exemple les invites pour générer des images de texte, en leur assignant des codes de capacité plus élevée.

English

Text-to-image (T2I) diffusion models have demonstrated impressive image generation capabilities. Still, their computational intensity prohibits resource-constrained organizations from deploying T2I models after fine-tuning them on their internal target data. While pruning techniques offer a potential solution to reduce the computational burden of T2I models, static pruning methods use the same pruned model for all input prompts, overlooking the varying capacity requirements of different prompts. Dynamic pruning addresses this issue by utilizing a separate sub-network for each prompt, but it prevents batch parallelism on GPUs. To overcome these limitations, we introduce Adaptive Prompt-Tailored Pruning (APTP), a novel prompt-based pruning method designed for T2I diffusion models. Central to our approach is a prompt router model, which learns to determine the required capacity for an input text prompt and routes it to an architecture code, given a total desired compute budget for prompts. Each architecture code represents a specialized model tailored to the prompts assigned to it, and the number of codes is a hyperparameter. We train the prompt router and architecture codes using contrastive learning, ensuring that similar prompts are mapped to nearby codes. Further, we employ optimal transport to prevent the codes from collapsing into a single one. We demonstrate APTP's effectiveness by pruning Stable Diffusion (SD) V2.1 using CC3M and COCO as target datasets. APTP outperforms the single-model pruning baselines in terms of FID, CLIP, and CMMD scores. Our analysis of the clusters learned by APTP reveals they are semantically meaningful. We also show that APTP can automatically discover previously empirically found challenging prompts for SD, e.g., prompts for generating text images, assigning them to higher capacity codes.