Diffusion: учимся обращать шум

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

Испортить легко, восстановить трудно — но по частям

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

Теперь попробуйте пройти путь обратно. Задача «по чистому шуму выдать фотографию» безнадёжна: правдоподобных ответов бесконечно много, и любая модель, обученная по среднеквадратичной ошибке, выдаст их среднее — размытое пятно. Но задача «по чуть-чуть испорченной фотографии убрать это чуть-чуть» решается прекрасно, потому что ответ почти единственный.

Идея диффузионной модели ровно в этом разрыве: разложить невозможный прыжок на тысячу возможных шажков. Каждый шаг — маленькое удаление шума; вся цепочка — генерация.

Работать будем на реальных данных, где всё можно посчитать до последней цифры: на наборе рукописных цифр load_digits из sklearn — $1797$ картинок $8\times8$, то есть векторов из $64$ чисел. Мы отложили $1400$ картинок в обучающую часть и $397$ в тестовую, а яркости перевели в отрезок $[-1,1]$. Маленький размер — не бедность, а роскошь: он позволяет выписать точный оптимальный шумоподавитель и увидеть механику диффузии без единой недоученной сети.

Прямой процесс: порча по расписанию

Пусть $x_0$ — изображение. На шаге $t$ к нему подмешивается гауссов шум:

$$q(x_t\mid x_{t-1}) =\mathcal N\!\left(\sqrt{1-\beta_t}\,x_{t-1},\;\beta_t I\right).$$

Множитель $\sqrt{1-\beta_t}$ поставлен не для красоты: он сжимает старый сигнал ровно настолько, чтобы полная дисперсия не разбухала. Обозначим $\alpha_t=1-\beta_t$ и $\bar\alpha_t=\prod_{s=1}^t\alpha_s$. Сумма независимых гауссовых величин снова гауссова, поэтому тысячу шагов можно проскочить одной формулой:

$$x_t=\sqrt{\bar\alpha_t}\,x_0+\sqrt{1-\bar\alpha_t}\,\varepsilon, \qquad \varepsilon\sim\mathcal N(0,I).$$

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

$$\operatorname{Var}(x_t)=\bar\alpha_t\cdot 1+(1-\bar\alpha_t)\cdot 1=1 .$$

Картинка тускнеет, но шкала не съезжает — и сеть на всех уровнях шума видит входы одного масштаба.

Формулу легко проверить измерением. Возьмём $400$ реальных цифр, зашумим их до уровня $t$ и посчитаем корреляцию между пикселями $x_t$ и $x_0$. Теория даёт

$$\rho_t=\frac{\sqrt{\bar\alpha_t}\,s}{\sqrt{\bar\alpha_t s^2+1-\bar\alpha_t}}, \qquad s^2=\operatorname{Var}(x_0)=0{,}565 .$$

Замер и формула совпадают: при $t=100$ — $0{,}976$ против $0{,}976$, при $t=300$ — $0{,}820$ против $0{,}822$, при $t=600$ — $0{,}476$ против $0{,}475$, при $t=900$ — $0{,}118$ против $0{,}117$. Обратите внимание: при $t=600$ коэффициент $\sqrt{\bar\alpha_t}=0{,}584$, а корреляция всего $0{,}476$ — потому что дисперсия наших картинок меньше единицы, и шум съедает сигнал быстрее номинала. Это первое напоминание, что «уровень шума» имеет смысл только относительно масштаба данных.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Расписание — это учебный план

Каждое обучающее сравнение выбирает случайный $t$, поэтому кривая $\bar\alpha_t$ определяет, какую долю бюджета сеть потратит на лёгкие, средние и безнадёжные задачи. Считать шаги бессмысленно, считать надо $\mathrm{SNR}$.

У линейного расписания $32{,}6\%$ всех шагов имеют $\mathrm{SNR}<0{,}01$: треть примеров — это почти чистый шум, где предсказывать нечего и градиенты почти не несут информации. У косинусного таких шагов $6{,}4\%$. Ровно поэтому косинусное расписание было предложено для картинок малого разрешения, где линейное успевает уничтожить сигнал слишком рано.

Один прыжок даёт среднее, а не образец

Что вообще должен выдавать идеальный шумоподавитель? Если минимизировать квадратичную ошибку, ответ известен из урока 20: оптимальный прогноз — условное среднее

$$\widehat x_0(x_t)=\mathbb E[x_0\mid x_t].$$

