DDPM: обратный процесс и целевая функция обучения

Прямой процесс дал нам диффузионный процесс — фиксированную цепочку, которая превращает изображение в гауссовский шум. Чтобы генерировать изображения, этот процесс нужно обратить, и это задача глубокой нейронной сети. В этой заметке мы выводим, чему обучается сеть. Путь тот же, что и для VAE: правдоподобие вычислительно неразрешимо (intractable), поэтому вместо отрицательного логарифма правдоподобия мы минимизируем его верхнюю оценку. Затем эта оценка шаг за шагом упрощается до обычной среднеквадратичной ошибки (mean squared error, MSE) на шуме. При первом чтении можно пропустить врезки (алгебру разложения на KL-члены, форму оценки 2015 года, дискретизованный декодер); сводка в конце собирает всю цепочку в шесть строк.

Обозначения. Прямой шаг добавляет шум с дисперсией $\beta_t \in (0, 1)$ на шаге $t = 1, \dots, T$. Обозначим

$$\alpha_t = 1 - \beta_t, \qquad \bar\alpha_t = \prod_{s=1}^{t} \alpha_s, \qquad \bar\alpha_0 = 1,$$

так что прямой шаг и переход в замкнутом виде из предыдущей заметки записываются как

$$\begin{gathered} q(x_t \mid x_{t-1}) = \mathcal N\!\big(\sqrt{\alpha_t}\,x_{t-1},\; \beta_t I\big), \\ q(x_t \mid x_0) = \mathcal N\!\big(\sqrt{\bar\alpha_t}\,x_0,\; (1 - \bar\alpha_t) I\big). \end{gathered}$$

Второй аргумент $\mathcal N(\mu, \sigma^2)$ — всегда дисперсия. Через $q$ обозначается фиксированный прямой процесс, а $q(x_0)$ — распределение данных, с которого он начинается; $p_\theta$ — обучаемый обратный процесс, $\theta$ — параметры нейронной сети, а $p(x_T) = \mathcal N(0, I)$ — распределение шума, с которого начинается генерация.

Обратный процесс#

Изображение $x_0$ находится на одном конце цепочки, а чистый шум $x_T$ — на другом. В прямом направлении шаг $q(x_t \mid x_{t-1})$ известен. В обратном нам нужна $p_\theta(x_{t-1} \mid x_t)$ — плотность вероятности, которую необходимо выучить, чтобы обратить процесс.

Рисунок

Прямой шаг известен, обратный нужно выучить

Что должны оптимизировать параметры $\theta$? Обычный ответ для генеративной модели — приписать реальным изображениям высокую вероятность. Мы минимизируем отрицательный логарифм правдоподобия, усреднённый по данным,

$$\mathbb{E}_{x_0 \sim q(x_0)}\big[-\log p_\theta(x_0)\big] \;\to\; \min_\theta ,$$

что равносильно максимизации правдоподобия — плотности, которую модель приписывает реальным изображениям. Вывод ниже проводится для одного изображения $x_0$; усреднение по изображениям мы вернём в конце.

Полные прямой и обратный процессы#

Полный прямой процесс — это совместное распределение всех зашумлённых версий изображения. По формуле умножения вероятностей (chain rule)

$$\begin{aligned} & q(x_1, \dots, x_T \mid x_0) \\ &\quad = q(x_1 \mid x_0)\, q(x_2 \mid x_1, x_0) \\ &\qquad \cdots q(x_T \mid x_{T-1}, \dots, x_0). \end{aligned}$$

Прямой процесс — марковская цепь: каждое следующее состояние зависит только от предыдущего, поэтому в каждом множителе остаётся только последнее условие:

$$\begin{aligned} & q(x_1, \dots, x_T \mid x_0) \\ &\quad = q(x_1 \mid x_0)\, q(x_2 \mid x_1) \cdots q(x_T \mid x_{T-1}). \end{aligned}$$

С сокращённой записью $x_{s:T} = (x_s, \dots, x_T)$, в которой $x_{1:T}$ — зашумлённые состояния, а $x_{0:T}$ — вся цепочка, полные прямой и обратный процессы имеют вид

$$\begin{gathered} \underbrace{q(x_{1:T} \mid x_0) = \prod_{t=1}^{T} q(x_t \mid x_{t-1})}_{\text{полный прямой процесс}} , \\[4pt] \underbrace{p_\theta(x_{0:T}) = p(x_T) \prod_{t=1}^{T} p_\theta(x_{t-1} \mid x_t)}_{\text{полный обратный процесс}} . \end{gathered}$$

Обратный процесс начинается с $p(x_T) = \mathcal N(0, I)$. У этого множителя нет параметров, и ему не нужны никакие знания о данных: это фиксированное стандартное нормальное распределение, одно и то же для любого набора данных. Всё обучение сосредоточено в обратных шагах $p_\theta(x_{t-1} \mid x_t)$.

Точный логарифм правдоподобия вычислительно неразрешим#

Чтобы получить точный логарифм правдоподобия изображения, возьмём совместное обратное распределение и проинтегрируем по всем промежуточным состояниям. Это маргинализация — получение маргинального распределения из совместного:

$$-\log p_\theta(x_0) = -\log \int p_\theta(x_{0:T})\, dx_{1:T} .$$

Можно ли им воспользоваться? В DDPM каждый обратный шаг — гауссиана, среднее которой вычисляется по $x_t$ нейронной сетью (точная форма приведена ниже). Сеть нелинейна по $x_t$, поэтому, в отличие от прямой цепочки линейных гауссовских шагов, композиция обратных шагов не гауссовская, и у интеграла нет замкнутого вида. Это то же препятствие, что и для декодера VAE. Остаётся численное интегрирование, поэтому посмотрим, что означает этот интеграл. Каждое значение $x_{1:T}$ — один путь от сэмпла шума $x_T$ к изображению $x_0$, и интеграл складывает вероятности всех таких путей. Для изображений у каждого $x_t$ столько измерений, сколько пикселей, а таких состояний $T$ (сотни или тысячи шагов), так что интеграл берётся по $T \times (\text{число пикселей})$ измерениям. Вычислить все эти пути мы на самом деле не можем.

Рисунок

Интеграл правдоподобия суммирует по всем путям от шума к изображению

Верхняя оценка через неравенство Йенсена#

Приём тот же, что и в VAE. Умножим и разделим на полный прямой процесс $q(x_{1:T} \mid x_0)$, который мы знаем и из которого умеем сэмплировать (это произведение гауссиан, всюду положительное, поэтому деление допустимо):

$$\begin{aligned} & -\log p_\theta(x_0) \\ &\quad = -\log \int q(x_{1:T} \mid x_0)\, \frac{p_\theta(x_{0:T})}{q(x_{1:T} \mid x_0)}\, dx_{1:T} \\ &\quad = -\log \mathbb{E}_{q(x_{1:T} \mid x_0)}\!\left[ \frac{p_\theta(x_{0:T})}{q(x_{1:T} \mid x_0)} \right]. \end{aligned}$$

Логарифм вогнут, поэтому неравенство Йенсена $\log \mathbb{E}[Y] \ge \mathbb{E}[\log Y]$ переносит его внутрь математического ожидания и превращает равенство в оценку:

$$\begin{aligned} -\log p_\theta(x_0) &\le -\mathbb{E}_{q(x_{1:T} \mid x_0)}\!\left[ \log \frac{p_\theta(x_{0:T})}{q(x_{1:T} \mid x_0)} \right]. \end{aligned}$$

