VA-VAE

发表于 更新于 分类于 阅读次数： 本文字数： 969 阅读时长 ≈ 4 分钟

Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models[1]

作者主要是来自华科的Jingfeng Yao等人。论文引用[1]:

Time

  • 2025.Mar

Key Words

  • aligned latent space with visual tokenizers

总结

  1. 用Transformer架构的Latent diffusion models擅长生成高质量的images，然而，最近的研究表明了这个两阶段设计的优化困境：增加visual tokenizers中的per-token feature dimension的时候，提高了reconstruction 的质量，它需要很大的diffusion models何更多的training iterations来实现相当的生成性能。因此，现有的systems经常遇到次优化的解决方式，要么是由于tokenizer的information loss造成的visual artifacts，要么是由于大量的计算导致难以收敛，作者argue，这个dilemma源于学习没有约束的高维latent spaces的困难，为了解决这个，作者提出了，当训练visual tokenizers的时候，将latent space和预训练的vision foundation models进行对齐。作者提出了VA-VAE(Vision foundation model Aligned VAE)，极大地扩展了latent diffusion models的reconstruction-generation 的边界，使得DiTs在高维latent spaces快速收敛，为了利用VA-VAE的全部的potential，作者构建了一个增强版的DiT，用了改进的训练策略和架构设计，称之为LightningDiT。

  2. latent diffusion model利用连续值的VAE，或者visual tokenizer来压缩视觉信号，然后降低高分辨率图像生成的计算需求，这些visual tokenizer的性能，特别是压缩和重建的能力，决定了整体系统的有效性。一个直接的增强reconstruction能力的方法是增加visual tokens的feature dimension，有效地扩充了latent representation的information capacity，最近，一些有影响力的text-to-image工作探索了更高维的tokenizers，这些tokenizers提供了detail reconstruction，使得finer generative quality。 然而，随着研究的进行，一个reconstruction和generation之间的优化困境显示出来了，特别地，当增加token feature dimension的时候，提高了tokenizer的reconstruction accuracy，但是生成性能退化了。当前，两个常见的策略是：第一个用scaling up model的参数，这需要更多的训练资源；第二个是限制tokenizer的reconstruction capacity，为了更快地收敛。这个妥协的reconstruction 质量限制了生成性能的上限，导致生成结果的不完美。两个方法都涉及trade-offs。

  3. 在本文中，作者提出了一个直接的方法来解决这个困境，作者从自回归中生成中得到启发，增加离散值VAEs的codebook size导致low codebook utilization，通过可视化不同feature dimensions的latent space 分布，作者观察到，更高维度的tokenizers会一些less spread-out latent representations，这个分析表明：优化困难源于从零训练的学习无约束的高维latent space的困难。为了解决这个问题，作者开发了一个vision foundation model为连续的VAEs指导优化策略。作者的key finding展示了：学习vision foundation model guided的latent representation增强了高维tokenizer的生成性能，同时保留了它们原始的重建能力。 主要的contribution是Vision Foundation model alignment Loss(VF Loss)，是一个即插即用的module，在tokenizer训练期间，将latent representations和预训练的vision foundation models对齐，用预训练的vision foundation models来初始化VAE encoders被证明是无效的，可能是由于latent representation 快速地偏离了它的初始state来优化reconstruction，作者发现，一个仔细涉及的joint reconstruction和alignment loss是重要的。作者的alignment loss是精心涉及的，来正则化高维latent spaces，不需要限制它们的capacity。首先，作者强制elelment-wise 和pair-wise相似度，来确保feature space中的global 和local structures的综合的regularization。其次，作者在similarity cost中引入了一个margin，在alignment中提供controlled flexibility，作者研究了不同vision foundation models的影响。