Для нашей выборки это среднее выписывается точно. Если считать распределением данных саму обучающую выборку из $N=1400$ картинок, то

$$\mathbb E[x_0\mid x_t]=\sum_{i=1}^{N}w_i\,x^{(i)},\qquad w_i\propto\exp\!\left(-\frac{\bigl\|x_t-\sqrt{\bar\alpha_t}\,x^{(i)}\bigr\|^2} {2(1-\bar\alpha_t)}\right).$$

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

$$n_{\mathrm{eff}}=\exp\!\left(-\sum_i w_i\log w_i\right).$$

Замер вдоль расписания: при $t=50$ и $t=200$ получается $n_{\mathrm{eff}}=1{,}0$ — шум ещё не скрыл, какая именно это была картинка. При $t=400$ уже $2{,}7$, при $t=600$ — $9{,}3$, при $t=800$ — $405$, а при $t=1000$ — все $1400$.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Вот и ответ на вопрос «почему нельзя убрать весь шум одним махом». Из чистого шума условное среднее равно просто среднему по выборке: наш замер даёт расстояние $0{,}0003$ между «одним прыжком» из $x_{1000}$ и средней картинкой обучающего набора. Одношаговая модель, обученная по квадратичной ошибке, честно выдаёт размытую кляксу — и это не её недоработка, а правильный ответ на неправильно поставленный вопрос.

Что именно учит сеть

Перепишем условное среднее через шум. Из формулы прямого процесса

$$\varepsilon=\frac{x_t-\sqrt{\bar\alpha_t}\,x_0}{\sqrt{1-\bar\alpha_t}},$$

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

$$\mathcal L_{\mathrm{simple}} =\mathbb E_{x_0,t,\varepsilon} \bigl\|\varepsilon-\varepsilon_\theta(x_t,t)\bigr\|_2^2 .$$

Эта простота не куплена нестрогостью. Вариационная нижняя оценка правдоподобия (та самая, что в уроке 85 обслуживала автокодировщик) для диффузии сводится ровно к той же сумме квадратов, но с весами по времени:

$$\mathcal L_{\mathrm{ELBO}}=\sum_{t} \frac{\beta_t^{2}}{2\sigma_t^{2}\alpha_t(1-\bar\alpha_t)}\, \mathbb E\bigl\|\varepsilon-\varepsilon_\theta(x_t,t)\bigr\|^2+\mathrm{const}.$$

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

Шум, плотность и формула Тweedie

У предсказанного шума есть второй, более глубокий смысл. Плотность зашумлённых данных $q_t$ — это свёртка распределения данных с гауссианой, и градиент её логарифма (его называют score) связан с оптимальным предсказанием шума жёстким тождеством:

$$\nabla_{x_t}\log q_t(x_t) =-\frac{\varepsilon^\star(x_t,t)}{\sqrt{1-\bar\alpha_t}} .$$

Эквивалентная запись — формула Тweedie, связывающая условное среднее с градиентом плотности:

$$\mathbb E[x_0\mid x_t] =\frac{x_t+(1-\bar\alpha_t)\nabla_{x_t}\log q_t(x_t)}{\sqrt{\bar\alpha_t}} .$$

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

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Обратный шаг

Обратную цепочку моделируют гауссовыми переходами

$$p_\theta(x_{t-1}\mid x_t)=\mathcal N\bigl(\mu_\theta(x_t,t),\,\Sigma_t\bigr),$$

а при $\varepsilon$-параметризации среднее выписывается явно:

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

Дисперсию шага обычно берут равной дисперсии истинного апостериорного перехода $q(x_{t-1}\mid x_t,x_0)$:

$$\sigma_t^2=\frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\,\beta_t .$$

Стартуем с $x_T\sim\mathcal N(0,I)$ и идём $T,T-1,\dots,1$; на последнем шаге шум не добавляют.

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

Лаборатория обращения шума

В лаборатории данные — две моды на плоскости (дуга и пятно), а score не приближённый, а точный: это score сглаженного эмпирического распределения показанных точек. Значит, всё, что вы увидите, — свойства самой схемы, а не недоучившейся сети.

Первое, что стоит сделать, — сравнить расписания по $\mathrm{SNR}$, а не по номеру шага. Поставьте $t=400$: у косинусного расписания сигнала ещё около двух третей дисперсии, у линейного — меньше пятой части, хотя номер шага один и тот же.

