Практика: DDPM с нуля

Ветка в двух формулах

$$\htmlData{k=train}{\mathcal{L} = \mathbb{E}\big\|\varepsilon_\theta(x_t, t) - \varepsilon\big\|^2}, \qquad \htmlData{k=sample}{dx = \Big[f - \tfrac12 g^2 \nabla\log p_t\Big]dt}$$

Слева — урок 060, справа — урок 030. Между ними ровно одно связующее утверждение: предсказанный шум и есть score с точностью до множителя. Всё остальное в ветке — расписания, guidance, потоки — это варианты того, как выбрать путь и как по нему пройти.

Что реализовать

1. Прямой процесс и расписания

Начните с того, что не требует обучения вовсе.

• реализуйте alpha_bar(t) для линейного и косинусного расписаний;
• нарисуйте $x_t$ для одного изображения при $t = 0, 0.25, 0.5, 0.75, 1$;
• посчитайте, при каком $t$ отношение сигнал/шум равно единице.

Ожидание из урока 070: $0.26$ для линейного и $0.50$ для косинусного. Если получилось иначе — ошибка в накоплении произведения, и её лучше найти сейчас, чем после суток обучения.

линейноекосинусное

данные при t

сигнал/шум, log-шкала

время t

0.26

ᾱ

0.4967

√ᾱ

0.7047

√(1−ᾱ)

0.7095

сигнал/шум

0.987

равенство при t

0.259

Шум преобладает. Если ваш выбор t для отладки лежит здесь, ожидать осмысленной реконструкции не стоит.

2. Обучение

Модель — маленький U-Net на MNIST. Обязательные детали:

• $t$ подаётся как эмбеддинг, синусоидальный (блок 6, урок 110) или через $\log$ SNR;
• на каждый пример один случайный $t$ и один $\varepsilon$;
• цель — упрощённая, без весов ELBO (урок 060).

Проверка, которую стоит сделать до обучения: при случайной инициализации лосс обязан быть около единицы. Причина в том, что цель — MSE до $\varepsilon \sim \mathcal{N}(0, I)$, а необученная сеть предсказывает примерно нуль; значит ошибка равна дисперсии шума. Значение сильно другое означает ошибку в масштабировании входа.

3. Три сэмплера

Один и тот же обученный набор весов, три способа пройти обратный путь:

сэмплер	формулировка	шагов
DDPM	стохастическая	$1000$
DDIM	детерминированная, каждый $k$-й шаг	$20{-}100$
Хойн по probability-flow ОДУ	детерминированная, второй порядок	$20{-}50$

Из урока 030: третий имеет смысл только потому, что второй детерминированный. Проверьте это измерением — сравните качество Хойна и Эйлера при равном числе вызовов сети, а не при равном числе шагов (урок 020).

4. Пять измерений

Как и в практике блока 6, работающая генерация ничего не доказывает. Проверьте утверждения ветки на своей реализации.

Измерение 1: замкнутая форма совпадает с цепью. Прогоните цепь $q(x_t\mid x_{t-1})$ сто раз и сравните выборочные среднее и дисперсию с $\sqrt{\bar\alpha}x_0$ и $1-\bar\alpha$. Ожидание — совпадение в пределах статистической ошибки. Это проверяет самый фундаментальный кирпич.

Измерение 2: score и $\varepsilon$ — одно и то же. Возьмите предсказание сети, переведите его в score по $s = -\varepsilon/\sqrt{1-\bar\alpha}$, затем восстановите $x_0$ по формуле Tweedie. Сравните с прямым переводом $\varepsilon \to x_0$. Ожидание: побитовое совпадение с точностью арифметики.

Измерение 3: DDIM детерминирован. Дважды сгенерируйте из одного и того же стартового шума. Ожидание: одинаковые изображения. У DDPM — разные. Если DDIM даёт разные, где-то остался случайный член.

Измерение 4: инверсия работает. DDIM обратим: пройдите от шума к изображению, затем обратно к шуму, и сравните с исходным шумом. Ожидание: расхождение мало и растёт при уменьшении числа шагов — это ошибка интегрирования, а не ошибка модели.

Измерение 5: guidance уменьшает разнообразие. Сгенерируйте по пятьдесят образцов при $w = 1, 3, 8$ и посчитайте попарное расстояние внутри каждого набора. Ожидание из урока 080: монотонное падение. Заодно посмотрите на долю пикселей, вышедших за допустимый диапазон, — она растёт.

5. Flow matching для сравнения

Та же архитектура, другая цель:

$\mathcal{L} = \mathbb{E}\big\|v_\theta(x_t, t) - (x_1 - x_0)\big\|^2, \qquad x_t = (1-t)x_0 + t x_1$

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

Ожидание из урока 090, сформулированное осторожно: путей прямее, шагов меньше, но не «один шаг». Если хотите увидеть настоящее спрямление, добавьте minibatch OT (урок 110): решите задачу назначения внутри батча и составляйте пары по ней, а не случайно.

flow matchingдиффузия

шагов сэмплера

4

длина пути

2.702

прямое расстояние

2.702

кривизна

×1

ошибка ломаной

0.0e+0

Прямой путь: число шагов перестаёт влиять на геометрию.

Куда смотреть, если не получается

симптом	что проверять первым
лосс на первом шаге сильно не единица	нормировку входа, масштаб $\varepsilon$
сэмплы — шум	знак в обратном шаге, порядок обхода $t$
сэмплы — серое пятно	подачу $t$ в сеть (модель усредняет по всем уровням шума)
DDIM даёт разные результаты из одного шума	оставшийся случайный член
сэмплы приемлемы при 1000 шагов и мусор при 50	использование стохастической формулы вместо ОДУ
guidance ломает картинку при $w > 5$	выход за диапазон, нужна обрезка

Третья строка — самая коварная и стоит объяснения. Если $t$ до сети не доходит (забытый эмбеддинг, потерянный аргумент), обучение не падает: сеть выучивает средний по всем $t$ шумоподавитель. Лосс при этом заметно выше оптимального, но снижается, а генерация даёт размытое нечто. Проверяется за минуту: подайте одно и то же $x_t$ при двух разных $t$ и убедитесь, что выход изменился.

Чем закончить

Прогоните все пять измерений на обученной модели. Особенно второе и третье: они проверяют понимание того, что предсказание шума, score и восстановление $x_0$ — один объект, и что детерминированное сэмплирование действительно детерминировано. Красивые картинки получаются и при ошибках; эти проверки — нет.