Improving Function Space Flow Matching with Kernel Optimal Transport
Abstract
Generative models for function-valued data, such as time series and solutions to partial differential equations, must learn distributions over infinite-dimensional spaces rather than over finite-dimensional vectors. Functional Flow Matching (FFM) is a recent extension of Flow Matching to this setting, learning a velocity field whose flow transports a Gaussian prior to the data distribution. However, FFM inherits a structural limitation from standard Flow Matching: in each training batch, prior and data samples are paired independently, so the conditional bridge between them must simultaneously traverse the shared global structure of the dataset and instance-specific residuals. This limitation is more consequential in function space than in finite dimensions: directly formulating optimal transport on function spaces is technically delicate, and a flat Euclidean surrogate ignores the function-space geometry that distinguishes function-valued data. We propose kernel Functional Flow Matching (kFFM), which replaces the independent pairing by entropic optimal transport in a kernel-induced Hilbert space via the Hilbert Sinkhorn Divergence, leaving the FFM neural-operator architecture unchanged. Theoretically, we prove the HSD objective is well-posed on Banach ambient spaces with bounded kernels, derive compact-metric approximation bounds against quadratic-cost OT with an explicit kernel-cost mismatch term, and quantify the gap between the function-space objective and the truncated representation computed on a grid, at a rate governed by data regularity. Empirically, kFFM consistently improves MMD-RBF distributional matching over FFM, diffusion, adversarial, and finite-dimensional OT baselines on time-series and PDE benchmarks, with function-space-aware kernels emerging as the strongest choices on every PDE and path-valued benchmarks.