21. Flow Matching

Diffusion works, and it is a strange way to have arrived at something that works. There is a stochastic differential equation, a variational bound, a noise schedule with a name, a discrete chain of a thousand steps, and three interchangeable parameterizations. The method is not complicated — the derivation is.

Flow matching¹ arrives at the same place from the other direction and the setup fits in three lines. Draw a noise sample and a data sample. Draw the straight line between them. Train a network to predict the direction of travel at a random point on that line. That is the whole method.

Here \(x_0 \sim \mathcal{N}(0, I)\) is the noise, \(x_1\) a sample of the data, \(t\) a time drawn uniformly in \([0, 1]\) and \(v_\theta\) the network, with weights \(\theta\). CFM stands for conditional flow matching, because the target \(x_1 - x_0\) is that of one drawn pair. Time runs from noise (\(t = 0\)) to data (\(t = 1\)), the opposite of chapter 20.

No schedule, no bound, no chain. And — this is the part worth being clear about — it is not a different family of model. Diffusion is flow matching with a curved interpolant. Once you see that, chapter 20 and this one collapse into one idea.

It is also not a historical curiosity. Since 2024 almost every open image and video generator whose training objective is published — Stable Diffusion 3, FLUX, Wan, HunyuanVideo, Qwen-Image — trains with exactly this loss, and so do speech models and the action heads of robot policies. Where it is used has the list, with dates.

Slides

Click the slides and use ← → to move through them; F goes full screen and P opens the presenter view, with the notes. The three simulators of this page and three of its diagrams run inside the slides. Open the slides in a new tab

Moving particles

You have particles distributed as \(p_0 = \mathcal{N}(0, I)\). You want them distributed like your data, \(p_1\), by time \(t = 1\). Learn a velocity field \(v_\theta(x, t)\) that says which way each particle should move at each instant, and integrate.

The panel has the two pictures this chapter keeps comparing. Training pairs is what the loss above sees: each noise sample joined to an independently drawn data point by a straight line. The learned flow is what a network trained on those lines converges to, computed exactly here (the data are a mixture of small Gaussians, so nothing needs training). Animate both.

There is an obvious problem with training this, and it is the one the whole field was stuck on. The velocity field you need is a marginal: at position \(x\) and time \(t\), it must average over every data point that could plausibly have produced \(x\),

You cannot compute that for real data. You cannot even write it down.

The theorem is easiest to see at a single point. With two data points, exactly two training lines pass through each \((x, t)\), and everything can be computed by hand:

Two consequences are visible in the particle panel and are worth stating, because both surprise people the first time:

• The training loss does not go to zero. Many pairs pass near the same \((x, t)\) with different directions; the network can only output their average, and the spread around it is a floor. On the eight-Gaussian target of the lab the floor is 3.73 (computed exactly), and the trained network ends at 3.76, a few hundredths above it. A loss of 3.76 here is a model that has converged, not one that failed.
• The learned paths cannot cross. An ODE has one velocity per point and time, so two particles never meet. The flow therefore re-pairs noise and data: each particle lands on the data, but not on the point it was paired with in training. In the panel the training pairs travel 1.23 on average and the flow 0.51.

The same two facts in the lab's code: one step of train_2d on three pairs in 2-D, then Euler sampling from the noise of pair 1. The network is replaced by the field it converges to, exact here because the data are two points, \(A = (2, 0)\) and \(B = (-2, 0)\). Pair 1 joins the noise \((1, 1)\) to \(B\): with the network already ideal it carries the whole loss, and the flow then takes that noise to \(A\). Change \(t\), the number of Euler steps or the case:

Diffusion is a special case

Take any pair of functions \(a(t)\), \(b(t)\) and define the interpolant⁸

Every method in this module is a choice of \(a\) and \(b\):

Here \(\bar\alpha_t\) is the noise schedule of chapter 20 read backwards, growing from 0 at the noise to 1 at the data. With the cosine schedule it is \(\bar\alpha_t = \sin^2(\pi t/2)\), so \(a = \sin(\pi t/2)\) and \(b = \cos(\pi t/2)\). OT stands for optimal transport, the name Lipman et al. give the straight path. One pair under both interpolants (pick \(t\)):

The conditional velocity is \(a'(t)x_1 + b'(t)x_0\) in every case, and the marginal is its conditional expectation. One framework, one loss, one code path. With a Gaussian \(p_0\) the two views are exactly interchangeable: a model trained one way can be sampled the other, and on the linear path the deterministic DDIM sampler is the Euler step of the flow⁹ (on other schedules it is an Euler step in rescaled variables). The remaining differences are the weighting of the loss across \(t\) and the shape of the trajectory — which is what the next panel measures.

What "straighter paths" actually means

The standard claim is that flow matching gives straight trajectories and therefore needs fewer integration steps. It is half true, and the half that is false is worth knowing.

Everything in that panel is exact — the target is a Gaussian mixture, so both interpolants' marginal velocities have closed forms and nothing is trained. Compare the two and:

