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

2021 год в машинном обучении ознаменовался мультимодальностью — активно развиваются нейросети, работающие одновременно с изображениями, текстами, речью, музыкой. Правит балом, как обычно, OpenAI, но, несмотря на слово «open» в своём названии, не спешит выкладывать модели в открытый доступ. В начале года компания представила нейросеть DALL-E, генерирующую любые изображения размером 256?256 пикселей по текстовому описанию. В качестве опорного материала для сообщества были доступны статья на arxiv и примеры в блоге.

С момента выхода DALL-E к проблеме активно подключились китайские исследователи: открытый код нейросети CogView позволяет решить ту же проблему — получать изображения из текстов. Но что в России? Разобрать, понять, обучить — уже, можно сказать, наш инженерный девиз. Мы нырнули с головой в новый проект и сегодня рассказываем, как создали с нуля полный пайплайн для генерации изображений по описаниям на русском языке.

В проекте активно участвовали команды SberAI, SberDevices, Самарского университета, AIRI и SberCloud.

Мы обучили две версии модели разного размера и дали им имена великих российских абстракционистов — Василия Кандинского и Казимира Малевича:

1. ruDALL-E Kandinsky (XXL) с 12 миллиардами параметров;

2. ruDALL-E Malevich (XL), содержащая 1,3 миллиарда параметров.

Некоторые версии наших моделей доступны в open source уже сейчас:

1. ruDALL-E Malevich (XL) [GitHub, HuggingFace]

2. Sber VQ-GAN [GitHub, HuggingFace]

3. ruCLIP Small [GitHub, HuggingFace]

4. Super Resolution (Real ESRGAN) [GitHub, HuggingFace]

Две последние модели встроены в пайплайн генерации изображений по тексту (об этом расскажем ниже).

Версии моделей ruDALL-E Malevich (XL), ruDALL-E Kandinsky (XXL), ruCLIP Small, ruCLIP Large, Super Resolution (Real ESRGAN) также скоро будут доступны в DataHub.

Обучение нейросети ruDALL-E на кластере Christofari стало самой большой вычислительной задачей в России: 

1. Модель ruDALL-E Kandinsky (XXL) обучалась 37 дней на 512 GPU TESLA V100, а затем ещё 11 дней на 128 GPU TESLA V100 — всего 20 352 GPU-дней;

2. Модель ruDALL-E Malevich (XL) обучалась 8 дней на 128 GPU TESLA V100, а затем еще 15 дней на 192  GPU TESLA V100 — всего 3 904 GPU-дня.

Таким образом, суммарно обучение обеих моделей заняло 24 256 GPU-дней.

Разберём возможности наших генеративных моделей.

«Озеро в горах, а рядом красивый олень пьёт воду» — генерация ruDALL-E Malevich (XL)

Почему Big Tech изучает генерацию изображений 

Долгосрочная цель нового направления — создание «мультимодальных» нейронных сетей, которые изучают концепции в нескольких модальностях, в первую очередь в текстовой и визуальной областях, чтобы лучше понимать мир. 

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

1. Возможность точно описать желаемое — и получить персонализированное изображение, которое раньше не существовало.

2. В любой момент создавать необходимое количество licence-free иллюстраций в неограниченном объеме.

Первые очевидные применения генерации изображений:

• Фото-иллюстрации для статей, копирайтинга, рекламы. Можно автоматически (а значит — быстрее и дешевле) создавать иллюстрации к статьям, генерировать концепты для рекламы по описанию.

  Примеры: 

«Лиса в лесу»

«Орел сидит на дереве, вид сбоку»

«Автомобиль на дороге среди красивых гор»

• Векторные иллюстрации, свободные от лицензии фотостоков, тоже можно генерировать бесконечно:

  «Векторная иллюстрация с розовыми цветами»

• Визуализации дизайна интерьеров — можно проверять свои идей для ремонта, играться с цветовыми решениями, формами и светом:

  «Шикарная гостиная с зелеными креслами», прам как с сайта мебели

  «Современное кресло фиолетового цвета»

