Latent diffusion model

Диффузионные модели были впервые представлены в 2015 году как метод построения моделей для генерации сложных вероятностных распределений, используя методы неравновесной термодинамики и диффузии^[4]. В то же время была предложена реализация на Theano^[5].

В 2019 году была предложена модель noise conditional score network (NCSN) или score-matching с динамикой Ланжевена (SMLD)^[6]. К статье прилагался программный пакет на PyTorch, опубликованный на GitHub^[7].

В 2020 году опубликована работа^[8], где была предложена Denoising Diffusion Probabilistic Model (DDPM), улучшающая прежние методы за счёт вариационного вывода. В публикации приводился программный пакет на TensorFlow, выпущенный на GitHub^[9]. Впоследствии реализована на PyTorch другим автором^[10][11].

20 декабря 2021 года статья о LDM появилась на arXiv^[12], а репозитории Stable Diffusion^[13] и LDM^[14] были опубликованы на GitHub. Существенная информация по Stable Diffusion v1 была добавлена только 10 августа 2022 года^[15].

Все версии Stable Diffusion (SD) с 1.1 по XL являются конкретными реализациями архитектуры LDM.

SD 1.1–1.4 выпущены CompVis в августе 2022 года (версии 1.0 не существует). SD 1.1 представляла собой LDM, обученную на датасете laion2B-en, и далее дообучалась до 1.2 на более эстетичных изображениях. Версии 1.2 была дообучена до 1.3, 1.4 и 1.5, с уменьшением текстового кондиционирования на 10 % для улучшения guidance^[16][17]. Версия SD 1.5 была выпущена компанией RunwayML в октябре 2022 года^[17].

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

LDM состоит из вариационного автокодировщика (VAE), модифицированного U-Net и текстового энкодера.

Энкодер VAE сжимает изображение из пространства пикселей в компактное латентное пространство, отражающее семантическое содержание изображения. К этому латентному представлению применяется гауссов шум в процессе прямой (forward) диффузии. Блок U-Net, построенный на базе ResNet, постепенно устраняет шум, двигаясь в обратном (reverse) направлении, чтобы восстановить скрытое представление. После этого декодер VAE восстанавливает изображение из латентного пространства в пиксельное^[3].

Шаг денойзинга может быть обусловлен строкой текста, изображением или другой модальностью. Закодированные условные данные подаются в U-Net через механизм cross-attention^[3]. Для текстового кондиционирования используется зафиксированный, предобученный текстовый энкодер CLIP, преобразующий запросы в эмбеддинги^[2].

Вначале вариационный автокодировщик (VAE) обучается на выборе изображений. Энкодер принимает изображение на вход и выдаёт его низкоразмерное латентное представление, которое затем используется в качестве входа для U-Net. После обучения энкодер применяется для кодирования изображений в латентные представления, а декодер — для обратного преобразования латентных в изображения.

Пусть энкодер и декодер VAE обозначаются ${\displaystyle E}$ и ${\displaystyle D}$.

Для кодирования RGB-изображения его три канала делятся на максимальное значение, образуя тензор ${\displaystyle x}$ размера ${\displaystyle (3,512,512)}$ с элементами ${\displaystyle [0,1]}$. Кодировка представляет собой ${\displaystyle 0.18215\times E(2x-1)}$ размером ${\displaystyle (4,64,64)}$, где 0.18215 — гиперпараметр, выбранный авторами для приблизительного беления вектора (unit variance). Обратно, для латентного тензора ${\displaystyle y}$ декодированное изображение вычисляется как ${\displaystyle (D(y/0.18125)+1)/2}$, обрезается по диапазону ${\displaystyle [0,1]}$^[18][19].

В исходной реализации^[2] энкодер — это сверточная нейронная сеть (CNN) с единственным механизмом self-attention ближе к выходу (выдаёт тензор ${\displaystyle (8,H/8,W/8)}$ — объединённые матожидание и дисперсия; после обучения используется только матожидание). Декодер также CNN с self-attention и преобразует тензор ${\displaystyle (4,H/8,W/8)}$ в изображение ${\displaystyle (3,H,W)}$.

Базовый блок U-Net принимает:

• Латентный массив изображения, полученный от энкодера VAE размером ${\displaystyle ({\text{каналы}},{\text{ширина}},{\text{высота}})}$. Обычно — ${\displaystyle (4,64,64)}$.
• Вектор временной эмбеддинга (timestep), показывающий степень зашумления: ${\displaystyle t=0}$ — отсутствие шума, ${\displaystyle t=100}$ — сильный шум.
• Последовательность эмбеддингов модальности, отражающих дополнительные условия (например, текстовые токены, закодированные через CLIP или последовательности Vision Transformer для генерации в определённом стиле).

Каждый проход через U-Net выдаёт предсказанный вектор шума, который масштабируется и вычитается из латентного массива, постепенно убирая шум. Процедура повторяется до конца графика шумов (noise schedule), итог обрабатывается VAE-декодером.

Структурно U-Net идентичен классической модели: сначала слои с уменьшением размерности, затем с увеличением. В SD 1.5 backbone дополнен модулями для работы с эмбеддингами. Например, слой уменьшения размерности:

• Латентный массив и time-embedding проходят через ResBlock:
  • Массив — через свёрточный слой.
  • Time-embedding — через однослойный персептрон, затем результат прибавляется к предыдущему (широковещание на все пиксели).
  • Затем новый свёрточный слой и ещё один time-embedding.
• Латентный массив и последовательность эмбеддингов — через SpatialTransformer (аналог decoder-а Transformer без causal masking):
  • В блоках cross-attention массив служит query–последовательностью (по 1 вектору на пиксель), эмбеддинги — и как ключи, и как значения.
  • Если нет внешней последовательности эмбеддингов, cross-attention превращается в self-attention^[20].

Псевдокод:

def ResBlock(x, time, residual_channels):
    x_in = x
    time_embedding = feedforward_network(time)
    x = concatenate(x, residual_channels)
    x = conv_layer_1(activate(normalize_1(x))) + time_embedding
    x = conv_layer_2(dropout(activate(normalize_2(x))))
    return x_in + x

def SpatialTransformer(x, cond):
    x_in = x
    x = normalize(x)
    x = proj_in(x)
    x = cross_attention(x, cond)
    x = proj_out(x)
    return x_in + x

def unet(x, time, cond):
    residual_channels = []
    for resblock, spatialtransformer in downscaling_layers:
        x = resblock(x, time)
        residual_channels.append(x)
        x = spatialtransformer(x, cond)

    x = middle_layer.resblock_1(x, time)
    x = middle_layer.spatialtransformer(x, time)
    x = middle_layer.resblock_2(x, time)

    for resblock, spatialtransformer in upscaling_layers:
        residual = residual_channels.pop()
        x = resblock(concatenate(x, residual), time)
        x = spatialtransformer(x, cond)

    return x

Детальное описание есть в источниках^[21].^[22]

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

• Процесс прямой диффузии: дано изображение ${\displaystyle x_{0}}$, к нему последовательно добавляется гауссов шум по заданному графику, образуя последовательность латентных переменных ${\displaystyle x_{1:T}}$.
• Обратный процесс: начиная с ${\displaystyle x_{T}}$ (шумовое изображение), модель учится восстанавливать/предсказывать шум на каждом шаге, чтобы обрести восстановленное ${\displaystyle x_{0}}$.

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

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

Подробнее — см. статью диффузионная модель.