Score matching и динамика Ланжевена

Заметка о DDPM пришла к модели, оценив сверху правдоподобие цепочки зашумления и упрощая оценку, пока она не стала регрессией на шум. Эта заметка приходит к той же модели с другой стороны. Вместо плотности учится её градиент — score (функция score); вместо декодера сэмплы получаются движением по score с небольшим добавленным шумом — динамикой Ланжевена (Langevin dynamics). Два пути встречаются: сеть, предсказывающая шум в DDPM, с точностью до известного множителя и есть score зашумлённых данных.

Обозначения. $p(x)$ — плотность данных на $\mathbb R^d$, $s(x) = \nabla_x \log p(x)$ — её score; $s_\theta(x)$ — модель. $q_\sigma(\tilde x) = \int p(x)\,\mathcal N(\tilde x; x, \sigma^2 I)\,dx$ — данные после добавления гауссова шума масштаба $\sigma$. $\nabla \cdot s$ — след якобиана $s$.

Зачем нужен score#

Многие модели задают плотность лишь с точностью до константы, как модели, основанные на энергии (energy-based models), из одной из предыдущих заметок: $p_\theta(x) = e^{-E_\theta(x)} / Z_\theta$ с энергией $E_\theta$ и нормировкой $Z_\theta = \int e^{-E_\theta(x)}\,dx$, которую в большой размерности не вычислить. Методу максимального правдоподобия нужны $\log Z_\theta$ и её градиент. Score — нет:

$$\nabla_x \log p_\theta(x) = -\nabla_x E_\theta(x) - \underbrace{\nabla_x \log Z_\theta}_{=\,0} .$$

$Z_\theta$ зависит от $\theta$, но не от $x$, поэтому исчезает. Score гауссианы $\mathcal N(\mu, \sigma^2 I)$ равен $-(x - \mu)/\sigma^2$: вектор, указывающий обратно к среднему и тем длиннее, чем дальше точка. В общем случае score указывает в сторону большей плотности и несёт всю информацию о $p$ там, где $p$ положительна и гладка на связной области, например на всём $\mathbb R^d$: две такие плотности с одинаковым score всюду отличаются постоянным множителем, который фиксирует нормировка. На двух разделённых компонентах score задаёт каждую часть лишь с точностью до своей константы и ничего не говорит о том, как масса делится между ними; гауссовское сглаживание снимает проблему, потому что зашумлённая плотность положительна на всём $\mathbb R^d$. Поэтому, когда score параметризован через энергию, $s_\theta = -\nabla_x E_\theta$, его подгонка обучает модель, основанную на энергии, без её нормировочной константы и без сэмплов из модели (модели, основанные на энергии). Произвольная сеть $s_\theta: \mathbb R^d \to \mathbb R^d$ — другой объект: векторное поле, якобиан которого не обязан быть симметричным, так что оно не обязано быть градиентом какой-либо скалярной функции. Это приближённый score, за которым не стоит никакая энергия, и именно так его используют диффузионные модели.

Подгонка score: трюк Хюваринена#

Естественная функция потерь — дивергенция Фишера (Fisher divergence), средний квадрат расстояния между двумя score по данным:

$$J(\theta) = \tfrac12\, \mathbb E_{p}\big\lVert s_\theta(x) - \nabla_x \log p(x)\big\rVert^2 .$$

Ей нужен истинный score, который неизвестен. Hyvärinen (2005) убирает его интегрированием по частям. В одном измерении раскроем квадрат; перекрёстное слагаемое равно

$$\mathbb E_p\big[s_\theta(x)\, (\log p)'(x)\big] = \int s_\theta(x)\, p'(x)\,dx = \big[s_\theta\, p\big]_{-\infty}^{\infty} - \int s_\theta'(x)\, p(x)\,dx = -\mathbb E_p\big[s_\theta'(x)\big],$$

где использовано $p\,(\log p)' = p'$ и предположено, что $s_\theta(x)\,p(x) \to 0$ на бесконечности. Квадрат истинного score от $\theta$ не зависит. В $d$ измерениях тот же шаг по каждой координате даёт

