Введение в диффузионные модели

Введение в диффузионные модели

Предыстория

За последние пару лет мир генеративных моделей совершил огромный скачок благодаря диффузионным моделям.

Они успешно применяются в различных областях. Например, при генерации:

• изображений высокого качества

• видео

• музыки и даже музыки со словами

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

Что такое диффузионные модели и откуда появилась идея этой прорывной технологии? Давайте разбираться 🙂

В этом обзоре мы познакомимся с понятием диффузионных моделей, рассмотрим его с точки зрения score matching’а, вариационных автокодировщиков и стохастических дифференциальных уравнений.

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

Optimal transport problem 🏎

Transport Mapping Problem

Задача оптимального транспорта имеет непосредственное отношение к генеративным моделям. Как перевести одно распределение в другое, еще и сделать это оптимальным способом? И что значит оптимальным? Давайте обсудим эту задачу подробнее 🙂

Есть два эмпирических распределения \( \pi_0 \) и \( \pi_1 \in \mathbb{R}^d \). Необходимо построить такую транспортировку (транспортную карту) \( T \): \( \mathbb{R}^d \rightarrow \mathbb{R}^d \), что

Поскольку возможно построить бесконечно много таких транспортных карт $\mathrm{\backslash (\; T\; \backslash )}$ , обычно пытаются решить задачу оптимального транспорта через поиск $\mathrm{\backslash (\; T\; \backslash )}$ , которое минимизирует транспортную стоимость:

Интуицию данной задачи любят объяснять через землекопа, которому надо одну кучу земли превратить в другу. Перетаскивание земли занимает какое-то количество сил землекопа и поэтому он хочет потратить как можно меньше усилий для этого. Две кучки земли — два распределения, транспортная карта — наш землекоп, а стоимость — это усилия землекопа.

Тем не менее решение задачи оптимального транспорта довольно сложное, особенно для многомерных задач. Также важным вопросом является выбор транспортной стоимости \( c \).

В случае генеративных моделей \( T \) чаще всего параметризуется нейронной сетью и обучается в стиле алгоритмов GAN или максимизацией правдоподобия (MLE) (например, так обучаются вариационные автокодировщики VAE). Но GAN страдают от мод коллапса, а VAE дает слишком нечеткие и некачественные результаты.

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

Такую популярность диффузионные модели получили, в первую очередь, из-за качества и разнообразия генерации. Однако качественные результаты получаются при генерации за большое количество шагов и требуют много времени по сравнению с другими генеративными моделями (например, GAN).

Отсюда мы получаем Generative Learning Trilemma: ни одна генеративная модель не удовлетворяет сразу 3-м пунктам — разнообразию, качеству и скорости.

Диффузия 🧪

Диффузионные модели можно рассматривать с разных сторон: как модели со скрытыми латентными переменными (идейно близко к VAE), с позиции score matching’а или стохастических дифференциальных уравнений (SDE). В этой части мы обсудим все три подхода. Но для начала давайте посмотрим, как строится процесс диффузии.

Диффузионный процесс состоит из двух частей:

• forward process — постепенное добавление шума к распределению объектов до тех пор, пока итоговое распределение не станет неотличимым от случайного шума (чаще всего работают с гауссовским шумом, но недавно была представлена работа, показывающая, как можно работать и с другими типами шумов). Важный фактор — процесс марковский, и каждый следующий этап зашумления зависит только от предыдущего
• backward process — зеркальное отражение прямого процесса. Начав с чистого шума, модель постепенно удаляет шум и приближает данные к исходным чистым данным

Прямой процесс (forward process)

Пусть у нас есть сэмпл \( x_0 \sim q(x) \) из исходного распределения данных. Тогда прямой процесс диффузии — постепенное добавление небольшого количества шума, чаще всего Гауссовского (далее работаем с ним), в течение \( T \) шагов, в результате которых получаем последовательность зашумленных сэмплов \( x_1, x_2 …, x_T \) и в итоге — гауссовский шум с нулевым средним и единичной дисперсией. Количество добавляемого шума регулируется variance scheduler \( \{\beta_t \in (0,1)\}^T_t \). Он показывает, как много мы оставляем от исходной картинки и берём именно шума на каждом шаге.

Тогда зашумлённое распределение на шаге \( t \) можно записать следующим образом:

Помимо распределений нам бы хотелось иметь возможность смотреть или выполнять какие-то операции с элементами этих распределений, то есть сэмплировать \( x_t \).

Этого можно добиться с помощью трюка репараметризации:

Обратный процесс (backward process)

Если мы сможем обернуть данный процесс — получим генеративную модель, которая возьмет сэмпл из Гауссовского шума с нулевым средним и единичной дисперсией \( x_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \) и восстановит элемент из исходного распределения.

