Содержание курса 26 / 33

26 / 33 · Диффузионные модели

DDPM: прямой процесс

Как шумоподавляющие диффузионные модели намеренно разрушают данные. В наивной цепочке, которая просто добавляет шум, дисперсия неограниченно растёт, а шаг с масштабированием √(1 − β) x + √β ε даёт переход q(xₜ | x₀) в замкнутом виде, который стремится к N(0, I); на практике β меняется по линейному, квадратичному, сигмоидному или косинусному расписанию.

На этой странице 8

Заметка о моделях, основанных на энергии, закончилась на цене этого подхода: модель, основанную на энергии (energy-based model), обучают с короткими смещёнными цепочками MCMC, которые запускаются внутри каждого обновления параметров, поэтому выученная модель зависит от сэмплера не меньше, чем от данных. Диффузионные модели выносят сэмплирование из обучения. Они фиксируют простой процесс зашумления, который превращает данные в гауссовский шум, и учатся обращать его по одному малому шагу за раз. Каждый обратный шаг близок к гауссовскому, поэтому сеть выучивает его регрессией на зашумлённых сэмплах, которые фиксированный процесс даёт бесплатно. Оказывается, эта сеть выучивает score зашумлённых данных — градиент энергии со знаком минус и направление, по которому движется сэмплирование динамикой Ланжевена (заметка о score matching). Зашумлённые латентные переменные сохраняют размерность данных, архитектура сети ничем не ограничена, а платой становится нижняя оценка правдоподобия вместо самого правдоподобия. DDPM (Denoising Diffusion Probabilistic Models, шумоподавляющие диффузионные вероятностные модели; Ho, Jain and Abbeel, 2020) — модель, которая довела качество изображений у диффузионных моделей до уровня, сравнимого с GAN. Эта заметка посвящена её первой половине — прямому процессу (forward process): фиксированной процедуре, которая постепенно превращает изображение в гауссовский шум. Вторая половина — нейронная сеть, которая выполняет эту процедуру в обратном направлении, — описана в следующей заметке.

Обозначения. x0x_0 — точка данных (изображение), а xtx_t — её зашумлённая версия после tt шагов, t=1,…,Tt = 1, \dots, T. Через qq обозначается фиксированный прямой процесс (процесс зашумления): q(x0)q(x_0) — распределение данных, а q(xt∣xt−1)q(x_t \mid x_{t-1}) — один шаг зашумления. pθ(xt−1∣xt)p_\theta(x_{t-1} \mid x_t) — обучаемый обратный шаг с параметрами сети θ\theta. ε∼N(0,1)\varepsilon \sim \mathcal N(0, 1) — стандартный гауссовский шум, который на каждом шаге сэмплируется заново, а второй аргумент N(μ,σ2)\mathcal N(\mu, \sigma^2) — это дисперсия. Мы работаем в одном измерении; для изображений каждый пиксель зашумляется независимо и 11 превращается в единичную матрицу II, так что целевое распределение — N(0,I)\mathcal N(0, I).

Немного истории#

  • В 2015 году, когда лучшими в генерации изображений были GAN, вышла статья Sohl-Dickstein et al. (2015) Deep Unsupervised Learning using Nonequilibrium Thermodynamics. В ней появилась идея диффузионных моделей: разрушить структуру данных медленным процессом зашумления, а затем научиться обращать этот процесс.
  • Пять лет спустя DDPM упростила параметризацию и функцию потерь и достигла качества изображений, которое может соперничать с GAN. Её целевая функция обучения по-прежнему выводится из вариационной нижней оценки правдоподобия — того же инструмента, что и в VAE.
  • Вскоре после этого диффузионные модели обошли GAN на стандартных бенчмарках генерации изображений (Dhariwal and Nichol, 2021).

Задача#

Мы не знаем распределение реальных изображений, но хотим сэмплировать из него новые изображения. В обозначениях диффузионных моделей это распределение записывается как q(x0)q(x_0), потому что с него начинается прямой процесс qq.

Идея диффузии — связать q(x0)q(x_0) с распределением, из которого мы умеем сэмплировать. Возьмём реальное изображение x0x_0 — исходную фотографию — и добавим немного шума, чтобы получить x1x_1, затем ещё немного, чтобы получить x2x_2, и так далее, пока xTx_T не станет чистым шумом. Каждый шаг прямого процесса фиксирован и не требует обучения. Модель — обучаемый обратный шаг pθ(xt−1∣xt)p_\theta(x_{t-1} \mid x_t) — идёт в другую сторону: по xtx_t она строит xt−1x_{t-1}, на один шаг ближе к изображению. Чтобы сгенерировать изображение, мы сэмплируем чистый шум xTx_T и применяем модель TT раз.

x0изображение  ⟶  x1  ⟶  ⋯  ⟶  xTшум,xt−1  ←  pθ(xt−1∣xt)    xt.\begin{gathered} \underset{\text{изображение}}{x_0} \;\longrightarrow\; x_1 \;\longrightarrow\; \cdots \;\longrightarrow\; \underset{\text{шум}}{x_T}, \\ x_{t-1} \;\xleftarrow{\;p_\theta(x_{t-1} \mid x_t)\;}\; x_t . \end{gathered}
Рисунок

Прямая цепочка добавляет шум, модель идёт обратно

