Latent diffusion model

Вначале вариационный автокодировщик (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]