Diffusion Sampling from Flow Matching Models

Flow matching models learn the time-dependent vector field that transports samples from a known distribution to samples from a desired distribution . Let denote the density of the random variable . Assuming affine interpolation between the particle endpoints and , and a Gaussian source distribution , we can show that a trained flow-matching model can be sampled stochastically using a family of diffusion SDEs.

Let’s use scalar interpolation schedules satisfying and :

We assume that and are differentiable and define the schedule determinant . We call the schedules nondegenerate at time when .

Conditioning on the endpoint gives a Gaussian conditional path:

The pathwise velocity is

For interior times at which , conditioning additionally on determines the source sample:

The corresponding conditional vector field is therefore

Assuming is positive definite, the Gaussian conditional path also gives us a closed form for the conditional score. We define

The conditional training objectives recover the corresponding marginal fields through posterior averaging:

Both posterior averages can be expressed using the posterior mean of the data endpoint, which we call the denoiser:

Taking conditional expectations of the affine path and its velocity gives

Eliminating the posterior mean of gives

Similarly, posterior averaging of the conditional Gaussian score gives the marginal score

and therefore

Thus, at interior times where , , and , the marginal velocity, denoiser, and marginal score are equivalent parameterizations:

The velocity can be learned with the conditional flow matching loss, while the score can be learned with the conditional score matching loss. The denoiser can likewise be learned by regressing directly onto with a squared error loss.

Now that we have a score function, for any choice of diffusion coefficient yielding a well-posed SDE, we can compute the drift for a corresponding diffusion process.

To see why this SDE preserves the same probability path as the original flow ODE, consider its probability current:

Here we used . The Fokker-Planck equation for the SDE therefore reduces to the same continuity equation as the ODE:

When initialized from the same , the ODE and every valid SDE in this family consequently have the same marginal distribution at every time . They do not have the same individual trajectories: the ODE trajectory is deterministic given its initial state, whereas the SDE trajectory remains stochastic.

These noise-to-data dynamics can be viewed as the reverse-time dynamics of a corresponding forward noising process, allowing us to construct stochastic paths transporting to exactly as we would with a diffusion model. Note how the choice of has an impact on the drift term since depends on which, in turn, depends on . Therefore we can think of the additional dependent terms in the drift as corrections to the ODE drift in the stochastic setting.

The animations below compare the deterministic flow ODE with isotropic, anisotropic, and state-dependent SDEs that share its marginal probability path. The arrows can display the ODE velocity , the SDE drift , or the correction between them. JAX code for simulations is available in this gist.