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

📄 arXiv: 2504.10483v3 📥 PDF

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

分类: cs.CV, cs.LG

发布日期: 2025-04-14 (更新: 2025-10-22)

💡 一句话要点

提出REPA-E，通过表征对齐损失实现VAE与潜在扩散Transformer的端到端训练。

🎯 匹配领域: 支柱四：生成式动作 (Generative Motion)

关键词: 潜在扩散模型 变分自编码器 端到端训练 表征对齐 图像生成

📋 核心要点

  1. 现有方法无法有效端到端训练VAE和潜在扩散模型，导致性能下降。
  2. 提出REPA-E，利用表征对齐损失，实现VAE和扩散模型的联合优化。
  3. REPA-E显著加速训练，并提升VAE的潜在空间结构和生成性能，达到SOTA。

📝 摘要（中文）

本文旨在解决一个根本问题：是否可以端到端地训练潜在扩散模型以及变分自编码器（VAE）tokenizer？ 传统的深度学习经验表明，如果可能，端到端训练通常是首选。然而，对于潜在扩散Transformer，观察到使用标准扩散损失进行VAE和扩散模型的端到端训练是无效的，甚至会导致最终性能下降。我们证明，虽然扩散损失无效，但可以通过表征对齐（REPA）损失来解锁端到端训练——允许在训练过程中联合调整VAE和扩散模型。尽管其简单性，所提出的训练方法（REPA-E）显示出卓越的性能；分别比REPA和vanilla训练方法加速扩散模型训练超过17倍和45倍。有趣的是，我们观察到使用REPA-E进行端到端调整也改进了VAE本身；从而改善了潜在空间结构和下游生成性能。在最终性能方面，我们的方法创造了新的state-of-the-art；在ImageNet 256 x 256上，使用和不使用无分类器指导的情况下，分别实现了1.12和1.69的FID。

🔬 方法详解

问题定义：论文旨在解决潜在扩散模型中VAE tokenizer和扩散模型无法有效进行端到端训练的问题。现有方法在端到端训练时，使用标准扩散损失会导致性能下降，无法充分发挥端到端训练的优势。这限制了模型性能的进一步提升，也增加了训练的复杂性。

核心思路：论文的核心思路是引入表征对齐（REPA）损失，以替代或补充标准扩散损失，从而实现VAE和扩散模型的有效端到端训练。REPA损失旨在对齐VAE编码器输出的潜在表征和扩散模型学习到的表征，从而促进二者之间的信息传递和协同优化。

技术框架：REPA-E的整体框架包括一个VAE（包含编码器和解码器）和一个潜在扩散Transformer。首先，图像通过VAE编码器映射到潜在空间。然后，扩散模型在这个潜在空间中进行训练，学习如何从噪声中生成潜在表征。最后，VAE解码器将潜在表征解码回图像空间。关键在于，REPA损失在训练过程中被用来对齐VAE编码器输出的潜在表征和扩散模型学习到的表征。

关键创新：最重要的技术创新点是REPA损失的应用，它使得VAE和扩散模型能够在端到端训练中协同优化。与传统方法仅依赖扩散损失不同，REPA损失直接促进了VAE和扩散模型之间的信息交互，从而避免了性能下降的问题。

关键设计：REPA损失的具体形式未知（论文中未明确给出，需要参考相关REPA论文）。关键在于如何有效地计算和应用REPA损失，以确保潜在表征的对齐。此外，VAE和扩散模型的网络结构、训练参数以及损失函数的权重等也需要进行仔细调整，以获得最佳的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

REPA-E训练方法显著加速了扩散模型的训练，相比REPA和vanilla训练方法分别提升了17倍和45倍。同时，REPA-E还提升了VAE的性能，改善了潜在空间结构。在ImageNet 256 x 256数据集上，REPA-E取得了SOTA结果，使用和不使用无分类器指导的情况下，FID分别为1.12和1.69。

🎯 应用场景

该研究成果可广泛应用于图像生成、图像编辑、视频生成等领域。通过端到端训练，可以提升生成模型的效率和质量，降低训练成本。此外，该方法还可以应用于其他需要联合训练多个模块的生成模型，具有重要的实际价值和潜在的未来影响。

📄 摘要（原文）

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.