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

作者：Xingjian Leng, Jaskirat Singh, Yunzhong Hou, Zhen-Chang Xing, Saining Xie, Liang Zheng · 发表于：IEEE International Conference on Computer Vision · 年份：2025 · DOI：10.1109/iccv51701.2025.01697 · 被引用次数：204 · 研究领域：Computer Science

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-toend 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 $17 \times$ and $45 \times$ 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 classifierfree guidance on ImageNet $256 \times 256$. Code is available at https://end2end-diffusion.qithub.io.