Title: Dynamic Conditional Optimal Transport through Simulation-Free Flows
Venue: NeurIPS 2024
Accepted: yes
Keywords: flow matching

optimal transport

generative models

conditional generation
Abstract: We study the geometry of conditional optimal transport (COT) and prove a dynamic 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.

Dynamic Conditional Optimal Transport through Simulation-Free Flows

Abstract

We study the geometry of conditional optimal transport (COT) and prove a dynamic 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.

1 Introduction

Many fundamental tasks in machine learning and statistics may be posed as modeling a conditional distribution ν ( u | y ) where obtaining an analytical representation of ν ( u | y ) is impractical. While sampling-based approaches such as Markov Chain Monte Carlo (MCMC) methods are useful, they suffer from several limitations. First, MCMC requires numerous likelihood evaluations, rendering it prohibitively expensive in scientific and engineering applications where the likelihood is determined by an expensive numerical simulator. Second, MCMC must be run anew for every observation y , which is impractical in applications such as Bayesian inverse problems [Dashti and Stuart, 2013] and generative modeling [Mirza and Osindero, 2014]. These limitations motivate the need for a likelihood-free [Cranmer et al., 2020] and amortized [Amos et al., 2023] approach. While methods like ABC [Beaumont, 2010] and variational inference [Blei et al., 2017] partially address these challenges, they are difficult to scale to high dimensions or have limited flexibility.

Recently, generative models such as normalizing flows [Papamakarios et al., 2019, 2021], GANs [Ramesh et al., 2022], and diffusion models [Sharrock et al., 2022] have shown promise in amortized and likelihood-free inference. These models may be viewed in the framework of measure transport [Baptista et al., 2020], where samples u ∼ η ( u ) from a tractable source distribution are transformed by a mapping T ( y, u ) such that the transformed samples are approximately distributed as ν ( u | y ) . One way to achieve this is through triangular mappings [Baptista et al., 2020, Spantini et al., 2022], where a joint source distribution η ( y, u ) is transformed by a mapping of the form T : ( y, u ) ↦→ ( T Y ( y ) , T U ( y, u )) . Under suitable assumptions, if T transforms the source η ( y, u ) into the target ν ( y, u ) , then T U ( y, -) couples the conditionals η ( u | y ) and ν ( u | y ) .

Typically, such a map T is not unique [Wang et al., 2023], and a natural idea is thus to regularize the transport and search for an admissible mapping that is in some sense optimal. In other words, learning a conditional sampler may be phrased as finding a conditional optimal transport (COT) map. While there exists some work on learning COT maps, these approaches often rely on a difficult adversarial optimization problem [Baptista et al., 2020, Hosseini et al., 2023, Bunne et al., 2022, Ray et al., 2022] or simulating from the model during training [Baptista et al., 2023, Wang et al., 2023]. In this work, we propose a conditional generative model for likelihood-free inference based on a dynamic formulation of conditional optimal transport. Specifically, our contributions are as follows:

We develop a general theoretical framework for dynamic conditional optimal transport in separable Hilbert spaces. Our framework is applicable in infinite-dimensional spaces, enabling applications in function space Bayesian inference. In Section 4, we study the conditional Wasserstein space P µ p ( Y × U ) and show that this space admits constant speed geodesics between any two measures. In Section 5, we characterize the absolutely continuous curves of measures in P µ p ( Y × U ) via the continuity equation and triangular vector fields. As a consequence, we obtain conditional generalizations of the McCann interpolants [McCann, 1997] and the Benamou-Brenier Theorem [Benamou and Brenier, 2000].

In Section 6, we propose COT flow matching (COT-FM), a simulation-free flow-based model for conditional generation. This model directly leverages our theoretical framework, where we learn to model a path of measures interpolating between an arbitrary source and target distribution via a geodesic in the conditional Wasserstein space.

In Section 7, we demonstrate our method on several challenging conditional generation tasks. We apply our method to two Bayesian inverse problems - one arising from the Lotka-Volterra dynamical system, and an infinite-dimensional problem arising from the Darcy Flow PDE. Our method shows competitive performance against recent COT methods.

2 Related Work

Conditional Optimal Transport. Conditional Optimal Transport (COT) remains relatively underexplored in both machine learning and related fields. Recent approaches learn static COT maps via input convex networks [Bunne et al., 2022, Wang et al., 2023] or normalizing flows [Wang et al., 2023]. In addition, there have been a number of heuristic approaches to conditional simulation through W-GANs [Sajjadi et al., 2017, Adler and Öktem, 2018, Kim et al., 2022, 2023], for which Chemseddine et al. [2023] provide a rigorous basis. Closely related to our work are those which employ triangular plans [Carlier et al., 2016, Trigila and Tabak, 2016], which have been modeled through GANs in Euclidean spaces [Baptista et al., 2020] and function spaces [Hosseini et al., 2023]. In contrast, we use a novel dynamic formulation of COT, which we model through a generalization of flow matching [Lipman et al., 2022, Albergo et al., 2023b, Liu et al., 2022]. This allows us to use flexible architectures while avoiding the difficulties of training GANs [Arora et al., 2018].

Simulation-Free Continuous Normalizing Flows. Flow matching [Lipman et al., 2022] (and the closely related stochastic interpolants [Albergo et al., 2023a] and rectified flows [Liu et al., 2022]) are a class of methods for building continuous-time normalizing flows in a simulation-free manner. Notably, these works do not approximate an optimal transport between the source and target measures. Pooladian et al. [2023] and Tong et al. [2023] propose instead to couple the source and target distributions via optimal transport, leading to marginally optimal paths. In this work, we study an extension of these techniques for conditional generation.

While some works [Davtyan et al., 2023, Gebhard et al., 2023, Isobe et al., 2024, Wildberger et al., 2024] have applied flow matching for conditional generation, these approaches do not employ COT. Notably, these approaches are limited to the finite-dimensional setting, whereas our method adds to the growing literature on function-space generative models [Hosseini et al., 2023, Kerrigan et al., 2023, 2024, Lim et al., 2023, Franzese et al., 2024]. Concurrent work by Chemseddine et al. [2024] develops the foundation of dynamic COT with applications to flow matching, and concurrent work by Barboni et al. [2024] develop the theory of dynamic COT for the purposes of studying infinitely deep ResNets. However, these works focus on the Euclidean setting, whereas our methods are applicable in general separable Hilbert spaces.

3 Background and Notation

Let X,X ′ represent arbitrary separable Hilbert spaces, equipped with the Borel σ -algebra. We use P ( X ) to represent the space of Borel probability measures on X , and P p ( X ) ⊆ P ( X ) to represent the subspace of measures having finite p th moment. If η ∈ P ( X ) is a probability measure on X and T : X → X ′ is measurable, then the pushforward measure T # η ( -) = η ( T -1 ( -)) is a probability measure on X ′ . Maps of the form e.g. π X : X × X ′ → X represent the canonical projection.

We assume that we have two separable Hilbert spaces of interest. The first, Y , is a space of observations, and the second, U , is a space of unknowns. These spaces may be of infinite dimensions, but a case of practical interest is when Y and U are finite dimensional Euclidean spaces. We will consider the product space Y × U , equipped with the canonical inner product obtained via the sum of the inner products on Y and U , under which the space Y × U is also a separable Hilbert space. Let η ∈ P ( Y × U ) be a joint probability measure. The measures π Y # η ∈ P ( Y ) and π U # η ∈ P ( U ) obtained via projection are the marginals of η . We use η y ∈ P ( U ) to represent the measure obtained by conditioning η on the value y ∈ Y . By the disintegration theorem [Bogachev and Ruas, 2007, Chapter 10], such conditional measures exist and are essentially unique, in the sense that there exists a Borel set E ⊆ Y with π Y # η ( E ) = 0 , and the η y are unique for y / ∈ E .

3.1 Static Conditional Optimal Transport

In conditional optimal transport, we are given a target measure ν ∈ P ( Y × U ) and some source measure η ∈ P ( U ) . We seek a transport map T : Y × U → U such that, for any given y ∈ Y , the mapping T ( y, -) : U → U transforms the source distribution η into the conditional distribution ν y , i.e., T ( y, -) # η = ν y . In a sense, T can be thought of as a collection of transport maps indexed by y ∈ Y . If such a T were available, by drawing samples u 0 ∼ η and transforming them, one would obtain samples T ( y, u 0 ) ∼ ν y . Solving this transport problem for each fixed y is expensive at best, or impossible when only has a single (or no) samples ( y, u ) ∼ ν for any given y . Thus, one must leverage information across different observations y . To that end, recent work has focused on the notion of triangular mappings T : Y × U → Y × U [Hosseini et al., 2023, Baptista et al., 2020] of the form T ( y, u ) = ( T Y ( y ) , T U ( T Y ( y ) , u )) for some T Y : Y → Y and T U : Y × U → U . Triangular mappings are of interest as they allow us to obtain conditional couplings from joint couplings.

Proposition 1 (Theorem 2.4 [Baptista et al., 2020], Prop. 2.3 [Hosseini et al., 2023]) Suppose η, ν ∈ P ( Y × U ) and T : Y × U → Y × U is triangular. If T # η = ν , then T U ( T Y ( y ) , -) # η y = ν T Y ( y ) for π Y # η -almost every y .

In many scenarios of practical interest, the source measure η and the target measure ν have the same Y -marginals. We will henceforth make this assumption, and use µ = π Y # η = π Y # ν to represent this marginal. In this case, we may take T Y to be the identity mapping, so that the conclusion of Proposition 1 simplifies to T U ( y, -) # η y = ν y for µ -almost every y . We note that in situations where such an assumption does not hold, one may simply preprocess the source measure η via an invertible mapping T Y satisfying [ T Y ] # [ π Y # η ] = π Y # ν [Hosseini et al., 2023, Prop 3.2].

Given a source and target measures η, ν ∈ P µ ( Y × U ) and a cost function c : ( Y × U ) 2 → R , the conditional Monge problem seeks to find a triangular mapping solving

The conditional Monge problem also admits a relaxation under which one only considers couplings whose Y -components are almost surely equal. To that end, for η, ν ∈ P µ p ( Y × U ) we define the set of triangular couplings Π Y ( η, ν ) to be the couplings of η and ν that almost surely fix the Y -components,

In other words, a triangular coupling γ ∈ Π Y ( η, ν ) has samples ( y 0 , u 0 , y 1 , u 1 ) ∼ γ such that y 0 = y 1 almost surely. The conditional Kantorovich problem seeks a triangular coupling solving

Hosseini et al. [2023] prove the existence of minimizers to the conditional Kantorovich and Monge problems under very general assumptions. Moreover, optimal couplings to the conditional Kantorovich problem induce optimal couplings for µ -almost every conditional measure. We refer to Appendix B and Hosseini et al. [2023] for further details.

4 Conditional Wasserstein Space

Motivated by our discussion on triangular transport maps, we introduce the conditional Wasserstein spaces, consisting of joint measures with finite p th moments and having fixed Y -marginals µ . Interestingly, Gigli [2008, Chapter 4] studies the same space for the purposes of constructing geometric tangent spaces in the usual Wasserstein space.

Definition 1 (Conditional Wasserstein Space)

Suppose µ ∈ P ( Y ) is given and 1 ≤ p < ∞ . The conditional p -Wasserstein space is

We now equip P µ p ( Y × U ) with a metric W µ p , the conditional Wasserstein distance. Intuitively, the conditional Wasserstein distance measures the usual Wasserstein distance between all of the conditional distributions in expectation under the fixed Y -marginal µ .

