Как шумоподавляющие диффузионные модели намеренно разрушают данные. В наивной цепочке, которая просто добавляет шум, дисперсия неограниченно растёт, а шаг с масштабированием √(1 − β) x + √β ε даёт переход q(xₜ | x₀) в замкнутом виде, который стремится к N(0, I); на практике β меняется по линейному, квадратичному, сигмоидному или косинусному расписанию.
Заметка о моделях, основанных на энергии, закончилась на цене этого подхода: модель, основанную на энергии (energy-based model), обучают с короткими смещёнными цепочками MCMC, которые запускаются внутри каждого обновления параметров, поэтому выученная модель зависит от сэмплера не меньше, чем от данных. Диффузионные модели выносят сэмплирование из обучения. Они фиксируют простой процесс зашумления, который превращает данные в гауссовский шум, и учатся обращать его по одному малому шагу за раз. Каждый обратный шаг близок к гауссовскому, поэтому сеть выучивает его регрессией на зашумлённых сэмплах, которые фиксированный процесс даёт бесплатно. Оказывается, эта сеть выучивает score зашумлённых данных — градиент энергии со знаком минус и направление, по которому движется сэмплирование динамикой Ланжевена (заметка о score matching). Зашумлённые латентные переменные сохраняют размерность данных, архитектура сети ничем не ограничена, а платой становится нижняя оценка правдоподобия вместо самого правдоподобия. DDPM (Denoising Diffusion Probabilistic Models, шумоподавляющие диффузионные вероятностные модели; Ho, Jain and Abbeel, 2020) — модель, которая довела качество изображений у диффузионных моделей до уровня, сравнимого с GAN. Эта заметка посвящена её первой половине — прямому процессу (forward process): фиксированной процедуре, которая постепенно превращает изображение в гауссовский шум. Вторая половина — нейронная сеть, которая выполняет эту процедуру в обратном направлении, — описана в следующей заметке.
Обозначения.x0 — точка данных (изображение), а xt — её зашумлённая версия после t шагов, t=1,…,T. Через q обозначается фиксированный прямой процесс (процесс зашумления): q(x0) — распределение данных, а q(xt∣xt−1) — один шаг зашумления. pθ(xt−1∣xt) — обучаемый обратный шаг с параметрами сети θ. ε∼N(0,1) — стандартный гауссовский шум, который на каждом шаге сэмплируется заново, а второй аргумент N(μ,σ2) — это дисперсия. Мы работаем в одном измерении; для изображений каждый пиксель зашумляется независимо и 1 превращается в единичную матрицу I, так что целевое распределение — 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.
Идея диффузии — связать q(x0) с распределением, из которого мы умеем сэмплировать. Возьмём реальное изображение x0 — исходную фотографию — и добавим немного шума, чтобы получить x1, затем ещё немного, чтобы получить x2, и так далее, пока xT не станет чистым шумом. Каждый шаг прямого процесса фиксирован и не требует обучения. Модель — обучаемый обратный шаг pθ(xt−1∣xt) — идёт в другую сторону: по xt она строит xt−1, на один шаг ближе к изображению. Чтобы сгенерировать изображение, мы сэмплируем чистый шум xT и применяем модель T раз.
где β — дисперсия шума. Здесь важен квадратный корень: дисперсия cε равна c2, поэтому шум с дисперсией β — это βε. Следующий шаг делает то же самое с x1: q(x2∣x1)=N(x1,β).
Теперь цепочку можно сократить. Два шага добавляют два независимых шума, а дисперсии независимых гауссовских величин складываются:
где ε=(ε1+ε2)/2 — снова стандартная нормальная величина. Всякий раз, когда ниже несколько шумов объединяются в один, ε обозначает объединённую стандартную нормальную величину в этом смысле, так что такие равенства выполняются по распределению.
Продолжая так же, после t шагов получаем
xt=x0+tβε,q(xt∣x0)=N(x0,tβ).
Это полезное сокращение: от x0 можно перейти к любому xt за одно сэмплирование. Но посмотрим на два параметра:
среднее не меняется: оно навсегда остаётся равным x0;
дисперсия неограниченно растёт: tβ→∞.
Поэтому цепочка никогда не приходит к N(0,1): дисперсия растёт без ограничений, и нет фиксированного предельного распределения, с которого могла бы начаться генерация. Влияние изображения ослабевает лишь относительно растущего разброса. Такой процесс называют диффузией с растущей дисперсией (variance-exploding) (Song et al., 2021, §3.4). Модели, построенные на нём, останавливаются на большом конечном уровне шума σmax и начинают сэмплирование из N(0,σmax2) (Song et al., 2021, §2.1).
каким бы ни было начальное изображение. Для этого нужен новый способ переходить от шага t−1 к шагу t; в частности, шаг должен каждый раз менять среднее, притягивая его к 0. В общем случае Sohl-Dickstein et al. допускают любое марковское переходное ядро Tπ(xt∣xt−1;βt), которое оставляет неизменным простое целевое распределение π, π(y)=∫Tπ(y∣y′;β)π(y′)dy′ (2015, §2.1, формулы 1–2): биномиальное ядро для бинарных данных, гауссовское — для вещественных. Для вещественных значений будем искать шаг вида xt=axt−1+βε. Гауссовское переходное ядро диффузионных моделей, которое ввели Sohl-Dickstein et al. (2015, таблица App.1) и сохранила DDPM (Ho, Jain and Abbeel, 2020, формула 2), умножает предыдущую точку на a=1−β:
Отдельные шаги и сокращённые переходы за один скачок для двух цепочек. Обе на каждом шаге добавляют новый гауссовский шум ε ~ N(0, 1), и в обеих t шагов сворачиваются в один гауссовский скачок q(xₜ | x₀), так что xₜ можно сэмплировать напрямую. Без масштабирования (сверху) скачок оставляет среднее равным x₀, а его дисперсия неограниченно растёт. Умножение на √(1 − β) на каждом шаге (снизу) уменьшает среднее и держит условную дисперсию 1 − α̅ₜ ниже 1, так что цепочка приближается к N(0, 1). Второй аргумент N — дисперсия.
Почему именно 1−β? Новый шум не зависит от xt−1, поэтому дисперсия подчиняется рекуррентному соотношению Var(xt)=a2Var(xt−1)+β. Неподвижная точка этого соотношения равна 1, дисперсии целевого распределения, только если a2+β=1. Значит, a=1−β — единственный положительный масштаб, при котором единичная дисперсия остаётся равной 1, и он же притягивает к 1 любую другую дисперсию (см. врезку ниже). При таком a переход от x0 к любому xt принимает особенно простой замкнутый вид, в котором сумма квадратов двух коэффициентов равна 1 (Ho, Jain and Abbeel, 2020, формула 4):
Если положить αˉ0=1, формула охватывает и случай t=0. Значит, xt для любого шага можно получить прямо из x0, за одно сэмплирование, не выполняя t предшествующих шагов.
Если β∈(0,1), то αˉt=(1−β)t→0 при t→∞. Тогда член с изображением αˉtx0→0, а коэффициент шума 1−αˉt→1, так что шумовой член стремится просто к ε с дисперсией 1:
q(xt∣x0)t→∞N(0,1).
Например, при постоянном β=0.02:
t
αˉt=0.98t
коэффициент среднего αˉt
стандартное отклонение шума 1−αˉt
стандартное отклонение шума в наивной цепочке tβ
1
0.980
0.990
0.141
0.141
10
0.817
0.904
0.428
0.447
50
0.364
0.603
0.797
1.000
100
0.133
0.364
0.931
1.414
200
0.0176
0.133
0.991
2.000
При малых t две цепочки почти совпадают. Затем разброс наивной цепочки превышает целевой и продолжает расти, а у DDPM он приближается к 1 снизу. При конечном T член с изображением αˉTx0 мал лишь относительно масштаба x0; это одна из причин, по которым изображения перед диффузией масштабируют в [−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.
То же происходит со всем распределением данных. Его маргинальное распределение на шаге t равно q(xt)=∫q(xt∣x0)q(x0)dx0. Чтобы увидеть, что оно тоже стремится к N(0,1), запишем xt=αˉtx0+1−αˉtε, где x0∼q(x0) не зависит от ε. Первый член — фиксированная случайная величина, умноженная на множитель, который стремится к 0, поэтому он стремится к 0 при любом распределении данных, а второй член стремится к ε; следовательно, xt→N(0,1) по распределению. В одном измерении двугорбая плотность расплывается, её горбы движутся к 0 и сливаются (между t=200 и t=300 для данных на рисунке ниже), и в конце остаётся одна колоколообразная кривая с центром в 0. Конечный батч сэмплов по пути выглядит зубчатым, но сама плотность остаётся гладкой: это свёртка плотности данных с гауссианой. Для смеси гауссиан маргинальное распределение остаётся смесью и записывается точно:
Обе компоненты притягиваются к среднему 0 и дисперсии 1, поэтому сливаются в одно стандартное нормальное распределение.
Рисунок
Распределение данных растворяется в N(0, 1)
Каждая точка данных зашумляется по одному и тому же правилу, поэтому вся одномерная плотность данных q(x₀) — смесь двух гауссиан — превращается в маргинальное распределение q(xₜ) = Σₖ πₖ N(√α̅ₜ μₖ, α̅ₜσₖ² + 1 − α̅ₜ). Кривые — точные плотности; столбцы — гистограмма 600 сэмплов, пропущенных через те же шаги, поэтому конечный батч выглядит зубчатым. Используется линейное расписание β из DDPM от 10⁻⁴ до 0.02 (T = 1000).
На практике β не постоянна, а меняется от шага к шагу: β1,…,βT; это расписание шума (noise schedule). Вывод выше проходит без изменений, если на шаге t заменить (1−β) на (1−βt):
Обычно используют сокращение αt=1−βt, так что αˉt=∏i≤tαi. В статье DDPM берутся T=1000 шагов и линейное расписание от β1=10−4 до βT=0.02 (Ho, Jain and Abbeel, 2020, §4). Позже обнаружилось, что косинусное расписание, которое разрушает информацию равномернее во времени, лучше работает на низком разрешении (Nichol and Dhariwal, 2021, §3.2). Малые шаги в начале оставляют ранние xt близкими к изображению; более крупные шаги потом доводят дело до конца. При линейном расписании DDPM αˉt примерно равно 0.5 при t≈259 (ниже 0.5 оно опускается при t=260), падает ниже 0.01 при t=674 и заканчивается на αˉ1000≈4.0×10−5, так что на практике x1000 — просто шум.
Рисунок
Линейное расписание β и соответствующее α̅ₜ
Сверху: βₜ линейно растёт от 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): квадрат коэффициента при изображении, делённый на дисперсию шума; для данных с единичной дисперсией это отношение двух дисперсий. Для линейного расписания оно падает примерно от 104 при t=1 до 4×10−5 при t=T. Быстрые сэмплеры из заметки о латентной диффузии работают с его логарифмом.
Линейное, квадратичное, сигмоидное и косинусное расписания#
Линейное расписание — лишь один из нескольких вариантов (их обзор — у Weng, 2021). Обозначим через ut=(t−1)/(T−1)∈[0,1] положение в цепочке. При тех же крайних значениях β1=10−4 и βT=0.02 три распространённые формы задаются через βt:
Квадратичное расписание линейно интерполирует амплитуду шума βt, поэтому его шаги дольше остаются малыми. Сигмоидное расписание плоское на обоих концах и меняется в середине; поскольку sigmoid(±6) не равна в точности 0 или 1, его βt на самом деле меняются от 1.5×10−4 до 0.01995. Косинусное расписание (Nichol and Dhariwal, 2021, §3.2, формула 17) задаёт не шаги, а αˉt и выводит шаги из него:
со сдвигом s=0.008, выбранным так, чтобы амплитуда шума на первом шаге была чуть меньше ширины 1/127.5 одного уровня яркости пикселя, и с βt, обрезанными сверху на 0.999. Без обрезки f(T)=0 дало бы βT=1; с ней αˉ1000≈2.4×10−9.
Рисунок
Четыре расписания шума: как убывает α̅ₜ
α̅ₜ для четырёх расписаний при T = 1000: линейного, квадратичного и сигмоидного с β от 10⁻⁴ до 0.02 и косинусного с s = 0.008. Точки отмечают, где α̅ₜ впервые опускается ниже 0.5: t = 260, 419, 440 и 497. Начиная с t = 674, линейное расписание оставляет меньше 1% дисперсии сигнала; косинусное убывает почти линейно в середине цепочки и медленно у её концов.
При T=1000:
расписание
αˉt<0.5 начиная с
αˉt<0.01 начиная с
αˉ1000
линейное
t=260
t=674
4.0×10−5
квадратичное
t=419
t=851
7.3×10−4
сигмоидное
t=440
t=721
4.0×10−5
косинусное
t=497
t=936
2.4×10−9
Линейное расписание разрушает изображение быстрее всех: половина дисперсии сигнала исчезает к t=260, а начиная с t=674 остаётся меньше 1%, так что последняя треть цепочки работает почти с чистым шумом. Nichol and Dhariwal сочли это расточительным при 32×32 и 64×64: модель, обученная с линейным расписанием, почти не проигрывает в FID, если пропустить первые 20% обратного процесса — его самый зашумлённый конец. Их косинусное расписание построено так, чтобы αˉt убывало почти линейно в середине и мало менялось у обоих концов; оно пересекает 0.5 близко к T/2. На ImageNet 64×64 (T=4000, Lsimple, 200K итераций) оно улучшило FID с 31.3 до 27.0, а правдоподобие — с 3.77 до 3.68 бита на размерность (bits per dimension; их таблица 1). Квадратичное расписание показывает обратный риск: оно заканчивается на αˉ1000≈7.3×10−4, так что в xT сохраняется доля αˉT≈0.027 изображения, и q(xT∣x0) лежит дальше от N(0,I), с которого начинается сэмплирование, чем у трёх других расписаний.
Ноутбук по диффузионным моделям из репозитория HSE-LAMBDA Deep Generative Models (ссылка закреплена на конкретном коммите; ячейки нумеруются с 1, считая все ячейки подряд) основан на руководстве по диффузионным моделям ACIDS-IRCAM. Сохранённых выходов ячеек в нём нет, поэтому он документирует процедуры, а не результаты.
В ноутбуке взята игрушечная постановка: двумерный «швейцарский рулет» (Swiss roll) и T=100 шагов. Первая демонстрация запускает пошаговую цепочку с постоянным β=0.0005, для которого αˉ100≈0.951. Остальная часть ноутбука использует сигмоидное расписание с началом 10−5 и концом 10−2, для которого αˉ100≈0.605; поскольку сигмоида от ±6 не равна в точности 0 или 1, фактические βt меняются примерно от 3.5×10−5 до 1.0×10−2. Ноутбук сэмплирует xt напрямую по формуле в замкнутом виде (q_sample, alphas_bar_sqrt, one_minus_alphas_bar_sqrt). Ни одно из игрушечных расписаний не подходит близко к αˉ=0: коэффициент среднего αˉ100 всё ещё около 0.98 и 0.78 соответственно, так что x100 несёт значительную часть сигнала данных. Реальным моделям нужно намного больше шагов или более крупные β.
Ячейки 4–6 сэмплируют двумерный «швейцарский рулет», который используется во всём ноутбуке.
Ячейки 8–10: текстовая ячейка 8 вводит общее марковское переходное ядро Tπ(xt∣xt−1;βt) из Sohl-Dickstein et al. и затем выбирает гауссовское. forward_process реализует xt=1−βtxt−1+βtε, где β — дисперсия, а β — амплитуда шума, и запускает его на 100 шагов с постоянным β=0.0005. Возвращаемый список начинается с чистых данных под индексом 0.
Ячейка 12: make_beta_schedule с линейным, квадратичным и сигмоидным расписаниями, записанными выше (по умолчанию T=1000, начало 10−5, конец 10−2; ноутбук вызывает её с T=100). При этих крайних значениях и T=100 значение αˉ100 равно 0.605 для линейного и сигмоидного расписаний и 0.707 — для квадратичного. Косинусное расписание не реализовано.
Ячейки 13–16: сэмплер q_sample по формуле в замкнутом виде с сигмоидным расписанием, который использует alphas_bar_sqrt и one_minus_alphas_bar_sqrt. Вспомогательная функция extract(a, t, x) выбирает коэффициент a[t] для шага каждого примера и меняет его форму так, чтобы он распространялся (broadcasting) на размерности данных. Индекс коэффициента i в коде соответствует t=i+1 в формулах, поэтому подпись графика x0 уже относится к первому зашумлённому маргинальному распределению. Это сэмплы, взятые напрямую из q(xt∣x0) для каждого t, а не одна коррелированная траектория.
Обратные сэмплеры (ячейки 24 и 47) начинают с N(0,I), но при сигмоидном расписании q(x100∣x0)=N(0.78x0,0.40I), так что начальное распределение не совпадает с концом прямой цепочки. Обратный процесс, целевая функция обучения и сэмплер — в ячейках 17–48; они описаны в следующей заметке.
Теперь у нас есть диффузионный процесс, который превращает любое изображение в шум, на практике неотличимый от N(0,I); для расписания DDPM αˉT≈4×10−5. Чтобы генерировать, этот процесс нужно обратить, и делается это с помощью глубокой нейронной сети. Как поставить задачу для этой сети и обучить её — тема заметки DDPM: обратный процесс и целевая функция обучения.