Марковские цепи и прямой процесс

Что мы собираемся построить

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

Прямой процесс задаётся руками, без обучения — и это его главное свойство. Он марковский:

$$\htmlData{k=markov}{q(x_t \mid x_{t-1})} = \mathcal{N}\big(x_t;\ \sqrt{1-\htmlData{k=beta}{\beta_t}}\, x_{t-1},\ \htmlData{k=beta}{\beta_t} I\big)$$

Марковость означает, что вся история сжата в текущее состояние. Это не упрощение ради удобства: именно из неё выйдет вся дальнейшая математика.

Обратите внимание на выбор коэффициентов: множитель у сигнала — $\sqrt{1-\beta_t}$, у шума — $\sqrt{\beta_t}$, и сумма их квадратов равна единице. Значит если $\operatorname{Var}(x_{t-1}) = 1$, то и $\operatorname{Var}(x_t) = 1$. Такой процесс называют сохраняющим дисперсию, и это не косметика: без нормировки масштаб данных уезжал бы вместе с шумом, а сеть должна работать в одном диапазоне на всех шагах.

Тысяча шагов в одну формулу

Наивно, чтобы получить $x_{100}$, нужно сделать сто шагов. Но композиция гауссиан — снова гауссиана, и произведение сворачивается.

Обозначим $\alpha_t = 1 - \beta_t$ и $\bar\alpha_t = \prod_{s\le t}\alpha_s$. Тогда

$$q(x_t \mid x_0) = \mathcal{N}\big(x_t;\ \sqrt{\htmlData{k=abar}{\bar\alpha_t}}\, x_0,\ (1 - \htmlData{k=abar}{\bar\alpha_t}) I\big), \qquad x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon$$

Проверено численно: пять шагов с $\beta = [0.1, 0.15, 0.2, 0.05, 0.3]$ из $x_0 = 2$ по четырёхсот тысячам траекторий дают среднее $1.2745$ и дисперсию $0.5919$; замкнутая форма предсказывает $1.2759$ и $0.5930$.

Это самое важное практическое следствие во всей теме. Обучение не требует прогона цепи: чтобы получить обучающий пример на шаге $t$, достаточно взять $x_0$, взять один $\varepsilon$ и посчитать одно выражение. Без этого диффузия была бы неприменима — тысяча последовательных шагов на каждый пример в батче.

Заметьте, что $\bar\alpha_t$ полностью описывает состояние процесса. Ни $t$, ни отдельные $\beta$ дальше не понадобятся — только одно число.

Что видит модель

Потяните время и посмотрите на облако и на отношение сигнал/шум.

линейноекосинусноеквадратичное

данные при t

сигнал/шум, log-шкала

время t

0.3

ᾱ

0.3956

√ᾱ

0.629

√(1−ᾱ)

0.7774

сигнал/шум

0.655

равенство при t

0.259

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

Два наблюдения, которые пригодятся дальше:

1. Облако сжимается к нулю и расширяется до единичной дисперсии. Множитель $\sqrt{\bar\alpha}$ стягивает данные к началу координат, $\sqrt{1-\bar\alpha}$ добавляет разброс. При $\bar\alpha \to 0$ остаётся $\mathcal{N}(0, I)$ независимо от того, чем были данные.
2. Время входит только через $\bar\alpha$. Поэтому «шаг 700 из 1000» — не информация; информация — это $\bar\alpha_{700}$, и именно её (или эквивалентное ей число) подают в сеть.

Обратный процесс: почему он вообще существует

Прямой процесс задан. Нужен обратный: $q(x_{t-1} \mid x_t)$. Он существует по правилу Байеса, но не выражается в замкнутой форме — для этого нужна маргинальная $q(x_t)$, то есть распределение всех данных, которого мы не знаем.

Здесь и появляется ключевой факт, из которого работает весь метод:

При малом $\beta_t$ обратный шаг $q(x_{t-1}\mid x_t)$ приближённо гауссов.

Это не тавтология. Прямой шаг гауссов по построению; обратный — вообще говоря нет. Но при малом шаге он становится гауссовым с точностью до $O(\beta_t^2)$, а значит его можно параметризовать сетью, предсказывающей всего два числа на компоненту: среднее и дисперсию.

Отсюда сразу видна цена и причина существования тысячи шагов. Приближение тем точнее, чем меньше $\beta$; меньше $\beta$ — больше шагов, чтобы дойти до шума. Тысяча шагов в DDPM — это не запас на всякий случай, а плата за гауссовость обратного ядра.

Обратное ядро при известном $x_0$

Одна вещь всё-таки считается точно, и на ней стоит обучение. Если знать $x_0$, то обратный шаг гауссов не приближённо, а ровно:

$q(x_{t-1} \mid x_t, x_0) = \mathcal{N}\big(x_{t-1};\ \tilde\mu_t(x_t, x_0),\ \tilde\beta_t I\big)$

$\tilde\mu_t = \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t}x_t + \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1-\bar\alpha_t}x_0, \qquad \tilde\beta_t = \frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t$

Формулы выглядят громоздко, а получаются механически: правило Байеса для гауссиан плюс приведение показателя к квадрату (блок 3, урок 090). Разбирать их наизусть не нужно; важно структурное наблюдение — среднее есть выпуклая комбинация $x_t$ и $x_0$, то есть обратный шаг тянет текущее состояние в сторону чистых данных, и вес этой тяги задан расписанием.

Именно поэтому в обучении фигурирует $x_0$ (или, что то же самое, $\varepsilon$): цель обучения — приблизить это точное ядро, а не угадывать неизвестную маргинальную плотность. Урок 060 покажет, что предсказывать $\varepsilon$, предсказывать $x_0$ и предсказывать score — три записи одной задачи.

Итог

• Прямой процесс задан руками, марковский, сохраняет дисперсию.
• Он сворачивается в одну формулу: $x_t = \sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\varepsilon$, и без этого обучение было бы непрактично.
• Всё состояние процесса — одно число $\bar\alpha_t$.
• Обратный шаг гауссов приближённо, с точностью $O(\beta^2)$, — отсюда требование малых шагов и их количество.
• Обратный шаг при известном $x_0$ гауссов точно, и на нём стоит обучение.