Flow Matching for Generative Modeling

Abstract

We introduce a new paradigm for generative modeling built on Continuous Normalizing Flows (CNFs), allowing us to train CNFs at unprecedented scale. Specifically, we present the notion of Flow Matching (FM), a simulation-free approach for training CNFs based on regressing vector fields of fixed conditional probability paths. Flow Matching is compatible with a general family of Gaussian probability paths for transforming between noise and data samples — which subsumes existing diffusion paths as special instances. Furthermore, Flow Matching opens the door to training CNFs with other, non-diffusion probability paths. An instance of particular interest is using Optimal Transport (OT) displacement interpolation to define the conditional probability paths.

Summary

Flow Matching introduces a simulation-free framework for training Continuous Normalizing Flows by directly regressing the velocity field of a target probability path. The key insight is that the intractable marginal Flow Matching objective ${\mathcal{L}}_{\text{FM}}={\mathbb{E}}_{t,p_t(x)}\Vert v_t(x)-u_t(x){\Vert }^2$ can be replaced by the tractable Conditional Flow Matching (CFM) objective ${\mathcal{L}}_{\text{CFM}}={\mathbb{E}}_{t,q(x_1),p_t(x|x_1)}\Vert v_t(x)-u_t(x|x_1){\Vert }^2$, which has identical gradients (Theorem 2). This allows training without ODE simulation, using only per-sample conditional probability paths.

The paper defines a general family of Gaussian conditional probability paths $p_t(x|x_1)=\mathcal{N}(x|{\mu }_t(x_1),{\sigma }_t(x_1{)}^{2}I)$ with corresponding conditional velocity fields $u_t(x|x_1)=\frac{{\sigma }_t^{\prime }}{{\sigma }_t}(x-{\mu }_t)+{\mu }_t^{\prime }$ (Theorem 3). Two key instances are: (1) Diffusion paths recovering VE/VP-SDE probability paths; (2) OT paths with ${\mu }_t=t{x}_1$ and ${\sigma }_t=1-(1-{\sigma }_{\min })t$, producing straight-line trajectories with constant-direction velocity fields. The OT conditional velocity $u_t(x|x_1)=(x_1-(1-{\sigma }_{\min })x)/(1-(1-{\sigma }_{\min })t)$ is time-constant in direction, making it simpler to learn than the time-varying diffusion score function.

FM with OT paths achieves state-of-the-art results: FID 2.99 on CIFAR-10, FID 5.02 on ImageNet-32, and FID 20.9 on ImageNet-128, outperforming score matching and DDPM while requiring ~60% fewer function evaluations.

Key Contributions

• Conditional Flow Matching objective providing unbiased, simulation-free training of CNFs (Theorem 2)
• General family of Gaussian conditional probability paths parameterized by ${\mu }_t,{\sigma }_t$ (Theorem 3)
• OT displacement interpolation as a conditional probability path, producing straight trajectories
• Subsumes diffusion paths (VE, VP) as special cases while enabling non-diffusion paths
• Empirical demonstration that OT paths are faster to train, faster to sample, and produce better samples than diffusion paths

Methodology

Training minimizes ${\mathcal{L}}_{\text{CFM}}={\mathbb{E}}_{t,q(x_1),p(x_0)}\Vert v_{\theta }({\psi }_t(x_0))-\frac{d}{dt}{\psi }_t(x_0){\Vert }^2$ where ${\psi }_t(x)={\sigma }_{t}x+{\mu }_t$ is the conditional flow map. For OT paths: ${\psi }_t(x)=(1-(1-{\sigma }_{\min })t)x+t{x}_1$, giving loss $\mathbb{E}\Vert v_{\theta }({\psi }_t(x_0))-(x_1-(1-{\sigma }_{\min })x_0){\Vert }^2$. Sampling via ODE solver (dopri5) with tolerances 1e-5.

Key Findings

• OT paths produce constant-direction conditional VFs — simpler regression target than time-varying diffusion scores
• FM with diffusion paths is more stable than score matching despite equivalent objectives
• OT paths require ~60% fewer NFEs to reach the same error threshold as diffusion paths
• The conditional OT flow is the displacement map between two Gaussians — optimal per-sample but the marginal VF is NOT the global OT map

Important References

1. Score-Based Generative Modeling through Stochastic Differential Equations — SDE framework whose probability paths are subsumed as special cases
2. Flow Straight and Fast — Concurrent work (rectified flow) arriving at similar linear interpolation objectives
3. Building Normalizing Flows with Stochastic Interpolants — Concurrent work by Albergo & Vanden-Eijnden with a similar framework

Atomic Notes

• conditional flow matching
• conditional probability path

---

paper