No Todos los Prompts Son Iguales: Poda Basada en Prompts de Modelos de Difusión de Texto a Imagen

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

June 17, 2024

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

cs.AI

Resumen

Los modelos de difusión de texto a imagen (T2I) han demostrado capacidades impresionantes en la generación de imágenes. Sin embargo, su intensidad computacional impide que organizaciones con recursos limitados implementen modelos T2I después de ajustarlos con sus datos internos objetivo. Si bien las técnicas de poda ofrecen una solución potencial para reducir la carga computacional de los modelos T2I, los métodos de poda estática utilizan el mismo modelo podado para todas las indicaciones de entrada, ignorando los requisitos de capacidad variables de diferentes indicaciones. La poda dinámica aborda este problema utilizando una subred separada para cada indicación, pero impide el paralelismo por lotes en las GPU. Para superar estas limitaciones, presentamos la Poda Adaptativa Personalizada por Indicación (APTP), un novedoso método de poda basado en indicaciones diseñado para modelos de difusión T2I. El núcleo de nuestro enfoque es un modelo enrutador de indicaciones, que aprende a determinar la capacidad requerida para una indicación de texto de entrada y la dirige a un código de arquitectura, dado un presupuesto computacional total deseado para las indicaciones. Cada código de arquitectura representa un modelo especializado adaptado a las indicaciones asignadas a él, y el número de códigos es un hiperparámetro. Entrenamos el enrutador de indicaciones y los códigos de arquitectura utilizando aprendizaje contrastivo, asegurando que indicaciones similares se asignen a códigos cercanos. Además, empleamos transporte óptimo para evitar que los códigos colapsen en uno solo. Demostramos la efectividad de APTP podando Stable Diffusion (SD) V2.1 utilizando CC3M y COCO como conjuntos de datos objetivo. APTP supera a los baselines de poda de un solo modelo en términos de puntuaciones FID, CLIP y CMMD. Nuestro análisis de los clústeres aprendidos por APTP revela que son semánticamente significativos. También mostramos que APTP puede descubrir automáticamente indicaciones previamente identificadas como desafiantes para SD, por ejemplo, indicaciones para generar imágenes de texto, asignándolas a códigos de mayor capacidad.

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.