12.7.1 Flow Matching

Overview

DDPM trains a network to undo a noise schedule and samples by reversing that schedule step by step — a curved, stochastic walk through a thousand small denoisings. It works, and the Lesson 12.3.1 lessons go through the details, but the curve is the problem. Curved paths through high-dimensional space need a lot of small steps to follow accurately. Cut the step count and the trajectory drifts off the curve into regions the model can’t recover from.

Flow Matching, formalised by Lipman et al. in 2022 [1] and independently by Liu, Gong and Liu’s Rectified Flow paper [2], asks a deliberately naive question: what if the path from noise to data were a straight line? Define x_t = (1 - t) · x_0 + t · x_1 where x_0 is noise and x_1 is data. The derivative — the velocity along that line — is simply x_1 - x_0. Train a neural network to predict the velocity at every point in the interpolation, sample by integrating that velocity field with Euler steps from t = 0 to t = 1. With straight lines, fifty Euler steps is enough to land cleanly. With DDPM’s curves, fifty steps misses by a country mile.

This shift — same data, same U-Net, fundamentally different objective and sampling procedure — is what powers FLUX.1, Stable Diffusion 3, and most of the current generation of high-end text-to-image systems [3].

Learning objectives

1. Explain why Flow Matching needs fewer sampling steps than DDPM, in terms of trajectory curvature.
2. Derive the Conditional Flow Matching training loss in five lines: sample x_0, sample t, interpolate, target velocity = x_1 - x_0, MSE.
3. Implement an Euler ODE sampler that integrates a learned velocity field from noise to data.
4. Predict the quality vs. step-count trade-off (5, 10, 20, 50 steps) and recognise the regime where Flow Matching beats DDPM.

Quick start — load the trained model, integrate 50 Euler steps

Download the checkpoint from the GitHub Release. Sampling is a fifteen-line loop: noise in, 50 velocity-field updates, fabric pattern out.

import torch
from flow_model import SimpleFlowUNet

model = SimpleFlowUNet(in_channels=3, base_channels=64)
ckpt = torch.load('flow_matching_fabrics_latest.pt', map_location='cpu')
model.load_state_dict(ckpt['model_state_dict'])
model.eval()

x = torch.randn(4, 3, 64, 64)     # pure noise at t = 0
dt = 1.0 / 50

with torch.no_grad():
    for i in range(50):
        t = torch.full((4,), i / 50)
        v = model(x, t)            # predicted velocity at (x, t)
        x = x + dt * v             # Euler step

x = x.clamp(-1, 1)                 # fabric patterns

That’s the whole sampling procedure. No noise schedule, no posterior-variance bookkeeping, no separate clamping at each step. Train one network to predict dx/dt, then numerically integrate.

Core concepts

Concept 1 — Straight paths instead of stochastic curves

DDPM’s forward process adds Gaussian noise in tiny steps; the reverse process learns to undo each step. The trajectory the model has to follow at sampling time is a curve that wanders through high-dimensional space according to a noise schedule β_t. To integrate that curve accurately you need hundreds of small steps. Ho, Jain and Abbeel’s original DDPM paper [4] used 1,000 steps; DDIM [5] brought it down to ~250 without retraining; further accelerated samplers like DPM-Solver get into the tens but at a cost in flexibility.

Flow Matching takes the opposite design stance. Forget the noise schedule entirely. Define the simplest possible path from noise to data — a straight line — and learn to follow it. The interpolation x_t = (1 - t) · x_0 + t · x_1 is parameterised by t ∈ [0, 1]; the derivative is dx/dt = x_1 - x_0, constant along the line. A network trained to predict this constant at every (x_t, t) produces a velocity field whose integral curve is, by construction, that straight line.

This connects to a much older idea. Optimal transport, the mathematical framework Gaspard Monge sketched in 1781 [6], characterises the most efficient way to push one probability distribution onto another. Straight-line interpolation between samples corresponds exactly to one of the canonical optimal-transport paths — the displacement interpolant — which is why Flow Matching’s geometry is the right one in a precise sense, not just a convenient simplification. Tong et al. [7] generalise this with minibatch optimal-transport couplings for further trajectory straightening.

Concept 2 — The training objective in five lines

The Conditional Flow Matching loss is one of the simplest objectives in modern deep learning. Five lines per training step, no schedules, no warm-up tricks.

def flow_matching_loss(model, x_1):
    x_0 = torch.randn_like(x_1)                          # 1. noise
    t   = torch.rand(x_1.shape[0], 1, 1, 1)              # 2. random time per sample
    x_t = (1 - t) * x_0 + t * x_1                        # 3. straight-line interpolant
    v_target = x_1 - x_0                                 # 4. constant velocity along the line
    v_pred   = model(x_t, t.squeeze())                   # 5. network's prediction
    return F.mse_loss(v_pred, v_target)