Правая часть — верхняя оценка отрицательного логарифма правдоподобия. Это отрицательная нижняя оценка обоснованности (evidence lower bound, ELBO); иначе говоря, ELBO — нижняя оценка логарифма правдоподобия. В оценку входят только величины, которые мы умеем сэмплировать и вычислять. Минимизируя её, мы ограничиваем отрицательный логарифм правдоподобия сверху, но сама оценка и отрицательный логарифм правдоподобия — не одно и то же. Зазор равен KL-дивергенции (KL divergence) $D_{\mathrm{KL}}\big(q(x_{1:T} \mid x_0)\,\|\,p_\theta(x_{1:T} \mid x_0)\big)$, в точности как в VAE (насколько неточна оценка), и он тоже зависит от $\theta$. Обозначим оценку через $B$, а зазор через $G$. Тогда $-\log p_\theta(x_0) = B - G$, и шаг обучения меняет отрицательный логарифм правдоподобия на $\Delta B - \Delta G$. Поэтому само по себе уменьшение оценки ничего не гарантирует: если $B$ падает с 10 до 9, а $G$ сокращается с 3 до 1, отрицательный логарифм правдоподобия растёт с 7 до 8. Минимизация оценки улучшает правдоподобие, только если оценка падает сильнее, чем сокращается зазор; если зазор не меняется или растёт, правдоподобие улучшается не меньше, чем оценка.

Разложение оценки на KL-члены#

Выкладки приведены в приложении A статьи DDPM (Ho, Jain and Abbeel, 2020, формула 5) и кратко изложены во врезке ниже. Оценка превращается в сумму членов трёх видов:

$$\begin{aligned} & \mathbb{E}_q\Bigg[\, \underbrace{D_{\mathrm{KL}}\big(q(x_T \mid x_0) \,\|\, p(x_T)\big)}_{\text{не зависит от }\theta} \\ & + \sum_{t > 1} D_{\mathrm{KL}}\big(q(x_{t-1} \mid x_t, x_0) \,\|\, p_\theta(x_{t-1} \mid x_t)\big) \\ & \underbrace{-\, \log p_\theta(x_0 \mid x_1)}_{\text{реконструкция}} \,\Bigg]. \end{aligned}$$

Здесь $\mathbb{E}_q$ — усреднение по зашумлённым состояниям $x_{1:T} \sim q(x_{1:T} \mid x_0)$ одного изображения. Каждому члену нужны лишь некоторые из них: $t$-я KL-дивергенция зависит только от $x_t$, поэтому её среднее берётся по $x_t \sim q(x_t \mid x_0)$, а член реконструкции зависит только от $x_1$. При обучении мы также усредняем по изображениям $x_0 \sim q(x_0)$.

Разберём три члена по очереди.

• Первый член. $q(x_T \mid x_0)$ — фиксированный прямой процесс, а $p(x_T)$ — фиксированное стандартное нормальное распределение. Ничто здесь не зависит от $\theta$, поэтому для обучения это константа. (Кроме того, при удачно выбранном расписании шума, noise schedule, он очень мал: линейное расписание DDPM доходит до $\bar\alpha_{1000} \approx 4 \times 10^{-5}$, так что $q(x_T \mid x_0) \approx \mathcal N(0, I)$. Игрушечное расписание с $T = 100$ шагами и $\beta_t$, растущими от $10^{-5}$ до $10^{-2}$, заканчивается около $\bar\alpha_{100} \approx 0.6$ и до этого не доходит.)
• Последний член. $-\log p_\theta(x_0 \mid x_1)$ показывает, насколько вероятно, что модель восстановит изображение по чуть более зашумлённому $x_1$. В $x_1$ шума почти не осталось, поэтому в выводе мы этот член откладываем. Это упрощение: DDPM сохраняет его в виде дискретизированного декодера для 8-битных значений пикселей (Ho, Jain and Abbeel, 2020, §3.3, формула 13; врезка ниже), а в упрощённой функции потерь его приближает член регрессии шума при $t = 1$ (§3.4).
• Средние члены. Именно они выполняют основную работу. Каждый сравнивает «истинное апостериорное распределение» (true posterior) $q(x_{t-1} \mid x_t, x_0)$ с шагом обратной диффузии модели $p_\theta(x_{t-1} \mid x_t)$.

После упрощения целевая функция имеет вид

$$\begin{aligned} \mathbb{E}_q\Bigg[ \sum_{t > 1} D_{\mathrm{KL}}\Big( & \underbrace{q(x_{t-1} \mid x_t, x_0)}_{\text{истинное апостериорное распределение}} \\ & \Big\|\, \underbrace{p_\theta(x_{t-1} \mid x_t)}_{\text{обратный шаг}} \Big) \Bigg]. \end{aligned}$$

Истинное апостериорное распределение#

Три состояния и три условных распределения укладываются в один небольшой граф. $x_{t-1}$ зашумлено, $x_t$ зашумлено сильнее, а $x_0$ — исходное изображение.

Рисунок

Прямой шаг, шаг модели и истинное апостериорное распределение

Прямой шаг $q(x_t \mid x_{t-1})$ идёт вправо. Шаг модели $p_\theta(x_{t-1} \mid x_t)$ идёт влево и использует только $x_t$. Истинное апостериорное распределение $q(x_{t-1} \mid x_t, x_0)$ тоже идёт влево, но смотрит и на $x_t$, и на исходное изображение $x_0$.

У истинного апостериорного распределения есть выражение в замкнутом виде (выведено ниже). Зачем же тогда вообще нужна нейронная сеть? Затем, что для истинного апостериорного распределения нужно знать $x_0$. При генерации изображения у нас нет: его нужно получить, начав с одного лишь чистого шума. Поэтому сеть $p_\theta(x_{t-1} \mid x_t)$ должна подражать истинному апостериорному распределению, не видя $x_0$.

Итак, мы обучаем обратный шаг модели совпадать с истинным апостериорным распределением:

$$\begin{gathered} p_\theta(x_{t-1} \mid x_t) = \mathcal N\!\big(\mu_\theta(x_t, t),\; \sigma_\theta^2(x_t, t)\, I\big) \\ \big\updownarrow \\ q(x_{t-1} \mid x_t, x_0) = \mathcal N\!\big(\tilde\mu_t(x_t, x_0),\; \tilde\beta_t I\big). \end{gathered}$$

Почему гауссианы?

1. При малых шагах обратный шаг остаётся гауссовским. Для диффузии с малыми шагами обращение гауссовской марковской цепи тоже (приближённо) гауссовское: при $\beta_t \to 0$ обратный шаг $q(x_{t-1} \mid x_t)$ имеет тот же гауссовский функциональный вид, что и прямой (Feller, 1949; используется у Sohl-Dickstein et al., 2015, §2.2). При конечных шагах $q(x_{t-1} \mid x_t)$ не является в точности гауссовским, потому что это смесь по всем изображениям $x_0$ из данных. Апостериорное распределение при условии $x_0$, $q(x_{t-1} \mid x_t, x_0)$, в точности гауссовское при любом $\beta_t$, если $t \ge 2$. При $t = 1$ оно вырождается в точечную массу в $x_0$ ($\tilde\beta_1 = 0$), поэтому этот шаг учитывается членом реконструкции.
2. Так попросту проще. У KL-дивергенций между гауссианами есть замкнутые формулы, а сэмплировать из гауссианы тривиально.

