Non Tutti i Prompt Sono Uguali: Potatura Basata su Prompt dei Modelli di Diffusione Testo-Immagine

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

June 17, 2024

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

cs.AI

Abstract

I modelli di diffusione Text-to-Image (T2I) hanno dimostrato capacità impressionanti nella generazione di immagini. Tuttavia, la loro intensità computazionale impedisce alle organizzazioni con risorse limitate di implementare modelli T2I dopo averli ottimizzati sui loro dati target interni. Sebbene le tecniche di pruning offrano una potenziale soluzione per ridurre il carico computazionale dei modelli T2I, i metodi di pruning statico utilizzano lo stesso modello potato per tutti gli input, trascurando i requisiti di capacità variabili per prompt diversi. Il pruning dinamico affronta questo problema utilizzando una sotto-rete separata per ogni prompt, ma impedisce il parallelismo batch sulle GPU. Per superare queste limitazioni, introduciamo l'Adaptive Prompt-Tailored Pruning (APTP), un nuovo metodo di pruning basato sui prompt progettato per i modelli di diffusione T2I. Al centro del nostro approccio c'è un modello di routing dei prompt, che impara a determinare la capacità richiesta per un prompt di testo in input e lo indirizza a un codice di architettura, dato un budget computazionale totale desiderato per i prompt. Ogni codice di architettura rappresenta un modello specializzato adattato ai prompt assegnati, e il numero di codici è un iperparametro. Addestriamo il router dei prompt e i codici di architettura utilizzando l'apprendimento contrastivo, garantendo che prompt simili siano mappati su codici vicini. Inoltre, impieghiamo il trasporto ottimale per evitare che i codici collassino in uno solo. Dimostriamo l'efficacia di APTP potando Stable Diffusion (SD) V2.1 utilizzando CC3M e COCO come dataset target. APTP supera le baseline di pruning a modello singolo in termini di punteggi FID, CLIP e CMMD. La nostra analisi dei cluster appresi da APTP rivela che sono semanticamente significativi. Mostriamo anche che APTP può scoprire automaticamente prompt precedentemente considerati empiricamente difficili per SD, ad esempio prompt per generare immagini di testo, assegnandoli a codici di capacità superiore.

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.