Definition 2 (Conditional p -Wasserstein Distance)

Suppose

η, ν

is the conditional p -Wasserstein distance. W p is the usual Wasserstein distance for measures on U .

By Jensen's inequality we have W µ p ( η, ν ) ≥ E y ∼ µ [ W p ( η y , ν y )] . In the following, we show that the conditional Wasserstein distance is a well-defined metric as well as a few other metric properties. µ

Let 1 ≤ p < ∞ .

Proposition 2 (Some Properties of W p )

(a) W µ p is well-defined, finite, and equals the minimal conditional Kantorovich cost.

(b) W µ p is a metric on the space P µ p ( Y × U ).

(c) There does not exist C > 0 such that W µ p ( η, ν ) ≤ CW p ( η, ν ) for all η, ν ∈ P µ p ( Y × U ) .

Proposition 2(c, d) together shows that one should expect the topology generated by W µ p to be stronger than the unconditional distance W p . Here, we note that Gigli [2008] and Chemseddine et al. [2023] previously showed that W µ p is a metric through an equivalence with restricted couplings. Our approach builds on the results of Hosseini et al. [2023] and is somewhat more direct, and hence our proofs may be of independent interest.

For the sake of concreteness, we include an example where the conditional 2 -Wasserstein distance may be explicitly computed. Here, the necessary calculations follows from the fact that the conditional distributions of a multivariate are again Gaussian, and Gaussian distributions admit a closed-form expression for the usual unconditional 2 -Wasserstein distance.

Example: Gaussian Measures. Suppose Y = U = R and that η, ν ∈ P µ p ( Y × U ) are Gaussians

It follows that µ = π Y # η = π Y # ν = N (0 , 1) . As η y , ν y are Gaussians, their W 2 distance admits a closed form and we can directly compute the expectation in Equation (5) to obtain W µ, 2 2 ( η, ν ) = 2(1 -√ 1 -ρ 2 ) . This is zero if and only if ρ = 0 , i.e. η = ν . However, π U # η = π U # ν = N (0 , 1) and W 2 ( π U # η, π U # ν ) = 0 regardless of ρ . Moreover, the unconditional distance is W 2 2 ( η, ν ) = 2 ( 2 - √ 1 -ρ - √ 1 + ρ ) , from which it is easy to verify that W 2 ( η, ν ) ≤ W µ 2 ( η, ν ) . See Appendix C for a similar derivation which applies to arbitrary Gaussians.

Conditional Wasserstein Space as a Geodesic Space. We now turn our attention to the geodesics in P µ p ( Y × U ) . In particular, we show that there exists a constant speed geodesic between any two measures in P µ p ( Y × U ) , generalizing a similar result in the unconditional setting [Santambrogio, 2015, Theorem 5.27]. Moreover, we show that under suitable regularity assumptions, solutions to the conditional Monge problem (1) induce constant speed geodesics. Our motivation for studying geodesics in P µ p ( Y × U ) is practical - in Section 6, we show how one can model geodesics in P µ p ( Y × U ) in order to obtain a conditional flow-based model whose paths are easy to integrate.

A curve is a continuous function γ · : I → P µ p ( Y × U ) where I = ( a, b ) ⊆ R is any open interval of finite length. A curve is absolutely continuous if there exists m ∈ L 1 (( a, b )) such that

If ( γ t ) is an absolutely continuous curve, then its metric derivative

exists for almost every t ∈ ( a, b ) , and, moreover, we almost surely have | γ ′ | ( t ) ≤ m ( t ) pointwise for any m satisfying Equation (7) [Ambrosio et al., 2005, Theorem 1.1.2]. A curve ( γ t ) is called a constant speed geodesic if for all a < s ≤ t < b , we have W µ p ( γ s , γ t ) = | t -s | W µ p ( γ a , γ b ) . It is straightforward to show that every constant speed geodesic is absolutely continuous.

For any η, ν ∈ P µ p ( Y × U ) , there exists a constant speed geodesic between η and ν .

When an optimal triangular coupling γ ⋆ ∈ Π Y ( η, ν ) is induced by an injective triangular map T ⋆ , we may recover a constant speed geodesic in P µ p ( Y × U ) , generalizing the McCann interpolant [McCann, 1997] to the conditional setting. We refer to Proposition 4 for sufficient conditions on η, ν under which such a T ⋆ exists. Informally, samples from ( y 0 , u 0 ) ∼ η flow in a straight path at a constant speed to their destination T ⋆ ( y 0 , u 0 ) .

Theorem 2 (Conditional McCann Interpolants)

Fix η, ν ∈ P µ p ( Y × U ) . Suppose T ⋆ ( y, u ) = ( y, T ⋆ U ( y, u )) is an injective triangular map solving the conditional Monge problem (1) . Define the maps T t : Y × U → Y × U for 0 ≤ t ≤ 1 via T t = (1 -t ) I + tT ⋆ , and define the curve of measures γ t = [ T t ] # η ∈ P γ p ( Y × U ) . Then,

(a) ( γ t ) is absolutely continuous and a constant speed geodesic between η, ν

(b) The vector field v t ( T ⋆ t ( y, u )) = (0 , T ⋆ U ( y, u ) -u ) generates the path γ t , in the sense that ( γ t , v t ) solve the continuity equation (9) .

5 Conditional Benamou-Brenier Theorem

In this section, we prove a characterization of the absolutely continuous curves in P µ p ( Y × U ) . As a corollary, we obtain a conditional generalization of the Benamou-Brenier Theorem [Benamou and Brenier, 2000], giving us a dynamic characterization of the conditional Wasserstein distance. Roughly speaking, all such curves are generated by a vector field on Y × U which has zero velocity in the Y component. This is natural, as all measures in P µ p ( Y × U ) have a fixed Y -marginal µ . Such a vector field can be informally seen as tangent to a curve of measures, and is the dynamic analogue of the triangular maps discussed in Section 3. More formally, given an open interval I ⊆ R , a time-dependent Borel vector field v : I × Y × U → Y × U is said to be triangular if there exists a Borel vector field v U : I × Y × U → U such that v t ( y, u ) = ( 0 , v U t ( y, u ) ) .

Continuity Equation. We introduce some necessary background which allows us to link vector fields to curves of measures. The continuity equation ∂ t γ t + div ( v t γ t ) = 0 describes the evolution of a measure γ t which flows along a given vector field v t [Ambrosio et al., 2005, Chapter 8]. This equation must be understood in the sense of distributions, i.e. for every φ in a space of test functions,

Weconsider cylindrical test functions φ ∈ Cyl ( Y × U × I ) , i.e. of the form φ ( y, u, t ) = ψ ( π d ( y, u ) , t ) where π d : Y × U → R d maps ( y, u ) ↦→ ( ⟨ ( y, u ) , e 1 ⟩ , . . . , ⟨ ( y, u ) , e d ⟩ ) and { e 1 , e 2 , . . . , e d } is any orthonormal family in Y × U . In the finite dimensional setting, one may take φ ∈ C ∞ c ( Y × U ) to be smooth and compactly supported [Ambrosio et al., 2005, Remark 8.1.1].

In Appendix E, we prove Lemma 1, which is key in proving Theorem 4 below. Informally, Lemma 1 states that if the weak continuity equation (9) is satisfied for a joint distribution and triangular vector field, then the continuity equation is also satisfied for the corresponding conditional distributions and U components of the vector field.

Lemma 1 (Triangular Vector Fields Preserve Conditionals)

Suppose v t ( y, u ) = (0 , v U t ( y, u )) is triangular and that ( γ t ) ⊂ P µ p ( Y × U ) is a path of measures such that ( v t , γ t ) satisfy the continuity equation in the sense of distributions. Then, it follows that for µ -almost every y ∈ Y , we have ∂ t γ y t + ∇· ( v U t ( y, -) γ y t ) = 0 .

We note that having v t be triangular is sufficient, but certainly not necessary, for the conditional continuity equation to almost surely hold. For instance, the vector field in R d that rotates N (0 , I ) about the origin is not triangular yet preserves all conditional distributions.

Absolutely Continuous Curves. In this section, we state our characterization of absolutely continuous curves in P µ p ( Y × U ) . Informally, given such a curve, Theorem 3 provides us with a triangular vector field which generates the curve, in the sense that the pair solve the continuity equation.

Theorem 3 (Absolutely Continuous Curves in P µ p ( Y × U ) ) Let I ⊂ R be an open interval, and suppose γ t : I → P µ p ( Y × U ) is an absolutely continuous in the W µ p metric with | γ ′ | ( t ) ∈ L 1 ( I ) . Then, there exists a Borel vector field v t ( y, u ) such that

(a) v t is triangular

(b) v t ∈ L p ( γ t , Y × U ) and ∥ v t ∥ L p ( γ t ,Y × U ) ≤ | γ ′ | ( t ) for a.e. t

(c) ( v t , γ t ) solve the continuity equation in the sense of distributions.

Conversely, we show in Theorem 4 that if the pair ( γ t , v t ) solve the continuity equation and v t is triangular, then the curve ( γ t ) is absolutely continuous and | γ ′ | ( t ) ≤ ∥ v t ∥ L p ( γ t ,Y × U ) . The main technique of this result is to study the collection of conditional continuity equations (which is feasible by Lemma 1) and to apply the converse of Ambrosio et al. [2005, Theorem 8.3.1]. In this setting, the infinite-dimensional result is obtained via a finite-dimensional approximation argument.

Theorem 4 (Continuous Curves Generated by Triangular Vector Fields) Suppose that γ t : I → P µ p ( Y × U ) is narrowly continuous and ( v t ) is a triangular vector field such that ( γ t , v t ) solve the continuity equation with ∥ v t ∥ L p ( γ t ,Y × U ) ∈ L 1 ( I ) . Then, γ t : I → P µ p ( Y × U ) is absolutely continuous in the W µ p metric and | γ ′ | ( t ) ≤ ∥ v t ∥ L p ( µ , Y × U ) for almost every t .

As a corollary of Theorem 3 and Theorem 4, we obtain a conditional version of the Benamou-Brenier theorem [Benamou and Brenier, 2000]. Once we have our characterization of absolutely continuous curves provided by these theorems, the proof of Theorem 5 largely follows the unconditional case (see e.g. Ambrosio et al. [2005, Chapter 8]), but we include it for the sake of completeness.

Theorem 5 (Conditional Benamou-Brenier)

Let 1 < p < ∞ . For any η, ν ∈ P µ p ( Y × U ) , we have

6 COT Flow Matching

We have thus far seen that the COT problem (3) admits a dynamic formulation by Theorem 5, where one may take the underlying vector fields to be triangular. We use these results to design a principled model for conditional generation based on flow matching [Lipman et al., 2022, Albergo et al., 2023b, Liu et al., 2022, Tong et al., 2023, Pooladian et al., 2023]. We hereafter use the squared-distance cost (i.e. p = 2 ).

Y

Checkerboard

True

COT-FM (Ours)

PCP-Map

COT-Flow

Conditional

U

Y

Circles

U

U

U

True

COT-FM

PCP-Map

COT-Flow

--一--

Flow Matching. We assume that we have access to samples z 0 = ( y 0 , u 0 ) ∼ η ( y 0 , u 0 ) ∈ P µ p ( Y × U ) from a source measure, and samples z 1 = ( y 1 , u 1 ) ∼ ν ( y 1 , u 1 ) ∈ P µ p ( Y × U ) from a target measure. Let z = ( z 0 , z 1 ) ∼ ρ ( z 0 , z 1 ) ∈ Π( η, ν ) be any coupling of the source and target measure. We specify a collection of measures and vector fields on Y × U via

