Яндекс представил новую линейку диффузионных моделей YandexART 2.5

Привет, Хабр! Сегодня наша команда компьютерного зрения представляет две новые модели. Базовая модель YandexART 2.5 уже доступна всем пользователям Шедеврума и в Чате с Алисой. Более мощную Pro-версию можно попробовать в новой подписке Шедеврум Про. Модели научились самостоятельно подбирать подходящие стиль и оформление изображения и стали лучше генерировать надписи на латинице. Новые возможности будут полезны, например, для создания книжных обложек, логотипов или художественных иллюстраций.

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

Как удалось улучшить качество генераций YandexART

По результатам SbS-тестирования, обновлённая линейка превосходит Midjourney 6.1 по всем нашим критериям, а также выигрывает в сравнениях или находится в паритете с такими мировыми лидерами, как Ideogram, Dall-E 3 и Flux. 

Кратко подсвечу основные изменения, которые и привели к столь значимым для нас и пользователей результатам. 

Прежде всего, мы расширили обучающий датасет, приблизив его к 1 млрд пар картинок и описаний к ним. Описания к изображениям подготовила разработанная в Яндексе визуально-языковая модель (VLM). Она умеет распознавать даже незначительные объекты на изображении, за счёт чего её тексты получаются подробными и детализированными. Благодаря обучению на таких данных YandexART удаётся точнее следовать промпту. Кроме того, VLM впервые применили и для оценки качества работы нейросети. Визуально-лингвистическая модель проверяла множество изображений на соответствие заданным промптам. Например, действительно ли на картинке присутствовали все указанные пользователем элементы.

Также заменили автокодировщик VAE с 4-канального на 16-канальный. То есть увеличили в 4 раза размер латентного пространства, в котором работает диффузионная модель. Это повышает её разрешающую способность для работы с большим числом мелких деталей. 

Ещё применили «супирование». Грубо говоря, проводим fine-tuning модели несколько раз и усредняем веса получившихся моделей. Более подробно об этом планируем рассказать сообществу Хабра в отдельной статье.

Кроме того, в качестве метода дистилляции мы использовали вариацию multistep consistency distillation. Данный подход хорош тем, что сравнительно прост в реализации, стабилен в обучении и сохраняет разнообразие модели, в отличие от альтернативных подходов (адверсариальной и distribution matching-дистилляции). В YandexART 2.0 удалось достигнуть ускорения генерации в 4 раза по сравнению с исходной моделью. В YandexART 2.5 мы смогли ускорить генерацию уже в 9 раз! При этом ключевым фактором для сохранения релевантности и бездефектности было увеличение classifier free guidance при дистиллировании в меньшее количество шагов.

Ну и небольшой лайфхак в завершение. В процессе обучения модели YandexART 2.5 использовался bfloat16, сильно выросли активации и стали не влезать в диапазон значений fp16, из-за чего при инференсе в этом формате модель выдавала в большинстве случаев чёрные картинки. Мы обнаружили, что основной рост активаций происходит в FFN-слоях трансформерных блоков и простая обрезка активаций по абсолютной величине позволяет добиться успешной генерации в fp16 без потери качества.

Наиболее заметные улучшения в генерации картинок

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

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

Также модели генерируют надписи на латинице на 30% лучше в сравнении с предыдущей моделью.

К слову, мы с коллегами из направления компьютерного зрения Яндекса регулярно делимся личным опытом и обзорами полезных статей в ТГ-канале CV Time. Заглядывайте на огонёк.