Второе — уменьшить число шагов обратной цепочки до двух. Все сорок образцов сваливаются в одну моду, разнообразие обнуляется. При трёх шагах доля моды «дуга» подскакивает до двух третей вместо честных $58\%$, и лишь начиная примерно с пяти шагов пропорции мод приходят в норму. Это и есть цена ускорения — не размытость, а перекос распределения.

Третье — включить условие и крутить guidance. При $w=1$ образцы честно попадают в нужный класс. При $w\geqslant4$ все сорок точек стягиваются к одной и той же обучающей точке: соответствие условию идеально, разнообразия нет вовсе. Ползунок guidance — это не ручка «качество», а ручка компромисса.

Порядок рождения деталей предсказуем

Посмотрите на любой ролик генерации: сначала возникает общая композиция, потом крупные формы, и только в конце — фактура. Это не метафора, а арифметика отношения сигнал/шум по направлениям.

Разложим данные по главным компонентам с дисперсиями $\lambda_j$. В направлении $j$ прямой процесс оставляет сигнал $\bar\alpha_t\lambda_j$ при шуме $1-\bar\alpha_t$, то есть собственное отношение

$$\mathrm{SNR}_j(t)=\frac{\bar\alpha_t\lambda_j}{1-\bar\alpha_t}.$$

Направление «просыпается» в обратном ходе там, где $\mathrm{SNR}_j$ переходит единицу, а это происходит тем раньше, чем больше $\lambda_j$:

$$\bar\alpha^\star_j=\frac{1}{1+\lambda_j}.$$

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Практический вывод: если правки нужны в глобальной композиции, вмешиваться надо рано, на больших $t$; если в текстуре — поздно. Именно на этом устроены приёмы редактирования, о которых пойдёт речь в уроке 87.

Сколько шагов достаточно

Число шагов можно уменьшить, если отказаться от случайности. Детерминированный солвер DDIM использует ту же сеть, но переходит сразу на выбранный уровень $s<t$:

$$\widehat x_0=\frac{x_t-\sqrt{1-\bar\alpha_t}\,\varepsilon_\theta(x_t,t)} {\sqrt{\bar\alpha_t}},$$ $$x_s=\sqrt{\bar\alpha_s}\,\widehat x_0+\sqrt{1-\bar\alpha_s}\, \varepsilon_\theta(x_t,t).$$

Из одного и того же $x_T$ разные $K$ дают сравнимые пары, поэтому качество ускорения честно мерить расстоянием до эталона.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Отсюда правило гигиены: название «$50$ шагов» ничего не описывает. Нужно указывать солвер, сетку шагов, порядок метода и параметр стохастичности — иначе два «одинаковых» замера сравнивать нельзя. Экономика этого выбора — продолжение разговора о стоимости инференса: пользователю важна задержка одного запроса, серверу — пропускная способность пакета, и оптимум у них разный.

Условие и guidance: ручка компромисса

Условие $c$ — класс, текст или маска — подаётся в шумоподавитель наряду с $t$. Обучают одну сеть на два режима, случайно заменяя часть условий пустым символом $\varnothing$. Тогда при генерации доступны оба прогноза, и их комбинируют:

$$\widehat\varepsilon=\varepsilon_\theta(x_t,t,\varnothing) +w\bigl[\varepsilon_\theta(x_t,t,c)-\varepsilon_\theta(x_t,t,\varnothing)\bigr].$$

При $w=0$ это безусловная генерация, при $w=1$ — обычная условная, при $w>1$ — экстраполяция за условный прогноз. В терминах плотностей комбинация означает выборку из

$$\tilde p(x\mid c)\;\propto\;p(x)\, \Bigl(\frac{p(x\mid c)}{p(x)}\Bigr)^{w} \;\propto\;p(x)\,p(c\mid x)^{w},$$

то есть возведение правдоподобия условия в степень $w$ — заострение, ровно как температура в софтмаксе.

Рисунок шире экрана — проведите по немуОткрыть целиком ↗

Это и есть главный урок про guidance: за него всегда платят разнообразием, а после некоторого $w$ перестают получать даже соответствие. В нашем замере оптимум по соответствию — $w=2$, но уже там разнообразие упало на пятую часть от безусловного.

Латентная диффузия: где считать дешевле

В пикселях считать дорого: сеть должна обрабатывать десятки тысяч чисел на каждом из полусотни шагов. Латентная диффузия переносит процесс в сжатое пространство: автокодировщик кодирует изображение в латент $z_0$ меньшего размера, вся цепочка идёт там, и лишь один раз декодер возвращает пиксели. Уменьшение стороны в $8$ раз сокращает число чисел в $64$ раза.

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