Замкнутый вид. Пусть $t \ge 2$. По теореме Байеса $q(x_{t-1} \mid x_t, x_0) \propto q(x_t \mid x_{t-1})\, q(x_{t-1} \mid x_0)$. Оба множителя гауссовские по $x_{t-1}$. У первого показатель экспоненты равен $-(x_t - \sqrt{\alpha_t}\,x_{t-1})^2/(2\beta_t) = -\alpha_t\,(x_{t-1} - x_t/\sqrt{\alpha_t})^2/(2\beta_t)$ (для каждой координаты), так что как функция $x_{t-1}$ он имеет среднее $x_t/\sqrt{\alpha_t}$ и обратную дисперсию (precision) $\alpha_t/\beta_t$. Второй равен $\mathcal N\big(\sqrt{\bar\alpha_{t-1}}\,x_0,\; 1 - \bar\alpha_{t-1}\big)$. Выделяя полный квадрат в произведении, получаем, что обратные дисперсии складываются, а среднее — это среднее двух средних, взвешенное по обратным дисперсиям:

$$\begin{aligned} \frac{1}{\tilde\beta_t} &= \frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar\alpha_{t-1}} = \frac{\alpha_t (1 - \bar\alpha_{t-1}) + \beta_t}{\beta_t (1 - \bar\alpha_{t-1})} \\ &= \frac{1 - \bar\alpha_t}{\beta_t (1 - \bar\alpha_{t-1})}, \\ \tilde\mu_t &= \tilde\beta_t \left( \frac{\sqrt{\alpha_t}}{\beta_t}\, x_t + \frac{\sqrt{\bar\alpha_{t-1}}}{1 - \bar\alpha_{t-1}}\, x_0 \right). \end{aligned}$$

Числитель упрощается: $\alpha_t(1 - \bar\alpha_{t-1}) + \beta_t = \alpha_t - \bar\alpha_t + \beta_t = 1 - \bar\alpha_t$, где использованы $\bar\alpha_t = \alpha_t \bar\alpha_{t-1}$ и $\alpha_t + \beta_t = 1$. Переходя от обратной дисперсии к дисперсии и раскрывая скобки в среднем, получаем

$$\begin{aligned} \tilde\mu_t(x_t, x_0) &= \frac{\sqrt{\bar\alpha_{t-1}}\,\beta_t}{1 - \bar\alpha_t}\, x_0 \\ &\quad + \frac{\sqrt{\alpha_t}\,(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t}\, x_t, \\ \tilde\beta_t &= \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t}\, \beta_t . \end{aligned}$$

Среднее истинного апостериорного распределения — взвешенная смесь изображения и текущего зашумлённого состояния.

Сопоставление двух гауссиан: KL превращается в L2#

Нам по-прежнему пришлось бы предсказывать и $\mu_\theta$, и $\sigma_\theta$. DDPM упрощает задачу дальше: $\sigma_\theta$ фиксируется расписанием $\sigma_t$, то есть не обучается, и обучается только среднее. Ho et al. используют $\sigma_t^2 = \beta_t$ или $\sigma_t^2 = \tilde\beta_t$ и сообщают, что результаты похожи (Ho, Jain and Abbeel, 2020, §3.2); позже Nichol and Dhariwal (2021), §3.1, обучают дисперсию, которая интерполирует между этими двумя вариантами.

При фиксированных дисперсиях KL-дивергенция между двумя гауссианами зависит от $\theta$ только через расстояние между средними. Для $q = \mathcal N(\tilde\mu_t, \tilde\beta_t I)$ и $p_\theta = \mathcal N(\mu_\theta, \sigma_t^2 I)$ в $d$ измерениях KL-дивергенция между двумя одномерными гауссианами, просуммированная по $d$ координатам, даёт

$$\begin{aligned} D_{\mathrm{KL}}(q \,\|\, p_\theta) &= \frac{1}{2\sigma_t^2}\,\big\|\tilde\mu_t - \mu_\theta\big\|^2 \\ &\quad + \underbrace{\frac d2\left(\frac{\tilde\beta_t}{\sigma_t^2} - 1 - \log\frac{\tilde\beta_t}{\sigma_t^2}\right)}_{\text{const; }0\text{ при }\sigma_t^2 = \tilde\beta_t} . \end{aligned}$$

Тот же результат следует из общей формулы для двух гауссиан в $d$ измерениях с положительно определёнными ковариационными матрицами,

$$\begin{aligned} & D_{\mathrm{KL}}\big(\mathcal N(\mu_q, \Sigma_q) \,\|\, \mathcal N(\mu_p, \Sigma_p)\big) \\ &\quad = \tfrac12\Big[\operatorname{tr}(\Sigma_p^{-1}\Sigma_q) - d + \ln\frac{\det\Sigma_p}{\det\Sigma_q} \\ &\qquad\quad + (\mu_p - \mu_q)^\top \Sigma_p^{-1} (\mu_p - \mu_q)\Big], \end{aligned}$$

при $\Sigma_q = \tilde\beta_t I$ и $\Sigma_p = \sigma_t^2 I$: след даёт $d\,\tilde\beta_t/\sigma_t^2$, логарифм отношения определителей — $d \ln(\sigma_t^2/\tilde\beta_t)$, а квадратичная форма — $\lVert\tilde\mu_t - \mu_\theta\rVert^2/\sigma_t^2$.

Таким образом, KL-член на самом деле — функция потерь L2 $\frac{1}{2\sigma_t^2}\|\tilde\mu_t - \mu_\theta\|^2$ плюс константа. Для одной пары $(x_t, x_0)$ L2-часть обращалась бы в ноль при $\mu_\theta$, равном $\tilde\mu_t$, а весь KL-член исчез бы, только если вдобавок $\sigma_t^2 = \tilde\beta_t$. Но $\tilde\mu_t$ зависит от $x_0$, которого сеть не видит, а одному и тому же $x_t$ соответствует много изображений. Поскольку $\mu_\theta$ видит только $x_t$, лучшее, что она может сделать, — выдать среднее $\tilde\mu_t$ по всем изображениям $x_0$, совместимым с $x_t$, потому что условное среднее минимизирует ожидаемую квадратичную ошибку:

$$\mu_\theta^*(x_t, t) = \mathbb{E}_q\big[\tilde\mu_t(x_t, x_0) \mid x_t\big],$$

где усреднение идёт по $x_0 \sim q(x_0 \mid x_t)$. По формуле полного математического ожидания (tower property) это также $\mathbb{E}_q[x_{t-1} \mid x_t]$ — среднее истинного обратного шага $q(x_{t-1} \mid x_t)$. Поэтому в оптимуме KL-член не равен нулю: остаток отражает, насколько $x_t$ не определяет изображение однозначно.

Рисунок

При равных дисперсиях KL-дивергенция — функция потерь L2 для средних

Итак:

$$\mathbb{E}_q\!\left[ \sum_{t > 1} \frac{1}{2\sigma_t^2}\, \big\| \tilde\mu_t - \mu_\theta(x_t, t) \big\|^2 \right] \;\to\; \min_\theta .$$

От средних к шуму#

Среднее истинного апостериорного распределения всё ещё содержит $x_0$. Исключим его с помощью репараметризации (reparameterization) перехода в замкнутом виде, как в VAE:

$$\begin{gathered} x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \varepsilon, \quad \varepsilon \sim \mathcal N(0, I) \\ \Longrightarrow\quad x_0 = \frac{1}{\sqrt{\bar\alpha_t}}\big(x_t - \sqrt{1 - \bar\alpha_t}\, \varepsilon\big). \end{gathered}$$

