Understand Flow
Matching

Chapter 1: Continuous Normalizing Flows

In Chapter 0, we showed particles flowing from noise to data along paths. That's a powerful visual intuition — but to actually build this, we need to answer a precise question: how do we tell a particle where to go at each moment?

The answer is a velocity field. At every point in space and every moment in time, the velocity field says "move in this direction at this speed." If you've ever seen a weather map with wind arrows, that's a velocity field — each arrow tells air molecules where to flow. Our velocity field does the same thing for data particles.

The math tool for "follow a velocity field" is a differential equation — specifically, an ordinary differential equation (ODE). An ODE just says: "the rate of change of x equals some function of x and time." If you've done basic calculus, you've seen this as dx/dt = f(x, t). That's all an ODE is — a rule for how something changes over time.

This gives us the mathematical backbone of flow matching, called a Continuous Normalizing Flow (CNF):

Read this as: "starting from a noise sample x(0), follow the velocity field v from time t=0 to t=1, and you arrive at a data sample x(1)." The velocity field v_θ is parameterized by a neural network — the θ are the learnable weights. Training teaches the network what velocity to assign at each (x, t) so that the flow transforms noise into realistic data.

The network takes in x_t (the particle's current position, shape [B, C, H, W] — a batch of images) and the scalar time t, and outputs the velocity v (same shape [B, C, H, W]). Same input/output shape as a diffusion denoiser — but instead of predicting noise, it predicts direction of motion.

Chapter 2: Optimal Transport

We now have a framework: a velocity field v(x, t) that guides particles from noise to data. But we haven't said anything about which paths those particles should take. A velocity field could send particles along wild, looping curves — or along straight lines. Both get from noise to data, but one is far more efficient.

Optimal transport (OT) is the math of "what's the cheapest way to move mass from one pile to another?" Think of it like a shipping problem: you have warehouses (noise) and stores (data), and you want to minimize total distance traveled. The answer, for our case, turns out to be beautifully simple: straight paths with constant velocity.

Why does straightness matter? Because straighter paths can be integrated with fewer numerical steps. A perfectly straight path needs just one Euler step. Curved paths need many steps to follow accurately.

Why Straight Paths Need Fewer Steps

Consider Euler integration: x_t+dt = x_t + v(x_t, t) · dt. If the true velocity is constant (straight path), then a single Euler step with dt=1 gives the exact answer. No numerical error at all. In practice, the velocity field isn't perfectly constant (different noise-data pairs have different velocities), so we need a few steps. But "a few" means 10-50, not 1000.

Contrast with diffusion: the variance-preserving path curves through high-dimensional space. The velocity changes direction at every point, so you need many small steps to track the curve accurately. This is the fundamental geometric reason flow matching is faster.

Chapter 3: Conditional Flow Matching

We know the ideal path for each particle: a straight line from its specific noise point z to its specific data point x₁. But here's the problem: during generation, the network sees an intermediate point x_t and has to decide where to send it — without knowing which x₁ it's headed toward. The "true" velocity at x_t would require averaging over every possible data point it might be going to. That's impossibly expensive to compute (mathematicians call this intractable — meaning "can't be computed in practice").

Conditional Flow Matching (CFM) is the elegant trick that makes training possible anyway. Instead of computing that impossible average, we train on individual pairs. For each training pair (noise z, data x₁), we know exactly where the particle should go: from z to x₁ in a straight line. The conditional velocity for this specific pair is simply u_t = x₁ - z (the direction from noise to data, constant along the path).

The beautiful mathematical insight (proven by Lipman et al., 2023): if you train the network to match these per-pair velocities across many random pairs, the network automatically learns the correct average velocity at every point. You never have to compute the intractable average explicitly — gradient descent finds it for you.

Here's the concrete difference this makes. Without CFM, the target velocity at a point x_t would require averaging over all possible noise-data pairs that pass through x_t — an intractable integral. With CFM, for each training sample we just compute x₁ - z (a single subtraction). The magic is that gradient descent over many such samples converges to the correct marginal velocity field.

Chapter 4: Training

Training a flow matching model is arguably even simpler than training a diffusion model. For each training step:

1. Sample noise z ~ N(0, I)
2. Sample a data point x₁ from the dataset
3. Sample a random time t ~ U(0, 1)
4. Compute x_t = (1-t)z + t x₁
5. Predict velocity: v̂ = v_θ(x_t, t)
6. Loss = || v̂ - (x₁ - z) ||²

The Complete Training Loop

Here it is in full. Compare this to the diffusion training loop — notice what's missing:

python
for x_1 in dataloader:                     # 1. Sample data
    x_0 = torch.randn_like(x_1)           # 2. Sample noise ~ N(0, I)
    t = torch.rand(B, 1, 1, 1)            # 3. Sample t ~ Uniform[0, 1]

    x_t = (1 - t) * x_0 + t * x_1          # 4. Linear interpolation
    target = x_1 - x_0                     # 5. Target velocity (constant!)

    v_hat = network(x_t, t.squeeze())      # 6. Predict velocity
    loss = F.mse_loss(v_hat, target)       # 7. MSE loss

    loss.backward()
    optimizer.step()

python
def cfm_loss(model, x_1, noise_fn=torch.randn_like):
    B = x_1.shape[0]

    # 1. Sample noise (same shape as data)
    z = noise_fn(x_1)                      # [B, C, H, W]

    # 2. Sample time uniformly in [0, 1]
    t = torch.rand(B, 1, 1, 1, device=x_1.device)  # [B,1,1,1] for broadcasting

    # 3. Linear interpolation: x_t = (1-t)*z + t*x_1
    x_t = (1 - t) * z + t * x_1             # [B, C, H, W]

    # 4. Target velocity = endpoint - startpoint
    target = x_1 - z                         # [B, C, H, W]

    # 5. Network predicts velocity at (x_t, t)
    v_pred = model(x_t, t.squeeze())        # [B, C, H, W]

    # 6. MSE between predicted and target velocity
    return F.mse_loss(v_pred, target)

Chapter 5: Sampling

To generate a sample, start from noise z ~ N(0, I) and integrate the learned ODE forward from t=0 to t=1. The simplest method is Euler integration: take N evenly spaced steps, at each step nudging x by the predicted velocity times the step size.

Because the paths are approximately straight, Euler's method works well even with few steps. Higher-order solvers (Midpoint, RK4) give better accuracy for the same step count, but even plain Euler with 20-50 steps produces excellent results.

The Sampling Loop

python
# Generate an image from pure noise
x = torch.randn(1, 4, 64, 64)  # start from noise (latent space)
N = 20                             # number of Euler steps
dt = 1.0 / N

for i in range(N):
    t = i / N                     # current time: 0, 0.05, 0.10, ...
    v = network(x, t)              # predict velocity at current position
    x = x + v * dt                 # take one Euler step
# x is now a sample from the data distribution

Compare this to DDPM's 1000 steps with stochastic noise injection at each step. Flow matching sampling is an ODE (deterministic), not an SDE (stochastic). Same noise in → same image out. Every time.

Chapter 6: Reflow & Distillation

Even though flow matching paths are straighter than diffusion, they're not perfectly straight. Paths from different noise-data pairs can cross each other, forcing the network to learn a curved velocity field to avoid collisions. Reflow straightens paths further.

The procedure is elegant: (1) Take a trained flow model. (2) Sample noise z ~ N(0, I), run the model forward to get the corresponding data point x₁ = ODE(z). Now you have a (z, x₁) pair that the model actually maps. (3) Retrain a new model on straight lines between these pairs. Because the pairs are already coupled (the model mapped z to x₁), the straight-line approximation is much closer to the true path.

python
# Reflow: straighten paths by retraining on coupled pairs
# Step 1: Generate coupled (noise, data) pairs from trained model
z = torch.randn(10000, 4, 64, 64)
x_1 = ode_solve(old_model, z, t=0, t_end=1)  # run the trained model
pairs = list(zip(z, x_1))                        # store the pairs

# Step 2: Retrain on straight lines between these pairs
for z_i, x_1_i in pairs:
    t = torch.rand(1)
    x_t = (1-t) * z_i + t * x_1_i        # straight line between COUPLED pair
    loss = mse(new_model(x_t, t), x_1_i - z_i)   # velocity = x_1 - z

Each iteration makes paths straighter. After 2-3 reflow iterations, paths are nearly straight enough for 1-step generation.

Distillation

Distillation takes a different approach: train a student model to mimic the teacher in fewer steps. Progressive distillation halves the step count repeatedly (64 → 32 → 16 → 8 → 4 → 2 → 1). Combined with reflow, this yields high-quality 1-4 step models.

Chapter 7: SD3 / Flux

Flow matching has gone from theory to production. Both Stable Diffusion 3 (Stability AI) and Flux (Black Forest Labs) use rectified flow matching as their core framework, combined with a new architecture: the MMDiT (Multimodal Diffusion Transformer).

MMDiT Architecture

MMDiT replaces the U-Net with a Transformer. Both the noisy latent patches and the text tokens are processed as separate streams that interact through joint attention layers. This bidirectional interaction gives the text genuine influence over image generation.

From Theory to Production Numbers

Flux.1-schnell generates a 1024×1024 image in 4 steps. Let that sink in. DDPM needed 1000 steps for 256×256. That's a 250× reduction in steps at 16× higher resolution. The combination of flow matching (straight paths), rectified flow (straightened further), and distillation (student mimics teacher in fewer steps) made this possible.

Chapter 8: Flow vs Diffusion

Flow matching and diffusion are closely related — in fact, diffusion can be seen as a special case of flow matching with a particular (non-straight) path choice. But the differences matter in practice.

Why Flow Matching Is Winning

Four concrete reasons flow matching is becoming the default for new systems:

1. Simpler math. No noise schedule design (β_t linear vs cosine vs sigmoid). No ᾱ_t cumulative products. Just linear interpolation.
2. Faster sampling. Straight paths need 10-50 Euler steps instead of DDPM's 1000. Even compared to DDIM (50 steps), flow matching matches quality with fewer steps.
3. Deterministic by default. ODE, not SDE. Same seed = same output. Reproducibility is free.
4. This is what production uses. Stable Diffusion 3, Flux, π0 (robot policy) — all flow matching. The industry has voted.

Conditional Flow Matching

Conditioning works exactly as in diffusion: inject class labels or text embeddings into the network. The velocity field becomes v(x_t, t, c) where c is the conditioning signal. Classifier-free guidance applies identically — train with condition dropout, at inference compute v_guided = v_uncond + w · (v_cond - v_uncond).

When to Use Which?