[OpenAI DALL·E] Zero-Shot Text-to-Image Generation Aditya Rame | DeepLearning ru
[OpenAI DALL·E] Zero-Shot Text-to-Image Generation
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, Ilya Sutskever
Статья: https://arxiv.org/abs/2102.12092
Пост в блоге: https://openai.com/blog/dall-e/
Код (официальный, но неполный): https://github.com/openai/dall-e
Код (неофициальный, но более полный): https://github.com/lucidrains/DALLE-pytorch
Обученная модель от Сбера: https://github.com/sberbank-ai/ru-dalle
В последнее время появилось много мультимодальных трансформеров, и хочется их поразбирать. Нельзя обойти при этом наиболее известные модели, одна из которых — DALL·E от OpenAI.
DALL·E прошумел уже довольно давно, в самом начале года, но статья и все детали реализации стали доступны не сразу (в официальной репе так вообще лежит только обученный dVAE, который лишь часть от DALL·E).
Собственно, что такое DALL·E? Это декодер трансформера, который авторегрессионно умеет генерить изображения, “продолжая” заданное текстовое описание и, возможно, начальную часть изображения.
В оригинальном посте было заявлено, что это 12-миллиардная версия GPT-3, обученная на парах картинок и их описаний, но реальность сильно сложнее.
Начать стоит с того, что процесс обучения двухэтапный. На первом этапе обучается дискретный VAE (dVAE), сжимающий входную картинку 256*256 в картиночные токены на сетке 32*32, принимающие 8192 возможных значения каждый (то есть размер словаря). Это нужно для того, чтобы уменьшить количество токенов, с которыми должен дальше оперировать трансформер, и даёт уменьшение размера контекста в 192 раза. Часть высокочастотного сигнала в изображении при этом, ожидаемо, теряется, мелкие детали становятся плохо- или неразличимы, но в целом качество восстановления картинки после такого dVAE вполне достойно. Именно отсюда получается та самая своеобразная гладкость генеримых DALL·E картинок и проблемы с мелкими деталями, соответственно не для любого типа изображений DALL·E подходит, либо же процесс требует переобучения dVAE на более специальные кейсы. В dVAE никаких трансформеров нет, это свёрточный резнет.
Второй этап — это собственно трансформер. 256 кодирующих текст BPE-токенов (словарь размера 16384) конкатенируются с 1024 картиночными токенами, полученными от dVAE, и эта последовательность авторегрессионно продолжается. Сам трансформер это 64-слойный sparse transformer от тех же OpenAI (https://arxiv.org/abs/1904.10509). Он хитрый, в нём три типа шаблонов внимания: 1) text-to-text классический masked механизм как в авторегрессионных языковых моделях по типу GPT, где текущий токен не имеет права заглядывать в будущие токены; 2) image-to-text, где каждый токен изображения смотрит на все токены текста; и 3) image-to-image, где используются паттерны из sparse transformer с вниманием по строкам, столбцам и более хитрыми свёрточными шаблонами, которые применяются только в последнем self-attention слое. При этом все три типа заведены в одну единственную операцию внимания, потому что это работало лучше, чем три отдельные операции.
Авторы сначала экспериментировали на маленькой модели в 1.2B параметров и датасете в 3.3M пар картинка-текст, а потом отскейлили это на 12B модель и датасет в 250M пар (который пришлось собрать). Собственно этот скейлинг и составляет главные идейную и техническую часть работы. Идейно — показать, что обученная на большом датасете модель большого размера демонстрирует интересное поведение, а технически — заставить всё это работать, потому что вылезает много инженерных челленджей, не проявляющихся на меньших масштабах.
По признанию авторов, самая сложная часть проекта была в том, чтобы обучить большую модель на 16-битных числах (вместо обычных 32-битных). Было много всяких нестабильностей, поэтому авторы придумали хитрые методы масштабирования градиентов для защиты от underflow. Другой челлендж — распределённое обучение, где использовался тонко настроенный вариант PowerSGD. За этими (на самом деле важными, если хотите повторить) техническими деталями обращайтесь в Appendix.