$$J(\theta) = \mathbb E_p\Big[\nabla \cdot s_\theta(x) + \tfrac12 \lVert s_\theta(x)\rVert^2\Big] + \text{const},$$

математическое ожидание по данным, в которое входит только модель.

Разобранный пример. Подгоним гауссов score $s_\theta(x) = -a(x - \mu)$ с $a = 1/\sigma^2$ к данным $\{1, 2, 4, 5\}$. Тогда $s_\theta' = -a$, а целевая функция — $\mathbb E\big[\tfrac12 a^2 (x - \mu)^2 - a\big]$. Приравнивание производной по $\mu$ к нулю даёт $\mu = \bar x = 3$; производная по $a$ даёт $a\,\mathbb E(x - \mu)^2 = 1$, так что $\sigma^2 = \frac14(4 + 1 + 1 + 4) = 2.5$. Score matching восстанавливает выборочные среднее и дисперсию — оценки максимального правдоподобия для гауссианы, — ни разу не коснувшись нормировочной константы.

Цена — слагаемое со следом. Для сети $s_\theta: \mathbb R^d \to \mathbb R^d$ величина $\nabla \cdot s_\theta$ требует $d$ обратных проходов, по одному на диагональный элемент якобиана: годится для $d = 2$, невозможно для изображения с $d = 3 \cdot 64 \cdot 64$.

Две целевые функции подешевле#

Срезанный score matching (sliced score matching, Song et al., 2019) проецирует на случайные направления $v$ с $\mathbb E[v v^\top] = I$:

$$J_{\text{SSM}}(\theta) = \mathbb E_{p}\,\mathbb E_{v}\Big[v^\top \nabla s_\theta(x)\, v + \tfrac12 \big(v^\top s_\theta(x)\big)^2\Big] .$$

В среднем по $v$ это целевая функция выше, потому что $\mathbb E[v^\top A v] = \mathrm{tr}(A)$ и $\mathbb E[(v^\top s)^2] = \lVert s\rVert^2$. Одно произведение якобиана на вектор заменяет $d$ обратных проходов.

Денойзинговый score matching (denoising score matching, Vincent, 2011) обходится в функции потерь без производных модели по входу — без якобиана $\nabla_x s_\theta$ и без следа, — подгоняя score зашумлённых данных; градиент по параметрам $\nabla_\theta$ для обучения по-прежнему нужен, как для любой функции потерь. Возьмём $x \sim p$ и возмутим его любым ядром $q(\tilde x \mid x)$ с известным score; зашумлённые данные имеют плотность $q(\tilde x) = \int q(\tilde x \mid x)\,p(x)\,dx$. Целевая функция

$$J_{\text{DSM}}(\theta) = \tfrac12\, \mathbb E_{x \sim p,\ \tilde x \sim q(\tilde x \mid x)}\big\lVert s_\theta(\tilde x) - \nabla_{\tilde x} \log q(\tilde x \mid x)\big\rVert^2$$

отличается от дивергенции Фишера относительно $q(\tilde x)$ лишь на константу (Song and Ermon, 2019, формула 2). Раскроем оба квадрата. Слагаемые $\lVert s_\theta\rVert^2$ имеют одинаковое математическое ожидание, слагаемые без $\theta$ — константы, а перекрёстные слагаемые совпадают, потому что $q(\tilde x)\,\nabla \log q(\tilde x) = \nabla q(\tilde x) = \int \nabla_{\tilde x} q(\tilde x \mid x)\,p(x)\,dx$, так что

$$\mathbb E_{q(\tilde x)}\big[s_\theta^\top \nabla \log q(\tilde x)\big] = \mathbb E_{p(x)\,q(\tilde x \mid x)}\big[s_\theta^\top \nabla_{\tilde x} \log q(\tilde x \mid x)\big] .$$

