Диффузия на трансформерах

Введение

Недавно OpenAI представили миру диффузионную модель SORA. Она способна генерировать видео потрясающего качества на основе картинки, текстового промпта или даже видео.

Конечно, деталей, архитектуру, процесс обучения нам открывать не стали. Зато сослались на некоторые полезные работы. В этом обзоре мы и поговорим о трансформерной архитектуре, которую применили в SORA (кстати, Stability выложили статью по StableDiffusion 3, которая тоже будет трансформерной).

Наш обзор скорее верхнеуровневый и без страшной математики 🙂

Сначала мы вспомним, что такое Vision Transformer и как он работает, поговорим о Masked Autoencoders для маскированных данных. А затем перейдем к статьям, где было предложено поменять стандартную архитектуру UNet на трансформерную, и узнаем, к какому результату это все привело.

Transformers for images

Для начала вспомним идею трансформерной архитектуры для задачи компьютерного зрения (CV). Она пришла в CV из NLP, оказалась очень полезной и часто используемой в самых различных задачах.

ViT

📄 Paper

📎 GitHub

Первая ключевая работа, на основе которой было построено множество других — это ViT, или Visual Transformer.

Авторы статьи постарались как можно точнее сохранить идею NLP transformers, применив ее к изображениям. Обычно трансформеры принимают на вход последовательность токенов (векторов). Но у нас есть картинка размером \( W \times H \times C \). Что можно сделать, чтобы превратить наше изображение в такую последовательность?

Давайте возьмем картинку и разрежем ее на одинаковые патчи меньшего размера — \( x^p_i \in \mathbb{R}^{P \times P \times C} \), затем проведем операцию Flatten и превратим патч в вектор размером \( P^2 \cdot C \). Тогда количество патчей будет равно \( N = \frac{HW}{P^2} \). Далее при помощи линейного преобразования проецируем это в некоторый вектор эмбеддинга \( z_i = x_i^p\mathbf{E} \), где \( E \in \mathbb{R}^{(P^2C) \times D} \)— матрица преобразования.

Еще нам нужно передать информацию о структуре изображения. Также, как и в задачах NLP, кодируем позиции токенов (positional embeddings), которые складываются с эмбеддингами патчей и подаются в финальный трансформер. Для более подробного ознакомления с ViT рекомендуем вам почитать этот пост 🙂

Эта работа очень важна: она получила широкое распространение, много различных продолжений и улучшений. Мы не будем подробно останавливаться на каждом, но отметим некоторые из них.

BEiT: BERT Pre-Training of Image Transformers

📄 Paper

📎 GitHub

Трансформерная архитектура часто используется для дополнительного предобучения в различных задачах. Например, работа BEiT предлагает способ предобучения без учителя ViT с помощью masked image modeling (MIM).

Изображение разбивается на патчи — \( x \in \mathbb{R}^{H \times W \times C} \rightarrow x^p \in \mathbb{R}^{N \times P^2C}, N=HW / P^2 \), некоторые из них маскируются. Затем каждый патч токенизируется — ставится в соответствие некоторому вектору из словаря:

где \( \mathcal{V} = \{1, …, |\mathcal{V}|\} \) — словарь размером 8192 (размер из статьи). Такой токенайзер можно выучить, используя дискретный вариационный автокодировщик dVAE. Подробнее об этом можно почитать в статье Dalle.

Обучение BEiT аналогично обучению вариационного автокодировщика, а именно — с помощью минимизации ELBO (Evidence lower bound). Она, в свою очередь, раскладывается на две составляющие:

• реконструкционная часть показывает качество реконструированного изображения
  • прогоняем изображение через токенайзер и получаем условное распределение на токены изображения \( q_{\phi}(z|x) \)
  • \( p_{\psi}(x|z) \) : декодируем изображение на основе входных токенов изображения
• KL-дивергенция
  • сравниваем распределение, полученное из токенайзера \( q_{\phi}(z|x) \) и нашего энкодера \( p_{\theta}(z|\hat{x}_i) \)

Таким образом, имеем следующее выражение для функции потерь:

Обучение проходит в два этапа: сначала минимизируем реконструкционную часть \( \mathbb{E}_{z_i \sim q_{\phi}(z|x_i)}\Big[\log p_{\psi}(x_i|z_i)\Big] \), а затем фиксируем токенайзер и декодер и обучаем только часть BEiT Encoder.

Результаты

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

