Improving Function Space Flow Matching with Kernel Optimal Transport

2026-09-29

AI Summary

This study addresses the inefficient transport paths and lack of geometric structure caused by independent endpoint pairing in generative models over function spaces. To this end, we propose Kernel Flow Matching, which replaces random pairing with entropic optimal transport under the Hilbert-Sinkhorn divergence. By introducing a kernel-induced cost, the method enables infinite-dimensional optimal transport to improve function distribution learning. Theoretically, we establish target boundedness and discretization invariance while isolating irreducible error terms. Empirically, our approach significantly outperforms baselines such as FFM on time series forecasting and PDE benchmarks, and its effectiveness is further validated in modeling turbulent Navier-Stokes equations.

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.

Problem

• Function-valued data
• Functional Flow Matching
• Optimal transport
• Infinite-dimensional spaces
• Generative models

Innovation

• Functional Flow Matching
• Kernel Optimal Transport
• Hilbert Sinkhorn Divergence
• Neural Operator
• Discretization Invariance