Об архитектурном сжатии моделей диффузии для генерации изображений из текста

On Architectural Compression of Text-to-Image Diffusion Models

May 25, 2023

Авторы: Bo-Kyeong Kim, Hyoung-Kyu Song, Thibault Castells, Shinkook Choi

cs.AI

Аннотация

Выдающиеся результаты генерации текста в изображение (T2I) моделей Stable Diffusion (SDM) сопровождаются значительными вычислительными затратами. Для решения этой проблемы современные исследования в области эффективных SDM сосредоточены на сокращении количества шагов выборки и использовании квантования сети. В отличие от этих подходов, данное исследование подчеркивает потенциал классического архитектурного сжатия для универсального синтеза T2I, представляя модели SDM с удаленными блоками и дистилляцией знаний (BK-SDMs). Мы удаляем несколько остаточных и внимательных блоков из U-Net SDM, достигая более чем 30% сокращения количества параметров, операций умножения-сложения (MAC) на шаг выборки и задержки. Мы проводим предварительное обучение на основе дистилляции с использованием всего 0,22 миллиона пар данных LAION (менее 0,1% от полного набора обучающих пар) на одном GPU A100. Несмотря на ограниченные ресурсы обучения, наши компактные модели способны имитировать оригинальную SDM, извлекая пользу из перенесенных знаний, и демонстрируют конкурентоспособные результаты по сравнению с более крупными моделями с миллиардами параметров на тесте MS-COCO с нулевым обучением. Кроме того, мы показываем применимость наших легковесных предварительно обученных моделей в персонализированной генерации с использованием тонкой настройки DreamBooth.

English

Exceptional text-to-image (T2I) generation results of Stable Diffusion models (SDMs) come with substantial computational demands. To resolve this issue, recent research on efficient SDMs has prioritized reducing the number of sampling steps and utilizing network quantization. Orthogonal to these directions, this study highlights the power of classical architectural compression for general-purpose T2I synthesis by introducing block-removed knowledge-distilled SDMs (BK-SDMs). We eliminate several residual and attention blocks from the U-Net of SDMs, obtaining over a 30% reduction in the number of parameters, MACs per sampling step, and latency. We conduct distillation-based pretraining with only 0.22M LAION pairs (fewer than 0.1% of the full training pairs) on a single A100 GPU. Despite being trained with limited resources, our compact models can imitate the original SDM by benefiting from transferred knowledge and achieve competitive results against larger multi-billion parameter models on the zero-shot MS-COCO benchmark. Moreover, we demonstrate the applicability of our lightweight pretrained models in personalized generation with DreamBooth finetuning.