Они показали: такое предобучение может заметно улучшить качество в различных задачах, например, в задаче классификации.

Это не последняя работа об эффективном предобучении на трасформерах (можно также почитать статью DINO или DINOV2). Но она наглядно показывает пользу такого предобучения.

ViTMAE

📄 Paper

📎 GitHub

А что делать, если наши данные “дырявые”? Пусть у нас есть патчи изображения, но они “повреждены”. Можно ли выучить модель так, чтобы она исправила “проблемные” части и восстановила полную картинку?

Именно эту задачу решает Masked Autoencoders (MAE). Он восстанавливает изображение на основе частично предоставленных данных, сначала кодируя его в некоторый латентный вектор, а затем декодируя обратно. В отличие от BEiT, здесь используется непрерывное латентное представление, полученное с помощью энкодера.

Обычно автокодировщики имеют симметричную структуру: архитектура декодера полностью повторяет архитектуру энкодера, за исключением UpSampling и DownSampling операций.

Авторы данной работы, наоборот, предлагают сделать ее ассиметричной. Это и помогает энкодеру работать с неполными данными (с некоторой маской) и декодировать полное изображение.

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

MAE Encoder

В качестве энкодера берется обычный ViT, однако применяется он только по отношению к видимым патчам — кодируются только неудаленные части изображения.

MAE Decoder

В декодер же подаются не только закодированные патчи из энкодера, но и маскированные патчи. Каждый маскированный токен — общий для всех изображений обучаемый вектор, который указывает на наличие пропущенного патча (его и нужно предсказать). Еще ко всем токенам добавляются эмбеддинги расположения — теперь становится понятным, где именно находятся эти маскированные токены.

В качестве функции потерь авторы предлагают использовать MSE, но при этом дополнительно проводить нормализацию пикселей в патче.

Результаты

Авторы показали: даже при сильном маскировании возможно восстановить изображение. Но стоит отметить: чем больше маска, тем хуже результат, картинки получаются слишком замыленными.

Еще они проверили, как влияет расположение маски (ее тип). Авторы использовали задачу классификации:

• сначала предобучали нейронную сеть методом ViTMAE, описанным выше;
• затем решали задачу классификации с помощью:
  • дообучения (fine-tuning)
  • linear probing

Для двух сценариев качество предсказания задачи классификации для различных типов маски приведено ниже. Можно заметить: при большой степени маскированности (75%) качество сильно падает, особенно на методе linear-probing.

Случайная маска при этом дает максимально качество при данных условиях.

Таким образом, ViTMAE может активно применяться в задачах предобучения.

Transformers for diffusion models

Итак, трансформерная архитектура показала свое качество и результативность даже на “неплотных” данных, где часть информации является утерянной.

Тем не менее долгое время у диффузионных моделей оставалась архитектура UNeT. Лишь позже появилось несколько работ, использующих трансформерную архитектуру. Появилось и снова забылось… до тех пор, пока на сцену не вышли SORA и Stable Diffusion3. Давайте же рассмотрим эти работы 🙂

Scalable Diffusion Models with Transformers

📄 Paper

📎 GitHub

Здесь авторы предлагают модель DiT, или Diffusion Transformers. Она базируется в первую очередь на архитектуре ViT.

На вход в DiT подается картинка / латент в случае латентной диффузии, которые затем перегоняются в последовательность токенов и кодируются с помощью positional encoding:

Важный вопрос — вопрос кондишнинга. Авторы рассматривают несколько возможных вариантов:

In-context conditioning

Просто суммируем эмбеддинг времени и кондишнинг, а затем добавляем к последовательности токенов картинки, как и cls token.

Cross-attention block

Конкатенируем эмбеддинг времени и кондишнинг, а затем добавляем дополнительный блок в трансформер — cross-attention, работающий аналогично трансформерному блоку в UNet архитектуре диффузионной модели.

Adaptive layer norm (adaLN) block

Вспоминаем ганы и используем вместо обычных нормализующих слоев — adaptive layer norm (adaLN), следовательно, получаем shift и scale из кондишнинга.

adaLN-Zero block

Финальный вариант происходящего (= adaLN block, но с хорошей инициализацией) 🙂

Авторы вводят дополнительный scaling параметр \( \alpha \), который используется для residual connections (суммирование входа и выхода трансформерного блока). Основное отличие от adaLN — специальная начальная инициализация residial blocks, чтобы на первом шаге они выполняли роль identity функции. Такая инициализация способствует более стабильному обучению, что было показано, например, тут.

