Flow Matching: Theory & Training

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Flow Matching Framework
5. Probability Paths
6. The Marginal Vector Field
7. Conditional Flow Matching (CFM)
8. Gaussian Probability Paths
9. Optimal Transport Path
10. Connection to Score Matching
11. Training Algorithm
12. Common Pitfalls
13. Summary
14. Exercises

Learning Objectives

1. Define a probability path and the associated vector field.
2. Derive the flow matching objective from the continuity equation.
3. Prove that conditional flow matching (CFM) has the same gradient as marginal FM.
4. Construct Gaussian and OT probability paths.
5. Connect flow matching to score-based diffusion models.

Notation

• $p_t$ — probability density at time $t \in [0, 1]$
• $u_t(\mathbf{x})$ — vector field generating $p_t$
• $\psi_t$ — flow map: $d\psi_t/dt = u_t(\psi_t)$
• $p_t(\mathbf{x}|z)$ — conditional path given endpoint $z$

Core Intuition

We want to train a neural ODE $v_\theta(\mathbf{x}, t)$ that transports noise $p_0 = \mathcal{N}(0, \mathbf{I})$ to data $p_1 = p_{\text{data}}$. The naive approach (simulate ODE during training) is expensive. Flow matching avoids simulation entirely: define a simple interpolation path between noise and data, compute the velocity along this path analytically, and regress the neural network to match it.

The Flow Matching Framework

Goal: Find $v_\theta$ such that the ODE $dx/dt = v_\theta(x, t)$ generates a flow from $p_0$ to $p_1$.

Continuity equation: Any valid probability path $p_t$ satisfies:

If we can find a vector field $u_t$ that generates our desired path, we can train $v_\theta$ to match it:

Problem: We don't know $p_t$ or $u_t$ for the marginal (unconditional) path.

Probability Paths

A probability path is a time-indexed family of distributions $\{p_t\}_{t \in [0,1]}$ with:

• $p_0 = \mathcal{N}(0, \mathbf{I})$ (noise).
• $p_1 = p_{\text{data}}$ (data).

Examples:

• Linear interpolation (OT path): $p_t = [(1-t)\text{Id} + t\text{T}]_\# p_0$ where $T$ is the OT map.
• VP-SDE path: $p_t(\mathbf{x}|x_1) = \mathcal{N}(\sqrt{\bar{\alpha}_t}x_1, (1-\bar{\alpha}_t)\mathbf{I})$.
• Gaussian path: $p_t(\mathbf{x}|x_1) = \mathcal{N}(\mu_t(x_1), \sigma_t^2\mathbf{I})$.

The Marginal Vector Field

Given a conditional vector field $u_t(\mathbf{x}|z)$ that generates the conditional path $p_t(\mathbf{x}|z)$, the marginal vector field is:

This marginal field is intractable to compute directly — but we don't need it for training.

Conditional Flow Matching (CFM)

Key theorem (Lipman et al., 2023): The conditional flow matching loss:

has the same gradients with respect to $\theta$ as the marginal flow matching loss $\mathcal{L}_{\text{FM}}$.

Proof sketch: The cross-term $\langle v_\theta, u_t \rangle$ integrated over $p_t(\mathbf{x}|z)q(z)$ equals the same cross-term integrated over the marginal $p_t(\mathbf{x})u_t(\mathbf{x})$ (by definition of the marginal vector field). The $\|v_\theta\|^2$ term is the same (both average over $p_t$). The $\|u_t\|^2$ term differs but is independent of $\theta$.

Consequence: We can train using only conditional paths (easy to sample from) instead of the intractable marginal path.

Gaussian Probability Paths

The most common choice: conditional Gaussian path from noise to a data point $x_1$:

with boundary conditions $\mu_0 = 0, \sigma_0 = 1$ (noise) and $\mu_1 = x_1, \sigma_1 \approx 0$ (data).

Linear interpolation (simplest):

Sampling: $\mathbf{x}_t = tx_1 + (1-t)\epsilon, \quad \epsilon \sim \mathcal{N}(0, \mathbf{I})$.

Conditional velocity:

More directly for the linear path:

Training loss:

Optimal Transport Path

Random coupling ($x_0, x_1$ independent): Straight lines may cross → curved velocity field needed.

OT coupling: Match $x_0$ to $x_1$ via optimal transport → non-crossing paths → straighter trajectories.

Mini-batch OT: Within each batch, solve a linear assignment:

Cost: $O(B^3)$ per batch (Hungarian algorithm) or $O(B^2)$ (Sinkhorn).

Effect: 2-3x fewer sampling steps needed; better FID at few steps.

Connection to Score Matching

Score matching (diffusion): Predict noise $\epsilon_\theta(\mathbf{x}_t, t) \approx \epsilon$.

Flow matching: Predict velocity $v_\theta(\mathbf{x}_t, t) \approx x_1 - x_0$.

Relationship: For the VP-SDE path:

So $v_\theta$ and $\epsilon_\theta$ are linear reparameterizations of each other. Same model, different parameterization.

Advantage of velocity parameterization: More natural for straight-line paths (OT); velocity is bounded while score $\nabla\log p_t$ can diverge.

Training Algorithm

For each training step:
    1. Sample x_1 ~ p_data (data batch)
    2. Sample x_0 ~ N(0, I) (noise batch)
    3. [Optional] Apply mini-batch OT to pair (x_0, x_1)
    4. Sample t ~ U[0, 1]
    5. Compute x_t = (1-t)*x_0 + t*x_1
    6. Compute target velocity: u = x_1 - x_0
    7. Loss = ||v_theta(x_t, t) - u||^2
    8. Gradient step on theta

Sampling (inference):

• Solve ODE: $dx/dt = v_\theta(x, t)$ from $t=0$ to $t=1$.
• Use Euler, midpoint, or adaptive solver.
• Fewer steps than diffusion (5-50 vs 20-1000).

Common Pitfalls

Pitfall 1. Using random coupling without OT for few-step generation. Without OT, paths cross and the velocity field becomes multi-valued at intermediate times — requiring many steps to resolve.

Pitfall 2. Confusing the conditional and marginal objectives. CFM trains on conditional velocities but the model learns the MARGINAL velocity (which correctly generates the marginal $p_1$).

Pitfall 3. Setting $\sigma_1 = 0$ exactly. This makes the conditional path degenerate at $t=1$. Use $\sigma_1 = 10^{-5}$ or stop integration slightly before $t=1$.

Summary

• Flow matching trains continuous flows without simulating ODEs during training.
• CFM uses conditional paths (easy to sample) with identical gradients to the marginal objective.
• Linear Gaussian path: $x_t = (1-t)x_0 + tx_1$; velocity target: $x_1 - x_0$.
• OT coupling straightens trajectories for faster sampling.
• Equivalent to diffusion under different path/parameterization choices.
• Used in Stable Diffusion 3, Flux, DALL-E 3.

Exercises

Exercise 1. Prove that CFM and marginal FM have the same gradient (expand both losses and show the $\nabla_\theta$ terms match).

Exercise 2. For the linear path $x_t = (1-t)x_0 + tx_1$: derive the time-dependent SNR and compare to the VP-SDE schedule.

Exercise 3. Implement mini-batch OT for a batch of 64 samples in 2D. Compare the resulting trajectories (with/without OT) visually.

Exercise 4. Show that for the VP-SDE path, the flow matching velocity is $v_t = \dot{\alpha}_t x_1/\alpha_t - \dot{\sigma}_t\sigma_t\nabla\log p_t$.

Exercise 5. Compute the number of Euler steps needed to achieve FID less than 10 on CIFAR-10 for linear FM vs VP-SDE FM.