Posterior Augmented Flow Matching

📄 arXiv: 2605.00825v1 📥 PDF

作者: George Stoica, Sayak Paul, Matthew Wallingford, Vivek Ramanujan, Abhay Nori, Winson Han, Ali Farhadi, Ranjay Krishna, Judy Hoffman

分类: cs.CV

发布日期: 2026-05-01

🔗 代码/项目: GITHUB

💡 一句话要点

提出后验增强Flow Matching，解决高维图像生成中Flow Collapse问题

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

关键词: Flow Matching 生成模型 后验增强 图像生成 Flow Collapse 重要性采样 高维数据 向量场

📋 核心要点

  1. Flow Matching在高维图像生成中面临训练信号稀疏和高方差的问题，导致Flow Collapse，泛化能力差。
  2. PAFM通过引入目标补全的后验分布，用后验期望代替单目标监督，聚合多个合理轨迹的信息。
  3. 实验表明，PAFM在不同模型和数据集上显著优于FM，FID50K指标提升高达3.4，计算开销可忽略。

📝 摘要（中文）

Flow Matching (FM) 训练一个时变向量场，将样本从简单的先验分布传输到复杂的数据分布。然而，对于高维图像，每个训练样本仅监督单个轨迹和中间点，产生极其稀疏和高方差的训练信号。这种欠约束的监督会导致flow collapse，即学习到的动态过程记忆特定的源-目标配对，将不同的输入映射到过于相似的输出，从而无法泛化。我们引入了后验增强Flow Matching (PAFM)，这是FM的一个理论推广，它用给定中间状态和条件下的有效目标补全的近似后验的期望来代替单目标监督。PAFM将这个难处理的后验分解为（i）假设终点下中间状态的似然和（ii）条件下的终点先验概率，并使用重要性采样方案来构建多个候选目标的混合。我们证明PAFM产生原始FM目标的无偏估计，同时通过聚合每个中间状态的多个合理延续轨迹的信息，显著降低训练期间的梯度方差。最后，我们表明PAFM在不同的模型规模（SiT-B/2和SiT-XL/2）、不同的架构（SiT和MMDiT）以及类和文本条件基准（ImageNet和CC12M）上，相比FM提高了高达3.4 FID50K，而计算开销可以忽略不计。

🔬 方法详解

问题定义：Flow Matching方法在高维图像生成任务中，由于每个训练样本仅监督一条轨迹上的一个点，导致训练信号极其稀疏和高方差。这种欠约束的监督使得模型容易记忆特定的源-目标配对，从而产生Flow Collapse现象，即模型将不同的输入映射到相似的输出，泛化能力较差。现有方法难以有效解决高维图像生成中的Flow Collapse问题。

核心思路：PAFM的核心思路是用目标补全的后验分布的期望来代替Flow Matching中的单目标监督。具体来说，对于给定的中间状态和条件，PAFM不是简单地监督到一个目标，而是考虑多个可能的“补全”目标，并计算这些目标在后验分布下的期望。这样可以聚合来自多个合理延续轨迹的信息，从而降低梯度方差，提高训练的稳定性。

技术框架：PAFM的整体框架包括以下几个主要步骤：1) 给定一个中间状态和条件，生成多个候选目标；2) 计算每个候选目标在后验分布下的概率；3) 使用重要性采样方法，根据后验概率对候选目标进行加权；4) 使用加权后的候选目标来监督向量场的训练。PAFM的关键在于如何有效地近似后验分布，以及如何使用重要性采样来降低方差。

关键创新：PAFM最重要的技术创新在于用后验增强的方式来解决Flow Matching中的欠约束监督问题。与传统的Flow Matching方法只使用单个目标进行监督不同，PAFM利用多个可能的“补全”目标，从而聚合了更多的信息，降低了梯度方差，提高了训练的稳定性。这种后验增强的思想可以应用于其他类似的生成模型中。

关键设计：PAFM的关键设计包括：1) 后验分布的近似：PAFM将后验分布分解为似然和先验两部分，分别进行建模。似然部分可以使用神经网络来学习，先验部分可以使用预训练的模型或者其他先验知识。2) 重要性采样：PAFM使用重要性采样来降低方差。具体来说，PAFM根据后验概率对候选目标进行加权，从而使得训练更加关注于那些更可能的“补全”目标。3) 损失函数：PAFM使用标准的Flow Matching损失函数，但是将目标替换为加权后的候选目标。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，PAFM在ImageNet和CC12M数据集上，使用SiT和MMDiT架构，以及不同模型规模（SiT-B/2和SiT-XL/2）下，相比于Flow Matching方法，FID50K指标提升高达3.4。这表明PAFM能够有效缓解Flow Collapse问题，提高生成模型的性能，并且计算开销可以忽略不计。

🎯 应用场景

PAFM具有广泛的应用前景，可以应用于图像生成、图像编辑、视频生成等领域。尤其是在高分辨率图像生成和复杂场景生成方面，PAFM能够有效缓解Flow Collapse问题，提高生成质量和泛化能力。此外，PAFM的后验增强思想也可以推广到其他生成模型中，例如扩散模型和GANs。

📄 摘要（原文）

Flow matching (FM) trains a time-dependent vector field that transports samples from a simple prior to a complex data distribution. However, for high-dimensional images, each training sample supervises only a single trajectory and intermediate point, yielding an extremely sparse and high-variance training signal. This under-constrained supervision can cause flow collapse, where the learned dynamics memorize specific source-target pairings, mapping diverse inputs to overly similar outputs, failing to generalize. We introduce Posterior-Augmented Flow Matching (PAFM), a theoretically grounded generalization of FM that replaces single-target supervision with an expectation over an approximate posterior of valid target completions for a given intermediate state and condition. PAFM factorizes this intractable posterior into (i) the likelihood of the intermediate under a hypothesized endpoint and (ii) the prior probability of that endpoint under the condition, and uses an importance sampling scheme to construct a mixture over multiple candidate targets. We prove that PAFM yields an unbiased estimator of the original FM objective while substantially reducing gradient variance during training by aggregating information from many plausible continuation trajectories per intermediate. Finally, we show that PAFM improves over FM by up to 3.4 FID50K across different model scales (SiT-B/2 and SiT-XL/2), different architectures (SiT and MMDiT), and in both class and text conditioned benchmarks (ImageNet and CC12M), with a negligible increase in the compute overhead. Code: https://github.com/gstoica27/PAFM.git.