where C is any trace-class covariance operator [Da Prato and Zabczyk, 2014]. As is standard in flow matching [Lipman et al., 2022, Kerrigan et al., 2024], we obtain from Equations (10) a marginal measure γ t ( y, u ) and vector field v t ( y, u ) satisfying the continuity equation via

This marginal path ( γ t ) 1 t =0 interpolates between the source measure ( t = 0 ) and a smoothed version of the target measure ( t = 1 ). To transform source samples from η into target samples from ν , we seek to learn the intractable vector field v t ( y, u ) with a model v θ ( t, y, u ) by minimizing the loss 1

which has the same θ -gradient as the MSE loss to the true vector field u t ( y, u ) [Tong et al., 2023].

COT Flow Matching. In the preceding section, ρ ( z ) may be an arbitrary coupling between η and ν . Motivated by Proposition 1, we will choose ρ to be a COT coupling. Under sufficient regularity conditions (see Appendix B), this COT plan will be induced by a triangular map. In turn, Theorem 2 gives us that this triangular map is generated by a triangular vector field of the form (10). Thus, we parametrize our model u θ to also be triangular. Moreover, we recover the optimal dynamic transport given in Theorem 5 as Tr ( C ) → 0 by a pointwise application of [Tong et al., 2023, Proposition 3.4].

Given a collection of samples { z i 0 , z i 1 } n i =1 drawn from η and ν , we approximate a conditional optimal coupling ρ using standard numerical techniques with the cost function c ϵ ( y 0 , u 0 , y 1 , u 1 ) = | y 1 -y 0 | 2 + ϵ | u 1 -u 0 | 2 for some 0 < ϵ ≪ 1 . Intuitively, such a cost penalizes mass transfer along the Y dimension, which is precisely the constraint sought in the COT problem (3). As ϵ ↓ 0 , we recover the true optimal triangular map [Carlier et al., 2010, Hosseini et al., 2023]. The COT coupling

can either be precomputed for small datasets or computed on each minibatch drawn during training. While the use of minibatches is a computational necessity, we find that surprisingly small batch sizes still yields accurate approximations of the true COT mapping using our COT-FM method. See Appendix G.

After training, we obtain a learned triangular vector field v θ ( t, y, u ) . Given an arbitrary fixed y ∈ Y , we may approximately sample from the target ν ( u | y ) by sampling u 0 ∼ η ( u 0 | y ) and numerically solving the corresponding flow equation ∂ t ( y, u t ) = v θ ( t, y, u t ) with initial condition ( y, u 0 ) .

Source Measure. Our framework is agnostic to the choice of source measure η , allowing for flexibility in the modeling process. The main requirement is that the Y -marginals of the source η and target η must match. In some scenarios, this is trivially satisfied. If one is interested in using a source distribution which is simply random noise, one may take η ( y 0 , u 0 ) = π Y # ν ( y 0 ) ⊗ η U ( u 0 ) to be the product of two independent distributions where η U is arbitrary, e.g. Gaussian noise.

7 Experiments

We now illustrate our methodology (COT-FM) on a variety of conditional simulation tasks. We compare our method against several competitive baselines, namely PCP-Map [Wang et al., 2023], COT-Flow [Wang et al., 2023], and WaMGAN [Hosseini et al., 2023]. These baselines are chosen as they reflect current state-of-the-art approaches to learning COT maps. We additionally compare against flow matching [Lipman et al., 2022, Wildberger et al., 2024] without COT, i.e. where the coupling between the source and target measures is the independent coupling ρ ( z 0 , z 1 ) = η ⊗ ν . This baseline serves as an ablation for the COT component of our model.

Table 2: Statistical distances between MCMC and posterior samples u ∼ ν ( u | y ) for each method on the LV dataset. Average results ± one standard deviation reported across five test sets.

Overall, our method (COT-FM) typically outperforms these baselines across the diverse and challenging set of tasks we consider. We find that PCP-Map [Wang et al., 2023] is a strong baseline, but we emphasize that this model relies on the use of an input-convex neural network [Amos et al., 2017] and it is hence unclear how to adapt this method to e.g. images. Appendix F contains further details and results for all of our experiments. 2

2D Synthetic Data. We first consider synthetic distributions where Y = U = R . Our source measure is taken to be the independent product η ( y, u ) = π Y # ν ⊗N (0 , 1) . We plot ground-truth joint distributions and samples for two datasets in Figure 1. See Appendix F for additional results. Samples from our method (COT-FM) closely match those from the ground-truth distribution, whereas samples from PCP-Map and COT-Flow [Wang et al., 2023] can produce samples in regions of zero support under the ground-truth distribution. In Table 1, we provide a quantitative analysis, where we measure the W 2 and MMD distances between the generated and ground-truth joint distributions. This is motivated by Proposition 1, as triangular maps which couple the joint distributions necessarily couple the conditional distributions. Our method outperforms the baselines across all metrics.

Lotka-Volterra (LV) Dynamical System. Here we estimate parameters of the Lotka-Volterra (LV) model given only noisy observations of its solution. The LV model has parameters u = ( α, β, γ, δ ) ∈ R 4 ≥ 0 and a pair of coupled nonlinear ODEs of the form

