Не все подсказки равнозначны: обрезка на основе подсказок моделей диффузии текста в изображение

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

June 17, 2024

Авторы: Alireza Ganjdanesh, Reza Shirkavand, Shangqian Gao, Heng Huang

cs.AI

Аннотация

Модели диффузии текста в изображение (T2I) продемонстрировали впечатляющие возможности по генерации изображений. Тем не менее, их вычислительная интенсивность мешает организациям с ограниченными ресурсами развертывать модели T2I после их донастройки на внутренних целевых данных. В то время как методы обрезки предлагают потенциальное решение для снижения вычислительной нагрузки моделей T2I, статические методы обрезки используют одну и ту же обрезанную модель для всех входных запросов, не учитывая различные требования к емкости различных запросов. Динамическая обрезка решает эту проблему, используя отдельную подсеть для каждого запроса, но она препятствует параллелизму по пакетам на графических процессорах. Для преодоления этих ограничений мы представляем метод Адаптивной Обрезки Под Запрос (APTP), новый метод обрезки, разработанный для моделей диффузии T2I. Центральным элементом нашего подхода является модель маршрутизатора запросов, которая учится определять необходимую емкость для входного текстового запроса и направлять его к коду архитектуры, учитывая общий желаемый бюджет вычислений для запросов. Каждый код архитектуры представляет собой специализированную модель, адаптированную к присвоенным ей запросам, и количество кодов является гиперпараметром. Мы обучаем маршрутизатор запросов и коды архитектуры с использованием контрастного обучения, обеспечивая отображение похожих запросов в близкие коды. Кроме того, мы используем оптимальный транспорт, чтобы предотвратить схлопывание кодов в один. Мы демонстрируем эффективность APTP путем обрезки Stable Diffusion (SD) V2.1 с использованием наборов данных CC3M и COCO. APTP превосходит базовые методы обрезки одиночных моделей по показателям FID, CLIP и CMMD. Наш анализ кластеров, выученных APTP, показывает, что они имеют семантический смысл. Мы также показываем, что APTP может автоматически обнаруживать ранее эмпирически найденные сложные запросы для SD, например, запросы для генерации текстовых изображений, присваивая им коды более высокой емкости.

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.