Nem Todos os Prompts São Iguais: Poda Baseada em Prompts de Modelos de Difusão Texto-para-Imagem

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

Resumo

Modelos de difusão texto-para-imagem (T2I) demonstraram capacidades impressionantes de geração de imagens. No entanto, sua intensidade computacional impede que organizações com recursos limitados implantem modelos T2I após ajustá-los em seus dados internos de destino. Embora técnicas de poda ofereçam uma solução potencial para reduzir a carga computacional dos modelos T2I, métodos de poda estática utilizam o mesmo modelo podado para todos os prompts de entrada, ignorando as necessidades variáveis de capacidade de diferentes prompts. A poda dinâmica aborda essa questão utilizando uma sub-rede separada para cada prompt, mas impede o paralelismo em lote em GPUs. Para superar essas limitações, introduzimos a Poda Adaptativa Personalizada por Prompt (APTP), um novo método de poda baseado em prompt projetado para modelos de difusão T2I. Central à nossa abordagem é um modelo de roteamento de prompt, que aprende a determinar a capacidade necessária para um prompt de texto de entrada e o direciona para um código de arquitetura, dado um orçamento computacional total desejado para os prompts. Cada código de arquitetura representa um modelo especializado adaptado aos prompts atribuídos a ele, e o número de códigos é um hiperparâmetro. Treinamos o roteador de prompt e os códigos de arquitetura usando aprendizado contrastivo, garantindo que prompts semelhantes sejam mapeados para códigos próximos. Além disso, empregamos transporte ótimo para evitar que os códigos colapsem em um único. Demonstramos a eficácia da APTP podando o Stable Diffusion (SD) V2.1 usando CC3M e COCO como conjuntos de dados de destino. A APTP supera as linhas de base de poda de modelo único em termos de pontuações FID, CLIP e CMMD. Nossa análise dos clusters aprendidos pela APTP revela que eles são semanticamente significativos. Também mostramos que a APTP pode descobrir automaticamente prompts previamente identificados como desafiadores para o SD, por exemplo, prompts para gerar imagens de texto, atribuindo-os a códigos de maior capacidade.

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.