Diffusion Models Beat GANs on Image Synthesis

Введение

В этой статье авторы показывают, что диффузионные модели могут обеспечить наилучшее качество генерации изображений. Они достигают этого на безусловном синтезе изображений (unconditional image synthesis): находят наилучшую архитектуру с помощью серии экспериментов. Для условного синтеза изображений (conditional image synthesis) они дополнительно улучшают качество с помощью classifier guidance.

GANs vs Diffusion Models

За последние несколько лет генеративные модели приобрели способность генерировать человеческий естественный язык, бесконечные высококачественные синтетические изображения и разнообразную музыку, написанную людьми.

На момент выхода статьи GAN’ы (Generative Adversarial Networks) являлись state-of-the-art в генерации изображений по таким метрикам как FID, Inception Score и Precision. Однако GAN’ы тяжело обучать, они коллапсируют без тщательного подбора гиперпараметров и регуляризации (подробнее об этом можно почитать тут). Более того, GAN’ы захватывают меньше diversity, чем state-of-the-art likelihood-based модели (основанные на правдоподобии). Примеры моделей: первая, вторая, третья.

Вероятностные (likelihood) модели обладают высоким diversity и скейлятся легче, чем GAN’ы, но не достигают того же высокого качества генераций.

Однако диффузионные модели все еще уступают GAN’ам на сложных датасетах — LSUN и Imagenet.

Разрыв между моделями диффузии и GAN’ами авторы объясняют двумя факторами:

1. Архитектуры GAN’ов тщательно изучены и усовершенствованы;
2. GAN’ы способны пожертвовать разнообразием ради реалистичности (diversity-fidelity trade-off), создать образцы высокого качества, но не охватить все распределение.

Авторы стремятся привнести эти преимущества в диффузионные модели, сначала улучшая их архитектуру, а затем разрабатывая схему diversity-fidelity trade-off. Благодаря таким улучшениям авторы достигают state-of-the-art, превосходя GAN’ы по нескольким показателям и наборам данных.

Диффузия

Если верхнеуровнево: диффузия моделирует образец из распределения, обращая вспять процесс постепенного зашумления. В частности процесс семплирования начинается с шума \( x_T \) и постепенно создает менее шумные семплы (reverse diffusion, где \( p_{\theta}(x_{t-1}|x_t) \) — желаемая модель с параметрами \( {\theta} \)) \( x_{T-1} \), \( x_{T-2} \),… до достижения конечного образца \( x_0 \). Процесс зашумления называется прямым, или forward diffusion (\( q(x_t|x_{t-1}) \)). Каждый временной шаг \( t \) соответствует определенному уровню шума, и \( x_t \) можно рассматривать как смесь сигнала \( x_0 \) с некоторым шумом, где отношение сигнал/шум определяется временным шагом \( t \).

В оставшейся части статьи авторы предполагают, что шум был получен из diagonal Gaussian distribution (многомерное нормальное распределение с variance-covariance матрицей, где ненулевые значения — только диагональные элементы). Это хорошо работает для изображений и упрощает различные выводы.

Основы denoising diffusion probabilistic models (DDPM)

Forward Diffusion

Пусть есть истинное распределение данных \( q(x_0) \) . Зададим марковский процесс шумообразования \( q \), который зашумляет данные из \( q(x_0) \) и создает зашумленные выборки \( x_1, …, x_T \), которые принадлежат распределениям \( q(x_1), …, q(x_T) \) соответственно. Если у \( T \) большие значения, то \( q(x_T) \) сходится к гауссову шуму.

На каждой итерации процесса зашумления добавляется гауссов шум [в теории бывают разные шумы, почитать про них можно тут] с некоторой дисперсией, которая задается коэффициентом \( \beta_t \) и расписанием (schedule). Простыми словами — мы берем изображение, сэмплируем нормальный шум (такого же размера, как и наше изображение) с дисперсией \( \beta_t \) и нулевым средним, складываем такой шум с картинкой, умноженной на \( \sqrt{1-\beta_t} \).

Для получения \( x_t \) необходимо пробежаться по всем \( t \), но это очень медленно. У такого семплирования есть свойство: можно явно выразить \( q(x_t|x_0) \) в закрытой форме. Это дает возможность быстро семплировать \( x_t \) из \( x_0 \).

