Posterior Augmented Flow Matching

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. Code: https://github.com/gstoica27/PAFM.git.

• The paper introduces Posterior-Augmented Flow Matching to densify supervision in generative models by aggregating posterior target candidates.
• It leverages importance sampling over multiple candidate targets to significantly reduce gradient variance and improve training stability.
• Empirical results demonstrate enhanced performance on ImageNet-1K and CC12M, achieving up to a 3.4 FID improvement at marginal computational cost.

Posterior Augmented Flow Matching: Dense Supervision for Flow-Based Generative Models

Introduction and Motivation

Flow Matching (FM) has become a highly effective framework for training continuous-time generative models, defining a vector field that smoothly transports samples from a tractable prior to a complex target distribution. In practical high-dimensional settings, however, the FM learning signal is fundamentally sparse: each training step supervises a single interpolant between a given source and target, and assigns a flow only for a single intermediate point along that trajectory. This can hinder generalization and induce flow collapse, where the model overfits to specific trajectories rather than generalizing to the true distribution of plausible completions for any intermediate.

The Posterior-Augmented Flow Matching (PAFM) formulation addresses this core limitation by replacing one-to-one supervision with posterior-weighted, expectation-based aggregation over multiple compatible targets for each intermediate state and condition. This modification produces denser, lower-variance gradients, guiding the model to learn a vector field aligned with the full set of plausible flow completions.

Figure 1: Comparison of FM (left) and PAFM (right): FM provides a sparse, one-to-one supervision; PAFM aggregates supervision over the posterior of compatible targets, yielding denser, more coherent plausible flows.

Methodological Advances

From Sparse to Posterior-Aggregated Supervision

PAFM formulates the learning objective such that for each intermediate point $z_t$ and condition $y$, rather than only supervising with the original paired target $z$, the model aggregates over the full (intractable) posterior $p(z'|z_t, y)$ of possible targets consistent with $z_t$ and $y$. This leads to the following effective loss:

$$\mathcal{L}^{(PAFM)}(\theta) = \mathbb{E}_{t, y, z_t}\ \mathbb{E}_{z' \sim p(z'|z_t, y)} \left\|f_\theta(z_t|t, y) - v(z_t|z')\right\|^2$$

Estimating this posterior exactly is infeasible in high-dimensional data, but PAFM factorizes the posterior into the likelihood of the intermediate under a target candidate and the prior over the candidate under the condition. Practical implementation uses importance sampling to construct weighted mixtures from a candidate pool.

The authors prove that PAFM is, in expectation, an unbiased estimator of the standard FM objective. However, critically, PAFM allows $K > 1$ candidate targets per intermediate, and theoretical analysis (Theorem 2) establishes that the gradient variance of PAFM is lower-bounded by a factor proportional to the empirical Effective Sample Size (ESS) of the candidate set, which is strictly greater than that of single-pair FM in all non-degenerate settings.

Integration with Modern Flow Models

PAFM is compatible with established architecture and trajectory-shaping improvements, including rectified flow and scalable Transformer-based generative models (e.g., SiT, MMDiT), without modification of the underlying interpolant.

PAFM can utilize diverse strategies for target candidate selection:

• Nearest neighbors in latent space via FAISS indexing,
• Random augmentations (e.g., multi-scale image crops),
• Stochastic sampling within the VAE latent posterior for each image.

This flexibility enables straightforward adoption across conditional/unconditional, class-based, and text-conditioned generation tasks.

Empirical Results

Robustness in Low-Data Regimes and Flow Field Accuracy

PAFM demonstrates enhanced robustness to data sparsity and improved capacity to learn accurate vector fields relative to FM—both qualitative and quantitative gains are illustrated on canonical low-dimensional generative tasks.

Figure 2: PAFM produces distributions and velocity fields that are more robust to data sparsity and high-dimensional ambiguity than standard FM.

Figure 3: Under varying data sparsity, PAFM learns the true target distribution better than FM—as seen by KDE on generated samples.

Scaling to Realistic High-Dimensional Data

Experiments on ImageNet-1K and CC12M using SiT and MMDiT architectures, across both class- and text-conditioned regimes, reveal that PAFM consistently improves performance. Notably, PAFM reduces FID50K by up to 3.4 points relative to FM at comparable computational cost (overhead of only 6.6% in throughput, <1% in memory, and unchanged FLOPs per step). The improvement is robust across variants of candidate selection, model scale, and conditioning strategy.

• On ImageNet-1K class-conditional SiT-B/2, PAFM achieves lower FID and sFID at all neighborhood sizes.
• On CC12M text-to-image, PAFM reduces FID50K from 10.37 (FM) to 9.45.

Augmentation and sampling alternatives (random crops or VAE perturbations) further validate that PAFM's theoretical benefit is agnostic to candidate pool generation, as long as plausible targets are included.

Gradient Variance and Training Stability

Empirical measurements confirm that PAFM dramatically reduces mini-batch gradient variance compared to FM—mean reduction factor of %%%%8$z_t$9%%%% for SiT-B/2 on ImageNet, in line with theoretical predictions, yielding more stable and reliable optimization trajectories.

Implications and Future Directions

The PAFM approach generalizes the FM objective and provides a practical, theoretically-motivated pathway to denser supervision, critical for high-dimensional generative modeling where true likelihood paths are highly underdetermined by single-pairings. This has several implications:

• Training Stability and Generalization: Lower variance and denser gradient signals mitigate flow collapse and promote learning of globalized, generalizable mappings.
• Versatility: PAFM is compatible with arbitrary target selection strategies and is architecturally agnostic, allowing future work to optimize candidate sets for specific downstream tasks, modalities, or data regimes.
• Potential for Broader Application: The posterior-weighted mixture principle could extend to other trajectory-based or path-matching objectives beyond rectified flow and even inform contrastive, representation, and alignment learning objectives.

Open questions include optimal design of importance-weighted candidate pools, adaptation for large-scale video or multimodal data, and integration with other variance reduction or regularization techniques.

Conclusion

Posterior-Augmented Flow Matching addresses the sparse supervision bottleneck in generative flow-based models by aggregating supervision across a posterior-weighted candidate set for each intermediate, rather than a single target. It is mathematically equivalent to FM in expectation but achieves significantly lower gradient variance, enhanced robustness to data sparsity, and empirically superior generation metrics in class- and text-conditioned image synthesis across standard benchmarks—all at marginal computational cost. PAFM is likely to have lasting influence on the design of scalable, stable, and high-fidelity continuous-time generative models.