Variational Rectified Flow Matching

📄 arXiv: 2502.09616v1 📥 PDF

作者: Pengsheng Guo, Alexander G. Schwing

分类: cs.LG, cs.CV

发布日期: 2025-02-13

💡 一句话要点

提出变分校正流匹配，通过建模多模态速度向量场提升生成模型性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 生成模型 校正流匹配 变分推断 多模态学习 速度向量场

📋 核心要点

  1. 传统校正流匹配使用均方误差损失训练，导致学习到的速度向量场无法捕捉真实的多模态特性。
  2. 变分校正流匹配通过学习和采样多模态流方向，能够更好地建模复杂的数据分布。
  3. 实验表明，变分校正流匹配在多个数据集上取得了优异的生成效果，验证了其有效性。

📝 摘要（中文）

本文研究了变分校正流匹配（Variational Rectified Flow Matching），该框架通过建模多模态速度向量场来增强经典的校正流匹配。在推理时，经典的校正流匹配通过沿速度向量场积分求解常微分方程，将样本从源分布“移动”到目标分布。在训练时，通过线性插值源分布和目标分布中随机抽取的一对耦合样本来学习速度向量场。这导致了在同一位置指向不同方向的“真值”速度向量场，即速度向量场是多模态/模糊的。然而，由于训练使用标准的均方误差损失，因此学习到的速度向量场平均了“真值”方向，而不是多模态的。相比之下，变分校正流匹配学习并从多模态流方向中采样。我们在合成数据、MNIST、CIFAR-10和ImageNet上展示了变分校正流匹配带来了引人注目的结果。

🔬 方法详解

问题定义：论文旨在解决传统校正流匹配在学习速度向量场时无法有效建模多模态分布的问题。现有方法使用均方误差损失，导致学习到的向量场是真值方向的平均，丢失了数据分布的复杂性，限制了生成模型的性能。

核心思路：论文的核心思路是引入变分推断，将速度向量场建模为一个概率分布，从而能够学习和采样多模态的流方向。通过学习速度向量场的分布，模型可以更好地捕捉数据中的不确定性和多样性。

技术框架：整体框架包括以下几个主要步骤：1) 从源分布和目标分布中采样数据点对；2) 使用神经网络建模速度向量场的条件分布；3) 使用变分推断方法学习该分布的参数，例如使用ELBO（Evidence Lower Bound）作为损失函数；4) 在推理阶段，从学习到的速度向量场分布中采样，并使用ODE求解器生成样本。

关键创新：最重要的技术创新在于将变分推断引入到校正流匹配框架中，从而能够学习和采样多模态的速度向量场。与传统方法直接回归一个平均的速度向量不同，该方法学习的是一个速度向量的分布，能够更好地捕捉数据中的不确定性和多样性。

关键设计：关键设计包括：1) 使用神经网络参数化速度向量场的条件分布，例如使用高斯混合模型或变分自编码器；2) 使用ELBO作为损失函数，鼓励模型学习一个能够解释观测数据的速度向量场分布；3) 在推理阶段，使用随机采样或确定性采样方法从学习到的分布中获取速度向量，并使用ODE求解器生成样本。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，变分校正流匹配在合成数据、MNIST、CIFAR-10和ImageNet等数据集上取得了显著的性能提升。与传统的校正流匹配方法相比，该方法能够生成更高质量、更多样化的样本，验证了其有效性。

🎯 应用场景

该研究成果可应用于图像生成、图像编辑、数据增强等领域。通过建模多模态速度向量场，可以生成更加真实、多样化的图像，提高生成模型的性能和泛化能力。此外，该方法还可以应用于其他生成任务，例如音频生成、文本生成等。

📄 摘要（原文）

We study Variational Rectified Flow Matching, a framework that enhances classic rectified flow matching by modeling multi-modal velocity vector-fields. At inference time, classic rectified flow matching 'moves' samples from a source distribution to the target distribution by solving an ordinary differential equation via integration along a velocity vector-field. At training time, the velocity vector-field is learnt by linearly interpolating between coupled samples one drawn from the source and one drawn from the target distribution randomly. This leads to ''ground-truth'' velocity vector-fields that point in different directions at the same location, i.e., the velocity vector-fields are multi-modal/ambiguous. However, since training uses a standard mean-squared-error loss, the learnt velocity vector-field averages ''ground-truth'' directions and isn't multi-modal. In contrast, variational rectified flow matching learns and samples from multi-modal flow directions. We show on synthetic data, MNIST, CIFAR-10, and ImageNet that variational rectified flow matching leads to compelling results.