REPA-E: Sbloccare i VAE per l'ottimizzazione end-to-end con la diffusione latente dei Transformer

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

April 14, 2025

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

cs.AI

Abstract

In questo articolo affrontiamo una questione fondamentale: "Possiamo addestrare modelli di diffusione latente insieme al tokenizzatore di autoencoder variazionale (VAE) in modo end-to-end?" La saggezza tradizionale del deep learning suggerisce che l'addestramento end-to-end sia spesso preferibile quando possibile. Tuttavia, per i trasformatori di diffusione latente, si osserva che l'addestramento end-to-end sia del VAE che del modello di diffusione utilizzando la loss standard di diffusione è inefficace, causando persino un degrado delle prestazioni finali. Dimostriamo che, sebbene la loss di diffusione sia inefficace, l'addestramento end-to-end può essere sbloccato attraverso la loss di allineamento della rappresentazione (REPA) -- consentendo sia al VAE che al modello di diffusione di essere regolati congiuntamente durante il processo di addestramento. Nonostante la sua semplicità, la ricetta di addestramento proposta (REPA-E) mostra prestazioni notevoli; accelerando l'addestramento del modello di diffusione di oltre 17x e 45x rispetto alle ricette REPA e vanilla, rispettivamente. È interessante notare che osserviamo come la regolazione end-to-end con REPA-E migliori anche il VAE stesso; portando a una struttura dello spazio latente migliorata e a prestazioni di generazione a valle superiori. In termini di prestazioni finali, il nostro approccio stabilisce un nuovo stato dell'arte; raggiungendo un FID di 1.26 e 1.83 con e senza guida classifier-free su ImageNet 256 x 256. Il codice è disponibile all'indirizzo 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.