Важная особенность диффузионного процесса: если \( \beta_t \) — достаточно маленькое значение, то \( q(x_{t-1}|x_t) \) тоже будет Гауссовским. Однако мы не можем так просто оценить это распределение, поскольку для оценки нам нужно все распределение исходного датасета. Поэтому нашей задачей становится аппроксимация с помощью какой-либо модели \( p_{\theta} \):

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

Связь диффузионных моделей с VAE

А как же обучается такая модель? Для того, чтобы это понять, необходимо вспомнить, что такое модели со скрытыми переменными, а точнее вариационный автокодировщик.

В этом разделе мы рассмотрим диффузионные модели как модели со скрытыми переменными и подробно разберём их связь с VAE.

Но прежде чем переходить к диффузионным моделям, давайте вспомним, что такое и как обучался VAE.

VAE

Вариационный автокодировщик — генеративная модель из двух частей:

• энкодер кодирует входное изображение в латентное пространство;
• декодер декодирует вектор, засэмплированный из латентного пространства.

Таким образом, у нас есть:

• распределение \( p(x) \) — неизвестное нам распределение изображений;
• распределение \( p(z) \) — латентное распределение, которое мы априори считаем нормальным с нулевым средним и единичной дисперсией \( p(z) = \mathcal{N}(0,I) \);
• \( p(z|x): \) условное распределение — вероятность z при условии x. Оно будет аппроксимироваться нашим энкодером \( q(z|x, \phi) \) и мы будем считать его гауссовым с неизвестным средним и дисперсией.
• \( p(x|z) \) условное распределение — вероятность получения \( x \) из латентного \( z \). Его мы уже аппроксимируем с помощью нашего декодера \( p(x|z, \theta) \).

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

Давайте теперь разберемся, как обучается вариационный автокодировщик. Для этого рассмотрим задачу максимизации правдоподобия:

Логарифм правдоподобия теперь представляется через сумму двух слагаемых:

• KL-дивергенцию распределения латентных переменных и условного распределения латентных переменных при наших данных;
• добавления ELBO (evidence lower bound).

Отметим: KL-дивергенция у нас неотрицательная и принимает минимальное значение в нуле при совпадении практически везде \( q(z) \) и \( p(z|x) \).

То есть максимизируя ELBO, мы будем также максимизировать и логарифм правдоподобия, поскольку \( ELBO \leq \log p(x| \theta) \).

Рассмотрим оптимизируемую нами ELBO. Мы говорим про автокодировщик, значит, распределение \( q(z) = q(z|x, \phi) \). Это то, что получается при кодировании энкодером.

Тогда

Итак, у нас есть две составляющих: реконструкция и регуляризация.

• Реконструкционный лосс показывает нам, что будет, если мы возьмем какой-то элемент из нашей выборки \( x_i \), прогонем его через энкодер, получим распределение с параметрами \( \mu_i, \sigma_i \), засэмплируем из полученного распределения латент \( z_i \), а потом декодируем его обратно в \( \hat{x}_i \). Таким образом, эта величина представляет, насколько сильно результат будет отличаться от исходного сэмпла. Чем ближе получится значение, тем выше будет реконструкционное слагаемое
• Регуляризационный член отвечает за приближение нашим энкодером латентного распределения

В чем проблема автокодировщика, и почему при сэмплировании возникают замыленные (нечеткие) и недостаточно качественные результаты?

Реконструкционный лосс стремится свести всё просто к максимизации правдоподобия, в том время как регуляризация старается приблизить нормальное априорное латентное распределение и наше параметрически заданное. В рeзультате при обучении латентное распределение не покрывается полностью, поэтому у нас не получается достаточно точно сгенерировать сэмпл (изображение) из желаемого распределения.

Но при чём здесь диффузия?

Диффузия решает проблему недостаточно качественного покрытия, поскольку во время обучения любой элемент из train выборки переходит в одно и то же распределение, в случае Гауссовского шума — в нормальное. Поэтому при инференсе мы сэмплируем из нормального распределения и получаем хорошие изображения.

Обучение

Как обучить такую модель? Выше мы упоминали, что диффузионную модель можно рассматривать как модель со скрытыми переменными [см. рисунок ниже]:

\( x_0 \) — observed varaibles

\( x_1, … x_T \) — hidden varaibles

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

В качестве обозначений введем \( \mu_{\theta} \) — непосредственно среднее искомого распределения, \( \epsilon_{\theta} \) — предсказание шума (\( \epsilon \)-параметризация).

Итоговую функцию потерь для \( \epsilon \)-параметризации можно записать следующим образом:

\( \mathbb{E}_{x_0, \epsilon}\Big[\frac{1 — \alpha_t)^2}{2\alpha_t(1 — \bar{\alpha}t)\|\Sigma{\theta}\|^2_2}\|\epsilon_t — \epsilon_{\theta}(x_t, t)\|^2\Big] \) где \( \epsilon_t \) — добавленный на шаге t шум; \( \epsilon_{\theta}(x_t, t) \approx \nabla_{x_t} \log q(x_t) \) — предсказанный шум, \( \bar{\alpha}t = \prod{i=1}^t \alpha_i = \prod_{i=1}^t (1 — \beta_i) \)

