arXiv:2504.10483cs.CVcs.LG2025-04ICCV被引 200

提出REPA-E方法,实现VAE与扩散模型端到端联合训练

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

论文配图:REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers
图 1 · 摘自论文原文
  • 用表示对齐损失替代扩散损失,实现编码器与扩散模型协同优化
  • 训练速度提升17倍以上,图像生成质量达新基准(FID 1.12)
  • 适合追求高效训练和高质量生成的扩散模型研究者

本文探讨一个根本性问题:能否将潜在扩散模型与变分自编码器(VAE)编码器进行端到端联合训练?传统深度学习观点认为端到端训练更优,但实验发现,使用标准扩散损失联合训练时,性能反而下降。我们证明,虽然扩散损失无效,但通过引入表示对齐(REPA)损失,可成功解锁端到端训练。所提出的训练方案(REPA-E)虽简单,却效果显著:相比REPA和原始训练,扩散模型训练速度分别提升17倍和45倍以上。有趣的是,该方法还提升了VAE自身表现,优化了潜在空间结构,进而改善下游生成效果。在ImageNet 256×256上,无分类器引导时达到FID 1.69,有引导时达FID 1.12,刷新当前最优水平。代码已开源。

原文摘要 · Abstract (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.12 and 1.69 with and without classifier-free guidance on ImageNet 256 x 256. Code is available at https://end2end-diffusion.github.io.

扩散模型端到端训练VAE图像生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。