Diffusion Models

Диффузионные модели — это модели генеративного ИИ, которые создают данные, обучаясь обращать постепенный процесс зашумления. При обучении изображения или другие данные искажаются гауссовым шумом за множество небольших шагов, а нейронная сеть учится предсказывать шум на каждом шаге. Для генерации модель начинает со случайного шума и пошагово удаляет его, пока не появится новый образец. Яша Золь-Дикштейн и коллеги предложили этот подход в 2015 году в статье «Глубокое обучение без учителя с использованием неравновесной термодинамики», а Джонатан Хо, Аджай Джайн и Питер Аббил сделали его конкурентоспособным в 2020 году с помощью метода вероятностных диффузионных моделей с устранением шума (DDPM), который установил новый рекорд FID — 3,17 на CIFAR-10. В 2021 году исследователи OpenAI показали, что диффузионные модели превосходят GANs в синтезе изображений ImageNet. Сегодня диффузионные модели и родственные им методы сопоставления потоков лежат в основе систем генерации изображений по тексту, таких как Stable Diffusion, генераторов видео, редакторов изображений и научных моделей.

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

Прямой процесс не обучается, а задается заранее. На каждом из T шагов к данным добавляется немного гауссова шума; в DDPM используется T = 1 000, а дисперсия шума линейно растет от 0.0001 до 0.02 на шаг. Поскольку сумма гауссовых шумов также является гауссовой, обучающее изображение можно напрямую преобразовать в любой шаг t по замкнутой формуле, xt = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * noise, где alpha_bar_t — произведение 1 - beta за первые t шагов. При обучении многократно выбираются изображение, случайный шаг и новый шум, формируется xt, а сеть учится предсказывать этот шум с помощью среднеквадратичной ошибки — упрощенной целевой функции DDPM. Этот код применяет прямой процесс к фотографии с расписанием DDPM:

import urllib.request

import numpy as np
import torch
from PIL import Image

# DDPM's linear schedule: 1,000 steps whose noise variance rises from 0.0001 to 0.02
betas = torch.linspace(1e-4, 0.02, 1000)
alpha_bar = torch.cumprod(1 - betas, dim=0)

# 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((256, 256))
x0 = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1) / 127.5 - 1

# Jump straight to step t in closed form: x_t = sqrt(alpha_bar) * x0 + sqrt(1 - alpha_bar) * noise
frames = []
for t in (0, 99, 299, 599, 999):
    xt = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * torch.randn_like(x0)
    frames.append(((xt.clamp(-1, 1) + 1) * 127.5).byte().permute(1, 2, 0).numpy())
    print(f"step {t + 1}: image weight {alpha_bar[t].sqrt():.3f}, noise weight {(1 - alpha_bar[t]).sqrt():.3f}")
Image.fromarray(np.hstack(frames)).save("forward_diffusion.jpg")

Вес изображения снижается с 1.000 на шаге 1 до 0.947 на шаге 100, 0.630 на шаге 300, 0.161 на шаге 600 и 0.006 на шаге 1 000, а вес шума растет до 1.000. На сохраненном изображении forward_diffusion.jpg автобус сначала четкий, затем зернистый, потом едва заметный и наконец исчезает.

При генерации процесс идет в обратном направлении. Он начинается с чистого гауссова шума; на каждом шаге сеть предсказывает шум в текущем образце, алгоритм решения удаляет его часть, а немного очищенный результат передается на следующий шаг. Текстовая подсказка или метка класса поступает в сеть как условие, часто через перекрестное внимание. Безклассификаторное управление, предложенное Джонатаном Хо и Тимом Салимансом, обучает одну сеть как с условием, так и без него, а во время генерации экстраполирует от безусловного предсказания в сторону условного. Более высокий коэффициент управления, например используемое по умолчанию значение 7.5 в конвейерах Stable Diffusion 1.x, точнее следует подсказке за счет снижения разнообразия.

Ускорение генерации#

Для генерации методом DDPM нужно 1 000 вычислений сети на изображение, поэтому значительная часть прогресса в этой области связана с сокращением их количества:

• Улучшенные алгоритмы генерации: DDIM повторно использует обученную DDPM с детерминированным алгоритмом генерации и создает высококачественные образцы в 10–50 раз быстрее по реальному времени. Более поздние решатели ОДУ сокращают типичную генерацию примерно до 20–50 шагов.
• Латентная диффузия: модель латентной диффузии устраняет шум в сжатом представлении, полученном автоэнкодером, а не в пикселях. Автоэнкодер Stable Diffusion преобразует изображение размером 512×512×3 в латентное представление 64×64×4, содержащее в 48 раз меньше значений.
• Дистилляция: модели согласованности учатся напрямую преобразовывать любой зашумленный образец в чистый, что позволяет генерировать за один шаг; латентные модели согласованности применяют ту же идею к Stable Diffusion, сокращая генерацию до 2–4 шагов.
• Более прямые траектории: сопоставление потоков и выпрямленный поток обучаются на прямых траекториях от шума к данным, по которым решатели могут пройти за меньшее число шагов.

Сама сеть устранения шума перешла от сверточной архитектуры U-Net, использовавшейся в DDPM и Stable Diffusion 1.x, к диффузионному трансформеру (DiT), который применяется в большинстве современных генераторов изображений и видео.

Диффузионные модели и другие генеративные модели#

• Генеративно-состязательные сети (GANs): GAN генерирует образец за один прямой проход, однако состязательное обучение генератора и дискриминатора может быть нестабильным и приводить к схлопыванию мод, когда генератор охватывает лишь часть данных. Обучение диффузионных моделей представляет собой стабильную регрессию и полнее охватывает распределение данных, но требует множества шагов генерации.
• Вариационные автоэнкодеры (VAEs): VAE декодирует латентный вектор в образец за один проход, но полученные образцы обычно более размыты. Модели латентной диффузии используют автоэнкодер в стиле VAE только для сжатия изображений, а генерацию выполняет диффузионная модель.
• Авторегрессионные модели: языковые и визуальные авторегрессионные модели генерируют по одному токену или масштабу за раз, опираясь на предыдущие данные. Диффузионная модель на каждом шаге параллельно уточняет все части образца, а диффузионные языковые модели применяют тот же процесс устранения шума к тексту.

Примеры применения в реальных условиях#

• Генерация и редактирование изображений: модели генерации изображений по тексту создают изображения на основе подсказок; заливка изображений заново генерирует замаскированную область с учетом окружения; ControlNet задает структуру с помощью карт границ, глубины или позы, а диффузионные модели увеличения разрешения выполняют сверхразрешение.
• Генерация видео: модели генерации видео устраняют шум в сжатых латентных представлениях сразу множества кадров, поэтому каждый кадр генерируется вместе с соседними.
• Синтетические обучающие данные: сгенерированные изображения восполняют пробелы в наборах данных для компьютерного зрения, например добавляют редкие дефекты, необычные погодные условия или нетипичные ракурсы. В наборе данных на платформе Ultralytics действие «Создать похожие изображения» создает от 1 до 16 новых изображений за пакет на основе исходного, а выбранные изображения добавляются в набор данных без разметки и готовы к аннотированию.
• Робототехника: диффузионные политики создают последовательности действий робота путем устранения шума, используя изображения с камеры в качестве условия.
• Наука: AlphaFold 3 предсказывает структуры комплексов белков, нуклеиновых кислот и небольших молекул с помощью диффузионного модуля, устраняющего шум в исходных координатах атомов и заменяющего модуль структуры AlphaFold 2.