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

2024-06-17

arXiv.org

AI Summary

Text-to-image diffusion models incur substantial computational overhead; static pruning overlooks prompt-specific semantic variations, while dynamic pruning compromises GPU batch-level parallelism. To address this, we propose Adaptive Prompt-Tuned Pruning (APTP), the first method to introduce a prompt routing mechanism coupled with architecture-code representations, jointly optimized via contrastive learning and optimal transport constraints. APTP enables semantic-aware, capacity-controllable dynamic subnet selection—automatically identifying challenging prompts (e.g., those requiring text rendering) and assigning them higher-capacity subnetworks. Evaluated on CC3M and COCO, APTP significantly outperforms static pruning baselines, achieving consistent improvements in FID, CLIP Score, and CMMD. The learned subnet cluster exhibits clear semantic grouping, offering both inference efficiency and architectural interpretability.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Prompt Engineering / Prompting
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Web Mining and Content Analysis: Large pretrained models with web data
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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.

Problem

• Reduces computational burden of T2I models
• Adapts pruning to varying prompt requirements
• Maintains batch parallelism on GPUs

Innovation

• Adaptive Prompt-Tailored Pruning (APTP)
• Prompt router model determines capacity
• Contrastive learning for prompt-code mapping