すべてのプロンプトが平等ではない：テキストから画像への拡散モデルのプロンプトベースプルーニング

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モデルの計算負荷を軽減する可能性を秘めていますが、静的プルーニング手法ではすべての入力プロンプトに対して同じプルーニングされたモデルを使用するため、異なるプロンプトの能力要件の違いを見落としています。動的プルーニングはこの問題に対処するために各プロンプトに別々のサブネットワークを利用しますが、GPU上でのバッチ並列処理を妨げます。これらの制限を克服するため、我々はT2I拡散モデル向けの新しいプロンプトベースのプルーニング手法であるAdaptive Prompt-Tailored Pruning（APTP）を提案します。我々のアプローチの中核となるのは、プロンプトルーターモデルです。このモデルは、入力テキストプロンプトに必要な能力を決定し、プロンプト全体の計算予算を考慮してアーキテクチャコードにルーティングすることを学習します。各アーキテクチャコードは、それに割り当てられたプロンプトに特化したモデルを表し、コードの数はハイパーパラメータです。プロンプトルーターとアーキテクチャコードは、類似したプロンプトが近いコードにマッピングされるように、コントラスティブ学習を用いて訓練されます。さらに、コードが単一のものに崩壊するのを防ぐために最適輸送を利用します。我々は、CC3MとCOCOをターゲットデータセットとしてStable Diffusion（SD）V2.1をプルーニングすることでAPTPの有効性を実証します。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.