REPA-E: Desbloqueando VAE para ajuste de extremo a extremo con difusión latente en Transformers

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

April 14, 2025

Autores: Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhenchang Xing, Saining Xie, Liang Zheng

cs.AI

Resumen

En este artículo abordamos una pregunta fundamental: "¿Podemos entrenar modelos de difusión latente junto con el tokenizador de auto-codificadores variacionales (VAE) de manera integral?" La sabiduría tradicional del aprendizaje profundo sugiere que el entrenamiento integral es preferible cuando es posible. Sin embargo, para los transformadores de difusión latente, se observa que el entrenamiento integral tanto del VAE como del modelo de difusión utilizando la pérdida estándar de difusión es ineficaz, incluso causando una degradación en el rendimiento final. Demostramos que, aunque la pérdida de difusión es ineficaz, el entrenamiento integral puede habilitarse mediante la pérdida de alineación de representaciones (REPA), permitiendo que tanto el VAE como el modelo de difusión se ajusten conjuntamente durante el proceso de entrenamiento. A pesar de su simplicidad, la receta de entrenamiento propuesta (REPA-E) muestra un rendimiento notable; acelerando el entrenamiento del modelo de difusión en más de 17x y 45x en comparación con las recetas de entrenamiento REPA y estándar, respectivamente. Curiosamente, observamos que el ajuste integral con REPA-E también mejora el VAE en sí; conduciendo a una estructura mejorada del espacio latente y a un mejor rendimiento en la generación posterior. En términos de rendimiento final, nuestro enfoque establece un nuevo estado del arte; logrando un FID de 1.26 y 1.83 con y sin guía libre de clasificador en ImageNet 256 x 256. El código está disponible en 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.