Обычно ядро гауссовское, $\tilde x = x + \sigma\varepsilon$ с $\varepsilon \sim \mathcal N(0, I)$, что даёт плотность $q_\sigma$ из обозначений. Его score равен $\nabla_{\tilde x} \log \mathcal N(\tilde x; x, \sigma^2 I) = -(\tilde x - x)/\sigma^2 = -\varepsilon/\sigma$, и целевая функция принимает вид

$$J_{\text{DSM}}(\theta) = \tfrac12\, \mathbb E_{x,\,\varepsilon}\Big\lVert s_\theta(x + \sigma\varepsilon) + \frac{\varepsilon}{\sigma}\Big\rVert^2 .$$

Почему это работает, можно увидеть и по-другому: регрессия с квадратичной ошибкой учит условное среднее, минимизатор — $s^*(\tilde x) = \mathbb E[-\varepsilon/\sigma \mid \tilde x]$, и это условное среднее в точности равно $\nabla \log q_\sigma(\tilde x)$ — факт, известный как формула Твиди (Tweedie’s formula).

Разобранный пример. Данные $x \sim \mathcal N(0, 1)$ и $\sigma = 0.5$. Тогда $\tilde x \sim \mathcal N(0, 1.25)$, а его score равен $-\tilde x / 1.25 = -0.8\,\tilde x$. Напрямую: $\varepsilon$ и $\tilde x$ совместно гауссовы с $\mathrm{Cov}(\varepsilon, \tilde x) = \sigma = 0.5$, так что $\mathbb E[\varepsilon \mid \tilde x] = 0.5\,\tilde x / 1.25 = 0.4\,\tilde x$ и $\mathbb E[-\varepsilon/\sigma \mid \tilde x] = -0.8\,\tilde x$: цель регрессии в среднем — score зашумлённых данных.

Цена денойзингового score matching в том, что он оценивает score $q_\sigma$, а не $p$. Малое $\sigma$ держит их близкими, но делает цели $\varepsilon/\sigma$ большими и шумными; большое $\sigma$ размывает данные. Выход, о котором ниже, — использовать сразу много значений $\sigma$.

Сэмплирование динамикой Ланжевена#

Имея score, динамика Ланжевена порождает сэмплы:

$$x_{k+1} = x_k + \frac{\eta}{2}\, s(x_k) + \sqrt{\eta}\, z_k, \qquad z_k \sim \mathcal N(0, I) .$$

Снос поднимается по плотности, шум не даёт цепочке схлопнуться в моды. При малом шаге $\eta$ и большом числе шагов распределение $x_k$ приближается к $p$; при условиях регулярности оно становится в точности равным $p$, когда $\eta \to 0$, а число шагов неограниченно растёт (Song and Ermon, 2019, §2.2, со ссылкой на Welling and Teh, 2011). Без шума обновление — градиентный подъём по $\log p$, и он заканчивается в моде, а не в сэмпле.

Разобранный пример. Для $p = \mathcal N(0, 1)$ score $s(x) = -x$, и обновление — $x_{k+1} = (1 - \eta/2)\, x_k + \sqrt\eta\, z_k$. Его стационарная дисперсия $v$ удовлетворяет $v = (1 - \eta/2)^2 v + \eta$, так что

$$v = \frac{\eta}{1 - (1 - \eta/2)^2} = \frac{1}{1 - \eta/4} .$$

При $\eta = 0.1$ цепочка устанавливается на дисперсии 1.026 вместо 1, при $\eta = 0.5$ — на 1.143. Смещение от дискретизации имеет порядок $\eta$, и меньшие шаги покупают точность ценой большего числа шагов; поправка Метрополиса–Гастингса убрала бы смещение, но потребовала бы вычислять плотность, а ведь score нужен был именно затем, чтобы этого избежать.

Уменьшение шага по ходу цепочки. Смещение исчезает и тогда, когда размер шага убывает вдоль цепочки, но только если он убывает медленно. Welling and Teh (2011), формула 2, требуют

$$\sum_k \eta_k = \infty, \qquad \sum_k \eta_k^2 < \infty,$$