Compare line-for-line with DDPM’s loss: no β_t schedule, no ᾱ_t cumulative-product computation, no posterior-variance term, no learned variance head. Just a uniform t, a linear interpolation, and an MSE on the velocity.

Training dynamics are markedly different from GANs and noticeably calmer than DDPM. Loss decreases monotonically (no oscillation, no mode collapse), and the curve flattens cleanly when the model is done learning. The model used in this lesson reached final loss 0.198 from an initial 0.29 over 100,000 steps in ~1.7 hours on an RTX 5070Ti.

Concept 3 — Euler integration: turning a velocity field into samples

A trained model gives you v_θ(x, t), an approximation of the velocity field that flows from noise to data. Sampling is solving the ODE dx/dt = v_θ(x, t) with x(0) ~ N(0, I). Euler’s method is the textbook starting point:

x_{t + Δt} ≈ x_t + Δt · v_θ(x_t, t)

With Δt = 1/50, fifty steps cover the interval [0, 1]. The U-Net backbone (the same architecture DDPM uses; Ronneberger, Fischer and Brox’s original 2015 U-Net [8] is the family) handles the per-step velocity prediction; the integration loop wraps around it.

@torch.no_grad()
def sample_flow(model, n_samples, n_steps=50, shape=(3, 64, 64)):
    x = torch.randn(n_samples, *shape)
    dt = 1.0 / n_steps
    for i in range(n_steps):
        t = torch.full((n_samples,), i / n_steps)
        v = model(x, t)
        x = x + dt * v
    return x.clamp(-1, 1)

Higher-order integrators (midpoint, RK4, Heun) need fewer steps for the same numerical accuracy but cost two or four U-Net forward passes per step. For Flow Matching with n ≥ 20 Euler steps the gain is marginal; for n ≤ 10 it’s substantial. The library implementations (torchcfm, Meta’s flow-matching repo) ship with both options.

Exercises

python exercise1_generate.py

MODIFY II.

Three explorations: step count, trajectory visualisation, velocity field

exercise2_explore.py — three explorations in one script

bash Copy

python exercise2_explore.py

Goal 1 — Step count comparison. Sample the same starting noise at 5, 10, 20, 50, and 100 steps. The script saves a side-by-side grid.

Goal 2 — Trajectory visualisation. Save intermediate x_t at evenly-spaced t values to see the noise → image transformation step by step (Figure 3).

Goal 3 — Velocity field magnitude. Visualise ||v_θ(x_t, t)|| at four different t values to see how the velocity field’s strength changes through the flow.

What to expect — step count

At 5 steps, Euler’s accumulating error produces patterns with visible blockiness and washed-out colour. At 10 steps you have something usable but soft. By 20 steps the output looks essentially indistinguishable from 50 unless you compare carefully. Past 50, the picture stops improving because Euler’s discretisation error on a near-straight path has already dropped below what the model’s own approximation error contributes.

What to expect — trajectory

Pure noise at t=0 resolves into structure progressively. Around t ≈ 0.3 colour zones emerge; around t ≈ 0.6 geometric structure becomes recognisable; the final t = 1.0 step sharpens the details. The transformation is gradual and monotonic in a way DDPM trajectories — which look like denoising — visually aren’t.

What to expect — velocity field

Velocity magnitude is highest at early t and decreases monotonically. This is the geometric counterpart of the training observation that the displacement x_1 - x_0 is what the model is predicting: when t is small, x_t is close to noise and far from data, so the predicted velocity (toward the target) has to be large. As t increases the model is already close to the target and the velocity scales down.

python exercise3_train.py --verify    # confirms 1,059 images found
python exercise3_train.py --train     # starts the 100k-step run

Summary

Common pitfalls

• Too few sampling steps. Below ~10 Euler steps you’ll see blockiness and washed-out colour. The break-even with DDPM is around 20.
• Clamping x to [-1, 1] during integration — the velocity field is defined on the entire real line; clamping mid-flow distorts the trajectory.
• Confusing DDPM’s discrete t ∈ {0, ..., T-1} with Flow Matching’s continuous t ∈ [0, 1]. Feeding integer-valued t to a Flow Matching model produces noise.
• Forgetting model.eval() before sampling. BatchNorm in training mode is the usual culprit when “the same trained model gives different results.”
• Using Euler when a higher-order integrator would help — for n_steps < 10, switch to midpoint or RK4 to recover quality without retraining.

References