🔬 Метод

🏢 Архитектура

С архитектурной точки зрения ничего фундаментально нового:

📌 Qwen-2.5VL (7B) в качестве энкодера. Данный выбор позволяет из коробки подавать одновременно текст и изображения, что необходимо для задач редактирования и персонализации. Энкодер заморожен.
📌 В качестве VAE за основу берут энкодер+декодер от коллег из Wan. Однако ее качество не полностью удовлетворяет авторов, поэтому они дообучают декодер при замороженном энкодере. Причем используют только взвешенную сумму MSE и perceptual loss. Адверсариальный лосс якобы скорее мешает. В данные для обучения заливают специально побольше текстов.
📌 Диффузионный трансформер - это просто MMDiT 🤖 а-ля SD3/FLUX.

Из примечательного - интересное решение по учету позиций картиночных и текстовых эмбедов под названием Multimodal Scalable RoPE (MSRoPE). В наивном подходе, каждому токену задавался бы просто позиционный индекс (сверху вниз, слева направо для патчей картинок и токену согласно позиции в тексте). Однако при таком подходе одинаковым образом бы кодировалась большая картинка, и меньшая картинка + текст. Ребята из Квена же предлагают все сэмплы выстраивать в диагональ. И текстовые токены это 1x1 элемент на диагонали.

📊 Данные

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

Следуя довольно стандартной практике, сначала обучают на низком разрешении (256px), а затем постепенно повышают до 640px-1328px.

Уже довольно на ранних стадиях отфильтровывают данные с чрезмерно низкой и высокой насыщенностью или контрастностью.

Для captioning используют Qwen-2.5VL (72B?), который выдает не только описание, но и другие метаданные - стиль, тип, водные знаки в виде некой json-ины.

Чтобы улучшить качество генерации текста в датасет заливают большое количество текстов 📝, в том числе и синтетически сгенерированных (некий шаблон вставляется на фон отпимальным образом, и если не шакален - принимается). Кроме того, на основе презентаций пробуют учить на разные сложные layout.

🧠 Процедура обучения

Диффузионный процесс - уже общепринятый flow-matching без танцев с бубнами.

Анализируя эффективность обучения, авторы заметили, что чекпоинтинг активаций сильно замедляет обучение, поэтому отключили его, и использовали только распределенные оптимизаторы. Вместо FSDP используют tensor-параллелизм поверх transformer engine.

Дабы улучшить общее качество и алайнмент, модельку прогоняют через DPO на ручной человеческой разметке и GRPO. Flow-matching ODE - детерминированный процесс, и дабы иметь разнообразие, на этапе GRPO используется стохастическое сэмплирование.