[论文解读] Flow Matching in Latent Space

Quan Dao, Hao Phung+2 more|arXiv (Cornell University)|Jul 17, 2023

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

本文提出潜在流匹配（LFM），一种新颖的框架，将流匹配应用于预训练自编码器的潜在空间，以提升高分辨率图像生成的计算效率和可扩展性。通过在压缩的潜在表征上训练流匹配，并结合无分类器指导实现条件生成，LFM 在图像修复、语义到图像生成以及类别条件生成等任务上达到最先进性能，在 CelebA-HQ 256 上 FID 分数分别达到 4.09 和 26.3，同时保持快速采样和更低的训练成本。

ABSTRACT

Flow matching is a recent framework to train generative models that exhibits impressive empirical performance while being relatively easier to train compared with diffusion-based models. Despite its advantageous properties, prior methods still face the challenges of expensive computing and a large number of function evaluations of off-the-shelf solvers in the pixel space. Furthermore, although latent-based generative methods have shown great success in recent years, this particular model type remains underexplored in this area. In this work, we propose to apply flow matching in the latent spaces of pretrained autoencoders, which offers improved computational efficiency and scalability for high-resolution image synthesis. This enables flow-matching training on constrained computational resources while maintaining their quality and flexibility. Additionally, our work stands as a pioneering contribution in the integration of various conditions into flow matching for conditional generation tasks, including label-conditioned image generation, image inpainting, and semantic-to-image generation. Through extensive experiments, our approach demonstrates its effectiveness in both quantitative and qualitative results on various datasets, such as CelebA-HQ, FFHQ, LSUN Church & Bedroom, and ImageNet. We also provide a theoretical control of the Wasserstein-2 distance between the reconstructed latent flow distribution and true data distribution, showing it is upper-bounded by the latent flow matching objective. Our code will be available at https://github.com/VinAIResearch/LFM.git.

研究动机与目标

• 解决像素空间流匹配模型计算成本高且采样缓慢的问题。
• 通过将流匹配迁移至预训练自编码器的潜在空间，实现高效、可扩展的高分辨率图像合成。
• 将流匹配扩展至条件生成任务，包括类别条件图像生成、图像修复和语义到图像合成。
• 将无分类器指导引入流匹配，以提升条件生成质量。
• 在理论上界定潜在空间中重构分布与真实数据分布之间的 Wasserstein-2 距离。

提出的方法

• 在预训练自编码器的潜在空间中应用流匹配，使用瓶颈表征作为生成建模的潜在流形。
• 设计一个条件速度场网络，输入为潜在码、条件（如类别标签、掩码或图像）和时间，以预测流动力学。
• 通过同时训练条件流和无条件流，将无分类器指导集成到速度场网络中，实现灵活的条件采样。
• 使用自适应 ODE 求解器（如 Runge-Kutta）在潜在空间中实现高效且精确的采样。
• 联合训练潜在流匹配网络与一个条件嵌入头（如布局到图像任务），使用流匹配目标。
• 理论分析表明，重构分布与真实数据分布之间的 Wasserstein-2 距离，由潜在流匹配目标上界控制。

实验结果

研究问题

• RQ1 潜在空间中的流匹配是否能显著降低计算成本，同时保持高图像质量？
• RQ2 潜在流匹配在图像修复和语义到图像合成等多样化条件图像生成任务中，泛化能力如何？
• RQ3 无分类器指导能否被有效适配至流匹配，以提升条件生成性能？
• RQ4 潜在流匹配目标与真实数据分布之间的 Wasserstein-2 距离在理论上存在何种关系？
• RQ5 自编码器架构的选择如何影响潜在流匹配的性能？

主要发现

• 在 CelebA-HQ 256 图像修复基准上，LFM 达到 4.09 的 FID 分数，优于大多数现有方法，接近 MAT 的 SOTA 分数 2.94。
• 在语义到图像任务中，LFM 在 CelebA-HQ 256 上达到 26.3 的 FID 分数，展现出强大性能，且仅使用简单架构，无需专门设计。
• 该方法仅通过 50 步自适应 ODE 求解器即可实现快速采样，推理速度具有竞争力。
• LFM 中的无分类器指导带来显著性能提升，在相同设置下将 FID 从 16.71 降低至 8.58。
• 理论分析证实，重构分布与真实数据分布之间的 Wasserstein-2 距离，由潜在流匹配目标上界控制，验证了方法的统计一致性。
• 尽管模型规模小于 LDM，LFM 在引入无分类器指导后仍取得更优的性能提升，表明其在更大架构下具备良好的可扩展潜力。