Stream-level Flow Matching with Gaussian Processes

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.

Flow matching (FM) is a family of fast and powerful methods for training deep generative models that learn how to transform draws from a simple source distribution into those from a highly complex target distribution. While the effectiveness of FM has been demonstrated in a range of applications, there are a few technical limitations that can lead to a decay in sample quality and occasionally lead to generation of outliers that do not accurately reflect the characteristics of the target distribution. We recognize one of the key limitations related to how the paths connecting draws from the source to the target are constructed in existing FM algorithms, and propose an extension that incorporates Gaussian processes to model more flexible, stochastic paths. The properties of Gaussian processes make the resulting algorithm more robust, enable integration of multiple correlated distributions, all while retaining the computational efficiency of FM algorithms.