Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models
[Статья][Код]

Введение

Сжимать можно не только матрицы весов LLMов, но много еще что-то другое не без пользы для человечества.

Более-менее все современные SOTA или около того диффузионные модели работают в латентном пространстве, с меньшим пространственным разрешением.

От автокодировщика (тандема 🚲 из сети, кодирующей из исходного пространства данных в латентное, называемой энкодером, и декодирующей обратно - декодера) хочется одновременно высокой степени сжатия и качественной реконструкции одновременно.

Наиболее популярный пространственный фактор сжатия - 8x (Stable Diffusion, FLUX, DALLE-3 (по всей видимости)). При таком сжатии (особенно если еще и поднять число каналов в латентном пространстве до 16, скажем) реконструкции почти не отличаются от оригинала. Однако карты признаков остаются все еще достаточно большими при подаче картинок высокого разрешения, и требуется немало вычислений.

Отсюда возникает желание двигаться в сторону больших факторов сжатия при сохранении качества. И авторы данной статьи предлагают несколько интересных архитектурных решений.

Метод

По сути в предложенной статье две ключевых идеи:

1️⃣ Skip-connection (Residual Autoencoding):
- При downsample патчи 2x2 сворачиваются в 4 канала, а зачем группы по 2 канала усредняются (ребята изобрели average pooling?)
- При upsample 4 подряд идущих канала разворачиваются в 2x2, каждый канал дублируется
2️⃣ Процедура обучения
- Предобучение на низком разрешении только на реконструкцию (без адверсариального лосса)
- Дообучение на низком разрешении верхних слоев декодера с GAN лоссом (для получения резких 🔪 деталей)
- Дообучение глубоких слоев энкодера и декодера на реконструкцию на высокое разрешение (дабы адаптировать латенты под high-resolution c другим распределением частот / плотностью обьектов)