Диффузия всё увереннее вытесняет GANы и авторегрессионные модели в ряде задач цифровой обработки изображений. Это не удивительно, ведь диффузия обучается проще, не требует сложного подбора гиперпараметров, min-max оптимизации и не страдает нестабильностью обучения. А главное, диффузионные модели демонстрируют state-of-the-art результаты почти на всех генеративных задачах — генерации картинок по тексту, генерация звуков, видео и даже 3D!

К сожалению, большинство работ в области text-to-something сосредоточены только на английском и китайском языках. Чтобы исправить эту несправедливость, мы решили создать мультиязычную text-to-image диффузионную модель Kandinsky 2.0, которая понимает запросы более чем на 100 языках! И главное, на русском ;) Подробности — под катом.

Модель Kandinsky 2.0 уже доступна на HuggingFace и GitHub, а попробовать генерацию можно на сайтах FusionBrain и rudalle.ru. Над этим проектом работали исследователи из SberAI и SberDevices при научной поддержке учёных из Института искусственного интеллекта AIRI.

Обязательно почитайте про наши предыдущие text-to-image модели, если вы не очень знакомы с темой синтеза изображений по текстовым описаниям:

• ruDALL-E

• Kandinsky

Что такое диффузия?

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

1. Сначала нужно итеративно испортить изображение, например, с помощью зашумления (прямой диффузионный процесс)

2. На втором этапе необходимо научиться итеративно восстанавливать изображение (обратный диффузионный процесс)

Впервые диффузионные модели были описаны в статье 2015 года Deep Unsupervised Learning using Nonequilibrium Thermodynamics — и, как видно из названия статьи, тут к её описанию пытались подойти через математический аппарат термодинамики. Отсюда взялось и название.

Диффу́зия — процесс перемешивания вещества, приводящий к выравниванию распределения.

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

Впервые превосходство диффузионных моделей над GAN'ами продемонстрировали в 2021 году в статье Diffusion Models Beat GANs on Image Synthesis. Авторы предложили и первый способ управления генерацией (conditioning), который назвали classifier guidance — использование градиентов от дополнительного классификатора для генерации объектов нужного класса (например, собак). Само же управление осуществляется через механизм Adaptive GroupNorm, то есть через предсказание коэффициентов нормализации.

Эту статью можно считать поворотным моментом, после которого многие переключились на исследование диффузии для различных задач генеративного ИИ. Каждые несколько недель начали появляться новые статьи о диффузии для superresolution, image inpainting, audio generation, text-to-video, text-to-3D и даже motion generation! Мы же сосредоточимся на задаче text-to-image — генерации картинок по текстовому описанию.

Text-to-image диффузия

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

Чтобы генерировать изображение на основе текста, нужно чтобы этот текст как-то учитывался во время генерации (text-conditioning). В статье OpenAI про модель GLIDE авторы предложили адаптировать метод classifier-free-guidance для текста. На этом подходе будут основаны все последующие text-to-image диффузионные модели, поэтому расскажем о нём чуть подробнее:

1. Текст кодируется в список эмбеддингов некоторым заранее определённым энкодером (в GLIDE — это был обучаемый трансформер).

2. Эмбеддинги проецируются на размерность слоя attention в UNet.

3. Затем они конкатенируются с контекстом UNet на каждом слое (управление через внимание).

4. CLS эмбеддинг подаётся на вход Adaptive Group Normalization (управление через нормировку).

5. Во время обучения текст на вход модели подаётся не всегда, а только с некоторой вероятностью

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

За счёт такого подхода удаётся управлять диффузией не только при помощи меток классов (кот, собака, машина и тд.), но и по полноценным текстам! Причём очень качественно.

В дальнейшем генерацию по тексту удалось значительно улучшить за счёт использования замороженных предобученных текстовых энкодеров и механизма каскадного повышения разрешения (Imagen). Оказалось, что совершенно незачем обучать текстовую часть text-to-image моделей — если взять замороженную T5-xxl, то качество картинок и понимание текста получается намного выше, а ресурсов для обучения требуется сильно меньше. Также в этой статье предложили улучшение сэмплирования — динамический трешхолдинг, который позволил повысить качество генерации при средних и высоких значениях guidance scale (варьируемый параметр, отвечающий за степень близости генерируемых сэмплов к тексту). Трешхолдинг — ограничение значений пикселей на каждом шаге диффузии, и если такое ограничение не использовать, то значения пикселей могут выйти далеко за пределы допустимого диапазона, что негативно скажется на генерации.

В статье про Latent Diffusion авторы показали, что на самом деле и картиночную часть не нужно обучать (по крайней мере не полностью). Если взять сильный автоэнкодер для картинок (VQ-VAE или KL-VAE) в качестве визуального декодера и пытаться диффузией генерировать эмбеддинги из его латентного пространства, а не само изображение, то обучение пойдёт ещё быстрее. Недавно выпущенная модель Stable Diffusion также основана на этом подходе.

