🔬 Метод

Чтобы моделировать на уровне пикселей, предлагают двухуровневый каскад из patch-level сети, обрабатывающей патчи, и pixel-level сети, работающей на пикселях.

Patch-level сеть - это просто трансформер обрабатывающие патчи (в основном 16x16), как токены. Она принимает на вход либо метку класса, либо текст.

Более примечательна pixel-level сеть. Она использует pixel-level модуляцию (домножение на скаляр и сдвиг) на основе выхода patch-level сети. Чтобы параметры модуляции были специфичны для каждого пикселя через линейный слой карту признаков патчей разворачивают в p x p (p - размер патча).

Полный аттеншен между всеми пикселями будет слишком дорогим, поэтому перед вниманием токены группируются снова в патчи (p x p), и затем снова разворачиваются в исходную последовательность.

Обычно берут pixel-level сетку неглубокой и неширокой (2-4 слоя) со скрытой размерностью 16, поэтому гонять ее даже на высокоразмерных изображениях недорого.

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

Обучают как class-condition, так и text-2-image модель. Чтобы улучшить метрики и сходимость используют REPA с признаками DINOv2.

Модель выбивает неплохие FID-ы, опережая все прошлые подходы на пиксельной диффузии, но недотягивая несколько до SOTA FID-арасов.

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

Уменьшение патчей не сильно улучшает качество, особенно для больших моделей, при этом кратно делая процедуру дороже, поэтому патчи 16x16 выбирают как более-менее оптимальные.

Затем обучают T2I MMDiT-модель с Gemma-2 энкодером размером 1.3B параметров и она выдает довольно неплохие метрики. Обучают на 26М пар, не очень долго, поэтому ожидать SOTA не приходится, но для таких ресурсов вполне достойное качество.

💡 Выводы

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