REPA-E: Desbloqueando VAEs para Ajuste Fim a Fim com Difusão Latente em Transformers

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

April 14, 2025

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

cs.AI

Resumo

Neste artigo, abordamos uma questão fundamental: "Podemos treinar modelos de difusão latente juntamente com o tokenizador de autoencoder variacional (VAE) de maneira end-to-end?" A sabedoria tradicional do aprendizado profundo sugere que o treinamento end-to-end é frequentemente preferível quando possível. No entanto, para transformadores de difusão latente, observa-se que o treinamento end-to-end tanto do VAE quanto do modelo de difusão usando a perda de difusão padrão é ineficaz, chegando a causar uma degradação no desempenho final. Mostramos que, embora a perda de difusão seja ineficaz, o treinamento end-to-end pode ser habilitado por meio da perda de alinhamento de representação (REPA) — permitindo que tanto o VAE quanto o modelo de difusão sejam ajustados conjuntamente durante o processo de treinamento. Apesar de sua simplicidade, a receita de treinamento proposta (REPA-E) apresenta um desempenho notável; acelerando o treinamento do modelo de difusão em mais de 17x e 45x em comparação com as receitas de treinamento REPA e convencional, respectivamente. Curiosamente, observamos que o ajuste end-to-end com REPA-E também melhora o próprio VAE; resultando em uma estrutura de espaço latente aprimorada e em um melhor desempenho de geração downstream. Em termos de desempenho final, nossa abordagem estabelece um novo estado da arte; alcançando FID de 1,26 e 1,83 com e sem orientação livre de classificador no ImageNet 256 x 256. O código está disponível em 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.