Диффузионная цепочка от исходной фотографии до чистого шумаЧетыре пиксельные плитки в ряд: x0 — исходная фотография солнца и холма, x1 с небольшим шумом, x2 с бо́льшим шумом, затем многоточие, затем xT — чистый шум. Прямые стрелки с подписью «зашумить» идут слева направо. Изогнутые стрелки сверху идут справа налево и подписаны p тета от x0 при условии x1, p тета от x1 при условии x2 и p тета от xT минус 1 при условии xT: это обучаемый обратный шаг. Длинная стрелка снизу подписана: прямой процесс q от xt при условии xt минус 1, фиксирован; длинная стрелка сверху — обратный процесс, обучаемая модель.обратный процесс: обучаемая модель pθ (на шаг назад)x0исходноефотоx1немногошумаx2большешумаxTчистыйшум…зашумитьзашумитьpθ(x0 | x1)pθ(x1 | x2)pθ(xT−1 | xT)прямой процесс q(xt | xt−1): фиксирован, только добавляет шумДиффузионная цепочка от исходной фотографии до чистого шумаЧетыре пиксельные плитки в столбец: x0 — исходная фотография солнца и холма, x1 с небольшим шумом, x2 с бо́льшим шумом, многоточие, затем xT — чистый шум. Стрелки с подписью «зашумить» направлены вниз. Изогнутые синие стрелки справа направлены вверх и подписаны p тета от x0 при условии x1, p тета от x1 при условии x2 и p тета от xT минус 1 при условии xT: это обучаемый обратный шаг, который возвращается на шаг назад.x0исходное фотоx1немного шумаx2больше шумаxTчистый шум⋮зашумитьзашумитьpθ(x0 | x1)pθ(x1 | x2)pθ(xT−1 | xT)синим — обучаемая модель, шаг назад
Стрелки вперёд — фиксированный шаг зашумления q(xₜ | xₜ₋₁); изогнутые стрелки — обучаемый обратный шаг pθ(xₜ₋₁ | xₜ), который берёт xₜ и возвращается на шаг назад. Плитки — настоящие сэмплы xₜ = √α̅ₜ x₀ + √(1 − α̅ₜ) ε для α̅ = 1, 0.8, 0.5 и ≈ 0 (4 × 10⁻⁵, конец расписания DDPM) с одной фиксированной реализацией шума ε.

Первая попытка: просто добавить шум#

Простейший шаг добавляет шум с дисперсией β\beta:

x1=x0+β ε,ε∼N(0,1)⟹q(x1∣x0)=N(x0,β),\begin{gathered} x_1 = x_0 + \sqrt{\beta}\,\varepsilon, \quad \varepsilon \sim \mathcal N(0, 1) \\ \Longrightarrow\quad q(x_1 \mid x_0) = \mathcal N(x_0, \beta), \end{gathered}

где β\beta — дисперсия шума. Здесь важен квадратный корень: дисперсия c εc\,\varepsilon равна c2c^2, поэтому шум с дисперсией β\beta — это β ε\sqrt\beta\,\varepsilon. Следующий шаг делает то же самое с x1x_1: q(x2∣x1)=N(x1,β)q(x_2 \mid x_1) = \mathcal N(x_1, \beta).

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

x2=x1+β ε2=x0+β ε1+β ε2=x0+2β ε,q(x2∣x0)=N(x0,2β),\begin{aligned} x_2 &= x_1 + \sqrt\beta\,\varepsilon_2 \\ &= x_0 + \sqrt\beta\,\varepsilon_1 + \sqrt\beta\,\varepsilon_2 \\ &= x_0 + \sqrt{2\beta}\,\varepsilon, \\ q(x_2 \mid x_0) &= \mathcal N(x_0, 2\beta), \end{aligned}

где ε=(ε1+ε2)/2\varepsilon = (\varepsilon_1 + \varepsilon_2)/\sqrt2 — снова стандартная нормальная величина. Всякий раз, когда ниже несколько шумов объединяются в один, ε\varepsilon обозначает объединённую стандартную нормальную величину в этом смысле, так что такие равенства выполняются по распределению.

Продолжая так же, после tt шагов получаем

xt=x0+tβ ε,q(xt∣x0)=N(x0,tβ).\begin{gathered} x_t = x_0 + \sqrt{t\beta}\,\varepsilon, \\ q(x_t \mid x_0) = \mathcal N(x_0, t\beta). \end{gathered}

Это полезное сокращение: от x0x_0 можно перейти к любому xtx_t за одно сэмплирование. Но посмотрим на два параметра:

  • среднее не меняется: оно навсегда остаётся равным x0x_0;
  • дисперсия неограниченно растёт: tβ→∞t\beta \to \infty.

Поэтому цепочка никогда не приходит к N(0,1)\mathcal N(0, 1): дисперсия растёт без ограничений, и нет фиксированного предельного распределения, с которого могла бы начаться генерация. Влияние изображения ослабевает лишь относительно растущего разброса. Такой процесс называют диффузией с растущей дисперсией (variance-exploding) (Song et al., 2021, §3.4). Модели, построенные на нём, останавливаются на большом конечном уровне шума σmax⁡\sigma_{\max} и начинают сэмплирование из N(0,σmax⁡2)\mathcal N(0, \sigma_{\max}^2) (Song et al., 2021, §2.1).

