REPA-E: Het ontgrendelen van VAE voor end-to-end afstemming met latente diffusie 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

Samenvatting

In dit artikel gaan we in op een fundamentele vraag: "Kunnen we latentie-diffusiemodellen samen met de variational auto-encoder (VAE) tokenizer end-to-end trainen?" Traditionele diep-lerenwijsheid stelt dat end-to-end training vaak de voorkeur verdient wanneer mogelijk. Echter, voor latentie-diffusietransformers wordt waargenomen dat end-to-end training van zowel VAE als diffusiemodel met standaard diffusieverlies niet effectief is, en zelfs leidt tot een verslechtering van de uiteindelijke prestaties. We laten zien dat hoewel diffusieverlies niet effectief is, end-to-end training mogelijk wordt gemaakt door de representation-alignment (REPA) loss — waardoor zowel VAE als diffusiemodel gezamenlijk kunnen worden afgesteld tijdens het trainingsproces. Ondanks zijn eenvoud toont het voorgestelde trainingsrecept (REPA-E) opmerkelijke prestaties; het versnelt de training van diffusiemodellen met meer dan 17x en 45x ten opzichte van REPA en standaard trainingsrecepten, respectievelijk. Interessant genoeg observeren we dat end-to-end afstemming met REPA-E ook de VAE zelf verbetert; wat leidt tot een verbeterde latentieruimtestructuur en betere downstream-generatieprestaties. Wat betreft de uiteindelijke prestaties stelt onze aanpak een nieuwe state-of-the-art; we behalen een FID van 1.26 en 1.83 met en zonder classifier-free guidance op ImageNet 256 x 256. Code is beschikbaar op 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.