Flow Matching

Train a velocity field to transport noise to data along straight interpolation paths, a simulation-free objective that generalizes diffusion and samples well in few steps.

Equation

Linear (rectified/OT) path between noise $x_0\sim \mathcal{N}(0,I)$ and data $x_1$:

$x_t=(1-t)x_0+t{x}_1,t\in [0,1]$

The conditional velocity is constant along each path:

$u_t(x_t\mid x_0,x_1)=\frac{d{x}_t}{dt}=x_1-x_0$

Objective (conditional flow matching), whose gradient matches that of the intractable marginal objective:

$\mathcal{L}={\mathbb{E}}_{t,x_0,x_1}\left[\Vert v_{\theta }(x_t,t)-(x_1-x_0){\Vert }^2\right]$

Sampling integrates the ODE $\dot{x}=v_{\theta }(x,t)$ from $t=0$ to $1$.

From scratch

import torch
 
def fm_loss(model, x1):                                 # x1: (b, ...) data
    x0 = torch.randn_like(x1)                           # noise
    t = torch.rand(x1.shape[0], device=x1.device)
    tb = t.view(-1, *[1] * (x1.dim() - 1))              # (b, 1, 1, 1)
    xt = (1 - tb) * x0 + tb * x1
    return (model(xt, t) - (x1 - x0)).pow(2).mean()     # target velocity
 
@torch.no_grad()
def euler_sample(model, shape, steps=25, device='cpu'):
    x = torch.randn(shape, device=device)               # t = 0
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((shape[0],), i * dt, device=device)
        x = x + model(x, t) * dt                        # explicit Euler
    return x                                            # t = 1, data

Notes

• Equivalent to diffusion under a reparameterization: the linear path corresponds to a specific variance-preserving-like schedule, and $v$-style targets relate the two. The practical difference is schedule and loss weighting, not the model class.
• Straighter paths mean lower discretization error, so Euler with 10 to 30 steps is usually enough; rectified flow reflows the model on its own samples to straighten further, reaching 1 to 4 steps.
• Sampling timestep distribution matters: logit-normal $t$ sampling (emphasizing mid $t$) noticeably beats uniform at scale.
• Sign/direction conventions vary ($t=0$ noise vs $t=0$ data); check which end is which before porting a sampler.