Transfusion: Predict the Next Token and Diffuse Images with One Multi-Modal Model
[ни кода, ни весов]

Выше мы разобрали основные подходы к мультимодальной генерации и предположили, что работа с непрерывными представлениями картинок - наиболее перспективный подход. Авторы данной работы перешли от слов к делу и обучили непрерывный аналог Chameleon

Модель
Авторы обучают линейку трансформеров с LLaMA 2 архитектурой и токенизатором на совместное решение задачи next token prediction и диффузионного расшумления патчей картинок

Помимо основного трансформера нам потребуются:
➡️Непрерывный VAE для получения латентов патчей картинок
➡️Линейный слой или U-Net encoder/decoder для отображения латентов в эмбеды генеративного транформера и обратно

Вся система кроме VAE учится end-to-end на смеси картиночных и текстовых данных общим размеров 2Т токенов. Для выделения картинок используются <BOI>/<EOI> токены, которые переводят авторегрессионную модель в диффузионный режим. Помимо механизма семплирования, перевод режима влияет на аттеншен. Для текстов используется causal attn, для картинок - intra-image attn и, судя по результатам, это очень важная деталь

Результат
➡️Transfusion круче Chameleon во всех постановках, включая text-to-text. Интуиция - непрерывной модели проще даётся text-to-image и больше капасити модели остаётся, например, для text-to-text
➡️Transfusion скейлится круче Chameleon на диапазоне размера моделей от 0.16В до 7В параметров
➡️Использование inductive bias из U-Net помогает даже на таком масштабе. Для отображения в эмбеды стоит использовать его. Усложняет он не сильно (~3.8% от общего числа параметров)
➡️Модель легко дообучается на редактирование изображений используя всего 8к семплов

(Неочевидные) минусы
➖Очень много форвардов на генерацию изображения: нужно не только сгенерировать картинку по латентным патчам (~256 форвардов на 256х256 картинку), но еще и сделать N шагов диффузионного семплирования всей моделью на каждый патч
➖Систему учат с нуля на 2Т токенах и это достаточно дорого. Как будет работать в режиме файнтюна в статье не исследовали