Flow matching and rectified flows

1. Intuition

A straight-line recipe. Take a data point $x_0$ (say, 2.0) and a noise sample $x_1$ (say, $-0.5$). The point $x_t = (1 - t)x_0 + t x_1$ slides along a straight line from data at $t = 0$ to noise at $t = 1$. Its velocity is constant: $x_1 - x_0 = -2.5$. Training asks a network $v_\theta(x, t)$ to predict that velocity from $x_t$ and $t$ alone.

Lines cross, so the network learns averages. Many (data, noise) pairs have lines passing through the same $(x, t)$, heading in different directions. The best squared-error prediction is their average velocity. This average field is still exactly right in one sense: if you start from noise and follow it, the particles end up distributed like the data. But the paths it produces curve, because the average direction changes along the way.

Sampling is solving an ODE. Start at $x_1 \sim \mathcal{N}(0, 1)$ at $t = 1$ and step backwards: $x \leftarrow x - \Delta t \cdot v_\theta(x, t)$. With straight paths, one big step would land exactly. With curved paths, big steps cut corners.

One step lands on the mean. At $t = 1$, $x_t$ is pure noise and says nothing about which data point it was paired with, so the averaged velocity points at the data mean. One Euler step from any noise sample lands near the mean of the data. For two modes at $-2$ and $+2$, that's 0: a spike in the empty gap.

Reflow straightens. Run the trained model with many small steps to map each noise sample to a data sample. These pairs are no longer random: nearby noise goes to nearby data, so their lines rarely cross. Train a second model on those pairs, and its paths are nearly straight. Now one step works.

2. Visualization

Synthetic data: a 50/50 mixture of Gaussians at $-2$ and $+2$ with standard deviation 0.3. The velocity network is a 2-hidden-layer tanh MLP with 64 units trained in NumPy. Learned paths (blue) can't cross, because an ODE has one velocity per point, so they curve where the straight training lines (gray) would cross.

Try it (predict first, then run the script and check):

1. Predict the mean of 1-step samples from the first model, then check the "frac |x|>1" column.
2. Predict how the 1-step error changes after reflow, and whether 4 steps after reflow is much better than 1.
3. Shift the data so both modes sit at $+1$ and $+5$. Predict where 1-step samples land.
4. Train the first model with steps=500 instead of 3000. Predict whether 32-step sampling still looks good.

3. The math

Symbols

Watch out

Papers disagree on direction. Lipman et al. (2023) put noise at $t = 0$ and data at $t = 1$; rectified flow as used in SD3 puts data at $t = 0$. The math is identical with $t \to 1 - t$ and a sign flip on $v$. Check the convention before reusing code or a time-sampling scheme.

Interpolant and loss

This is conditional flow matching: the target uses the specific pair. Training is simulation-free: no ODE is solved during training.

What the network converges to

The minimizer of a squared loss is the conditional mean:

Following $dx/dt = v^*$ backward from $\mathcal{N}(0, I)$ produces exactly $p_{\text{data}}$ at $t = 0$ (Lipman et al. 2023; Liu et al. 2022). The loss never reaches zero, because different pairs share the same $(x_t, t)$; the leftover is the variance of $x_1 - x_0$ given $x_t$.

At $t = 1$, $x_t = x_1$, and $x_0$ is independent of $x_1$, so $v^*(x, 1) = x - \mathbb{E}[x_0]$. One Euler step of size 1 from $x_1$ gives $x_1 - v^*(x_1, 1) = \mathbb{E}[x_0]$: every sample lands on the data mean.

Sampling

Relation to diffusion

Diffusion with a Gaussian path writes $x_t = \alpha_t x_0 + \sigma_t \epsilon$. Rectified flow is the choice $\alpha_t = 1 - t$, $\sigma_t = t$, with $\epsilon = x_1$. The three common prediction targets carry the same information:

So a velocity model is a diffusion model with a particular schedule and output parameterization. What changed in practice: straight interpolants, a simple constant target, and sampling as a deterministic ODE.

Rectified flow (reflow)

1. Train $v^{(1)}$ with random pairs $(x_0, x_1)$.
2. For many $x_1$, solve the ODE with $v^{(1)}$ (many small steps) to get $\hat{x}_0 = \Phi^{(1)}(x_1)$.
3. Train $v^{(2)}$ with the same loss on the coupled pairs $(\hat{x}_0, x_1)$.

The coupled pairs come from a non-crossing flow, so their straight lines rarely cross, and the new conditional average is close to each line's own velocity. Liu et al. show reflow never increases transport cost and makes paths straighter; one or two rounds usually suffice before distillation.

Worked example: one training pair