Цель и шаг с масштабированием#

Мы хотим, чтобы выполнялось

q(xt∣x0)→t→∞N(0,1),q(x_t \mid x_0) \xrightarrow[t\to\infty]{} \mathcal N(0, 1),

каким бы ни было начальное изображение. Для этого нужен новый способ переходить от шага t−1t-1 к шагу tt; в частности, шаг должен каждый раз менять среднее, притягивая его к 00. В общем случае Sohl-Dickstein et al. допускают любое марковское переходное ядро Tπ(xt∣xt−1;βt)T_\pi(x_t \mid x_{t-1}; \beta_t), которое оставляет неизменным простое целевое распределение π\pi, π(y)=∫Tπ(y∣y′;β) π(y′) dy′\pi(y) = \int T_\pi(y \mid y'; \beta)\,\pi(y')\,dy' (2015, §2.1, формулы 1–2): биномиальное ядро для бинарных данных, гауссовское — для вещественных. Для вещественных значений будем искать шаг вида xt=a xt−1+β εx_t = a\,x_{t-1} + \sqrt\beta\,\varepsilon. Гауссовское переходное ядро диффузионных моделей, которое ввели Sohl-Dickstein et al. (2015, таблица App.1) и сохранила DDPM (Ho, Jain and Abbeel, 2020, формула 2), умножает предыдущую точку на a=1−βa = \sqrt{1-\beta}:

xt=1−β  xt−1+β  ε,q(xt∣xt−1)=N ⁣(1−β  xt−1,  β).\begin{gathered} x_t = \sqrt{1-\beta}\;x_{t-1} + \sqrt{\beta}\;\varepsilon, \\ q(x_t \mid x_{t-1}) = \mathcal N\!\big(\sqrt{1-\beta}\;x_{t-1},\; \beta\big). \end{gathered}
Рисунок

Шаг за шагом или сразу от x₀ к xₜ

Наивная цепочка зашумления и цепочка DDPM с переходами за один скачокДва ряда узлов в кружках x0, x1, x2, многоточие и xt, соединённых стрелками шагов. Верхний ряд, наивная цепочка: каждый шаг добавляет корень из бета, умноженный на эпсилон; длинная дуга от x0 к xt подписана q от xt при условии x0 равно N от x0 и t бета, а красное предупреждение говорит, что среднее остаётся x0, а дисперсия t бета стремится к бесконечности, поэтому цепочка никогда не приходит к N(0, 1). Нижний ряд, DDPM: каждый шаг умножает на корень из 1 минус бета и добавляет корень из бета, умноженный на эпсилон; длинная дуга подписана q от xt при условии x0 равно N от корня из альфа с чертой t, умноженного на x0, и 1 минус альфа с чертой t, а зелёная пометка говорит, что альфа с чертой t равно 1 минус бета в степени t и стремится к нулю, поэтому цепочка приближается к N(0, 1).Наивная: xₜ = xₜ₋₁ + √β εx0x1x2xt…+ √β ε+ √β εq(xₜ | x₀) = N(x₀, tβ)среднее x₀, дисперсия → ∞: N(0, 1) недостижимоDDPM: xₜ = √(1 − β) xₜ₋₁ + √β εx0x1x2xt…× √(1 − β)+ √β ε× √(1 − β)+ √β εq(xₜ | x₀) = N(√α̅ₜ x₀, 1 − α̅ₜ)α̅ₜ = (1 − β)ᵗ → 0, поэтому → N(0, 1)Наивная цепочка зашумления и цепочка DDPM с переходами за один скачокДва ряда узлов в кружках x0, x1, x2, многоточие и xt, соединённых стрелками шагов. Верхний ряд, наивная цепочка: каждый шаг добавляет корень из бета, умноженный на эпсилон; длинная дуга от x0 к xt подписана q от xt при условии x0 равно N от x0 и t бета, а красное предупреждение говорит, что среднее остаётся x0, а дисперсия t бета стремится к бесконечности, поэтому цепочка никогда не приходит к N(0, 1). Нижний ряд, DDPM: каждый шаг умножает на корень из 1 минус бета и добавляет корень из бета, умноженный на эпсилон; длинная дуга подписана q от xt при условии x0 равно N от корня из альфа с чертой t, умноженного на x0, и 1 минус альфа с чертой t, а зелёная пометка говорит, что альфа с чертой t равно 1 минус бета в степени t и стремится к нулю, поэтому цепочка приближается к N(0, 1).Наивная: xₜ = xₜ₋₁ + √β εx0x1x2xt…q(xₜ | x₀) = N(x₀, tβ)среднее x₀, дисперсия → ∞: N(0, 1) недостижимоDDPM: xₜ = √(1 − β) xₜ₋₁ + √β εx0x1x2xt…q(xₜ | x₀) = N(√α̅ₜ x₀, 1 − α̅ₜ)α̅ₜ = (1 − β)ᵗ → 0, поэтому → N(0, 1)
Отдельные шаги и сокращённые переходы за один скачок для двух цепочек. Обе на каждом шаге добавляют новый гауссовский шум ε ~ N(0, 1), и в обеих t шагов сворачиваются в один гауссовский скачок q(xₜ | x₀), так что xₜ можно сэмплировать напрямую. Без масштабирования (сверху) скачок оставляет среднее равным x₀, а его дисперсия неограниченно растёт. Умножение на √(1 − β) на каждом шаге (снизу) уменьшает среднее и держит условную дисперсию 1 − α̅ₜ ниже 1, так что цепочка приближается к N(0, 1). Второй аргумент N — дисперсия.

