Latent Diffusion Model (LDM)

Латентная диффузионная модель (LDM) — это диффузионная модель, которая генерирует данные, устраняя шум в сжатом представлении предварительно обученного автоэнкодера, а не в исходных пикселях. Робин Ромбах, Андреас Блаттманн, Доминик Лоренц, Патрик Эссер и Бьёрн Оммер представили этот подход в работе «Синтез изображений высокого разрешения с помощью латентных диффузионных моделей», представленной на CVPR 2022. Автоэнкодер удаляет мелкие детали пикселей, моделирование которых требует больших затрат и мало влияет на внешний вид изображения, поэтому диффузионная сеть сосредоточивает ресурсы на содержании и композиции. Обучение и генерация поэтому обходятся намного дешевле, чем диффузия в пространстве пикселей. Благодаря этому Stable Diffusion, самая известная LDM, генерирует изображения на одном пользовательском GPU. Большинство современных генераторов изображений, видео и аудио в области генеративного ИИ используют ту же двухэтапную архитектуру.

Как работает латентная диффузия#

На первом этапе сжимаются изображения. Кодировщик преобразует изображение размером H×W×3 в латентное представление размером H/f×W/f×c, где f — коэффициент уменьшения разрешения, а c — число латентных каналов. Декодер преобразует латентные представления обратно в пиксели. Ромбах и его коллеги обучали этот автоэнкодер с перцептивной функцией потерь и состязательной целевой функцией на основе патчей, как у дискриминатора GAN, чтобы реконструкции сохраняли резкость. Латентное представление регуляризовали либо с помощью небольшой KL-штрафной функции, направляющей его к стандартному нормальному распределению, либо с помощью векторного квантования. Коэффициенты уменьшения разрешения от 4 до 16 обеспечивали баланс скорости и качества: обучение в пространстве пикселей (f = 1) шло медленно, а при f = 32 сжатие было настолько сильным, что качество изображений снижалось. В автоэнкодере Stable Diffusion используются f = 8 и 4 канала. Этот код кодирует фотографию, декодирует её и оценивает потери после такого цикла:

import urllib.request

import numpy as np
import torch
from diffusers import AutoencoderKL
from PIL import Image

# Stable Diffusion's autoencoder, the first stage of a latent diffusion model
vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")

# Load a photo and scale its pixels to [-1, 1]
image = Image.open(urllib.request.urlopen("https://ultralytics.com/images/bus.jpg")).convert("RGB").resize((480, 640))
x = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1)[None] / 127.5 - 1

# Encode to the latent that diffusion runs on, then decode back to pixels
with torch.no_grad():
    latent = vae.encode(x).latent_dist.mode()
    recon = vae.decode(latent).sample.clamp(-1, 1)

psnr = 10 * torch.log10(4 / torch.mean((recon - x) ** 2))
print(f"image {tuple(x.shape)} -> latent {tuple(latent.shape)}: {x.numel() // latent.numel()}x fewer values")
print(f"reconstruction PSNR: {psnr:.1f} dB")
Image.fromarray(((recon[0] + 1) * 127.5).byte().permute(1, 2, 0).numpy()).save("bus_reconstructed.jpg")

Он выводит image (1, 3, 640, 480) -> latent (1, 4, 80, 60): 48x fewer values и reconstruction PSNR: 23.0 dB. Сохранённое bus_reconstructed.jpg сохраняет автобус, людей и улицу, но мелкая надпись на борту автобуса искажается: латентное представление из 4 каналов не может сохранить такие детали.

На втором этапе обычная диффузионная модель обучается на этих латентных представлениях. При обучении изображение кодируется, к нему на случайном шаге добавляется гауссовский шум, а сеть устранения шума — в исходной работе это U-Net, а в большинстве новых моделей — диффузионный трансформер, — обучается предсказывать этот шум. В моделях с выпрямленным потоком сеть предсказывает скорость, с которой шум переходит в чистое латентное представление. Условия вводятся через слои перекрёстного внимания, позволяющие латентному представлению учитывать эмбеддинги текстового запроса, метки класса или семантической карты. Генерация начинается со случайного латентного шума, который шаг за шагом очищается с учётом заданного условия, после чего декодер запускается один раз. Таким образом, каждый ресурсоёмкий шаг выполняется на небольшом латентном представлении.

От Stable Diffusion до современных моделей#

• Stable Diffusion: Выпущенная в августе 2022 года, она распространила подход LDM на генерацию изображений 512×512 по текстовому описанию. Модель использует латентное представление размером 64×64×4, U-Net с 860 миллионами параметров и текстовый кодировщик CLIP; скачать веса мог любой желающий.
• Stable Diffusion 3: Перешла на автоэнкодер с 16 каналами, повысив PSNR реконструкции с 25.12 дБ при 4 каналах до 28.62 дБ, и заменила U-Net мультимодальным диффузионным трансформером, обученным с помощью выпрямленного потока. Модели FLUX.1 от Black Forest Labs также устраняют шум в латентных представлениях с 16 каналами с помощью трансформера выпрямленного потока.
• Видео и аудио: Генераторы видео сжимают ролики в пространстве и времени. OpenAI описала Sora как модель, которая сжимает видео в латентное пространство меньшей размерности и разбивает его на пространственно-временные патчи. AudioLDM и Stable Audio от Stability AI применяют тот же подход к латентным представлениям спектрограмм или звуковых волн.
• Модели согласованности в латентном пространстве: Дистиллируют LDM, сокращая число шагов генерации с 20–50 до 2–4, благодаря чему становится возможна почти мгновенная генерация.

LDM и связанные понятия#

• Диффузионные модели в пространстве пикселей: Выполняют добавление и устранение шума для каждого пикселя. LDM — это тоже диффузионная модель; различие лишь в том, где выполняется этот процесс.
• Вариационные автоэнкодеры: VAE выполняет генерацию за один проход, декодируя выбранное латентное представление, но получаемые изображения обычно размыты. LDM использует автоэнкодер в стиле VAE только для сжатия, а генерацию выполняет диффузионная модель.
• Stable Diffusion: Одно семейство LDM от Stability AI, CompVis и Runway, поэтому любая модель Stable Diffusion — это LDM, а к LDM также относятся FLUX, генераторы видео и аудиомодели.

Применение#

LDM лежат в основе генерации изображений по тексту, восстановления изображений, где в первоначальной статье был достигнут новый на тот момент лучший результат, сверхвысокого разрешения и генерации изображений по макету на основе семантических карт. В компьютерном зрении они также создают синтетические данные: редкие дефекты, необычные погодные условия или нечастые ракурсы, которые трудно запечатлеть на фото. У сгенерированных изображений нет меток, поэтому функция Ultralytics auto_annotate размечает папку с ними, обнаруживая объекты с помощью модели YOLO26 и обводя каждую найденную рамку моделью сегментации Meta SAM. Функция записывает метки сегментации в формате YOLO для обучения.