Flow Matching for Generative Modeling

Chapter 1: Continuous Normalizing Flows

A Continuous Normalizing Flow (CNF) defines a generative model by moving samples along a time-dependent vector field. Think of it as a river: at time t = 0 you drop particles into noise, and the flow carries them to data at t = 1.

The ODE Formulation

A CNF is defined by a vector field v_t : R^d → R^d that generates a flow φ_t via the ODE:

Starting from φ₀(x) = x (identity at t = 0), the flow map φ_t transforms samples from the source distribution p₀ (noise) to the target distribution p₁ (data). At each time t, there is an intermediate distribution p_t = [φ_t]_# p₀ — the pushforward of p₀ through φ_t.

The Continuity Equation

The densities p_t satisfy the continuity equation:

This is a conservation law: probability mass is neither created nor destroyed, it just flows. The vector field v_t generates the probability path p_t. Given v_t, you can compute p_t by solving this PDE (or equivalently, by simulating the ODE and pushing samples forward).

Sampling from a CNF

To generate a sample: draw x₀ ~ p₀ = N(0, I), then solve the ODE forward from t = 0 to t = 1 using an ODE solver (Euler, RK45, etc.):

With a learned v_θ, this gives you a sample from the model's approximation of q (the data distribution).

Chapter 2: The Training Problem

We have a beautiful generative model: learn v_θ, integrate it to generate samples. But how do we train v_θ? This is where CNFs ran into a wall for years.

The Naive Approach: Flow Matching (Intractable)

The most natural objective is to match our learned vector field to the true one that generates the data distribution:

where u_t is the vector field that generates a probability path p_t interpolating between p₀ and the data distribution q.

Problem 1: We don't know u_t. It's defined implicitly as whatever vector field generates p_t via the continuity equation. There's no closed-form expression for the marginal u_t(x).

Problem 2: Even if we could compute u_t, sampling from p_t requires simulating the ODE forward from t = 0 — which means running an ODE solver inside the training loop. This is the same issue that plagued original CNF training (FFJORD): backpropagating through an ODE solver is expensive, memory-intensive, and numerically unstable.

Prior Solutions and Their Limitations