(ниже для самых смелых мы приводим доказательство)

Заметим, что по сути мы смотрим взвешенную сумму в зависимости от шага \( t \) разности предсказанного и реального шума. Время при этом сэмплируется из равномерного распределения для обучения.

Формула функции потерь для диффузионного процесса [любителям математики]

Как и VAE, будем оптимизировать Variationsl Lower Bound:

\( \log p(obs| \theta) \geq\int q(hidd|obs, \phi) \log\frac{p(obs, hidd| \theta)}{q(hidd|obs, \phi)}d \text{} hidd \)

\( log p(x_0| \theta) \geq \int q(x_1, x_2, … x_T| x_0) \log \frac{p(x_0, x_1, … x_T| \theta)}{q(x_1, … x_T| x_0)}dx_1…dx_T \) =

=\( \int q(x_1, x_2, …, x_T|x_0) \log p(x_0|x_1)dx_1dx_2…dx_T \) +

\( \int q(x_1, …, x_T|x_0) \log \) \( \frac{p(x_1|x_2,\theta)p(x_2|x_3,\theta) … p(x_{T-1}|x_T,\theta)p(x_T|\theta)}{q(x_1|x_0)q(x_2|x_1, x_0)…q(x_T|x_{T-1}, x_0)}dx_1dx_2…dx_T \) =

= \( \int q(x_1|x_0) \log p(x_0|x_1)dx_1 \underbrace{\int q(x_2|x_0)dx_2}_{=0} …\underbrace{\int q(x_T|x_0)dx_T}_{=0} \) +

\( \sum_{t=2}^T \int q(x_t|x_0) q(x_{t-1}|x_t,x_0) \log \frac{p(x_{t-1}|x_t, \theta)}{q(x_{t-1}|x_t,x_0)}dx_{t-1}dx_t \) +

\( \int q(x_T|x_0) \log \frac{p(x_T)}{q(x_T|x_0)}dx_T = \int q(x_1|x_0) \log p(x_0|x_1)dx_1 \) —

\( \sum_{t=2}^T \int q(x_t|x_0) D_{KL}(q(x_{t-1}|x_t,x_0)|p(x_{t-1}|x_t, \theta)) \) —