Достигается это за счет reparameterization trick [подробнее про него тут, страница 4], расписывания формул для \( x_t \) и \( x_{t-1} \) и следующего факта: сумма двух независимых нормальных случайных величин остается нормальной там, где среднее значение является суммой двух средних, а дисперсия представляет собой сумму двух дисперсий.

Итак, мы научились из чистой картинки семплировать шумные версии и делать переход из нее в гауссов шум. Такой переход и называется forward diffusion process. Было бы круто уметь его обращать вспять — возвращаться из шума в чистую картинку. Тут возникает проблема: теоретически такое обращение не посчитать. Но! Если что-то мы не можем посчитать, значит, это что-то можем выучить. Следовательно, будем учить модель обращать вспять такой процесс зашумления.

Reverse Diffusion

Задача обратного процесса диффузии (\( q(x_{t-1}|x_{t}) \)) — при зашумленном семпле на шаге \( t \) получить семпл для шага \( t-1 \) (менее зашумленный). Мы не можем семплировать из \( q(x_{t-1}|x_{t}) \) напрямую, так как для этого необходима генеральная совокупность объектов \( x \). Здесь на помощь приходит нейронная сеть.

Если мы хотим семплировать из распределения \( q(x_0) \), то мы можем сначала семплировать из \( q(x_T) \), а затем выполнять обратные шаги \( q(x_{t−1}|x_t) \), пока не достигнем \( x_0 \).

Обучение