whose solution p ( t ) = ( p 1 ( t ) , p 2 ( t )) ∈ R 2 ≥ 0 represents the number of prey and predator species at time t ∈ [0 , T ] . Following Alfonso et al. [2023], we assume p (0) = (30 , 1) and that log( u ) ∼ N ( m, 0 . 5 I ) with m = ( -0 . 125 , -3 , -0 . 125 , -3) . Given parameters u ∈ R 4 ≥ 0 , we simulate Equation (13) for t ∈ { 0 , 2 , . . . , 20 } to obtain a solution z ( u ) ∈ R 22 ≥ 0 . An observation y ∈ R 22 ≥ 0 is obtained by the addition of log-normal noise, i.e. log( y ) ∼ N (log( z ( u ) , 0 . 1 I ) . We thus may simulate many ( y, u ) pairs from the target measure for training.

As a benchmark, we follow the settings of Alfonso et al. [2023] and choose parameters u = (0 . 83 , 0 . 041 , 1 . 08 , 0 . 04) to generate a single observation y as described above. In Figure 2, we plot a histogram of 10 , 000 samples from the posterior ν ( u | y ) of COT-FM.

Since the ground-truth posterior is intractable, we compare against differential evolution Metropolis MCMC [Braak, 2006]. Samples from our method qualitatively resemble those from MCMC, and the posterior mode is typically close to the true unknown u (shown in red). Our method is quantitatively closest to the MCMC samples in the W 2 metric, and competitive in the MMDmetric (Table 2).

Darcy Flow Inverse Problem. Here we consider an infinitedimensional Bayesian inverse problem from the 2D Darcy flow PDE. The setting is adapted from Hosseini et al. [2023]. We opt to compare against WaMGAN [Hosseini et al., 2023], as this is currently the only other extant amortized function-space COT method, and FFM [Kerrigan et al., 2023] as a function-space flow matching ablation.

0

.

50

0

.

83

1

.

25

α

0

0.041

0.1

β

0.8

1.08

1.2

γ

0.02

0.04

0.06

δ

MCMC

COT-FM

The Darcy flow PDE is an elliptic equation on a smooth domain Ω ⊆ R d which relates a permeability field exp( u ) , a pressure field ρ , and a source term f via -div exp( u ) ∇ ρ = f on Ω subject to ρ = 0 on ∂ Ω . Our goal is to recover the permeability u from noisy measurements y of the pressure ρ . Both the unknown u and observations y are functions and thus infinite-dimensional. To define our target measure, we specify a prior ν ( u ) = N (0 , C ) with a Matérn kernel C of lengthscale ℓ = 1 / 2 and ν = 3 / 2 . Given u ∼ η ( u ) , the Darcy flow PDE is solved numerically

[Alnæs et al., 2015] to obtain a solution ρ ( u ) observed at some finite but arbitrary number of points { x 1 , . . . , x n } ⊂ R 2 . An observation y ( u ) is obtained by adding Gaussian noise to each observation, i.e. y ( u ) ∼ N ( ρ ( u ) , σ 2 I ) where σ = 2 . 5 × 10 -2 . We implement all models via a Fourier Neural Operator [Li et al., 2020], allowing us to work with arbitrary discretizations, as required by the functional nature of this problem.

We provide an illustration in Figure 3. As the true posterior is intractable, we compare against preconditioned Crank-Nicolson (pCN) [Cotter et al., 2013], a function-space MCMC method. In Figure 3, we plot the mean posteriors obtained from the various methods. Qualitatively, both COT-FM and FFM are good approximations to pCN, while WaMGAN has visual artifacts. However, the MSE

True (

u

)

Pressure (

ρ

)

Observed (

y

)

pCN

MSE: 21.27

COT-FM (Ours)

MSE: 22.83

FFM

MSE: 25.52

WaMGAN

MSE: 2.21

MSE: 3.49

MSE: 3.63

MSE: 4.40

MSE: 4.66

MSE: 6.63

between our method and the pCN mean is lower than that of FFM. Table 3 provides a quantitative comparison between the methods on a test set of 5,000 samples, where we measure MSE and CRPS [Hersbach, 2000]. We compare the ensemble mean of 10 samples against the true u value as running pCN for each observation is prohibitively expensive. COT-FM outperforms FFM and WaMGAN in terms of MSE and is on-par with FFM in terms of CRPS. See Appendix F for further details.

8 Conclusion

We analyze conditional optimal transport from a geometric and dynamical point of view. Our analysis culminates in the characterization of absolutely continuous curves of measures in a conditional Wasserstein space, resulting in a conditional analog of the Benamou-Brenier Theorem.

We use these result to build on the framework of triangular transport and flow matching to develop simulation-free methods for conditional generative models. Our methods are applicable across a wide class of problems, and we demonstrate our methodology on several challenging inverse problems.

Limitations and Broader Impacts. A limitation COT-FM is that computing the full COT plan can be expensive for large datasets, necessitating the use of minibatch approximations potentially resulting in sub-optimal plans. While this approximation does not limit the practical applicability of our method, an interesting challenge is to characterize the precise relationship between this minibatch approximation and the full COT plan. Moreover, computing the COT plan incurs a small additional computational cost compared to standard flow matching. As with all generative models, a potential negative impact is the potential for disinformation through generated samples being purported as real.

Acknowledgments and Disclosure of Funding

This research was supported by the Hasso Plattner Institute (HPI) Research Center in Machine Learning and Data Science at the University of California, Irvine, by the National Science Foundation under award 1900644, and by the National Institutes of Health under award R01-LM013344.

References

Oriol Abril-Pla, Virgile Andreani, Colin Carroll, Larry Dong, Christopher J Fonnesbeck, Maxim Kochurov, Ravin Kumar, Junpeng Lao, Christian C Luhmann, Osvaldo A Martin, et al. PyMC: A modern, and comprehensive probabilistic programming framework in python. PeerJ Computer Science , 9:e1516, 2023.

Jonas Adler and Ozan Öktem. Deep Bayesian inversion. arXiv preprint arXiv:1811.05910 , 2018.

Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 , 2023a.

Michael S Albergo, Mark Goldstein, Nicholas M Boffi, Rajesh Ranganath, and Eric Vanden-Eijnden. Stochastic interpolants with data-dependent couplings. arXiv preprint arXiv:2310.03725 , 2023b.

Jason Alfonso, Ricardo Baptista, Anupam Bhakta, Noam Gal, Alfin Hou, Isa Lyubimova, Daniel Pocklington, Josef Sajonz, Giulio Trigila, and Ryan Tsai. A generative flow for conditional sampling via optimal transport. arXiv preprint arXiv:2307.04102 , 2023.

Martin Alnæs, Jan Blechta, Johan Hake, August Johansson, Benjamin Kehlet, Anders Logg, Chris Richardson, Johannes Ring, Marie E Rognes, and Garth N Wells. The FEniCS project version 1.5. Archive of numerical

software , 3(100), 2015.

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: In Metric Spaces and in the Space of Probability Measures . Springer Science & Business Media, 2005.

Luigi Ambrosio, Alberto Bressan, Dirk Helbing, Axel Klar, Enrique Zuazua, Luigi Ambrosio, and Nicola Gigli. A user's guide to optimal transport. Modelling and Optimisation of Flows on Networks: Cetraro, Italy 2009, Editors: Benedetto Piccoli, Michel Rascle , pages 1-155, 2013.

Brandon Amos, Lei Xu, and J Zico Kolter. Input convex neural networks. In International Conference on Machine Learning , pages 146-155. PMLR, 2017.

Brandon Amos et al. Tutorial on amortized optimization. Foundations and Trends in Machine Learning , 16(5): 592-732, 2023.

Sanjeev Arora, Andrej Risteski, and Yi Zhang. Do GANs learn the distribution? Some theory and empirics. In International Conference on Learning Representations , 2018.

Ricardo Baptista, Bamdad Hosseini, Nikola B Kovachki, and Youssef Marzouk. Conditional sampling with monotone GANs: from generative models to likelihood-free inference. arXiv preprint arXiv:2006.06755 , 2020.

Ricardo Baptista, Youssef Marzouk, and Olivier Zahm. On the representation and learning of monotone triangular transport maps. Foundations of Computational Mathematics , pages 1-46, 2023.

Raphaël Barboni, Gabriel Peyré, and François-Xavier Vialard. Understanding the training of infinitely deep and wide ResNets with conditional optimal transport. arXiv preprint arXiv:2403.12887 , 2024.

Mark A Beaumont. Approximate Bayesian computation in evolution and ecology. Annual review of ecology, evolution, and systematics , 41:379-406, 2010.

Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the Monge-Kantorovich mass transfer problem. Numerische Mathematik , 84(3):375-393, 2000.

David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

Vladimir Igorevich Bogachev and Maria Aparecida Soares Ruas. Measure Theory , volume 2. Springer, 2007.

Cajo JF Ter Braak. A Markov chain Monte Carlo version of the genetic algorithm differential evolution: easy bayesian computing for real parameter spaces. Statistics and Computing , 16:239-249, 2006.

Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised training of conditional monge maps. Advances in Neural Information Processing Systems , 35:6859-6872, 2022.

Guillaume Carlier, Alfred Galichon, and Filippo Santambrogio. From knothe's transport to Brenier's map and a continuation method for optimal transport. SIAM Journal on Mathematical Analysis , 41(6):2554-2576, 2010.

Guillaume Carlier, Victor Chernozhukov, and Alfred Galichon. Vector quantile regression: An optimal transport approach. The Annals of Statistics , 44(3):1165 - 1192, 2016. doi: 10.1214/15-AOS1401. URL https: //doi.org/10.1214/15-AOS1401 .

Jannis Chemseddine, Paul Hagemann, and Christian Wald. Y-Diagonal couplings: Approximating posteriors with conditional Wasserstein distances. arXiv preprint arXiv:2310.13433 , 2023.

Jannis Chemseddine, Paul Hagemann, Christian Wald, and Gabriele Steidl. Conditional Wasserstein distances with applications in Bayesian OT flow matching. arXiv preprint arXiv:2403.18705 , 2024.

L. Cotter, G. O. Roberts, A. M. Stuart, and D. White. MCMC methods for functions: Modifying old algorithms to make them faster. Statistical Science , 28(3):424 - 446, 2013.

Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-based inference. Proceedings of the National Academy of Sciences , 117(48):30055-30062, 2020.

Giuseppe Da Prato and Jerzy Zabczyk. Stochastic Equations in Infinite Dimensions . Cambridge University Press, 2014.

Masoumeh Dashti and Andrew M Stuart. The Bayesian approach to inverse problems. arXiv preprint arXiv:1302.6989 , 2013.

Aram Davtyan, Sepehr Sameni, and Paolo Favaro. Efficient video prediction via sparsely conditioned flow matching. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 2326323274, 2023.

Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Z. Alaya, Aurélie Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamonjy, Ievgen Redko, Antoine Rolet, Antony Schutz, Vivien Seguy, Danica J. Sutherland, Romain Tavenard, Alexander Tong, and Titouan Vayer. POT: Python optimal transport. Journal of Machine Learning Research , 22(78):1-8, 2021. URL http://jmlr.org/papers/v22/20-451.html .

Giulio Franzese, Giulio Corallo, Simone Rossi, Markus Heinonen, Maurizio Filippone, and Pietro Michiardi. Continuous-time functional diffusion processes. Advances in Neural Information Processing Systems , 36, 2024.

Timothy D Gebhard, Jonas Wildberger, Maximilian Dax, Daniel Angerhausen, Sascha P Quanz, and Bernhard Schölkopf. Inferring atmospheric properties of exoplanets with flow matching and neural importance sampling. arXiv preprint arXiv:2312.08295 , 2023.

Nicola Gigli. On the geometry of the space of probability measures in Rn endowed with the quadratic optimal transport distance . PhD thesis, Scuola Normale Superiore, 2008.

Hans Hersbach. Decomposition of the continuous ranked probability score for ensemble prediction systems. Weather and Forecasting , 15(5):559-570, 2000.

Bamdad Hosseini, Alexander W Hsu, and Amirhossein Taghvaei. Conditional optimal transport on function spaces. arXiv preprint arXiv:2311.05672 , 2023.

Noboru Isobe, Masanori Koyama, Kohei Hayashi, and Kenji Fukumizu. Extended flow matching: a method of conditional generation with generalized continuity equation. arXiv preprint arXiv:2402.18839 , 2024.

Gavin Kerrigan, Justin Ley, and Padhraic Smyth. Diffusion generative models in infinite dimensions. In Proceedings of The 26th International Conference on Artificial Intelligence and Statistics , volume 206 of Proceedings of Machine Learning Research , pages 9538-9563. PMLR, 2023.

Gavin Kerrigan, Giosue Migliorini, and Padhraic Smyth. Functional flow matching. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics , volume 238 of Proceedings of Machine Learning Research , pages 3934-3942, 2024.

Young-geun Kim, Kyungbok Lee, and Myunghee Cho Paik. Conditional wasserstein generator. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2022.

Young-geun Kim, Kyungbok Lee, Youngwon Choi, Joong-Ho Won, and Myunghee Cho Paik. Wasserstein geodesic generator for conditional distributions. arXiv preprint arXiv:2308.10145 , 2023.

Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self-normalizing neural networks. Advances in neural information processing systems , 30, 2017.

Nikola B. Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M. Stuart, and Anima Anandkumar. Neural operator: Learning maps between function spaces. CoRR , abs/2108.08481, 2021.

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895 , 2020.

Jae Hyun Lim, Nikola B Kovachki, Ricardo Baptista, Christopher Beckham, Kamyar Azizzadenesheli, Jean Kossaifi, Vikram Voleti, Jiaming Song, Karsten Kreis, Jan Kautz, et al. Score-based diffusion models in function space. arXiv preprint arXiv:2302.07400 , 2023.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations , 2022.

Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003 , 2022.

Robert J McCann. A convexity principle for interacting gases. Advances in Mathematics , 128(1):153-179, 1997.

Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv preprint arXiv:1411.1784 , 2014.

George Papamakarios, David Sterratt, and Iain Murray. Sequential neural likelihood: Fast likelihood-free inference with autoregressive flows. In The 22nd international conference on artificial intelligence and statistics , pages 837-848. PMLR, 2019.

George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research , 22(57): 1-64, 2021.

Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research , 12:2825-2830, 2011.

Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304.14772 , 2023.

Poornima Ramesh, Jan-Matthis Lueckmann, Jan Boelts, Álvaro Tejero-Cantero, David S Greenberg, Pedro J Gonçalves, and Jakob H Macke. Gatsbi: Generative adversarial training for simulation-based inference. arXiv preprint arXiv:2203.06481 , 2022.

Deep Ray, Harisankar Ramaswamy, Dhruv V Patel, and Assad A Oberai. The efficacy and generalizability of conditional gans for posterior inference in physics-based inverse problems. arXiv preprint arXiv:2202.07773 , 2022.

Mehdi SM Sajjadi, Bernhard Scholkopf, and Michael Hirsch. Enhancenet: Single image super-resolution through automated texture synthesis. In Proceedings of the IEEE international conference on computer vision , pages 4491-4500, 2017.

Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY , 55(58-63):94, 2015.

Louis Sharrock, Jack Simons, Song Liu, and Mark Beaumont. Sequential neural score estimation: Likelihoodfree inference with conditional score based diffusion models. arXiv preprint arXiv:2210.04872 , 2022.

Alessio Spantini, Ricardo Baptista, and Youssef Marzouk. Coupling techniques for nonlinear ensemble filtering. SIAM Review , 64(4):921-953, 2022.

Alexander Tong, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based generative models with minibatch optimal transport. In ICML Workshop on New Frontiers in Learning, Control, and Dynamical Systems , 2023.

Giulio Trigila and Esteban G Tabak. Data-driven optimal transport. Communications on Pure and Applied Mathematics , 69(4):613-648, 2016.

Cédric Villani et al. Optimal Transport: Old and New , volume 338. Springer, 2009.

Zheyu Oliver Wang, Ricardo Baptista, Youssef Marzouk, Lars Ruthotto, and Deepanshu Verma. Efficient neural network approaches for conditional optimal transport with applications in Bayesian inference. arXiv preprint arXiv:2310.16975 , 2023.

Jonas Wildberger, Maximilian Dax, Simon Buchholz, Stephen Green, Jakob H Macke, and Bernhard Schölkopf. Flow matching for scalable simulation-based inference. Advances in Neural Information Processing Systems , 36, 2024.

Appendix: Table of Contents

A Optimal Transport

We provide here a brief and informal overview of optimal transport in the standard unconditional setting. For more details, we refer to the standard references of Villani et al. [2009], Santambrogio [2015], Ambrosio et al. [2005]. Let X be a separable metric space and fix a cost function c : X × X → R ∪ { + ∞} . Suppose we have two Borel measures η, ν ∈ P ( X ) . The Monge problem seeks to find a measurable transport map T : X → X minimizing the expected cost of transport, i.e. corresponding to the optimization problem

This optimization problem is challenging, though, as it involves a nonlinear constraint and the set of feasible maps may be empty. In contrast, the Kantorovich problem is a relaxation which seeks to find an optimal coupling γ ∈ Π( η, ν ) , i.e. a probability distribution over X × X with marginals η, ν , which solves

Under fairly weak conditions (e.g., the cost is lower semicontinuous and bounded from below [Ambrosio et al., 2013, Theorem 2.5]), minimizers to the Kantorovich problem are guaranteed to exist. If the cost function is c ( x 0 , x 1 ) = | x 0 -x 1 | p for some 1 < p < ∞ , under sufficient regularity conditions on η a solution T ⋆ to the Monge problem is guaranteed to exist and, moreover, the coupling γ ⋆ = ( I, T ⋆ ) # η is optimal for the Kantorovich problem. See Ambrosio et al. [2013, Chapter 2] and Ambrosio et al. [2005, Theorem 6.2.10].

Wasserstein Space. In the special case that c ( x 0 , x 1 ) = | x 0 -x 1 | p for 1 ≤ p < ∞ , and η, ν ∈ P p ( X ) , the Kantorovich problem admits a finite-cost solution. The cost of such an optimal coupling is the p -Wasserstein distance

which, as the name suggests, is a metric on the space P p ( X ) [Ambrosio et al., 2005, Section 7.1] [Santambrogio, 2015, Section 5.1]. The Wasserstein distance admits a dynamical formulation via the Benamou-Brenier theorem [Benamou and Brenier, 2000]. Namely, the p -Wasserstein distance can be obtained by finding a time-dependent vector field transforming η to ν across time t ∈ [0 , 1] with minimal energy:

Here, we constrain our minimization problem over the set of measures and vector fields ( γ t , v t ) interpolating between η and ν , satisfying a continuity equation (see Section 5). In Section 4, we study a generalization of the Wasserstein distances for conditional optimal transport problems. In particular, Theorem 5 provides a generalization of the Benamou-Brenier theorem to the conditional setting which recovers a conditional Wasserstein distance.

B Conditional Optimal Transport

This section contains additional discussion regarding the static COT problem, supplementing Section 3.1. We refer to Hosseini et al. [2023], Baptista et al. [2020], and Chemseddine et al. [2023] for further results and details.

Given a source and target measures η, ν ∈ P µ p ( Y × U ) , the conditional Monge problem seeks to find a triangular mapping solving

The conditional Monge problem also admits a relaxation under which one only considers couplings whose Y -components are almost surely equal. To that end, we consider the subset C ⊂ ( Y × U ) 2 whose Y components are identical, i.e.,

and we define the set of ( Y ) -restricted probability measures R Y ⊂ P ( ( Y × U ) 2 ) such that every γ ∈ R Y is concentrated on C . In other words, if γ ∈ R Y , then samples ( y 0 , u 0 , y 1 , u 1 ) ∼ γ have y 0 = y 1 almost surely. In addition, for any η, ν ∈ P ( Y × U ) , we define the set of triangular couplings Π Y ( η, ν ) to be the probability measures in R Y whose marginals are η and ν , i.e.

The conditional Kantorovich problem seeks a triangular coupling γ ⋆ solving

Hosseini et al. [2023] prove the existence of minimizers to the conditional Kantorovich and Monge problems under very general assumptions. Moreover, optimal couplings to the conditional Kantorovich problem induce optimal couplings for µ -almost every conditional measure. Assuming sufficient regularity assumptions on the conditional measures, unique solutions to the conditional Monge problem exist. We restate these results here for the sake of completeness.

Fix η, ν ∈ P µ ( Y × U ) . Suppose the cost function c is continuous, inf c > -∞ , and there exists a finite cost coupling γ ∈ Π Y ( η, ν ) . Then, the conditional Kantorovich problem admits a minimizer γ ⋆ . Moreover, γ ⋆,y 0 ( y 1 , u 0 , u 1 ) = ˆ γ ⋆,y 0 ( u 0 , u 1 ) δ ( y 1 -y 0 ) where for µ -almost every y the measure γ ⋆,y is an optimal coupling for η y , ν y under the cost c y ( u 0 , u 1 ) = c ( y, u 0 , y, u 1 )

Proposition 4 (Prop 3.8 [Hosseini et al., 2023])

Fix 1 < p < ∞ and η, ν ∈ P µ p ( Y × U ) . Suppose c ( y 0 , u 0 , y 1 , u 1 ) = | u 0 -u 1 | p . If η y assign zero measure to Gaussian null sets for µ -almost every y , then there is a unique solution T ⋆ to the conditional Monge problem, and γ ⋆ = ( I, T ⋆ ) # η is the unique solution to the conditional Kantorovich problem. If ν y also assign zero measure to Gaussian null sets for µ -almost every y , then T ⋆ is injective η -almost everywhere.

C Closed-Form Conditional Wasserstein Distance for Gaussian Measures

In this section, we provide additional details and results regarding the closed-form conditional Wasserstein distance for Gaussian distributions. See Section 4 in the main paper.

Suppose Y = R d and U = R d ′ are Euclidean spaces (of possibly different dimensions), and that η, ν ∈ P µ p ( Y × U ) are Gaussians of the form

where m ∈ R d , m η u , m ν u ∈ R d ′ , and the (block) covariance matrices are Σ ∈ R d × d , Λ η , Λ ν ∈ R d × d ′ , and Σ η u , Σ ν u ∈ R d ′ × d ′ .

This form is chosen to ensure that η and ν have equal Y -marginals. It follows that µ = π Y # η = π Y # ν = N ( m, Σ) . Let

We have that the conditionals η y , ν y are available in closed-form:

Thus, for any fixed y , we use the known closed-form unconditional Wasserstein distance to obtain

We now take an expectation over y ∼ µ = N ( m, Σ) to compute W µ, 2 2 . Observe that R ( y -m ) ∼ N (0 , R Σ R T ) and that E y ∼ µ [ | R ( y -m ) | 2 ] = Tr ( R Σ R T ) . Thus,

This form, perhaps unsurprisingly, closely resembles the unconditional Wasserstein distance between two Gaussians, except for the presence of an additional Tr ( R Σ R T ) term. Note that when η, ν have uncorrelated Y, U components, we precisely recover W 2 2 ( π U # η, π U # ν ) as one may expect. As a special case of interest, if Y = U = R and

then we obtain as a special case of Equation (26) that W µ, 2 2 ( η, ν ) = 2(1 -√ 1 -ρ 2 ) . This is zero if and only if ρ = 0 , i.e. η = ν .

D Proofs: Section 4

In this section, we provide detailed proofs of our claims in Section 4, regarding the metric properties of the conditional Wasserstein space.

D.1 Metric Properties

We first note that W µ p ( η, ν ) may be viewed as the minimal value of the constrained Kantorovich problem in Equation (3) when one takes the cost to be the metric on the space Y × U . Similar results, relating the conditional Wasserstein distance to triangular couplings, have appeared previously, but our proof is independent of these prior works [Chemseddine et al., 2023, Gigli, 2008].

Proposition 5 (Equivalent Formulation of the Conditional Wasserstein Distance) Fix η, ν ∈ P µ p ( Y × U ) and 1 ≤ p < ∞ . Then, W µ p ( η, ν ) is well-defined, finite, and

where W µ,p p ( η, ν ) represents the p -th power of the conditional p -Wasserstein distance.

Proof. The cost function d p is clearly continuous and non-negative, and hence by Proposition 3 it suffices to exhibit a finite-cost coupling γ ∈ Π Y ( η, ν ) between η and ν . Indeed, take the conditionally independent coupling

which is clearly in Π Y ( η, ν ) . We then have that

Hence, Equation (30) admits a minimizer γ ⋆ ∈ Π Y ( η, ν ) . By Proposition 3, this minimizer may be taken to have the form γ ⋆ = γ ⋆,y 1 ( u 0 , u 1 ) δ ( y 1 -y 0 ) µ ( y 1 ) where γ ⋆,y 1 ( u 0 , u 1 ) is µ ( y 1 ) -almost surely an optimal coupling between η y 1 , ν y 1 for the cost | u 1 -u 0 | p . Thus,

Here, we emphasize that the µ -almost sure uniqueness of the disintegrations of η, ν along Y result in a well-defined expression.

Moreover, if η ∈ P µ p ( Y × U ) it follows that η y ∈ P p ( U ) for µ -a.e. y , because

Thus all considered p -Wasserstein distances on U are finite.

We now proceed to prove several metric properties of our distance.

Proposition 2 (Some Properties of W µ p )

Let 1 ≤ p < ∞ .

(a) W µ p is well-defined, finite, and equals the minimal conditional Kantorovich cost.

(b) W µ p is a metric on the space P µ p ( Y × U ).

(c) There does not exist C > 0 such that W µ p ( η, ν ) ≤ CW p ( η, ν ) for all η, ν ∈ P µ p ( Y × U ) .

(d) For all η, ν ∈ P µ p ( Y × U ) , W p ( π U # η, π U # ν ) ≤ W µ p ( η, ν ) and W p ( η, ν ) ≤ W µ p ( η, ν ) .

Proof. Part (a). This is simply a restatement of Proposition 5.

Part (b). Fix η, ν, ρ ∈ P µ p ( Y × U ) . Since W p is a metric on P p ( U ) , we immediately obtain the symmetry of W µ p . Moreover, we have that W µ p ( η, ν ) = 0 if and only if η y = ν y for µ -almost every y . Thus, if W µ p ( η, ν ) = 0 and E ⊆ Y × U is Borel measurable,

which shows that η = ν . Here, E y = { u | ( y, u ) ∈ E } is the y -slice of E . Conversely, if η = ν , then η y = ν y up to a µ -null set by the essential uniqueness of disintegrations. Thus, W µ p ( η, ν ) = 0 if and only if η = ν .

By Minkowski's inequality and the triangle inequality for W p on P p ( U ) , we see

̸

Part (c). We provide a counterexample. Fix any u 0 = 0 ∈ U and y 0 , y 1 ∈ Y such that y 0 = y 1 . Define µ = 1 2 ( δ y 0 + δ y 1 ) . Set u k = ( k +1) u 0 for k = 1 , 2 , . . . and for each k , define two measures on Y × U by

̸

It is clear that

Moreover, as k →∞ we have W µ p ( µ k , ν k ) →∞ but W p p ( ν k , η k ) remains bounded. See Figure 4.

Part (d). First, the unconditional distance W p ( η, ν ) may be obtained via an unrestricted coupling in Π( η, ν ) , i.e. the set of all joint measures on Y × U having marginals η, ν . Since Π( η, ν ) ⊇ Π Y ( η, ν ) , by part (a) we see that W p ( η, ν ) ≤ W µ p ( η, ν ) .

Let γ ⋆ ( y 0 , u 0 , y 1 , u 1 ) = γ ⋆,y 1 ( u 0 , u 1 ) δ ( y 1 -y 0 ) µ ( y 1 ) be an optimal γ ⋆ ∈ Π Y ( η, ν ) . We claim that γ ( u 0 , u 1 ) := ∫ Y γ ⋆,y ( u 0 , u 1 ) d µ ( y ) couples π U # η and π U # ν . Let π 0 : ( u 0 , u 1 ) ↦→ u 0 be the projection onto the first coordinate of U × U . Observe that for µ -almost every y , we have that γ ⋆,y ∈ Π( η y , ν y ) is optimal, and, in particular, π 0 # γ ⋆,y = η y . Fix an arbitrary φ ∈ C b ( U ) . We then have

Y

y

0

y

1

u

0

u

k

U

Thus π U # γ = π U # η . A similar argument shows that for the map π 1 : ( u 0 , u 1 ) ↦→ u 1 we have π 1 # γ = π U # ν , so that γ ∈ Π( π U # η, π U # ν ) .

Now, as γ ⋆,y 1 ( u 0 , u 1 ) ∈ Π( η y 1 , ν y 1 ) is µ -almost surely optimal in the usual Wasserstein sense,

since γ ∈ Π( π U # η, π U # ν ) is a coupling but potentially sub-optimal.

D.2 Geodesics

We now study the geodesics in the space P µ p ( Y × U ) .

For any η, ν ∈ P µ p ( Y × U ) , there exists a constant speed geodesic between η and ν .

Proof. Write λ t : ( Y × U ) 2 → Y × U for the linear interpolant

Let γ ⋆ ∈ Π Y ( η, ν ) be an optimal restricted coupling, and consider the path of measures in P p ( Y × U ) given by

Step one: We check that for each 0 ≤ t ≤ 1 , we have γ t ∈ P µ p ( Y × U ) . That is, we need to check that for all Borel A ⊆ Y , we have γ t ( A × U ) = µ ( A ) . Indeed, recall that restricted measures are concentrated on the set C (see Equation (19)). Thus,

i.e. γ t ( A × Y ) = µ ( A ) as claimed.

Step two: We show that W µ p ( γ t , γ s ) = | t -s | W µ p ( η, ν ) . Set γ s t := ( λ t , λ s ) # γ ⋆ for 0 ≤ s < t ≤ 1 . We claim γ s t ∈ Π Y ( γ t , γ s ) . Indeed, we have π 1 , 2 # γ s t = γ t because for all Borel A ⊆ Y × U ,

Indeed, if c = ( y, u 0 , y, u 1 ) ∈ C , then

Thus γ s t ∈ Π Y ( η, ν ) as claimed. Now, we have

Conversely, an application of the previous inequality and the triangle inequality show that for 0 ≤ s ≤ t ≤ 1 ,

Rearranging the previous inequality implies | t -s | W µ p ( η, ν ) ≤ W µ p ( γ s , γ t ) for all s, t ∈ [0 , 1] , and hence W µ p ( γ t , γ s ) = | t -s | p W µ p ( η, ν ) .

Theorem 2 (Conditional McCann Interpolants)

Fix η, ν ∈ P µ p ( Y × U ) . Suppose T ⋆ ( y, u ) = ( y, T ⋆ U ( y, u )) is an injective triangular map solving the conditional Monge problem (1) . Define the maps T t : Y × U → Y × U for 0 ≤ t ≤ 1 via T t = (1 -t ) I + tT ⋆ , and define the curve of measures γ t = [ T t ] # η ∈ P γ p ( Y × U ) . Then,

(a) ( γ t ) is absolutely continuous and a constant speed geodesic between η, ν

(b) The vector field v t ( T ⋆ t ( y, u )) = (0 , T ⋆ U ( y, u ) -u ) generates the path γ t , in the sense that ( γ t , v t ) solve the continuity equation (9) .

Proof. Consider the function w t : Y × U → U given by

and note this is precisely w t ( y, u ) = ∂ t T ⋆ t ( y, u ) . Define the vector field

For any φ ∈ Cyl ( Y × U ) , we have

̸

An analogous calculation shows that π 3 , 4 # γ s t = γ s , so that γ s t ∈ Π( γ t , γ s ) . We now check that γ s t ∈ R Y ( Y × U ) . Indeed, suppose E ⊆ Y × U is a Borel set such that E ∩ C = ∅ . In other words, for every ( y 0 , u 0 , y 1 , u 1 ) ∈ E we have y 0 = y 1 . Set D := ( λ t , λ s ) -1 ( E ) . We claim D ∩ C = ∅ , so that

which shows that ( γ t , v t ) solve the continuity equation.

Now, note that for 0 ≤ a ≤ b ≤ 1 , we have

In particular, ∫ 1 0 ∥ v t ∥ L p ( γ t ,Y × U ) d t < ∞ and so by Theorem 4 ( γ t ) is absolutely continuous. A similar calculation shows that ( b -a ) W µ p ( η, ν ) = W µ p ( γ b , γ a ) = ∫ b a | γ ′ ( t ) | , where the last line follows from the absolute continuity of γ t . Thus, ∥ v t ∥ L p ( γ t ,Y × U ) = | γ ′ | ( t ) for almost every t ∈ [0 , 1] by Lebesgue differentiation.

E Proofs: Section 5

In this section, we provide proofs of all claims made in Section 5.

E.1 Continuity Equation

We begin with a lemma that is used in the proof of Theorem 4. Informally, a solution to the continuity equation with a triangular vector field will result in the conditional measures almost surely satisfying the continuity equation as well.

Lemma 1 (Triangular Vector Fields Preserve Conditionals) Suppose v t ( y, u ) = (0 , v U t ( y, u )) is triangular and that ( γ t ) ⊂ P µ p ( Y × U ) is a path of measures such that ( v t , γ t ) satisfy the continuity equation in the sense of distributions. Then, it follows that for µ -almost every y ∈ Y , we have ∂ t γ y t + ∇· ( v U t ( y, -) γ y t ) = 0 .

Proof. Fix any φ ∈ Cyl ( U × I ) . Suppose ψ ∈ Cyl ( Y ) is given, and note that ψ ( y ) φ ( u, t ) ∈ Cyl ( Y × U × I ) . As ( v t , γ t ) solve the continuity equation, it follows from the triangular structure of v t that upon testing against ψφ we have

Because ψ ( y ) ∈ Cyl ( Y ) , it is of the form ρ ( π ( y )) where π : Y → R k for some k ≥ 1 and ρ ∈ C ∞ c ( R k ) . Taking ρ to be a sequence of smooth approximations to the indicator function of an arbitrary rectangle E = E 1 × E 2 ×··· × E k ⊆ R k , we see

As Y is separable, the Borel σ -algebra on Y is generated by the cylinder sets, i.e. those which are precisely of the form π -1 ( E ) for some finite-dimensional rectangle E . We have thus shown that for an arbitrary Borel measurable set E ⊆ Y ,

From this, it follows that

E.2 Absolutely Continuous Curves

We now proceed to prove the main results of this section. First, we introduce some preliminary notions. We define the map j q : L q ( γ, Y × U ) → L p ( γ, Y × U ) for 1 /p +1 /q = 1 via

̸

which is the Fréchet differential of the convex functional 1 q ∥ w ∥ q L q ( γ,Y × U ) . A straightforward calculation shows that this map satisfies

Theorem 3 (Absolutely Continuous Curves in P µ p ( Y × U ) ) Let I ⊂ R be an open interval, and suppose γ t : I → P µ p ( Y × U ) is an absolutely continuous in the W µ p metric with | γ ′ | ( t ) ∈ L 1 ( I ) . Then, there exists a Borel vector field v t ( y, u ) such that

(a) v t is triangular

(b) v t ∈ L p ( γ t , Y × U ) and ∥ v t ∥ L p ( γ t ,Y × U ) ≤ | γ ′ | ( t ) for a.e. t

(c) ( v t , γ t ) solve the continuity equation in the sense of distributions.

Proof. Assume without loss of generality that | γ ′ | ( t ) ∈ L ∞ ( I ) and that I = (0 , 1) [Ambrosio et al., 2005, Lemma 1.1.4, Lemma 8.1.3]. Fix any φ ∈ Cyl ( Y × U ) . For s, t ∈ I there exists an optimal triangular coupling γ st ∈ Π Y ( γ s , γ t ) . By Hölder's inequality,

It follows that t ↦→ γ t ( φ ) is absolutely continuous. We can introduce the upper semicontinuous and bounded map

̸

For | h | sufficiently small, choose any optimal coupling γ ( s + h ) h ∈ Π Y ( γ s + h , γ s ) and note that

If t is a point of metric differentiability for t ↦→ γ t , note that γ ( t + h ) t → ( I, I ) # γ t narrowly, where I is the identity map on Y × U . Moreover, since γ t ∈ P µ p ( Y × U ) , it follows that on the diagonal we have that almost surely H ( y 0 , u 0 , y 0 , u 1 ) = ι ( |∇ u φ ( y 0 , u 0 ) | . Thus,

Taking Q = Y × U × I and γ = ∫ γ t d t , fix any φ ∈ Cyl ( Q ) . We have that

An application of Fatou's Lemma, Equation (78), and Hölder's inequality gives us

for any interval J ⊂ I with supp φ ⊂ J × Y × U .

Fix the subspace

and denote by V its L q ( γ, Y × U × I ) closure. Define the linear functional L : V → R via

and note that Equation (81) implies that L is a bounded linear functional on V . Thus (by HahnBanach and the fact that V ⊆ V is dense) we may uniquely extend L to V . We thus have a convex minimization problem

which admits the unique solution w such that j q ( w ) -L = 0 . In particular, the estimate (81) shows that the above functional is coercive and hence admits a minimizer which we may obtain via its differential as a consequence of convexity. Thus, we obtain a triangular vector field v = j q ( w ) such that for all φ ∈ Cyl ( Q ) ,

This precisely shows that ( v t , γ t ) is a triangular distributional solution to the continuity equation.

Now, choose any interval J ⊂ I and choose a sequence η k ∈ C ∞ c ( J ) , with 0 ≤ η k ≤ 1 and η k → ✶ J as k →∞ . Moreover choose a sequence ( ∇ u φ n ) ⊂ V converging to w = j p ( v ) in L q ( γ, Q ) . Our previous calculations give

Taking k →∞ we see that

and since J ⊂ I was arbitrary, we conclude

We now prove, in some sense, a converse of the previous theorem.

Theorem 4 (Continuous Curves Generated by Triangular Vector Fields)

Suppose that γ t : I → P µ p ( Y × U ) is narrowly continuous and ( v t ) is a triangular vector field such that ( γ t , v t ) solve the continuity equation with ∥ v t ∥ L p ( γ t ,Y × U ) ∈ L 1 ( I ) . Then, γ t : I → P µ p ( Y × U ) is absolutely continuous in the W µ p metric and | γ ′ | ( t ) ≤ ∥ v t ∥ L p ( µ , Y × U ) for almost every t .

Proof. We first assume that U is finite dimensional. Our strategy is to check the hypotheses necessary for Ambrosio et al. [2005, Theorem 8.3.1] to hold for µ -almost every y , followed by an application of this theorem. By Lemma 1, for µ -almost every y we have that ( γ y t , v U t ( y, -)) solve the continuity equation distributionally on I × U .

By Jensen's inequality (and the assumption p ≥ 1 ) we see

Since the first term is finite, it follows that

Now Ambrosio et al. [2005, Lemma 8.1.2] shows that for µ -almost every y we have that ( γ y t ) admits a narrowly continuous representative (˜ γ y t ) with ˜ γ y t = γ y t for almost every t . It follows from Ambrosio et al. [2005, Theorem 8.3.1] that for any t 1 ≤ t 2 in I , we have

where the second line follows as ˜ γ y t = γ y t for almost every t .

Let ˜ γ t = ∫ Y ˜ γ y t d µ ( y ) be the measure obtained via marginalizing over the Y -variables. Taking an expectation over y ∼ µ , the previous inequality shows us that

Now, note that t 1 is almost surely a Lebesgue point of the right-hand side and ˜ γ t 1 = γ t 1 . Taking t 2 → t 1 along a sequence where ˜ γ t 2 = γ t 2 shows us that

for almost every t ∈ I .

In the case that U is infinite dimensional, fix any y ∈ Y such that Lemma 1 holds (which is of full measure) and fix a countable orthonormal basis ( e k ) for U . Set π d : U → R d to be the projection operator for this basis, i.e. u ↦→ ( ⟨ u, e 1 ⟩ , . . . , ⟨ u, e d ⟩ ) . We consider the collection of finite dimensional conditional measures γ d,y t = π d # γ y t . By the same argument in Ambrosio et al. [2005,

Theorem 8.3.1], there exists a vector field v d,y t on R d such that ( γ d,y t , v d,y t ) solve the continuity equation and

It follows from the finite-dimensional case above that for almost every t 1 ≤ t 2 , we have

Let ˆ γ y,d t = ( π d ) ⋆ # γ y,d t where ( π d ) ⋆ : R d → U maps z ↦→ ∑ d k =1 z k e k . As d → ∞ we have ˆ γ d,y t → γ y t narrowly for all t ∈ I . Since ( π d ) ⋆ is an isometry, Ambrosio et al. [2005, Lemma 7.1.4] shows that

Now, integration with respect to d µ ( y ) yields

Taking t 2 → t 1 shows that for almost every t we have

Together, Theorem 3 and Theorem 4 give us a dynamical interpretation of the conditional Wasserstein distance. The following result is a conditional analogue of the well-known Benamou-Brenier Theorem [Benamou and Brenier, 2000]. Here, we note that the following proof follows the standard proof closely - the main legwork in obtaining this conditional generalization is through the previous two theorems.

Theorem 5 (Conditional Benamou-Brenier)

Let 1 < p < ∞ . For any η, ν ∈ P µ p ( Y × U ) , we have

Proof. Write M for the infimum on the right-hand side.

First, suppose that ( v t , µ t ) are admissible and ∫ 1 0 ∥ v t ∥ L p ( µ t ) < ∞ . It follows from Theorem 4 that ( γ t ) is an absolutely continuous curve in P µ p ( Y × U ) and ∥ v t ∥ L p ( µ t ,Y × U ) ≥ | γ ′ | ( t ) . Thus,

Conversely, by Theorem 1 there exists a constant speed geodesic ( γ t ) ⊂ P µ p ( Y × U ) connecting η and ν . Recall that constant speed geodesics are absolutely continuous. By Theorem 3, there exists a Borel triangular vector field v t such that ( v t , γ t ) solve the continuity equation, and moreover ∥ v t ∥ L p ( µ t ,Y × U ) ≤ | γ ′ | ( t ) . In fact, because ( v t , γ t ) solve the continuity equation, Theorem 4 yields that ∥ v t ∥ L p ( µ t ,Y × U ) = | γ ′ | ( t ) .

Since γ t is a constant speed geodesic in P µ p ( Y × U ) , it follows that | µ ′ | ( t ) = W µ p ( η, ν ) for almost every t ∈ (0 , 1) . Hence,

Thus, W µ,p p ( η, ν ) = M as desired.

F Experiment Details

In this section, we provide additional details regarding all of our experiments, as well as additional results not contained within the main paper. All models can be trained on a single GPU with less than 24 GB of memory, and our experiments were parallelized over 8 such GPUs on a local server. We first describe our setting for the 2D and Lotka-Volterra experiments, as these share a similar setup. Details for the Darcy flow inverse problem are described in the corresponding section.

Models. For FM and COT-FM, our model architecture is an MLP with SeLU activations [Klambauer et al., 2017]. Time conditioning is achieved by concatenating the time variable as an input to the network. The covariance operator C chosen in the path of measures in Equation (10) is taken to be C = σ 2 I where σ is a hyperparameter.

Our implementation of FM is adapted from the torchcfm package Tong et al. [2023], available under the MIT License. For PCP-Map and COT-Flow, we adapt the open-source implementations from Wang et al. [2023], available under the MIT License.

Training and Model Selection. Hyperparameter tuning of the PCP-Map and COT-Flow models was performed directly using the code of Wang et al. [2023], essentially implementing grid-search with an early stopping procedure. We refer to the paper and codebase of Wang et al. [2023] for further details. For COT-FM and FM, we perform a random grid search over 100 hyperparameter settings using the grid described in Table 4. For all model types, we select the best model used to generate the results in the paper as the training checkpoint that resulted in the lowest W 2 error to the joint target distribution on a held-out validation set. For training, we use the Adam optimizer where we only tune the learning rate, leaving all other settings as their defaults in pytorch .

F.1 2D Synthetic Data

Data Generation. This experiment consists of four 2D synthetic datasets, where Y = U = R . The datasets moons, circles, swissroll are available through scikit-learn [Pedregosa et al., 2011]. The moons dataset is generated with noise= 0 . 05 followed by standard scaling with a mean of m = (0 . 5 , 0 . 25) and standard deviation of σ = (0 . 75 , 0 . 25) . The circles dataset is generated with factor=0.5 and noise=0.05 . The swissroll dataset is generated with noise=0.75 , followed by projection to the first two coordinates and re-scaling by a factor of 12 . All other unstated parameters are left as their default values. We use the code available from Hosseini et al. [2023] to generate the checkerboard dataset. For all datasets, we generate a training set (i.e., samples from the target distribution) of 20 , 000 samples and 1 , 000 held-out validation samples for model selection. Means and standard deviations in Table 1 are reported across five independent testing sets of 5,000 samples for the best representative of each model type.

In COT-FM, to generate samples from the source distribution, we sample an additional 20 , 000 points from the target distribution and keep only the Y coordinates. This ensures that the source and target have equal Y marginals. During training, standard Gaussian noise N (0 , 1) is sampled for the U coordinate of these source points at each minibatch.

We use minibatch COT couplings [Tong et al., 2023] in this experiment as computing the full COT plan was prohibitively expensive in terms of memory usage. However, we note that we use large batch sizes, meaning that the COT plan we find in this way should not be too far from optimal. All couplings are computed using the POT Python package [Flamary et al., 2021].

F.2 Lotka-Volterra Dynamical System

Data Generation. We adopt the settings of Alfonso et al. [2023] for this experiment. As described in the main paper, we assume p (0) = (30 , 1) and that log( u ) ∼ N ( m, 0 . 5 I ) with m = ( -0 . 125 , -3 , -0 . 125 , -3) . Given parameters u ∈ R 4 ≥ 0 , we simulate Equation (13) for t ∈ { 0 , 2 , . . . , 20 } to obtain a solution z ( u ) ∈ R 22 ≥ 0 . An observation y ∈ R 22 ≥ 0 is obtained by the addition of log-normal noise, i.e. log( y ) ∼ N (log( z ( u ) , 0 . 1 I ) . We thus may simulate many ( y, u ) pairs from the target measure for training.

We generate a training set of 10 , 000 ( y, u ) pairs using the procedure described above and a held-out validation set of 10 , 000 ( y, u ) pairs for model selection. Means and standard deviations in Table 2 are reported across five independent testing sets of 5,000 samples for the best representative of each model type. Figure 2 and Figures 9, 8, 7, 10 show 10 , 000 samples from each model, as well as 10 , 000 samples from the differential evolution Metropolis MCMC sampler [Braak, 2006] after a

Y

Checkerboard

True

COT-FM (Ours)

PCP-Map

COT-Flow

FM

Y

Moons

Y

Circles

U

Y

Swissroll

U

U

U

U

00

burn-in of 50 , 000 samples. This is implemented through the PyMC Python package [Abril-Pla et al., 2023].

For COT-FM we use the full COT couplings, i.e. without minibatches. This is available to use due to the smaller size of the training set used in this experiment. The COT couplings are computed in the same way as the previous section, and as described in Section 7.

Checkerboard

COT-FM (Ours)

PCP-Map

COT-Flow

FM

Moons

Circles

Swissroll

0

.

50

0

.

75

1

.

00

1

.

25

α

α

0

.

0

0

.

1

0

.

50

1

.

25

β

0

.

00

0

.

05

0

.

10

β

0

.

7

1

.

3

0

.

50

1

.

25

γ

0

.

7

1

.

3

0

.

0

0

.

1

0

.

75

1

.

00

1

.

25

γ

0

.

02

0

.

06

0

.

50

1

.

25

δ

0

.

02

0

.

06

0

.

0

0

.

1

0

.

02

0

.

06

0

.

7

1

.

3

0

.

02

0

.

04

0

.

06

δ

Lotka-Volterra Samples: COT-FM (Ours)

0

.

50

0

.

75

1

.

00

1

.

25

α

α

0

.

0

0

.

1

0

.

50

1

.

25

β

0

.

00

0

.

05

0

.

10

β

0

.

7

1

.

3

0

.

50

1

.

25

γ

0

.

7

1

.

3

0

.

0

0

.

1

0

.

75

1

.

00

1

.

25

γ

0

.

02

0

.

06

0

.

50

1

.

25

δ

0

.

02

0

.

06

0

.

0

0

.

1

0

.

02

0

.

06

0

.

7

1

.

3

0

.

02

0

.

04

0

.

06

δ

Lotka-Volterra Samples: PCP-Map

Lotka-Volterra Samples: COT-Flow

0

.

50

0

.

75

1

.

00

1

.

25

α

α

0

.

0

0

.

1

0

.

50

1

.

25

β

0

.

00

0

.

05

0

.

10

β

0

.

7

1

.

3

0

.

50

1

.

25

γ

0

.

7

1

.

3

0

.

0

0

.

1

0

.

75

1

.

00

1

.

25

γ

0

.

02

0

.

06

0

.

50

1

.

25

δ

0

.

02

0

.

06

0

.

0

0

.

1

0

.

02

0

.

06

0

.

7

1

.

3

0

.

02

0

.

04

0

.

06

δ

0

.

50

0

.

75

1

.

00

1

.

25

α

α

0

.

0

0

.

1

0

.

50

1

.

25

β

0

.

00

0

.

05

0

.

10

β

0

.

7

1

.

3

0

.

50

1

.

25

γ

0

.

7

1

.

3

0

.

0

0

.

1

0

.

75

1

.

00

1

.

25

γ

0

.

02

0

.

06

0

.

50

1

.

25

δ

0

.

02

0

.

06

0

.

0

0

.

1

0

.

02

0

.

06

0

.

7

1

.

3

0

.

02

0

.

04

0

.

06

δ

Lotka-Volterra Samples: FM

F.3 Inverse Darcy Flow

Dataset. The training and test datasets are generated following the same procedure as Hosseini et al. [2023]: pressure fields u are sampled from a Gaussian process with Matérn kernel having ν = 3 / 2 and lengthscale ℓ = 1 / 2 , on a regular 40 × 40 grid. The parameters are then exponentiated and used to simulate the permeability fields p from the forward model F solving the Darcy flow PDE, using FEniCS [Alnæs et al., 2015]. Stochasticity arises from adding Gaussian noise to the permeability fields, obtaining y = F ( u ) + ϵ, ϵ ∼ N (0 , σ 2 I ) . For our experiments we observe y on a 100 × 100 grid, and we use σ = 2 . 5 × 10 -2 . We note that this level of noise is quite considerable, as it accounts for roughly 60% of the variability in the y . Figure 11 showcases a data point for reference. Our source and target training sets contain 1 × 10 4 samples each, and our test set comprises 5 × 10 3 samples. We remark that although y and u are observed on a grid their resolution does not need to be fixed, allowing for training at different resolutions.

Permeability field u

Pressure field p

Observedy

Models. In order to make learning feasible in infinite-dimensional Hilbert spaces, we adapt the architecture of a Fourier Neural Operator (FNO) [Li et al., 2020] from the neuraloperator package [Kovachki et al., 2021] to accommodate for conditioning information observed at an arbitrary resolution. We do so by introducing a projection layer mapping the conditioning information to match the hidden channels of the input lifting block, and a pooling operation to project to the input dimensions. The two are then concatenated and passed through an FNOBlock mapping from (2 × hidden_channels ) × input_dim to hidden_channels × input_dim , before following the original architecture. For all of the models in consideration, we fix the architecture to be have hidden_channels = 64 , projection_channels = 256 , and 32 Fourier modes. We train each model for 1500 epochs, and hyperparameters for each architecture are selected as follows:

WaMGAN[Hosseini et al., 2023]: using an adaptation to the FNO architecture of the original code 3 , we perform a grid search as detailed in Table 5. We found the training procedure to be rather unstable, and for this reason we checkpoint the model every 100 epochs and report the results for the best performing model at its best checkpoint. We found this to be a model with learning rate 1 × 10 -4 , 2 full critic iterations, and monotone penalty of 1 × 10 -3 . The gradient penalty parameter did not seem to significantly affect performance on the test set, and was set to 5.

FFM [Kerrigan et al., 2024]: the learning rate is fixed to 5 × 10 -4 , and the covariance operator C is set to match that of the prior, but rescaled by a factor of σ = 1 × 10 -3 . We use the code from the original repository 4 .

COT-FFM: we set ϵ = 1 × 10 -5 in the cost function used to build the COT plan. The learning rate and C are chosen to be the same as FFM. In order to build COT couplings, we take the source measure to be the product measure π Y # η ×N (0 , C ) . Approximate couplings are obtained on minibatches of size 256 .

It should be noted that in any scenario where the source and the target U -marginals are identical, using the OT coupling would yield the identity mapping as the optimal vector field minimizing (12). Hence, the OT-CFM model [Tong et al., 2023] is inapplicable here.

Sampling. The resulting amortized sampler, denoted for simplification by the mapping ( y, u 0 ) ↦→ u 1 = ˜ T U ( y, u 0 ) , will parameterize an approximate posterior measure. Notice that, in contrast to classical variational inference techniques, no distributional assumptions are made about the approximate posterior. In turn, integrals are obtained numerically by Monte Carlo sampling K samples from the prior, resulting in the approximation

G Minibatch COT

In this section, we perform an additional experiment demonstrating that our COT-FM is able to obtain good approximations of the true COT map even when trained on relatively small minibatches.

We work on Y = U = R and use a standard Gaussian source and a Gaussian target distribution with covariance ρ = 0 . 75 . These are chosen so that we may evaluate, in closed-form, the true conditional 2 -Wasserstein distance, as detailed in Section C. We then train our COT-FM method using various batch sizes, and measure the resulting model's conditional 2 -Wasserstein distance by sampling 10 , 000 points from the source distribution, flowing each source point along the model's learned vector field, and computing the resulting squared distance to the corresponding terminal point.

In Figure 12, we plot the resulting deviation from the true value of W µ, 2 2 as a function of batch size. We see that even at a relatively small batch size of 16 , the resulting error in the (squared) distances is less than 10% of the true value ( ≈ 0 . 678 ). While this experiment is only feasible on synthetic data (as we must compute the true distance), it nonetheless demonstrates that even with small batches we may recover the true distance.

1024

512

256

128

64

32

16

8

4

0.05

0.00

0.05

0.10

0.15

0.20

True CWD 2  - Model CWD 2

Transport Distance Error

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims and contributions of our work are clearly stated in the introduction and abstract. In addition, the introduction clearly points to the sections in the paper where the claims are discussed in more detail.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations of the work are discussed in Section 8.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All formal claims in the paper have their assumptions clearly stated. Proofs of all theoretical claims are available in the Appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Our experimental settings are clearly described in Section 7 and further details, including hyperparameter choices, are provided in Appendix F. We will release our code upon publication of the paper.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We will release our code as open-source upon acceptance of the paper for publication.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: These details are provided in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Error bars are provided for all quantitative results. Explicit descriptions of how these are calculated are included as well.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: These details are provided in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the ethics guidelines and conform to all requirements.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The potential societal impacts are discussed in Section 8.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We believe that our paper does not pose any such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Any assets used are properly cited throughout the paper. We provide addiitonal details in Appendix F.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: Our paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces COT-FM, a generalization of the Flow Matching model for conditional generation. Specifically, this paper investigates the Conditional Wasserstein Space, a space of joint probability measures on $Y \\times U$ with fixed $Y$-mariginals $\\mu$. This paper proves that an absolutely continuous path in the Conditional Wasserstein Space can be generated by a triangular vector field. Based on this characterization, the COT-FM is proposed as a Flow Matching model that employs a triangular vector field on the $Y \\times U$.\n\nSTRENGTHS:\n- This paper presents theoretical analysis of the Conditional Wasserstein Space, such as the characterization of the absolutely continuous path and the conditional generalization of the Benamou-Brenier Theorem.\n- This paper proposes a Flow Matching model for conditional generation.\n- This paper is easy to follow.\n\nWEAKNESSES:\n- Whether the COT-FM can recover the dynamic optimal transport requires further clarification.\n- Please see the Questions Section below.\n\nQUESTIONS:\n- I would like to clarify the connection between Section (4, 5) and Section 6. It appears that Section (4,5) establish the existence of an absolutely continuous path between two arbitrary measures in the Conditional Wasserstein Space, which can be generated by a triangular vector field. In this context, Section 6 introduces a triangular vector field parametrization to the Flow Matching model. Hence, Section (4,5) justify the triangular parametrization of the standard Flow Matching model within the Conditional Wasserstein Space. Is this correct?\n- I am curious whether COT-FM can recover the dynamic optimal transport within the Conditional Wasserstein Space, as mentioned in Line 256. Proposition 3.4 in [Tong et al., 2023] addresses the standard Wasserstein Space case. Could you provide clarification on how this applies to the Conditional Wasserstein Space?\n- Table 1 presents the W2 and MMD distances between the joint distributions. Could you also provide the W2 and MMD results between the conditional distributions? \n- I am curious about the significance of minibatch optimal coupling for COT-FM, in Lines 257-265. Could you provide the COT-FM results using independent coupling?\n\nLIMITATIONS:\n-\tThe authors addressed the potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a theory for conditional optimal transport (as defined by the authors), followed by numerical simulations. Among their contributions, the authors put forth theory for the geometry of the conditional Wasserstein space (where analogous quantities of e.g., the McCann interpolation, hold). This is due to the geometry given by triangular vector fields that are studied. Their proposed algorithm is based off Flow Matching [Lipman et al. 2023], and achieves strong numerical performance against the other baseline algorithms.\n\nSTRENGTHS:\nThis paper has many strengths! It is well-written, fits nicely in the conference format without omitting many details, and has appropriate experiments. The proposed methodology is also quite elegant, and circumvents many issues other methods face.\n\nWEAKNESSES:\nN/A :)\n\nQUESTIONS:\nI am using this space for comments and suggestions as well as questions.\n\n- Is there a clear way to choose the $\\epsilon$ parameter for the COT Flow Matching? Any heuristics whatsoever? This appears to be a bottleneck to making this methodology fully practical\n- Convergence of the $\\epsilon\\to 0$ limit of the proposed OT map for the twisted cost is originally due to Carlier et al (2010) --- I would argue that the recent results by Hosseini et al. (2023) are extensions of this older result.\n- When citing flow matching throughout the draft, it would be equitable to also cite Liu et al. (2023) alongside Lipman et al. 2023 and Albergo et al. (2023). Same goes for Pooladian et al. (2023) --- should be cited alongside Tong et al. (2023) (in e.g., Section 6)\n- Is equation (9) not due to the original flow matching papers?\n- For equation (6): I have never heard anyone say the equation should be \"understood distributionally\". Maybe consider \"in the sense of distributions\" \n- Stylistic comment: Maybe omit \"unconditional\" from the title of Appendix A? This is not really used  \n\n@article{carlier2010knothe,\n  title={From Knothe's transport to Brenier's map and a continuation method for optimal transport},\n  author={Carlier, Guillaume and Galichon, Alfred and Santambrogio, Filippo},\n  journal={SIAM Journal on Mathematical Analysis},\n  volume={41},\n  number={6},\n  pages={2554--2576},\n  year={2010},\n  publisher={SIAM}\n}\n\n@article{liu2022flow,\n  title={Flow straight and fast: Learning to generate and transfer data with rectified flow},\n  author={Liu, Xingchao and Gong, Chengyue and Liu, Qiang},\n  journal={arXiv preprint arXiv:2209.03003},\n  year={2022}\n}\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper characterizes dynamical conditional optimal transport (COT). It generalizes the Benamou-Brenier theorem to dynamical COT. The authors then propose conditional flow matching and apply it to synthetic data.\n\nSTRENGTHS:\n- The paper successfully extends the Benamou-Brenier theorem to the context of dynamical conditional optimal transport. \n- The paper is easy to follow.\n\nWEAKNESSES:\n- The paper appears to be extremely similar to [1] both in theoretically and empirically. Particularly, Theorem 18, which discusses a Benamou-Brenier-like formula and it applies COT with flow matching as like this work. I would like to ask authors to discuss the difference with result in [1].\n\n- The paper discusses dynamical COT in Sections 4-5. However, in Section 6, the authors propose the Conditional OT Flow Matching (COT-FM) method. This method solves dynamic COT only when the given joint coupling is the solution of COT. In other words, optimal coupling should be given for COT-FM algorithm to solve dynamic COT problem. In OT literature, the most of the application aims to find optimal coupling (rather than given), hence, this algorithm can be applied only in very restricted situation. Thus, the application as an OT method is extremely limited.\n\n- Moreover, in the conducted experiments, the given pairs are not the solutions to COT (only mini-batch sense). Therefore, the experiments seem to address conditional generation rather than conditional optimal transport. It is unclear if the experimental settings are appropriate for the subject of the paper. Moreover, the experiments were conducted on very small datasets.\n\n[1] Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching (arxiv, v1 released in March, 2024)\n\nQUESTIONS:\n- In line 51, it is discussed that \"COT-FM ... interpolates between an arbitrary source and target distribution via a geodesic in the conditional Wasserstein space\". Does it mean that FM model learn geodesics?\n\nLIMITATIONS:\nThe limitation is discussed in Weakness section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This work first extends conditional optimal transport theory to the dynamical setting. Then a flow-matching model is proposed to approximate these flows with a simulation free training objective. This is then applied in several conditional generation tasks including two Bayesian inverse problems. Triangular optimal transport maps are used in combination with the dynamic Brenier-Benamou formulation of standard optimal transport.

STRENGTHS:
- This work combines ideas from conditional optimal transport and optimal transport flow matching to create a new approach for Bayesian inverse problems and likelihood-free inference.
- The results effectively demonstrate how this approach can be used in a variety of settings.
- The work is well written and fairly clear

WEAKNESSES:
- Only applied to relatively low dimensional problems.
- It might be useful to clearly highlight the power and generalization of this work over a “simple” conditional optimal transport formulation which simply conditions on a single-class variable. e.g. https://github.com/atong01/conditional-flow-matching/blob/main/examples/images/conditional_mnist.ipynb
    
    While it is clear from a deep enough reading I suggest the authors might want to highlight these differences for potentially broader appeal. This might be done through an algorithm box or other presentation.

QUESTIONS:
It would be good to know empirically how far from optimal the learned transport maps are, as this is a known limitation of miqibatch-based approaches. 

Comment: I would suggest citing work on rectified flows as a concurrent invention of flow matching. 
https://arxiv.org/abs/2209.03003

LIMITATIONS:
Adequately discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-13