Введение
В этом параграфе мы снова попробуем решить задачу генерации, когда нам дана выборка объектов из распределения $x_0 \sim q(x)$, и хотим научиться генерировать новые объекты из распределения , которых нет в нашей выборке.
Вероятно, вы уже знакомы с другими генеративными моделями, например VAE или GAN-ы. Здесь же мы познакомим вас с еще одним видом генеративных моделей: диффузионные модели, которые стали крайне популярны в последнее время благодаря своему высокому качеству генерации объектов из заданного распределения. В общий чертах, они работают следующим образом: берем шум из $\mathcal{N}(0, I)$ и шаг за шагом удаляем компоненты шума до тех пор, пока не получим объект $x_0$ из распределения, см. иллюстрацию ниже.
Более детально
Для детального понимания стоит объяснить, что такое прямой и обратный диффузионные процессы. Прямой процесс заключается в постепенном зашумлении картинки с помощью распределения $q$, а обратный, наоборот, в расшумлении с помощью распределения $p$. Их можно схематично изобразить следующим образом:
Прямой диффузионный процесс определяется как апостериорное распределение $q(x_{1:T}|x_0)$. Это распределение также является Марковской цепочкой, которая постепенно добавляет гауссовский шум к объекту $x_0$. На каждом шаге шум добавляется с различной магнитудой, которая определяется расписанием дисперсий ${\beta_1, ... , \beta_T}$. При правильном выборе расписания в пределе по числу шагов $T$ мы должны сойтись к шуму из $\mathcal{N}(0, I)$. В качестве распределений $q$ берут нормальные распределения:
$$ q(x_t | x_{t - 1}) := \mathcal{N}(x_t; \sqrt{1 - \beta_t}x_{t - 1}, \beta_tI), \ \ \ \ \ \ \ q(x_{1:T}|x_0) = \prod_{t = 1}^T q(x_t | x_{t - 1}) $$
Теперь перейдем к обратному процессу и к самой диффузионной модели.
Диффузионная модель - это вероятностная модель с латентными переменными вида $p_\theta(x_0) := \int p_\theta(x_{0:T}) dx_{1:T}$, где промежуточные состояния $x_1, ..., x_T$ соответствуют зашумленным объектам, a $x_0$ - объект из распределения. Совместное распределение $p_\theta(x_{0:T})$ называет обратным диффузионным процессом, который представляет собой Марковскую цепочку из гауссовских распределений $p_\theta(x_{i-1}|x_{i})$:
$$ p(x_{0:T}) = p(x_0) \prod_{t = 1}^Tp_{\theta}(x_{t-1}|x_t) \ \ \ \ \ \ \ \ \ p_\theta(x_{T})=\mathcal{N}(x_T | 0, I) $$
$$ p_{\theta}(x_{t - 1}|x_t):= \mathcal{N}(x_{t - 1}; \mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t)) $$
Таким образом, обратный процесс параметризуется моделью $\theta$, которая по зашумленному объекту $x_t$ и шагу $t$ предсказывает среднее $\mu_{\theta}(x_t, t)$ и дисперсию $\Sigma_{\theta}(x_t, t)$.
Обучение диффузионной модели
Диффузионный модели обучаются, максимизируя вариационную нижнюю оценку (ELBO) логарифма правдоподобия $\log p_{\theta}(x_0)$. По тому же принципу обучаются VAE, с тем лишь отличием, что у диффузионных моделей другая форма модели с латентными переменными. Итак, давайте выведем ELBO для диффузии:
Комментарий
Если вы знакомы с VAE, то вывод $L_{VLB}$ должен быть вам понятен, однако ниже приведен вывод с помощью неравенства Йенсена
$$ \begin{aligned} L_\text{CE} &= - \mathbb{E}_{q(\mathbf{x}_0)} \log p_\theta(\mathbf{x}_0) \\ &= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T} \Big) \\ &= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \int q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} d\mathbf{x}_{1:T} \Big) \\ &= - \mathbb{E}_{q(\mathbf{x}_0)} \log \Big( \mathbb{E}_{q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)} \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \Big) \\ &\leq - \mathbb{E}_{q(\mathbf{x}_{0:T})} \log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})} \\ &= \mathbb{E}_{q(\mathbf{x}_{0:T})}\Big[\log \frac{q(\mathbf{x}_{1:T} \vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})} \Big] = L_\text{VLB} \end{aligned} $$
Теперь вернемся к распределению $q(x_t | x_{t - 1})$. Для того чтобы получить $x_t$, придется итеративно получать $x_1, ..., x_{t - 1}$. Однако это можно сделать более эффективно благодаря нормальным распределениям. Для этого обозначим $\alpha_t := 1- \beta_t$ и $\bar{\alpha}_t:= \prod_{i = 1}^t\alpha_i$, тогда
$$ q(x_t | x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)I) $$
Формальный вывод этого факта
$$ \begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\mathbf{z}_{t-1}; \text{ где } \mathbf{z}_{t-1}, \mathbf{z}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\ &= \sqrt{\alpha_t}(\sqrt{\alpha}_{t - 1}x_{t - 2} + \sqrt{1 - \alpha_{t - 1}}\mathbf{z}_{t - 2}) + \sqrt{1 - \alpha_t} \mathbf{z}_{t-1} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\mathbf{z}}_{t-2}; \text{ где } \bar{\mathbf{z}}_{t-2} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \ \ {(*)} \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{z} \\ q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \end{aligned} $$(*) Пояснение ко второму переходу. У нас выходит
$$ \sqrt{\alpha_t(1 - \alpha_{t - 1})}z_{t - 2} + \sqrt{1 - \alpha_t}z_{t - 1} \\ = \sqrt{\alpha_t(1 - \alpha_{t - 1}) + (1 - \alpha_t)}\bar{z}_{t- 2} \\ = \sqrt{1 - \alpha_t\alpha_{t - 1}}\bar{z}_{t - 2}; \text{ где } z_{t - 1},z_{t - 2},\bar{z}_{t - 2} \sim \mathcal{N}(0, I) $$
Тогда $L_{VLB}$ может быть переписано как
Долгий вывод
Серым в скобках комментарий к последующему переходу.
$$ \begin{aligned}L_\text{VLB} &= \mathbb{E}_{q(\mathbf{x}_{0:T})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \quad \textit{\color{gray}{(расписываем совместное распределение)}}\\&= \mathbb{E}_q \Big[ \log\frac{\prod_{t=1}^T q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{ p_\theta(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) } \Big] \ \ \ \textit{\color{gray}{(берем логарифм)}} \\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=1}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} \Big] \textit{\color{gray}{(отщепляем члены суммы)}} \\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(*)}}\\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \Big( \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(лог произведения раскрываем)}}\\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(от второй суммы останется только 1ый и последний член)}} \\&= \mathbb{E}_q \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \textit{\color{gray}{(комбинируем 1 и 3 член, 3 и 4 член)}}\\&= \mathbb{E}_q \Big[ \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \Big] \\&= \mathbb{E}_q [\underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))}_{L_{t-1}} \underbrace{- \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} ]\end{aligned} $$Пояснение (*). Пользуемся тем, что у нас Марковский процесс, и теоремой Байеса:
$$ q(x_t | x_{t - 1}) = q(x_t|x_{t - 1}, x_0) = \frac{q(x_{t - 1}| x_t, x_0)q(x_t | x_0)}{q(x_{t - 1}|x_0)} $$
Таким образом во время обучения, на каждой итерации параллельно оптимизируются случайные член $L_t$ с помощью градиентного спуск (сэмлируем $t \sim U{1,...,T}$). Поскольку все распределения нормальные, то KL между ними можно выписать в явной форме (см. ниже).
Формула KL между двумя нормальными
$$ \begin{split} KL(\mathcal{N}_1 \ || \ \mathcal{N}_2) = \frac{1}{2}\bigg(Tr(\Sigma^{-1}_2 \Sigma_1) + (\mu_2 - \mu_1)^T \Sigma^{-1}_2 (\mu_2 - \mu_1) + \[1.5ex] + \ln \frac{det(\Sigma_2)}{det(\Sigma_1)} - d \bigg) \end{split} $$
Если $\Sigma_1 = \sigma_1I, \ \Sigma_2 = \sigma_2I$
$$ \begin{split} KL(\mathcal{N}_1 \ || \ \mathcal{N}_2) = \frac{1}{2}\bigg(\frac{d\sigma_1}{\sigma_2} + \frac{1}{\sigma_2}|\mu_2 - \mu_1|^2 + \[1.5ex] + \ln \frac{\sigma_2}{\sigma_1}\bigg) \end{split} $$
Осталось только выписать $q(x_{t - 1}| x_t, x_0)$ . Мы знаем, поскольку у нас все распределения нормальные, то и $q(x_{t - 1}| x_t, x_0)$ будет нормальным.
Обозначим
Вывод $q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)$
Применим формулу Байеса и распишем. Тут мы просто пытаемся понять, как будут выглядеть среднее и дисперсия, выделяя квадратичную форму в показателе экспоненты
$$ q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \propto $$$$ \propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) $$$$ = \exp \Big(-\frac{1}{2} \big(\frac{\mathbf{x}_t^2 - 2\sqrt{\alpha_t} \mathbf{x}_t \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \alpha_t} \color{red}{\mathbf{x}_{t-1}^2} }{\beta_t} + \frac{ \color{red}{\mathbf{x}_{t-1}^2} \color{black}{- 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0} \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \bar{\alpha}_{t-1} \mathbf{x}_0^2} }{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big)= $$$$ = \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t+\frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)} \mathbf{x}_{t-1} \color{black}{ + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)} $$Далее перепишем красные и синие выражения в более красивой форме
$$ \color{red}{\tilde{\beta}_t} = 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) = 1/(\frac{\alpha_t - \bar{\alpha}_t + \beta_t}{\beta_t(1 - \bar{\alpha}_{t-1})}) = \boxed{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} $$$$ \color{blue}{\tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0)} = (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) $$$$ = (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t $$$$ = \boxed{\frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0} $$
Другой лосс. Предсказываем шум
В прошлой подсекции наша модель предсказывала среднее и дисперсию нормального распределения. Давайте зафиксируем $\Sigma_{\theta}(x_t, t) = \sigma^2_tI$. Обычно берут $\sigma^2_t = \beta_t$ или $\sigma^2_t = \tilde{\beta}_t = \frac{1 - \bar{\alpha}_{t - 1}}{1 - \bar{\alpha}_t}\beta_t.$ Тогда $L_{t - 1}$ из предыдущей секции можно переписать как
Это первый момент, как меняется функционал, если мы не хотим предсказывать $\Sigma_{\theta}(x_t, t)$, а фиксируем её.
Теперь вспомним, что $q(x_t | x_0) = \mathcal{N}(x_t;\sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t)I)$, но благодаря тому, что у нас гауссовское распределение, это можно переписать в виде
$$ x_t(x_0, \epsilon) = \sqrt{\bar{\alpha}_t} x_0 + (1-\bar{\alpha}_t)\epsilon, \ \ \ \epsilon \sim \mathcal{N}(0, I) $$
Выразим отсюда $x_0$ и получим, что $x_0 = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \sqrt{1 - \bar{\alpha}_t}\epsilon)$, тогда подставим это выражение в формулу для $\tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0)$ (из подсекции «Вывод $q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)$») и получим
$$ \tilde{\boldsymbol{\mu}}(\mathbf{x}_t, \mathbf{x}_0) = \frac{1}{\sqrt{\bar{\alpha}_t}}(x_t - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\epsilon) $$
Теперь скажем, что наша модель будет предсказывать $\epsilon$. И просто будем «подставлять» его в выражение для $\tilde{\mu}$ выше. Обозначим предсказание модели как $\epsilon_{\theta}(x_t, t)$ — предсказанный шум $\epsilon$. Тогда лосс $L_t$ превратиться в
Тем не менее лосс можно еще больше упростить и просто обучать с помощью MSE на $\epsilon$.
Итак, алгоритмы обучения и сэмплирования выглядят вот так (на картинке $z:= \epsilon$).
Стоит отметить, что важным недостатком диффузионных моделей является низкая скорость сэмплирования. Согласно Song et al. 2020: «Требуется 20 часов на генерацию 50 тысяч картинок размера 32х32, используя DDPM, и меньше минуты, используя GAN» (Nvidia 2080 Ti GPU). Тем не менее, в данном направлении был достигнут значительный прогресс и в целом проблема медленного сэмплирования была частично решена: Jiaming Song et al. (2021), Kong & Ping (2021), Bond-Taylor et al. (2021)
Давайте зафиксируем, какие функции потерь можно использовать. Для всех них справедлив тот факт, что мы сэмплируем шаг равномерно во время обучение $t \sim U{1,...,T})$ и оптимизируем соответствующий $L_{t}$.
- Оптимизируя член из суммы $L_{VLB}$. Это KL дивергенция между двумя нормальными распределениями
- При фиксированной дисперсии $\Sigma_{\theta}$ можно оптимизировать взвешенную MSE между средними нормальных распределений
- При фиксированной дисперсии и при предсказании шума с помощью взвешенной MSE. Или просто MSE. $L^{simple}_t$ является самым популярным вариантом, который на практике дает лучшие результаты.
Выбор расписания $\beta_t$
Расписание является гиперпараметром, основными требованиями на который являются невозрастание $(\beta_1 \leq ... \leq \beta_T)$ и чтобы прямой процесс сходился к $\mathcal{N}(0, I)$ в пределе по $T$. Второе может гарантироваться тем, что $\bar{\alpha}_t \to 0$. Вспомним,
Однако на практике оно также проверяется, чтобы $D_{KL}(q(x_T | x_0) || \mathcal{N}(0, I))$ было близко к 0.
Также стоит упомянуть, что обычно берут $T = 1000$. Но также важно помнить про требования выше, ведь расписание шума непосредственно зависит от $T$.
Чаще всего используют линейное расписание, где $\beta_1 = 10^{-4}, \ \beta_T = 0.02$. У данных констант нет никакой мотивации, кроме той, которая описана выше. Они были предложены в Ho et al. (2020).
В Nichol & Dhariwal (2021) было предложено косинусное расписание, которое помогло диффузионным моделям достичь лучшего NLL (negative loglikelihood):
Авторы обнаружили, что линейное расписание плохо работает на картинках 64х64 и меньше. А именно, последнии шаги прямого прохода были шумными и малоинформатиыными (просто зашумляем шум еще больше):
Также они обнаружили, что если обучать модель с линейным расписанием только на 80% первых шагов, то модель не становится сильно хуже, что подтверждает неиформативность последних шагов. Далее, они подобрали расписание так, чтобы $\sqrt{\bar{\alpha}_t}$ убывало линейно на большей части отрезка (от 0 до $T$) и почти не менялось рядом с 0 и $T$. Разницу в $\sqrt{\bar{\alpha}}_t$ для разных расписаний можно увидеть на картинке ниже:
Детали
Также они ограничивают $\beta_t$ числом 0.999, чтобы в конце процесса не было проблем с численной устойчивостью. Коэффициент $s$ используется, чтобы $\beta_t$ не были слишком малы рядом с нулем. Он равен 0.008. Такое число было выбрано так, чтобы «$\sqrt{\beta_0}$ была немного меньше, чем размер бина одного пикселя, то есть $1/127.5$»
Classifier guidance
В Nichol & Dhariwal (2021) был предложен метод условной генерации, который повышает качество генерируемых картинок, при этом уменьшая их разнообразие. Для этого предобучается «шумный» классификатор на зашумленных картинках, то есть $p_{\phi}(y |x_t)$. Затем он используется во время сэмплирования, корректируя предсказанное среднее на $\nabla_x \log p_{\phi}(y|x_t)$. В Nichol & Dhariwal (2021) (Секция 4.1) показывают, что данная добавка позволяет превратить распределение $p_\theta(x_{i-1}|x_{i})$ в $p_\theta(x_{i-1}|x_{i}, y)$. Важно, что исходная диффузионная модель никак не меняется, что делает трюк еще более привлекательным. Алгоритм сэмплирования можно видеть на картинке ниже. Коэффициент $s$ отвечает за силу guidance.
Мотивация
У генеративной модели GAN есть способ, который позволяет «балансировать» между разнообразием картинок и их качеством — truncation trick. Он заключается в сэмплировании латентного вектора truncated normal distibution. Данный трюк был хорошо описан и исследован в статье про BigGAN. Поэтому в диффузионных моделях тоже хотелось бы иметь метод, который позволяет балансировать между качеством и разнообразием. Авторы предложили classifier guidance, сравнили его с truncation trick и показали, что их метод строго лучше.
Classifier-free guidance
Ho & Salimans (2021) предложили метод, в котором guidance достигается без использования дополнительной модели, поскольку это достаточно затратно. Для этого они обучали условную модель $\epsilon_{\theta}(x_t | y)$, у которой во время обучения реальная метка $y$ заменялась с какой-то фиксированной вероятностью (10%) на пустую метку ($y=\emptyset$). Это по сути позволяет нам обучать безусловную модель $\epsilon_{\theta}(x_t)$ одновременно с условной $\epsilon_{\theta}(x_t | y)$Тогда во время сэмплирования делаем так, чтобы предсказание немного менялось в сторону $\epsilon_{\theta}(x_t | y)$, а именно:
Мотивация этой формулы следовала из формулы Байеса:
$$ p(y | x_t) \propto \frac{p(x_t | y) }{p(x_t)} \ \implies \log p(y | x_t) \propto \log p(x_t | y) - \log p(x_t) \ \implies \nabla_{x_t} \log p(y | x_t) \propto \nabla_{x_t} \log p(x_t | y) - \nabla_{x_t} \log p(x_t) \ \implies \nabla_{x_t} \log p(y | x_t) \propto \epsilon(x_t | y) - \epsilon(x_t) $$
Тогда мы можем просто подставить $\nabla_{x_t} \log p(y | x_t)$ в формулу для classifier guidance из предыдущей подсекции и получить желаемое равенство с точностью до коэффициента $s$.
Овервью ключевых работ на сегодняшний день
- Jonathan Ho et al. «Denoising diffusion probabilistic models.» arxiv Preprint arxiv:2006.11239 (2020)
Основная работа, в которой диффузионные модели (Denoising Diffusion Probabilistic Models, DDPMs) были применены для генерации картинок. Параграф в основном построен на ней.
- Jiaming Song et al. «Denoising diffusion implicit models.» arxiv Preprint arxiv:2010.02502 (2020)
Одна из первых попыток ускорить генерацию объектов. Идея следущая: давайте изменим прямой диффузионный процесс так, чтобы используя предобученную DDPM, приближать новый обратный процесс за меньшее число шагов.
Чтобы не обучать новую модель, нам нужен прямой диффузионный процесс, у которого будет такая же (суррогатная) функция потерь, а обратный процесс все еще останется Марковским. Оказалось, что существует целое семейство не-Марковских прямых процессов, удовлетворяющих этим требования. Это семейство имеет следующий вид:
$$ q_\sigma (x_{1:T}|x_0):= q_\sigma (x_{T}|x_0)\prod_{t=2}^{T} q_\sigma (x_{t-1}|x_t,x_0), $$
где $q_\sigma (x_{T}|x_0)= \mathcal{N}(\sqrt{\alpha_t}x_0, (1 - \alpha_t)I)$ и для всех $t>1,$
$$ q_\sigma (x_{t-1}|x_t,x_0)= \mathcal{N}(\sqrt{\alpha_{t-1}}x_0+\sqrt{1-\alpha_{t-1}-\sigma_t^2 }\cdot\frac{x_t-\sqrt{\alpha_{t}}x_0}{\sqrt{1-\alpha_t}},\sigma_t^2 I) $$
Среднее было выбрано так, чтобы $q_{\sigma}(x_t | x_0) = \mathcal{N}(\sqrt{\alpha_t}x_0, (1 - \alpha_t)I)$ для всех $t$. (см. Лемму 1 в Приложении B к статье). То есть важно лишь то, чтобы маргинальное распределение $q_{\sigma}(x_t | x_0)$ не менялось по сравнению с обычным Марковским случаем. Прямой процесс может быть получен с помощью теоремы Байеса:
$$ q_\sigma (x_t|x_{t-1},x_0)=\frac{q_\sigma (x_{t-1}|x_t,x_0)q_\sigma (x_t|x_0)}{q_\sigma (x_{t-1}|x_0)} $$
Тут $\sigma$ контролирует степень стохастичности прямого процесса. Можно заметить, что в отличии от исходного диффузионного процесса, предложенный прямой процесс больше не является Марковским, так как каждый $x_t$ теперь зависит и от $x_{t-1}$ и от $x_0$. Схематично, это можно изобразить как на картинке справа. (Слева исходный диффузионный процесс для сравнения)
Заметка
Авторы обращают внимание, что функция потерь в DDPM зависит от $q(x_t|x_0)$, а не от $q(x_1{:}x_T | x_0)$ напрямую. Это означает, что нам нужно выбрать любой другой прямой диффузионный процесс, у которого $q(x_t|x_0)$ остались те же.
Далее, мы можем переписать обратный процесс в данном виде:
$$ x_{t-1}=\sqrt{\alpha_{t-1}};\underbrace{\frac{x_t-\sqrt{1-\alpha_{t}}\epsilon_\theta^{(t)}(x_t)}{\sqrt{\alpha_t}}}{"predicted:x_0"}+\underbrace{\sqrt{1-\alpha{t-1}-\sigma_t^2 }\cdot \epsilon_\theta^{(t)}(x_t)}{"direction;pointing;to;x_t"}+\underbrace{\sigma_t\epsilon_t}{random:noise} $$
Заметим, что при $\sigma_t = \sqrt{(1 - \alpha_{t - 1})(1 - \alpha_t)}\sqrt{1 - \alpha_t / \alpha_{t - 1}}$ прямой процесс становится марковским, а обратный как у DDPM (обычное сэмплирование, описанное в основной секции). При $\sigma_t = 0$ процесс сэмплирования становится детерминистичным (данный способ и называется DDIM). Ускорение сэмплирования достигается засчет использования лишь какого-то подмножества шагов ($0 \leq \tau_1 \leq ... \leq \tau_S \leq T, \ \ \ S < T$). Также одним из плюсов детерминистичного сэмплирования является возможность делать семантическую интерполяцию в латентном пространстве (как у GANов).
- Alex Nichol & Prafulla Dhariwal. «Improved denoising diffusion probabilistic models» arxiv Preprint arxiv:2102.09672 (2021)
Улучшение DDPM, в котором был предложен новое расписание шума, что улучшило NLL. Также был изучен вариант, в котором дисперсия $\Sigma_{\theta}(x_t, t)$ предсказывается моделью.
- Prafula Dhariwal & Alex Nichol. «Diffusion Models Beat GANs on Image Synthesis.» arxiv Preprint arxiv:2105.05233 (2021).
Статья, в которой показывается, что DDPM могут генерировать более качественные картинки по сравнению с GANами. Также был предложен метод conditional сэмплирования. Для этого предобучается классификатор на зашумленных сэмплах, а во время сэмплирования среднее нормального распределения «корректируется» на градиент классификатора.
- Jacob Austin et al. «Structured Denoising Diffusion Models in Discrete State-Spaces». arXiv:2107.03006 (2021)
Диффузионные модели на дискретных данных (например, текст). Вместо нормальных распределений используются категориальные. Также была обобщена мультиномиальная диффузия с помощью «матриц перехода», которые задают способ зашумления дискретных данных.
Более подробно: у нас есть $x_t \in \{1, ..., K\}$ — дискретная величина на всех шагах диффузии, тогда для каждого шага $t$ определена **матрица прямого перехода $Q_t$** такая, что $[Q_t]_{ij} = q(x_t = j| x_{t - 1} = i)$. То есть строки матрицы суммируются в единицу. Тогда если обозначить через $\mathbf{x}_t \in \mathbb{R}^K$ one-hot-закодированную версию $x_t$, то прямой процесс можно описать через категориальные распределения:
Как и в нормальных распределениях, можем выписать
Поскольку тут нет такой хорошей параметризации через $\epsilon$, как у нормальных распределений, то единственный способ обучать — с помощью KL дивергенции (членами$L_{VLB}$). Остается только понять, как выбирать $Q_t$. Помимо того, чтобы сумма в каждой строчке была один, требуется, чтобы $\bar{Q}_t$ сходилось (при $t \to \infty$) к равномерному распределению в каждой строчке (аналог нормального шума). За конкретными примерами стоит обратиться к статье.
Опишем работу метода GLIDE. Стоит задача генерировать картинки по заданному текстовому описанию. Для этого используется classifier-free guided diffusion model или CLIP. Это два разных варианта модели, которые авторы сравнивают. В первом случае модель обуславливается на эмбеддинги текста, которые были получены из обучаемого трансформера. Во втором случае guidance осуществляется за счет $\nabla_{x_t} \langle f(x_t), g(c) \rangle$ (это по сути градиент лосса метода CLIP) . Тут $f$ — это картиночный энкодер (на зашумленных картинках), а $g$ — это энкодер текстового входа. В целом, авторы получили, что classifier-free guidance генерирует более качественные картинки.
Способ описать диффузионные модели через стохастические дифференциальные уравнения.
- What are Diffusion Models?. Прекрасный блог от Lilian Weng (OpenAI).