Continuum-marginal optimal transport: a mesh-free kernel method

Abstract: In this paper we study continuum-marginal optimal transport. Given a time-continuous family of probability marginals, the problem is to recover the minimum-energy velocity field whose flow reproduces every marginal. This problem is the continuum limit of the classical two-marginal Benamou--Brenier formulation, and also the deterministic limit of the Nelson problem of stochastic optimal transport. We propose a practical mesh-free solver for this problem. The weak continuity equation is embedded in a reproducing kernel Hilbert space, yielding a sample-only objective that requires no spatial discretization. The velocity is parametrized by any linear-in-parameters dictionary or neural network, and is optimized by mini-batch stochastic methods. Synthetic experiments confirm that the method achieves accurate drift recovery and marginal consistency. The same computational framework also applies to the stochastic Nelson problem.

• The paper introduces a novel mesh-free RKHS method that accurately recovers minimum-energy velocity fields under continuous marginal constraints.
• It employs a kernel-based penalty formulation that enforces the continuity equation, balancing kinetic energy minimization with marginal consistency.
• Numerical experiments on synthetic and real high-dimensional data reveal superior drift estimation compared to traditional two-marginal approaches.

Continuum-Marginal Optimal Transport: A Mesh-Free Kernel Method

Introduction and Problem Formulation

This paper (2604.24226) presents a comprehensive study of the continuum-marginal optimal transport (OT) problem, a natural generalization of the Benamou–Brenier two-marginal dynamic OT to the case when the entire time-continuous family of probability marginals $\{\mu_t\}_{t \in [0, T]}$ is prescribed, not solely the endpoints. The objective is to recover the unique minimum-energy velocity field $u$ whose flow reproduces each marginal, minimizing the quadratic kinetic energy functional subject to the strong requirement that the continuity equation holds at every $t$.

The continuum-marginal OT paradigm connects classical Monge transport, dynamic OT, and the deterministic limit of stochastic Nelson OT (Schrödinger bridge with marginals fixed at all times). The uniqueness and gradient structure of the minimizer, as established by Mikami, guarantee the existence of a potential $\psi$ such that $u^*(t, x) = \nabla_x \psi(t, x)$.

The feasible set under those all-time constraints is strictly smaller than in two-marginal OT, and contains dynamical information not recoverable from endpoints alone. The approach addresses practical needs in fields where intermediate densities are observed but individual paths not tracked (e.g., particle image velocimetry, crowd monitoring, density-functional molecular dynamics).

Methodology: Mesh-Free RKHS Embedding and Sample-Only Objective

The central computational innovation is the mesh-free kernel-based approach for enforcing marginal consistency. The continuity equation constraint is embedded in a reproducing kernel Hilbert space (RKHS); the violation of the weak continuity equation is quantified as an RKHS residual norm $\|H_t^u\|^2$, which admits a sample-only expansion, requiring only kernel evaluations and values of $u$ at sample points.

Figure 1: Experiment~1—learned drift $\hat u(t,x)$ (solid orange) vs.\ true $u^* = 2$ (dashed black) for the Gaussian translation model at multiple time slices.

The velocity field can be parameterized by linear dictionaries or by neural networks and is optimized via stochastic gradient methods, completely obviating the need for spatial discretization. The method’s penalty formulation trades off kinetic energy against violation of the continuity equation, with convergence to the true minimizer in $L^2$ as the penalty $u$0 grows.

Detailed derivations are provided for the sample estimator of the penalty, including closed-form kernel expansions and U-statistics to achieve empirical unbiasedness with $u$1 bias. The approach generalizes to the stochastic Nelson problem, where the Fokker–Planck equation replaces the continuity equation and a Laplacian term is added to the RKHS operator.

Numerical Results: Drift Recovery and Marginal Consistency

A diverse suite of synthetic and real data experiments demonstrates the accuracy and robustness of the estimator.

Synthetic Validation (1D and 2D Modeling)

The method is validated on 1D Gaussian translation, roundtrip (with coinciding endpoints but nontrivial intermediate drift), and bimodal merging flows, including nonlinear drifts not representable by affine models.

Figure 2: Experiment~1—Marginal verification, histograms of ODE-simulated particles (blue) vs.\ true marginals (black) at selected time slices.

Comparisons to baseline methods (Waddington-OT, flow matching, multi-marginal Monge methods) show that the all-time method offers superior velocity field recovery, especially as the number of snapshots increases, where baseline drift estimation becomes noisy and unstable.

For nonlinear flows, neural architectures achieve up to $u$2 lower drift MSE than affine models, confirming the RKHS penalty as an efficacious supervisory signal for drift learning.

Model Expressiveness and Identifiability

Analysis of identifiability valleys in higher dimensions indicates that, while drift grid MSE may underreport estimator quality due to divergence-free perturbations, marginal reconstruction metrics ($u$3, sliced $u$4, MMD) remain tightly controlled, and the variational objective actually penalizes only observationally distinguishable errors.

Figure 3: Experiment~2 (Roundtrip)—drift profile at $u$5 as a function of $u$6. All-time method (orange) recovers nontrivial drift; two-marginal solution (dashed green) is flat.

Figure 4: Experiment~2—Marginal densities at multiple time points. All-time method achieves consistent tracking; two-marginal baseline fails for intermediate times.

Figure 5: Experiment~2—Comparison of drift recovery and marginal tracking between all-time OT and Waddington-OT as the number of snapshots increases.

High-Dimensional and Real-Data Application

Scaling is demonstrated up to $u$7 in the context of single-cell trajectory inference (scRNA-seq) for the embryoid body development dataset. The mesh-free estimator achieves marginal accuracy comparable to Waddington-OT, but yields a globally defined, continuous velocity field.

Figure 6: Experiment~6—Prediction of held-out day-15 cloud in EB scRNA-seq, first two PCA components. All-time OT method overlays predicted cloud, achieving marginal proximity to true.

Stochastic Extension

The computational estimator generalizes straightforwardly to the stochastic Nelson problem: the RKHS operator is augmented to accommodate Laplacian contributions, and SDE simulation replaces ODE propagation. Numerical verification (with closed-form optimal drift) confirms accurate learning of the drift and marginal consistency, with residual error controlled via penalty weight.

Figure 7: Stochastic 1d Gaussian—learned drift $u$8 vs.\ true $u$9, demonstrating mesh-free recovery in the stochastic regime.

Practical and Theoretical Implications

This mesh-free kernel method eliminates dependence on spatial grids, scaling gracefully with dimensionality and offering domain-agnostic application in velocity recovery tasks, including fluid dynamics and crowd modeling. The neural-network parametrization provides universal function approximation within this sample-only framework.

Theoretical guarantees extend to strong convergence in the deterministic case. The methodology bridges generative modeling (all-time constraints vs. endpoint-only constraints), informs improvements in flow matching, and offers new directions for scalable kernel approximations and convergence analysis.

Conclusion

The mesh-free RKHS penalty method for continuum-marginal optimal transport offers a practical, theoretically grounded approach to reconstructing minimum-energy velocity fields from fully observed marginal flows. The method achieves consistent drift recovery and marginal consistency across diverse settings, handles nonlinear and high-dimensional dynamics, and has direct applicability to real-data scenarios. Future work should address adaptive penalty selection, scalable kernel approximations, rigorous convergence rate analysis, and further stochastic extensions.