🔬 Метод

Напомним, что RAE кодирует одновременно высокочастотную информацию, необходимую для реконструкции и сематническую. RAE инициализуется неким семантическим энкодером (SigLIP/WebSSL) и дообучается на реконструкцию + дополнительные лоссы.

В данной работе рассматривают в качестве базовых тушек SigLIP2/WebSSL и обучают RAE на реконструкцию + LPIPS + Грам лосс + адверсариальный лосс. Удается добиться качества реконструкции между SDXL VAE и FLUX VAE. Для реконструкции текстов (внезапно) важно долить текстовых данных.

Используют MetaQuery в архитектуре - последовательность обучаемых токенов той же длины, что и выход RAE (256 токенов для изображений 224x224).

Для оригинального RAE для лучших результатов было важно зашумлять латенты и делать широкую DDТ голову. Для T2I оказывается, что данные изменения не дают особого прироста если учить достаточно долго достаточно большие модели.

А вот сдвиг расписания (flow shift) остается довольно важным.

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

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

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

Основная масса экспериментов проводится на 1.5B Квене в качестве текстового энкодера и 2.4B диффузионной модели.

RAE сходится значительно быстрее к тем же метрикам (в 4 раза по GenEval, в 4.6 раз по DPG Bench) и выходят на лучшее качество при заданном бюджете.

Дальнейший скейлинг модели по размеру не приводит к значительному улучшению качества (вероятно, потому что упираются в данные).

Интересной фичой подхода с RAE является способность оценивать свои генерации, тем самым делая возможным test-time scaling по типу Best-of-N.

💡 Выводы

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