Подставим это в $\tilde\mu_t$. Ключевое тождество — $\bar\alpha_t = \alpha_t \bar\alpha_{t-1}$, так что $\sqrt{\bar\alpha_{t-1}}/\sqrt{\bar\alpha_t} = 1/\sqrt{\alpha_t}$. После приведения подобных коэффициент при $x_t$ равен

$$\begin{aligned} &\frac{\beta_t}{\sqrt{\alpha_t}\,(1 - \bar\alpha_t)} + \frac{\sqrt{\alpha_t}\,(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t} \\ &\quad= \frac{\beta_t + \alpha_t - \bar\alpha_t}{\sqrt{\alpha_t}\,(1 - \bar\alpha_t)} = \frac{1 - \bar\alpha_t}{\sqrt{\alpha_t}\,(1 - \bar\alpha_t)} = \frac{1}{\sqrt{\alpha_t}}, \end{aligned}$$

где использовано $\beta_t + \alpha_t = 1$. Коэффициент при $\varepsilon$ равен $-\frac{\beta_t}{\sqrt{\alpha_t}\,(1-\bar\alpha_t)}\sqrt{1 - \bar\alpha_t} = -\frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1 - \bar\alpha_t}}$. Следовательно,

$$\tilde\mu_t = \frac{1}{\sqrt{\alpha_t}}\left( x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}}\, \varepsilon \right).$$

Денойзеру (модели шумоподавления) доступны $x_t$ и $t$, но не $\varepsilon$. Поэтому дадим среднему модели в точности такую же форму и поручим сети предсказывать только недостающую часть — шум (Ho, Jain and Abbeel, 2020, формула 11):

$$\mu_\theta(x_t, t) = \frac{1}{\sqrt{\alpha_t}}\left( x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}}\, \varepsilon_\theta(x_t, t) \right).$$

Теперь и $\tilde\mu_t$, и $\mu_\theta(x_t, t)$ выражены только через $x_t$ и шум. В разности $x_t$ сокращается, остаётся $\tilde\mu_t - \mu_\theta = \frac{\beta_t}{\sqrt{\alpha_t}\sqrt{1 - \bar\alpha_t}}\big(\varepsilon_\theta - \varepsilon\big)$, и функция потерь принимает вид

$$\begin{aligned} L = \mathbb{E}_q\Bigg[ & \sum_{t > 1} \frac{\beta_t^2}{2\sigma_t^2\, \alpha_t\, (1 - \bar\alpha_t)} \\ & \cdot \big\| \varepsilon - \varepsilon_\theta(x_t, t) \big\|^2 \Bigg]. \end{aligned}$$

Это формула 12 из Ho, Jain and Abbeel (2020), просуммированная по $t$. Здесь $\varepsilon_\theta(x_t, t)$ — предсказание шума сетью. Как и в случае средних, сеть не может восстановить тот конкретный $\varepsilon$, который породил $x_t$, потому что не видит $x_0$. Оптимум — $\varepsilon_\theta^*(x_t, t) = \mathbb{E}[\varepsilon \mid x_t]$, средний шум, согласующийся с $x_t$; это тот же оптимум, что и $\mu_\theta^*$, поскольку при фиксированном $x_t$ величина $\tilde\mu_t$ — аффинная функция $\varepsilon$. Денойзер — это предсказатель шума.

Последнее упрощение: случайный шаг для каждого сэмпла#

Суммировать все $T - 1$ членов для каждого обучающего изображения слишком дорого: это стоит $T - 1$ вычислений сети на изображение. Вместо этого выберем случайный шаг для каждого сэмпла. Поскольку $\sum_{t=2}^{T} f(t) = (T - 1)\, \mathbb{E}_{t \sim \mathrm{Uniform}\{2, \dots, T\}}\big[f(t)\big]$, одно равномерно выбранное $t$ даёт несмещённую оценку суммы, делённой на $T - 1$, а постоянный множитель не меняет точку минимума. Усредняя также по изображениям $x_0 \sim q(x_0)$ и шумам $\varepsilon \sim \mathcal N(0, I)$ при $t \sim \mathrm{Uniform}\{2, \dots, T\}$, получаем итоговую функцию потерь

$$\begin{gathered} \mathbb{E}_{x_0,\, t,\, \varepsilon}\!\left[ \frac{\beta_t^2}{2\sigma_t^2\, \alpha_t\, (1 - \bar\alpha_t)}\, \big\| \varepsilon - \varepsilon_\theta(x_t, t) \big\|^2 \right], \\ x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \varepsilon . \end{gathered}$$

Парные шаги. Оценка остаётся несмещённой, как бы ни зависели друг от друга шаги одного батча, если каждый из них по отдельности распределён равномерно. Приём снижения дисперсии из открытого кода DDIM (Song, Meng and Ermon, 2021) выбирает $t$ для половины батча, а второй половине даёт зеркальные шаги $t' = T + 1 - t$ (когда $t$ пробегает $1, \dots, T$), тоже распределённые равномерно. Пусть $f(t)$ — ожидаемая функция потерь на шаге $t$. Зеркальная пара оценивает среднее $f$ с дисперсией $\tfrac12\big[\operatorname{Var} f(t) + \operatorname{Cov}\big(f(t), f(t')\big)\big]$, а два независимых выбора — с дисперсией $\tfrac12 \operatorname{Var} f(t)$. Если $f$ монотонна по $t$, то $f(t)$ и $f(T + 1 - t)$ меняются в противоположных направлениях, ковариация по неравенству Чебышёва для сумм (Chebyshev’s sum inequality) отрицательна, и пара выигрывает. Игрушечный пример: $T = 4$ и $f(t) = 1/t$ — функция потерь, наибольшая на наименее зашумлённом шаге. Два независимых выбора дают дисперсию $0.0423$, зеркальная пара — $0.0109$, примерно в $3.9$ раза меньше. Функция потерь реальной сети на шаге монотонна по $t$ лишь приблизительно, поэтому выигрыш на практике проверяется только экспериментом.

Шаг обучения. Одному шагу обучения цепочка вообще не нужна: берём изображение, шаг и шум, переходим прямо к $x_t$ за один скачок и решаем задачу регрессии для шума.

Рисунок

Один шаг обучения: сэмплировать, перейти, предсказать шум

Отбрасывание веса — сознательное перевзвешивание, а не косметическое изменение. При линейном расписании из статьи ($T = 1000$, $\beta_t$ от $10^{-4}$ до $0.02$) и $\sigma_t^2 = \beta_t$ вес упрощается до $\beta_t / (2\alpha_t(1 - \bar\alpha_t))$. Он равен $0.273$ при $t = 2$, падает до минимума $0.0050$ при $t = 349$ и снова растёт до $0.0102$ при $t = T$. Истинная оценка придаёт шагу $t = 2$ примерно в $55$ раз больший вес, чем шагу в минимуме. $L_{\text{simple}}$ взвешивает все шаги одинаково, что смещает усилия к более трудным, более зашумлённым шагам. Ho et al. обнаружили, что это даёт лучшие сэмплы ценой немного худшего правдоподобия (§3.4 и таблица 1: FID 3.17 против 13.51 для истинной оценки, но $3.75$ против $3.70$ бита на размерность).

Рисунок

Вес из оценки отдаёт предпочтение почти чистым шагам

Сэмплирование: цепочка в обратном направлении#

