Diffusion Transformer (DiT)

Diffusion Transformer (DiT) — это диффузионная модель, в которой для подавления шума используется трансформер, а не свёрточная сеть U-Net. Модель разбивает зашумлённый латентный образ изображения на патчи, обрабатывает их как последовательность токенов с помощью self-attention, как Vision Transformer (ViT) обрабатывает изображение, и предсказывает шум, который нужно удалить на каждом шаге подавления шума. Уильям Пиблс и Сайнинг Се представили её в работе Scalable Diffusion Models with Transformers (ICCV 2023): DiT-XL/2 достиг показателя расстояния Фреше—Инцепшн (FID) 2,27 на ImageNet 256×256 с условием по классу, превзойдя предыдущий лучший результат 3,60 у латентной диффузии. В статье показано, что качество изображений стабильно растёт вместе с вычислительными затратами DiT. С тех пор трансформеры, подавляющие шум, заменили U-Net в большинстве крупных генераторов изображений и видео, включая Stable Diffusion 3, FLUX.1 и Sora от OpenAI.

Как работает Diffusion Transformer#

DiT работает в латентном пространстве вариационного автокодировщика (VAE), предварительно обученного для Stable Diffusion. Его энкодер уменьшает изображение в 8 раз, поэтому изображение 256×256×3 преобразуется в латентный образ 32×32×4. Слой разбиения на патчи делит зашумлённый латентный образ на патчи размером p×p и линейно преобразует каждый из них в токен; размеры патчей 2, 4 и 8 дают соответственно 256, 64 и 16 токенов. Синусно-косинусные позиционные эмбеддинги указывают, откуда взялся каждый токен, последовательность обрабатывается набором блоков трансформера, а заключительный линейный слой преобразует каждый токен обратно в патч с предсказанными шумом и дисперсией, после чего патчи собираются в латентный образ исходной формы. Каждый блок получает условия в виде шага диффузии и метки класса либо текстового запроса в модели преобразования текста в изображение. Генерация повторяет предсказание, начиная с чистого шума и заканчивая последним шагом, а декодер VAE преобразует очищенный латентный образ в пиксели:

В статье сравниваются четыре способа передать условие в каждый блок: дополнительные токены в последовательности, cross-attention, адаптивная нормализация слоя (adaLN) и adaLN-Zero. Вариант adaLN-Zero показал наименьший FID и оказался самым вычислительно эффективным. Небольшая сеть обрабатывает сумму эмбеддингов шага диффузии и класса и предсказывает сдвиг и масштаб, применяемые после каждой нормализации слоя, а также вентиль, регулирующий масштаб каждой остаточной ветви. Вентили изначально равны нулю, поэтому обучение каждого блока начинается с тождественной функции. Этот фрагмент кода на PyTorch разбивает латентный образ на патчи и применяет ветвь внимания одного блока:

import torch
from torch import nn

torch.manual_seed(0)
# A 256x256 image becomes a 4x32x32 latent; 2x2 patches turn it into 256 tokens of width 384 (DiT-S/2)
latent = torch.randn(1, 4, 32, 32)
tokens = nn.Conv2d(4, 384, kernel_size=2, stride=2)(latent).flatten(2).transpose(1, 2)

# adaLN-Zero: the timestep and class embedding regresses a shift, scale, and gate for the attention branch
condition = torch.randn(1, 384)
modulation = nn.Linear(384, 3 * 384)
nn.init.zeros_(modulation.weight)
nn.init.zeros_(modulation.bias)
shift, scale, gate = modulation(condition).unsqueeze(1).chunk(3, dim=-1)

# Self-attention branch of a DiT block, modulated by the condition and scaled by the gate
attention = nn.MultiheadAttention(384, num_heads=6, batch_first=True)
h = nn.functional.layer_norm(tokens, (384,)) * (1 + scale) + shift
out = tokens + gate * attention(h, h, h)[0]
print(f"tokens {tuple(tokens.shape)}, block is identity at initialization: {torch.equal(out, tokens)}")

