45倍加速+最新SOTA！VAE与扩散模型迎来端到端联合训练：REPA-E让VAE自我进化！

文章主旨：

通过REPA-E框架实现VAE与扩散模型的端到端联合训练，显著提升生成性能与训练效率。

关键要点：

• 提出REPA Loss代替传统扩散损失，实现VAE与扩散模型的目标对齐，解决两阶段训练目标不一致性。
• 训练效率提升45倍，生成质量刷新当前最佳记录，FID从5.9降至4.07。
• 端到端训练改善VAE隐空间结构，使其适配扩散模型并提升生成细节。
• 技术细节包括批归一化层、自适应隐空间优化及正则化损失的设计。
• 实验验证REPA-E在不同模型规模、设置下的性能提升及泛化性。

内容结构：

1. 解决的问题：

现有隐空间扩散模型采用两阶段训练方式，导致VAE与扩散模型的优化目标不一致，限制生成性能。尝试直接联合训练时传统扩散损失失效，甚至导致性能下降。

2. 提出的方案：

通过表示对齐损失（REPA Loss）代替传统扩散损失，协调VAE与扩散模型的训练目标，实现端到端联合优化。

3. 应用的技术：

• 表示对齐损失：对齐隐空间表示的分布，提高生成性能。
• 端到端梯度传播：动态调整VAE的隐空间结构。
• 自适应隐空间优化：平衡VAE的重建能力与扩散模型的训练需求。

4. 达到的效果：

• 训练步数减少45倍，生成质量显著提升（ImageNet FID刷新记录）。
• 改善不同类型VAE的隐空间结构，提升下游任务性能。
• 端到端训练实现统一优化，提升模型扩展性与稳健性。

5. 实验与分析：

• 定量评估：REPA-E在训练速度与性能上超越传统方法，80轮训练FID达到4.07。
• 定性评估：生成图像结构更合理，质量更优。
• 泛化性分析：对不同模型规模、编码器、VAE架构均表现出一致性能提升。
• 组件消融实验：分析批归一化、stop-grad操作等的关键作用。
• 端到端训练对VAE影响：改善隐空间结构，使其适应生成任务需求。

文章总结：

REPA-E框架通过端到端训练释放了VAE潜力，显著提升隐空间扩散模型的训练效率与生成性能，推动了相关研究的发展。

警惕，90%企业困在「需求沼泽」里

了解更多

AI生成未来

查看原文：45倍加速+最新SOTA！VAE与扩散模型迎来端到端联合训练：REPA-E让VAE自我进化！

文章来源：

AI生成未来

扫码关注公众号