При обучении цепочка не запускается, а при генерации — запускается. Начнём с чистого шума $x_T \sim \mathcal N(0, I)$ и применим выученный обратный шаг $T$ раз. Каждый шаг сэмплирует из $p_\theta(x_{t-1} \mid x_t) = \mathcal N(\mu_\theta(x_t, t), \sigma_t^2 I)$, используя выведенную выше форму среднего через шум, с новым шумом $z \sim \mathcal N(0, I)$ при $t > 1$ и $z = 0$ при $t = 1$ (Ho, Jain and Abbeel, 2020, алгоритм 2):

$$\begin{aligned} x_{t-1} = {} & \frac{1}{\sqrt{\alpha_t}}\left( x_t - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}}\, \varepsilon_\theta(x_t, t) \right) \\ & + \sigma_t z . \end{aligned}$$

Каждый шаг вычитает оценку шума, полученную сетью, масштабирует результат и добавляет немного нового шума; последний шаг возвращает среднее без дополнительного шума. Поэтому генерация стоит $T$ вычислений сети, по одному на шаг, и это главный практический недостаток DDPM. Более поздние сэмплеры, например DDIM (Song, Meng and Ermon, 2021), используют тот же обученный $\varepsilon_\theta$ с гораздо меньшим числом шагов.

Сеть шумоподавления#

В выводе $\varepsilon_\theta(x_t, t)$ рассматривается как функция двух аргументов. На практике это одна сеть для всех шагов — U-Net, которая получает шаг в виде эмбеддинга, прибавляемого внутри каждого блока.

Рисунок

Денойзер DDPM: U-Net, получающая на вход шаг t

1. Зашумлённое изображение $x_t$ входит в энкодер. Каждый уровень вдвое уменьшает разрешение; на $16 \times 16$ DDPM добавляет self-attention.
2. Шаг $t$ превращается в вектор синусов и косинусов $e(t)$.
3. Небольшой MLP переводит $e(t)$ в $m(t)$, и каждый остаточный блок прибавляет к своим каналам свою линейную проекцию $A_k\,m(t)$.
4. Декодер уровень за уровнем вдвое увеличивает разрешение; обходные связи (skip connections) приносят признаки энкодера того же размера.
5. Выход $\varepsilon_\theta(x_t, t)$ имеет форму $x_t$ и сравнивается с истинным шумом $\varepsilon$.

Одна сеть для всех шагов#

Отдельная сеть для каждого шага означала бы $T = 1000$ сетей, каждая из которых обучалась бы на тысячной доле обновлений. Соседние шаги решают почти одну и ту же задачу: $x_{t-1}$ и $x_t$ отличаются одним малым шагом: умножением на $\sqrt{1 - \beta_t} \ge 0.99$ и шумом с дисперсией $\beta_t \le 0.02$. Поэтому DDPM разделяет параметры между шагами и сообщает сети, на каком шаге она находится (Ho, Jain and Abbeel, 2020, §4). Шаг нужен потому, что один и тот же вход требует разных выходов при разных уровнях шума.

Пример. Для данных $\mathcal N(0, 0.5^2)$ в одном измерении $\varepsilon$ и $x_t$ совместно гауссовские с ковариацией $\sqrt{1 - \bar\alpha_t}$, поэтому оптимальное предсказание равно

$$\mathbb E[\varepsilon \mid x_t] = \frac{\sqrt{1 - \bar\alpha_t}}{0.25\,\bar\alpha_t + 1 - \bar\alpha_t}\, x_t .$$

При линейном расписании и $x_t = 1$ оно равно $0.17$ при $t = 10$ ($\bar\alpha_t = 0.998$), где $x_t$ — в основном изображение, и $1.02$ при $t = 500$ ($\bar\alpha_t = 0.079$), где $x_t$ — в основном шум.

U-Net#

Выход $\varepsilon_\theta(x_t, t)$ имеет форму входа — одно значение шума на каждый пиксель и канал, — так что сеть решает задачу плотного предсказания (dense prediction), как при сегментации. Чтобы отличить шум от детали изображения в одном пикселе, ей нужен контекст издалека и точность в самом пикселе. U-Net (Ronneberger, Fischer and Brox, 2015) даёт и то и другое. Энкодер уровень за уровнем вдвое уменьшает разрешение, так что грубые уровни видят всё изображение; декодер вдвое увеличивает его обратно; а обходные связи передают признаки энкодера на каждом разрешении уровню декодера того же размера, где они конкатенируются с признаками после повышения разрешения, так что мелкие детали доходят до выхода. Сети score, зависящие от уровня шума (noise-conditional score networks), используют вариант U-Net по той же причине (Song and Ermon, 2019, §4.1 и прил. A.3).

Сеть DDPM (Ho, Jain and Abbeel, 2020, приложение B):

• базовая сеть (backbone) немаскированной PixelCNN++ — U-Net из остаточных блоков Wide ResNet;
• групповая нормализация (group normalization; Wu and He, 2018) вместо нормализации весов (weight normalization): статистики считаются по группам каналов внутри одного изображения и поэтому, в отличие от батч-нормализации, не зависят от других изображений батча;
• четыре разрешения карт признаков для изображений $32 \times 32$ (от $32 \times 32$ до $4 \times 4$) и шесть для $256 \times 256$, по два остаточных блока на разрешение;
• self-attention на разрешении $16 \times 16$, между свёрточными блоками;
• шаг, который прибавляется в каждом остаточном блоке через синусоидальный эмбеддинг из трансформера;
• 35.7 млн параметров для CIFAR-10, 114 млн для LSUN и CelebA-HQ при $256 \times 256$ и около 256 млн для более крупного варианта модели для LSUN Bedroom.

Как сообщить сети шаг#

DDPM кодирует $t$ так же, как трансформер кодирует позицию (позиционное кодирование):

$$\begin{gathered} e(t)_{2i} = \sin(\omega_i t), \qquad e(t)_{2i+1} = \cos(\omega_i t), \\ \omega_i = 10000^{-2i/d}, \qquad i = 0, \dots, d/2 - 1 . \end{gathered}$$

Небольшой MLP превращает код в вектор признаков, $m(t) = W_2\,\mathrm{SiLU}(W_1\, e(t))$ с $\mathrm{SiLU}(z) = z \cdot \mathrm{sigmoid}(z)$, а у каждого остаточного блока $k$ есть своё линейное отображение $A_k$ в его число каналов $C_k$. Результат прибавляется к картам признаков блока как одно смещение на канал, одинаковое во всех пикселях:

$$h_k[c, i, j] \;\leftarrow\; h_k[c, i, j] + \big(A_k\, m(t)\big)_c .$$

Почему бы не подать $t$ просто числом? Скаляр входит в первый слой как ещё один вход, и сети приходится строить из него всю зависимость от $t$. Синусоидальный код — вектор, компоненты которого вращаются со скоростями от $1$ до примерно $10^{-4}$ радиана за шаг, поэтому их линейные комбинации могут следить и за быстрыми, и за медленными изменениями $t$, а близкие шаги получают близкие коды. Порядок компонент — вперемешку или сначала все синусы — не важен: следующий слой линеен и может переставить свои входы.

Другие способы сообщить сети шаг или другое условие:

