Kernel optimal transport improves function space flow matching accuracy
Improving Function Space Flow Matching with Kernel Optimal Transport
Machine Learning
Summary
Modeling data that are functions, like time series or physical simulations, is challenging because these data live in infinite-dimensional spaces. The researchers enhanced a method called Functional Flow Matching, which transports a simple starting distribution to match complex function data, by pairing samples more meaningfully using kernel-based optimal transport. This new approach, called kernel Functional Flow Matching, better respects the geometry of function data and leads to more accurate generative models. They tested it on various benchmarks including time series and fluid simulations, where it outperformed prior methods.
What this means in practice
- •For simulation engineers: Generate realistic function-valued simulation outputs such as solutions of PDEs with improved accuracy using kernel-based flow matching.
- •For time series analysts: Model and generate complex time series data by leveraging function-space-aware transport to enhance sample quality and distribution matching.
Authors
Fred Xu, Thomas Markovich, Barbora Barancikova, Yizhou Sun
Abstract
Generative models for function-valued data, such as time series and solutions of partial differential equations, must learn distributions over infinite-dimensional spaces. Functional Flow Matching (FFM) extends Flow Matching to this setting, learning a velocity field whose flow transports a Gaussian prior to the data distribution, but it inherits the independent endpoint pairing of standard Flow Matching: in each batch, prior and data samples are matched arbitrarily, so the conditional bridge must traverse both the shared global structure of the dataset and instance-specific residuals. In function space this is harder to fix than in finite dimensions, since optimal transport (OT) on function spaces is delicate to formulate and a flat Euclidean surrogate ignores the geometry that distinguishes function-valued data. We propose kernel Functional Flow Matching (kFFM), which replaces the independent pairing by entropic OT under a kernel-induced cost, the coupling underlying the Hilbert Sinkhorn Divergence (HSD), leaving the FFM neural-operator architecture unchanged. We prove that the kernel cost and the HSD objective are uniformly bounded and well-posed on Banach ambient spaces, derive an error decomposition against quadratic-cost OT on compact metric spaces that isolates an irreducible kernel-cost mismatch term, and prove a discretization-invariance bound whose rate is governed by Sobolev regularity. Empirically, kFFM improves distributional matching over FFM, diffusion, adversarial, and finite-dimensional OT baselines on time-series and PDE benchmarks, with significant paired-seed gains over FFM and improvements that persist under non-kernel and physics-based diagnostics, including a turbulent Navier-Stokes benchmark. Bounded kernel costs already outperform raw $L^2$ Sinkhorn, and function-space-aware kernels (signature, Sobolev RBF) give further gains on rough or path-valued data.