Генерация изображений — диффузионные модели

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

Тип: Сборка Языки: Python Предварительные требования: Фаза 4, урок 07 (U-Net), фаза 1, урок 06 (Вероятность), фаза 3, урок 06 (Оптимизаторы) Время: ~75 минут

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

• Вывести прямой процесс зашумления x_0 -> x_1 -> ... -> x_T и объяснить, почему замкнутая формула q(x_t | x_0) работает для любого t
• Реализовать цель обучения в стиле DDPM, регрессирующую добавленный на каждом шаге шум, и сэмплер, который проходит от чистого шума к изображению
• Построить U-Net с условием по времени (достаточно маленькую для обучения на CPU), предсказывающую шум для любого временного шага
• Объяснить разницу между сэмплированием DDPM и DDIM и случаи применения каждого метода (урок 23 подробно рассматривает сопоставление потоков и rectified flow)

Проблема

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

Помимо устойчивости обучения, именно итеративная структура диффузии открывает всё, что делает современная генерация изображений: обусловливание текстом, inpainting, редактирование изображений, сверхразрешение, управляемый стиль. Каждый шаг цикла сэмплирования — место, куда можно внедрить новое ограничение. Этот крючок — причина, по которой Stable Diffusion, Imagen, DALL-E 3, Midjourney и каждая управляемая модель изображений, которой вы будете пользоваться, основаны на диффузии.

В этом уроке строится минимальная DDPM: прямое зашумление, обратное устранение шума, цикл обучения. Следующий урок (Stable Diffusion) подключает её к производственной системе с VAE, текстовым кодировщиком и classifier-free guidance.

Концепция

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

Возьмём изображение x_0. Добавим небольшое количество гауссовского шума, чтобы получить x_1. Добавим ещё немного, чтобы получить x_2. Продолжим T шагов, пока x_T почти невозможно будет отличить от чистого гауссовского шума.

q(x_t | x_{t-1}) = N(x_t; sqrt(1 - beta_t) * x_{t-1},  beta_t * I)

beta_t — малое расписание дисперсии, обычно линейное от 0.0001 до 0.02 на T=1000 шагов. Каждый шаг немного уменьшает сигнал и добавляет свежий шум.

Замкнутый переход

Добавление шума по одному шагу — цепь Маркова, но математику можно свернуть: вы можете сэмплировать x_t непосредственно из x_0 за один шаг.

Define alpha_t = 1 - beta_t
Define alpha_bar_t = prod_{s=1..t} alpha_s

Then:
  q(x_t | x_0) = N(x_t; sqrt(alpha_bar_t) * x_0,  (1 - alpha_bar_t) * I)

Equivalently:
  x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
  where epsilon ~ N(0, I)

Это единственное уравнение — вся причина практичности диффузии. Во время обучения вы выбираете случайное t, сэмплируете x_t непосредственно из x_0 и обучаете модель за один шаг — не требуется моделировать всю цепь Маркова.

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

Прямой процесс фиксирован. Обратный процесс p(x_{t-1} | x_t) — то, чему учится нейронная сеть. Диффузионные модели не предсказывают x_{t-1} напрямую; они предсказывают шум epsilon, добавленный на шаге t, а математика выводит из него x_{t-1}.

Потеря обучения

Для каждого шага обучения:

1. Сэмплируйте реальное изображение x_0.
2. Сэмплируйте временной шаг t равномерно из [1, T].
3. Сэмплируйте шум epsilon ~ N(0, I).
4. Вычислите x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon.
5. Предскажите epsilon_theta(x_t, t) сетью.
6. Минимизируйте || epsilon - epsilon_theta(x_t, t) ||^2.

Вот и всё. Нейронная сеть учится предсказывать шум на любом временном шаге. Функция потерь — MSE. Нет состязательной игры, коллапса или осцилляций.

Сэмплер (DDPM)

Чтобы генерировать, начните с x_T ~ N(0, I) и идите назад по одному шагу.

for t = T, T-1, ..., 1:
    eps = model(x_t, t)
    x_{t-1} = (1 / sqrt(alpha_t)) * (x_t - (beta_t / sqrt(1 - alpha_bar_t)) * eps) + sqrt(beta_t) * z
    where z ~ N(0, I) if t > 1, else 0
return x_0

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

Почему 1000 шагов

Прямое расписание шума выбирается так, чтобы каждый шаг добавлял ровно столько шума, что обратный шаг был почти гауссовским. При слишком малом числе шагов обратный шаг далёк от гауссовского, и сеть не может хорошо его моделировать. При слишком большом числе шагов сэмплирование становится дорогим с убывающей отдачей. T=1000 с линейным расписанием — настройка DDPM по умолчанию.

DDIM: сэмплирование в 20 раз быстрее

Обучение остаётся тем же. Меняется сэмплирование. DDIM (Song и др., 2020) задаёт детерминированный обратный процесс, который пропускает временные шаги без переобучения. Сэмплирование DDIM за 50 шагов даёт качество, близкое к 1000-шаговой DDPM. Каждая производственная система использует DDIM или ещё более быстрый вариант (DPM-Solver, Euler ancestral).