Как мы поняли ранее, наша сеть должна выучить обратный процесс \( q(x_{t-1}|x_t) \). Как выглядит это \( q(x_{t-1}|x_t)\( ? Считаем, что \( q(x_{t-1}|x_t)\( тоже имеет гауссово распределение с некими параметрами \( \mu \) и \( \Sigma \). Таким образом, наша модель будет принимать на вход \( x_t \) и предсказывать \( \mu \) и \( \Sigma \) для \( q(x_{t-1}|x_t) \). Однако, чтобы ее научить такое предсказывать, надо иметь сами \( \mu \) и \( \Sigma \). Получить их можно, если расписать \( q(x_{t-1}|x_t, x_0) \) через теорему Байеса (\( x_0 \) здесь добавляется для того, чтобы обратный процесс был отслеживаемым (tractable)). Все выводы можно посмотреть в этой статье.

Мы способны выразить \( q(x_t|x_0) \) в закрытой форме, значит, мы можем подставить в уравнение выше функцию от \( x_t \) вместо значения \( x_0 \):

Также можно обучать предсказывать значение стандартной нормальной величины (\( \varepsilon_{\theta} \)). Тогда средние будут выражены в таком виде:

Наша модель:

\( \Sigma_{\theta}$ фиксируем, выучиваем только $\mu_{\theta} \).

Функция потерь — дивергенция Кульбака-Лейблера между reverse-процессом и выходом модели на шаге \( t \):

\( D_{KL}(q(x_{t-1}|x_t, x_0) || p_{\theta}(x_{t-1}|x_t)) \)

\( \argmin_{\theta}D_{KL}(q(x_{t-1}|x_t, x_0)) || p_{\theta}(x_{t-1}|x_t)\\=\argmin_{\theta}D_{KL}(\mathcal{N}(x_{t-1}; \mu_{q}, \Sigma_{q}(t))||\mathcal{N}(x_{t-1};\mu_{\theta}, \Sigma_q(t)))\\=\argmin_{\theta}\frac{1}{2}[\log\frac{\Sigma_q(t)}{\Sigma_q(t)}-d+tr(\Sigma_q(t)^{-1}\Sigma_q(t))+(\mu_{\theta} — \mu_{q})^{T}\Sigma_q(t)^{-1}(\mu_{\theta}-\mu_q)]\\=\argmin_{\theta}\frac{1}{2}[\log{1}-d+d+(\mu_{\theta} — \mu_{q})^{T}\Sigma_q(t)^{-1}(\mu_{\theta}-\mu_q)]\\=\argmin_{\theta}\frac{1}{2}[(\mu_{\theta} — \mu_{q})^{T}\Sigma_q(t)^{-1}(\mu_{\theta}-\mu_q)]\\=\argmin_{\theta}\frac{1}{2}[(\mu_{\theta} — \mu_{q})^{T}(\sigma_q^{2}(t)\mathcal{I})^{-1}(\mu_{\theta}-\mu_q)]\\=\argmin_{\theta}\frac{1}{2\sigma_q^{2}(t)}[||\mu_{\theta}-\mu_q||_{2}^2] \)

Итоговая функция потерь — MSE (mean squared error) между векторами средних \( \mu_{\theta} \) и \( \mu_{q} \), поделенная на нормализационную константу. Вспоминая reparameterization trick и забывая нормализационную константу, мы получаем упрощенную функцию потерь — MSE между шумом истинным и предсказанным.

В итоге алгоритм обучения выглядит следующим образом:

1. Берем \( x_0 \) из \( q(x_0) \);
2. Сэмплируем \( t \sim U(1, …, T) \);
3. Сэмплируем \( \varepsilon \sim \mathcal{N}(0, \mathcal{I}) \);
4. Делаем градиентный шаг \( ∇_{\theta}L_{simple} \);
5. Повторяем шаги 2-5, пока не сойдется функция потерь.

Алгоритм сэмплирования:

1. Сэмплируем \( x_T \sim \mathcal{N}(0, \mathcal{I}) \);
2. Для каждого \( t=T,…,1 \) делаем шаги 3 и 4;
3. \( z \sim \mathcal{N}(0, \mathcal{I}) \) if \( t>1 \), else \( z=0 \);
4. \( x_{t-1}=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{1-\alpha_t}{\sqrt{1-\overline\alpha_t}}\varepsilon_{\theta}(x_t, t))+\sigma_tz \), (\( \sigma_t=\beta_t \));
5. Возвращаем \( x_0 \).

Больше математики можно найти в следующих источниках:

• appendix в оригинальной статье;
• статья DDPM;
• статья Understanding Diffusion Models: A Unified Perspective.

Улучшения DDPM

Давайте поговорим о том, что авторы решили добавить к диффузии из статьи DDPM, чтобы улучшить качество генерации.

1. Выучиваем дисперсию

Разумные верхние и нижние границы дисперсии ограничены (подробности можно найти в оригинальной статье). На рисунке ниже бета с волной — нижняя граница, без волны — верхняя.

Такой диапазон тяжело предсказывать напрямую, поэтому вместо выучивания дисперсии напрямую выучивают \( v \):

Важно отметить: \( L_{simple} \) не зависит от дисперсии, поэтому лосс модифицируют в \( L_{hybrid} = L_{simple} + \lambda L_{vlb} \), где \( \lambda = 0.001 \), а \( L_{vlb} \) получается из evidence lower bound:

Авторы взяли это себе на вооружение 😉

2. Используем DDIM

Идея следующая: переопределить процесс диффузии как немарковский (теперь текущее состояние всегда зависит от нулевого, начального), генерализировать формулу для обратного процесса диффузии, занулить дисперсию. Так мы быстрее будем семплировать.

Авторы решают применять DDIM, когда количество шагов семплирования меньше 50.

Метрики

Авторы используют Precision или IS для оценки fidelity, Recall — для оценки diversity.

Улучшения архитектуры

Авторами DDPM было предложено использовать UNet в качестве модели диффузии. От классического UNet’а он отличается:

• наличием слоев внимания;
• приемом на вход не только x, но и шага диффузии t (для t применяется positional encoding, как в статье Attention Is All You Need).

Прочитать и разобрать поэтапно работу такой UNet модели можно тут. Там же вы найдете процесс обучения и семплирования (крайне рекомендуем к ознакомлению 🙂).

Авторы статьи, которую мы разбираем, исследовали следующие архитектурные изменения:

Для всех сравнений в этой секции они использовали модель, обученную на ImageNet 128×128 с batch size’ом 256, а также 250 шагов для семплирования.

Рескейл

Если кратко — рескейл не улучшает ситуацию, поэтому авторы его далее не используют.

Проверка голов

Больше голов в multi-head attention или меньше каналов для каждой головы — оба варианта улучшают FID.

Увеличение глубины

Увеличение глубины (res=4, это количество блоков на каждом разрешении) авторы тоже решили не брать (на графике видно, что времени на обучение требуется больше). 64 канала для каждой головы в multi-head attention — лучшие показатели и, следовательно, выбор авторов.

Итог

BigGAN up/down sampling — это просто более навороченная и глубокая версия блока:

Adaptive Group Normalization

В DDPM было так (код тут):

То есть conv -> norm -> adding time embedding.

В Adaptive Group Normalization стало так: conv -> norm -> scale with time embedding -> adding class embedding (conditional image generation).

Если мы можем использовать эмбеддинги для времени \( t \) — можем использовать и class эмбеддинги, которые получаем как выход линейного слоя с num_classes входами.

Classifier Guidance

Важно также изучить различные способы добавления меток классов в модели диффузии. Авторы уже включили информацию о классах в слои нормализации. Здесь они исследуют другой подход: использование классификатора \( p(y|x) \) для улучшения генерации.

Как этого достичь? Авторы показывают один из способов: предварительно обученная диффузионная модель может быть обусловлена использованием градиентов классификатора. В частности мы можем обучить классификатор \( p(y|x_t) \) на зашумленных изображениях \( x_t \), а затем применить градиенты \( ∇_{x_t}\log(p(y|x_t)) \), чтобы направить процесс диффузионной выборки к произвольной метке класса y.

Расписываем градиент (по \( x_t \)) log-likelihood через Байеса, получаем два градиента: unconditional (диффузионная модель) и conditional (классификатор), \( ∇_{x_t}\log(p(x_t)) \)_ и \( ∇_{x_t}\log(p(y|x_t)) \) соответственно.

Алгоритм 1 — classifier guided семплирование.

Алгоритм 2 — classifier guided DDIM семплирование.

Это достаточно сложные концепты: понять, почему и как это работает можно, если почитать главы 4.1 и 4.2 статьи или сразу посмотреть Classifier-Free Diffusion Guidance.

Scaling Classifier Gradients

Классификатор обучен на ImageNet. Его архитектура представляет собой downsampling часть модели UNet.

Классификаторы обучаются на том же распределении шума, что и соответствующая диффузионная модель + также добавлены случайные кропы для уменьшения переобучения. После обучения классификатор внедряется в процесс семплирования.

Авторы провели эксперименты со скейлингом градиентов при семплировании:

На основе данных экспериментов видим:

• Scale — это trade-off между precision и recall;
• Guidance улучшает как условную, так и безусловную генерацию.

Выводы

Чтобы оценить улучшенную архитектуру модели при безусловной генерации изображений, авторы обучили отдельные диффузионные модели на трех классах LSUN: спальня, лошадь и кошка.

Чтобы оценить classifier guidance, авторы обучили модели условной диффузии на наборе данных ImageNet с разрешением 128×128, 256×256 и 512×512.

Диффузионные модели позволяют получить наилучший FID для каждой задачи и наилучший sFID для всех задач, кроме одной. Благодаря улучшенной архитектуре авторы обладают современной генерацией изображений на LSUN и ImageNet 64×64.

Для ImageNet с более высоким разрешением наблюдается следующее: classifier guidance позволяет моделям существенно превосходить лучшие GAN’ы. Такие модели обеспечивают качество генерации, аналогичное GAN’ам, и сохраняют более высокий охват распределения (diversity), измеренный по recall. Они могут это осуществлять всего за 25 шагов диффузии.

Полезные ссылки

1. Оригинальная статья
2. Гитхаб-репозиторий авторов
3. Мод коллапс GAN’ов
4. Подбор гиперпараметров для обучения GAN’ов (и про BigGAN блок)
5. SotA модели генерации, основанные на правдоподобии [на момент выхода статьи]: первая, вторая, третья
6. Статья Denoising Diffusion Probabilistic Models (DDPM)
7. Виды шумов (кроме гауссова)
8. Статья с математикой диффузионных моделей (в частности, про reparameterization trick)
9. Evidence lower bound (ELBO)
10. Статья про Denoising Diffusion Implicit Models (DDIM)
11. Дивергенция Кульбака-Лейблера
12. Mean Squared Error
13. Реализация DDPM
14. Статья про classifier-free guidance
15. Статья про датасет LSUN