REPA-E: Раскрытие потенциала VAE для сквозной настройки с использованием латентной диффузии и трансформеров

REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers

April 14, 2025

Авторы: Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, Liang Zheng

cs.AI

Аннотация

В данной работе мы рассматриваем фундаментальный вопрос: "Можно ли обучать латентные диффузионные модели совместно с токенизатором на основе вариационного автокодировщика (VAE) в сквозном режиме?" Традиционная мудрость глубокого обучения гласит, что сквозное обучение предпочтительно, когда это возможно. Однако для латентных диффузионных трансформеров наблюдается, что сквозное обучение как VAE, так и диффузионной модели с использованием стандартной диффузионной функции потерь неэффективно и даже приводит к ухудшению итоговой производительности. Мы показываем, что хотя диффузионная функция потерь неэффективна, сквозное обучение становится возможным благодаря функции потерь выравнивания представлений (REPA) — это позволяет совместно настраивать как VAE, так и диффузионную модель в процессе обучения. Несмотря на свою простоту, предложенный метод обучения (REPA-E) демонстрирует впечатляющую производительность; ускоряя обучение диффузионной модели более чем в 17 и 45 раз по сравнению с методами REPA и стандартного обучения соответственно. Интересно, что мы наблюдаем, что сквозная настройка с использованием REPA-E также улучшает сам VAE; это приводит к улучшению структуры латентного пространства и повышению производительности генерации. С точки зрения итоговой производительности наш подход устанавливает новый эталон; достигая FID 1.26 и 1.83 с использованием и без использования классификатор-независимого руководства на данных ImageNet 256 x 256. Код доступен по адресу https://end2end-diffusion.github.io.

English

In this paper we tackle a fundamental question: "Can we train latent diffusion models together with the variational auto-encoder (VAE) tokenizer in an end-to-end manner?" Traditional deep-learning wisdom dictates that end-to-end training is often preferable when possible. However, for latent diffusion transformers, it is observed that end-to-end training both VAE and diffusion-model using standard diffusion-loss is ineffective, even causing a degradation in final performance. We show that while diffusion loss is ineffective, end-to-end training can be unlocked through the representation-alignment (REPA) loss -- allowing both VAE and diffusion model to be jointly tuned during the training process. Despite its simplicity, the proposed training recipe (REPA-E) shows remarkable performance; speeding up diffusion model training by over 17x and 45x over REPA and vanilla training recipes, respectively. Interestingly, we observe that end-to-end tuning with REPA-E also improves the VAE itself; leading to improved latent space structure and downstream generation performance. In terms of final performance, our approach sets a new state-of-the-art; achieving FID of 1.26 and 1.83 with and without classifier-free guidance on ImageNet 256 x 256. Code is available at https://end2end-diffusion.github.io.