Stream-level flow matching with Gaussian processes

TL;DR

This paper introduces a novel stream-level flow matching method using Gaussian processes to model latent stochastic paths, improving sample quality and variance reduction in continuous normalizing flows.

Contribution

It extends conditional flow matching by incorporating Gaussian process streams, enabling better variance control and handling correlated data in flow-based models.

Findings

Reduces variance in marginal vector field estimation.

Improves sample quality in image and time series applications.

Effectively models correlated data with Gaussian process streams.

Abstract

Flow matching (FM) is a family of training algorithms for fitting continuous normalizing flows (CNFs). Conditional flow matching (CFM) exploits the fact that the marginal vector field of a CNF can be learned by fitting least-squares regression to the conditional vector field specified given one or both ends of the flow path. In this paper, we extend the CFM algorithm by defining conditional probability paths along ``streams'', instances of latent stochastic paths that connect data pairs of source and target, which are modeled with Gaussian process (GP) distributions. The unique distributional properties of GPs help preserve the ``simulation-free" nature of CFM training. We show that this generalization of the CFM can effectively reduce the variance in the estimated marginal vector field at a moderate computational cost, thereby improving the quality of the generated samples under common metrics. Additionally, adopting the GP on the streams allows for flexibly linking multiple correlated training data points (e.g., time series). We empirically validate our claim through both simulations and applications to image and neural time series data.