• Visual ART — источник визуальных концепций, соединений различных признаков и абстракций:

  «Темная энергия»

«Кот на Луне»

«Кошка которая сделана из белого облака»

«Енот с пушкой»

«Красивое озеро на закате»

Архитектура модели

В основе архитектуры DALL-E — так называемый трансформер, он состоит из энкодера и декодера. Общая идея состоит в том, чтобы вычислить embedding по входным данным с помощью энкодера, а затем с учетом известного выхода декодировать вектор для сопоставления с ожидаемым выходом. В трансформере энкодер и декодер состоят из ряда идентичных блоков, которые основываются на механизме Attention.

Механизм Self-attention — основа архитектуры трансформера. Он позволяет модели понять, какие фрагменты входных данных важны и насколько важен каждый фрагмент входных данных для других фрагментов. Концептуально слой Self-attention позволяет обновлять элементы входной последовательности за счёт включения глобальной информации обо всех элементах последовательности. Как и LSTM-модели, трансформер позволяет естественным образом моделировать связи «в долгую». Однако, в отличие от LSTM-моделей, он подходит для распараллеливания и, следовательно, эффективных реализаций.

Первым шагом при вычислении Self-attention является создание трёх векторов для каждого входного вектора энкодера (для каждого элемента входной последовательности). То есть для каждого элемента мы создаем векторы Query, Key и Value. Эти векторы получаются путем перемножения embedding’а и трех матриц, которые мы получаем в процессе обучения. Далее мы используем полученные векторы для формирования Self-attention-представления каждого embedding’а, что дает возможность оценить возможные связи в элементах входных данных, а также определить степень «полезности» каждого элемента.

Трансформер также характеризует наличие словаря. Каждый элемент словаря — это токен. В зависимости от модели размер словаря может меняться. Таким образом, входные данные сначала превращаются в последовательность токенов, которая далее конвертируется в embedding с помощью энкодера. Для текста используется свой токенизатор, для изображения сначала вычисляются low-level-фичи, а затем в скользящем окне вычисляются визуальные токены. Применение механизма Attention позволяет извлечь контекст из входной последовательности токенов в ходе обучения. Следует отметить, что для обучения трансформера требуются большие объёмы (желательно «чистых») данных, о которых мы расскажем ниже.

Как устроен ruDALL-E

Цель архитектуры — обучение трансформера авторегрессивно моделировать токены текста и изображения как единый поток данных. Однако использование пикселей непосредственно в качестве признаков изображений потребует чрезмерного количества памяти, особенно для изображений с высоким разрешением. Чтобы не учить только краткосрочные зависимости между пикселями и текстами, а делать это более высокоуровнево, обучение модели проходит в 2 этапа:

1. Предварительно сжатые изображения с разрешением 256х256 поступают на вход автоэнкодера (мы обучили свой SBER VQ-GAN, улучшив метрики для генерации по некоторым доменам, и об этом как раз рассказывали тут, причем также поделились кодом), который учится сжимать изображение в матрицу токенов 32х32. Фактор сжатия 8 позволяет восстанавливать изображение с небольшой потерей качества: см. котика ниже.

2. Трансформер учится сопоставлять токены текста (у ruDALL-E их 128) и 32?32=1024 токена изображения (токены конкатенируются построчно в последовательность). Для токенизации текстов использовался токенизатор YTTM.

   Исходный и восстановленный котик

Трудности, возникшие при обучении и генерации:

1. На данный момент в открытом доступе нет кода модели DALL-E от OpenAI. Публикация описывает её общими словами, но обходит вниманием некоторые важные нюансы реализации. Мы взяли наш собственный код для обучения ruGPT-моделей и, опираясь на описание, а также попытки воспроизведения opensource-сообщества и китайских исследователей, написали свой код DALL-E-модели. Он включает такие детали, как позиционное кодирование блоков картинки, свёрточные и координатные маски Attention-слоёв, общее представление эмбеддингов текста и картинок, взвешенные лоссы для текста и картиночной части, dropout-токенизатора.

