Diffusion Models Course, юнит 1: диффузионная модель с нуля

Перевод ноутбука «Diffusion Models from Scratch» курса Hugging Face: игрушечная диффузия на MNIST — своя порча шумом, минимальная U-Net, обучение и сэмплирование — и сравнение с DDPM: расписание шума, предсказание шума, условие на шаг, сэмплеры.

Это перевод второго ноутбука юнита 1 курса Diffusion Models Course. Первая часть юнита — введение и библиотека Diffusers, карта курса — в README оригинала.

Ноутбук делает то же, что практика урока 9.1.2, только сначала нарочно «наивно»: шум смешивается с картинкой линейно, а модель предсказывает чистое изображение. Потом автор сравнивает это с DDPM и объясняет, зачем нужны расписание шума, предсказание шума и условие на шаг. Интересно, что наивная схема «смешиваем картинку с шумом по прямой и идём маленькими шагами» — по сути идея flow matching, о котором урок 9.1.4. Свёртки, из которых собрана U-Net, — в уроке 3.5.1. Весь ноутбук идёт на одной GPU за несколько минут, а игрушечная часть — даже на CPU.

Диффузионные модели с нуля

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

Мы разберём:

процесс порчи (добавление шума к данным);
что такое U-Net и как реализовать предельно минимальную её версию с нуля;
обучение диффузионной модели;
теорию сэмплирования.

Затем сравним наши версии с реализацией DDPM в diffusers и разберём:

улучшения по сравнению с нашей мини-U-Net;
расписание шума DDPM;
различия в цели обучения;
условие на временной шаг;
подходы к сэмплированию.

Ноутбук довольно подробный, и его спокойно можно пропустить, если глубокое погружение «с нуля» вас не увлекает!

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

Подготовка и импорты

%pip install -q diffusers

import torch
import torchvision
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from diffusers import DDPMScheduler, UNet2DModel
from matplotlib import pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')

Вывод:

Using device: cuda

Данные

Проверять всё будем на совсем маленьком датасете — MNIST. Если хотите дать модели задачу чуть сложнее, ничего больше не меняя, подставьте torchvision.datasets.FashionMNIST — он подходит как прямая замена.

dataset = torchvision.datasets.MNIST(root="mnist/", train=True, download=True, transform=torchvision.transforms.ToTensor())

train_dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

x, y = next(iter(train_dataloader))
print('Input shape:', x.shape)
print('Labels:', y)
plt.imshow(torchvision.utils.make_grid(x)[0], cmap='Greys');

Вывод:

Input shape: torch.Size([8, 1, 28, 28])
Labels: tensor([1, 9, 7, 3, 5, 2, 1, 4])

Каждое изображение — рисунок цифры в оттенках серого размером 28×28 пикселей со значениями от 0 до 1.

(Метки у вас будут другие: загрузчик перемешивает данные.)

Процесс порчи

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

Наверное, нужен простой способ управлять силой порчи. Что если принимать параметр amount — сколько шума добавить, — и делать так:

noise = torch.rand_like(x)

noisy_x = (1-amount)*x + amount*noise

При amount = 0 получаем вход без изменений. Если amount дорастает до 1, получаем шум без следа входа x. Смешивая вход с шумом таким образом, мы оставляем результат в том же диапазоне (от 0 до 1).

Реализовать это довольно легко (только следите за формами тензоров, чтобы не обжечься на правилах broadcasting):

def corrupt(x, amount):
  """Портит вход `x`, смешивая его с шумом в пропорции `amount`"""
  noise = torch.rand_like(x)
  amount = amount.view(-1, 1, 1, 1) # меняем форму, чтобы сработал broadcasting
  return x*(1-amount) + noise*amount

(Зачем view(-1, 1, 1, 1): amount — это вектор из 8 чисел, по одному на картинку, а x имеет форму [8, 1, 28, 28]. Форма [8, 1, 1, 1] говорит PyTorch: «одно число на картинку, растяни его на все пиксели». Правила broadcasting — в уроке 0.3.2.)

И посмотрим на результат, чтобы убедиться, что всё работает как ожидалось:

# Рисуем входные данные
fig, axs = plt.subplots(2, 1, figsize=(12, 5))
axs[0].set_title('Input data')
axs[0].imshow(torchvision.utils.make_grid(x)[0], cmap='Greys')

# Добавляем шум
amount = torch.linspace(0, 1, x.shape[0]) # слева направо -> больше порчи
noised_x = corrupt(x, amount)

# Рисуем зашумлённую версию
axs[1].set_title('Corrupted data (-- amount increases -->)')
axs[1].imshow(torchvision.utils.make_grid(noised_x)[0], cmap='Greys');

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

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

Модель

Нам нужна модель, которая принимает зашумлённые изображения 28 пикселей и выдаёт предсказание той же формы. Популярный выбор здесь — архитектура U-Net. Изобретённая для сегментации медицинских изображений, U-Net состоит из «сужающегося пути», на котором данные сжимаются, и «расширяющегося пути», на котором они снова разворачиваются до исходного размера (как в автокодировщике), а также обходных связей, по которым информация и градиенты проходят напрямую между уровнями.

В некоторых U-Net на каждом этапе сложные блоки, но для игрушечной демонстрации мы построим минимальный пример: он принимает одноканальное изображение и пропускает его через три свёрточных слоя на пути вниз (down_layers на схеме и в коде) и три на пути вверх, с обходными связями между слоями путей вниз и вверх. Для уменьшения возьмём max pooling, для увеличения — nn.Upsample, а не обучаемые слои, как в более сложных U-Net. Вот примерная архитектура с числом каналов на выходе каждого слоя.

(В оригинале — схема: вход 1 канал → свёртка 32 канала → max pool → 64 канала → max pool → 64 канала, затем вверх: 64 канала → upsample + обходная связь → 32 канала → upsample + обходная связь → 1 канал на выходе.)

flowchart LR
I["вход 1×28×28"] --> C1["conv → 32×28×28"]
C1 --> P1["maxpool → 32×14×14"] --> C2["conv → 64×14×14"]
C2 --> P2["maxpool → 64×7×7"] --> C3["conv → 64×7×7"]
C3 --> U1["conv → 64×7×7"]
U1 --> S1["upsample → 64×14×14<br/>+ выход C2"] --> U2["conv → 32×14×14"]
U2 --> S2["upsample → 32×28×28<br/>+ выход C1"] --> U3["conv → 1×28×28"]

BasicUNet: формы тензоров на каждом шаге (рисунок переводчика)

Вот как это выглядит в коде:

class BasicUNet(nn.Module):
    """Минимальная реализация U-Net."""
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()
        self.down_layers = torch.nn.ModuleList([
            nn.Conv2d(in_channels, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
        ])
        self.up_layers = torch.nn.ModuleList([
            nn.Conv2d(64, 64, kernel_size=5, padding=2),
            nn.Conv2d(64, 32, kernel_size=5, padding=2),
            nn.Conv2d(32, out_channels, kernel_size=5, padding=2),
        ])
        self.act = nn.SiLU() # функция активации
        self.downscale = nn.MaxPool2d(2)
        self.upscale = nn.Upsample(scale_factor=2)

    def forward(self, x):
        h = []
        for i, l in enumerate(self.down_layers):
            x = self.act(l(x)) # через слой и функцию активации
            if i < 2: # для всех слоёв, кроме третьего (последнего) на пути вниз:
              h.append(x) # сохраняем выход для обходной связи
              x = self.downscale(x) # уменьшаем перед следующим слоем

        for i, l in enumerate(self.up_layers):
            if i > 0: # для всех слоёв, кроме первого на пути вверх
              x = self.upscale(x) # увеличиваем
              x += h.pop() # достаём сохранённый выход (обходная связь)
            x = self.act(l(x)) # через слой и функцию активации

        return x

Проверим, что форма выхода совпадает с формой входа, как мы и ждём:

net = BasicUNet()
x = torch.rand(8, 1, 28, 28)
net(x).shape

Вывод:

torch.Size([8, 1, 28, 28])

В этой сети чуть больше 300 000 параметров:

sum([p.numel() for p in net.parameters()])

Вывод:

(Проверка от переводчика: у свёртки $5 \times 5$ из $C_{in}$ каналов в $C_{out}$ — $C_{out} \cdot C_{in} \cdot 25 + C_{out}$ параметров. Вниз: $832 + 51\,264 + 102\,464$ , вверх: $102\,464 + 51\,232 + 801$ ; сумма — ровно 309 057.)

Можете поэкспериментировать: поменять число каналов в каждом слое или подставить другие архитектуры.

Обучаем сеть

Так что же именно должна делать модель? Здесь тоже есть разные подходы, но для демонстрации выберем простую постановку: по испорченному входу noisy_x модель должна выдать свою лучшую догадку о том, как выглядел исходный x. Сравнивать её с настоящим значением будем среднеквадратичной ошибкой.

Теперь можно попробовать обучить сеть:

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

Смело меняйте код и пробуйте добиться лучшего результата!

# Загрузчик данных (можно поиграть с размером батча)
batch_size = 128
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Сколько раз пройти по данным?
n_epochs = 3

# Создаём сеть
net = BasicUNet()
net.to(device)

# Функция потерь
loss_fn = nn.MSELoss()

# Оптимизатор
opt = torch.optim.Adam(net.parameters(), lr=1e-3)

# Сохраняем значения loss, чтобы потом посмотреть
losses = []

# Цикл обучения
for epoch in range(n_epochs):

    for x, y in train_dataloader:

        # Берём данные и готовим испорченную версию
        x = x.to(device) # данные на GPU
        noise_amount = torch.rand(x.shape[0]).to(device) # выбираем случайные уровни шума
        noisy_x = corrupt(x, noise_amount) # создаём зашумлённый x

        # Получаем предсказание модели
        pred = net(noisy_x)

        # Считаем loss
        loss = loss_fn(pred, x) # насколько выход близок к настоящему «чистому» x?

        # Обратное распространение и обновление параметров:
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Сохраняем loss на потом
        losses.append(loss.item())

    # Печатаем среднее значение loss за эпоху:
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')

# Смотрим на кривую loss
plt.plot(losses)
plt.ylim(0, 0.1);

Вывод:

Finished epoch 0. Average loss for this epoch: 0.026736
Finished epoch 1. Average loss for this epoch: 0.020692
Finished epoch 2. Average loss for this epoch: 0.018887

(Что такое цикл обучения и зачем opt.zero_grad() — в уроке 3.3.1. На RTX 4080 три эпохи по 60 000 картинок занимают меньше минуты.)

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

#@markdown Смотрим на предсказания модели по зашумлённым входам:

# Берём данные
x, y = next(iter(train_dataloader))
x = x[:8] # только первые 8 — чтобы удобно рисовать

# Портим с разной силой
amount = torch.linspace(0, 1, x.shape[0]) # слева направо -> больше порчи
noised_x = corrupt(x, amount)

# Получаем предсказания модели
with torch.no_grad():
  preds = net(noised_x.to(device)).detach().cpu()

# Рисуем
fig, axs = plt.subplots(3, 1, figsize=(12, 7))
axs[0].set_title('Input data')
axs[0].imshow(torchvision.utils.make_grid(x)[0].clip(0, 1), cmap='Greys')
axs[1].set_title('Corrupted data')
axs[1].imshow(torchvision.utils.make_grid(noised_x)[0].clip(0, 1), cmap='Greys')
axs[2].set_title('Network Predictions')
axs[2].imshow(torchvision.utils.make_grid(preds)[0].clip(0, 1), cmap='Greys');

(#@markdown — служебная пометка Colab, на работу кода не влияет.)

Видно, что при небольших уровнях шума предсказания довольно хорошие! Но когда уровень очень высок, модели остаётся мало информации, и к amount=1 она выдаёт размытую кашу, близкую к среднему по датасету, — так она «страхуется», не зная, как должен выглядеть результат…

Сэмплирование

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

А что если начать со случайного шума, посмотреть на предсказание модели, но сдвинуться к нему лишь немного — скажем, на 20% пути? Получится очень шумное изображение, в котором, возможно, уже намечена структура, и его можно снова подать в модель и получить новое предсказание. Надежда в том, что новое предсказание чуть лучше первого (ведь стартовая точка чуть менее шумная), и с ним можно сделать ещё один маленький шаг.

Повторим несколько раз и (если всё пойдёт хорошо) получим изображение! Вот этот процесс всего за 5 шагов: слева — вход модели на каждом этапе, справа — предсказанное очищенное изображение. Обратите внимание: модель предсказывает очищенное изображение уже на шаге 1, но мы сдвигаем x лишь на часть пути. За несколько шагов проявляются и уточняются структуры, пока не получится итоговый результат.

#@markdown Стратегия сэмплирования: делим процесс на 5 шагов и каждый раз проходим 1/5 пути:
n_steps = 5
x = torch.rand(8, 1, 28, 28).to(device) # начинаем со случайного
step_history = [x.detach().cpu()]
pred_output_history = []

for i in range(n_steps):
    with torch.no_grad(): # при инференсе градиенты не нужны
        pred = net(x) # предсказываем очищенный x0
    pred_output_history.append(pred.detach().cpu()) # сохраняем выход модели для графика
    mix_factor = 1/(n_steps - i) # насколько сдвигаемся к предсказанию
    x = x*(1-mix_factor) + pred*mix_factor # проходим часть пути
    step_history.append(x.detach().cpu()) # сохраняем шаг для графика

fig, axs = plt.subplots(n_steps, 2, figsize=(9, 4), sharex=True)
axs[0,0].set_title('x (model input)')
axs[0,1].set_title('model prediction')
for i in range(n_steps):
    axs[i, 0].imshow(torchvision.utils.make_grid(step_history[i])[0].clip(0, 1), cmap='Greys')
    axs[i, 1].imshow(torchvision.utils.make_grid(pred_output_history[i])[0].clip(0, 1), cmap='Greys')

Доля сдвига растёт: на шагах 0, 1, 2, 3, 4 это $1/5, 1/4, 1/3, 1/2, 1$ . Так каждый шаг убирает примерно одинаковую долю исходного шума: после первого шага от шума остаётся $4/5$ , после второго — $4/5 \cdot 3/4 = 3/5$ , затем $2/5$ , $1/5$ и на последнем шаге — ноль, и $x$ совпадает с последним предсказанием. Это та же идея, что в сэмплировании flow matching методом Эйлера: идём к цели равными «порциями пути».

Можно разбить процесс на больше шагов в надежде получить изображения получше:

#@markdown Больше результатов: 40 шагов сэмплирования
n_steps = 40
x = torch.rand(64, 1, 28, 28).to(device)
for i in range(n_steps):
  noise_amount = torch.ones((x.shape[0], )).to(device) * (1-(i/n_steps)) # от высокого к низкому
  with torch.no_grad():
    pred = net(x)
  mix_factor = 1/(n_steps - i)
  x = x*(1-mix_factor) + pred*mix_factor
fig, ax = plt.subplots(1, 1, figsize=(12, 12))
ax.imshow(torchvision.utils.make_grid(x.detach().cpu(), nrow=8)[0].clip(0, 1), cmap='Greys')

(Переменная noise_amount здесь считается, но нигде не используется: наша сеть не знает уровня шума. Это задел на «условие на шаг», о котором речь пойдёт ниже.)

Не очень, но некоторые цифры узнаются! Можно поэкспериментировать: обучать дольше (например, 10 или 20 эпох), подстроить конфигурацию модели, скорость обучения, оптимизатор и так далее. И не забывайте, что fashionMNIST — замена в одну строку, если хотите датасет посложнее.

Сравнение с DDPM

В этом разделе посмотрим, чем наша игрушечная реализация отличается от подхода в другом ноутбуке (Введение в Diffusers), основанного на статье DDPM.

Мы увидим, что:

UNet2DModel из diffusers чуть продвинутее нашей BasicUNet;
процесс порчи устроен иначе;
цель обучения другая: модель предсказывает шум, а не очищенное изображение;
модель получает условие на количество шума через временной шаг: t передаётся в метод forward дополнительным аргументом;
доступно много разных стратегий сэмплирования, которые должны работать лучше нашей упрощённой версии выше.

Со времени выхода статьи DDPM предложено много улучшений, но этот пример, надеюсь, показывает, какие проектные решения вообще бывают. Разобравшись с ним, вы можете с удовольствием погрузиться в статью «Elucidating the Design Space of Diffusion-Based Generative Models» (Karras et al., 2022), где все эти компоненты изучены подробно и даны новые рекомендации, как добиться наилучшего качества.

Если всё это слишком технично или пугает — не волнуйтесь! Остаток ноутбука можно пропустить или отложить на дождливый день.

U-Net

У UNet2DModel из diffusers несколько улучшений по сравнению с нашей простой U-Net выше:

GroupNorm применяет групповую нормализацию ко входам каждого блока;
слои Dropout для более гладкого обучения;
несколько слоёв ResNet в каждом блоке (если layers_per_block не равно 1);
внимание (обычно только в блоках с низким разрешением);
условие на временной шаг;
блоки уменьшения и увеличения с обучаемыми параметрами.

(Нормализация и остаточные связи — в уроках 3.4.1 и 3.4.2.)

Создадим UNet2DModel и посмотрим на неё:

model = UNet2DModel(
    sample_size=28,           # целевое разрешение изображения
    in_channels=1,            # число входных каналов, 3 для цветных изображений
    out_channels=1,           # число выходных каналов
    layers_per_block=2,       # сколько слоёв ResNet в каждом блоке U-Net
    block_out_channels=(32, 64, 64), # примерно как в нашей простой U-Net
    down_block_types=(
        "DownBlock2D",        # обычный блок уменьшения на ResNet
        "AttnDownBlock2D",    # блок уменьшения на ResNet с пространственным self-attention
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",      # блок увеличения на ResNet с пространственным self-attention
        "UpBlock2D",          # обычный блок увеличения на ResNet
      ),
)
print(model)

Вывод (начало):

UNet2DModel(
  (conv_in): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (time_proj): Timesteps()
  (time_embedding): TimestepEmbedding(
    (linear_1): Linear(in_features=32, out_features=128, bias=True)
    (act): SiLU()
    ...

Как видите, здесь происходит побольше всего! И параметров у неё заметно больше, чем у BasicUNet:

sum([p.numel() for p in model.parameters()]) # 1,7 млн против ~309 тыс. у BasicUNet

Вывод:

Можно повторить обучение, показанное выше, взяв эту модель вместо исходной. Модели нужно передавать и x, и временной шаг (здесь я всегда передаю t=0, чтобы показать, что всё работает и без условия на шаг, и чтобы код сэмплирования оставался простым; но можно попробовать подавать (amount*1000), чтобы получить шаг, эквивалентный силе порчи). Изменённые строки помечены #<<<, если хотите их рассмотреть.

#@markdown Пробуем UNet2DModel вместо BasicUNet:

# Загрузчик данных (можно поиграть с размером батча)
batch_size = 128
train_dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# Сколько раз пройти по данным?
n_epochs = 3

# Создаём сеть
net = UNet2DModel(
    sample_size=28,  # целевое разрешение изображения
    in_channels=1,  # число входных каналов, 3 для цветных изображений
    out_channels=1,  # число выходных каналов
    layers_per_block=2,  # сколько слоёв ResNet в каждом блоке U-Net
    block_out_channels=(32, 64, 64),  # примерно как в нашей простой U-Net
    down_block_types=(
        "DownBlock2D",  # обычный блок уменьшения на ResNet
        "AttnDownBlock2D",  # блок уменьшения на ResNet с пространственным self-attention
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",  # блок увеличения на ResNet с пространственным self-attention
        "UpBlock2D",   # обычный блок увеличения на ResNet
      ),
) #<<<
net.to(device)

# Функция потерь
loss_fn = nn.MSELoss()

# Оптимизатор
opt = torch.optim.Adam(net.parameters(), lr=1e-3)

# Сохраняем значения loss, чтобы потом посмотреть
losses = []

# Цикл обучения
for epoch in range(n_epochs):

    for x, y in train_dataloader:

        # Берём данные и готовим испорченную версию
        x = x.to(device) # данные на GPU
        noise_amount = torch.rand(x.shape[0]).to(device) # выбираем случайные уровни шума
        noisy_x = corrupt(x, noise_amount) # создаём зашумлённый x

        # Получаем предсказание модели
        pred = net(noisy_x, 0).sample #<<< всегда шаг 0, добавили .sample

        # Считаем loss
        loss = loss_fn(pred, x) # насколько выход близок к настоящему «чистому» x?

        # Обратное распространение и обновление параметров:
        opt.zero_grad()
        loss.backward()
        opt.step()

        # Сохраняем loss на потом
        losses.append(loss.item())

    # Печатаем среднее значение loss за эпоху:
    avg_loss = sum(losses[-len(train_dataloader):])/len(train_dataloader)
    print(f'Finished epoch {epoch}. Average loss for this epoch: {avg_loss:05f}')

# Рисуем loss и несколько сэмплов
fig, axs = plt.subplots(1, 2, figsize=(12, 5))

# Loss
axs[0].plot(losses)
axs[0].set_ylim(0, 0.1)
axs[0].set_title('Loss over time')

# Сэмплы
n_steps = 40
x = torch.rand(64, 1, 28, 28).to(device)
for i in range(n_steps):
  noise_amount = torch.ones((x.shape[0], )).to(device) * (1-(i/n_steps)) # от высокого к низкому
  with torch.no_grad():
    pred = net(x, 0).sample
  mix_factor = 1/(n_steps - i)
  x = x*(1-mix_factor) + pred*mix_factor

axs[1].imshow(torchvision.utils.make_grid(x.detach().cpu(), nrow=8)[0].clip(0, 1), cmap='Greys')
axs[1].set_title('Generated Samples');

Вывод:

Finished epoch 0. Average loss for this epoch: 0.018925
Finished epoch 1. Average loss for this epoch: 0.012785
Finished epoch 2. Average loss for this epoch: 0.011694

Это выглядит заметно лучше наших первых результатов! Можно подстроить конфигурацию U-Net или обучать дольше, чтобы получить ещё лучшее качество.

Процесс порчи

Статья DDPM описывает процесс порчи, который на каждом «временном шаге» добавляет немного шума. Имея $x_{t-1}$ для какого-то шага, следующую (чуть более зашумлённую) версию $x_t$ получаем так:

q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I}) \quad q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod^T_{t=1} q(\mathbf{x}_t \vert \mathbf{x}_{t-1})

То есть берём $x_{t-1}$ , умножаем на $\sqrt{1 - \beta_t}$ и добавляем шум, масштабированный $\beta_t$ . Это $\beta$ задано для каждого $t$ по некоторому расписанию и определяет, сколько шума добавляется на шаге. Нам не обязательно проделывать эту операцию 500 раз, чтобы получить $x_{500}$ , — есть другая формула, которая даёт $x_t$ для любого $t$ сразу из $x_0$ :

$q(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, \sqrt{(1 - \bar{\alpha}_t)} \mathbf{I}), \quad \text{где } \bar{\alpha}_t = \prod_{i=1}^T \alpha_i \text{ и } \alpha_i = 1-\beta_i$

(Примечание переводчика: в этой формуле две опечатки. Дисперсия равна $(1 - \bar\alpha_t)\mathbf{I}$ , без корня — корень стоит в стандартном отклонении, то есть в множителе перед шумом; и произведение берётся до $t$ , а не до $T$ . Как в правильной записи в первой части юнита.)

Математические обозначения всегда выглядят пугающе! К счастью, планировщик всё это делает за нас (раскомментируйте следующую ячейку, чтобы посмотреть код). Можно построить графики $\sqrt{\bar{\alpha}_t}$ (подписан как sqrt_alpha_prod) и $\sqrt{(1 - \bar{\alpha}_t)}$ (подписан как sqrt_one_minus_alpha_prod), чтобы увидеть, как вход ( $x$ ) и шум масштабируются и смешиваются на разных шагах:

#??noise_scheduler.add_noise

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
plt.legend(fontsize="x-large");

Сначала зашумлённый x — это в основном x (sqrt_alpha_prod ≈ 1), но со временем вклад x падает, а доля шума растёт. В отличие от нашего линейного смешивания x и шума по amount, здесь шум нарастает сравнительно быстро. Посмотрим на это на данных:

#@markdown Смотрим на процесс зашумления DDPM при разных временных шагах:

# Зашумляем батч изображений, чтобы увидеть эффект
fig, axs = plt.subplots(3, 1, figsize=(16, 10))
xb, yb = next(iter(train_dataloader))
xb = xb.to(device)[:8]
xb = xb * 2. - 1. # переводим в (-1, 1)
print('X shape', xb.shape)

# Показываем чистые входы
axs[0].imshow(torchvision.utils.make_grid(xb[:8])[0].detach().cpu(), cmap='Greys')
axs[0].set_title('Clean X')

# Добавляем шум планировщиком
timesteps = torch.linspace(0, 999, 8).long().to(device)
noise = torch.randn_like(xb) # << NB: randn, а не rand
noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
print('Noisy X shape', noisy_xb.shape)

# Показываем зашумлённую версию (с обрезкой значений и без)
axs[1].imshow(torchvision.utils.make_grid(noisy_xb[:8])[0].detach().cpu().clip(-1, 1),  cmap='Greys')
axs[1].set_title('Noisy X (clipped to (-1, 1)')
axs[2].imshow(torchvision.utils.make_grid(noisy_xb[:8])[0].detach().cpu(),  cmap='Greys')
axs[2].set_title('Noisy X');

Вывод:

X shape torch.Size([8, 1, 28, 28])
Noisy X shape torch.Size([8, 1, 28, 28])

Ещё одно отличие: версия DDPM добавляет шум из гауссова (нормального) распределения (среднее 0, стандартное отклонение 1, из torch.randn), а не равномерный шум от 0 до 1 (из torch.rand), как в нашей исходной функции corrupt. И вообще разумно нормировать обучающие данные. В другом ноутбуке вы увидите в списке преобразований Normalize(0.5, 0.5), которое переводит данные изображения из (0, 1) в (−1, 1), — для наших целей этого «достаточно». В этом ноутбуке мы так не делали, но ячейка визуализации выше добавляет нормировку, чтобы масштаб и картинки были точнее.

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

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

noise = torch.randn_like(xb) # << NB: randn, а не rand
noisy_x = noise_scheduler.add_noise(x, noise, timesteps)
model_prediction = model(noisy_x, timesteps).sample
loss = mse_loss(model_prediction, noise) # цель — шум

Может показаться, что предсказывать шум (по которому можно вычислить, как выглядит очищенное изображение) — то же самое, что сразу предсказывать очищенное изображение. Так почему предпочитать одно другому — только ради математического удобства?

Оказывается, здесь есть ещё одна тонкость. При обучении мы считаем loss на разных (случайно выбранных) временных шагах. Разные цели приводят к разному «неявному взвешиванию» этих loss: предсказание шума даёт больший вес низким уровням шума. Можно выбрать более сложную цель, чтобы изменить это «неявное взвешивание loss». Или выбрать расписание шума, при котором будет больше примеров с высоким уровнем шума. Или пусть модель предсказывает «скорость» $v$ — сочетание изображения и шума, зависящее от уровня шума (см. «Progressive Distillation for Fast Sampling of Diffusion Models»). Или пусть модель предсказывает шум, но loss масштабируется множителем, зависящим от количества шума, — исходя из теории (см. «Perception Prioritized Training of Diffusion Models») или из экспериментов о том, какие уровни шума наиболее информативны для модели (см. «Elucidating the Design Space of Diffusion-Based Generative Models»). Коротко: выбор цели влияет на качество модели, и исследования того, какой вариант «лучший», продолжаются.

Сейчас предпочтительный подход — предсказывать шум (иногда его обозначают epsilon или eps), но со временем, вероятно, библиотека будет поддерживать и другие цели, и их будут применять в разных ситуациях.

Из $x_t = \sqrt{\bar\alpha_t}\,x_0 + \sqrt{1-\bar\alpha_t}\,\varepsilon$ любая из трёх величин выражается через две другие: зная предсказанный шум $\hat\varepsilon$ , получаем $\hat x_0 = (x_t - \sqrt{1-\bar\alpha_t}\,\hat\varepsilon) / \sqrt{\bar\alpha_t}$ . Ошибка в $x_0$ равна ошибке в $\varepsilon$ , умноженной на $\sqrt{1-\bar\alpha_t}/\sqrt{\bar\alpha_t}$ : при малом шуме этот множитель мал, при большом — огромен. Поэтому одна и та же MSE по шуму и по картинке по-разному распределяет «внимание» обучения между уровнями шума. Цель $v$ — середина между ними; а в flow matching (урок 9.1.4) модель предсказывает скорость $x_1 - x_0$ на прямом пути от шума к данным (в обозначениях урока $x_0$ — шум, $x_1$ — данные, наоборот по сравнению с DDPM). После 2022 года именно v-предсказание и flow matching стали стандартом в больших моделях (Stable Diffusion 3, FLUX и др.).

Условие на временной шаг

UNet2DModel принимает и x, и временной шаг. Шаг превращается в эмбеддинг и подаётся в модель в нескольких местах.

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

(Как именно: число $t$ превращается в вектор синусоидальным эмбеддингом — тем же приёмом, что позиционное кодирование в трансформере (урок 4.5.2), — затем проходит через небольшой MLP (time_embedding в выводе print(model) выше) и добавляется к признакам в каждом блоке ResNet.)

Сэмплирование

Если есть модель, которая оценивает шум в зашумлённом входе (или предсказывает очищенную версию), как получить новые изображения?

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

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

Насколько большой шаг делать? Иначе говоря, какому «расписанию шума» следовать?
Использовать для шага только текущее предсказание модели (как DDPM, DDIM и многие другие)? Вызывать модель несколько раз, чтобы оценить градиенты высших порядков для большего и более точного шага (методы высших порядков и некоторые дискретные решатели ОДУ)? Или хранить историю прошлых предсказаний, чтобы точнее выбрать текущий шаг (линейные многошаговые и «предковые» (ancestral) сэмплеры)?
Добавлять ли дополнительный шум (иногда его называют churn), чтобы сделать сэмплирование стохастичнее (случайнее), или держать его полностью детерминированным? Многие сэмплеры управляют этим параметром (например, eta у сэмплеров DDIM), чтобы пользователь мог выбрать.

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

Выводы

Надеюсь, это помогло взглянуть на диффузионные модели немного под другим углом.

Этот ноутбук написал для курса Hugging Face Джонатан Уитакер; он пересекается с версией из его собственного курса «The Generative Landscape». Загляните туда, если хотите увидеть этот простой пример, расширенный условиями на шум и на класс. Вопросы и ошибки можно присылать через issues на GitHub или в Discord. Можно также написать в Twitter @johnowhitaker.

Словарик

English	Русский
corruption process	процесс порчи (зашумления)
uniform / Gaussian noise	равномерный / гауссов (нормальный) шум
broadcasting	broadcasting (растягивание размерностей)
max pooling, upsample	max pooling, увеличение размера
skip connection	обходная связь
training objective	цель обучения
noise prediction (epsilon), v-prediction	предсказание шума (эпсилон), предсказание скорости $v$
implicit loss weighting	неявное взвешивание loss
timestep conditioning	условие на временной шаг
sampler, ODE solver	сэмплер, решатель ОДУ
ancestral sampler, churn	«предковый» сэмплер, дополнительный шум при сэмплировании
deterministic / stochastic	детерминированный / стохастический

Перевод выполнен по лицензии оригинала (Apache-2.0); автор оригинала — Джонатан Уитакер (Jonathan Whitaker) для Hugging Face. Нашли неточность — сверяйтесь с оригиналом по ссылке выше. · обновлено 2026-10-04