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

1. Идея: испортить и научиться чинить

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

Вероятностная модель шумоподавляющей диффузии (англ. Denoising Diffusion Probabilistic Model, DDPM; Соль-Дикстейн и соавт., 2015; существенно улучшена Хо и соавт., 2020) состоит из двух процессов:

• прямой диффузионный процесс постепенно добавляет шум к изображению; после последнего шага остаётся чистый шум — выборка из изотропного нормального распределения \(N(0, \mathbf{I})\); этот процесс задан формулой и не обучается;
• обратный диффузионный процесс учится обращать каждый шаг зашумления; применённый к случайному шуму, он порождает новое изображение.

Ключевое наблюдение: если известно более зашумлённое изображение \(x_t\) и сам шум \(\epsilon_t\), то менее зашумлённое \(x_{t-1}\) получается простым вычитанием. Значит, для обратного процесса достаточно нейросети, оценивающей шум между двумя соседними шагами.

2. Прямой диффузионный процесс

Начав с реального изображения \(x_0\) из распределения данных \(q(x_0)\), прямой процесс на каждом шаге добавляет небольшую порцию гауссовского шума, порождая последовательность всё более зашумлённых версий \(x_0, x_1, \ldots, x_T\). Условное распределение каждого шага — нормальное:

Набор \(\{\beta_t \in (0,1)\}_{t=1}^{T}\) называется расписанием дисперсий (англ. variance schedule) — он задаёт, сколько шума добавляется на каждом шаге. Типичный выбор — линейный рост от \(\beta_1 = 10^{-4}\) до \(\beta_T = 0{,}02\) при \(T = 1000\) шагов (Хо и соавт., 2020). Расписание подобрано так, что \(q(x_T | x_0) \approx N(0, \mathbf{I})\): в конце от изображения не остаётся ничего.

Чтобы получить зашумлённую версию шага \(t\), не нужно выполнять \(t\) шагов подряд. Обозначив

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

С ростом \(t\) коэффициент при изображении \(\sqrt{\bar\alpha_t}\) падает от 1 к 0, а коэффициент при шуме растёт от 0 к 1 — сигнал плавно замещается шумом:

3. Обратный диффузионный процесс

Цель генерации — выборка из \(q(x_{t-1} | x_t)\): взяв чистый шум \(x_T \sim N(0,\mathbf{I})\) и пройдя все шаги в обратном порядке, получить объект из распределения данных. Но истинное \(q(x_{t-1} | x_t)\) вычислить нельзя — оно требует знания распределения всех данных. Решение — аппроксимировать его нейросетью. Если шаги \(\beta_t\) малы, обратное распределение тоже близко к нормальному, поэтому сети достаточно предсказывать его среднее:

4. Чему обучать сеть: предсказанию шума

Формально модель обучается максимизировать правдоподобие данных — минимизировать \(E[-\log p_\Theta(x_0)]\). Напрямую эта величина не вычисляется, поэтому, как и в VAE (конспект 9), минимизируют её вариационную границу, которая распадается в сумму KL-дивергенций по шагам процесса. После подстановки нормальных распределений (10.1) и (10.4) каждое слагаемое сводится к разнице между истинным шумом и шумом, предсказанным сетью. Хо и соавторы показали, что взвешивающие коэффициенты можно просто отбросить — качество только улучшается, — и остаётся поразительно простая функция ошибки:

Словами: взять случайный шаг \(t\), зашумить изображение по ядру (10.3) случайным шумом \(\epsilon\), попросить сеть \(\epsilon_\Theta\) угадать этот шум по зашумлённому изображению и номеру шага — и штрафовать квадратом отклонения. Это обычная регрессия (конспект 2, формула (2.2)) — никакой состязательности, никакого балансирования двух сетей.

При генерации предсказанный шум используется для вычисления среднего (10.4):

где случайная добавка \(\sigma_t z\) присутствует на всех шагах, кроме последнего. Оба алгоритма целиком — обучение и генерация — умещаются в несколько строк:

5. Реализация на Keras

В источнике готового кода нет — ниже минимальная реализация DDPM для MNIST, написанная по формулам этого конспекта. Расписание дисперсий и зашумление:

import tensorflow as tf
import numpy as np

T = 1000
beta = np.linspace(1e-4, 0.02, T).astype('float32')   # расписание дисперсий
alpha_bar = np.cumprod(1.0 - beta)                       # формула (10.2)

def q_sample(x0, t, eps):
    # зашумление сразу до шага t — диффузионное ядро (10.3)
    a = tf.gather(alpha_bar, t)[:, None, None, None]
    return tf.sqrt(a) * x0 + tf.sqrt(1.0 - a) * eps

Сеть \(\epsilon_\Theta(x_t, t)\) получает зашумлённое изображение и номер шага; простейший способ подать номер шага — дополнительным каналом (в серьёзных реализациях используют синусоидальные эмбеддинги, как в трансформерах):

from tensorflow.keras import layers, Model