$x_0 = 2.0$, $x_1 = -0.5$, $t = 0.25$.

Going back to $t = 0$ in one step along this pair's own line: $x_t - t\,v = 1.375 - 0.25 \times (-2.5) = 1.375 + 0.625 = 2.0$, which recovers $x_0$. The trained network won't output exactly $-2.5$ here, because other pairs also pass through $(1.375, 0.25)$; it outputs their average.

Worked example: one Euler step from noise

Data mean $\mathbb{E}[x_0] = 0.5 \times (-2) + 0.5 \times 2 = 0$. For a noise sample $x_1 = 1.3$: $v^*(1.3, 1) = 1.3 - 0 = 1.3$, so one step gives $1.3 - 1 \times 1.3 = 0$. The script's trained network gives $v(1, 1) = +1.114$ against the theoretical $+1$, close enough that all 1-step samples land within 1 of zero.

4. Implementation

The full script is code/13-deep-architectures/flow_matching_scratch.py (NumPy only, about 3 seconds on a laptop CPU). The MLP class and Adam are the same pattern as in the GAN script.

Build it from scratch

Step 1: features for $v(x, t)$. The network sees position and time. Smooth functions of $t$ (sine and cosine) make it easier to represent velocity fields that change shape over time.

def features(x, t):
    # [n, 4]: position, time, and two Fourier features of time
    return np.stack([x, t, np.sin(np.pi * t), np.cos(np.pi * t)], axis=1)

Predict: if you dropped $t$ from the features entirely, what velocity would the network have to learn, and could it be right at both $t = 0.1$ and $t = 0.9$?

Step 2: the training step. Sample a time per example, build the point on the line, and regress onto the constant velocity. The gradient of mean squared error, 2 * (pred - target) / n, goes into the hand-written backward pass.

def fm_train_step(net, x0, x1):
    t = rng.random(len(x0))                  # t ~ Uniform(0, 1)
    xt = (1 - t) * x0 + t * x1               # a point on the straight line
    target = x1 - x0                         # its constant velocity
    pred = net(xt, t)
    loss = np.mean((pred - target) ** 2)
    net.adam(net.backward(2 * (pred - target) / len(x0)))
    return loss

Predict: the final loss is 2.301, not 0. Is the model underfitting? (Hint: section 3, "what the network converges to".)

Step 3: random pairing. For the first model, data and noise are drawn independently: train(net, lambda n: (sample_data(n), rng.standard_normal(n))). The pairing function is the only thing reflow will change.

def train(net, pairs_fn, steps=3000, batch=256):
    for _ in range(steps):
        x0, x1 = pairs_fn(batch)
        loss = fm_train_step(net, x0, x1)
    return loss

Predict: with independent pairs, how often do two training lines cross?

Step 4: Euler sampling. Start at noise, step from $t = 1$ down to $t = \Delta t$. Each step evaluates the network once, so n_steps is the cost.

def sample(net, x1, n_steps, keep_path=False):
    x, dt = x1.copy(), 1.0 / n_steps
    path = [x.copy()]
    for i in range(n_steps):
        t = np.full_like(x, 1.0 - i * dt)    # integrate from t = 1 (noise) to t = 0 (data)
        x = x - dt * net(x, t)               # Euler step
        path.append(x.copy())
    return (x, np.array(path)) if keep_path else x

Predict: with n_steps=2, at which two times is the network evaluated?

Step 5: reflow. Map 20,000 noise samples through the first model with 64 steps, then train a fresh model on those coupled pairs. Indexing both arrays with the same idx keeps each pair together.

    z_pairs = rng.standard_normal(20000)
    x_pairs = sample(net, z_pairs, 64)                    # deterministic ODE couples z -> x
    net2 = MLP()

    def coupled(n):
        idx = rng.integers(0, len(z_pairs), n)
        return x_pairs[idx], z_pairs[idx]

    train(net2, coupled)

Predict: if you shuffled x_pairs before training net2, what would reflow do?

Use the library

There's no one-line flow-matching layer; the training loop is short enough to write in PyTorch directly. Production models use the same loss with a transformer for $v_\theta$, latents instead of pixels, a conditioning input, and non-uniform time sampling.

import torch

def fm_loss(model, x0, cond):                     # x0: [B, C, H, W] latents
    x1 = torch.randn_like(x0)
    t = torch.sigmoid(torch.randn(x0.shape[0], device=x0.device))   # logit-normal t, as in SD3
    tt = t.view(-1, 1, 1, 1)
    xt = (1 - tt) * x0 + tt * x1
    return torch.nn.functional.mse_loss(model(xt, t, cond), x1 - x0)

