Диффузионная модель

Коротко

Definition

Диффузионная модель — это генеративная модель, которая учится создавать данные, начиная с шума и постепенно удаляя из него шум до получения осмысленного объекта.

Диффузионные модели особенно известны в генерации изображений, но сама идея применима и к другим типам данных.

Модель состоит из двух процессов:

1. Прямой процесс — постепенно зашумляет реальные данные.
2. Обратный процесс — учится постепенно убирать шум и восстанавливать данные.

Интуитивно:

$$\text{real data}\to \text{noise}$$

а затем:

$$\text{noise}\to \text{generated data}$$

Интуиция

Диффузионную модель можно понимать как обучение «машины денойзинга».

Во время обучения мы берём реальный объект, добавляем к нему шум и просим нейросеть предсказать, какой шум был добавлен.

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

Пример для изображения:

1. Берём настоящее изображение.
2. Постепенно добавляем гауссовский шум, пока изображение не станет почти чистым шумом.
3. Обучаем нейросеть предсказывать добавленный шум.
4. Для генерации стартуем со случайного шума.
5. Пошагово удаляем шум.
6. Получаем новое изображение.

Формальное описание

Диффузионная модель задаёт последовательность скрытых состояний:

$$x_0,x_1,x_2,\dots ,x_T$$

где:

• $x_0$ — исходные реальные данные;
• $x_t$ — данные после $t$ шагов зашумления;
• $x_T$ — почти чистый гауссовский шум.

Прямой процесс обычно фиксирован и не обучается.

Обратный процесс параметризуется нейросетью, например U-Net:

$${ϵ}_{\theta }(x_t,t)$$

Эта сеть получает зашумленный объект $x_t$ и timestep $t$, а затем предсказывает шум, который нужно убрать.

Входы и выходы

Во время обучения вход модели:

• зашумленный объект $x_t$;
• timestep $t$;
• иногда условие $c$, например текстовый prompt, класс или другой conditioning signal.

Выход модели чаще всего:

• предсказанный шум ${ϵ}_{\theta }(x_t,t)$.

В conditional diffusion output может зависеть от условия:

$${ϵ}_{\theta }(x_t,t,c)$$

Например, в text-to-image модели условием является текстовое описание изображения.

Прямой процесс

Прямой процесс постепенно добавляет гауссовский шум к исходным данным.

На каждом шаге $t$ объект $x_{t-1}$ превращается в более шумный объект $x_t$:

$$x_t=\sqrt{1-{\beta }_t}\cdot x_{t-1}+\sqrt{{\beta }_t}\cdot ϵ$$

где:

• $ϵ\sim \mathcal{N}(0,I)$ — случайный гауссовский шум;
• ${\beta }_t$ — величина добавляемого шума на шаге $t$;
• $x_t$ — объект после очередного шага зашумления.

Расписание ${\beta }_t$ называют noise schedule.

Note

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

Прыжок на произвольный timestep

Ключевое свойство прямого процесса: можно получить $x_t$ напрямую из $x_0$, не вычисляя все промежуточные шаги.

Обозначим:

$${\alpha }_t=1-{\beta }_t$$

и:

$${\bar{\alpha }}_t=\prod _{i=1}^t{\alpha }_i$$

Тогда:

$$x_t=\sqrt{{\bar{\alpha }}_t}\cdot x_0+\sqrt{1-{\bar{\alpha }}_t}\cdot ϵ$$

Это сильно ускоряет обучение: можно случайно выбрать timestep $t$, сразу получить зашумленный объект $x_t$ и обучать модель предсказывать добавленный шум.

Пример прямого процесса

Пусть исходный объект — 2D-вектор:

$$x_0=[0.8,-0.6]$$

Пусть для первого шага:

$${\beta }_1=0.1$$

Тогда:

$${\alpha }_1=1-0.1=0.9$$

и:

$${\bar{\alpha }}_1=0.9$$

Сэмплируем шум:

$$ϵ=[0.5,1.2]$$

Используем формулу прыжка:

$$x_1=\sqrt{0.9}\cdot [0.8,-0.6]+\sqrt{0.1}\cdot [0.5,1.2]$$

Приближённо:

$$x_1=0.95\cdot [0.8,-0.6]+0.32\cdot [0.5,1.2]$$ $$x_1=[0.76,-0.57]+[0.16,0.38]=[0.92,-0.19]$$

Вектор $x_1$ — это исходный объект с небольшим добавлением шума.

Непрерывное время и SDE

Дискретный процесс можно рассматривать как приближение непрерывного процесса.

Если число шагов $T$ стремится к бесконечности, а размер шага становится очень маленьким, процесс можно описывать стохастическим дифференциальным уравнением:

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

где:

• $f(x,t)dt$ — drift, детерминированная часть движения;
• $g(t)dw$ — diffusion, стохастическая часть с гауссовским шумом;
• $dw$ — винеровский процесс, или броуновское движение.

Это непрерывное описание полезно для связи diffusion models со score-based generative models и flow matching.