например, $\eta_k \propto (b + k)^{-\gamma}$ с $\gamma \in (0.5, 1]$. Первое условие позволяет цепочке уйти от старта на любое расстояние; второе заставляет шаги, а с ними и ошибку дискретизации, стремиться к нулю. Геометрическое убывание $\eta_k = \eta_0 \rho^k$ нарушает первое условие: сумма шагов равна $\eta_0/(1 - \rho)$, и цепочка замирает, не успев перемешаться. В гауссовском примере начальная точка сохраняется в $x_K$ с множителем $\prod_k (1 - \eta_k/2)$. При $\eta_0 = 0.01$ и $\rho = 0.9$ сумма шагов не больше $0.1$, и после десяти шагов множитель всё ещё равен $0.968$: цепочка едва сдвинулась со старта.

Температура. Некоторые сэмплеры масштабируют шум температурой $\tau$, $x_{k+1} = x_k + \frac\eta2\, s(x_k) + \tau\sqrt\eta\, z_k$. Поскольку $\frac\eta2 \nabla \log p = \frac{\tau^2\eta}{2} \nabla \log p^{1/\tau^2}$, это обычный сэмплер для плотности, пропорциональной $p^{1/\tau^2}$, с шагом $\tau^2\eta$. Поэтому температура ниже $1$ концентрирует сэмплы около мод. В гауссовском примере стационарная дисперсия становится равной $v = \tau^2/(1 - \eta/4)$: при $\tau = 0.5$ и $\eta = 0.1$ это $0.256$, близко к дисперсии $0.25$ распределения $p^4 \propto \mathcal N(0, 1/4)$.

Почему одного уровня шума мало#

Две проблемы появляются, как только данные сложнее одного сгустка (Song and Ermon, 2019).

Score неизвестен там, где нет данных. Score matching взвешивает ошибку плотностью $p(x)$, так что в областях низкой плотности выученный score почти не обучен, и цепочка, стартовавшая там, — а там стартует любая цепочка из шума, — следует произвольным стрелкам. Когда данные лежат около маломерного многообразия, как изображения, почти всё пространство — такая область, а score самих данных вне многообразия даже не определён.

Динамика Ланжевена не умеет балансировать разделённые моды. Возьмём $p = 0.8\,\mathcal N(-4, 1) + 0.2\,\mathcal N(4, 1)$. Около каждой моды score — это score одной гауссианы, и от весов он не зависит; веса проявляются только в промежутке, где score в $x = 0$ равен $4(1 - 2 \cdot 0.8) = -2.4$, но плотность там около $0.00013$, и цепочки туда редко заходят. Цепочка падает в ближайшую к старту моду и остаётся в ней. С точным score 4000 цепочек, стартовавших равномерно на $[-8, 8]$, после 2000 шагов размера 0.05 оказываются в левой моде в доле 0.52 — близко к доле стартовавших слева, а не к истинным 0.8.

Шум помогает с обеими проблемами. Зашумлённая плотность $q_\sigma$ всюду положительна, так что её score определён всюду, а при большом $\sigma$ зашумлённые обучающие точки разбросаны по гораздо более широкой области, и сеть обучается там, где начинаются цепочки из шума; обучается она, впрочем, только там, куда действительно попадают зашумлённые сэмплы. При большом $\sigma$ плотность $q_\sigma$ к тому же заполняет промежуток между модами, так что веса становятся видимыми для цепочки. Насколько это помогает, зависит от выбора: уровни шума должны перекрываться, каждому нужно достаточно шагов подходящего размера, а обученный score добавляет собственную ошибку аппроксимации. В примере со смесью ниже используется точный score.

Рисунок

Шум заполняет пустоту между модами

Score, зависящий от уровня шума, и отжиг#

Song and Ermon (2019) обучают одну сеть сразу для многих уровней шума. Выбирается геометрическая последовательность $\sigma_1 > \sigma_2 > \dots > \sigma_L$, сеть получает уровень на вход (noise-conditional score network, NCSN), и минимизируется взвешенная сумма денойзинговых целевых функций:

