Flow matching: turning noise into shapes

Watch the whole method first, then take it apart

The video shows the method on a flat shape in under two minutes. The rest of the page explains each part, with numbers you can check.

A straight path joins each noise draw to an example

The noise is a cloud of random points. Each coordinate of a noise draw x₀ is an independent draw from the standard bell curve, with mean 0 and spread 1. In two dimensions the cloud is round and densest at the center. The examples are points on a shape. A 3D point cloud works the same way, with three coordinates per point.

Pair a noise draw x₀ with an example x₁. The straight path between them is

x_t = (1 − t) x₀ + t x₁,   0 ≤ t ≤ 1.

At t = 0 the point is at the noise draw; at t = 1 it is at the example. Its velocity, the rate of change of x_t with time, is the same all along the path:

velocity = x₁ − x₀.

So the point moves at a steady speed and covers the whole path in one unit of time.

Your turn. With the same pair, where is the point at t = 0.5, and what is its velocity there?

Training: predict the path's velocity from the position and the time

The network, written v_θ(x, t), takes a position and a time and returns a velocity of the same size as the position. θ (theta) stands for all of its adjustable numbers, its weights. The training target is the velocity of the straight path, x₁ − x₀. The loss is the squared error between prediction and target, averaged over many draws:

L(θ) = average of ‖ v_θ(x_t, t) − (x₁ − x₀) ‖²

Here ‖…‖² is the squared length: the sum of the squared coordinates. The average runs over examples x₁ from the training data, noise draws x₀ from the standard bell curve, and times t drawn evenly from 0 to 1. One training step does this:

1. Take a minibatch of examples x₁.
2. Draw one noise draw x₀ for each example, with the same shape.
3. Draw one time t for each example.
4. Compute x_t = (1 − t) x₀ + t x₁.
5. Run the network to get the predicted velocity v_θ(x_t, t).
6. Compute the loss: the average squared error against x₁ − x₀.
7. Take one gradient-descent step to change θ so the loss goes down.

For a point cloud, an example x₁ is the whole cloud, say 1,000 points with 3 coordinates each. The noise draw has the same shape, so every point gets its own noise point, and the network returns one velocity per point.

Many paths pass through each spot at the same time, so the network learns their average

Pick one spot and one time. Many training pairs have a path that passes through that spot at that time, and their velocities differ. The network sees only the spot and the time, so it cannot tell those pairs apart.

Squared error is smallest when a prediction equals the average of its targets (see loss). For the targets 1 and 3, for example, the prediction 2 gives squared errors 1 + 1 = 2, and the prediction 1 gives 0 + 4 = 4. So the best possible network predicts, at each spot and time, the average velocity of all paths through it. This averaged field is what training learns.

Three things follow.

• The loss never reaches zero. The targets at one spot and time disagree with each other, and no single prediction can match them all.
• The field's paths are curved. A particle that follows the averaged field does not follow any single straight path. Its path bends, because two particles can never be at the same place at the same time (see velocity fields).
• The field still carries noise onto data. Start a cloud of particles from the noise at t = 0 and let each follow the averaged field. At every time, the cloud is spread exactly like the points x_t of the training pairs. At t = 1, that is the spread of the examples. This is the central result of flow matching.

Generation: start from fresh noise and take Euler steps

To make a new example, draw a fresh noise draw and follow the learned field from t = 0 to t = 1 with Euler's method. With n steps of size h = 1/n, each step is

x_t+h = x_t + h × v_θ(x_t, t),   t = 0, h, 2h, …, 1 − h.

Each step runs the network once, so n steps cost n network runs. A new noise draw gives a new example. Few steps are cheap but drift off the curved paths; many steps follow them closely.

The program below trains a small network on points of a ring with radius 2, then generates with 1, 4 and 40 steps. Predict first: where do the points land after just one step?

flow_matching_ring.py (PyTorch, CPU, about 5 seconds)

import torch
import torch.nn as nn

torch.manual_seed(0)                       # the same random draws on every run

def data(n):
    # Training examples: points on a ring of radius 2.
    a = 2 * torch.pi * torch.rand(n)       # n random angles from 0 to 2π
    return 2 * torch.stack([a.cos(), a.sin()], dim=1)

# The network reads (x, y, t) and returns a velocity (vx, vy).
net = nn.Sequential(nn.Linear(3, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
opt = torch.optim.Adam(net.parameters(), lr=0.001)   # Adam, an improved gradient descent

for step in range(3000):
    x1 = data(256)                         # data examples
    x0 = torch.randn_like(x1)              # noise, one draw per example
    t = torch.rand(256, 1)                 # a random time per example
    xt = (1 - t) * x0 + t * x1             # the point on each straight path
    # torch.cat puts t beside xt, so each row reads (x, y, t).
    v = net(torch.cat([xt, t], dim=1))     # the predicted velocity
    loss = ((v - (x1 - x0)) ** 2).mean()   # squared error to the target x1 - x0
    opt.zero_grad()                        # clear the old gradients
    loss.backward()
    opt.step()                             # one step of Adam
print(f"final loss {loss.item():.2f}")

# @torch.no_grad() works like "with torch.no_grad():" around the whole function.
@torch.no_grad()
def generate(n, steps):
    x = torch.randn(n, 2)                  # fresh noise at t = 0
    for k in range(steps):
        t = torch.full((n, 1), k / steps)  # the time k / steps, once per point
        x = x + (1 / steps) * net(torch.cat([x, t], dim=1))   # one Euler step
    return x

for steps in [1, 4, 40]:
    r = generate(2000, steps).norm(dim=1)  # each point's distance from the center
    # {steps:2d} pads the step count to 2 characters, so the columns line up.
    print(f"{steps:2d} steps: distance from center {r.mean():.2f} on average, spread {r.std():.2f}")

final loss 2.03
 1 steps: distance from center 0.18 on average, spread 0.10
 4 steps: distance from center 1.91 on average, spread 0.30
40 steps: distance from center 2.00 on average, spread 0.09

One step lands every point near the center, the average example, as the worked example predicted. Four steps reach the ring but blur it. Forty steps give a sharp ring of radius 2. The final loss stays near 2, far from zero, because the targets at each spot disagree. Your numbers may differ slightly on another machine.

Flow matching is a close relative of diffusion models

Diffusion models also start from bell-curve noise and remove it over many steps. The two families differ in four choices. Diffusion models usually blend noise and data along a curved schedule rather than a straight line. Their networks often predict the noise that was added rather than a velocity. They weight the times in the loss differently. And many diffusion samplers add a little fresh noise at every step.

With bell-curve noise, the two are closely linked. A prediction of the noise can be converted into a velocity and back, and the way to sample a diffusion model without fresh noise also follows a velocity field. So the families share most of their mathematics, but a model trained one way is not the same model as one trained the other way. "Rectified flow" is the same straight-path method, published at about the same time under another name.