2. Из-за огромных вычислительных требований эффективно обучать модель можно только в режиме точности fp16. Это в 5-7 раз быстрее, чем обучение в классическом fp32, и она занимает меньше места. Но ограничение точности представления чисел повлекло за собой множество сложностей для такой глубокой архитектуры:

   a) иногда встречающиеся очень большие значения внутри сети приводят к вырождению лосса в Nan и прекращению обучения;

   b) при малых значениях learning rate, помогающих избежать проблемы а), сеть перестает улучшаться и расходится из-за большого числа нулей в градиентах.

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

   Для распределенного обучения на нескольких DGX мы используем DeepSpeed, как и в случае с ruGPT-3.

3.  Сбор данных и их фильтрация: безусловно, когда мы говорим об архитектуре, нововведениях и других технических тонкостях, нельзя не упомянуть такой важный аспект как данные. Как известно, для обучения трансформеров их должно быть много, причем «чистых». Под «чистотой» мы понимали в первую очередь хорошие описания, которые потом нам придётся переводить на русский язык, и изображения с отношением сторон не хуже 1:2 или 2:1, чтобы при кропах не потерять содержательный контент изображений.

   Первым делом мы взялись за те данные, которые использовали OpenAI (в статье указаны 250 млн. пар) и создатели CogView (30 млн пар): Conceptual Captions, YFCC100m, данные Википедии, ImageNet. Затем мы добавили датасеты OpenImages, LAION-400m, WIT, Web2M и HowTo как источник данных о деятельности людей, и другие, которые покрывали бы интересующие нас домены. Ключевыми доменами стали люди, животные, знаменитости, интерьеры, достопримечательности и пейзажи, различные виды техники, деятельность людей, эмоции.

   После сбора и фильтрации данных от слишком коротких описаний, маленьких изображений и изображений с непригодным отношением сторон, а также изображений, слабо соответствующих описаниям (мы использовали для этого англоязычную модель CLIP), мы выполнили перевод всех описаний на русский язык с помощью ruGPT3-модели и в итоге сформировали широкий спектр данных для обучения — более 150 млн. пар. 

   Кривая обучения ruDALL-E XXL: как видно, обучение несколько раз приходилось возобновлять после ошибок и уходов в Nan.

   

4. Сложность выбора оптимальных режимов генерации для разных объектов и доменов. В ходе исследования генерации объектов мы начали с доказавших свою полезность в NLP-задачах подходов Nucleus Sampling и Top-K sampling, которые ограничивают пространство токенов, доступных для генерации. Эта тема хорошо исследована в применении к задачам создания текстов, но для изображений общепринятые настройки генерации оказались не самыми удачными. Серия экспериментов помогла нам определить приемлемые диапазоны параметров, но также указала на то, что для разных типов желаемых объектов эти диапазоны могут очень существенно отличаться. И неправильный их выбор может привести к существенной деградации качества получившегося изображения. Вопрос автоматического выбора диапазона параметров по теме генерации остаётся предметом будущих исследований.

Вот не совсем удачные генерации объектов на примере котиков, сгенерированные по запросу «котик с красной лентой»: 

Картинка 1 — у кота 3 уха; второй не вышел формой; третий немного не в фокусе. 

А вот «Автомобиль на дороге среди красивых гор». Автомобиль слева въехал в какую-то трубу, а справа — странноватой формы.

«Автомобиль на дороге среди красивых гор»

Пайплайн генерации изображений

Сейчас генерация изображений представляет из себя пайплайн из 3 частей: генерация при помощи ruDALL-E — ранжирование результатов с помощью ruCLIP — и увеличение качества и разрешения картинок с помощью SuperResolution.

При этом на этапе генерации и ранжирования можно менять различные параметры, влияющие на количество генерируемых примеров, их отбор и абстрактность. 