Почему именно 1−β\sqrt{1-\beta}? Новый шум не зависит от xt−1x_{t-1}, поэтому дисперсия подчиняется рекуррентному соотношению Var⁡(xt)=a2Var⁡(xt−1)+β\operatorname{Var}(x_t) = a^2 \operatorname{Var}(x_{t-1}) + \beta. Неподвижная точка этого соотношения равна 11, дисперсии целевого распределения, только если a2+β=1a^2 + \beta = 1. Значит, a=1−βa = \sqrt{1-\beta} — единственный положительный масштаб, при котором единичная дисперсия остаётся равной 11, и он же притягивает к 11 любую другую дисперсию (см. врезку ниже). При таком aa переход от x0x_0 к любому xtx_t принимает особенно простой замкнутый вид, в котором сумма квадратов двух коэффициентов равна 11 (Ho, Jain and Abbeel, 2020, формула 4):

xt=αˉt x0+1−αˉt  ε,αˉt=(1−β)t,\begin{gathered} x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\;\varepsilon, \\ \bar\alpha_t = (1-\beta)^t, \end{gathered} q(xt∣x0)=N ⁣(αˉt x0,  1−αˉt).q(x_t \mid x_0) = \mathcal N\!\big(\sqrt{\bar\alpha_t}\,x_0,\; 1-\bar\alpha_t\big).

Если положить αˉ0=1\bar\alpha_0 = 1, формула охватывает и случай t=0t = 0. Значит, xtx_t для любого шага можно получить прямо из x0x_0, за одно сэмплирование, не выполняя tt предшествующих шагов.

Почему это приводит к N(0, 1)#

Если β∈(0,1)\beta \in (0, 1), то αˉt=(1−β)t→0\bar\alpha_t = (1-\beta)^t \to 0 при t→∞t \to \infty. Тогда член с изображением αˉt x0→0\sqrt{\bar\alpha_t}\,x_0 \to 0, а коэффициент шума 1−αˉt→1\sqrt{1-\bar\alpha_t} \to 1, так что шумовой член стремится просто к ε\varepsilon с дисперсией 11:

q(xt∣x0)→t→∞N(0,1).q(x_t \mid x_0) \xrightarrow[t\to\infty]{} \mathcal N(0, 1).

Например, при постоянном β=0.02\beta = 0.02:

ttαˉt=0.98t\bar\alpha_t = 0.98^tкоэффициент среднего αˉt\sqrt{\bar\alpha_t}стандартное отклонение шума 1−αˉt\sqrt{1-\bar\alpha_t}стандартное отклонение шума в наивной цепочке tβ\sqrt{t\beta}
10.9800.9900.1410.141
100.8170.9040.4280.447
500.3640.6030.7971.000
1000.1330.3640.9311.414
2000.01760.1330.9912.000

При малых tt две цепочки почти совпадают. Затем разброс наивной цепочки превышает целевой и продолжает расти, а у DDPM он приближается к 11 снизу. При конечном TT член с изображением αˉT x0\sqrt{\bar\alpha_T}\,x_0 мал лишь относительно масштаба x0x_0; это одна из причин, по которым изображения перед диффузией масштабируют в [−1,1][-1, 1] (следующая заметка).

Рисунок

Растущая дисперсия против сохранения дисперсии

Коэффициент среднего и стандартное отклонение шума q(xt | x0) в зависимости от tДва графика по t от 0 до 200 при бета, равном 0.02. Зелёная штриховая линия на высоте 1 отмечает целевое стандартное отклонение N(0, 1). Слева или сверху — наивная цепочка с растущей дисперсией: красный коэффициент среднего постоянно равен 1, то есть среднее не меняется, а синее стандартное отклонение шума, корень из t бета, пересекает 1 при t, равном 50, и продолжает расти до 2. Справа или снизу — цепочка DDPM, сохраняющая дисперсию: красный коэффициент среднего, корень из альфа с чертой t, убывает от 1 примерно до 0.13, а синее стандартное отклонение шума, корень из 1 минус альфа с чертой t, растёт к 1 и выходит там на постоянный уровень.Наивная: растущая дисперсия00.511.52050100150200tDDPM: дисперсия сохраняется00.511.52050100150200tσ шума √(tβ)дисперсия растётнеограниченносреднее не меняетсяσ цели = 1среднее: √α̅ₜ → 0σ шума √(1 − α̅ₜ) → 1Коэффициент среднего и стандартное отклонение шума q(xt | x0) в зависимости от tДва графика по t от 0 до 200 при бета, равном 0.02. Зелёная штриховая линия на высоте 1 отмечает целевое стандартное отклонение N(0, 1). Слева или сверху — наивная цепочка с растущей дисперсией: красный коэффициент среднего постоянно равен 1, то есть среднее не меняется, а синее стандартное отклонение шума, корень из t бета, пересекает 1 при t, равном 50, и продолжает расти до 2. Справа или снизу — цепочка DDPM, сохраняющая дисперсию: красный коэффициент среднего, корень из альфа с чертой t, убывает от 1 примерно до 0.13, а синее стандартное отклонение шума, корень из 1 минус альфа с чертой t, растёт к 1 и выходит там на постоянный уровень.Наивная: растущая дисперсия00.511.52050100150200tDDPM: дисперсия сохраняется00.511.52050100150200tσ шума √(tβ)дисперсия растётнеограниченносреднее не меняетсяσ цели = 1среднее: √α̅ₜ → 0σ шума √(1 − α̅ₜ) → 1
Коэффициенты q(xₜ | x₀) при постоянном β = 0.02: xₜ = (коэффициент среднего)·x₀ + (стандартное отклонение шума)·ε. Наивная цепочка оставляет x₀ нетронутым, а стандартное отклонение шума √(tβ) неограниченно растёт и переходит через 1 (2.0 при t = 200). Цепочка DDPM умножает x₀ на √α̅ₜ (0.13 при t = 200), а стандартное отклонение шума √(1 − α̅ₜ) растёт к 1 (0.991 при t = 200) — стандартному отклонению целевого N(0, 1). Зелёная штриховая линия отмечает целевое стандартное отклонение 1; на левом графике она проходит поверх горизонтальной линии среднего, которая тоже находится на уровне 1.