Для базового понимания достаточно помнить: дискретная диффузия — это пошаговое зашумление, а SDE — его непрерывный аналог.

Обратный процесс

Обратный процесс должен восстановить данные из шума:

$$x_T\to x_{T-1}\to \cdots \to x_0$$

Проблема: прямой процесс известен, а обратный нужно выучить.

Модель должна понять, как из зашумленного объекта $x_t$ получить менее зашумленный $x_{t-1}$.

На практике нейросеть часто учат предсказывать шум:

$${ϵ}_{\theta }(x_t,t)$$

После этого предсказанный шум используется, чтобы сделать шаг денойзинга.

Score function

Для обращения процесса полезна score function:

$${\nabla }_x\log p_t(x)$$

Это градиент логарифма плотности распределения данных на timestep $t$.

Интуитивно score function показывает, в какую сторону нужно сдвинуть зашумленную точку, чтобы она стала более похожа на данные из текущего распределения.

Можно представить плотность данных как ландшафт. Score function указывает направление подъёма к областям с большей вероятностью.

Диффузионная модель может быть понята как модель, которая учится оценивать это направление.

Почему предсказание шума связано со score matching

В прямом процессе $x_t$ получается из $x_0$ добавлением гауссовского шума:

$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$$

Для гауссовского распределения score имеет удобное выражение:

$${\nabla }_{x_t}\log p(x_t\mid x_0)=-\frac{ϵ}{\sqrt{1-{\bar{\alpha }}_t}}$$

То есть score пропорционален добавленному шуму $ϵ$.

Поэтому обучение модели предсказывать шум и обучение модели предсказывать score тесно связаны.

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

Timestep embedding

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

Поэтому timestep $t$ подают в сеть как embedding.

Часто используют sinusoidal или harmonic embeddings, похожие на positional encodings в трансформерах.

Этот embedding затем добавляется или подмешивается в блоки нейросети, например в residual blocks внутри U-Net.

Иначе одна и та же сеть не сможет отличить слегка зашумленный объект от почти полного шума.

Как обучается

Типовой training loop:

1. Взять реальный объект $x_0$.
2. Случайно выбрать timestep $t$.
3. Сэмплировать шум $ϵ\sim \mathcal{N}(0,I)$.
4. Получить зашумленный объект:

$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ$$

5. Подать $x_t$ и $t$ в нейросеть.
6. Получить предсказание шума ${ϵ}_{\theta }(x_t,t)$.
7. Посчитать MSE между настоящим и предсказанным шумом.
8. Обновить параметры модели через Обратное распространение ошибки.

Функция потерь

Частая функция потерь в DDPM-подобных моделях — MSE между реальным шумом и предсказанным шумом:

$$\mathcal{L}={\mathbb{E}}_{t,x_0,ϵ}\left[{\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t)\Vert }^2\right]$$

Здесь:

• $x_0$ — реальный объект;
• $t$ — timestep;
• $ϵ$ — настоящий добавленный шум;
• ${ϵ}_{\theta }$ — шум, предсказанный моделью.

Пример:

• вход сети: $x_1=[0.92,-0.19]$ и $t=1$;
• истинный шум: $ϵ=[0.5,1.2]$;
• предсказанный шум: ${ϵ}_{\theta }=[0.2,0.9]$.

MSE:

$$(0.5-0.2{)}^2+(1.2-0.9{)}^2=0.09+0.09=0.18$$

Этот loss используется для обновления весов нейросети.

Генерация

После обучения можно генерировать новые объекты.

Процедура sampling:

1. Сэмплировать чистый гауссовский шум:

$$x_T\sim \mathcal{N}(0,I)$$

2. Для $t=T,T-1,\dots ,1$:

   • подать $x_t$ и $t$ в модель;
   • получить предсказанный шум ${ϵ}_{\theta }(x_t,t)$;
   • вычислить менее зашумленный $x_{t-1}$.

3. Получить финальный объект $x_0$.

Один из вариантов обратного шага:

$$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}\left(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_{\theta }(x_t,t)\right)+{\sigma }_{t}z$$

где:

• $z\sim \mathcal{N}(0,I)$ — дополнительный шум;
• ${\sigma }_t$ — уровень шума на обратном шаге.

В детерминированных сэмплерах дополнительный шум может быть равен нулю.

Пример генерации

Пусть модель уже обучена, и нужно сгенерировать новый 2D-вектор. Пусть $T=2$.

Старт:

$$x_2=[-1.5,0.8]$$

Шаг $t=2$:

• модель предсказывает шум:

$${ϵ}_{\theta }(x_2,2)=[-1.3,0.5]$$

• по обратной формуле получаем менее зашумленный вектор:

$$x_1\approx [0.4,0.7]$$

Шаг $t=1$:

• модель предсказывает:

$${ϵ}_{\theta }(x_1,1)=[0.2,0.9]$$

• получаем:

$$x_0\approx [0.85,-0.55]$$