Код выводит tokens (1, 256, 384), block is identity at initialization: True. 256 токенов шириной 384 соответствуют самой маленькой модели DiT-S/2, а нулевой вентиль пропускает вход без изменений, пока обучение его не изменит. Полный блок DiT предсказывает ещё один сдвиг, масштаб и вентиль для своей ветви MLP.

Масштабирование по числу токенов и вычислительным затратам#

В статье обучили модели четырёх размеров — от DiT-S (12 слоёв, ширина 384) до DiT-XL (28 слоёв, ширина 1 152), — каждую с размерами патчей 8, 4 и 2. Для всех 12 моделей FID соответствовал вычислительным затратам прямого прохода в Gflops, независимо от того, росли ли затраты из-за большей глубины или ширины сети либо из-за увеличения числа токенов. Уменьшение размера патча вдвое увеличивает число токенов и вычислительные затраты в четыре раза, почти не меняя число параметров. В DiT-XL/2 675 миллионов параметров, не считая 84-миллионного VAE, а один прямой проход при разрешении 256×256 требует 119 Gflops. При разрешении 512×512 модель обрабатывает 1 024 токена из латентного образа 64×64×4, затрачивая 525 Gflops, и достигает FID 3,04. Поскольку self-attention сопоставляет каждый токен со всеми остальными, увеличение разрешения и длительности видео повышает его стоимость квадратично. Ultralytics YOLO26 избегает этих затрат: механизм внимания используется только на этапе с масштабом 1/32, где для изображения 640×640 есть 400 позиций, а на всех остальных этапах применяются свёртки, обеспечивающие детекцию в реальном времени.

DiT, U-Net и ViT: сравнение#

• U-Net для подавления шума: U-Net в DDPM и Stable Diffusion 1.x и 2.x — это свёрточные энкодеры-декодеры, работающие с несколькими разрешениями, с пропускными соединениями между соответствующими уровнями и механизмом внимания на низких разрешениях. DiT использует одно разрешение токенов, применяет глобальное внимание в каждом блоке и имеет меньше архитектурных особенностей, связанных с изображениями, поэтому масштабируется так же, как трансформеры в больших языковых моделях.
• Vision Transformer: ViT кодирует изображение для таких задач, как классификация. DiT повторно использует его эмбеддинги патчей и блоки трансформера, добавляет условия шага диффузии через adaLN-Zero и выдаёт предсказание для каждого патча.
• Латентная диффузия и сопоставление потоков: Латентная диффузионная модель — это общий подход к подавлению шума в латентном пространстве VAE, а диффузия или сопоставление потоков — это целевая функция обучения. DiT — сеть для подавления шума в рамках этого подхода, а современные генераторы объединяют архитектуры DiT с целевыми функциями сопоставления потоков.

DiT для преобразования текста в изображения и видео#

• Мультимодальный DiT (MMDiT): Stable Diffusion 3 обрабатывает токены изображений и текста с помощью отдельных весов для каждой модальности и объединяет их в одной операции внимания, благодаря чему информация передаётся в обоих направлениях между запросом и изображением. Это улучшило типографику и следование запросам, а Stable Diffusion 3.5 Large масштабирует эту архитектуру примерно до 8 миллиардов параметров.
• FLUX.1: FLUX.1 от Black Forest Labs — это трансформер rectified flow с 12 миллиардами параметров для генерации изображений по тексту.
• Видео: В техническом отчёте Sora от OpenAI описан диффузионный трансформер, который сжимает видео в латентное пространство и делит их на пространственно-временные патчи, поэтому клипы разной длительности, разрешения и соотношения сторон преобразуются в последовательности токенов разной длины. Для генерации видео крупнейшая модель Movie Gen от Meta использует трансформер с 30 миллиардами параметров и контекстом из 73 000 видеотокенов, соответствующим 16 секундам видео при 16 кадрах в секунду.