REPA-E
参考链接
相关工作
前置知识
传统端到端的调优是无效的
当使用扩散模型的损失直接反向传播到 VAE 调优,会出来以下现象:
- 扩散模型会让 VAE 的隐空间结构变得更简单,方差变得更低。
- 一个更简单的隐空间结构创造了一个更简单的去噪问题,但性能却下降了。
面临问题与核心想法
REPA 通过引入一个外部指导特征提取模型,增强模型提取特征的能力。
隐扩散模型中 REPA 正则化存在上限。
传统的 REPA 可实现的最大对齐分数,受到了 VAE 隐空间特征的瓶颈限制,而对齐分数越高,说明学习到的特征越好,模型也就越好。
能否以端到端的方式,与潜在扩散模型一起训练 VAE 分词器?
也就是说,让 VAE 也参与到微调中来,这样, REPA 的上限就提高了。
但是传统端到端的调优是无效的。因此不能使用 diffusion loss. 为此,本模型改用了 REPA Loss.
阅读前疑问
怎么使用 REPA loss 来进行端到端的训练。
贡献
模型和方法
传统方法先训练好一个 VAE, 而后冻结 VAE, 训练 Diffusion Model.
朴素的端到端方法直接使用 Diffusion Loss, 进行 VAE 和 LDM 的联合训练。
本文提出的方法中, VAE 可以进行微调,扩散模型的可以训练。
本文使用 Diffusion Loss 来训练扩散模型,但是 Diffusion Model 在 VAE 之前进行停梯度操作。另一方面,则使用 REPA Loss 来微调 VAE.
另一方面,在这个 REPA-E 模型中, VAE 与扩散模型之间有一个 BatchNorm 块。
传统的 SD-VAE 中,使用训练数据计算一个方差统计量,然后进行 normalization, 从而减少计算量。但这样就需要检查所有的计算数据。
训练
\(\mathcal{L}_{REG}\) 分三个部分, GAN Loss, KL 散度,以及重建损失。这是 VAE 的损失,用来确保训练过程中不会影响 VAE 的重建性能。与 SD-VAE 一样,使用三个损失。
实验
结论
收获
本文中使用两个损失来训练两个不同的模块的操作很有意思,感觉比较少见。