Nicht alle Aufforderungen sind gleich: Aufforderungsbasiertes Beschneiden von Text-zu-Bild-Diffusionsmodellen

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

June 17, 2024

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

cs.AI

Zusammenfassung

Text-zu-Bild (T2I) Diffusionsmodelle haben beeindruckende Bildgenerierungsfähigkeiten gezeigt. Ihre hohe Rechenintensität hindert jedoch ressourcenbeschränkte Organisationen daran, T2I-Modelle nach dem Feintuning auf ihren internen Ziel-Daten einzusetzen. Während Beschneidungstechniken eine potenzielle Lösung bieten, um die Rechenlast von T2I-Modellen zu reduzieren, verwenden statische Beschneidungsmethoden dasselbe beschnittene Modell für alle Eingabeaufforderungen, ohne auf die unterschiedlichen Kapazitätsanforderungen verschiedener Aufforderungen einzugehen. Die dynamische Beschneidung behebt dieses Problem, indem für jede Aufforderung ein separates Teilnetzwerk verwendet wird, aber sie verhindert die Batch-Parallelität auf GPUs. Um diese Einschränkungen zu überwinden, stellen wir Adaptive Prompt-basierte Beschneidung (APTP) vor, eine neuartige auf Aufforderungen basierende Beschneidungsmethode, die für T2I-Diffusionsmodelle entwickelt wurde. Zentral für unseren Ansatz ist ein Aufforderungsrouter-Modell, das lernt, die erforderliche Kapazität für eine Eingabetext-Aufforderung zu bestimmen und sie einem Architekturcode zuzuweisen, basierend auf einem insgesamt gewünschten Rechenbudget für Aufforderungen. Jeder Architekturcode repräsentiert ein spezialisiertes Modell, das auf die ihm zugewiesenen Aufforderungen zugeschnitten ist, und die Anzahl der Codes ist ein Hyperparameter. Wir trainieren den Aufforderungsrouter und die Architekturcodes unter Verwendung von kontrastivem Lernen, um sicherzustellen, dass ähnliche Aufforderungen nahe beieinander liegenden Codes zugeordnet werden. Darüber hinaus verwenden wir den optimalen Transport, um zu verhindern, dass die Codes in einen einzigen zusammenfallen. Wir zeigen die Wirksamkeit von APTP, indem wir Stable Diffusion (SD) V2.1 unter Verwendung von CC3M und COCO als Ziel-Datensätze beschneiden. APTP übertrifft die Einzelmodell-Beschneidungs-Baselines in Bezug auf FID, CLIP und CMMD-Werte. Unsere Analyse der von APTP gelernten Cluster zeigt, dass sie semantisch sinnvoll sind. Wir zeigen auch, dass APTP automatisch zuvor empirisch herausfordernde Aufforderungen für SD entdecken kann, z. B. Aufforderungen zur Generierung von Textbildern, indem sie diesen höheren Kapazitätscodes zuweist.

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.