Итоговый $x_0$ — новый сгенерированный объект, похожий на данные из обучающего распределения.

Архитектура

В диффузионных моделях для изображений часто используют U-Net.

Почему U-Net удобен:

• encoder постепенно сжимает spatial resolution;
• decoder восстанавливает resolution;
• skip connections сохраняют пространственные детали;
• residual blocks помогают обучению;
• timestep embedding подмешивается в блоки;
• attention blocks помогают учитывать глобальный контекст.

Связанные блоки:

• Свёрточный слой;
• Skip connection;
• Батч-нормализация;
• Attention;
• Трансформер.

В современных text-to-image моделях U-Net или transformer-like denoiser часто получает не только $x_t$ и $t$, но и conditioning, например текстовый embedding.

Conditional diffusion

Диффузионная модель может быть условной.

Тогда генерация зависит от условия $c$:

$${ϵ}_{\theta }(x_t,t,c)$$

Условие может быть:

• текстовым prompt;
• классом;
• изображением;
• маской;
• low-resolution изображением;
• свойствами объекта.

Примеры:

• text-to-image generation;
• image inpainting;
• super-resolution;
• class-conditional generation;
• генерация молекул или материалов с заданными свойствами.

Гиперпараметры

Основные гиперпараметры:

• число diffusion steps $T$;
• noise schedule ${\beta }_t$;
• архитектура denoiser;
• размер timestep embedding;
• learning rate;
• batch size;
• optimizer;
• способ conditioning;
• sampling method;
• classifier-free guidance scale;
• число sampling steps;
• loss weighting.

Количество sampling steps влияет на качество и скорость генерации. Больше шагов обычно даёт лучшее качество, но медленнее.

Когда использовать

Диффузионные модели стоит использовать, если:

• нужна генерация сложных данных;
• важны качество и разнообразие сэмплов;
• есть достаточно данных и вычислений;
• задача связана с изображениями, аудио, видео или structured generation;
• требуется conditional generation;
• нужна современная альтернатива GAN и VAE.

Диффузионные модели часто обучаются стабильнее, чем GAN, и лучше покрывают распределение данных.

Когда не использовать

Диффузионная модель может быть плохим выбором, если:

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

Недостаток классических diffusion models — медленный sampling, потому что генерация требует многих последовательных шагов.

Метрики оценки

Для изображений часто используют:

• FID;
• Inception Score;
• CLIP-based metrics для text-to-image;
• human evaluation;
• diversity metrics;
• domain-specific metrics.

Также важно проверять:

• разнообразие;
• артефакты;
• соответствие условию;
• устойчивость генерации;
• редкие случаи;
• bias в данных.

Loss обучения не всегда напрямую соответствует субъективному качеству сгенерированных объектов.

Типичные ошибки понимания

Думать, что модель учится добавлять шум

Прямой процесс зашумления фиксирован и не обучается. Модель учится обратному процессу — удалять шум.

Путать $x_t$ и $ϵ$

$x_t$ — зашумленный объект. $ϵ$ — шум, который был добавлен к исходному объекту.

Модель часто предсказывает именно $ϵ$, а не сразу чистый $x_0$.

Думать, что все шаги прямого процесса нужно считать последовательно

Во время обучения можно напрямую получить $x_t$ из $x_0$ через формулу с ${\bar{\alpha }}_t$.

Считать diffusion тем же самым, что VAE

VAE учит latent distribution и decoder. Diffusion учит обратный процесс денойзинга из шума.

Игнорировать timestep embedding

Без timestep модель не знает, сколько шума в объекте, и не может корректно денойзить разные уровни шума.

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

Классический sampling может быть медленным. Для ускорения используют DDIM, distillation, fewer-step samplers и другие методы.

Минимальный пример

Пусть объект — число $x_0=1.0$.

Добавляем шум на timestep $t$:

$$x_t=0.8x_0+0.6ϵ$$

Если шум:

$$ϵ=-0.5$$

то:

$$x_t=0.8\cdot 1.0+0.6\cdot (-0.5)=0.5$$

Модель получает $x_t=0.5$ и timestep $t$.

Её задача — предсказать шум:

$${ϵ}_{\theta }(x_t,t)\approx -0.5$$

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

Практические замечания

Хороший workflow:

1. Начать с простой DDPM/DDIM реализации.
2. Проверить нормализацию данных.
3. Выбрать noise schedule.
4. Убедиться, что timestep embedding подаётся в модель.
5. Обучать модель предсказывать шум.
6. Визуально проверять samples во время обучения.
7. Сравнивать разные samplers.
8. Проверять diversity и mode coverage.
9. Для conditional generation проверять соответствие условию.
10. Не оценивать качество только по training loss.

Для изображений важно смотреть реальные сэмплы на разных этапах. Для научных объектов нужно дополнительно проверять валидность сгенерированных структур предметными метриками.

Связанные понятия

• Автоэнкодер
• Сопоставление потоков
• CNN
• Изображения
• Обновление параметров
• Затухание градиента