# Zero-out adaLN modulation layers in DiT blocks:
  for block in self.blocks:
      nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
      nn.init.constant_(block.adaLN_modulation[-1].bias, 0)

  # Zero-out output layers:
  nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
  nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
  nn.init.constant_(self.final_layer.linear.weight, 0)
  nn.init.constant_(self.final_layer.linear.bias, 0)

class DiTBlock(nn.Module):
    """
    A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
    """
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        approx_gelu = lambda: nn.GELU(approximate="tanh")
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x

class FinalLayer(nn.Module):
    """
    The final layer of DiT.
    """
    def __init__(self, hidden_size, patch_size, out_channels):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
        x = modulate(self.norm_final(x), shift, scale)
        x = self.linear(x)
        return x

Авторы дополнительно сравнивают качество различных способов кондишнинга и представляют график, показывающий, что adaLN-Zero сходится быстрее и лучше всех остальных вариантов.

Результаты

Авторы проводят ряд экспериментов на class-conditional генерации, изменяя размер трансформерной модели и патча. Ниже представлено изображение, где заметно улучшение качества генерации при увеличении размера модели и уменьшение размера патча:

Также отдельно показаны результаты генерации трансформерной диффузионной модели, обученной на ImageNet на различных разрешениях: \( 512 \times 5 \) и \( 256 \times 256 \).

Что при этом происходит с метриками?

Ниже представлено сравнение по нескольким метрикам:

• FID
• sFID
• IS
• Precision и Recall (задачи генерации по классу)

С чем сравнивается модель:

• StyleGAN-XL
• BigGan
• ADM
• LDM

Стоит отметить: предложенная авторами модель выигрывает по метрикам практически все модели со стандартной архитектурой UNet. Значит, архитектура трансформера действительно вносит заметный вклад в качество генерации модели 🙂.

All are Worth Words: A ViT Backbone for Diffusion Models

📄 Paper

📎 GitHub

Еще одна работа с идеей использования трансформерной архитектуры для диффузии (латентной диффузии). Авторы, как и в предыдущей статье, предлагают архитектуру трансформера, но добавляют при этом ряд важных изменений.

Основная идея: добавление long skip connections. Она заимствована из UNet архитектуры, которая позволяет передать знания об изображениях на более низкие уровни, а значит, облегчить и ускорить процесс обучения.

Таким образом, в этой работе используется важная особенность UNet, но без Upsampling и Downsampling операторов.

Результаты

В первую очередь авторы проводят эксперименты на нескольких датасетах с безусловной генерацией (CIFAR10. Celeba \( 64 \times 64 \)), а также class-conditional генерацию на ImageNet с различными разрешениями: \( 64 \times 64 \), \( 256 \times 256 \), \( 512 \times 512 \).

Помимо class-conditional генерации, авторы экспериментируют на text-to-image генерации:

Если посмотреть на метрики — нельзя сказать, что архитектура значительно улучшает генерацию. С предыдущей работой авторы никак не сравниваются, поэтому нельзя точно утверждать, что данная архитектура дает какое-то значительное улучшение по сравнению с предложенным DiT.

Diffusion Models as Masked Autoencoders

📄 Paper

Давайте перейдем к статье, где есть объединение двух основных идей: диффузионных моделей и ViTMAE. Теперь мы можем работать с пропуском данных в диффузионных моделях, что может помочь нам в различных задачах с “проблемными” данными.

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

На прямом процессе диффузии в данном случае зашумляется только маскированная часть.

Энкодер, аналогично ViTMAE, работает только с немаскированными патчами. Декодер, в свою очередь, пытается предсказывать зашумленные маскированные патчи на основе немаскированных. Авторы рассматривают три разных конфигурации такого декодера:

• Joint decoder

Конкатенируем последовательность выхода энкодера и зашумленные токены, а затем пропускаем через self-attention слой трансформера.

• Cross-self decoder

Имеем cross-attention блоки для обработки маскированных патчей и выход из энкодера, а затем self-attention блоки.

• Cross decoder

Работает почти так же, как и декодер выше, только без self-attention блоков: не добавляет связь между различными зашумленными токенами. Поскольку мы убираем довольно большую часть, эта архитектура становится более эффективной в вычислительном плане.

Авторы исследуют также влияние архитектуры декодера на качество inpainting и fine-tuned задачи классификации.

