Flow Matching: Simplified Generative Modeling

TL;DR

Flow Matching trains a continuous normalizing flow by regressing a velocity field on chosen conditional probability paths. The original paper includes diffusion paths and a conditional Gaussian optimal transport path with linear mean and standard deviation. The latter gives straight conditional trajectories, but the learned marginal trajectories need not be straight or globally optimal. Training requires no ODE simulation; generation still requires numerical integration. Sampling efficiency depends on the learned field, solver, and quality target, rather than a universal 10–50× reduction.

The Core Idea: Choose a Probability Path

Generative modeling asks how to transform simple noise into complex data. Flow Matching makes the probability path an explicit design choice, then trains a neural velocity field v_\theta(\mathbf{x},t) to generate it. Diffusion probability paths are among the supported choices; the conditional OT construction offers a different path with simple velocity targets.

Generation integrates the learned field from noise to a final sample:

The network sees position and time, not the destination data point. The pipeline diagram is an authored schematic; its fixed straight connections do not show measured model trajectories. An ODE solver’s error depends on the actual field, so straight conditional supervision alone does not permit arbitrary large steps.

Why Straight Paths Matter

For the original conditional OT path, each fixed noise–data pair follows a straight trajectory at constant velocity. This makes its regression target simple. After mixing over data points, the marginal field averages conditional velocities with position-dependent posterior weights. Its trajectories can curve even though every conditional training path is straight.

This authored geometric comparison contrasts fixed straight and curved paths, rather than executing diffusion or a learned flow. Its pixel lengths describe those drawings. Neither training variance nor sampling quality is measured. The original experiments compare specific objectives and solvers; they do not establish a universal step count or guarantee that Flow Matching avoids curves.

The Velocity Field

The velocity field v_\theta(\mathbf{x},t) takes a position and a time and returns a velocity. Unlike conditional supervision, generation does not have access to a particular target \mathbf{x}_1. The learned field must represent how the whole distribution moves.

The arrows above come from an authored analytic field, with hand-defined attraction toward two cluster centers. They are not learned network outputs. For an exact continuous ODE, suitable regularity and uniqueness conditions support an invertible flow. Continuity alone does not establish those conditions, and a finite Euler discretization does not guarantee invertibility or exact reverse recovery.

Training: Conditional Flow Matching

Conditional Flow Matching (CFM) turns fitting the marginal velocity field into regression on tractable conditional examples. Sections 3.1–3.2 establish the connection: Theorem 1 constructs the marginal field with a posterior-weighted mixture, and Theorem 2 gives equal expected parameter gradients for FM and CFM, assuming positive marginal densities.

For the conditional Gaussian OT path in Equations 20–23, draw independent noise and data samples, and a uniform time. Keep a positive residual noise scale:

At t=1 the path reaches \mathbf{x}_1+\sigma_{\min}\mathbf{x}_0, so its distribution approximates the data distribution with residual Gaussian noise. Setting sigma_min to zero gives the familiar interpolation only as a separate limiting case.

This authored training diagram illustrates that zero-noise limit. Its prediction arrow is hand-rotated, and its displayed error is normalized by the squared target magnitude; no network is trained in the diagram. The executable reference below instead uses positive sigma_min and performs real gradient updates.

The loss uses the squared velocity-vector norm, averaged over the batch:

This is velocity regression related to mean squared error; averaging over components as well would change the scale by a factor of the dimension. Conditional targets are constant along each of these paths, but differ between examples. Training samples the conditional path without an ODE solve. The learned field then generates a marginal trajectory using only state and time.

Implementation

The seven operations below keep conditional supervision separate from generation. They execute the positive-noise Gaussian construction in the original paper, then integrate a reduced learned field. Later minibatch pairing methods are discussed afterward.

def draw_endpoints(batch=16, dimension=2, seed=7):
    """Draw independent Gaussian noise, synthetic two-mode data, and time."""
    _positive_int(batch, "batch")
    _positive_int(dimension, "dimension")
    generator = torch.Generator(device="cpu").manual_seed(seed)
    x0 = torch.randn(batch, dimension, generator=generator)
    modes = 2 * torch.randint(0, 2, (batch, 1), generator=generator) - 1
    x1 = 2 * modes + 0.2 * torch.randn(batch, dimension, generator=generator)
    t = torch.rand(batch, 1, generator=generator)
    return x0, x1, t

def conditional_path(x0, x1, t, sigma_min=0.05):
    """Evaluate Eq. 22; at t=1 the endpoint is x1 + sigma_min*x0."""
    _pair(x0, x1)
    _time(x0, t)
    _sigma(sigma_min)
    a = 1 - (1 - sigma_min) * t
    return a * x0 + t * x1

def conditional_target(x0, x1, sigma_min=0.05):
    """Differentiate Eq. 22 along a fixed conditional trajectory."""
    _pair(x0, x1)
    _sigma(sigma_min)
    return x1 - (1 - sigma_min) * x0

class VelocityModel(nn.Module):
    """Reduced field v(x,t); the unknown target x1 is never an input."""

    def __init__(self, dimension=2, hidden=32):
        super().__init__()
        _positive_int(dimension, "dimension")
        _positive_int(hidden, "hidden")
        self.dimension = dimension
        self.layers = nn.Sequential(
            nn.Linear(dimension + 1, hidden),
            nn.Tanh(),
            nn.Linear(hidden, dimension),
        )

    def forward(self, x, t):
        """Map state [B,D] and time [B,1] to velocity [B,D]."""
        _state(x)
        _time(x, t)
        if x.shape[1] != self.dimension:
            raise ValueError("state dimension must match model")
        return self.layers(torch.cat((x, t), dim=1))

