Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Stream-level flow matching with Gaussian processes

Ganchao Wei, Li Ma

TL;DR

This work extends conditional flow matching (CFM) by introducing stream-level conditioning where latent streams are modeled with Gaussian processes (GPs), yielding GP-CFM. The approach preserves the simulation-free nature of CFM while smoothing the marginal vector field and reducing sampling variance, and it naturally accommodates multiple correlated observations, such as time series. Empirical results on synthetic data, MNIST, CIFAR-10, HWD+, and LFP demonstrate improved sample quality (lower $W_2$, $KID$, and $FID$) and smoother transformations, with covariate conditioning further enhancing performance. The GP-CFM framework is complementary to endpoint-based conditioning (e.g., OT-CFM) and provides a flexible, scalable tool for high-quality generative modeling of CNFs with structured training data.

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.

Stream-level flow matching with Gaussian processes

TL;DR

This work extends conditional flow matching (CFM) by introducing stream-level conditioning where latent streams are modeled with Gaussian processes (GPs), yielding GP-CFM. The approach preserves the simulation-free nature of CFM while smoothing the marginal vector field and reducing sampling variance, and it naturally accommodates multiple correlated observations, such as time series. Empirical results on synthetic data, MNIST, CIFAR-10, HWD+, and LFP demonstrate improved sample quality (lower , , and ) and smoother transformations, with covariate conditioning further enhancing performance. The GP-CFM framework is complementary to endpoint-based conditioning (e.g., OT-CFM) and provides a flexible, scalable tool for high-quality generative modeling of CNFs with structured training data.

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.
Paper Structure (25 sections, 3 theorems, 26 equations, 10 figures, 5 tables, 1 algorithm)

This paper contains 25 sections, 3 theorems, 26 equations, 10 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Notation
  3. Stream-level Flow Matching
  4. A Per-stream Perspective on Flow Matching
  5. Choice of the Stream Model
  6. Numerical Experiments
  7. Adjusting GP Variance for High Quality Samples
  8. Incorporating Multiple Correlated Training Observations
  9. Applications
  10. Variance Reduction
  11. Multiple Training Observations
  12. Conclusion
  13. Bayesian Decision Theoretic Perspective on (stream-level) Flow Matching Training
  14. Discussion on Per-stream Perspective on Flow Matching
  15. Derivation of joint conditional mean and covariance
  16. ...and 10 more sections

Key Result

Proposition 1

The marginal vector field over stream $u_t(x)$ generates the marginal probability path $p_t(x)$ from initial condition $p_0(x)$.

Figures (10)

  • Figure 1: GP streams reduce extrapolation by expanding coverage area. Generate samples of 2-Gaussian mixture from the standard Gaussian. Training observations are shown in red, generated samples in orange, and noise source samples in black. A. FM with straight conditional stream (e.g. I-CFM) may generate "leaky” or outlier samples due to extrapolation errors. The FM method with GP conditional stream has a broader coverage area. B. We train models with I-CFM and GP-I-CFM 100 times and calculate 2-Wasserstein (W2) distance. Among these 100 trained models, generate 1000 samples (orange) and streams (blue) for I-CFM and GP-I-CFM with largest W2 distance (worst case).
  • Figure 2: Change variance over time by tweaking the covariance kernel. Examples of conditional stream between two points, under different variance change scheme.
  • Figure 3: GP streams accommodate correlated points flexibly. A. Paired data with observations on t = 0.5 (red) and t = 1 (orange). B. The generated samples (red for t = 0.5 and orange for t = 1) and streams (blue) for I-CFMs. The I-CFMs contain two separate models trained by I-CFM, t = 0 (standard Gaussian noise) to t = 0.5 and t = 0.5 to t = 1. C. The generated samples for GP-I-CFM.
  • Figure 4: Further conditioning on the starting points helps with stream generation. A. Paired data with observations on three time points: t = 0 (black), t = 0.5 (red) and t = 1 (orange). The two stream cross regions are marked with light blue square. B. The generated samples and streams for GP-I-CFM (without covariate), where the initial points at $t = 0$ are generated from noise using a separate I-CFM. C. The generated samples and streams for GP-I-CFM with covariate using the same starting points, where the neural network is further conditioning on data at $t = 0$.
  • Figure 5: Application to MNIST dataset. We compare the performance of four algorithms (I-CFM, OT-CFM, GP-I-CFM and GP-OT-CFM) on fitting MNIST dataset. We fit the models 100 times for each, and evaluate the quality of the samples by KID and FID. The figures above show the historams of KID and FID.
  • ...and 5 more figures

Theorems & Definitions (6)

  • Proposition 1
  • proof
  • Proposition 2
  • proof
  • Proposition 3
  • proof