End-to-End Training for Unified Tokenization and Latent Denoising
[код и веса]

Современные LDM почти всегда двухступенчатые: сначала отдельно учат токенизатор (автоэнкодер), потом замораживают его и сверху учат генератор в зафиксированном латентном пространстве. Это удобно инженерно, но генеративная часть никак не влияет на то, каким получается латентный код и приходится подстраивать генератор под чужое пространство признаков.​

UNITE предлагает рассматривать токенизацию и генерацию как одну и ту же задачу вычисления латентов, просто в разных режимах: в одном случае модель видит исходное изображение почти целиком и должна выдать его компактное представление, в другом — стартует с шума и слабых подсказок и должна прийти к тому же типу представления. Важный шаг — один и тот же энкодер используется и как токенизатор, и как генератор в латентном пространстве, а обучающие сигналы от задачи реконструкции и от задачи генерации совместно формуют одно общее пространство признаков.​

За счёт этого авторы показывают, что можно обойтись одним обучающим запуском и одной моделью вместо каскада токенизатор + генератор, сохранив почти SOTA качество как по реконструкции, так и по генерации. Анализ промежуточных представлений показывает, что даже без внешних учителей токенизация и генерация естественно выравниваются: внутренние слои модели переиспользуются обеими задачами, а не расходятся в две несвязанные подсистемы.​

Интересно, что совместное обучение ведёт себя немного адверсариально: критерии реконструкции и генерации местами конфликтуют, но именно этот конфликт заставляет модель искать латентное пространство, которое одновременно и достаточно информативно для восстановления деталей, и достаточно устойчиво, чтобы по нему было удобно генерировать.

В итоге UNITE показывает, что можно тренировать end-to-end LDM с нуля, не разбивая на два этапа и не привлекая DINO-учителей. Это открывает путь к обучению генеративок в узкоспециализированных областях (робототехника, биология), где нет готовых VAE.