Posterior Augmented Flow Matching

1. Problem Statement: Sparse Supervision and Flow Collapse

The paper addresses a fundamental limitation in Flow Matching (FM), a class of continuous-time generative models that learn a time-dependent vector field to transport samples from a simple prior (e.g., Gaussian noise) to a complex data distribution.

• The Core Issue: In standard FM, training relies on sparse, one-to-one supervision. For any given intermediate state $z_t$ along a trajectory connecting a noise sample $\epsilon$ and a target data point $z$, the model is supervised to predict the velocity toward that single specific target $z$.
• High-Dimensional Reality: In high-dimensional spaces (like images), a single intermediate point $z_t$ could validly lie on many different trajectories leading to diverse, valid target images that satisfy the same condition (e.g., "a photo of a dog").
• Consequences:
  • High Variance: The gradient at any specific $z_t$ is estimated from a single noisy signal.
  • Flow Collapse: The model tends to memorize specific source-target pairings rather than learning the true continuous mapping. This leads to "flow collapse," where diverse inputs are mapped to overly similar or muddled outputs, failing to generalize.

2. Methodology: Posterior-Augmented Flow Matching (PAFM)

The authors propose PAFM, a theoretically grounded generalization of FM that replaces single-target supervision with an expectation over the posterior distribution of valid target completions.

Theoretical Formulation

Instead of minimizing the loss against a single target $z$, PAFM minimizes the expected loss over a set of plausible targets $\{z_j\}$ drawn from the posterior distribution $p(z_j | z_t, y)$, where $y$ is the condition.

• Objective: The new loss function aggregates gradients from multiple candidate targets, weighted by their posterior likelihood.
• Unbiased Estimator: The authors prove (Theorem 1) that PAFM is an unbiased estimator of the original FM objective. In expectation, $L_{PAFM} = L_{FM}$.
• Variance Reduction: Crucially, Theorem 2 proves that by aggregating information from $K > 1$ plausible trajectories per intermediate point, PAFM lower-bounds the gradient variance compared to standard FM. The variance reduction factor is related to the Effective Sample Size (ESS) of the sampling scheme.

Practical Implementation

Since the true posterior is intractable, PAFM factorizes the problem into two approximations using Importance Sampling:

1. Likelihood of the Intermediate: $p(z_t | z_j)$: How plausible is it that target $z_j$ generated the current intermediate $z_t$? This is computed using the known conditional probability path (Gaussian).
2. Condition Likelihood: $p(y | z_j)$: How well does target $z_j$ match the condition $y$?
   • Class-conditioned: Deterministic (1 if labels match, 0 otherwise).
   • Text-conditioned: Approximated using vision-language models (e.g., CLIP/SigLIP) or assumed uniform if candidates are pre-filtered.

Training Algorithm:
For each training step, the model samples an intermediate $z_t$. Instead of one target, it retrieves a set of $K$ candidate targets $\{z_j\}$. The model predicts the velocity toward a weighted average of these targets, where weights are proportional to $p(z_t | z_j) \cdot p(y | z_j)$.

3. Target Selection Strategies

The paper explores several practical strategies to construct the candidate set $\{z_j\}$:

1. Latent Space Nearest Neighbors (Primary): Using FAISS to retrieve $K$ nearest neighbors of the training image in the latent space. This ensures geometric proximity and semantic similarity.
2. Random Augmentations: Creating targets by applying random crops, resizes, and rotations to the source image. This increases spatial diversity.
3. VAE Moment Sampling: Sampling $K$ times from the Gaussian distribution (moment) of the VAE encoding of the source image. This provides small perturbations around the mean encoding.

4. Key Contributions

• Theoretical Generalization: PAFM is proven to be an unbiased generalization of FM that reduces gradient variance by aggregating supervision over the full posterior of compatible targets.
• Flow Collapse Mitigation: By providing a denser supervision signal, PAFM prevents the model from memorizing specific pairings, leading to more robust vector fields.
• Flexibility: The framework is agnostic to how candidates are selected, allowing integration with various data modalities and architectures.
• Efficiency: The method introduces negligible computational overhead because candidate sets can be pre-computed (e.g., via FAISS indices) and loaded during training.

5. Experimental Results

The authors evaluated PAFM across different scales, architectures, and datasets:

• Datasets:
  • ImageNet-1K (Class-conditioned): Tested on SiT-B/2 and SiT-XL/2 models.
  • CC12M (Text-to-Image): Tested on MMDiT architecture.
• Metrics: FID50K (Fréchet Inception Distance), sFID, Precision, Recall.
• Performance Gains:
  • ImageNet: PAFM improved FID by up to 3.4 points (e.g., SiT-XL/2 improved from 11.14 to 9.85 FID).
  • CC12M: PAFM reduced FID from 10.37 (FM) to 9.45.
  • Consistency: Improvements were observed across all model scales (SiT-B/2, SiT-XL/2) and architectures (SiT, MMDiT).
• Gradient Variance: Empirical analysis showed PAFM reduced mini-batch gradient variance by approximately 4x compared to FM, leading to more stable training dynamics.
• Computational Cost:
  • Throughput decreased by only 6.6%.
  • Peak memory usage increased by 0.4%.
  • FLOPs per step remained unchanged.

6. Significance

PAFM represents a significant step forward in generative modeling by addressing the "sparse supervision" bottleneck inherent in flow-based and diffusion models.

• Robustness: It fundamentally improves the stability of training high-dimensional generative models, reducing the risk of mode collapse and artifacts.
• Efficiency: It achieves state-of-the-art performance improvements with minimal additional compute, making it a highly practical drop-in replacement for standard Flow Matching objectives.
• Generalizability: The approach is applicable to both class-conditioned and text-to-image generation, suggesting it could benefit a wide range of continuous generative models beyond the specific architectures tested.