Flow Matching & Rectified Flows

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Continuous Normalizing Flows
5. The Flow Matching Objective
6. Conditional Flow Matching
7. Rectified Flows
8. Connection to Diffusion Models
9. Advantages Over Diffusion
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Define continuous normalizing flows and the transport equation.
2. Derive the flow matching training objective.
3. Explain why conditional flow matching makes training tractable.
4. Derive rectified flows as straight-line interpolations.
5. Compare flow matching vs diffusion in terms of sample quality and speed.

Notation

• $\psi_t: \mathbb{R}^d \to \mathbb{R}^d$ — flow (diffeomorphism) at time $t$
• $\mathbf{v}_\theta(\mathbf{x}, t)$ — learned velocity field
• $p_t$ — probability density at time $t$ (evolving)
• $p_0 = \mathcal{N}(0, \mathbf{I})$, $p_1 = p_{\text{data}}$

Core Intuition

Instead of adding and removing noise (diffusion), flow matching learns a velocity field that transports samples from noise ($t=0$) to data ($t=1$) along smooth paths. The simplest paths are straight lines between noise and data points. This leads to a simpler objective than diffusion, faster sampling (straighter paths need fewer integration steps), and a unified framework.

Continuous Normalizing Flows

A flow $\psi_t$ is defined by an ODE:

The density evolves according to the continuity equation:

Goal: Find $\mathbf{v}_t$ such that $p_0 = \mathcal{N}(0, \mathbf{I})$ is transported to $p_1 = p_{\text{data}}$.

The Flow Matching Objective

Ideal objective: Match the true velocity field $\mathbf{u}_t$ that generates the target flow:

Problem: We don't know $\mathbf{u}_t$ or $p_t$ (they depend on the optimal transport plan).

Conditional Flow Matching

Solution (Lipman et al., 2023): Condition on individual data points. Define a simple conditional flow for each data point $\mathbf{x}_1$:

This is a straight line from noise $\mathbf{x}_0$ to data $\mathbf{x}_1$.

Conditional velocity:

Training objective (equivalent to FM):

Simplicity: Sample noise $\mathbf{x}_0$, sample data $\mathbf{x}_1$, interpolate to get $\mathbf{x}_t$, predict the velocity (direction from noise to data). That's it.

Rectified Flows

Rectified flow (Liu et al., 2022): The straight-line interpolation $\mathbf{x}_t = (1-t)\mathbf{x}_0 + t\mathbf{x}_1$ defines a "rectified" flow.

Property: If the learned velocity field perfectly matches, then the ODE trajectories are straight lines — requiring only one Euler step for perfect generation.

Reflow procedure: Iteratively straighten trajectories:

1. Train flow matching model.
2. Generate pairs $(\mathbf{x}_0, \mathbf{x}_1)$ by integrating the learned ODE.
3. Retrain on these pairs (now the coupling is better than random).
4. Repeat: trajectories become straighter each iteration.

After 2–3 reflow iterations, 1–4 step generation becomes possible.

Connection to Diffusion Models

Diffusion as a special flow: The VP-SDE diffusion process defines a specific (curved) path:

This is NOT a straight line (nonlinear interpolation). The corresponding velocity involves the score function.

Flow matching uses straight lines (equation 4): simpler, straighter paths → fewer ODE steps needed.

Equivalence: Both can achieve the same final quality, but flow matching reaches it in fewer steps because trajectories are straighter (less curvature → less numerical error per step).

Advantages Over Diffusion

1. Simpler objective: Predict $\mathbf{x}_1 - \mathbf{x}_0$ instead of noise $\boldsymbol{\epsilon}$. No noise schedule needed.

2. Fewer sampling steps: Straight-line flows need 10–20 Euler steps for high quality (vs 50+ for diffusion).

3. No variance schedule tuning: The interpolation $t \in [0,1]$ is the only "schedule."

4. Unified framework: Natural extension to other source distributions (not just Gaussian noise), optimal transport couplings, and conditional generation.

Common Pitfalls

Pitfall 1. Assuming flow matching is always better. For very high-quality samples (FID < 2), well-tuned diffusion with many steps can still match or beat flow matching.

Pitfall 2. Using random coupling between $\mathbf{x}_0$ and $\mathbf{x}_1$. Random pairing creates crossing trajectories that are hard to learn. Optimal transport coupling (mini-batch OT) significantly improves quality.

Pitfall 3. Ignoring the simulation-free property. Flow matching is trained without simulating the ODE — each training step only needs a random $t$ and the interpolated point.

Summary

• Flow matching learns a velocity field transporting noise to data along straight lines.
• Conditional flow matching makes training tractable by conditioning on individual data points.
• Objective: Predict the velocity $\mathbf{x}_1 - \mathbf{x}_0$ at interpolated points.
• Rectified flows iteratively straighten trajectories for few-step generation.
• Simpler than diffusion: no noise schedule, fewer steps, unified framework.
• Used in: Stable Diffusion 3, Flux, modern image/video generation.

Exercises

Exercise 1. Verify that $\psi_t(\mathbf{x}_0|\mathbf{x}_1) = (1-t)\mathbf{x}_0 + t\mathbf{x}_1$ satisfies $\psi_0 = \mathbf{x}_0$ and $\psi_1 = \mathbf{x}_1$.

Exercise 2. Derive the conditional velocity field for the VP-SDE interpolation $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_1 + \sqrt{1-\bar{\alpha}_t}\mathbf{x}_0$.

Exercise 3. Prove that the marginal flow matching objective (equation 3) equals the conditional flow matching objective (equation 6) in expectation.

Exercise 4. Explain why straight-line flows need fewer Euler steps than curved diffusion paths.

Exercise 5. Design a reflow procedure for a trained model and estimate the quality after 1, 2, and 3 reflow iterations.