def cfm_loss(model, x0, x1, t, sigma_min=0.05):
    """Use Eq. 23's squared vector norm, then average over the batch."""
    xt = conditional_path(x0, x1, t, sigma_min)
    target = conditional_target(x0, x1, sigma_min)
    prediction = model(xt, t)
    if (
        not isinstance(prediction, torch.Tensor)
        or prediction.shape != target.shape
        or prediction.dtype != target.dtype
        or prediction.device != target.device
        or not torch.isfinite(prediction).all()
    ):
        raise ValueError("field must return finite matching [B,D] velocities")
    return (prediction - target).square().sum(dim=1).mean()


def update_field(model, optimizer, x0, x1, t, sigma_min=0.05):
    """Perform a real optimizer step and report the preceding gradient norm."""
    optimizer.zero_grad(set_to_none=True)
    loss = cfm_loss(model, x0, x1, t, sigma_min)
    loss.backward()
    gradients = [p.grad for p in model.parameters()]
    if any(g is None or not torch.isfinite(g).all() for g in gradients):
        raise ValueError("all parameter gradients must be finite")
    grad_norm = torch.sqrt(sum(g.square().sum() for g in gradients))
    optimizer.step()
    return float(loss.detach()), float(grad_norm)

@torch.no_grad()
def euler_integrator(field, x0, steps=32):
    """Integrate dx/dt=field(x,t) with pedagogical Euler step h=1/steps."""
    _state(x0)
    _positive_int(steps, "steps")
    x = x0.clone()
    trajectory = [x.clone()]
    h = 1 / steps
    for k in range(steps):
        t = x.new_full((x.shape[0], 1), k * h)
        velocity = field(x, t)
        if (
            not isinstance(velocity, torch.Tensor)
            or velocity.shape != x.shape
            or velocity.dtype != x.dtype
            or velocity.device != x.device
            or not torch.isfinite(velocity).all()
        ):
            raise ValueError(
                "field must return finite matching [B,D] velocities"
            )
        x = x + h * velocity
        _state(x)
        trajectory.append(x.clone())
    return torch.stack(trajectory)

@torch.no_grad()
def compare_paths(model, x0, x1, sigma_min=0.05, steps=32):
    """Keep conditional training paths separate from field-only sampling."""
    _positive_int(steps, "steps")
    conditional = torch.stack(
        [
            conditional_path(
                x0, x1, x0.new_full((len(x0), 1), k / steps), sigma_min
            )
            for k in range(steps + 1)
        ]
    )
    generated = euler_integrator(model, x0, steps)
    return {"conditional_training": conditional, "generated": generated}

python -m pip install -r requirements-flowmatching.txt
python flowmatching.py --smoke --seed 7

Optimal Transport: Better Pairings

The original conditional Gaussian displacement map is optimal between its two conditional Gaussian distributions. Section 4.1 explicitly warns that this does not make the marginal field a global optimal transport solution. Independent noise–data pairing in the reference above is enough for that original construction.

Later work by Tong et al. studies minibatch transport couplings between sampled noise and data. Hungarian assignment or Sinkhorn methods optimize a specified transport objective within a finite batch; costs and approximation behavior depend on batch size and solver. Nearest-neighbor assignment does not generally solve the global assignment problem, and optimized pairings need not form parallel trajectories.

This authored comparison sorts points by vertical coordinate for an OT-like pairing. It is not a Hungarian or Sinkhorn implementation. The displayed lengths and crossings are measurements of the drawing; no training variance is measured. The preserved atong01/conditional-flow-matching repository is a later CFM/OT implementation associated with this line of work, rather than provenance for the original paper’s implementation.

How Flow Matching Compares

Key Takeaways

1. Conditional geometry and marginal generation differ. Straight Gaussian conditional paths provide simple targets, while learned marginal trajectories may curve and need not be globally optimal.

2. Training is velocity regression. With positive sigma_min the target is \mathbf{x}_1-(1-\sigma_{\min})\mathbf{x}_0. Theorem 2 relates conditional and marginal objectives through their expected parameter gradients, not equality of each sampled loss.

3. Pairing is a separate choice. The original construction uses independent draws. Later minibatch OT can alter training couplings, but does not guarantee parallel paths or improvements for every setting.

4. Generation requires a numerical solver. Exact continuous-flow properties require appropriate ODE assumptions; a small Euler demonstration does not prove accuracy, invertibility, or sample quality.

5. Diffusion paths remain included. Flow Matching provides a simulation-free training objective for a family of probability paths, including diffusion and conditional OT paths.

Impact and Legacy

Flow Matching helped make simulation-free training of continuous normalizing flows practical at scale. Subsequent systems combine related velocity-regression ideas with new architectures, conditioning, and sampling methods. Stable Diffusion 3’s research paper uses rectified flow transformers. SDXL Turbo’s Adversarial Diffusion Distillation paper describes diffusion distillation, so it should not be attributed to the original Flow Matching method.

Applications now extend across image, video, audio, molecular, and other generative modeling domains. Their quality and speed depend on the full training and sampling system; the two-dimensional mechanics here do not reproduce those systems or establish their performance.

The lasting connection is between tractable conditional regression and a distribution-level generative field. Understanding that connection also clarifies why conditional Gaussian OT, later minibatch OT coupling, and learned sampling trajectories must be evaluated separately.

Related Reading

• Attention Is All You Need — The transformer architecture that underlies the neural networks used in Flow Matching
• Deep Residual Learning — Residual connections central to the U-Net and DiT architectures used in flow-based generators
• DINO — Self-supervised visual features often used as perceptual losses for training flow-based image generators
• DINOv2 — Universal visual features that complement flow-based generation with strong discriminative representations