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.

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, the last four chapters collapse into one idea.

Moving particles

You have particles distributed as \(p_0 = \mathcal{N}(0, I)\). You want them distributed like your data 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.

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. You cannot even write it down.

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\):

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. 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 learned linear-interpolant paths come out more curved than the diffusion ones.
• Past about 32 steps the two interpolants converge to the same error. The entire argument is about the few-step regime.

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. But the coupling is now one the model can realize with a straight line, and a straight line is integrated exactly by one Euler step.

Tick after one reflow in the panel and the curvature goes to zero and the error goes to zero at every step count. That is the mechanism behind one-step and few-step generation, and it is why SD3 and FLUX are described as rectified-flow models rather than merely flow-matching ones.

Where it is used

Flow matching is not a niche alternative any more; it is what the current generation of models is built on.

flowchart LR
    A["Text<br/>(prompt)"] --> B["Text encoders<br/>CLIP + T5-XXL"]
    N["Noise<br/>z₀ ~ N(0,I)"] --> C
    B --> C["Diffusion Transformer<br/>rectified flow"]
    C -->|"ODE: 20–50 steps"| D["Latent z₁"]
    D --> E["VAE decoder"]
    E --> F["Image 1024×1024"]

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.
       At equal cost it beats Euler once the field has any curvature at all."""
    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

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.
3. Diffusion is flow matching with a curved interpolant. \(a(t) = \sqrt{\bar\alpha_t}\) instead of \(a(t) = t\). One framework covers both.
4. "Straight paths" applies to the conditional path. The marginal path — the one the model learns — is not straight, and may be more curved than diffusion's.
5. The real advantage is the objective: simpler, unweighted, no schedule to get wrong. That is why SD3 and FLUX adopted it.
6. Reflow delivers actual straightness by re-coupling each noise sample with the output it already produces. This is what makes one-step generation possible.
7. Consistency models, shortcut models and MeanFlow all attack the same target: collapse the trajectory into fewer evaluations.
8. Past ~32 steps the interpolant stops mattering. Every argument here is about the few-step regime.