• Масштаб и сдвиг. Эмбеддинг задаёт масштаб и сдвиг для каждого канала после нормализации, $h \leftarrow y_s(t) \odot \mathrm{GroupNorm}(h) + y_b(t)$: адаптивная групповая нормализация (adaptive group normalization; Dhariwal and Nichol, 2021, §3.1). Это диффузионный аналог адаптивной нормализации по экземпляру (adaptive instance normalization, AdaIN), с помощью которой StyleGAN вводит свой вектор стиля (Karras, Laine and Aila, 2019, формула 1). На ImageNet $128 \times 128$ она улучшила FID с $15.08$ (сложение и затем групповая нормализация, как в DDPM) до $13.06$ (их таблица 3). Сети score, зависящие от уровня шума, используют ту же идею с одним обучаемым масштабом и сдвигом на каждый уровень шума (заметка о score matching).
• Дополнительный входной канал. Добавить к изображению $t/T$ как постоянный канал. Это самый простой вариант, но шаг тогда входит только в первую свёртку, и его приходится проносить через все последующие слои.
• Обучаемый вектор на каждый шаг. Таблица с одним обучаемым вектором на шаг, который прибавляется к скрытому слою или умножается на него, работает для небольших сетей и коротких цепочек. Но у $T$ векторов нет ничего общего: каждый обучается только на примерах, выпавших на его шаг.
• Другие условия. Метку класса или усреднённый (pooled) эмбеддинг текста можно спроецировать в размерность эмбеддинга шага и прибавить к нему, как это делает Imagen (Saharia et al., 2022, §2.5); целая последовательность текстовых векторов входит через перекрёстное внимание (cross-attention; латентная диффузия).

Почему пиксели масштабируют в [−1, 1]#

DDPM линейно отображает значения пикселей $0, \dots, 255$ в $[-1, 1]$, чтобы обратный процесс работал со входами согласованного масштаба, начиная со стандартного нормального априорного распределения (Ho, Jain and Abbeel, 2020, §3.3). Это важно по трём причинам:

1. Цепочка должна заканчиваться в априорном распределении. Прямой процесс убирает изображение только сжатием, $x_T = \sqrt{\bar\alpha_T}\,x_0 + \sqrt{1 - \bar\alpha_T}\,\varepsilon$, и при расписании DDPM $\sqrt{\bar\alpha_T} = 0.0064$. Для значений в $[-1, 1]$ остаток изображения не больше $0.0064$ — пренебрежимо мало рядом с шумом со стандартным отклонением $1$. Для исходных значений в $[0, 255]$ он доходил бы до $1.6$: в $x_T$ всё ещё было бы видно изображение, и сэмплирование из $\mathcal N(0, I)$ начиналось бы не с того распределения.
2. Каждый шаг видит входы одного размера. По рекуррентному соотношению для дисперсии $\operatorname{Var}(x_t) = \bar\alpha_t \operatorname{Var}(x_0) + 1 - \bar\alpha_t$, и эта величина лежит между $\operatorname{Var}(x_0)$ и $1$, если $\operatorname{Var}(x_0) \le 1$, что гарантируют значения в $[-1, 1]$. Тогда одна сеть может служить всем шагам; для исходных значений пикселей дисперсия падала бы вдоль цепочки от тысяч до $1$.
3. Данные центрированы, как и априорное распределение, а у дискретизированного декодера интервалы шириной $2/255$ и открытые крайние интервалы у $\pm 1$.

Чтобы показать изображение, $x$ отображают обратно по формуле $255\,(x + 1)/2$ после обрезки до $[-1, 1]$. Зашумлённые состояния выходят за этот интервал постоянно: пиксель с $x_0 = 0.9$ на шаге с $\bar\alpha_t = 0.5$ имеет $x_t \sim \mathcal N(0.64,\ 0.5)$ и выходит за $[-1, 1]$ с вероятностью $0.31$. Без обрезки преобразование в 8-битные целые превращает такие значения в неправильные цвета.

Обучение на практике: усреднённые веса и обрезанные градиенты#

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

Экспоненциальное скользящее среднее весов (exponential moving average, EMA). Будем хранить вторую копию параметров $\bar\theta$ и обновлять её после каждого шага оптимизатора:

$$\bar\theta \;\leftarrow\; \mu\, \bar\theta + (1 - \mu)\, \theta .$$

Если развернуть это обновление на $n$ шагов, $\bar\theta$ окажется взвешенным средним прошлых значений параметров с весами $(1 - \mu)\,\mu^{n-k}$ плюс начальное значение с множителем $\mu^n$: примерно средним по последним $1/(1 - \mu)$ шагам. Обучение продолжается с $\theta$; сэмплы генерируются с $\bar\theta$. Причина — шум градиента: каждый пример приносит своё изображение, шаг и шум, поэтому значения $\theta$ дрожат вокруг хорошей области, а их среднее лежит ближе к её центру. DDPM использует $\mu = 0.9999$ (приложение B) — среднее примерно по $10\,000$ шагам. Nichol and Dhariwal (2021) перебрали этот коэффициент для модели, получающей метку класса, на ImageNet $64 \times 64$ и нашли лучшими $0.9999$ и $0.99995$ (их рисунок 17). Разобранный пример: начальное значение сохраняет вес $\mu^n$, равный $0.9999^{10\,000} = 0.37$ через десять тысяч шагов и $0.9999^{50\,000} = 0.0067$ через пятьдесят тысяч. Поэтому усреднёнными весами имеет смысл пользоваться, только когда обучение длится в несколько раз дольше окна усреднения. При $\mu = 0.9$ окно — около $10$ шагов, слишком мало, чтобы что-то заметно усреднить.

Обрезка градиентов (gradient clipping). Перед каждым шагом градиент $g$ масштабируют, если его норма превышает порог $c$: $g \leftarrow g \cdot \min(1, c/\lVert g\rVert)$ (заметка об обучении сетей). Направление сохраняется, а укорачиваются только редкие большие градиенты: при $c = 1$ градиент нормы $4$ делится на $4$, а градиент нормы $0.5$ остаётся как есть. Функция потерь батча зависит от того, какие шаги в него попали, и обрезка не даёт необычному батчу далеко сдвинуть веса.

Настройки DDPM (приложение B): Adam с шагом обучения $2 \times 10^{-4}$ ($2 \times 10^{-5}$ при $256 \times 256$), размер батча 128 для CIFAR-10 и 64 для более крупных изображений, дропаут (dropout) $0.1$ на CIFAR-10 и без дропаута на остальных наборах, случайные горизонтальные отражения, коэффициент EMA $0.9999$ и $T = 1000$. Модель для CIFAR-10 обучалась 800K шагов со скоростью 21 шаг в секунду на TPU v3-8, около 10.6 часа; сэмплирование батча из 256 изображений заняло 17 секунд.

Сводка: от правдоподобия к одному шагу обучения#

1. Оценка. $-\log p_\theta(x_0)$ не больше отрицательной ELBO; зазор — KL-дивергенция, поэтому уменьшение оценки само по себе не уменьшает отрицательный логарифм правдоподобия.
2. Апостериорные цели. Оценка распадается на KL-члены, которые сравнивают каждый обратный шаг $p_\theta(x_{t-1} \mid x_t)$ с гауссовским истинным апостериорным распределением $q(x_{t-1} \mid x_t, x_0)$, среднее которого $\tilde\mu_t$ известно в явном виде.
3. Взвешенная функция потерь на шуме. При фиксированных дисперсиях каждый KL-член — это L2-расстояние между средними плюс константа, а если записать $\mu_\theta$ через предсказатель шума, он становится $\frac{\beta_t^2}{2\sigma_t^2\, \alpha_t\, (1 - \bar\alpha_t)}\, \lVert \varepsilon - \varepsilon_\theta(x_t, t) \rVert^2$.
4. Перевзвешивание. $L_{\text{simple}}$ отбрасывает этот вес. Это другая цель, а не оценка: она смещает усилия к более шумным шагам и даёт лучшие сэмплы ценой немного худшего правдоподобия.
5. Один шаг обучения. Сэмплируем $x_0$, $t$ и $\varepsilon$, строим $x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\, \varepsilon$ и делаем градиентный шаг по $\lVert \varepsilon - \varepsilon_\theta(x_t, t) \rVert^2$.
6. Сэмплирование. Из $x_T \sim \mathcal N(0, I)$ применяем $x_{t-1} = \mu_\theta(x_t, t) + \sigma_t z$ для $t = T, \dots, 1$, без шума на последнем шаге.