def build_eps_model():
    img = layers.Input(shape=(28, 28, 1))
    t_in = layers.Input(shape=(), dtype='int32')          # номер шага t
    t_map = layers.Lambda(lambda t: tf.cast(t[:, None, None, None], 'float32')
                          / T * tf.ones((1, 28, 28, 1)))(t_in)
    x = layers.Concatenate()([img, t_map])               # t — дополнительным каналом

    # мини-U-Net: кодировщик
    c1 = layers.Conv2D(32, 3, padding='same', activation='relu')(x)
    p1 = layers.MaxPooling2D()(c1)
    c2 = layers.Conv2D(64, 3, padding='same', activation='relu')(p1)
    p2 = layers.MaxPooling2D()(c2)
    b  = layers.Conv2D(128, 3, padding='same', activation='relu')(p2)
    # декодировщик с конкатенациями (конспект 7, раздел 6)
    u2 = layers.Conv2D(64, 2, padding='same', activation='relu')(layers.UpSampling2D()(b))
    m2 = layers.Concatenate()([c2, u2])
    c3 = layers.Conv2D(64, 3, padding='same', activation='relu')(m2)
    u1 = layers.Conv2D(32, 2, padding='same', activation='relu')(layers.UpSampling2D()(c3))
    m1 = layers.Concatenate()([c1, u1])
    c4 = layers.Conv2D(32, 3, padding='same', activation='relu')(m1)
    out = layers.Conv2D(1, 1)(c4)                          # предсказанный шум
    return Model([img, t_in], out)

model = build_eps_model()
optimizer = tf.keras.optimizers.Adam(2e-4)

Шаг обучения — алгоритм 1 из рисунка выше, дословно по формуле (10.5):

@tf.function
def train_step(x0):
    n = tf.shape(x0)[0]
    t = tf.random.uniform((n,), 0, T, dtype=tf.int32)   # случайный шаг для каждого примера
    eps = tf.random.normal(tf.shape(x0))                 # истинный шум
    with tf.GradientTape() as tape:
        eps_pred = model([q_sample(x0, t, eps), t], training=True)
        loss = tf.reduce_mean(tf.square(eps - eps_pred)) # формула (10.5)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

Генерация — алгоритм 2, формула (10.6):

def sample(n):
    x = tf.random.normal((n, 28, 28, 1))               # x_T — чистый шум
    for t in reversed(range(T)):
        eps_pred = model([x, tf.fill((n,), t)], training=False)
        x = (x - beta[t] / np.sqrt(1 - alpha_bar[t]) * eps_pred) / np.sqrt(1 - beta[t])
        if t > 0:                                     # случайная добавка — кроме последнего шага
            x += np.sqrt(beta[t]) * tf.random.normal(tf.shape(x))
    return x

6. Архитектура: U-Net с вниманием

Сеть \(\epsilon_\Theta\) должна принимать изображение и возвращать «изображение» того же размера (карту шума) — это ровно та задача, для которой в конспекте 7 вводилась U-Net: сжимающий путь копит признаки, разжимающий восстанавливает разрешение, горизонтальные связи переносят детали. В диффузионных моделях используется модификация U-Net: в блоки добавлены остаточные связи (конспект 7, раздел 5) с батч-нормализацией и блоки многоголового внимания (англ. multi-head attention), а номер шага \(t\) подаётся в каждый блок через эмбеддинг:

Применения диффузионных моделей те же, что у других генеративных подходов, и шире: генерация изображений по тексту (DALL-E 2, Imagen), преобразование изображения в изображение, повышение разрешения, сегментация, достройка трёхмерных форм. Платой за качество служит скорость: для одного изображения требуется \(T\) последовательных проходов сети, тогда как GAN (конспект 11) генерирует за один проход.

Контрольные вопросы

• Почему прямой диффузионный процесс не требует обучения?
  Он полностью задан формулой (10.1): на каждом шаге к изображению добавляется гауссовский шум с известными параметрами из расписания дисперсий. Обучается только обратный процесс — предсказание добавленного шума.
• Что такое диффузионное ядро и зачем оно нужно при обучении?
  Распределение q(x_t|x_0), позволяющее получить зашумлённую версию любого шага t сразу из исходного изображения: x_t = √ᾱ_t·x₀ + √(1−ᾱ_t)·ε (формула (10.3)). При обучении для каждого примера берётся случайный шаг t, и ядро строит x_t за одну операцию вместо t последовательных шагов.
• К чему в итоге сводится функция ошибки диффузионной модели?
  К обычной регрессии: среднему квадрату отклонения между истинным шумом ε, добавленным по ядру (10.3), и шумом, предсказанным сетью по зашумлённому изображению и номеру шага (формула (10.5)). Вариационная граница с KL-дивергенциями после упрощений (λ_t = 1) сводится именно к этому.
• Почему в качестве сети предсказания шума используют U-Net?
  Вход и выход задачи — изображения одного размера (зашумлённое изображение → карта шума), а U-Net создана ровно для таких задач: сжимающий путь извлекает признаки, разжимающий восстанавливает разрешение, горизонтальные связи сохраняют детали (конспект 7, раздел 6).
• Чем генерация диффузионной моделью отличается от генерации декодировщиком VAE по скорости и почему?
  VAE генерирует за один проход декодировщика по случайному латентному вектору. Диффузионной модели нужно T последовательных проходов сети (обычно сотни–тысячи): каждый шаг вычитает лишь малую порцию шума по формуле (10.6).
• Зачем при генерации на каждом шаге, кроме последнего, добавляется случайное слагаемое σ_t·z?
  Обратный процесс — вероятностный: p(x_{t−1}|x_t) — распределение, а не точка, и случайная добавка реализует выборку из него, обеспечивая разнообразие результатов. На последнем шаге её убирают, чтобы вернуть чистое изображение без остаточного шума.