tgoop.com/quant_prune_distill/96
Last Update:
On Architectural Compression of Text-to-Image Diffusion Models
[Статья][Код]
Главным недостатком любимых всеми диффузионных моделей для генерации чего-либо является их скорость. Каждый процесс генерации представляет собой итеративное применение довольно увесистой модели. И потому есть два направления для ускорения:
1️⃣ Уменьшение количества шагов сэмплирования
2️⃣ Удешевление одного шага сэмплирования
Сегодняшний сюжет будет относиться ко второму направлению.
Постановка эксперимента
За основу берется Stable Diffusion v1.4. Как известно, данная модель состоит из UNet модели, генерирующей в латентном пространстве, текстового энкодера, и вариационного кодировщика, погружающего в скрытое пространство и из него. Но так как почти все вычисления приходятся на UNet целесообразно сжимать в первую очередь его.
UNet Stable Diffusion состоит из сверточных ResNet блоков и Transformer блоков.
Первое наблюдение, которое делают авторы, — оказывается, что из UNet Stable Diffusion можно вытащить среднюю часть (с наименьшим разрешением карт активаций) с довольно небольшой просадкой по качеству и метрикам без какого либо дообучения, что любопытно. Выходит, что данный участок сети не выполняет серьезной работы, либо сеть, как человеческая печень, продолжает функционировать в штатном режиме при удалении части от нее.
Далее авторы предлагают три варианта уменьшенной Stable Diffusion:
1️⃣ Base. Меньше блоков в верхних слоях.
2️⃣ Small. Без блоков в середине модели.
3️⃣ Tiny. Укорачивание на промежуточных разрешениях.
Обучать с нуля это хозяйство (оригинальная Stable Diffusion обучалась 150к GPU часов на A100) - удел богатых. И авторы прибегают к дистилляции.
Лосс состоит из трех компонент:
1️⃣ исходного denoising loss
2️⃣ mse с учителем
3️⃣ mse между активациями на выходах каждой стадии (перед понижением/повышением разрешения).
Существенная деталь, от которого зависит итоговое качество всей процедуры, — на каких данных следует дистиллировать в условиях ограниченного бюджета? Авторы используют подвыборку из 220k примеров из LAION-Aesthetics с наибольшими оценками по эстетике.
BY КПД
Share with your friend now:
tgoop.com/quant_prune_distill/96