REPA-E : Libérer le potentiel des VAE pour un réglage de bout en bout avec la diffusion latente des Transformers

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

April 14, 2025

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

cs.AI

Résumé

Dans cet article, nous abordons une question fondamentale : "Pouvons-nous entraîner des modèles de diffusion latente conjointement avec le tokeniseur de l'auto-encodeur variationnel (VAE) de manière end-to-end ?" La sagesse traditionnelle en apprentissage profond suggère que l'entraînement end-to-end est souvent préférable lorsque cela est possible. Cependant, pour les transformateurs de diffusion latente, il est observé que l'entraînement end-to-end du VAE et du modèle de diffusion en utilisant la perte de diffusion standard est inefficace, voire entraîne une dégradation des performances finales. Nous montrons que bien que la perte de diffusion soit inefficace, l'entraînement end-to-end peut être rendu possible grâce à la perte d'alignement des représentations (REPA) — permettant ainsi de régler conjointement le VAE et le modèle de diffusion pendant le processus d'entraînement. Malgré sa simplicité, la méthode d'entraînement proposée (REPA-E) montre des performances remarquables ; accélérant l'entraînement du modèle de diffusion de plus de 17x et 45x par rapport aux méthodes REPA et classiques, respectivement. Fait intéressant, nous observons que le réglage end-to-end avec REPA-E améliore également le VAE lui-même ; conduisant à une meilleure structure de l'espace latent et à de meilleures performances de génération en aval. En termes de performances finales, notre approche établit un nouvel état de l'art ; atteignant un FID de 1,26 et 1,83 avec et sans guidage sans classifieur sur ImageNet 256 x 256. Le code est disponible à l'adresse 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.