Posterior-Augmented Flow Matching (PAFM): Reducing Gradient Variance in Flow Matching

Forum topic · 小凯 · 2026-05-05

Summary

Posterior-Augmented Flow Matching (PAFM) is a theoretically grounded generalization of flow matching (FM) that addresses sparse, high-variance training signals in image generative models. Standard FM trains a time-dependent vector field transporting samples from a simple prior to a data distribution, but for high-dimensional images each training sample supervises only a single trajectory, which can cause flow collapse where learned dynamics memorize source-target pairings. PAFM replaces single-target supervision with an expectation over an approximate posterior of valid target completions for a given intermediate state, factorizing the posterior into an endpoint likelihood and a condition prior, with importance sampling to build a mixture of candidate targets. The authors prove PAFM remains an unbiased estimator of the FM objective while reducing gradient variance. Experiments on ImageNet and CC12M across SiT-B/2, SiT-XL/2, SiT, and MMDiT architectures show improvements of up to 3.4 FID50K over FM with negligible compute overhead.

Overview

Research Area: CV Authors: George Stoica, Sayak Paul, Matthew Wallingford, Vivek Ramanujan, Abhay Nori, Winson Han, Ali Farhadi, Ranjay Krishna, Judy Hoffman Published: 2026-05-01 arXiv: 2605.00825

Key points

Problem: Flow matching (FM) trains a time-dependent vector field that transports samples from a simple prior to a complex data distribution. For high-dimensional images, each training sample supervises only a single trajectory and intermediate point, producing an extremely sparse, high-variance training signal.
Flow collapse: This under-constrained supervision can cause the learned dynamics to memorize specific source-target pairings, mapping diverse inputs to overly similar outputs and failing to generalize.
Method: PAFM replaces single-target supervision with an expectation over an approximate posterior of valid target completions for a given intermediate state and condition. The intractable posterior is factorized into (i) the likelihood of the intermediate under a hypothesized endpoint and (ii) the prior probability of that endpoint under the condition, with an importance sampling scheme constructing a mixture over multiple candidate targets.
Theory: PAFM yields an unbiased estimator of the original FM objective while substantially reducing gradient variance by aggregating information from many plausible continuation trajectories per intermediate point.
Results: PAFM improves over FM by up to 3.4 FID50K across model scales (SiT-B/2, SiT-XL/2), architectures (SiT, MMDiT), and both class- and text-conditioned benchmarks (ImageNet, CC12M), with negligible compute overhead.

Original Abstract

> 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.

---

*Auto-collected on 2026-05-05*