Ноутбук#

Ноутбук по диффузионным моделям из репозитория HSE-LAMBDA Deep Generative Models (ссылка закреплена на конкретном коммите; ячейки нумеруются с 1, считая все ячейки подряд, включая текстовые; сохранённых выходов ячеек нет) обучает две версии обратного процесса на двумерном «швейцарском рулете» (Swiss roll) из 10 000 точек с $T = 100$ и сигмоидным расписанием от $10^{-5}$ до $10^{-2}$ ($\beta_1 \approx 3.5 \times 10^{-5}$, $\beta_{100} \approx 1.0 \times 10^{-2}$). Индексы в коде начинаются с 0, поэтому шаг кода t = 0 соответствует $t = 1$ в формулах.

• Ячейки 17–18 (повторены в ячейке 37): истинное апостериорное распределение. q_posterior_mean_variance использует коэффициенты $\beta_t\sqrt{\bar\alpha_{t-1}}/(1-\bar\alpha_t)$ и $(1-\bar\alpha_{t-1})\sqrt{\alpha_t}/(1-\bar\alpha_t)$ и дисперсию $\tilde\beta_t$, как в выводе выше. На шаге кода t = 0 дисперсия равна $\tilde\beta_1 = 0$ — это вырожденный случай замкнутого вида, и «обрезанный» (clipped) логарифм дисперсии всё равно берёт $\log 0 = -\infty$. В loss_likelihood_bound (ячейка 31) из-за этого KL-дивергенция становится бесконечной всякий раз, когда в батч попадает шаг кода t = 0, так что функция потерь, которую печатает ячейка 33, может быть равна inf; градиенты остаются конечными, потому что бесконечная часть не зависит от $\theta$. Этого можно избежать, если взять $\tilde\beta_2$ в качестве первого элемента.
• Ячейки 19–24: исходная диффузионная модель Sohl-Dickstein et al. Одна сеть служит всем шагам и предсказывает среднее и логарифм дисперсии. Её слои ConditionalLinear умножают выход линейного слоя на обучаемый вектор своего шага, $\gamma_t \odot (W h + b)$, где $\gamma$ — nn.Embedding(n_steps, 128), инициализированный равномерно на $[0, 1]$: это обучаемый вектор на каждый шаг, а не синусоидальный код. p_sample и p_sample_loop запускают обратную цепочку.
• Ячейки 25–29: оценка. Ячейка 26 записывает оценку как сумму логарифмов отношений по шагам из первой врезки. Затем ячейка 27 даёт ещё две формы, и обе не вполне верны: первая (априорный член, все $T$ KL-членов и одна энтропия прямого процесса) не совпадает со стандартной оценкой, а вторая берёт энтропийные члены формы 2015 года со знаками, которые у них в $K$ — нижней оценке логарифма правдоподобия, — и ставит их в функцию потерь, которая должна равняться $-K$, рядом с лишним априорным членом. normal_kl (ячейка 28) — KL-дивергенция между диагональными гауссианами, а entropy — энтропия гауссианы $\tfrac12(1 + \ln 2\pi) + \tfrac12\ln\sigma^2$ на одну размерность. compute_loss (ячейка 29) берёт знаки правильно: она возвращает $\big(T \cdot \mathrm{KL} + H_q(X_1 \mid X_0) - H_q(X_T \mid X_0) + H_p(X_T) - H_{\mathcal N(0,1)}\big)/\ln 2$, усреднённое по батчу и двум координатам, — оценку $-K$ по одному шагу минус энтропия стандартной гауссианы, в битах на размерность, то есть $-(K - L_{\text{null}})$ у Sohl-Dickstein. Множитель $T$ масштабирует один случайный шаг до суммы, в которой в оценке $T - 1$ членов ($t = 2, \dots, T$) и нет шага кода t = 0.
• Ячейки 30–33: обучение на случайных шагах. Ячейка 30 объясняет зеркальные шаги на примере $t = 5 \leftrightarrow 94$ в индексах кода ($t = 6 \leftrightarrow 95$ в формулах): loss_likelihood_bound выбирает 65 шагов, добавляет к ним парные n_steps - t - 1 и оставляет 128. Цикл в ячейке 33 использует Adam с шагом обучения $10^{-3}$, батчи по 128 (79 на один проход по 10 000 точек), 5001 проход и обрезает норму градиента на $1.0$.
• Ячейки 34–36: DDPM. ConditionalModelDDPM (три слоя ConditionalLinear) предсказывает шум, а p_sample_ddpm реализует описанный выше шаг сэмплирования с $\sigma_t = \sqrt{\beta_t}$. Эта функция добавляет новый шум на каждом шаге, включая последний, где алгоритм 2 шума не добавляет. В тексте ячейки 34 сказано, что обычно предпочитают $\tilde\beta_t$ за лучшие результаты; Ho et al. получили с обоими вариантами похожие результаты (§3.2).
• Ячейки 38–41: вариационная функция потерь с дискретизированным декодером на первом шаге. Ячейка 38 даёт KL-дивергенцию между гауссианами с полными ковариационными матрицами. discretized_gaussian_log_likelihood (ячейка 39) реализует формулу 13, приближая $\Phi$ функцией $\tfrac12\big(1 + \tanh\big(\sqrt{2/\pi}\,(x + 0.044715\,x^3)\big)\big)$ (наибольшая ошибка $1.8 \times 10^{-4}$) и применяя открытые крайние интервалы при $x < -0.999$ и $x > 0.999$. loss_variational (ячейка 41) усредняет KL-дивергенцию и член декодера по размерностям и делит на $\ln 2$, получая биты на размерность, а член декодера для шага кода t = 0 выбирает через torch.where. Эта функция потерь приведена для иллюстрации и никогда не обучается. В нынешнем виде она вызывает p_mean_variance с глобальной model — сетью из ячейки 33, предсказывающей среднее и дисперсию, — поскольку модель DDPM возвращает только оценку шума. Кроме того, декодер предполагает 8-битные пиксели в $[-1, 1]$, а «швейцарский рулет» непрерывен, и примерно у 31% его точек есть координата вне $[-1, 1]$, где правило для крайних интервалов приписывает им целый хвост распределения.
• Ячейки 42–43: noise_estimation_loss, невзвешенная $L_{\text{simple}}$, (noise - predicted_noise).square().mean(), с теми же зеркальными шагами. В тексте ячейки 42 $L_{\text{simple}}$ названа прямой минимизацией KL-членов; как показано выше, она их перевзвешивает.
• Ячейки 44–48: скользящее среднее весов и цикл обучения DDPM (Adam $10^{-3}$, батч 128, 1000 проходов, норма градиента обрезается на $1.0$). По умолчанию в классе EMA $\mu = 0.999$, но цикл использует $\mu = 0.9$ — окно примерно в 10 обновлений, короче одного прохода по данным. apply_shadow скопировал бы $\bar\theta$ в модель, но нигде не вызывается, так что сэмплирование идёт с необработанными весами. Ячейки с графиками (33 и 48) показывают x_seq[10*j] — обратную цепочку после $10j$ шагов (от $x_{90}$ до итогового сэмпла), но подписывают панели как $q(x_{100j})$: неверный масштаб, неверное направление и $q$ вместо $p_\theta$.