• FFJORD (2019): Train with maximum likelihood via the instantaneous change of variables formula. Requires trace estimation (Hutchinson's) + ODE simulation. Expensive.
• Score matching for diffusion (2020-2021): Sidestep the ODE issue by using a specific noising process (VP, VE SDE) where p_t(x | x₁) is known in closed form. But this locks you into diffusion-style curved paths.

Flow Matching's insight: we can get the simplicity of score matching (closed-form targets, no simulation) but for any probability path — not just diffusion paths.

Chapter 3: Conditional Paths

Here's the key insight that unlocks everything. Instead of thinking about the marginal probability path p_t (which is intractable), we think about conditional probability paths p_t(x | x₁) — the path from noise to a specific data point x₁.

The Conditioning Trick

Define a conditional probability path p_t(x | x₁) that:

• At t = 0: p₀(x | x₁) ≈ p₀(x) = N(0, I)   (starts at noise, approximately independent of x₁)
• At t = 1: p₁(x | x₁) = N(x₁, σ²_min I)   (concentrated around x₁, essentially a delta)

The simplest choice is a Gaussian conditional path:

where μ_t and σ_t are differentiable functions of t with boundary conditions μ₀ ≈ 0, σ₀ ≈ 1, μ₁ = x₁, σ₁ ≈ 0.

The Conditional Vector Field

For any Gaussian conditional path, there exists a unique vector field u_t(x | x₁) that generates it (via the continuity equation). For Gaussian paths, it has a closed-form expression:

This is exact. No approximation. No simulation needed. Given a sample x_t from p_t(x | x₁), we can compute the vector field analytically.

The Reparameterization

Sampling from p_t(x | x₁) is trivial via reparameterization:

For the simplest choice (linear interpolation, a.k.a. Optimal Transport path):

This is just linear interpolation between a noise sample and a data sample! And the conditional vector field becomes:

The target is simply the displacement from noise to data. Constant in time. A straight line.

Chapter 4: The Flow Matching Objective

We now have all the pieces. The conditional vector field u_t(x | x₁) is tractable. But we want to learn the marginal vector field u_t(x) that generates the marginal p_t. How does conditioning on x₁ help?

From Marginal to Conditional

The crucial mathematical result (Theorem 2 in the paper): the marginal probability path can be recovered by marginalizing over the conditional paths:

And the marginal vector field that generates this p_t is:

This is still intractable (it involves an integral over all data points). But here's the punchline:

The Conditional Flow Matching Loss

Define the Conditional Flow Matching (CFM) loss:

The Proof Sketch

The proof uses a simple trick. Expand L_FM:

The cross term ⟨v_θ(t,x), u_t(x)⟩ under E_{x ~ pt} can be rewritten by substituting the marginal expressions. After expanding p_t(x) = ∫ p_t(x|x₁)q(x₁)dx₁ and u_t(x) as its conditional mixture, the cross term becomes E_{x1 ~ q} E_{x ~ pt(x|x1)} ⟨v_θ, u_t(x|x₁)⟩ — exactly the cross term in L_CFM. The squared norms also match up to constants independent of θ.

The Training Algorithm

With OT conditional paths, training is beautifully simple:

# Flow Matching training step (OT path)
def fm_step(model, x1, optimizer):
    # x1: batch of data samples [B, d]
    t = torch.rand(B, 1)           # sample time
    x0 = torch.randn_like(x1)      # sample noise
    xt = (1 - t) * x0 + t * x1     # interpolate
    target = x1 - x0               # conditional vector field
    pred = model(t, xt)             # predict vector field
    loss = (pred - target).pow(2).mean()
    loss.backward()
    optimizer.step()

That's it. Four lines of math. No ODE solver. No score function. No noise schedule to tune.

Chapter 5: Optimal Transport Paths

Flow Matching works with any Gaussian conditional path. But which path should you choose? This is where the Optimal Transport (OT) path shines — and where FM decisively separates from diffusion.

The OT Conditional Path

The OT path uses linear interpolation:

For practical purposes with σ_min → 0:

This is the displacement interpolation from Optimal Transport theory. The sample moves along a straight line from x₀ to x₁ at constant speed.

Why "Optimal Transport"?

In OT theory, the optimal way to transport mass from one distribution to another (minimizing total displacement) is via straight lines. The conditional OT path x_t = (1-t)x₀ + tx₁ moves each particle along the shortest path between its noise sample and its data target. While this is only the true OT map when conditioning on a single x₁, the conditional paths still produce straighter marginal flows than diffusion paths.

Diffusion Path (VP-SDE) for Comparison

The Variance Preserving path used by DDPM:

where α_t = exp(−½ ∫₀^t β(s) ds) with some noise schedule β(t). This produces curved trajectories: the particle first drifts toward the origin (as noise is added) then curves toward x₁ (as signal emerges).

Why Straight Lines Win

• Easier to learn: The vector field for straight lines is nearly constant (x₁ - x₀). For curved paths, the vector field changes direction throughout, requiring the network to learn complex time-dependent behavior.
• Fewer integration steps: Straight trajectories have low curvature → Euler's method with large steps is accurate → fewer NFEs at inference (10 vs 100+).
• Better conditioning: The neural network sees the same target (x₁ - x₀) regardless of t, reducing variance in gradient estimates.

Chapter 6: Connection to Diffusion

Flow Matching doesn't replace diffusion — it generalizes it. Score matching for diffusion models is a special case of Conditional Flow Matching when you choose the diffusion probability path.

Score Matching as a Special Case

Recall: in diffusion models, we have a noising process:

The score function is ∇_x log p_t(x | x₁) = −(x_t − α_t x₁) / σ_t² = −ε / σ_t.

Denoising score matching trains: predict ε from x_t.

Now, the conditional vector field for the diffusion (VP) path is:

Substituting x = α_t x₁ + σ_t ε:

This is a linear combination of ε and x₁ — exactly what diffusion models predict (up to reparameterization). Training a network to predict this vector field is equivalent to training it to predict ε or x₁.

What FM Adds Beyond Diffusion

Relationship to Rectified Flows

Rectified Flows (Liu et al., 2022) independently proposed the same OT conditional path x_t = (1-t)x₀ + tx₁ and the same training objective. The key difference: Rectified Flows additionally proposed reflow — iteratively straightening trajectories by using the model's own predictions as new training pairs. Flow Matching provides the theoretical foundation (the equivalence theorem) while Rectified Flows provide the practical refinement technique.

Chapter 7: Results & Applications

Flow Matching isn't just a theoretical curiosity — it achieves state-of-the-art results and has become the foundation for modern generative models.

ImageNet Results

On ImageNet 64×64 (unconditional), comparing different probability paths with the same architecture:

*Score SDE uses 2000 NFEs and Predictor-Corrector sampling. OT-CFM uses only 110 NFEs with an adaptive ODE solver.

Training Stability

The paper shows that FM with OT paths has significantly more stable training dynamics compared to VP paths:

• Loss variance is lower (the target x₁ - x₀ is constant across time, reducing gradient noise)
• No need to tune noise schedules (linear, cosine, etc.) — there is no schedule
• The learned vector field is smoother (lower Lipschitz constant), making ODE integration more stable

Downstream Impact

Flow Matching has become the standard training paradigm for modern generative models:

• Stable Diffusion 3 / Flux (2024): Use rectified flow (= OT-CFM) as the training objective. The "diffusion" in the name is legacy — the math is Flow Matching.
• Speech synthesis (Voicebox, 2023): Flow Matching for text-to-speech, achieving state-of-the-art naturalness. The dMel decoder uses FM to generate mel spectrograms.
• Video generation (MovieGen, 2024): FM for text-to-video, leveraging the fast inference (fewer steps = less latency for video frames).
• Protein design (FrameFlow, 2023): FM on SE(3) for protein backbone generation.
• Robotics (pi0, 2024): Flow Matching as the action generation head in Vision-Language-Action models.

Chapter 8: Connections

Flow Matching sits at the intersection of several major ideas in generative modeling and beyond.

Related Papers in This Collection

• Diffusion Policy (Chi et al., 2023): Uses DDPM for robot action generation — the precursor to flow-based policies. FM provides the faster alternative. Read lesson →
• DiT (Peebles & Xie, 2023): Diffusion Transformers for image generation. Modern DiTs (SD3) use FM instead of DDPM training. Read lesson →
• Stable Diffusion 3 (Esser et al., 2024): The practical realization of FM + DiT at scale. Read lesson →
• pi0 (Black et al., 2024): Flow Matching as the action head in a VLA — the direct application of this paper to robotics. Read lesson →

The Intellectual Lineage

Key Takeaways

References

1. Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., Le, M. "Flow Matching for Generative Modeling." ICLR 2023. arXiv:2210.02747
2. Liu, X., Gong, C., Liu, Q. "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow." ICLR 2023. arXiv:2209.03003
3. Chen, R. T. Q., Rubanova, Y., Bettencourt, J., Duvenaud, D. "Neural Ordinary Differential Equations." NeurIPS 2018. arXiv:1806.07366
4. Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., Poole, B. "Score-Based Generative Modeling through Stochastic Differential Equations." ICLR 2021. arXiv:2011.13456
5. Esser, P. et al. "Scaling Rectified Flow Transformers for High-Resolution Image Synthesis." ICML 2024. arXiv:2403.03206