То же происходит со всем распределением данных. Его маргинальное распределение на шаге tt равно q(xt)=∫q(xt∣x0) q(x0) dx0q(x_t) = \int q(x_t \mid x_0)\,q(x_0)\,dx_0. Чтобы увидеть, что оно тоже стремится к N(0,1)\mathcal N(0, 1), запишем xt=αˉt x0+1−αˉt εx_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon, где x0∼q(x0)x_0 \sim q(x_0) не зависит от ε\varepsilon. Первый член — фиксированная случайная величина, умноженная на множитель, который стремится к 00, поэтому он стремится к 00 при любом распределении данных, а второй член стремится к ε\varepsilon; следовательно, xt→N(0,1)x_t \to \mathcal N(0, 1) по распределению. В одном измерении двугорбая плотность расплывается, её горбы движутся к 00 и сливаются (между t=200t = 200 и t=300t = 300 для данных на рисунке ниже), и в конце остаётся одна колоколообразная кривая с центром в 00. Конечный батч сэмплов по пути выглядит зубчатым, но сама плотность остаётся гладкой: это свёртка плотности данных с гауссианой. Для смеси гауссиан маргинальное распределение остаётся смесью и записывается точно:

q(x0)=∑kπk N(μk,σk2),q(xt)=∑kπk N ⁣(αˉt μk,  sk,t2),sk,t2=αˉtσk2+1−αˉt.\begin{aligned} q(x_0) &= \textstyle\sum_k \pi_k\,\mathcal N(\mu_k, \sigma_k^2), \\ q(x_t) &= \textstyle\sum_k \pi_k\,\mathcal N\!\big(\sqrt{\bar\alpha_t}\,\mu_k,\; s_{k,t}^2\big), \\ s_{k,t}^2 &= \bar\alpha_t\sigma_k^2 + 1 - \bar\alpha_t. \end{aligned}

Обе компоненты притягиваются к среднему 00 и дисперсии 11, поэтому сливаются в одно стандартное нормальное распределение.

Рисунок

Распределение данных растворяется в N(0, 1)

Маргинальная плотность xt при t = 0, 100, 300 и 1000Четыре небольших графика плотности, соединённые стрелками. При t = 0 у красной плотности данных два горба, примерно в минус 1.3 и плюс 1.25. При t = 100, альфа с чертой 0.90, горбы ниже, шире и ближе друг к другу. При t = 300, альфа с чертой 0.40, они слились в один несимметричный бугор. При t = 1000, альфа с чертой около 0, зелёная плотность — стандартное нормальное распределение N(0, 1). На каждом графике также показана зубчатая гистограмма 600 зашумлённых сэмплов.t = 0 (данные)α̅ = 1q(x₀)t = 100α̅ = 0.90t = 300α̅ = 0.40t = 1000α̅ ≈ 4.0 × 10⁻⁵N(0, 1)Маргинальная плотность xt при t = 0, 100, 300 и 1000Четыре небольших графика плотности, соединённые стрелками. При t = 0 у красной плотности данных два горба, примерно в минус 1.3 и плюс 1.25. При t = 100, альфа с чертой 0.90, горбы ниже, шире и ближе друг к другу. При t = 300, альфа с чертой 0.40, они слились в один несимметричный бугор. При t = 1000, альфа с чертой около 0, зелёная плотность — стандартное нормальное распределение N(0, 1). На каждом графике также показана зубчатая гистограмма 600 зашумлённых сэмплов.t = 0 (данные)α̅ = 1q(x₀)t = 100α̅ = 0.90t = 300α̅ = 0.40t = 1000α̅ ≈ 4.0 × 10⁻⁵N(0, 1)
Каждая точка данных зашумляется по одному и тому же правилу, поэтому вся одномерная плотность данных q(x₀) — смесь двух гауссиан — превращается в маргинальное распределение q(xₜ) = Σₖ πₖ N(√α̅ₜ μₖ, α̅ₜσₖ² + 1 − α̅ₜ). Кривые — точные плотности; столбцы — гистограмма 600 сэмплов, пропущенных через те же шаги, поэтому конечный батч выглядит зубчатым. Используется линейное расписание β из DDPM от 10⁻⁴ до 0.02 (T = 1000).

