cs.LGSep 29, 2026

Improving Function Space Flow Matching with Kernel Optimal Transport

Authors: Fred Xu, Thomas Markovich, Barbora Barancikova, Yizhou Sun

Organizations: University of California, Los Angeles · Block, Inc. · Imperial College London

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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 L2L^2 Sinkhorn, and function-space-aware kernels (signature, Sobolev RBF) give further gains on rough or path-valued data.

Figures & tables

DatasetkFFM (selected)CFM-OT( L2L^{2} )FFMDDPMDDO/NCSNGANOGain (%)
HestonkFFM (RBF/wn) : −1.30×10−4-1.30\times 10^{-4}−8.26×10−5-8.26\times 10^{-5}3.10×10−43.10\times 10^{-4}1.10×10−31.10\times 10^{-3}5.53×10−25.53\times 10^{-2}3.92×10−23.92\times 10^{-2}–
AEMETkFFM (Sig/gp) : −5.10×10−3-5.10\times 10^{-3}−4.11×10−3-4.11\times 10^{-3}−3.24×10−3-3.24\times 10^{-3}−3.99×10−3-3.99\times 10^{-3}2.44×10−12.44\times 10^{-1}5.24×10−15.24\times 10^{-1}–
KdVkFFM (Sob-RBF/gp) : −2.56×10−3-2.56\times 10^{-3}−1.32×10−3-1.32\times 10^{-3}−1.16×10−3-1.16\times 10^{-3}4.84×10−24.84\times 10^{-2}2.37×10−12.37\times 10^{-1}4.12×10−14.12\times 10^{-1}–
EconomykFFM (Sig/gp) : 2.89×10−32.89\times 10^{-3}8.46×10−38.46\times 10^{-3}7.30×10−37.30\times 10^{-3}9.97×10−39.97\times 10^{-3}1.57×10−11.57\times 10^{-1}8.98×10−18.98\times 10^{-1}60.4\mathbf{60.4}
Gene Expr.kFFM (Sig/gp) : 7.11×10−37.11\times 10^{-3}2.80×10−22.80\times 10^{-2}1.78×10−21.78\times 10^{-2}5.32×10−25.32\times 10^{-2}3.42×10−23.42\times 10^{-2}4.17×10−14.17\times 10^{-1}60.1\mathbf{60.1}
Stoch. KdVkFFM (Sob-RBF/gp) : 1.00×10−41.00\times 10^{-4}6.58×10−46.58\times 10^{-4}8.80×10−48.80\times 10^{-4}3.84×10−43.84\times 10^{-4}4.48×10−24.48\times 10^{-2}2.39×10−22.39\times 10^{-2}74.0\mathbf{74.0}
Table 1: Population-level distributional quality measured by MMD-RBF (lower is better; unbiased estimator, so small negative values occur). “kFFM (selected)” is the per-dataset configuration chosen on a held-out validation split by sliced Wasserstein distance (Appendix C.5 ); each cell is tagged (cost/base) with Sig = signature kernel, RBF = Euclidean–RBF kernel, Sob-RBF = Sobolev-RBF kernel, gp = Gaussian-process prior, wn = white noise. The kFFM and FFM columns are the matched-seed runs of Table 2 ( nn as given there); the remaining baselines are averaged over 10 seeds. “Gain (%)” is relative to the strongest non-kFFM baseline in the row and is omitted (–) where that baseline sits at the noise floor of the unbiased estimator (non-positive MMD), where ratios are not meaningful; Table 2 reports paired tests instead.
 Sliced- WWMarginal- W1W_{1}