\( D_{KL}(q(x_T|x_0)|p(x_T) \rightarrow max_{\theta} \)

Рассмотрим каждое слагаемое по отдельности:

• \( D_{KL}(q(x_T|x_0)|p(x_T) = 0 \), так как считаем, что конечное распределение для любого \( x_0 \) является нормальным с нулевым средним и единичной дисперсией;
• \( \int q(x_1|x_0) \log p(x_0|x_1)dx_1 \approx const \)
• остается суммирование KL-дивергенций, где для t-го уровня зашумления \( \mathcal{L}_t = D_{KL}(q(x_t|x_{t+1},x_0) \| p_{\theta}(x_t, x_{t+1})) \).

Поскольку оба распределения в KL-дивергенции являются нормальными, мы можем выразить ее через средние и дисперсии этих распределений:

\( D_{KL}(\mathcal{N}_0 \| \mathcal{N}_1) = \frac{1}{2}(tr(\Sigma_1^{-1}\Sigma_0) + (\mu_1 — \mu_0)^T\Sigma_1^{-1}(\mu_1 — \mu_0) — k + \ln (\frac{det\Sigma_1}{det \Sigma_0}) \)

Для начала разберемся с \( q(x_t|x_t, x_0) \). Воспользуемся формулой Байеса:

\( q(x_{t-1}|x_t,x_0) = q(x_t|x_{t-1}, x_0) \frac{q(x_{t-1}|x_0)}{q(x_t|x_0)} \).

А также тем фактом, что \( q(x_t|x_{t-1}) = \mathcal{N}(x_t| \sqrt{1-\beta_t}x_{t-1}, \beta_t \mathbf{I}) \)

Тогда можно получить, что \( q(x_{t-1}|x_t,x_0) \) — нормальное распределение со средним

\( \tilde{\mu}_t(x_t, x_0) = \frac{\sqrt{\alpha}_t(1 — \bar{\alpha})}{1 — \bar{\alpha}_t}x_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 — \bar{\alpha}_t}x_0 = \frac{1}{\alpha_t}(x_t — \frac{1 — \alpha_t}{\sqrt{1 — \bar{\alpha}_t}}\epsilon_t) \)

Тогда для \( p_{\theta}(x_t, x_{t+1}) \mu_{\theta}(x_t, t) = \frac{1}{\sqrt{\alpha_t}}\Big(x_t — \frac{1 — \alpha_t}{\sqrt{1 — \hat{\alpha_t}}}\epsilon_{\theta}(x_t,t)\Big)​ \)

И сэмпл из распределения \( x_{t-1} = \mathcal{N}\Big(x_{t-1};\frac{1}{\sqrt{\alpha_t}}(x_t — \frac{1 — \alpha_t}{\sqrt{1 — \hat{\alpha_t}}}\epsilon_{\theta}(x_t,t)), \Sigma_{\theta}(x_t,t)\Big)​ \)

В данном приближении будем считать, что дисперсия распределений одинаковая (?????), тогда:

\( \mathcal{L}_t = \mathbb{E}_{x_0, \epsilon}\Big[\frac{1}{2\| \Sigma_{\theta}(x_t, t)\|^2_2}\|\tilde{\mu}_t(x_t, x_0) — \mu_{\theta}(x_t, t)\|^2\Big] \) =

\( \mathbb{E}_{x_0, \epsilon}\Big[\frac{1}{2\| \Sigma_{\theta}(x_t, t)\|^2_2}\|\frac{1}{\sqrt{\alpha_t}}(x_t — \frac{1 -\alpha_t}{\sqrt{1-\bar{\alpha}}}\epsilon_t) — \frac{1}{\sqrt{\alpha_t}}(x_t — \frac{1 — \alpha_t}{\sqrt{1 — \bar{\alpha}_t}}\epsilon_{\theta}(x_t,t))\|^2\Big] \) =

= \( \mathbb{E}_{x_0, \epsilon}\Big[\frac{1 — \alpha_t)^2}{2\alpha_t(1 — \bar{\alpha}_t)\|\Sigma_{\theta}\|^2_2}\|\epsilon_t — \epsilon_{\theta}(x_t, t)\|^2\Big] \)

Таким образом, всю схему обучения можно представить так:

Однако эмпирически было показано, что упрощенная функция потерь для обучения работает лучше, чем полная:

Тогда алгоритм обучения и сэмплирования можно представить следующим образом:

Обучение

1. Сэмплируем изображение из нашего обучающего датасета
2. Сэмплируем время из равномерного распределения
3. Сэмплируем шум из нормального распределения
4. Зашумляем сэмпл из датасета до шага t, используя шум
5. С помощью нейронной сети делаем один шаг расшумления и считаем функцию потерь
6. Делаем шаг градиентного спуска

Инференс

1. Сэмплируем латент из нормального распределения
2. Делаем t шагов расшумления

Таким образом, мы снова пришли к DDPM, но только с точки зрения модели со скрытыми латентными переменными.

Давайте теперь посмотрим на примеры генерации:

Score-based models → от NCSN к DDPM

Мы обсудили с вами связь диффузионных моделей и VAE, а в этом разделе мы обсудим связь диффузионных и score-based моделей. Идейно диффузию в целом можно рассматривать как улучшение имеющихся алгоритмов score-matching’а.

Score-based models

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

Поэтому мы поговорим о другом типе моделей — score-based моделях и score matching’е, которые также можно использовать для генерации сэмплов из желаемого распределения.

Для начала поймем, что такое Energy-based модели.

Наша задача — моделировать плотность распределения изображений и сэмплировать оттуда объекты. Одним из удобных способов моделирования распределения являются Energy-based модели.

Они задаются потенциальной скалярной функцией энергии \( f_{\theta}(x) \). Плотность их распределения: \( p_{\theta}(x) = \frac{e^{-f_{\theta}(x)}}{Z_{\theta}} \), где \( Z_{\theta} = \int e^{-f_{\theta}(x)}dx \) — нормировочная константа. В этом случае высоковероятные сэмплы соответствуют маленьким значениям энергии, а маловероятные — высоким (получается некий потенциальный барьер).

Такие модели более устойчивы для обучения и имеют ряд преимуществ (подробнее можно почитать в статье). Но что будет, если мы попробуем оценить MLE этой модели?

И вот с логарифмом этой нормировочной константы непонятно, что делать 😟

Поэтому необходимо моделировать не саму функцию распределения \( p(x) \), а ее score функцию: \( \nabla_x \log p(x) \).

Следовательно, в случае energy-based модели мы получим простое выражение, избавившись от нормализующей константы:

Теперь мы можем обучать модель, минимизируя следующую ошибку:

Простой пример вычисления score функции для одной одномерной гауссианы

\( p(x) = \frac{1}{\sqrt{2 \pi \sigma^2}} e^{-\frac{(x-\mu)^2}{2\sigma^2}} \)

import torch
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

sns.set_style("darkgrid")


def normal_distr(x, mu_1, sigma_1,):
  '''
  Одна гауссиана:
  args:
  mu_1: среднее нормального распределения
  sigma_1: дисперсия нормального распределения
  
  '''
  normal = 1 / ((2*torch.pi)**(1/2) * sigma_1) * (-(x - mu_1)**2 / (2*sigma_1**2)).exp()
  
  return normal
  
  
x = np.linspace(-6, 6, 10000)

mu_1 = 2
sigma_1 = 0.8

plt.plot(x, normal_distr(torch.from_numpy(x), mu_1, sigma_1))
plt.plot()

Тогда:

\( s(x) = \nabla_x \log p(x) = \nabla_x (\log \frac{1}{\sqrt{2 \pi \sigma^2}} — \frac{(x-\mu)^2}{2\sigma^2}) = -\frac{(x-\mu)}{\sigma^2} \)

Давайте теперь напишем код для теоретического и численного расчета score функции:

def numeric_score(x,  mu_1, sigma_1):
  x = torch.tensor(x, requires_grad=True)
  out = normal_distr(x,  mu_1, sigma_1).log()
  out.backward(torch.ones_like(x))
  return (x.grad).numpy()

def theory_score(x,  mu_1, sigma_1):
  return -(x - mu_1) / sigma_1**2
  

x = np.linspace(-6, 6, 10000)

mu_1 = 2
sigma_1 = 0.8

numeric_score_res = numeric_score(x,  mu_1, sigma_1)
theory_score_res = theory_score(x,  mu_1, sigma_1)

fig, axs = plt.subplots(1, 3, figsize=(32,8))
axs[0].plot(x, numeric_score_res)
axs[0].set_title('Numeric score function')
axs[1].plot(x, theory_score_res)
axs[1].set_title('Theoretic score function')
axs[2].scatter(x, np.abs(theory_score_res -  numeric_score_res))
axs[2].set_title('Difference')

plt.show()

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

Немного интуиции

1. Величина векторного поля (score функции) принимает большие значения там, где изменения \( \log p(x) \) максимальны, значит, в окрестностях моды распределения score функция будет довольно мала.
2. В физике понятие “score функции” аналогично понятию “дрифт”: как диффузионные частички должны двигаться в сторону минимального энергетического состояния.

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

Динамика Ланжевена

Итак, мы знаем score функцию. Но у нас нет самого распределения. Как же тогда сэмплировать из него объекты?

Здесь на помощь приходит динамика Ланжевена — зная score функцию и начальное приближение, мы можем получить сэмпл из искомого распределения:

где \( z_i \sim \mathcal{N}(0,\mathbf{I}), i = 0, 1,…, K \)

Значит, у нас есть итеративная процедура, которая при стремлении размера шага \( \epsilon \) к нулю и увеличении до бесконечности количества шагов \( K \rightarrow \infty \) позволяет сойтись к искомому распределению.

Интересный факт: без слагаемого с шумом динамика Ланжевена напоминает уже давно знакомый градиентный спуск 🙂

Пример с двумя одномерными гауссианами и динамикой Ланжевена [для любителей питона]

Итак, рассмотрим соединение двух гауссиан:

import torch
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

sns.set_style("darkgrid")


def two_normal_distr(x, mu_1, sigma_1, a_1, mu_2, sigma_2, a_2):
	'''
	Функция, на выходе которой мы получаем смесь двух гауссиан
	'''
  normal_1 = a_1 * 1 / ((2*torch.pi)**(1/2) * sigma_1) * (-(x - mu_1)**2 / (2*sigma_1**2)).exp()
  normal_2 = a_2 * 1 / ((2*torch.pi)**(1/2) * sigma_2) * (-(x - mu_2)**2 / (2*sigma_2**2)).exp()
  return normal_1 + normal_2

# Визуализируем наши гауссианы  
x = np.linspace(-6, 6, 10000)

mu_1 = 3
sigma_1 = 0.8
a_1 = 0.8

mu_2 = -2.5
sigma_2 = 0.1
a_2 = 0.2

plt.plot(x, two_normal_distr(torch.from_numpy(x), mu_1, sigma_1, a_1, mu_2, sigma_2, a_2))
plt.show()

Получаем следующее:

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

import math

def langeven(x, step=0.05, use_noise=True):
  x = torch.tensor([x], requires_grad=True)
  out = two_normal_distr(x,  mu_1, sigma_1, a_1, mu_2, sigma_2, a_2).log()
  out.backward(torch.ones_like(x))
  x = x[0] + step * x.grad
  if use_noise:
    x += math.sqrt(2*step) * torch.randn(1)
  return x

И запустим на T=300 шагов.

• Без слагаемого с шумом:

from tqdm import tqdm

x = 0.
n_steps = 300
all_x = []
all_y = []

for i in tqdm(range(n_steps)):
  x = langeven(x, step=0.05, use_noise=False)
  y = two_normal_distr(x, mu_1, sigma_1, a_1, mu_2, sigma_2, a_2)
  all_x.append(x.item())
  all_y.append(y.item())
  
  
x = np.linspace(-6,
                6, 10000)
plt.plot(x, two_normal_distr(torch.from_numpy(x), mu_1, sigma_1, a_1, mu_2, sigma_2, a_2))
plt.scatter(all_x, all_y, c='r')
plt.plot()

• Со слагаемым с шумом:

from tqdm import tqdm

x = 0.
n_steps = 300
all_x = []
all_y = []

for i in tqdm(range(n_steps)):
  x = langeven(x, step=0.05, use_noise=True)
  y = two_normal_distr(x, mu_1, sigma_1, a_1, mu_2, sigma_2, a_2)
  all_x.append(x.item())
  all_y.append(y.item())

x = np.linspace(-6,
                6, 10000)
plt.plot(x, two_normal_distr(torch.from_numpy(x), mu_1, sigma_1, a_1, mu_2, sigma_2, a_2))
plt.scatter(all_x, all_y, c='r')
plt.plot()

Отметим: без добавления шума мы просто сходимся к моде распределения. Но добавив слагаемое со стохастикой, мы уже неплохо описываем окрестность моды.

Но у нас остается проблема такого обучения: откуда мы берем \( \nabla \log p(x) \)?

Тут на помощь приходят несколько возможных методов:

• sliced score mathcing;
• denoising score matching.

В нашем обзоре мы рассмотрим второй метод.

Denoising Score matching

Давайте добавим немного нормального шума в наши данные: \( \hat{x} = x + \sigma \epsilon \)

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

Оно несильно отличается от исходного, а значит, несильно отличается и его score функция.

Доказательство для любителей математики

\( \mathbb{E}_{q{_\sigma}(\hat{x}|x)p_{data}(x)}\Big[\|s_{\theta}(\hat{x}) — \nabla_{\hat{x}}\log q_{\sigma}(\hat{x}|x)\|^2_2\Big] = \mathbb{E}_{p_{\sigma}(\hat{x},x)}\Big[(\|s_{\theta}(\hat{x})\|^2 \) — \( 2<s_{\theta}(x)\nabla_{\hat{x}}\log q_{\sigma}(\hat{x}|x)> + \underbrace{\|\nabla_{\hat{x}}\log q_{\sigma}(\hat{x}|x)\|^2}_{const}\Big] \)

Заметим, что \( \mathbb{E}{p_{\sigma}(\hat{x},x)}(\|s_{\theta}(\hat{x})\|^2) = \mathbb{E}{p_{\sigma}(\hat{x})}(\|s_{\theta}(\hat{x})\|^2) + const \), ведь score функция не зависит от x. Она будет давать константный вклад в оптимизацию.

\( \mathbb{E}_{p_{\sigma}(\hat{x})}\Big[\|s_{\theta}(\hat{x}) — \nabla_{\hat{x}} \log p_{\sigma}(\hat{x})\|^2_2\Big] = \mathbb{E}_{p_{\sigma}(\hat{x})}\Big[\|s_{\theta}(\hat{x})\|^2 \) — 2\( <\( s_{\theta}(\hat{x}), \nabla_{\hat{x}} \log p_{\sigma}(\hat{x})> + const\Big]​ \)

Здесь первые слагаемые одинаковые. Осталось доказать равенство матожиданий скалярных произведений.

\( \mathbb{E}_{p_{\sigma}(\hat{x})}\Big[2<s_{\theta}(\hat{x}), \nabla_x \log p_{\sigma}(\hat{x})>\Big] =2 \int p_{\sigma}(\hat{x}) \)< \( s_{\theta}(\hat{x}), \frac{\nabla_{\hat{x}} p_{\sigma}(\hat{x})}{p_{\sigma}(\hat{x})}>d\hat{x} = 2 \int <s_{\theta}(\hat{x}), \nabla_{\hat{x}} p_{\sigma}(\hat{x})>d \hat{x} = 2\int \) <\( s_{\theta}(\hat{x}), \nabla_{\hat{x}}\int q_{\sigma}(\hat{x}|x)p_{data}(x)dx>d\hat{x} = \int_x \int_{\hat{x}} \)<\( s_{\theta}(\hat{x}), \int \nabla_{\hat{x}}q_{\sigma}(\hat{x}|x)p_{data}(x)dx>d\hat{x} = \int_x \int_{\hat{x}}p_{data}(x) \)<\( s_{\theta}(\hat{x}), \nabla_{\hat{x}}q_{\sigma}(\hat{x}|x)>d\hat{x} dx = \int_x \int_{\hat{x}}p_{data}(x) q_{\sigma}(\hat{x}|x) \)<\( s_{\theta}(\hat{x}), \nabla_{\hat{x}} \log q_{\sigma}(\hat{x}|x)>d\hat{x} dx = 2\mathbb{E}_{p_{\sigma}(\hat{x},x)} \)< \( s_{\theta}(x),\nabla_{\hat{x}}\log q_{\sigma}(\hat{x}|x)>​ \)

Таким образом, мы доказали эквивалентность двух задач оптимизации.

\( \nabla_{\hat{x}}\log q_{\sigma}(\hat{x}|x) \), в свою очередь, выразить несложно:

\( -\nabla_{\hat{x}}\log q_{\sigma}(\hat{x}|x) = \frac{(x + \sigma \epsilon — x)}{\sigma^2} = \frac{\epsilon}{\sigma} \)

Следовательно, чтобы обучать score функцию зашумленных данных, нам уже не нужно распределение реальных данных. Мы обучаем \( s_{\theta}(\tilde{x}) \), поскольку знаем, что \( q(\tilde{x}|x) = \mathcal{N}(x, \sigma^2 \mathbf{I}) \). И затем уже применяем динамику Ланжевена.

Но какие здесь могут быть проблемы?

• в многомерном пространстве обычно есть много мод, и при сэмплировании велика вероятность, что мы зайдем в одну моду, из которой уже не выйдем;
• проблема выбора \( \sigma \) для зашумления.

Score-based models c различным уровнем шума (NCSN)

Решение перечисленных проблем — взять несколько уровней шума и помаленьку увеличивать степень зашумления.

Теперь мы можем рассмотреть взвешенную сумму для разных уровней зашумления:

Получается, мы хотим выучить score функцию \( s_{\theta}(x.i) \) для каждого из уровней шума. А что с этим делать дальше?

Динамика Ланжевена запускается следующим образом:

• инициализируем \( x_0 \);
• для каждого уровня шума (начиная с самого большого) делаем \( T \) шагов динамики Ланжевена, используя значения \( s_{\theta} \) на этом уровне и размер шага \( \alpha_i \), который уменьшается при переходе на менее шумный уровень;
• конечное предсказание — инициализация для следующего уровня зашумления.

Ниже приведен алгоритм в качестве примера:

Пример работы сэмплирования:

Такая модель работает очень даже неплохо:

Описанный выше метод называется  Noise Conditional Score Network (NCSN). Авторы предлагали брать 10 различных уровней зашумления и делать на каждом из них по 100 шагов.

DDPM

Вот мы и подобрались к известному диффузионному процессу, который используется сегодня практически во всех диффузионных моделях. Сейчас мы сформулируем основное отличие denosing diffusion probabilistic model (DDPM) модели от NCSN. Авторы DDPM предлагают делать всего один шаг динамики Ланжевена, при этом брать гораздо больше уровней зашумления (1000). Это отличие привело нас к основной идее диффузионного процесса, описанного выше.

Подробнее про DDPM и все улучшения, предложенные в данной работе, можно почитать в этом посте.

Диффузия и SDE

При переходе от дискретных шагов по времени к непрерывным можно связать диффузионные модели со стохастическими дифференциальными уравнениями.

Зачем нам вообще такой взгляд на диффузию?

Оказывается, именно благодаря связи диффузии со SDE и ODE можно использовать все те различные сэмплеры для генерации.

В таком случае мы работаем с непрерывным диффузионным процессом \( \{x(t)\}^T_{t=0} \) , \( x(0) \sim p_0 \) и \( x(T) \sim p_T \). Тогда (без подробностей вывода) можно записать прямое и обратное SDE для процесса, характеризующегося функцией \( f(x,t) \) и виннеровским процессом \( w \).

Прямой процесс (forward SDE): данные → шум

У этого процесса следующий физический смысл. \( f(x,t) \) показывает, как элементы (частицы) будут двигаться в замкнутой системе при отсутствии каких-либо случайных эффектов. Например, мы уже говорили про градиентный спуск, который может быть частным случаем подобного движения. А \( g(t) \) характеризует амплитуду случайного блуждания.

Ниже показан пример постепенного зашумления данных в нормальное распределение (слева: изображение, справа: одномерное распределение):

Обратный процесс (reverse SDE): шум → данные

где \( \nabla \log p_t(x) \) — score функция \( p_t(x) \).

Ниже представлен обратный процесс — генерация данных из шума путем решения обратного SDE.

То есть прямой и обратный процесс выглядят следующим образом:

Чтобы решить обратное SDE, необходимо знать распределение, с которого мы стартуем \( p_T(x) \), а также score функцию \( \nabla_x \log p_t(x) \). Из формулировки диффузионного процесса мы можем взять \( p_T(x) = \mathcal{N}(0, I) \), а score функцию аппроксимировать с помощью нейронной сети \( s_{\theta}(x,t) \approx \nabla_x \log p_t(x) \) и обучить, минимизируя функцию потерь:

где \( U(0,T) \) — равномерное распределение по времени, \( \lambda(t) \) — неотрицательная весовая функция.

После того, как мы обучили score функцию, мы можем сгенерировать элемент, решая SDE и стартуя из \( x(T) \):

Как решать SDE?

Для решения обратного SDE есть специальные солверы. Один из самых простых методов — Euler–Maruyama.

Итеративная процедура решения выглядит следующим образом:

Есть и другие методы, например:

• Milstein method;
• Stochastic Runge-Kutta methods;
• Adaptive step-size SDE solvers.

А откуда берется ODE?

Представленному стохастическому дифференциальному уравнению можно сопоставить обыкновенное дифференциальное уравнение. Его решения, сэмплированные в точках t траектории, совпадают с решениями SDE (то есть совпадают маргинальные распределения \( \{p_t(x)\}_{t \in [0,T]} \)):

Из-за существования соответствующего ODE возникает возможность использовать различного рода сэмплирование для ускорения генерации диффузии, поскольку есть большое количество способов решения ODE.

Решение обратного ODE

Как мы уже заметили выше, для SDE можно поставить в соответствие ODE, для решения которого используются различные солверы. Некоторые из них мы и рассмотрим сейчас 🙂

Рассмотрим формулировку \( \frac{dx}{dt} = f(x, t) \).

Forward Euler Method

Один из самых простых методов: \( x_{\delta + t} = x_t + \delta f(x, t) \).

Runge-Kutta Method

Более аккуратный метод, позволяющий добиться лучшего качества решения, однако требующий больше вычислений. Это метод Рунге-Кутты 4-го порядка:

\( \begin{cases} k_1 = f(x_t, t), \\ k_2 = f(x_t + \frac{\delta}{2}k_1, t + \frac{\delta}{2}) \\ k_3 = f(x_t+ \frac{\delta}{2}k_2, t + \frac{\delta}{2}) \\ k_4 = f(x_t + \delta k_3, t + \delta) \\ x_{t+ \delta} = x_t + \frac{\delta}{6}(k_1 + 2k_2 + 2k_3 + k_4) \end{cases} \)

где \( \delta \) — просто шаг.

Хотя RK4 — самый популярный метод Рунге-Кутты, существуют методы более высокого порядка (например, пятого и шестого), которые обеспечивают большую точность за счет дополнительных вычислений функции.

VE и VP diffusion

И последнее, что мы сегодня обсудим — это Variance Exploding (VE) и Variance Preserving (VP) модели.

Эти методы влияют на процесс добавления шума к данным, который, в свою очередь, определяет, как обучаются и работают модели.

Variance Exploding (VE)

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

Тогда в случае \( N \rightarrow \infty \) можно записать \( dx = \sqrt{\frac{d[\sigma^2(t)]}{dt}}dw \)

Variance Preserving (VP)

Тут суммарная дисперсия данных и шума постоянна на протяжении всего процесса диффузии. Это позволяет контролировать влияние шума и поддерживать стабильное распределение данных. Модель DDPM — пример VP диффузии.

В случае DDPM пертурбации данных можно записать марковскую цепочку так:

В случае \( N \rightarrow \infty \) выражение выше сходится к \( dx = \underbrace{-\frac{1}{2}\beta(t)x}{f(x,t)}dt + \underbrace{\sqrt{\beta(t)}}{g(t)}dw \).

Доказательство

\( x(t + \Delta t) = \sqrt{1 — \beta(t + \Delta t)\Delta t}x(t) + \sqrt{\beta(t + \Delta t)\Delta t}z(t) \)

\( x(t + \Delta t) \approx (1 — \frac{1}{2}\beta(t + \Delta t)\Delta t)x(t) + \sqrt{\beta(t + \Delta t)\Delta t}z(t) \)

При \( \Delta t \rightarrow 0 \)

\( dx = \underbrace{-\frac{1}{2}\beta(t)x}_{f(x,t)}dt + \underbrace{\sqrt{\beta(t)}}_{g(t)}dw \)

Тогда для обратного SDE:

\( dx = -\beta(t)[\frac{x}{2} + \nabla_x \log p_t(x)]dt + \sqrt{\beta(t)}dw \)

Summary

Таким образом, при переходе к непрерывным шагам по времени диффузионный процесс (как прямой, так и обратный) можно представить через SDE. Более того, SDE сопоставимо с ODE. Оно сохраняет маргинальные распределения на каждом шаге, что позволяет перейти к быстрому сэмплированию. Для решения ODE можно применять различные солверы.

Заключение

В этой статье мы рассмотрели диффузионные модели с разных точек зрения, включая стохастические дифференциальные уравнения (SDE), score matching и вариационные автоэнкодеры (VAE).

Каждая концепция играет важна для понимания и применения диффузионных моделей. Надеемся, что вам стало немного понятней, какая математика стоит за столь мощным инструментом для генерации различных объектов в большом количестве задач 😊

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

• Лекции Дмитрия Ветрова:
  • Introduction part 1 (диффузионные модели и связь с VAE)
  • Introdiction part 2 (диффузионные модели и связь с SDE)
• Лекции BayesGroup:
  • Diffusion Models
• Лекции AIRI:
  • Введение в диффузионные модели, part 1
  • Введение в диффузионные модели, part 2
  • Введение в диффузионные модели, part 3
• Blog post 1
• Blog post 2