用后验期望增强流匹配,让图像生成更稳定多样。
Posterior Augmented Flow Matching

- 用多个可能的目标轨迹替代单一目标,提升训练信号质量
- 在不同模型和数据集上最高降低3.4点FID50K,效果显著
- 适合追求高质量图像生成的开发者与研究者
流匹配(FM)通过时间依赖向量场将简单先验分布映射到复杂数据分布。然而,对于高维图像,每个训练样本仅监督一条轨迹和一个中间点,导致训练信号极度稀疏且方差大,易引发流坍缩——即模型记忆特定输入-输出对,使多样输入生成相似输出,失去泛化能力。本文提出后验增强流匹配(PAFM),理论上推广了FM,将单目标监督替换为给定中间状态与条件下的有效目标补全近似后验期望。PAFM将不可计算的后验分解为:(i) 假设终点下中间状态的似然度,(ii) 该终点在条件下的先验概率,并通过重要性采样构建多候选目标混合。证明了PAFM能无偏估计原始FM目标,同时通过聚合每中间点的多个合理延续轨迹,显著降低训练梯度方差。实验显示,PAFM在不同模型尺度(SiT-B/2和SiT-XL/2)、架构(SiT和MMDiT)及条件类型(类别与文本)的基准(ImageNet和CC12M)上,最多提升3.4点FID50K,计算开销几乎不变。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。