На практике: расписание β#

На практике β\beta не постоянна, а меняется от шага к шагу: β1,…,βT\beta_1, \dots, \beta_T; это расписание шума (noise schedule). Вывод выше проходит без изменений, если на шаге tt заменить (1−β)(1-\beta) на (1−βt)(1-\beta_t):

q(xt∣xt−1)=N ⁣(1−βt xt−1,  βt),αˉt=∏i=1t(1−βi),q(xt∣x0)=N ⁣(αˉt x0,  1−αˉt).\begin{gathered} q(x_t \mid x_{t-1}) = \mathcal N\!\big(\sqrt{1-\beta_t}\,x_{t-1},\; \beta_t\big), \\ \bar\alpha_t = \textstyle\prod_{i=1}^{t} (1-\beta_i), \\ q(x_t \mid x_0) = \mathcal N\!\big(\sqrt{\bar\alpha_t}\,x_0,\; 1-\bar\alpha_t\big). \end{gathered}

Обычно используют сокращение αt=1−βt\alpha_t = 1-\beta_t, так что αˉt=∏i≤tαi\bar\alpha_t = \prod_{i \le t}\alpha_i. В статье DDPM берутся T=1000T = 1000 шагов и линейное расписание от β1=10−4\beta_1 = 10^{-4} до βT=0.02\beta_T = 0.02 (Ho, Jain and Abbeel, 2020, §4). Позже обнаружилось, что косинусное расписание, которое разрушает информацию равномернее во времени, лучше работает на низком разрешении (Nichol and Dhariwal, 2021, §3.2). Малые шаги в начале оставляют ранние xtx_t близкими к изображению; более крупные шаги потом доводят дело до конца. При линейном расписании DDPM αˉt\bar\alpha_t примерно равно 0.50.5 при t≈259t \approx 259 (ниже 0.50.5 оно опускается при t=260t = 260), падает ниже 0.010.01 при t=674t = 674 и заканчивается на αˉ1000≈4.0×10−5\bar\alpha_{1000} \approx 4.0\times10^{-5}, так что на практике x1000x_{1000} — просто шум.

Рисунок

Линейное расписание β и соответствующее α̅ₜ

Линейное расписание бета и альфа с чертой на 1000 шагахДва графика с общей осью времени t от 0 до 1000. Верхний график показывает бета t — прямую янтарную линию, которая растёт от 0.0001 до 0.02. Нижний график показывает альфа с чертой t, накопленное произведение 1 минус бета, — убывающую кривую, которая начинается в 1, сначала убывает медленно, пересекает 0.5 примерно при t, равном 259, опускается ниже 0.01 при t, равном 674, и при t, равном 1000, составляет около 4 на 10 в минус пятой степени.00.010.02βtлинейно, 10⁻⁴ → 0.0200.250.50.75102004006008001000tα̅t = ∏ (1 − βi)α̅ₜ ≈ 0.5 при t ≈ 259α̅ₜ < 0.01 с t = 674:x₀ почти исчез, xₜ ≈ εЛинейное расписание бета и альфа с чертой на 1000 шагахДва графика с общей осью времени t от 0 до 1000. Верхний график показывает бета t — прямую янтарную линию, которая растёт от 0.0001 до 0.02. Нижний график показывает альфа с чертой t, накопленное произведение 1 минус бета, — убывающую кривую, которая начинается в 1, сначала убывает медленно, пересекает 0.5 примерно при t, равном 259, опускается ниже 0.01 при t, равном 674, и при t, равном 1000, составляет около 4 на 10 в минус пятой степени.00.010.02βtлинейно, 10⁻⁴ → 0.0200.250.50.75102505007501000tα̅t = ∏ (1 − βi)α̅ₜ ≈ 0.5 при t ≈ 259α̅ₜ < 0.01 с t = 674:x₀ почти исчез, xₜ ≈ ε
Сверху: βₜ линейно растёт от 10⁻⁴ до 0.02. Снизу: α̅ₜ = ∏ (1 − βᵢ) сначала падает медленно (малые шаги оставляют ранние изображения узнаваемыми), составляет ≈ 0.5 при t ≈ 259, опускается ниже 0.01 при t = 674 и заканчивается на α̅₁₀₀₀ ≈ 4.0 × 10⁻⁵, так что x₁₀₀₀ — практически чистый шум N(0, I).

Расписание можно описать и через отношение сигнал/шум (signal-to-noise ratio), SNRt=αˉt/(1−αˉt)\mathrm{SNR}_t = \bar\alpha_t / (1 - \bar\alpha_t): квадрат коэффициента при изображении, делённый на дисперсию шума; для данных с единичной дисперсией это отношение двух дисперсий. Для линейного расписания оно падает примерно от 10410^4 при t=1t = 1 до 4×10−54\times10^{-5} при t=Tt = T. Быстрые сэмплеры из заметки о латентной диффузии работают с его логарифмом.

Линейное, квадратичное, сигмоидное и косинусное расписания#

