Dynamic Conditional Optimal Transport through Simulation-Free Flows

TL;DR

We address conditional sampling for complex or infinite-dimensional problems by developing a dynamic, geometric theory of conditional optimal transport (COT) that extends the Benamou–Brenier formulation. The authors integrate this framework with triangular transport maps and flow matching to produce a simulation-free, flow-based conditional generator (COT-FM) that operates in function spaces and other high-dimensional settings. Theoretical contributions include a conditional Wasserstein space with constant-speed geodesics, a conditional Benamou–Brenier theorem, and a flow-matching scheme that learns a triangular velocity field to interpolate between source and target measures. Empirically, the method competes with or surpasses state-of-the-art COT approaches on 2D synthetic tasks, a Lotka–Volterra inverse problem, and a Darcy flow inverse problem, highlighting its applicability to Bayesian inverse problems and amortized likelihood-free inference.

Abstract

We study the geometry of conditional optimal transport (COT) and prove a dynamical formulation which generalizes the Benamou-Brenier Theorem. Equipped with these tools, we propose a simulation-free flow-based method for conditional generative modeling. Our method couples an arbitrary source distribution to a specified target distribution through a triangular COT plan, and a conditional generative model is obtained by approximating the geodesic path of measures induced by this COT plan. Our theory and methods are applicable in infinite-dimensional settings, making them well suited for a wide class of Bayesian inverse problems. Empirically, we demonstrate that our method is competitive on several challenging conditional generation tasks, including an infinite-dimensional inverse problem.

Figures (12)

• Figure 1: Samples from the ground-truth joint target distribution and the various models. Samples from COT-FM more closely match the ground-truth distribution than the baselines. In the final column, we plot conditional KDEs for samples drawn conditioned on the $y$ value indicated by the dashed horizontal line. See Appendix \ref{['appendix:experiment_details']} for a larger figure and additional results.
• Figure 2: Sample KDEs on the Lotka-Volterra inverse problem. The red lines denote the true parameter values.
• Figure 3: Darcy flow illustration. A true permeability $u$ is shown, as well as the pressure field $\rho$ and its observed, noisy version $y$. We compare an ensemble average of posterior samples from the various methods against MCMC (pCN) 10.1214/13-STS421. COT-FM achieves the lowest MSE to pCN.
• Figure 4: The counterexample in Proposition \ref{['prop:properties_of_distance']}. The measure $\eta_k$ is shown in black and the measure $\nu_k$ is shown in white.
• Figure 5: Samples from the ground-truth joint target distribution and the various models for the 2D datasets. Samples from COT-FM more closely match the ground-truth distribution than the baselines. A common failure mode for the baselines is to generate samples from regions with zero support under the true data distributions. Table \ref{['tab:2d_table']} contains a quantitative evaluation.

Theorems & Definitions (22)

• definition 1: Conditional Wasserstein Space
• definition 2: Conditional $p$-Wasserstein Distance
• Theorem 1: $\P_p^\mu(Y \times U$) is a Geodesic Space
• Theorem 2: Conditional McCann Interpolants
• lemma 1: Triangular Vector Fields Preserve Conditionals
• Theorem 3: Absolutely Continuous Curves in $\P_p^\mu(Y \times U)$
• Theorem 4: Continuous Curves Generated by Triangular Vector Fields
• Theorem 5: Conditional Benamou-Brenier
• proof
• proof