🔬 Метод

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

За основу берут DiT из PixArt-α, но с рядом особенностей и модификаций:
📌 В качестве бейзлайна берут модель на 400М с MultiQueryAttention и уменьшенными FFN-ами. (не понял зачем MQA, если не используются KV-кэши)
Архитектуру DiT-a превращают в песочные часы. После нескольких блоков делают downsampling x2 по каждой из осей, и гоняют карты признаков меньшего размера, а затем ближе к концу снова повышают разрешение добавляя skip-connections как в UNet.
📌 Чтобы удешевить attention на высоком разрешении, делают следующее:
1) Сверткой сжимают ключи и значения
2) Attention не глобальный, а блочный, друг на друга внимаются только соседние токены в пределах блока.
📌Затем еще добавляют skip-connections в middle-block (после downsampling), увеличивают ffn на высоком разрешении и подбирают количество блоков до сжатия и после разжатия.

Архитектурный ablation проводят на ImageNet на основе валидационного лосса, а затем переносят на text-2-image модель.

Кроме того, учат не одну модель за раз, а суперсеть 🕸, включащую в себя подсети для работы на разных устройствах:
👶 Самую маленькую на 0.3B параметров, чтобы гонять на не high-end андроидах.
👦 Чуть побольше на 0.4B, для более high-end.
💪 И самую большую и сильную на 1.6B, для быстрого инференса на GPU.

На обучении все подсети сэмплируются из подсети и обучаются на flow-matching loss. Подсети еще дистиллируют на всей суперсети для лучшей сходимости и качества.

Кроме того, еще добавляют дистилляцию на выходы учителя (Qwen-Image), причем дистиллируют не только скорость, но и последнее скрытое состояние.

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

🧪 Эксперименты

Учат подсеть на 1.6B параметров с двумя подсетями. В качестве текстовых энкодеров берут TinyClip и Gemma3-4b-it.

Сначала преобучают на разрешении 256, а потом переходят на 1024.
0.4B моделька при инференсе в 4 шага отрабатывает на 16-ом айфоне на 1.7c.

По user study меньшая модель чуть лучше SANA-1.6M и чуть хуже SD3M, а большая на 1.6B параметров как FLUX-1-dev.

На автобенчах (DPG/GenEval) тоже неплохо.

💡 Выводы

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