Линейное расписание — лишь один из нескольких вариантов (их обзор — у Weng, 2021). Обозначим через ut=(t−1)/(T−1)∈[0,1]u_t = (t - 1)/(T - 1) \in [0, 1] положение в цепочке. При тех же крайних значениях β1=10−4\beta_1 = 10^{-4} и βT=0.02\beta_T = 0.02 три распространённые формы задаются через βt\beta_t:

линейное:βt=β1+(βT−β1) ut,квадратичное:βt=(β1+(βT−β1) ut)2,сигмоидное:βt=β1+(βT−β1) sigmoid⁡(12 ut−6).\begin{aligned} \text{линейное:}\quad & \beta_t = \beta_1 + (\beta_T - \beta_1)\,u_t, \\ \text{квадратичное:}\quad & \beta_t = \big(\sqrt{\beta_1} + (\sqrt{\beta_T} - \sqrt{\beta_1})\,u_t\big)^2, \\ \text{сигмоидное:}\quad & \beta_t = \beta_1 + (\beta_T - \beta_1)\,\operatorname{sigmoid}(12\,u_t - 6). \end{aligned}

Квадратичное расписание линейно интерполирует амплитуду шума βt\sqrt{\beta_t}, поэтому его шаги дольше остаются малыми. Сигмоидное расписание плоское на обоих концах и меняется в середине; поскольку sigmoid⁡(±6)\operatorname{sigmoid}(\pm 6) не равна в точности 00 или 11, его βt\beta_t на самом деле меняются от 1.5×10−41.5\times10^{-4} до 0.019950.01995. Косинусное расписание (Nichol and Dhariwal, 2021, §3.2, формула 17) задаёт не шаги, а αˉt\bar\alpha_t и выводит шаги из него:

αˉt=f(t)f(0),f(t)=cos⁡2 ⁣(t/T+s1+s⋅π2),βt=1−αˉtαˉt−1,\begin{gathered} \bar\alpha_t = \frac{f(t)}{f(0)}, \qquad f(t) = \cos^2\!\Big(\frac{t/T + s}{1 + s}\cdot\frac{\pi}{2}\Big), \\ \beta_t = 1 - \frac{\bar\alpha_t}{\bar\alpha_{t-1}} , \end{gathered}

со сдвигом s=0.008s = 0.008, выбранным так, чтобы амплитуда шума на первом шаге была чуть меньше ширины 1/127.51/127.5 одного уровня яркости пикселя, и с βt\beta_t, обрезанными сверху на 0.9990.999. Без обрезки f(T)=0f(T) = 0 дало бы βT=1\beta_T = 1; с ней αˉ1000≈2.4×10−9\bar\alpha_{1000} \approx 2.4 \times 10^{-9}.

Рисунок

Четыре расписания шума: как убывает α̅ₜ

Альфа с чертой для линейного, квадратичного, сигмоидного и косинусного расписанийЧетыре убывающие кривые альфа с чертой t на шагах от 0 до 1000, все начинаются в 1. Линейная падает раньше всех и пересекает 0.5 при t, равном 260; квадратичная пересекает 0.5 при 419, сигмоидная — при 440, косинусная — при 497 и дольше всех остаётся заметно выше нуля.00.250.50.75102004006008001000tα̅tлинейноеквадратичноесигмоидноекосинусноеАльфа с чертой для линейного, квадратичного, сигмоидного и косинусного расписанийЧетыре убывающие кривые альфа с чертой t на шагах от 0 до 1000, все начинаются в 1. Линейная падает раньше всех и пересекает 0.5 при t, равном 260; квадратичная пересекает 0.5 при 419, сигмоидная — при 440, косинусная — при 497 и дольше всех остаётся заметно выше нуля.00.250.50.75102505007501000tα̅tлинейноеквадратичноесигмоидноекосинусное
α̅ₜ для четырёх расписаний при T = 1000: линейного, квадратичного и сигмоидного с β от 10⁻⁴ до 0.02 и косинусного с s = 0.008. Точки отмечают, где α̅ₜ впервые опускается ниже 0.5: t = 260, 419, 440 и 497. Начиная с t = 674, линейное расписание оставляет меньше 1% дисперсии сигнала; косинусное убывает почти линейно в середине цепочки и медленно у её концов.

При T=1000T = 1000:

расписаниеαˉt<0.5\bar\alpha_t < 0.5 начиная сαˉt<0.01\bar\alpha_t < 0.01 начиная сαˉ1000\bar\alpha_{1000}
линейноеt=260t = 260t=674t = 6744.0×10−54.0 \times 10^{-5}
квадратичноеt=419t = 419t=851t = 8517.3×10−47.3 \times 10^{-4}
сигмоидноеt=440t = 440t=721t = 7214.0×10−54.0 \times 10^{-5}
косинусноеt=497t = 497t=936t = 9362.4×10−92.4 \times 10^{-9}

