Как обучать диффузионные модели: разбираемся на примере YandexART

НачалоСбор обучающих датасетовФильтрация изображенийФильтрация текстовРелевантность картинок и текстовОценка генеративной модели

Text-to-Image диффузионные модели умеют генерировать изображения по текстовому описанию и создавать совершенно новые концепты: например, картинку с астронавтами, которые скачут на лягушках по просторам Марса.

Руководитель группы нейросетевых технологий компьютерного зрения в Яндекс Поиске, Сергей Овчаренко, рассказывает о процессе обучения YandexART — модели, которая работает в приложении Шедеврум и других продуктах Яндекса.

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

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

Мы используем различные открытые источники данных для обучения модели, в том числе материалы, которые есть в Яндекс Картинках. Проблема в том, что из-за большого объёма данные Яндекс Картинок очень шумные, их нужно тщательно фильтровать.

У нас есть задача: из сотен миллиардов изображений отобрать классный красивый топ. Начинаем очищать изображения:

→ Фильтруем по размеру, убираем маленькие картинки
→ Исключаем дубликаты при помощи технологии, разработанной командой компьютерного зрения Яндекса
→ Убираем нежелательный контент, чтобы наша модель умела генерировать: с водяными знаками, запрещённый и категории 18+

Размер датасета уже сократился, но остаётся ещё достаточно много изображений. На этом этапе можно обучить модель, которая будет отбирать лучшие по эстетике и содержанию картинки. Потом использовать Fine-Grained IQA классификаторы, чтобы из достаточно хороших картинок отобрать наиболее визуально красивые.

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

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

Выбор подходящих текстов тоже проходит в несколько стадий:

→ Оставляем только интересующие нас языки
→ Проводим базовые фильтрации: используем эвристические фильтры, сортируем тексты по длине, энтропии; формируем блэк-листы, чтобы отбросить нежелательные слова

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

Остаются сложные кейсы, когда тексты уже хорошие, но мы уверены, что можно лучше. Подход, который позволяет получать интересные результаты, — использовать большие языковые модели. Например, можно попросить YandexGPT «причесать» тексты: перегенерировать, убрать дубли. Либо использовать LLM как дискриминативную модель, чтобы отсеять плохие тексты.

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

Наша задача — получить не слишком большие и не слишком маленькие согласованные тексты, которые дают возможность представить какую-то картинку.

Теперь у нас есть есть два множества: какие-то хорошие тексты, какие-то красивые картинки. Задача — объединить их и понять, действительно ли они связаны между собой.
Измерить семантическую близость между картинками и текстом помогает наша Image-to-text модель в сочетании с контрастивной картиночно-текстовой моделью CLIP ViT-bigG.

С этими моделями есть проблема: тексты вроде wallpaper или «красивая картинка» с точки зрения модели подойдут большинству картинок из интернета, потому что модели учились на этих данных. В результате в фильтрации будет большое количество общих текстов. В этих условиях определить глобальный порог близости, по которому мы будем фильтровать пары, очень сложно.

Мы придумали, как избавиться от глобального порога на релевантность и сравнивать текстовое и картиночное представление. Это работает так:

→ Взять примерно 100 тысяч случайных текстов и сравнить с картинкой
→ Отсортировать тексты по схожести с картинкой на три категории: совсем непохожие, общие и те, что хорошо описывают изображение

На этом этапе мы в целом были довольны результатами, но решили применить ещё data-driven подход:

→ Разметили семплы на три градации: хорошие, нормальные и плохие
→ Обучили классификатор, чтобы дофильтровать датасет до хорошего качества

Справа — идеальный семпл, который мы хотим видеть в своём датасете: красивая картинка и информативное описание

На такой разметке можно обучить классификатор и дофильтровать датасет до хорошего качества.

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

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

Все замеры мы делаем при помощи асессоров. Они смотрят на пары изображений и выбирают лучшее по критериям:

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

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

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

Поделитесь увиденным

Скопировать ссылку
ТелеграмВКонтакте