DatasetFFMkFFMFFMkFFM
Gene Expr. (Sig/gp)0.106±0.0060.106\pm 0.0060.087±0.006\mathbf{0.087\pm 0.006}0.074±0.0050.074\pm 0.0050.059±0.004\mathbf{0.059\pm 0.004}
Economy (Sig/gp)0.0291±0.00190.0291\pm 0.00190.0233±0.0012\mathbf{0.0233\pm 0.0012}0.0171±0.00160.0171\pm 0.00160.0137±0.0009\mathbf{0.0137\pm 0.0009}
Navier–Stokes (RBF/wn)0.3150.3150.081\mathbf{0.081}0.2390.2390.058\mathbf{0.058}
Stoch. NS (RBF/wn)1.061.060.365\mathbf{0.365}0.7950.7950.249\mathbf{0.249}
Table 3: Non-kernel metrics for kFFM (selected configuration in parentheses) vs. FFM on the runs of Table 2 ; lower is better. Sequence rows: mean ±\pm std over 10 shared seeds (Economy averages the population and GDP series within each seed), each comparison significant when paired by seed (Wilcoxon p=0.002p=0.002 ); Navier–Stokes rows are means. Appendix Table 16 adds the autocorrelation and log-spectrum errors.
Figure 1: Turbulent 2D Navier–Stokes ( ν=10−5\nu=10^{-5} , seed 1): pooled vorticity PDF (log scale) and radially averaged energy spectrum. The data PDF is bimodal and heavy-tailed; kFFM (selected) tracks it, DDO/NCSN collapses to a narrow near-Gaussian spike, and DDPM is unimodal with the wrong shape. The flow-matching spectra track the data across scales; DDO/NCSN’s is nearly flat.
Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 2: Qualitative results on AEMET; panels labeled Euclidean, RBF, and Signature are kFFM variants (kFFM-Euc, kFFM-RBF, kFFM-Sig) and Independent is FFM. Kernel OT better preserves the seasonal spread and amplitude of trajectories both at the original resolution and under temporal super-resolution.
Figure 3: Qualitative results on Navier–Stokes; panels labeled Euclidean and RBF are kFFM variants (kFFM-Euc, kFFM-RBF) and Independent is FFM. Kernel OT produces smoother and more coherent flow structures than independent coupling, and this advantage persists under spatial super-resolution.
QuantityMeasurement
End-to-end training overhead vs. FFM55 – 7%7\% at the largest trained resolution ( 1282128^{2} Navier–Stokes)
Peak coupling memory8080 – 151151 MB across benchmarks and b∈{64,…,512}b\in\{64,\dots,512\} ; 1.51.5 – 9%9\% of the FNO forward+backward peak (e.g. 116116 MB vs. 7.627.62 GB on Navier–Stokes 64264^{2} )
Per-batch coupling time11 – 8080 ms, about one FNO step
Bounded RBF vs. unbounded L2L^{2} Sinkhorn ( b=512b=512 )99 vs. 7373 ms per batch ( ≈8×\approx 8\times ); 1515 – 25%25\% faster end-to-end training
Signature-kernel memory0.630.63 GB at 100100 time points vs. 12.512.5 GB at 512512 ; the selection rule therefore assigns RBF/ L2L^{2} costs to very long sequences
Larger-model controlGiving FFM the +6%+6\% budget as extra epochs does not close the gap (loss curves plateau early)
Appendix
Table 15: Measured training cost of the kernel-OT coupling step across benchmarks. The O(b2)O(b^{2}) Sinkhorn solve is independent of model size and resolution.
DatasetMetricFFMkFFM-Sig
Gene Expr.Sliced- WW0.106±0.0060.106\pm 0.0060.087±0.006\mathbf{0.087\pm 0.006}
Marginal- W1W_{1}0.074±0.0050.074\pm 0.0050.059±0.004\mathbf{0.059\pm 0.004}
Autocorr. MSE ( ×10−5\times 10^{-5} )9.3±6.89.3\pm 6.83.3±3.6\mathbf{3.3\pm 3.6}
log-Spectrum MSE ( ×10−2\times 10^{-2} )1.76±0.421.76\pm 0.420.86±0.22\mathbf{0.86\pm 0.22}
EconomySliced- WW0.0291±0.00190.0291\pm 0.00190.0233±0.0012\mathbf{0.0233\pm 0.0012}
Marginal- W1W_{1}0.0171±0.00160.0171\pm 0.00160.0137±0.0009\mathbf{0.0137\pm 0.0009}
Appendix
Table 16: Non-kernel and structural metrics for kFFM-Sig vs. FFM on the sequence datasets (same prior and architecture; lower is better), computed on the matched runs of Table 2 : mean ±\pm std over 10 shared seeds (Economy averages the population and GDP series within each seed). Every comparison is individually significant when paired by seed (Wilcoxon p≤0.027p\leq 0.027 ; kFFM better on 9–10 of the 10 seeds on both datasets); the wide autocorrelation spread on Gene Expression is across-seed variation shared by both methods, which the pairing removes. These runs are distinct from the 10-seed runs behind Table 10 , so the autocorrelation values differ between the two tables; both favor kFFM.
Figure 4: Turbulent 2D Navier–Stokes ( ν=10−5\nu=10^{-5} ): vorticity samples on a shared color scale (for each method, the samples whose enstrophy is closest to the data median). kFFM (selected) and FFM reproduce the data’s large-scale vortical organization, DDO/NCSN produces near-noise fields, and DDPM coarse blobs. The corresponding pooled vorticity PDFs and energy spectra are in Figure 1 .
Figure 5: Entropic minibatch couplings ( b=64b=64 , ϵ=0.01\epsilon=0.01 , median-normalized costs as in training) for three synthetic scenarios (rows) and four costs (columns). Sources and targets are sorted so that the intended matching is the diagonal, and each panel reports the plan mass on intended matches (a uniform plan gives ≈0.14\approx 0.14 – 0.190.19 ). Top: shifted bumps matched by location; every cost recovers the matching. Middle: bumps where 40%40\% of samples carry large, independent nuisance bursts. Burst pairs dominate the median normalization of the unbounded L2L^{2} cost, collapsing clean-pair contrasts below ϵ\epsilon , so its plan degenerates to near uniform ( 0.160.16 ); the bounded RBF cost saturates the bursts and recovers the matching ( 0.580.58 ), invariant to burst amplitude (the conditioning effect of Theorem 1 ). Bottom: Brownian paths f=vW(t)f=vW(t) , v∼U[0.5,2]v\sim U[0.5,2] , matched by volatility. The expected L2L^{2} cost is separable, E∥f−g∥2=(v2+w2)∫t dt\mathbb{E}\|f-g\|^{2}=(v^{2}+w^{2})\int t\,dt , so it carries no assortative signal in expectation and its plan stays near uniform ( 0.250.25 ); the signature kernel (lead-lag) reads volatility through quadratic variation and concentrates on the intended matching ( 0.430.43 ). No single cost wins every scenario, which motivates the per-data-class selection rule of Appendix C.5 .
DatasetHestonAEMETKdVEconomyGene Expr.Stoch. KdVStoch. NSNavier–Stokes
Δ\Delta MMD-RBF (%)+32%\mathbf{+32\%}+53%\mathbf{+53\%}+8%\mathbf{+8\%}+52%\mathbf{+52\%}+64%\mathbf{+64\%}+225%\mathbf{+225\%}+7%\mathbf{+7\%}+9%\mathbf{+9\%}
Appendix
Table 18: Relative MMD-RBF improvement of kFFM over FFM at the end of the convergence curves of Figure 6 (higher is better; values above 100%100\% occur where the unbiased estimate for kFFM is negative). Both methods share the GP prior and FNO backbone, isolating the coupling as the only changed component.
Figure 6: Per-dataset target-metric convergence using MMD-RBF (lower is better) corresponding to the headline summary in Table 18 . The plot shows the best-so-far across-seed mean MMD-RBF over the final phase of training (on KdV, the running mean over the final phase, from epoch 125) for kFFM versus the original FFM under matched GP prior and FNO backbone, so the only change is the coupling. The kFFM curve uses the signature kernel on Stoch. KdV, Gene Expr., AEMET, Economy, and Heston, the Sobolev-RBF cost on KdV, the Euclidean–RBF cost on Stoch. NS, and the raw L2L^{2} cost on Navier–Stokes and Heston-Long. kFFM ends training at a better MMD-RBF on every dataset, with especially large gains on Stoch. KdV, Gene Expr., AEMET, Economy, and Heston. The additional Heston-Long panel is a long-horizon variant of Heston ( 10001000 time steps) used only in this diagnostic and in the sensitivity study; it is not one of the eight benchmark datasets.
Figure 7: Aggregate MMD-RBF sensitivity of kFFM hyperparameters across datasets; the spread is (max⁡−min⁡)/∣mean∣(\max-\min)/|\mathrm{mean}| over the tested settings of the hyperparameter in parentheses, and lower is less sensitive. Signature-kernel settings are nearly invariant on the tested sequence datasets, the Sinkhorn regularization is typically less influential than the RBF bandwidth, and the results depend more on ϵ\epsilon under the raw L2L^{2} (Euclidean) cost than under the bounded RBF cost (median spread of about 8%8\% vs. near zero).
Figure 8: Direct one-at-a-time MMD-RBF sweeps on representative datasets (mean and one standard deviation across seeds). Top: varying the RBF bandwidth σ\sigma at fixed Sinkhorn regularization ( ϵ=0.1\epsilon=0.1 ). Bottom: varying ϵ\epsilon at fixed bandwidth. All arms in this figure use the white-noise base measure: the solid curve (legend: CFM-OT(RBF)) is the Euclidean–RBF kernel coupling, the dashed line is CFM-OT( L2L^{2} ), and the dotted line (legend: No coupling) is independent coupling at the same base. Across the tested values the swept arm moves within its seed-to-seed variation.
Figure 9: Distribution of the pointwise metrics (mean MSE, variance MSE, log-spectrum error) on the two 2D PDE datasets, Navier–Stokes and Stoch. NS, as a function of the length scale of the Gaussian prior.
More in this topic

Flow Matching

Flow matching learns continuous-time vector fields that transport simple noise distributions into data for efficient generative modeling.

See all 656 papers

Explore similar work

CardsList

  1. Optimal-Transport-Guided Functional Flow Matching for Turbulent Field Generation in Hilbert Space

    Apr 7, 2026Li Kunpeng, Wan Chenguang, Qu Zhisong +5Flow-Based Generative ModelingFlow Matching

  2. Potential Matching Optimal Transport: Continuous Normalizing Flows for Exact pp-Wasserstein Dynamics

    Aug 6, 2026Lishuo Zhang, Ruizhi Huang, Yang Yu +1Flow MatchingContinuous Normalizing Flows

    School of Mathematical Sciences, Shanghai Jiao Tong University · Institute of Natural Sciences, MOE-LSC, Shanghai Jiao Tong University

  3. Probabilistic Geodesic Flow Matching on Location-Scale Families

    Sep 28, 2026Zeyuan Yu, Zhi Chang, Shiwei LanFlow MatchingGenerative Modeling

    School of Mathematical & Statistical Sciences Arizona State University Tempe, AZ 85287 USA