REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers
This work addresses a fundamental bottleneck in training latent diffusion models for image generation, offering significant efficiency and performance gains for researchers and practitioners in generative AI.
The paper tackles the problem of end-to-end training variational auto-encoders (VAEs) with latent diffusion transformers, which traditionally degrades performance, by introducing a representation-alignment (REPA) loss that enables effective joint tuning, resulting in a 17x to 45x speedup in training and achieving state-of-the-art FID scores of 1.12 and 1.69 on ImageNet 256x256.
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.12 and 1.69 with and without classifier-free guidance on ImageNet 256 x 256. Code is available at https://end2end-diffusion.github.io.