Ноутбук с домашним заданием 04_hw_diffusion из того же репозитория (нумерация ячеек та же) обучает DDPM на фотографиях лиц CelebA размером $32 \times 32$. Большая часть его кода — шаблон с пропусками для заполнения, но в нём сохранены выходы ячеек эталонного запуска на GPU Tesla T4; числа ниже взяты из этих выходов.

• Ячейки 9–11: данные. Батчи по 128 изображений, нормированных в $[-1, 1]$, как выше. Ячейка 10 печатает форму изображения $3 \times 64 \times 64$, хотя image_size = 32, а сетки в ячейках 11 и 26 показывают изображения $32 \times 32$: сохранённые выходы получены в разных запусках.
• Ячейки 13–14 повторяют прямой процесс, среднее через предсказание шума, два варианта $\sigma_t^2 = \beta_t$ или $\tilde\beta_t$ и алгоритмы 1 и 2 и называют в качестве сети U-Net или Efficient U-Net из Imagen. Efficient U-Net была разработана для ступеней сверхразрешения (super-resolution) в Imagen: она переносит параметры в блоки низкого разрешения, масштабирует обходные связи на $1/\sqrt2$ и меняет порядок операций понижения и повышения разрешения, чтобы ускорить прямой проход (Saharia et al., 2022, прил. B.1).
• Ячейка 16: класс DDPM хранит расписание (n_steps, min_beta, max_beta, оставлены пустыми), forward получает $x_t$ из $x_0$ за один шаг, а backward вызывает сеть, чтобы получить $\varepsilon_\theta(x_t, t)$.
• Ячейки 18–19. show_forward зашумляет батч до 25, 50, 75 и 100% от $T$. generate_new_images выполняет алгоритм 2, начиная с $x_T \sim \mathcal N(0, I)$, добавляет шум, только пока шаг кода t > 0, то есть не на последнем шаге, как и предписывает алгоритм 2, и записывает цепочку в GIF.
• Ячейки 20–21: U-Net. Текстовая ячейка перечисляет варианты подачи $t$ (AdaIN, как в StyleGAN, дополнительный канал, синусоидальное кодирование) и выбирает синусоидальное, добавляя, что так же можно прибавлять спроецированный эмбеддинг текста (здесь не используется). sinusoidal_embedding(n, d) заполняет замороженную таблицу nn.Embedding значениями $d = 100$ синусов и косинусов для каждого из $n = 1000$ шагов. _make_te — MLP вида Linear–SiLU–Linear, который переводит код в число каналов блока; forward прибавляет выход te1 к входному изображению, а блоки, оставленные пустыми, могут делать то же со своими отображениями _make_te.
• Ячейки 22–23: у эталонной модели 650 430 параметров; в это число входит замороженная таблица эмбеддингов.
• Ячейки 26–27: сетки прямого процесса на 25–100% от $T$ и 16 сэмплов необученной сети, которые представляют собой шум.
• Ячейки 29–31: обучение. training_loop реализует алгоритм 1 с nn.MSELoss, то есть $L_{\text{simple}}$, и сохраняет чекпоинт с наименьшей за эпоху функцией потерь на обучении. В сохранённом журнале 1272 батча по 128 на эпоху, что согласуется со 162 770 изображениями обучающей части CelebA. Функция потерь упала с $0.037$ после эпохи 1 до $0.033$ после эпохи 16; планировалось 200 эпох, а журнал обрывается в начале эпохи 17.
• Ячейки 33–36 загружают лучший чекпоинт и генерируют 100 лиц в GIF; сохранённых сэмплов нет.

Проблемы:

• show_images (ячейка 7) отображает $x$ в $255\,(x + 1)/2$ и приводит к uint8 без обрезки. Зашумлённые изображения выходят за $[-1, 1]$, поэтому сетки прямого процесса в ячейке 26 уже на 25% от $T$ показывают насыщенные, похожие на шум цвета; перед приведением типа нужно обрезать значения до $[-1, 1]$. Сколько осталось бы от лица, зависит от расписания, которое шаблон оставляет пустым; при линейном расписании DDPM $\bar\alpha_{250} = 0.52$, так что на 25% остаётся больше половины дисперсии сигнала.
• Класс набора данных в ячейке 10 возвращает изображения без меток, так что батч из загрузчика — тензор изображений, и batch[0] в show_first_batch, show_forward и training_loop выбирает первое изображение вместо батча. Шаблон взят от набора данных, который возвращает пары (изображение, метка); с этим набором нужно использовать сам батч.
• Ячейка 33 заново создаёт модель как DDPM(UNet(), n_steps=n_steps, device=device): min_beta и max_beta не передаются, поэтому используются значения по умолчанию из класса, а UNet() берёт свои 1000 шагов по умолчанию. Расписание, отличное от значений по умолчанию, придётся снова передать вручную.
• sinusoidal_embedding вычисляет $d$ частот $w_k = 10000^{-2k/d}$, $k = 0, \dots, d - 1$, но заполняет только $d$ значений. Кодирование из трансформера использует первые $d/2$ частот, каждую один раз в синусе и один раз в косинусе. Если же брать косинусы от второй половины частот, получаются частоты не больше $10^{-4}$, для которых $\cos(w_k t) \ge 0.995$ при всех $t \le 1000$: эти компоненты почти не меняются.

Дальше#

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

Литература#

• J. Ho, A. Jain, P. Abbeel. Denoising diffusion probabilistic models. NeurIPS, 2020.
• J. Sohl-Dickstein, E. A. Weiss, N. Maheswaranathan, S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. ICML, 2015.
• W. Feller. On the theory of stochastic processes, with particular reference to applications. Proceedings of the First Berkeley Symposium on Mathematical Statistics and Probability, 1949.
• J. Song, C. Meng, S. Ermon. Denoising diffusion implicit models. ICLR, 2021. Код: ermongroup/ddim (зеркальные шаги в runners/diffusion.py).
• A. Q. Nichol, P. Dhariwal. Improved denoising diffusion probabilistic models. ICML, 2021.
• C. Luo. Understanding diffusion models: A unified perspective. arXiv:2208.11970, 2022.
• P. Dhariwal, A. Nichol. Diffusion models beat GANs on image synthesis. NeurIPS, 2021. §3 (U-Net из DDPM) и §3.1, таблица 3 (адаптивная групповая нормализация).
• Y. Song, S. Ermon. Generative modeling by estimating gradients of the data distribution. NeurIPS, 2019. §4.1 и приложение A (архитектура сети score).
• O. Ronneberger, P. Fischer, T. Brox. U-Net: Convolutional networks for biomedical image segmentation. MICCAI, 2015.
• Y. Wu, K. He. Group normalization. ECCV, 2018.
• T. Karras, S. Laine, T. Aila. A style-based generator architecture for generative adversarial networks. CVPR, 2019. §2, формула 1 (AdaIN).
• C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. Denton, S. K. S. Ghasemipour, B. K. Ayan, S. S. Mahdavi, R. G. Lopes, T. Salimans, J. Ho, D. J. Fleet, M. Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. NeurIPS, 2022. Imagen: §2.5 (подача условия) и приложение B.1 (Efficient U-Net).
• HSE-LAMBDA, репозиторий курса Deep Generative Models: 02_DPM_Models.ipynb и 04_hw_diffusion.ipynb.