For sampling, diffusers ships FlowMatchEulerDiscreteScheduler, which SD3 and Flux pipelines use; num_inference_steps is $N$. The torchcfm package provides conditional flow matching variants, including minibatch optimal-transport pairing, another way to get straighter paths without a full reflow.

Run it

uv run code/13-deep-architectures/flow_matching_scratch.py

Key lines of output:

x0 = 2.0, x1 = -0.5, t = 0.25:  x_t = 1.375, target velocity = -2.500
one Euler step back to t = 0: x_t - t * v = 2.000 (recovers x0)
final training loss: 2.301 (cannot reach 0: one x_t has many possible x1 - x0)
v(x=0, t=1) = -0.011   v(x=1, t=1) = +1.114   (theory: v(x, 1) = x - E[x0] = x)
steps | W1 to data | frac |x|>1
    1 |      1.920 |       0.00
    2 |      0.687 |       0.72
    4 |      0.276 |       0.92
    8 |      0.147 |       0.97
   32 |      0.068 |       0.99

steps | W1 before | W1 after reflow
    1 |     1.920 |           0.081
    4 |     0.276 |           0.078

W1 is the 1-D Wasserstein distance to fresh data samples (mean gap between sorted samples; lower is better). Before reflow, error falls steadily with steps. After one reflow, 1 step is already as good as 32 steps of the original model.

5. Engineering

Why it took over. Training is a stable regression with a simple target. Straighter paths make deterministic ODE samplers efficient (20 to 30 steps instead of hundreds) and give distillation a better starting point. SD3 reported that rectified flow with logit-normal time sampling (more training on middle times, where prediction is hardest) beat other diffusion formulations at scale in their experiments.

Time sampling. Uniform $t$ wastes effort near $t = 0$ and $t = 1$, where the target is easy to predict. Logit-normal sampling, $t = \sigma(u)$ with $u \sim \mathcal{N}(m, s^2)$, concentrates on the middle. At high resolution, models also shift the schedule toward noisier times, because more pixels mean more redundant signal survives a given noise level.

Getting to few steps. Options, from cheapest to most involved: a better ODE solver (Heun, midpoint) at the same number of evaluations; reflow (costs a full generation pass over a large set of noise samples to build pairs); and distillation into a 1 to 4 step student (consistency-style or adversarial). Fast open models such as Flux.1-schnell are flow-matching models distilled to a few steps.

Guidance. Classifier-free guidance works on velocities just as on noise: $\tilde v = v_\varnothing + w(v_c - v_\varnothing)$ (see latent diffusion).

Beyond images. The same recipe trains audio, video, molecule, and robot-action generators: anything where you can draw samples and need a generator that maps noise to data.

Common mistakes

• Mixing time conventions (data at $t = 0$ versus $t = 1$) between training and sampling, which flips the velocity sign.
• Expecting a model trained with random pairs to sample well in 1 step: that lands on the data mean.
• Reading the training loss as sample quality. Its floor is the irreducible variance of crossing lines; compare sample metrics across step counts.
• Shuffling or mismatching pairs when building a reflow dataset, which throws away the coupling.
• Evaluating at only one step count; report quality versus number of function evaluations.

6. Knowledge check

Practice exercise. A velocity model is trained with $x_t = (1 - t)x_0 + t x_1$. At $t = 0.6$, $x_t = 0.8$, and the model predicts $v = -1.5$. What are its implied predictions $\hat{x}_0$ and $\hat\epsilon$? Then take one Euler step of size $\Delta t = 0.2$.

Implementation challenge. Replace Euler with Heun's method (predict with Euler, evaluate the velocity at the predicted point, average the two velocities) in sample. Compare W1 at an equal number of network evaluations: Euler with 8 steps versus Heun with 4 steps (also 8 evaluations).

Summary

• Flow matching regresses $v_\theta(x_t, t)$ onto $x_1 - x_0$ along straight lines $x_t = (1-t)x_0 + t x_1$; training needs no ODE solves.
• The network learns the average velocity of all lines through a point; following it from noise reproduces the data distribution.
• With random pairs, paths curve and one Euler step lands on the data mean; quality improves with more steps.
• Reflow retrains on the model's own (noise, sample) pairs, straightening paths so 1 to 2 steps work; distillation goes further.
• Velocity, noise, and data predictions are interchangeable; rectified flow is diffusion with a straight-line schedule. SD3 and Flux use it.

Next: Object detection and YOLO

Related: Latent diffusion and Stable Diffusion · GANs · Autoencoders, diffusion, and transfer learning · Model card: diffusion model

Go deeper, code-first (AI Engineering from Scratch): Diffusion Transformers & Rectified Flow

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Flow Matching & Rectified Flows.

Interview angle