Iterative Flow Matching -- Path Correction and Gradual Refinement for Enhanced Generative Modeling

TL;DR

This work analyzes flow matching as a transport-based approach for generative modeling and identifies why standard flow matching can yield hallucinations due to trajectory interpolation artifacts. It introduces two iterative refinements—end-path correction and gradual refinement—to progressively align generated samples with the target distribution $\pi_T$, with theoretical support via Wasserstein and KL-based bounds and practical validation on MNIST and CIFAR-10 using latent-space FM with RBF interpolation. The proposed framework can be integrated into diverse generative pipelines to improve robustness and sample fidelity at the cost of extra computation. Overall, iterative flow matching offers a principled, flexible path to tighten the gap between the learned and target distributions and reduce out-of-distribution samples in high-dimensional synthesis tasks.

Abstract

Generative models for image generation are now commonly used for a wide variety of applications, ranging from guided image generation for entertainment to solving inverse problems. Nonetheless, training a generator is a non-trivial feat that requires fine-tuning and can lead to so-called hallucinations, that is, the generation of images that are unrealistic. In this work, we explore image generation using flow matching. We explain and demonstrate why flow matching can generate hallucinations, and propose an iterative process to improve the generation process. Our iterative process can be integrated into virtually $\textit{any}$ generative modeling technique, thereby enhancing the performance and robustness of image synthesis systems.

Figures (8)

• Figure 1: Schematic of the standard flow matching technique (shown inside the dashed box) vs. our proposed iterative approaches, end-path correction (black trajectory), and gradual refinement (red trajectory). $\pi_0$ and $\pi_T$ represent the source and target distributions, respectively. A trajectory between two distributions represents a learned mapping parametrized by some $\theta$. The intermediate distributions represent the pushforward distributions obtained via the integration of the ODE in Equation \ref{['eq:ODE']} using the learned mappings $\theta$.
• Figure 1: Using flow matching to move the distribution of blue points to the distribution of red points. The system is trained on the magenta trajectories obtained by linear combinations of points from both mixtures. The integrated data is plotted in cyan, and its trajectories are plotted in black.
• Figure 1: Using Algorithm \ref{['alg:eppr']} to correct the final distribution obtained from the standard flow matching optimization.
• Figure 1: Generated MNIST samples from different iterations.
• Figure 2: Using prediction-correction flow matching to move between two Gaussian mixtures.