Можно заметить: есть заметный прирост качества cross-декодера на более шумных этапах.

Для задачи inpainting используется cross-self декодер, который показывает максимальное качество на метрике LPIPS.

Результаты

Такую модель можно использовать для различных задач, в частности для Inpainting, где вырезаны квадратные патчи.

А еще для генерации изображений, где маскировано 75% картинки:

и 90% картинки:

MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer

📄 Paper

📎 GitHub

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

Процесс обучения состоит из двух основных частей:

• latent masking operation — маскирование части изображения;
• ассиметричный диффузионный трансформер, который работает как обычная диффузионная модель, но на вход может принимать sparse данные.

Важная особенность архитектуры — Side-interpolater. Во время обучения энкодер работает с немаскированными токенами \( \hat{u} \), однако при инференсе ему приходится взаимодействовать со всеми токенами \( u \), ведь на этом этапе у нас нет маски. Таким образом, на выходе энкодера (входе декодера) у нас есть разная информация на этапе обучения и инференса. А с этим нужно что-то делать. Поэтому авторы предлагают небольшую сеточку, которая предскажет замаскированные токены на выходе энкодера в процессе обучения и исчезнет при инференсе.

На этапе инференса Side-interpolater заменяется position embedding оператором. Мы не будем здесь подробно останавливаться на архитектуре (но вы всегда можете обратиться к первоисточнику 🙂).

Такой подход помогает модели дополнительно изучать некоторую контекстную информацию из патчей и восстанавливать ее.

В итоге можно заметить, что обучение происходит быстрее, и алгоритм сходится лучше.

Еще авторы провели ряд экспериментов и проверили, какой уровень маскирования улучшает сходимость качественнее остальных:

DiffiT: Diffusion Vision Transformers for Image Generation

📄 Paper

📎 GitHub

Последняя работа, которую мы разберем — недавняя работа от NVIDIA. В ней авторы предлагают улучшенную архитектуру:

• новый self-attention модуль, который также зависит от времени;
• улучшенную transformer-based архитектуру;

Time-dependent Self-Attention

Трансформер на каждом шаге получает набор токенов изображения (латента) \( \{x_s\} \) и \( x_t \) — токен шага по времени. Авторы предлагают следующую cross-attention схему, где queries, keys и values зависят не только от токенов изображения, но и от токена времени:

где \( Q:=\{q_s\}, K:=\{k_s\}, V:=\{v_s\} \), d — scaling фактор.

DiffiT Block

Сам трансформерный блок, в свою очередь, состоит из двух частей:

где TMSA — Time-dependent Multi-head Self-Attention.

Итоговая архитектура — архитектура U-shaped encoder-decoder:

Она может использоваться как для пиксельного пространства, так и для латентного:

Для более подробного ознакомления с архитектурой рекомендуем также обратиться к первоисточнику.

Результаты

В итоге авторы получили SOTA результаты по генерации на нескольких различных датасетах. Они сравнивались в частности с архитектурами, которые мы упоминали выше (DiT и U-Vit) и обеспечили прирост на различных метриках.

Выводы

Итак, в нашей статье мы рассмотрели несколько работ, которые показывают эффективность использования трансформеров для диффузии. Это особенно хорошо заметно по последней из них.

Но все эти работы показывали эффективность на небольших разрешениях, чаще всего, на unconditional или class-conditional генерации, а все большие и часто используемые модели оставались на архитектуре UNet. Однако, как мы видим, крупные “игроки” OpenAI и Stability перешли к архитектуре на трансформерах и показали возможность очень хорошего качества. Кто знает, может, когда-нибудь мы посмотрим, что находится под капотом SORA (очень сомневаюсь) или Stable Diffusion 3 (очень надеюсь) 😊

Материалы

• Scalable Diffusion Models with Transformers

📄 Paper

📎 GitHub

• BEiT: BERT Pre-Training of Image Transformers

📄 Paper

📎 GitHub

• ViTMAE

📄 Paper

📎 GitHub

• Scalable Diffusion Models with Transformers

📄 Paper

📎 GitHub

• All are Worth Words: A ViT Backbone for Diffusion Models

📄 Paper

📎 GitHub

• Diffusion Models as Masked Autoencoders

📄 Paper

• MDTv2: Masked Diffusion Transformer is a Strong Image Synthesizer

📄 Paper

📎 GitHub

• DiffiT: Diffusion Vision Transformers for Image Generation

📄 Paper

📎 GitHub