• The conditional path — the line between one noise sample and one data sample — is straight by construction with the linear interpolant.
• The marginal path, which is what the model actually learns, is not. It has to average over every data point compatible with the current position, and independently paired noise and data cross each other constantly. On this target the linear-interpolant paths come out more curved than the diffusion ones: 111% longer than the straight line against 58%.
• Neither interpolant wins at every step count (at 2 steps diffusion's endpoint error is lower, at 8 the linear one's), and past about 32 steps both are below anything you would see. The entire argument is about the few-step regime.

On real images the measurements are mixed, and the weighting of the loss across \(t\) matters as much as the path. Lipman et al. counted the evaluations an adaptive solver needs: 122–142 with the linear path against 183–193 with the diffusion path with the same architecture, on CIFAR-10 and ImageNet 32 and 64¹ — 22 to 37% fewer, not an order of magnitude. Esser et al. trained 61 formulations for Stable Diffusion 3: the rectified-flow variants were the best when sampling with few steps, but with a uniform \(t\) rectified flow ranked below the classic \(\epsilon\)-prediction overall, and only with logit-normal timestep sampling (below) did it come out first⁷. The lab trains both paths on MNIST with a uniform \(t\) and the same budget, and the cosine diffusion path wins at every step count (FD 94.0 against 133 at 8 Euler steps). FD is the recipe of FID, the Fréchet distance between the features a network extracts from generated and from real images, with a digit classifier in place of the Inception network; lower is better.

Rectified flow: straightening on purpose

If the paths are not straight, straighten them³. Reflow is two steps:

1. Sample noise \(z\), integrate the trained ODE accurately, get \(x = \text{ODE}(z)\).
2. Retrain the model on the pairs \((z, x)\) that it produced itself, on the straight line between them.

The marginal distribution of \(x\) is unchanged — you are not altering what the model generates, only which noise is paired with which output. And the new pairing is one the model can nearly realize with straight lines: the flow's paths do not meet, so the lines between their endpoints rarely cross, and a straight line is integrated exactly by one Euler step. The same steps with six noise samples and the two data points from before:

Tick after an ideal reflow in the paths panel above: the curvature (how much longer the paths are than the straight line between their ends) and the error go to zero at every step count. That is the limit, a model that learned the new lines exactly. A trained reflow gets close: on the lab's eight Gaussians it drops the curvature from 88% to 0.1% and makes one step better than 8 of the original model; on MNIST one step goes from FD 2255 to 201. It is not free, though: each round needs a large set of pairs generated with the full ODE, the errors of the first model are baked into them, and on images one reflow left a visible gap until the training recipe was reworked (2-rectified flow on CIFAR-10 went from a one-step FID of 12.21 to 3.07 with a better initialization, timestep sampling and loss)¹¹. InstaFlow ran reflow plus distillation on Stable Diffusion for one-step text-to-image in 2023¹².

Fewer steps: from velocities to jumps

Reflow straightens the field so that a short integration is accurate. The other line of work stops integrating: it trains the network to predict where the trajectory ends up after a jump of any size, a flow map \(X_{s \to t}(x)\): the point at time \(t\) of the trajectory that passes through \(x\) at time \(s\)¹⁴. The instantaneous velocity is the limit of a jump of zero length, and a one-step sampler is the jump from 0 to 1.

The research methods (the first three) train one-step models from scratch on ImageNet; what ships in products in 2026 is still mostly distillation of a many-step flow-matching model, adversarial or distribution-matching. Both start from the field this chapter trains.

Training choices that matter

The loss is one line, but three choices around it changed between the first papers and the current models, and you will meet all three in any recent codebase:

• Where in \(t\) to train. Uniform \(t\) wastes effort near the ends, where the target is easy. SD3 samples \(t\) from a logit-normal, concentrated on the middle⁷, and HunyuanVideo does the same; the robot policies π0 and GR00T use a Beta distribution that favours the noisy end¹⁷.
• Shift with the resolution. More pixels or tokens make a given noise level easier to see through, so large images and videos move the timesteps toward the noise end, by a factor that grows with the number of tokens (SD3, LTX-Video). This is the timestep shift of the SNR panel in chapter 20.
• What the network outputs. Velocity is the default, but the network may also predict \(x_1\) (the clean data) or the noise, all convertible given \(t\). Li and He argue that predicting the clean data directly is what lets a plain Transformer generate in pixel space with large patches¹⁸ — a 2025 result that reopened a choice everyone thought settled.

Where it is used

Flow matching is not a niche alternative any more. Where a training objective is published, the current generation of generative models uses it — but not everywhere, and some of the best-known systems do not say.

Two things in that table are worth a second look. In robotics, flow matching won for a reason that is not image quality: π0 needs about 10 Euler steps for a chunk of 50 actions, fast enough for control at up to 50 Hz, and it can represent an action distribution with several modes, where a regression head would average them into an action that matches none¹⁷. In video, the pipeline is the same as for images — only the latent gains a time axis.