$$\mathcal L(\theta) = \sum_{i=1}^{L} \sigma_i^2\; \tfrac12\, \mathbb E_{x,\,\varepsilon}\Big\lVert s_\theta(x + \sigma_i\varepsilon, \sigma_i) + \frac{\varepsilon}{\sigma_i}\Big\rVert^2 .$$

Вес $\sigma_i^2$ делает слагаемые сравнимыми: цель $\varepsilon/\sigma_i$ имеет размер $1/\sigma_i$, так что каждое взвешенное слагаемое порядка единицы. Сэмплирование, отжиг динамики Ланжевена (annealed Langevin dynamics), делает несколько шагов Ланжевена на самом большом уровне шума, где цепочка движется свободно, затем продолжает со следующего уровня с того места, где остановилась, и так до самого малого, с размерами шага $\alpha_i = \epsilon\, \sigma_i^2/\sigma_L^2$, уменьшающимися вместе с шумом. Для изображений в статье $L = 10$ уровней от $\sigma_1 = 1$ до $\sigma_{10} = 0.01$, 100 шагов на уровень и $\epsilon = 2 \cdot 10^{-5}$.

Обучение держится на двух деталях. Сеть видит уровень. В NCSN Song and Ermon у каждого слоя нормализации есть обучаемые масштаб и сдвиг для каждого уровня — условная нормализация по экземпляру (conditional instance normalization; прил. A.1). В небольшой полносвязной сети уровень может просто выбирать обучаемый вектор, который прибавляется к каждому скрытому слою, $h \mapsto W h + b + e_i$. Диффузионные сети получают свой шаг теми же способами (заметка о DDPM). Один уровень на пример. Сумма по уровням состоит из $L$ слагаемых, и $\sum_i \lambda_i \ell_i(\theta) = L\, \mathbb E_{i \sim \mathrm{Uniform}\{1, \dots, L\}}\big[\lambda_i \ell_i(\theta)\big]$ для взвешенных функций потерь уровней $\lambda_i \ell_i$. Поэтому среднее выбранных функций потерь по уровням, по одному случайному уровню на обучающий пример, — несмещённая оценка $\mathcal L(\theta)/L$, то есть целевой функции с точностью до постоянного множителя $L$, который масштабирует градиент, но не сдвигает минимум (чтобы оценить сам $\mathcal L$, нужно умножить на $L$), — так же, как случайный шаг в DDPM. Например, четыре геометрических уровня от $1$ до $0.01$ — это $\sigma_i = 100^{-(i-1)/3}$: $1$, $0.215$, $0.046$ и $0.01$. Рисунок ниже показывает поля, которые такой сети нужно выучить на этих четырёх уровнях.

Рисунок

Поле score на четырёх уровнях шума

В примере со смесью отжиг по десяти уровням от $\sigma = 5$ до $\sigma = 0.1$, по 100 шагов на уровне, с точным score каждого $q_\sigma$ заканчивается долей 0.79 цепочек в левой моде против истинных 0.8.

Рисунок

Без шума Ланжевен теряет веса мод

Та же модель, что DDPM#

DDPM обучает сеть $\varepsilon_\theta(x_t, t)$ предсказывать шум в $x_t = \sqrt{\bar\alpha_t}\, x_0 + \sqrt{1 - \bar\alpha_t}\,\varepsilon$ (прямой процесс). Score этого ядра возмущения равен

$$\nabla_{x_t} \log q(x_t \mid x_0) = -\frac{x_t - \sqrt{\bar\alpha_t}\, x_0}{1 - \bar\alpha_t} = -\frac{\varepsilon}{\sqrt{1 - \bar\alpha_t}} ,$$

