Z-Image: An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer [1/3]

Сегодня начнём разбирать большую статью о новой генеративной модели Z-Image. В первом посте серии поговорим об инфраструктуре для обработки данных, во втором — об ошибках модели, в третьем — о её архитектуре.

Z-Image знаменита двумя вещами. Она:

- небольшая по современным меркам — около 6B параметров;
- весьма эффективно обучена: для тренировки использовали всего 314К GPU-часов на H800 GPU.

Пока авторы зарелизили только text-to-image-модель, но скоро обещают выкатить editing. Текущая версия уже впечатляет: иногда она ошибается в генерации текстов на изображении, но в среднем отлично справляется с этими и другими визуальными эффектами.

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

1. Data Profiling Engine. Базовая фильтрация по качеству изображений — отсеиваются картинки со слишком высоким уровнем шума или чрезмерным объёмом фона, с артефактами компрессии и т. д.

Эстетичность картинки оценивает отдельная модель: она создаёт дополнительный атрибут aesthetic score.

Что касается caption:

- Пары «изображение + текст» с низкой похожестью фильтрует специально дообученный CLIP. Он возвращает коэффициент того, насколько картинка соответствует тексту (по шкале от 0 до 1). Таким образом датасет очищается от самых не подходящих друг другу пар.

- Чтобы обогатить caption, генерируют три его версии: очень подробную, детальную и короткую. Все три получают с помощью простой VLM. Далее их используют для прогрессивного обучения: постепенно повышают разрешение изображений и длину caption, чтобы модель познакомилась с миром изображений, поняла, как их генерировать, и усовершенствовала свои навыки.

2. Cross-modal Vector Engine. Нестандартный механизм, который проверяет, насколько хорошо датасет покрывает всё множество реальных пар «картинка + текст».

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

3. World Knowledge Topological Graph. Для того чтобы оценить, насколько концепты реального мира представлены в датасете, используют необычный механизм. Строят граф: вершины — концепты, рёбра — их ссылки друг на друга. От маловажных концептов избавляются алгоритмом PageRank. Оставшимся точкам присваивают теги и группируют их в иерархические деревья. Такой граф позволяет балансировать датасет во время обучения, повышая вероятность у семплов, представляющих слаборепрезентованные концепты.

4. Active Curation Engine. Люди и VLM поочерёдно проверяют датасет, пока он не окажется достаточно хорошим.

Подробнее об ошибках Z-Image и том, как она устроена, расскажем в следующих постах.

Разбор подготовил ❣ Илларион Дмитриев
CV Time