What is not flow matching is just as instructive: AlphaFold 3 uses a standard diffusion module; CogVideoX trains with \(\mathbf{v}\)-prediction diffusion; Stable Audio Open with the \(\mathbf{v}\) objective; Sora and Veo 3 are described by their makers as diffusion models, without the objective being disclosed. Given the equivalence above the boundary is thin, and the label matters less than the interpolant, the target and the timestep weighting.

FLUX.1-dev at 1024 × 1024, with the shape of what passes between the parts. The velocity field of this chapter is the box in the middle, now also conditioned on the text, \(c\); everything around it is the latent-diffusion plumbing of chapter 20, and the Transformer itself is chapter 22. FLUX.2 keeps the scheme with a 32B Transformer and replaces the two text encoders with a 24B vision-language model¹⁹.

Sampling

The training gives you a velocity field. Turning it into a sample is ODE integration, and you may pick the solver — the same separation as in chapter 20.

import torch

@torch.no_grad()
def sample_euler(model, shape, n_steps=25, device='cuda'):
    """First order. One model evaluation per step."""
    x = torch.randn(shape, device=device)             # x at t = 0 is pure noise
    dt = 1.0 / n_steps
    for i in range(n_steps):
        t = torch.full((shape[0],), i * dt, device=device)
        x = x + dt * model(x, t)
    return x

@torch.no_grad()
def sample_heun(model, shape, n_steps=15, device='cuda'):
    """Second order: two evaluations per step, but far less error per step."""
    x = torch.randn(shape, device=device)
    dt = 1.0 / n_steps
    for i in range(n_steps):
        t  = torch.full((shape[0],), i * dt, device=device)
        t2 = torch.full((shape[0],), (i + 1) * dt, device=device)
        v1 = model(x, t)
        v2 = model(x + dt * v1, t2)                   # predictor, then correct
        x = x + dt * (v1 + v2) / 2
    return x

Count model evaluations, not steps: Heun with 8 steps costs what Euler with 16 does. On the lab's MNIST model Heun loses with few evaluations (FD 550 against Euler's 226 at 4) and wins from 16 on (89.7 against 99.3): a second-order step needs steps small enough for its correction to be right.

Lab: flow matching from scratch

The chapter's claims are measured in a Colab notebook that trains flow-matching models end to end, first on two dimensions, where the exact field is known, then on MNIST. It runs on the free tier: CPU is enough for sections 1–5, and the MNIST sections want a T4 GPU (Runtime → Change runtime type → T4 GPU); the reference run of every cell took 15 to 25 minutes on an Apple laptop GPU shared with other jobs, slower than a T4. Eight sections, each ending in a checkpoint with a check value:

1. the exact marginal field of a 2-D mixture, and the floor of the loss, computed;
2. a velocity MLP trained to that floor, and its distance to the exact field;
3. sampling: endpoint error and distance to the data against the number of steps, Euler against Heun at equal cost, and the curvature of the learned paths;
4. the diffusion interpolant with the same code, and which one is straighter;
5. one round of reflow, and what it does to the curvature and to one-step samples;
6. MNIST: a small U-Net trained with the same loss, scored by a Fréchet distance;
7. one step lands on the mean digit; Heun, and the diffusion interpolant, on images;
8. reflow on MNIST, and an open task: logit-normal timesteps, with a table to fill in.

Open the lab in Colab — save a copy to your Drive (File → Save a copy in Drive) to run and edit it.

Key takeaways

1. Flow matching learns a velocity field and samples by integrating an ODE. The loss is one MSE with no schedule and no bound.
2. The marginal velocity is intractable, but regressing the conditional one has the same gradient. That theorem is the whole method. A consequence: the loss has a floor above zero.
3. Diffusion is flow matching with a curved interpolant. \(a(t) = \sqrt{\bar\alpha_t}\) instead of \(a(t) = t\). One framework covers both; on the linear path, DDIM is exactly the Euler step of the flow.
4. "Straight paths" applies to the conditional path. The marginal path — the one the model learns — is not straight; it re-pairs noise and data so that paths never meet.
5. On images the linear path saves at most about a third of the solver steps, and only with the right weighting across \(t\); with a uniform \(t\) a good diffusion schedule can beat it, as in the lab. The bigger reasons for adopting it are the objective (simple, robust to the schedule) and the recipe around it: logit-normal timesteps, resolution shift.
6. Reflow straightens the field by retraining on the model's own pairs. It works, at the cost of generating those pairs. The "rectified flow" of SD3 and FLUX is the straight interpolant, not reflow.
7. One-step generation comes from learning jumps instead of velocities — consistency, shortcut, MeanFlow — and, in shipped models, from distillation (FLUX.1 [schnell], SD3.5 Turbo).
8. Since 2024 the objective is shared by image, video, speech and robot-action models: a different \(p_1\), the same line of code.