Skip to content

Flow Matching — Visual Guide

Flow matching (rectified flow) connects each noise sample to a data sample by a straight line and trains the network to predict the direction of that line, the velocity. Sampling follows the learned velocity field from noise (t = 0) to an ECG (t = 1) with a few Euler or Heun steps. Code: src/ecgen/processes/flow.py (RectifiedFlowProcess), family src/ecgen/families/continuous/module.py (FlowMatchingModule). Written reference: Diffusion and flow matching — Flow matching.

FlowMatchingModule is DiffusionModule with a different process. The backbone (UNet1d by default), the conditioning, classifier-free guidance and EMA are identical; see the Diffusion visual guide for those.

One training step

noise x₀ N(0, I) data x₁ (B, 12, 5000) + (1 − t) · x₀ t · x₁ t ~ logit-N(0, 1) x_t UNet1dx_t, 1000·t, c conditioner v̂ MSEtarget x₁ − x₀ schematic traces · the module, backbone and loss code are shared with diffusion
Hover over a block for details.

Straight paths and solver steps

The left panel shows training pairs: each noise point is joined to a data point by a straight line, and the network learns the slope x₁ − x₀ at every point of it. Because many lines pass through the same region, the network learns their average direction, so the paths it follows at sampling time (right) are curved. Euler steps follow them in straight segments; more steps or Heun stay closer to the exact path. Minibatch OT pairing (ot_coupling: true) shortens the training paths and makes them cross less.

Toy 2-D example: noise N(0, I), data = three Gaussian clusters, and the exact velocity field of that toy instead of a trained network. A real sample is a point in 12 × 5000 = 60 000 dimensions. The OT pairing here is a 2-opt approximation; the code uses scipy.optimize.linear_sum_assignment.

Which t is trained on

t_sampling: logit_normal (default) draws t = sigmoid(n), n ~ N(logit_mean = 0, logit_std = 1). It puts most training examples at intermediate t, where predicting the velocity is hardest, and few near pure noise or clean data.

How the loss is computed

# RectifiedFlowProcess.training_pair(x1, noise)
if self.ot_coupling and x1.shape[0] > 1:
    noise = noise[ot_pairing(x1, noise)]      # re-pair noise with data inside the batch
t = self.sample_t(x1.shape[0], x1.device)     # logit-normal in (0, 1)
xt = t * x1 + (1 - (1 - self.sigma_min) * t) * noise
target = x1 - (1 - self.sigma_min) * noise
# then, as for diffusion: loss = (weight * ((backbone(xt, 1000 * t, cond) - target) ** 2).mean()).mean(), weight = 1
\[ x_t = t\,x_1 + \big(1 - (1-\sigma_{\min})\,t\big)\,x_0,\qquad \frac{dx_t}{dt} = x_1 - (1-\sigma_{\min})\,x_0,\qquad \mathcal{L} = \mathbb{E}_{t,\,x_0,\,x_1}\,\big\|\hat v(x_t, t, c) - (x_1 - (1-\sigma_{\min})\,x_0)\big\|^2 \]
∂ℒ → UNet1d and conditioner data x₁ noise x₀ training_pair [x₀ ← x₀[OT perm]] t = sigmoid(N(0, 1)) x_t = t·x₁ + (1−t)·x₀ target = x₁ − x₀ weight = 1 (σ_min = 0) UNet1dx_t, 1000·t, c conditioner v̂ target x₁ − x₀ MSE per sample(B,) ℒ (mean) purple = gradient path · the target is a constant
Hover over a block.

Sampling

RectifiedFlowProcess.sample starts from x = noise at t = 0 and takes steps equal steps of size h = 1 / steps up to t = 1. Guidance is applied to the velocity with the same cfg_model_fn.

for i in range(steps):                     # t = i / steps
    d = v(x, t)                            # u + w (c - u) with guidance
    if solver == "heun" and i < steps - 1:
        x_euler = x + h * d
        x = x + h * 0.5 * (d + v(x_euler, t + h))
    else:                                  # Euler (and the last Heun step)
        x = x + h * d

The cost of a sample is the number of backbone calls. Every call is the same UNet1d pass, so the bars compare directly (defaults: sample_steps = 50 for diffusion; flow_matching_mimic.yaml uses 25 Heun steps):

samplernetwork calls per sample (1 call = 1 UNet1d pass on (n, 12, 5000)) flow · Euler 1010 flow · Euler 2525 flow · Heun 2549 diffusion · DDIM 5050 DDIM 50 + guidance100 DDIM 50 + RePaint246 → bar length ∝ calls (RePaint bar truncated) · fewer calls is faster; sample quality per step count is not shown here
Hover over a bar for how the count is made up.

Inpainting. Flows support inpaint="replace": before every step the observed entries are set to the point on the straight path from the initial noise to the observation, t·obs + (1 − t)·noise, and to obs itself at the end. RePaint is diffusion-only and raises an error here. Models trained with clean_observed use the clean observation instead. x_init / t_start start mid-path for SDEdit-style editing.