Линейное расписание разрушает изображение быстрее всех: половина дисперсии сигнала исчезает к t=260t = 260, а начиная с t=674t = 674 остаётся меньше 1%1\%, так что последняя треть цепочки работает почти с чистым шумом. Nichol and Dhariwal сочли это расточительным при 32×3232 \times 32 и 64×6464 \times 64: модель, обученная с линейным расписанием, почти не проигрывает в FID, если пропустить первые 20%20\% обратного процесса — его самый зашумлённый конец. Их косинусное расписание построено так, чтобы αˉt\bar\alpha_t убывало почти линейно в середине и мало менялось у обоих концов; оно пересекает 0.50.5 близко к T/2T/2. На ImageNet 64×6464 \times 64 (T=4000T = 4000, LsimpleL_{\text{simple}}, 200K итераций) оно улучшило FID с 31.331.3 до 27.027.0, а правдоподобие — с 3.773.77 до 3.683.68 бита на размерность (bits per dimension; их таблица 1). Квадратичное расписание показывает обратный риск: оно заканчивается на αˉ1000≈7.3×10−4\bar\alpha_{1000} \approx 7.3 \times 10^{-4}, так что в xTx_T сохраняется доля αˉT≈0.027\sqrt{\bar\alpha_T} \approx 0.027 изображения, и q(xT∣x0)q(x_T \mid x_0) лежит дальше от N(0,I)\mathcal N(0, I), с которого начинается сэмплирование, чем у трёх других расписаний.

Ноутбук#

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

В ноутбуке взята игрушечная постановка: двумерный «швейцарский рулет» (Swiss roll) и T=100T = 100 шагов. Первая демонстрация запускает пошаговую цепочку с постоянным β=0.0005\beta = 0.0005, для которого αˉ100≈0.951\bar\alpha_{100} \approx 0.951. Остальная часть ноутбука использует сигмоидное расписание с началом 10−510^{-5} и концом 10−210^{-2}, для которого αˉ100≈0.605\bar\alpha_{100} \approx 0.605; поскольку сигмоида от ±6\pm 6 не равна в точности 00 или 11, фактические βt\beta_t меняются примерно от 3.5×10−53.5\times10^{-5} до 1.0×10−21.0\times10^{-2}. Ноутбук сэмплирует xtx_t напрямую по формуле в замкнутом виде (q_sample, alphas_bar_sqrt, one_minus_alphas_bar_sqrt). Ни одно из игрушечных расписаний не подходит близко к αˉ=0\bar\alpha = 0: коэффициент среднего αˉ100\sqrt{\bar\alpha_{100}} всё ещё около 0.980.98 и 0.780.78 соответственно, так что x100x_{100} несёт значительную часть сигнала данных. Реальным моделям нужно намного больше шагов или более крупные β\beta.

  • Ячейки 4–6 сэмплируют двумерный «швейцарский рулет», который используется во всём ноутбуке.
  • Ячейки 8–10: текстовая ячейка 8 вводит общее марковское переходное ядро Tπ(xt∣xt−1;βt)T_\pi(x_t \mid x_{t-1}; \beta_t) из Sohl-Dickstein et al. и затем выбирает гауссовское. forward_process реализует xt=1−βt xt−1+βt εx_t = \sqrt{1-\beta_t}\,x_{t-1} + \sqrt{\beta_t}\,\varepsilon, где β\beta — дисперсия, а β\sqrt\beta — амплитуда шума, и запускает его на 100 шагов с постоянным β=0.0005\beta = 0.0005. Возвращаемый список начинается с чистых данных под индексом 0.
  • Ячейка 12: make_beta_schedule с линейным, квадратичным и сигмоидным расписаниями, записанными выше (по умолчанию T=1000T = 1000, начало 10−510^{-5}, конец 10−210^{-2}; ноутбук вызывает её с T=100T = 100). При этих крайних значениях и T=100T = 100 значение αˉ100\bar\alpha_{100} равно 0.6050.605 для линейного и сигмоидного расписаний и 0.7070.707 — для квадратичного. Косинусное расписание не реализовано.
  • Ячейки 13–16: сэмплер q_sample по формуле в замкнутом виде с сигмоидным расписанием, который использует alphas_bar_sqrt и one_minus_alphas_bar_sqrt. Вспомогательная функция extract(a, t, x) выбирает коэффициент a[t] для шага каждого примера и меняет его форму так, чтобы он распространялся (broadcasting) на размерности данных. Индекс коэффициента ii в коде соответствует t=i+1t = i + 1 в формулах, поэтому подпись графика x0x_0 уже относится к первому зашумлённому маргинальному распределению. Это сэмплы, взятые напрямую из q(xt∣x0)q(x_t \mid x_0) для каждого tt, а не одна коррелированная траектория.

Обратные сэмплеры (ячейки 24 и 47) начинают с N(0,I)\mathcal N(0, I), но при сигмоидном расписании q(x100∣x0)=N(0.78 x0,  0.40 I)q(x_{100} \mid x_0) = \mathcal N(0.78\,x_0,\; 0.40\,I), так что начальное распределение не совпадает с концом прямой цепочки. Обратный процесс, целевая функция обучения и сэмплер — в ячейках 17–48; они описаны в следующей заметке.

Дальше#

Теперь у нас есть диффузионный процесс, который превращает любое изображение в шум, на практике неотличимый от N(0,I)\mathcal N(0, I); для расписания DDPM αˉT≈4×10−5\bar\alpha_T \approx 4\times10^{-5}. Чтобы генерировать, этот процесс нужно обратить, и делается это с помощью глубокой нейронной сети. Как поставить задачу для этой сети и обучить её — тема заметки DDPM: обратный процесс и целевая функция обучения.

Литература#