Условие по времени

Сеть epsilon_theta(x_t, t) должна знать, на каком временном шаге она устраняет шум. Современные диффузионные модели внедряют t через синусоидальные временные встраивания (та же идея, что позиционное кодирование в трансформерах), добавляемые к картам признаков на каждом уровне U-Net.

t_embedding = sinusoidal(t)
feature_map += MLP(t_embedding)

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

Соберите

Шаг 1: Расписание шума

import torch

def linear_beta_schedule(T=1000, beta_start=1e-4, beta_end=2e-2):
    return torch.linspace(beta_start, beta_end, T)


def precompute_schedule(betas):
    alphas = 1.0 - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    return {
        "betas": betas,
        "alphas": alphas,
        "alphas_cumprod": alphas_cumprod,
        "sqrt_alphas_cumprod": torch.sqrt(alphas_cumprod),
        "sqrt_one_minus_alphas_cumprod": torch.sqrt(1.0 - alphas_cumprod),
        "sqrt_recip_alphas": torch.sqrt(1.0 / alphas),
    }

schedule = precompute_schedule(linear_beta_schedule(T=1000))

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

Шаг 2: Прямая диффузия (q_sample)

def q_sample(x0, t, noise, schedule):
    sqrt_a = schedule["sqrt_alphas_cumprod"][t].view(-1, 1, 1, 1)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"][t].view(-1, 1, 1, 1)
    return sqrt_a * x0 + sqrt_one_minus_a * noise

Однострочная замкнутая формула. t — пакет временных шагов, по одному на изображение в пакете.

Шаг 3: Небольшая U-Net с условием по времени

import torch.nn as nn
import torch.nn.functional as F
import math

def timestep_embedding(t, dim=64):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    emb = torch.cat([args.sin(), args.cos()], dim=-1)
    return emb


class TinyUNet(nn.Module):
    def __init__(self, img_channels=3, base=32, t_dim=64):
        super().__init__()
        self.t_mlp = nn.Sequential(
            nn.Linear(t_dim, base * 4),
            nn.SiLU(),
            nn.Linear(base * 4, base * 4),
        )
        self.t_dim = t_dim
        self.enc1 = nn.Conv2d(img_channels, base, 3, padding=1)
        self.enc2 = nn.Conv2d(base, base * 2, 4, stride=2, padding=1)
        self.mid = nn.Conv2d(base * 2, base * 2, 3, padding=1)
        self.dec1 = nn.ConvTranspose2d(base * 2, base, 4, stride=2, padding=1)
        self.dec2 = nn.Conv2d(base * 2, img_channels, 3, padding=1)
        self.time_proj = nn.Linear(base * 4, base * 2)

    def forward(self, x, t):
        t_emb = timestep_embedding(t, self.t_dim)
        t_emb = self.t_mlp(t_emb)
        t_proj = self.time_proj(t_emb)[:, :, None, None]

        h1 = F.silu(self.enc1(x))
        h2 = F.silu(self.enc2(h1)) + t_proj
        h3 = F.silu(self.mid(h2))
        d1 = F.silu(self.dec1(h3))
        d2 = torch.cat([d1, h1], dim=1)
        return self.dec2(d2)

Двухуровневая U-Net с условием по времени, внедрённым в узкое место. Для реальных изображений увеличивайте глубину и ширину.

Шаг 4: Цикл обучения

def train_step(model, x0, schedule, optimizer, device, T=1000):
    model.train()
    x0 = x0.to(device)
    bs = x0.size(0)
    t = torch.randint(0, T, (bs,), device=device)
    noise = torch.randn_like(x0)
    x_t = q_sample(x0, t, noise, schedule)
    pred = model(x_t, t)
    loss = F.mse_loss(pred, noise)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

Это весь цикл обучения. Никакой игры GAN, никакой специализированной функции потерь — один вызов MSE.

Шаг 5: Сэмплер (DDPM)

@torch.no_grad()
def sample(model, schedule, shape, T=1000, device="cpu"):
    model.eval()
    x = torch.randn(shape, device=device)
    betas = schedule["betas"].to(device)
    sqrt_one_minus_a = schedule["sqrt_one_minus_alphas_cumprod"].to(device)
    sqrt_recip_alphas = schedule["sqrt_recip_alphas"].to(device)

    for t in reversed(range(T)):
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        coef = betas[t] / sqrt_one_minus_a[t]
        mean = sqrt_recip_alphas[t] * (x - coef * eps)
        if t > 0:
            x = mean + torch.sqrt(betas[t]) * torch.randn_like(x)
        else:
            x = mean
    return x

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

Шаг 6: Сэмплер DDIM (детерминированный, ~20x быстрее)