Архитектура Kandinsky 2.0

В основе Kandinsky 2.0 лежит улучшенный подход Latent Diffusion (мы генерируем не изображения, а их латентные векторы) с несколькими существенными отличиями:

• Мы используем два мультилингвальных текстовых энкодера, эмбеддинги которых конкатенируются

• Больше UNet (1.2 млрд параметров)

• Динамический трешхолдинг в процессе сэмплирования

Чтобы модель стала по-настоящему мультиязычной, мы решили использовать мультилингвальный энкодер, причём сразу два: XLMR-clip и mT5-small. Таким образом, модель может понимать не только английский, русский, французский и немецкий, но даже монгольский, иврит и фарси! Всего 101 язык. Почему было решено использовать сразу две модели для кодирования текста? Потому что у этих моделей разные, но очень важные особенности — XLMR-clip видел картинки и даёт близкие эмбеддинги для разных языков, а mT5-small хорошо понимает сложные тексты. И наши предварительные тесты показали, что такая комбинация работает лучше всего, а так как обе модели имеют не очень много параметров (560M и 146M), то было решено использовать два энкодера сразу.

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

Как происходило обучение?

Обучение проходило на платформе ML Space c использованием мощностей суперкомпьютеров Christofari. Нам потребовалось 196 карт NVIDIA A100 с 80 Гб памяти на каждой из них. Всё обучение заняло 14 дней или 65 856 GPU-часов. Сначала 5 дней модель училась на разрешении 256х256, затем 6 дней на разрешении 512x512 и 3 дня на максимально чистых данных.

Для обучения использовались следующие гиперпараметры:

• Оптимизатор: AdamW с дефолтными параметрами

• OneCycleLR: max_iter=500k, max_lr=1e-4, pct_start=0.01

• Batch size: 48

• Linear noise scheduler: beta_start = 0.0001, beta_end = 0.02

• Diffusion steps: 1000

• Ограничение длины текстов: 77 токенов

• Разрешение: 256 -> 512

В качестве обучающих данных мы использовали комбинацию датасетов, которые мы предварительно отфильтровали по критериям наличия водяных знаков, низкого разрешения, низкого соответствия текстовому описанию по метрике CLIP-score:

• LAION-improved-aesthetics — 700 M

• LAION-aesthetics-multilang — 46 M

• ruDALLE-russian — 44 M

Дополнительные возможности модели

Для решения задач inpainting и outpainting мы обучаем дополнительную модель по аналогии с подходом из GLIDE: в качестве инициализации берём обученную Kandinsky 2.0 и добавляем к первой конволюции дополнительный канал для маски. После этого продолжаем обучение модели, но случайным образом закрашиваем часть каждой картинки и конкатенируем к ней соответствующую маску.

Мультиязычная генерация

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

Ниже приведены результаты генерации одного и того же запроса на разных языках. Например, для запроса на русском «человек с высшим образованием» в результатах генерации присутствуют исключительно белые мужчины, а для французского варианта “Photo d'une personne diplômée de l'enseignement supérieur” результаты более разнообразные. Хочется отметить, что только в русскоязычном варианте присутствуют грустные люди с высшим образованием.

А ещё с помощью этого подхода можно составить средний портрет «опасного человека» или «грабителя» для разных языковых культур и образ национальной кухни.

Заключение

Впереди у нас запланировано ещё много экспериментов с крупными языковыми моделями и другими подходами к процессу диффузии, но уже сейчас можно смело заявить, что Kandinsky 2.0 — это первая по-настоящему мультиязычная диффузионная модель! Генерацию и инпеинтинг, можно попробовать на сайтах FusionBrain и rudalle.ru, а также при помощи команды «Запусти художника» на умных устройствах Sber и в мобильном приложении Салют.

Авторы

Если вам интересно следить за новинками в DL, то обязательно подписывайтесь на каналы авторов Kandinsky 2.0 в Telegram:

• Градиентное погружение

• AbstractDL

Коллектив авторов: Арсений Шахматов, Антон Разжигаев, Владимир Архипкин, Александр Николич, Игорь Павлов, Андрей Кузнецов, Денис Димитров, Татьяна Шаврина, Сергей Марков.

Полезные ссылки

• Kandinsky 2.0 демо

• ruDALL-E демо

• Большая версия ruDALL-E, или Как отличить Кандинского от Малевича

• ruDALL-E: генерируем изображения по текстовому описанию, или Самый большой вычислительный проект в России

• Deep Unsupervised Learning using Nonequilibrium Thermodynamics

• Diffusion Models Beat GANs on Image Synthesis

• GLIDE: Towards Photorealistic Image Generation and Editing with Text-Guided Diffusion Models

• Classifier-Free-Guidance

• Imagen: Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding

• High-Resolution Image Synthesis with Latent Diffusion Models

• Stable Diffusion

• The Illustrated Stable Diffusion

• XLMR-clip

• mT5-small