Текстовое условие обычно подаётся через перекрёстное внимание, а сами представления текста берутся из моделей, родственных контрастивным текстовым энкодерам.

Непрерывный предел: обратное время

Если шаги устремить к нулю, зашумление перестаёт быть цепочкой картинок и становится непрерывным движением. Записывают его так (запись читается как «за маленькое время $dt$ точка сдвинулась на снос плюс случайный толчок»):

$$dx=f(x,t)\,dt+g(t)\,dW_t,$$

где $W_t$ — винеровский процесс, тот самый броуновский шум из урока 66. Замечательный факт: у такого уравнения существует уравнение обратного времени, и оно снова диффузионное —

$$dx=\bigl[f(x,t)-g(t)^2\nabla_x\log p_t(x)\bigr]dt+g(t)\,d\bar W_t .$$

Всё, чего не хватает для обращения времени, — то самое поле score. Более того, существует детерминированное уравнение потока вероятности

$$\frac{dx}{dt}=f(x,t)-\tfrac12 g(t)^2\nabla_x\log p_t(x),$$

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

Скороход: любое распределение — остановленное блуждание

Здесь уместно вспомнить советскую школу теории случайных процессов. Анатолий Владимирович Скороход (1930–2011) вместе с Иосифом Ильичом Гихманом построил систематическую теорию стохастических дифференциальных уравнений — того самого объекта, который сегодня описывает прямой и обратный процессы диффузионных моделей. Их «Стохастические дифференциальные уравнения» (Киев, 1968) стали одной из первых монографий, где диффузионный процесс определялся как решение уравнения со случайной правой частью, а не как набор частных примеров.

Но ещё ближе к нашему сюжету другая его теорема — о вложении, доказанная в начале 1960-х. Она утверждает поразительное: для любой случайной величины с нулевым средним и конечной дисперсией найдётся момент остановки $\tau$ броуновского движения такой, что

$$B_\tau\;\stackrel{d}{=}\;X,\qquad \mathbb E\,\tau=\operatorname{Var}X .$$

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

Когда модель просто помнит

Точный score эмпирического распределения — идеальный шумоподавитель для обучающей выборки, и именно поэтому он бесполезен как генератор. Мы прогнали цепочку с этим точным score: среднее расстояние до ближайшей обучающей картинки оказалось $0{,}000$ — модель воспроизвела обучающие цифры буква в букву. У модели с ограниченной ёмкостью то же расстояние равно $3{,}21$.

Вывод парадоксальный и важный: обобщение диффузионной модели — следствие того, что сеть не достигает оптимума функции потерь. Она сглаживает, потому что не может выучить $1400$ дельта-функций, и это сглаживание и есть генерация нового. Оттого и проверка на запоминание должна быть обязательной частью протокола: редкие обучающие объекты, близкие подсказки, поиск и точные совпадения, и обрезанных фрагментов. Причины запоминания обычно лежат в дубликатах корпуса, слишком узком условии или избыточном числе эпох — та же тема, что и в дедупликации предобучения.

Что мерить и чему не верить

Соберём протокол честной оценки диффузионной модели.

Во-первых, ошибка предсказания шума по уровням $\mathrm{SNR}$, а не в среднем: одна и та же величина ошибки при $t=100$ и при $t=900$ означает совершенно разные вещи для итоговой картинки. Во-вторых, генеративные метрики распределений (расстояние Фреше в признаках замороженного классификатора, точность и полнота) — обязательно на одинаковом числе образцов и одинаковой статистике реальных данных, иначе числа несравнимы. В-третьих, покрытие мод: доля классов, которые модель вообще порождает. В-четвёртых, время генерации отдельно для пакета из одного и из шестидесяти четырёх образцов.

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

Генерация как обращение шума

Диффузия раскладывает невозможную задачу на тысячу возможных. Прямой процесс известен точно и бесплатен; обучение сводится к одной регрессии на предсказание шума; предсказанный шум оказывается градиентом логарифма плотности; обратная цепочка идёт по этому полю, а её случайность отвечает за разнообразие. Расписание задаёт учебный план, число шагов — цену и точность, guidance — размен соответствия и разнообразия. Всё это части модели результата, а не «настройки по вкусу»: меняя их, вы меняете распределение, из которого берёте образец.

Задачи