REPA-E: Freischaltung von VAEs für End-to-End-Optimierung mit latenter Diffusion Transformern

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

April 14, 2025

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

cs.AI

Zusammenfassung

In dieser Arbeit gehen wir einer grundlegenden Frage nach: „Können wir latente Diffusionsmodelle zusammen mit dem Tokenizer des variationalen Autoencoders (VAE) end-to-end trainieren?“ Die traditionelle Weisheit des Deep Learnings besagt, dass end-to-end Training, wenn möglich, oft vorzuziehen ist. Für latente Diffusions-Transformer wurde jedoch beobachtet, dass das end-to-end Training sowohl des VAE als auch des Diffusionsmodells mit dem Standard-Diffusionsverlust ineffektiv ist und sogar zu einer Verschlechterung der endgültigen Leistung führt. Wir zeigen, dass, obwohl der Diffusionsverlust ineffektiv ist, das end-to-end Training durch den Representation-Alignment (REPA)-Verlust ermöglicht werden kann – wodurch sowohl der VAE als auch das Diffusionsmodell während des Trainingsprozesses gemeinsam optimiert werden können. Trotz seiner Einfachheit zeigt das vorgeschlagene Trainingsrezept (REPA-E) bemerkenswerte Leistung; es beschleunigt das Training des Diffusionsmodells um mehr als das 17-fache bzw. 45-fache im Vergleich zu REPA und herkömmlichen Trainingsrezepten. Interessanterweise beobachten wir, dass das end-to-end Tuning mit REPA-E auch den VAE selbst verbessert; dies führt zu einer verbesserten latenten Raumstruktur und einer besseren Generationsleistung in nachgelagerten Anwendungen. In Bezug auf die endgültige Leistung setzt unser Ansatz einen neuen State-of-the-art; er erreicht einen FID von 1,26 und 1,83 mit und ohne Classifier-Free Guidance auf ImageNet 256 x 256. Der Code ist verfügbar unter 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.