Niet Alle Prompts Zijn Gelijk: Prompt-gebaseerd Snoeien van Text-to-Image Diffusiemodellen

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

Samenvatting

Text-to-image (T2I) diffusiemodellen hebben indrukwekkende beeldgeneratiecapaciteiten getoond. Toch verhindert hun rekenintensiteit dat organisaties met beperkte middelen T2I-modellen kunnen inzetten na het finetunen op hun interne doeldataset. Hoewel pruningtechnieken een mogelijke oplossing bieden om de rekenlast van T2I-modellen te verminderen, gebruiken statische pruningmethoden hetzelfde geprunte model voor alle invoerprompts, wat de variërende capaciteitsbehoeften van verschillende prompts over het hoofd ziet. Dynamisch pruning lost dit probleem op door voor elke prompt een apart subnetwerk te gebruiken, maar dit belemmert batchparallelisme op GPU's. Om deze beperkingen te overwinnen, introduceren we Adaptive Prompt-Tailored Pruning (APTP), een nieuwe prompt-gebaseerde pruningmethode ontworpen voor T2I-diffusiemodellen. Centraal in onze aanpak staat een promptroutermodel, dat leert om de benodigde capaciteit voor een invoertekstprompt te bepalen en deze door te sturen naar een architectuurcode, gegeven een totaal gewenst rekenbudget voor prompts. Elke architectuurcode vertegenwoordigt een gespecialiseerd model dat is afgestemd op de prompts die eraan zijn toegewezen, en het aantal codes is een hyperparameter. We trainen de promptrouter en de architectuurcodes met contrastief leren, zodat vergelijkbare prompts worden toegewezen aan nabijgelegen codes. Verder gebruiken we optimaal transport om te voorkomen dat de codes samenvallen in één enkele code. We demonstreren de effectiviteit van APTP door Stable Diffusion (SD) V2.1 te prunen met CC3M en COCO als doeldatasets. APTP presteert beter dan de single-model pruning-baselines in termen van FID, CLIP en CMMD-scores. Onze analyse van de door APTP geleerde clusters laat zien dat ze semantisch betekenisvol zijn. We tonen ook aan dat APTP automatisch eerder empirisch gevonden uitdagende prompts voor SD kan ontdekken, zoals prompts voor het genereren van tekstbeelden, en deze toewijst aan codes met een hogere capaciteit.

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.