так что денойзинговый score matching на уровне $t$ регрессирует $s_\theta(x_t, t)$ на $-\varepsilon/\sqrt{1 - \bar\alpha_t}$. Запись $s_\theta(x_t, t) = -\varepsilon_\theta(x_t, t)/\sqrt{1 - \bar\alpha_t}$ превращает функцию потерь DSM в $\lVert\varepsilon_\theta - \varepsilon\rVert^2 / (1 - \bar\alpha_t)$, а вес $1 - \bar\alpha_t$, дисперсия шума, — это взвешивание $\sigma^2$ выше. Целевая функция DDPM $L_{\text{simple}}$ — денойзинговый score matching, зависящий от уровня шума, с дисперсией шума в роли веса, а шаг сэмплирования DDPM — шаг в духе Ланжевена с выученным score. Song et al. (2021) делают это соответствие точным в непрерывном времени: обе модели дискретизируют стохастическое дифференциальное уравнение, добавляющее шум, а генерация проходит его обращение во времени, в снос которого входит score $\nabla_x \log p_t(x)$.

Что изменилось#

	Правдоподобие (VAE, потоки)	Score matching
Что учится	плотность или её оценка	градиент логарифма плотности
Нормировочная константа	нужна или обходится архитектурой	не нужна никогда
Обучающий сигнал	логарифм правдоподобия или ELBO	дивергенция Фишера через след, случайные проекции или денойзинг
Сэмплирование	один проход декодера	много шагов Ланжевена
Слабое место	архитектура ограничивает модель	области без данных и разделённые моды, смягчается уровнями шума

Ограничения#

• Много шагов сэмплирования. Отжиг Ланжевена делает от сотен до тысяч вычислений сети на сэмпл; более быстрые сэмплеры — предмет дальнейших работ.
• Смещение от дискретизации. Нескорректированная динамика Ланжевена с конечным шагом сэмплирует немного не то распределение, как показывает дисперсия 1.026 вместо 1. Убывающий шаг убирает смещение, только если сумма шагов по-прежнему бесконечна.
• Расписание шума — проектное решение. Самый большой уровень должен перекрывать расстояния между модами, самый малый — быть ниже важных деталей, а соседние уровни — перекрываться.
• Правдоподобия бесплатно нет. Модель score не даёт $\log p(x)$ напрямую; его можно вычислить через ОДУ потока вероятности из Song et al. (2021) ценой её решения.

Итоги#

• Score $\nabla_x \log p(x)$ не зависит от нормировочной константы и там, где $p$ положительна на связной области, определяет $p$ с точностью до неё.
• Интегрирование по частям Хюваринена превращает дивергенцию Фишера в $\mathbb E_p[\nabla\cdot s_\theta + \frac12 \lVert s_\theta\rVert^2]$; срезанный score matching оценивает след случайными проекциями; денойзинговый score matching регрессирует на score ядра возмущения, $-\varepsilon/\sigma$ для гауссовского шума, и учит score зашумлённых данных.
• Динамика Ланжевена сэмплирует с помощью score; размер шага меняет смещение на число шагов, сумма убывающих шагов должна быть бесконечной, а температура $\tau$ даёт сэмплы из $p^{1/\tau^2}$ вместо $p$.
• Один уровень шума не справляется в пустых областях и между модами; score, зависящий от уровня шума, вместе с отжигом смягчает и то, и другое, если уровни и размеры шагов подходят к данным. Сеть получает уровень в каждом слое и обучается с одним случайным уровнем на пример.
• Предсказание шума в DDPM — это денойзинговый score matching, зависящий от уровня шума: $s_\theta = -\varepsilon_\theta / \sqrt{1 - \bar\alpha_t}$.

Ноутбук#

Ноутбук о score matching из репозитория HSE-LAMBDA Deep Generative Models работает с двумерным «швейцарским рулетом» (Swiss roll) и основан на руководстве по диффузионным моделям ACIDS-IRCAM (ячейка 1; номера ячеек считают все ячейки, включая текстовые, начиная с 1). Сохранённых выходов ячеек в нём нет, поэтому он описывает процедуру, а не результаты.

