Flow Matching (Rectified Flow) - Generative Models

Conditional Flow Matching (Rectified Flow)

What you'll build

The training target and loss of flow matching with straight paths, and an Euler sampler that integrates the learned velocity field from noise to data.

Background

Flow matching (Lipman et al., 2023) and rectified flow (Liu et al., 2023) train a network v_θ(x, t) to be a velocity field whose ODE carries noise to data. With straight-line paths between a noise sample and a data sample, training is plain regression, with no noise schedule or variance bookkeeping. Stable Diffusion 3 and Flux are trained this way.

Formula (training)

Sample noise x0 ~ N(0, I), a data point x1, and t ~ U(0, 1) per example:

x_t    = (1 − t) · x0 + t · x1           # point on the straight path
target = d x_t / dt = x1 − x0            # its constant velocity
loss   = mean( (v_pred − target)² )      # mean over every element of the batch

Here t = 0 is noise and t = 1 is data (some papers use the opposite convention). t has shape (B,): reshape it to (B, 1, ..., 1) so it broadcasts over the remaining dimensions.

Algorithm (sampling with Euler steps)

1. Start from x = x0 and set dt = 1 / n_steps.
2. For i = 0, ..., n_steps − 1: t = i · dt (a Python float), x ← x + dt · velocity_fn(x, t).
3. Return x.

If the learned field is exactly the straight-line velocity, a single Euler step lands on the data: that is the appeal of straight (rectified) paths.

Inputs

• x0, x1, v_pred: float arrays (B, ...) of the same shape
• t: float array (B,) with values in [0, 1]
• velocity_fn(x, t): callable returning an array shaped like x; n_steps: int ≥ 1

Expected return format

• flow_matching_loss(x0, x1, t, v_pred) returns a dict:
  • 'x_t': interpolated points, same shape as x0
  • 'target': target velocity x1 − x0, same shape as x0
  • 'loss': MSE between v_pred and target (Python float)
• euler_step(x_t, v, dt) returns x_t + dt · v
• euler_sample(x0, velocity_fn, n_steps) returns the final x (same shape as x0)

Constraints

NumPy only; the Euler loop is over steps, never over elements.

Follow-up

• How is this related to diffusion? (With Gaussian paths, flow matching recovers a probability-flow ODE of a diffusion model.)
• What does "reflow" do to make one-step sampling work for a real, imperfect model?

Input: x0 = [[0, 0], [1, 1]], x1 = [[2, 4], [3, -1]], t = [0.25, 0.5]
Output: x_t = [[0.5, 1.0], [2.0, 0.0]], target = [[2, 4], [2, -2]]

Input: euler_sample(np.array([2.0]), lambda x, t: -x, 4)
Output: array([0.6328])

import numpy as np

def flow_matching_loss(x0, x1, t, v_pred):
    """
    Conditional flow matching with straight paths.

    Args:
        x0: (B, ...) noise samples
        x1: (B, ...) data samples
        t: (B,) times in [0, 1]
        v_pred: (B, ...) predicted velocities at (x_t, t)

    Returns:
        dict with 'x_t', 'target' (same shape as x0) and 'loss' (float)
    """
    # Your code here
    pass


def euler_step(x_t, v, dt):
    """One Euler step: x_t + dt * v"""
    # Your code here
    pass


def euler_sample(x0, velocity_fn, n_steps):
    """Integrate dx/dt = velocity_fn(x, t) from t = 0 to 1 with n_steps Euler steps."""
    # Your code here
    pass