The Mechanics of Flow Matching

Given samples from a distribution with no closed form, we want to learn to sample from . We can model this problem as learning a transport between unpaired samples from a known distribution and . Then to sample from we just have to sample from and apply the transport.

Since samples from and do not have any natural pairing, learning a direct mapping between them is underdetermined. Additional structure is required to learn sufficiently complex . Some ways of supplying this structure are adversarial training, variational autoencoders, optimal transport objectives, and normalizing flows. These methods impose structure in different ways: adversarial methods learn by comparison against a discriminator, variational methods introduce latent-variable structure, and normalizing flows constrain the model to be an invertible transformation with tractable density change. Flow matching takes a different point of view: instead of learning a single map all at once, it specifies a probability path from to and learns the local vector field that transports mass along that path.

Setup

Formally, we define an evolution over time between samples and . This requires two choices: a coupling between and , and a time schedule. We assume the independent coupling

We also choose interpolation schedules and , such that

We can define a probability path such that

This implies a pathwise velocity given by

At time and point , the conditional vector field given endpoint can be written as the conditional expectation of the pathwise velocity:

We can write the associated marginal vector field that transports mass at time and point as

Using the independent coupling above, we have . Then

Thus, the marginal vector field is the posterior average of the conditional vector field over endpoints consistent with the current state . This converts the pathwise Lagrangian view of particle motion (i.e. ) to an Eulerian velocity field at a point in space (i.e. ). This marginal field is the transport field associated with the probability path , which connects to by the continuity equation. Take a smooth compactly supported test function . The expectation of at time is then given as:

Assuming enough regularity to interchange expectation and differentiation with respect to , we get

By the chain rule and then the law of iterated expectations:

This is precisely the weak form of the continuity equation.

Training

To learn the estimator we can minimize the squared error between the estimated and target vector fields. We call this the flow matching objective :

There is a problem here: we do not have target values of the vector field that we can use to minimize this loss. We also cannot compute such targets directly. Even though the conditional field can be computed in closed form for common choices of conditional probability path, such as Gaussian paths built from a Gaussian base distribution, we cannot average against the posterior over endpoints conditioned on current state (i.e. ) since it is generally intractable. We can set up a surrogate loss using the fact that pathwise target velocities are tractable.

Expanding the mean squared error , we get:

Using the Law of Iterated Expectations on the middle term we get:

Here is fully determined by and behaves as a constant when we condition on . This allows us to extract it from the inner expectation. Adding and subtracting we get:

Since the bracketed term is constant in , the gradient of is identical to that of . Therefore and have the same minimizers.

Sampling

Once we’ve learned the estimator of the target vector field, we can draw a sample from and integrate the learned ODE to obtain a sample from :

Using a simple Euler solver this would look like