• Ячейка 5: 10 000 точек швейцарского рулета с шумом 0.5; первая и третья координаты делятся на 10.
• Ячейки 7–12: стандартный score matching с двухслойной сетью на Softplus, точный якобиан, собранный одним обратным проходом на каждую выходную координату, и 2000 шагов Adam по полному батчу.
• Ячейки 13–14: plot_gradients рисует score стрелками на сетке $50 \times 50$ над $[-1.5, 2]^2$, перемасштабируя каждую до $v \log(1 + \lVert v\rVert)/\lVert v\rVert$ — это правило длины рисунка с полем score, — чтобы огромные значения score вдали от данных оставались читаемыми. Ячейки 19, 23, 27 и 32–33 используют тот же график.
• Ячейки 15–19: градиентный подъём по score (20 шагов по $10^{-3}$) против динамики Ланжевена (10 шагов, начиная с $\epsilon = 10^{-2}$ и умножая на 0.9 после каждого шага), оба из фиксированной точки $(1.5, -1.5)$; в тексте сказано, что старт берётся из $\mathcal N(0, I)$. Ячейка 17 ссылается на Welling and Teh (2011) как на источник сходимости динамики Ланжевена. У sample_langevin (ячейка 18) есть аргумент temperature, по умолчанию 1, который масштабирует шум. Геометрическое убывание шага в этой функции нарушает условие $\sum_k \eta_k = \infty$: десять шагов в сумме дают $0.065$, так что цепочка едва сдвигается со старта.
• Ячейки 20–23: срезанный score matching на первых 1000 точках. Случайные направления нормируются до единичной длины, так что $\mathbb E[v v^\top] = I/d$, и оба слагаемых усредняются по $d$ координатам: слагаемое с якобианом становится $v^\top \nabla s_\theta\, v / d$ с математическим ожиданием $\mathrm{tr}(\nabla s_\theta)/d^2$, а второе слагаемое — средним $(v_i\, s_i)^2/2$ вместо $(v^\top s)^2/2$ с математическим ожиданием $\lVert s_\theta\rVert^2/(2d^2)$. Оба несут одинаковый множитель $1/d^2$, так что целевая функция пропорциональна целевой функции Хюваринена и имеет тот же минимизатор.
• Ячейки 24–27: денойзинговый score matching. Текстовая ячейка 24 даёт общую целевую функцию для любого ядра возмущения и лишь затем гауссовскую; код использует $\sigma = 0.01$ и 5000 шагов. При таком уровне шума цели имеют размер $1/\sigma = 100$, а точки вне рулета почти не получают обучающего сигнала — это и есть сбой, который мотивирует следующие ячейки.
• Ячейки 28–33: сеть, зависящая от уровня шума, с четырьмя уровнями $1$, $0.215$, $0.046$ и $0.01$ и функцией потерь с весом $\sigma^2$ (anneal_power = 2). ConditionalLinear прибавляет обучаемый вектор уровня — nn.Embedding(4, 128), инициализированный равномерно на $[-0.1, 0.1]$, — к выходу каждого из двух скрытых линейных слоёв; это аддитивный вариант, описанный выше. Каждый из 5000 шагов по полному батчу выбирает один уровень на точку через torch.randint. Ячейка 32 рисует выученное поле score на каждом уровне. Отжиг динамики Ланжевена не реализован, а последний график приписывает каждой точке сетки случайный уровень шума, так что его стрелки смешивают четыре разных поля.

Дальше#

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

Литература#

• A. Hyvärinen. Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 6, 2005.
• P. Vincent. A Connection Between Score Matching and Denoising Autoencoders. Neural Computation, 23(7), 2011.
• Y. Song, S. Garg, J. Shi, S. Ermon. Sliced Score Matching: A Scalable Approach to Density and Score Estimation. UAI, 2019.
• Y. Song, S. Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. NeurIPS, 2019. Сети score, зависящие от уровня шума, и отжиг динамики Ланжевена.
• Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, B. Poole. Score-Based Generative Modeling through Stochastic Differential Equations. ICLR, 2021.
• B. Efron. Tweedie’s Formula and Selection Bias. Journal of the American Statistical Association, 106(496), 2011.
• M. Welling, Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. ICML, 2011. Формула 2 (условия на размер шага).
• G. O. Roberts, R. L. Tweedie. Exponential Convergence of Langevin Distributions and Their Discrete Approximations. Bernoulli, 2(4), 1996.