Пайплайн генерации изображений по тексту

GitHub + Colab

В Colab можно запускать инференс модели ruDALL-E Malevich (XL) с полным пайплайном: генерацией изображений, их автоматическим ранжированием и увеличением:

Рассмотрим его на примере с оленями выше:

Шаг 1. Установка и импорт библиотек

git clone https://github.com/sberbank-ai/ru-dalle

pip install -r ru-dalle/requirements.txt > /dev/null

from rudalle import get_rudalle_model, get_tokenizer, get_vae, get_realesrgan, get_ruclip

from rudalle.pipelines import generate_images, show, super_resolution, cherry_pick_by_clip

from rudalle.utils import seed_everything
seed_everything(42)

device = 'cuda'

Шаг 2. Генерация вариантов по тексту

text = 'озеро в горах, а рядом красивый олень пьет воду'

tokenizer = get_tokenizer()
dalle = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
vae = get_vae().to(device)

pil_images, _ = generate_images(text, tokenizer, dalle, vae, top_k=1024, top_p=0.99, images_num=24)

show(pil_images, 24)

Результат:

Генерация изображений по тексту

Шаг 3. Автоматическое ранжирование изображений, выбор лучших

ruclip, ruclip_processor = get_ruclip('ruclip-vit-base-patch32-v5') ruclip = ruclip.to(device)

top_images, _ = cherry_pick_by_clip(pil_images, text, ruclip, ruclip_processor, device=device, count=24)

show(top_images, 6)

Результат ранжирование ruCLIP-ом (топ6)

Можно заметить, что один из оленей получился достаточно «улиточным». На этапе генерации можно делать перебор гиперпараметров для получения наиболее удачного результата именно под ваш домен. Опытным путем мы установили, что параметры top_p и top_k контролируют степень абстрактности изображения. Их общие рекомендуемые значения: top_k=2048, top_p=0.995; top_k=1536, top_p=0.99; top_k=1024, top_p=0.99.

Шаг 4. Делаем SuperResolution

realesrgan = get_realesrgan('x4', device=device)
sr_images = super_resolution(top_images, realesrgan)
show(sr_images, 6)

Super Resolution версии генерации

Для запуска пайплайна с моделью ruDALL-E Kandinsky (XXL) или Malevich (XL) можно будет также использовать каталог моделей DataHub (ML Space Christofari).

Будущее мультимодальных моделей

Мультимодальные исследования становятся всё более популярны для самых разных задач: прежде всего, это CV + NLP (о первой такой модели для русского языка, ruCLIP, мы рассказали ранее), а также NLP + программный код, и даже много источников данных одновременно. 

В конкурсе Fusion Brain Challenge конференции AI Journey единой моделью нужно решить 4 задачи:

• С2С — перевод с Java на Python; 

• HTR — распознавание рукописного текста на фотографиях;

• Zero-shot Object Detection — детекция на изображениях объектов, заданных на естественном языке; 

• VQA — ответы на вопросы по картинкам. 

По условиям соревнования (которое продлится до 5 ноября) на общие веса нейросети должно приходиться как минимум 25% параметров! Совместное использование весов для разных задач делает модели более экономичными в сравнении с их мономодальными аналогами. Организаторами также был предоставлен бейзлайн решения, который можно найти на официальном GitHub. 

И пока команды соревнуются за первые места, а компании наращивают вычислительные мощности для обучения закрытых моделей, нашим интересом остается open source и большое сообщество. Будем рады вашим прототипам, неожиданным находкам, тестам и предложениям по улучшению моделей!

• Demo и галерея лучших изображений

• Github

• Telegram bot: @sber_rudalle_xl_bot

Коллектив авторов: @rybolos, @ollmer, @shonenkov, @kuznetsoff87, @oulenspeigel, @achertok, @da0c, @mboyarkin, @boomb0om

Источник: habr.com

Комментарии: