Reproduce “REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers” on OpenResearch

About this paper

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

REPA-E jointly adapts a VAE tokenizer and diffusion transformer using representation-alignment loss from a pretrained vision encoder, while stop-gradient prevents denoising loss from updating the VAE. On ImageNet 256×256, SiT-XL reaches gFID 4.07 after 80 epochs, versus REPA’s 7.90 at the same duration and 5.90 after 800 epochs. This supports fewer training epochs in tested settings, not wall-clock savings.

Key insights

• Direct backpropagation of the diffusion loss to the VAE, while simplifying the denoising task, severely degrades the VAE's latent space structure and thus the final image generation quality.
• A strong correlation exists between higher representation-alignment scores (measured between intermediate diffusion features and perceptual model features) and improved generation performance, offering an effective proxy for optimization.
• The maximum achievable representation alignment is limited by the fixed features provided by the pre-trained VAE, indicating that tuning the VAE's latent space can further enhance alignment and generative quality.

Summary by alphaXiv.