Шумоподавление — это и есть score

Тождество

Урок 040 закончился тупиком: обучать score можно, но след якобиана стоит $d$ проходов. Выход такой — если зашумление известно, score условного распределения выписывается явно.

Из урока 010: $x_t = \sqrt{\bar\alpha}\,x_0 + \sqrt{1-\bar\alpha}\,\varepsilon$, то есть $q(x_t\mid x_0) = \mathcal{N}(\sqrt{\bar\alpha}x_0,\ (1-\bar\alpha)I)$. Score гауссианы известен:

$$\nabla_{x_t}\log q(x_t \mid x_0) = -\frac{x_t - \sqrt{\bar\alpha}x_0}{1-\bar\alpha} = -\frac{\htmlData{k=eps}{\varepsilon}}{\sqrt{1-\bar\alpha}}$$

Второе равенство — подстановка определения $x_t$. Проверено численно: при $\bar\alpha = 0.6$, $x_0 = 1.3$, $\varepsilon = 0.7$ обе стороны дают $-1.1067971811$, разность в точности нуль.

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

Почему условный score годится вместо безусловного

Здесь легко пропустить главное. Нам нужен $\nabla\log q(x_t)$ — score маргинального распределения, а известен $\nabla\log q(x_t\mid x_0)$ — условного. Это разные вещи.

Тождество Винсента говорит, что минимизация по обеим целям даёт одно и то же решение:

$\mathbb{E}_{x_t}\big\|s_\theta(x_t) - \nabla\log q(x_t)\big\|^2 = \mathbb{E}_{x_0, x_t}\big\|s_\theta(x_t) - \nabla\log q(x_t\mid x_0)\big\|^2 + \text{const}$

Причина, если убрать выкладки: оптимум квадратичной регрессии — условное среднее (блок 3, урок 050). Обучаясь предсказывать условный score, сеть в оптимуме выдаёт его среднее по всем $x_0$, совместимым с данным $x_t$, — а это и есть маргинальный score.

Формулировка, которую стоит запомнить: сеть учится на величине, которую нельзя знать в момент генерации, и в оптимуме выдаёт величину, которую знать можно. Тот же приём, что в блоке 4 с ELBO, и он же — причина, по которой обучение диффузии выглядит подозрительно просто.

Цель обучения

Собираем всё вместе:

$\mathcal{L} = \mathbb{E}_{x_0,\, t,\, \varepsilon}\Big[\big\|\varepsilon_\theta\big(\sqrt{\bar\alpha_t}x_0 + \sqrt{1-\bar\alpha_t}\varepsilon,\ t\big) - \varepsilon\big\|^2\Big]$

Обычный MSE. Ни KL, ни следа якобиана, ни нормировочной константы, ни цепи — одна выборка $t$, один $\varepsilon$, один прогон сети. Именно эта простота, а не какая-то особая архитектура, объясняет, почему диффузия вытеснила предшественников.

Стоит отметить, что $\mathcal{L}$ — не совсем ELBO. Точная вариационная нижняя граница содержит веса, зависящие от $t$; DDPM их выбросил, и получившаяся «упрощённая» цель работает лучше. Интерпретация: веса ELBO переоценивают шаги с малым шумом, где задача проста и почти нечему учиться.

Три эквивалентные параметризации

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

предсказываем	обозначение	связь
шум	$\varepsilon_\theta$	базовая
чистые данные	$x_{0,\theta}$	$x_0 = \dfrac{x_t - \sqrt{1-\bar\alpha}\,\varepsilon}{\sqrt{\bar\alpha}}$
score	$s_\theta$	$s = -\dfrac{\varepsilon}{\sqrt{1-\bar\alpha}}$

Математически это одно и то же; численно — нет, и разница практическая:

• при $\bar\alpha \to 1$ (мало шума) деление на $\sqrt{1-\bar\alpha}$ взрывается, и $\varepsilon$-параметризация численно неудобна;
• при $\bar\alpha \to 0$ (много шума) деление на $\sqrt{\bar\alpha}$ взрывается у $x_0$-параметризации;
• отсюда $v$-параметризация ($v = \sqrt{\bar\alpha}\varepsilon - \sqrt{1-\bar\alpha}x_0$), которая ведёт себя прилично на обоих концах.

Поставьте время близко к нулю и к единице и посмотрите на $\sqrt{\bar\alpha}$ и $\sqrt{1-\bar\alpha}$ — видно, какое из делений опасно на каком конце.

косинусноелинейное

данные при t

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

время t

0.5

ᾱ

0.4938

√ᾱ

0.7027

√(1−ᾱ)

0.7114

сигнал/шум

0.976

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

0.496

Шум преобладает. В этом режиме предсказание x₀ численно ненадёжно: делить приходится на малое √ᾱ.

Формула Tweedie

Ещё одна запись того же факта, полезная тем, что переводит всё на язык статистики:

$\mathbb{E}[x_0 \mid x_t] = \frac{x_t + (1-\bar\alpha)\nabla\log q(x_t)}{\sqrt{\bar\alpha}}$

Проверено: при $\bar\alpha = 0.6$, $x_0 = 1.3$, $\varepsilon = 0.7$ подстановка score даёт ровно $1.3000000000$.

Читается так: апостериорное среднее чистых данных выражается через score зашумлённых. То есть score и оптимальный шумоподавитель — одно и то же, записанное разными буквами. Отсюда и название темы: сеть, обученная убирать шум, автоматически оценивает score, даже если её авторы про score не думали.

Итог

• $\nabla\log q(x_t\mid x_0) = -\varepsilon/\sqrt{1-\bar\alpha}$ — score условного распределения известен точно, потому что зашумление задали мы.
• Минимизация по условному score даёт в оптимуме маргинальный (тождество Винсента), потому что оптимум MSE — условное среднее.
• Цель обучения — обычный MSE по предсказанию $\varepsilon$: без следа, без цепи, без $Z$.
• Три параметризации ($\varepsilon$, $x_0$, score) эквивалентны алгебраически и различаются численной устойчивостью на концах диапазона.
• Формула Tweedie говорит то же самое иначе: оптимальный шумоподавитель и score — один объект.