@torch.no_grad()
def sample_ddim(model, schedule, shape, steps=50, T=1000, device="cpu", eta=0.0):
    model.eval()
    x = torch.randn(shape, device=device)
    alphas_cumprod = schedule["alphas_cumprod"].to(device)

    ts = torch.linspace(T - 1, 0, steps + 1).long()
    for i in range(steps):
        t = ts[i]
        t_prev = ts[i + 1]
        t_batch = torch.full((shape[0],), t, dtype=torch.long, device=device)
        eps = model(x, t_batch)
        a_t = alphas_cumprod[t]
        a_prev = alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0, device=device)
        x0_pred = (x - torch.sqrt(1 - a_t) * eps) / torch.sqrt(a_t)
        sigma = eta * torch.sqrt((1 - a_prev) / (1 - a_t) * (1 - a_t / a_prev))
        dir_xt = torch.sqrt(1 - a_prev - sigma ** 2) * eps
        noise = sigma * torch.randn_like(x) if eta > 0 else 0
        x = torch.sqrt(a_prev) * x0_pred + dir_xt + noise
    return x

eta=0 полностью детерминирован (тот же входной шум всегда даёт тот же результат). eta=1 воспроизводит DDPM.

Используйте

Для производственной работы используйте diffusers:

from diffusers import DDPMScheduler, UNet2DModel

unet = UNet2DModel(sample_size=32, in_channels=3, out_channels=3, layers_per_block=2)
scheduler = DDPMScheduler(num_train_timesteps=1000)

Библиотека поставляет готовые планировщики (DDPM, DDIM, DPM-Solver, Euler, Heun), настраиваемые U-Net, конвейеры для text-to-image и image-to-image, а также помощники тонкой настройки LoRA.

Для исследований k-diffusion (Katherine Crowson) содержит наиболее точные эталонные реализации и лучшие варианты сэмплирования.

Поставьте

Этот урок создаёт:

• outputs/prompt-diffusion-sampler-picker.md — промпт, выбирающий DDPM / DDIM / DPM-Solver / Euler по целевому качеству, бюджету задержки и типу обусловливания.
• outputs/skill-noise-schedule-designer.md — навык, создающий линейное, косинусное или сигмоидальное расписание beta по T и целевому уровню повреждения, а также диагностические графики отношения сигнал/шум во времени.

Упражнения

1. (Легко) Визуализируйте прямой процесс: возьмите одно изображение и постройте x_t при t in [0, 100, 250, 500, 750, 1000]. Убедитесь, что x_1000 выглядит как чистый гауссовский шум.
2. (Средне) Обучите TinyUNet на наборе synthetic-circles в течение 20 эпох и сэмплируйте 16 окружностей. Сравните сэмплирование DDPM (1000 шагов) и DDIM (50 шагов) — дают ли они похожие изображения при одном seed шума?
3. (Сложно) Реализуйте косинусное расписание шума (Nichol & Dhariwal, 2021): alpha_bar_t = cos^2((t/T + s) / (1 + s) * pi / 2). Обучите одну и ту же модель с линейным и косинусным расписаниями и покажите, что косинусное даёт лучшие образцы при малом числе шагов.

Ключевые термины

Термин	Как обычно говорят	Что это на самом деле означает
Прямой процесс	«Добавлять шум со временем»	Фиксированная цепь Маркова, за T шагов превращающая изображение в гауссовский шум
Обратный процесс	«Убирать шум шаг за шагом»	Выученное распределение, возвращающееся от шума к изображению
Предсказание epsilon	«Предсказать шум»	Цель обучения: epsilon_theta(x_t, t) предсказывает шум, добавленный на шаге t
Расписание beta	«Количество шума»	Последовательность T малых дисперсий, определяющих объём шума, поступающий на каждом шаге
alpha_bar_t	«Накопленный коэффициент сохранения»	Произведение (1 - beta_s) до момента t; чем больше t, тем меньше осталось сигнала
Сэмплер DDPM	«Предковый, стохастический»	Сэмплирует каждый x_{t-1} из его условного гауссовского распределения; 1000 шагов
Сэмплер DDIM	«Детерминированный, быстрый»	Переписывает сэмплирование как детерминированное ODE; 20–100 шагов с похожим качеством
Условие по времени	«Сообщить модели t»	Синусоидальное встраивание t, внедряемое в U-Net, чтобы сеть знала уровень шума

Дополнительные материалы

• Denoising Diffusion Probabilistic Models (Ho et al., 2020) — статья, сделавшая диффузию практичной и превзошедшая GAN по FID
• Improved DDPM (Nichol & Dhariwal, 2021) — косинусное расписание и v-параметризация
• DDIM (Song, Meng, Ermon, 2020) — детерминированный сэмплер, сделавший возможным инференс в реальном времени
• Elucidating the Design Space of Diffusion (Karras et al., 2022) — единый взгляд на каждый выбор при проектировании диффузии; текущая лучшая справочная работа

---

Источник: Image Generation — Diffusion Models 04.09 — Генерация изображений — GAN · Фаза 04 — Компьютерное зрение · 04.11 — Stable Diffusion · Полный каталог