Dynamic Conditional Optimal Transport through Simulation-Free Flows
Authors: Gavin Kerrigan
Giosue Migliorini
Padhraic Smyth
Venue: NeurIPS 2024 poster (2024)
Keywords: flow matching
optimal transport
generative models
conditional generation
Primary area: generative_models
TL;DR: We study the geometry of conditional optimal transport from a dynamical perspective, and use our theory to build conditional generative models.

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
Gavin Kerrigan
Giosue Migliorini
Department of Computer Science University of California, Irvine gavin k@uci. edu
Department of Statistics
University of California, Irvine
gmiglior@uci.er
Padhraic Smyth
Department of Computer Science University of California, Irvine smyth@ics uci.edu
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 gener- ative modeling. Our method couples arbitrary source distribution to specified target distribution through 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 fiintte-dimensional settings, making them well suited for wide class of Bayesian inverse problems. Empirically, we demonstrate that our method is competitive on several challenging conditional generation tasks, including infinite-dimen dimensioanll inverse problem.
Introduction
Many fundamental tasks in machine learning and statistics may posed as modeling a conditional distribution v(u y) where obtaining an analytical representation of v(u) | y) is impractical. While sampling-basee approaches such as Markov Chain Monte Carlo (MCMC) methods are useful, they suffer from several limitations. First, MCMC requires numerous likelihood evaluations, rendering prohibitively expensive in scientific and engineering applications where the likelihood is determined by expensive numerical simulator. Second, MCMC must be 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 likelihood-free [Cranmer et al., 2020) and amortized [Amos et al., 2023] approach. While methods like ABC [Beaumont, 2010] and variational inference [Blei al., 2017] partially address these challenges, they are difficult to scale to high dimensions or have limited flexibility.
al., 2019, 2021], GANs [Ramesh et al., 2022], and diffusion models [Sharrock et al., 2022] have shown promise amortized and Riklihooddfrre inference. These models may be viewed in the framework of measure transport [Baptista et al., 2020], where samples u ~ 17(4) from a tractable source distribution are transformed by mapping Tyy, u) such that the transformed samples are approximately distributed v(u y). One way to achieve this is through triangular mappings [Baptista et al., 2020, Spantini et al., 2022], where joint source distribution 17,," u) is transformed by mapping of the form T (y,u) (Ty (v), To (y, u)). Under suitable assumptions, transforms the source ny u) into the target v(y, u), then (v -) couples the conditionals 77(u y) and v(u y).
Recently, generative models such normalizing flows [Papamakarios
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Typically, such a map T' is not unique [Wang et al., 2023], and natural idea is thus to regularize the transport and search for an admissible mapping that 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 conditional generative model for fieeee inference based on dynamic formulation of conditional optimal transport. Specifically, our contributions are as follows:
1. We develop general theoretical framework for dynamic conditional optimal transport in separable Hilbert spaces. Our framework is applicable in infiniteddimensionall spaces, enabling applications in function space Bayesian inference. In Section 4, we study the conditional Wasserstein space TP 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 PP (Y x U) via the continuity equation and triangular vector fields. As consequence, we obtain conditional generalizations of the McCann interpolants [McCann, 1997] and the Benamou- Sennamou-Brenier Theorem [Benamou and Brenier, 2000].
2. 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 path of measures interpolating between an arbitrary source and target distribution via geodesic in the conditional Wasserstein space.
3. 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-dimensionn 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 under- explored 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 hcuristic 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 novel dynamic formulation of COT, which we model through 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].
Flow matching [Lipman et al., 2022] (and the closely related stochastic interpolants [Albergo et al., 2023a] and rectified flows [Liu et al., 2022]) are class of methods for building continuous- -itimuuuu--immeeee normalizing flows in a simulation-fis 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.
Simulation-Free Continuous Normalizing Flows.
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 inite-dimensional setting, whereas our method adds to the growing literature on function-space 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 Fuclidean setting, whereas our methods are applicable in general separable Hilbert spaces.
2

---

3 Background and Notation
Let X, represent arbitrary separable Hilbert spaces, equipped with the Borel o-algebra. We use P(X) to represent the space of Borel probability measures on X, and (x(X) C P(X) to represent the subspace of measures having finite pth moment. If P(X) is probability measure on X and
T: P + X' is measurable, then the pushforward measure 7 (-) ny ( is probability measure on X'. Maps of the form e.g. x² 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, 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 x U. equipped with the canonical inner product obtained via the sum of the inner products on Y and U, under which the space Y X U is also a separable Hilbert space. Let 7 E P(Y x U) be a joint probability measure. The measures T E P(Y) and o " E PPU) obtained via projection are the marginals of 1). We use 77" E P(U) to represent the measure obtained by conditioning 17 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 Borel set E C Y with m E)= = 0, and the rpM unique for y & E.
3.1 Static Conditional Optimal Transport
In conditional optimal transport, we are given a target measure V E P(Y x U) and some source measure 7) E PUU). We seek transport map T Y x U - U such that, for any given y E Y. the mapping Tyy, -) U - U transforms the source distribution into the conditional distribution vV, i... T(y,- +] Da_ In sense, T can be thought of collection of transport maps indexed by y E Y. such a T were available, by drawing samples to ~ and transforming them, one would obtain samples Tyy, uo) ~ 1², Solving this transport problem for each fixed y is expensive at best, or impossible when only has a single (or no) samples (y, u) ~ V 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 X + Y XxU U [Hosseini et al., 2023, Baptista et al., 2020] of the form T(y, u) (y)) To (Ty (y). u)) for some Ty mappings are of interest as they allow us to obtain conditional couplings from joint couplings.
Y and To Y U + U. Triangular
Proposition 1 (Theorem 2.4 [Baptista et al., 2020], Prop. 2.3 [Hosseini et al., 2023])
Suppose 7, V E P(Y x U) and T : Y U + Y x U is triangular. If Twm = v, then To (v), -) = (1)) for TH q-almsst every
In many scenarios of practical interest, the source measure ii and the target measure v have the same Y-marginals. We will henceforth make this assumption, and use μ = - = T v to represent this marginal. In this case, we may take Ty be the identity mapping, so that the conclusion of Proposition simplifies to T' (y, 7) n = ndd for p-almost every y. We note that in situations where such an assumption does not hold, one may simply preprocess the source measure r) via an invertible mapping Ty satisfying [Ty] | 7| = n v [Hosseini et al., 2023, Prop 3.2].
Given a source and target measures 17, v E PP(Y x U) and a cost function e: (Y x U)2 + R, the conditional Monge problem seeks to find a triangular mapping solving
Yxu inf { cly, u, Tly, u)) dny(v, u) Twn T: (v. u) + (v, wi)
(I)
The conditional Monge problem also admits a relaxation under which one only considers couplings whose Y ccomponentts are almost surely equal. To that end, for 17, v E (Y x U) we define the set of triangular couplings 1y (n) v) to be the couplings of 77 and v that almost surely fix the Y Y-components,
(n. = = (XxU3)) b y y=
(I, (1,
(2)
In other words, a triangular coupling 7 E Ty v) has samples (No, uo, V1, 11) ~ 7 such that yo = Vı almost surely. The conditional Kantorovich problem seeks a triangular coupling solving
inf { clyu, un 81, un dr un yı, 21) Y E Hl (n,
(3)
3

---

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 Kan- torovich problem induce optimal couplings for μ-almmstt every conditional measure. We refer to Appendix B and Hosseini 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 pth moments and having fixed Y-marginals μ.. Interestingly, Gigli [2008, Chapter 4] studies the same space for the purposes constructing geometric tangent spaces in the usual Wasserstein space.
Definition (Conditional Wasserstein Space)
Suppose μ E IP(Y) is given and 1 S < 3c. The conditional p- p-Wasserstein space is P x U) (7 E PP x XU) 7 = (d)
We now equip PY (Y U) with a metric WH", 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 μs.
Definition 2 (Conditional p- Wasserstein Distance)
Suppose T), V E (P x U) and S < oc. The function WE :P XU) X PK(Y x U) + R
p
W (17 v) (Eywu ( w²)) 1/8 = (f 12)
( dp(w))
1/p
(5)
the conditional p- Wasserstein distance. Wp is the usual Wasserstein distance for measures on U.
By Jensen's 's inequality we have WH' v) 2 Eynu (n, v.). In the following, we show that the conditional Wasserstein distance is a well-defined metric as well as few other metric properties. Proposition 2 (Some Properties of W
Let S p < co.
(a) Wht is well-defined, finite, and equals the minimal conditional Kantorovich cost.
p
(b) W is a metric on the space (Y U).
(c) There does not exist C > such that Wh v) S CW,(n, for all n, E PK U).
(d) For all n, E PY x U). W, ( S (n, and W (n, v) S WH,
Proposition 2(c, d) together shows that one should expect the topology generated by W to be stronger than the unconditional distance W,. Here, we note that Gigli [2008] and Chemseddine et al. [2023] previously showed that WH is 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 proof's 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 n, v E (Y U) are Gaussians
NO == 0.. f))
1
P
It follows that = Tل = TX = 1), As 17", Di are Gaussians, their W2 distance admits a closed form and we can directly compute the expectation in Equation (5) to obtain Wh (v. 2(1 - /1 - -22). This is zero if and only if p = 0. i.e. = v. However, = regardless of p. Moreover, the uncondi- (2 VI- p- vT+ from which it is easy to verify that (2(n,)) Wh'(n, v). See Appendix C for similar derivation which applies to arbitrary Gaussians.
U "
= # = (V 1) and W2(
tional distance is Wz(n,

---

Conditional Wasserstein Space Geodesic Space. We now turn our attention to the geodesics in I U.. In particular, we show that there exists constant speed geodesic between any two measures in Py 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 PPC (Y x U) is practical in Section 6, we show how one can model geodesics in ( x U) in order to obtain conditional flow-based model whose paths are easy to integrate.
U) generalizing a similar result in the unconditional setting [Santambrogio,
curve is a continuous function % : - (Y x U) where [a, b) C R is any open interval of finite length. A curve is absolutely continuous if there exists m L ((a, b)) such that
Wh (1s: - [ m(T) dr
va S b.
If (7)) is an absolutely continuous curve, then its metric derivative
h'l(t) lim W7( 76) [s-t]
(8)
exists for almost every t (a, b), and, moreover, for any m satisfying Equation (7) [Ambrosio et al., 2005, Theorem 1.1.2]. A curve (7t) is called a constant speed geodesic if for all < s S t < b. we have (1; = s|W (7er 76). It is straightforward to show that every constant speed geodesic is absolutely continuous.
almost surely have hn((t) S ma(t) pointwise
Theorem 1 (PP) (Y U) is Geodesic Space)
For any n) V E PR (Y U), there exists constant speed geodesic between and v.
When an optimal triangular coupling y E 1ly (17, 1/) is induced by an injective triangular map T we may recover constant speed geodesic in py x U), generalizing the McCann interpolant [McCann, 1997] to the conditional setting. We refer to Proposition for sufficient conditions on 17, v under which such T* exists. Informally, samples from (yu, un) r7 flow in straight path a constant speed to their destination T* (yo, Mn))
Theorem 2 (Conditional McCann Interpolants)
Fix q. V E IP (Y x U). Suppose T (y. u) (y. (y, u)) is an injective triangular map solving the conditional Monge problem (1). Define the maps T. Y x U Y x U for 0 < t $ via Te = =(1- t)I + tT*, and define the curve of measures " = [T E P) x U). Then,
(a) (1x) is absolutely continuous and constant speed geodesic between 1J v
(b) The vector field oe (T (y, u)) = (0, T7 (y, u) - u) generates the path tt in the sense that (7i, vx) solve the continuity equation (9).
5 Conditional Benamou- Theorem
In this section, we prove a characterization of the absolutely continuous curves PK X U). As 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 vector field on Y x U which has zero velocity in the Y component. This is natural, as all measures in Tp Y x U' have fixed -marginnal μ.. Such vector field can be informally seen 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 R, time-dependent Borel vector field 1 x Y U + Y x U is said to be triangular if there exists Borel vector field D x + U such that v, (y. u) = (0, of (y, u)).
Continuity Equation. We introduce some necessary background which allows us to link vector fields curves of measures. The continuity equation 0, + div(uay) = 0 describes the evolution of measure which flows along a given vector field [Ambrosio et al., 2005, Chapter 8). This equation must be understood in the sense of distributions, i.e. for every yp a space of test functions,
J. (Bissy, (, + (ve(y, u), V Venulyy u,t))) dy u) dt 0.
5

---

We consider cylindrical test functions & E XUXYYxxII,ii of the form p u, = ( (v. u). where nº : Y x U -> IR maps (y. u) - (((v, e)).,).... ((v. u), ea)) and fe1, eaf is any orthonormal family in Y U. In the finite dimensional setting, one may take ip E C U3 be smooth and compactly supported [Ambrosio et al., 2005, Remark 8.1.1].
In Appendix E, 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 (Triangular Vector Fields Preserve Conditionals)
Suppose (y, u) = (0, of (y, is triangular and that (1x) C PI x U) is path of measures such that (wi: 7()) satisfy the continuity equation in the sense of distributions. Then, it follows that for ju-almost every y E Y, we have d + +V. (v (y. )),) =0
We note that having be triangular is sufficient, but certainly not necessary. for the conditional continuity equation to almost surely hold. For instance, the vector field in Rd that rotates (0, about the origin is not triangular yet preserves all conditional distributions.
bbsolutely Continuous Curves. In this section, we state our characterization of absolutely contin- uous curves in PY x U). Informally, given such a curve, Theorem 3 provides us with triangular vector field which generates the curve, in the sense that the pair solve the continuity equation.
Theorem 3 (Absolutely Continuous Curves in PK U)) Let I C R be an open interval, and suppose "It - (Y xx U) is an absolutely continuous in the Wh metric with (n((()) E (T). Then, there exists Borel vector field vely. u) such that
(a) de is triangular
(b) vr LP(,,, U) and ||vell S h/1()) a.e. L.
(c) (Ut, Te)) solve the continuity equation in the sense of distributions.
Conversely, we show in Theorem 4 that if the pair (7e, vz) solve the continuity equation and i, is triangular, then the curve (1x) is absolutely continuous and ln'((t) S ||vell 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 diinite-dimensionall result is obtained via fiite-dimensional approximation argument.
Theorem 4 (Continuous Curves Generated by Triangular Vector Fields)
Suppose that "i= I + PY x U) is narrowly continuous and (ve) is a triangular vector field such that (7t: v.) solve the continuity equation with |[vell E L1). Then, ? 1 + P4 x U)
is absolutely continuous in the WHl metric and h'l(t) |[eell LF XUr for almost every t.
As corollary of Theorem 3 and Theorem 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 8J), but we include it for the sake of completeness.
obtain conditional version of the Benamou- Brenier
Theorem (Conditional
Ainamou-Brenier
Let 1 < p < co. For any T), v E PY x U), we have
( [1ollBpeuw) di (we, solve (9), 70 = 7771 = U, and U2 is triangular
VP.((1,))) = min (Tean)
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 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.c. p 2).

---

True
COT-FM (Ours)
PCP-Map
COT-Flow
Conditional
Checkerboar card
True
COT-FM
Circles
PCP-Map COT-Flow
A
Figure 1: Samples from the ground-truth joint target distribution and the various models. Samples from COT-FM more closely match the ground-trune distribution than the baselines. In the final column, we plot conditional KDEs for samples drawn conditioned on the y value indicated by the dashed horizontal line. See Appendix for larger figure and additional results.
Flow Matching. We assume that we have access to samples 20 = (yo, un) 7/0o uo) (Y x U) from source measure, and samples target measure. Let (=o, 21) p(z0, =1) E [[ v) be any coupling of the source and target measure. We specify a collection of measures and vector fields on Y U via
(wn ui) =(y1, un) PY x U) from
Tely a =) - N (y, a 121 + (1 - zo, C)
ve u =) = 21 - 20
(10)
where C is any trace-class covariance operator [Da Prato and Zabezyk, 2014]. As is standard in flow matching [Lipman et al., 2022, Kerrigan et al., 2024], we obtain from Equations (10) a marginal measure (y 2) and vector field (y u) satisfying the continuity equation via
7el., u) Sixxer> rel, u dp(z)
vely.u) = - ve dr u 2) drely, u)
dp(z).
(11)
This marginal path (1 interpolates between the source measure (t = 0) and a smoothed version of the target measure (t 1). To transform source samples from 7) into target samples from v, seek to learn the intractable vector field er (y, u) with a model v (4, y, u) by minimizing the loss'
(()) = (ee. || (1. v. u)
ve ||
(12)
which has the same 0-gradient as the MSE loss to the true vector field ue(y, u) [Tong et al., 2023].
COT Flow Matching. In the preceding section, p(z) may an arbitrary coupling between v and v. 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 triangular map. turn, Theorem gives that this triangular map is generated by triangular vector field of the form (10). Thus, parametrize our model va also be triangular. Moreover, we recover the optimal dynamic transport given in Theorem 5 as Tr(C') + 0 by a pointwise application of [TTng et al., 2023, Proposition 3.4).
Given collection of samples {af, = drawn from 17 and v, we approximate conditional optimal coupling p using standard numerical techniques with the cost function c. 4g, 8/, 21) = [y 20 + e|uu 24 for some e & 1. Intuitively, such cost penalizes mass transfer along the Y dimension, which is precisely the constraint sought in the COT problem (3). As J. O, we recover the true optimal triangular map [Carlier et al., 2010, Hosseini et al., 2023]. The COT coupling
Previous work has referred to this as the conditional flow matching loss [Tong et al., 2023J, which is not to be confused with the notion of conditioning that focus in this work.
7

---

Table 1: Distances between the ground-truth and generated joint distributions for the 2D datasets. Our method (COT-FM) obtains lower distances than the considered baselines. Average results + one standard deviation are reported across five test sets, with the lowest average distance in bold.
Checkerboard
Moons
Circles
Swissroll
Wa (1e-2)
(Ie3)
(le-2)
MMD (1e:3)
Wa (1e-2)
MMD (1e:3)
W2 (1e:2)
MMD (1e-3)
PCP-Map
6.2720.81
0.21
8.44
0.2240.10
6.1940.45
0.20+0.17
5.35+0.35
0.1620.13
COT-Flow
8,20=0.49
0,26
18. 18.49=2.22
1.32+0.79
10.04=1.00
0,24=0.22
6.47=0..6
0.19
FM
8.8
0.24
15.55
1.8540.2
7.03+9.17
0.45.0.11
8.18
0.5820.09
COT-FM (Ours)
4.69=1.00
0.17#0.13
6.50a141
0.13
5.56.0.43
0.20#a.04
4.64+1.6
0.15#0.19
can either be precomputed for small datasets or computed on each minibatch drawn during training. While the use of minibatches 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 learned triangular vector field (1, y, u). Given an arbitrary fixed E Y, we may approximately sample from the target v(u y) by sampling ug ~ niuo y) and numerically solving the corresponding flow equation (y, u4) o (,,y, y, u4) with initial condition (v, uo)-
Source Measure. Our framework is agnostic to the choice of source measure 17, allowing for flcxibility in the modeling process. The main requirement is that the Y Y-marginals of the source 17 and target 17 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 ny un) v(oo @ nu(uu) to be the product of two independent distributions where yo is arbitrary, e.g. Gaussian noise.
7 Experiments
We now illustrate our methodology (COT-FM) on 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 al., 2023]. These baselines are chosen as they reflect current sfatc-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.c. where the coupling between the source and target measures is the independent coupling p(zy.2) = n & v. This baseline serves as an ablation for the COT component of our model.
Table 2: Statistical distances between
MCMC and posterior samples u ~ v(u y) for each method on the LV dataset. Average results + one standard deviation reported across five test sets.
Wa (le-2)
MMD (1c-3)
PCP-Map COT-Flow
5.0410.0
2.67
4.86+1.1
0.83+0.50
FM COT-FM (Ours)
11.41+0.26
2.65+0.14
4.0220.08
0.9520.0
Overall, method (COT-FM) typically outperforms these baselines across the diverse and challeng- ing 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 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. ?
2D Synthetic Data. We first consider synthetic distributions where Y = U = R. Our source measure is taken to be the independent product 17y,, u) π @ NO, 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-trum 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 W2 and MMD distances between the generated and ground-truth joint distributions. This is motivated by Proposition I, as triangular maps which couple the joint distributions necessarily couple the conditional distributions. Our method outperforms the baselines across all metrics.
"ood for all of our experiments is available https: http://ttttuu...oaa.aaaoee hatpb//..ooaaaineeerraaa//eee/e
8

---

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 (a, 5, y, 6) E R2o and pair of coupled nonlinear ODBs of the form
dpi(t)
dp2(t) dt
dt = ap1 - 3p1p2
-2P2 + 8p1P2
(13)
whose solution p(t) = (pi (1) p2 E 1R20 represents the num- ber of prey and predator species at time f E [0. 77. Follow-
ing Alfonso al. [2023], we assume p(0) (30, 1) and that
log(u) ~ NN (m,, 0.51) with m = (-0.125,- -3, -0.125. -3). R
simulate Equation (13) for t 20. 20} obtain a solution z(u) R23 An ob- servation y E 2 is obtained by the addition of log-normal noise, i.c. log(y) ~ (llog(z(u), 0.1Z). We thus may simulate many (y. u) pairs from the target measure for training.
Given parameters
0.83
As benchmark, we follow the settings of Alfonso et al. [2023] and choose parameters a single observation y as described above. In Figure 2, we plot histogram of 10, 000 samples from the posterior v(u y) of COT-FM.
0.041
0.1
(0.83, 0.041, 1.08, 0.04) to generate
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 W2 metric, and competitive in the MMD metric (Table 2).
0.8
1.08
1.2
0.01
MCMC
COT-FM
Darcy Flow Inverse Problem. Here we consider an infinite- dimensional 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 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.
Figure 2: Sample KDEs on the Lotka Volterra inverse problem. The red lines denote the true pa- rameter values.
The Darcy flow PDE an elliptic equation on smooth domain Ω C Rd which relates a permeability field exp(u), pressure field p. and a source term via div exp (u) V an. Our goal is to recover the permeability u from noisy measurements y of the pressure p. Both the unknown and observations y are functions and thus infinite-dimensionaa To define our target measure, we specify a prior v(u) = AO C') with a Matérn kernel C' of lengthscale é = 1/2 and v = 3/2. Given u ~ ry(u,, the Darcy flow PDE is solved numerically [Alas et al., 2015] to obtain solution p(u) observed at some finite but arbitrary number of points (Z1...., 2n.) C IR?2 An observation y(u) is obtained by adding Gaussian noise to each observation, L.c. yuu) ~ (p(()), o² where = 2.5 10-2 We implement all models via Fourier Neural Operator [Li et al., 2020], allowing us to work with arbitrary discretizations, as required by the functional nature of this problem.
Table Predictive performance of the gener- ated samples on the Darcy flow inverse prob- lem. Average result + one standard deviation obtained on 5 test sets of 5,000 samples each.
on Ω subject to p = 0 on
MSE (1e-2)
CRPS (Ie-2)
WaMGAN
6.55=0.07
18. 18.7520.10
FFM
7.30=0.07
15.4720.06
CFM (Ours)
5.40
15.56.0.08
We provide an illustration Figure 3. As the true posterior is intractable, we compare against preconditioned Crank-Nicolson (pCN) [Cotter et al., 2013], 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 (p)
Observed (v)
PCIN
COT--M (Ours)
FFM
WaMGAN
MSE: 21.27
MSE: 22.83
MSE: 25.52
MSE: 2.21
MSE: 3:49
MSE: 3.63
MSE: 4.40
MSE: 4.66
MSE: 6.63
Figure 3: Darcy flow illustration. Several true permeability fields shown, well the pressure field and its observed, noisy version y. We compare an ensemble average of posterior samples from the various methods against MCMC (pCN) [Cotter et al., 2013). COT-FM achieves the lowest MSE to PCN. We note here that WaMGAN has clear visual artifacts despite achieving reasonable MSE and CRPS scores.
between our method and the pCN mean is lower than that of FFM. Table 3 provides quantitative comparison between the methods test set of 5,000 samples, where measure MSE and CRPS (Hersbach, 2000]. We compare the ensemble mean 10 samples against the true u value as running PCN for each observation 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.
Conclusion
We analyze conditional optimal transport from geometric and dynamical point of view. Our analysis culminates in the characterization of absolutely continuous curves of measures in conditional Wasserstein space, resulting conditional analog of the Benamou-Brees -Beenirrrrrr 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 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 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 small additional computational cost compared to standard flow matching. As with all generative models, 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,
10

---

References
Oriol Abril-Pla, Virgile Andreani, Colin Carroll, Larry Dong, Christopher Fonnesheck, Maxim Kochurov, Ravin Kumar, Jumpeng Lao, Christian Luhmann, Osvaldo A Martin, PyMC: modem,, and comprehensive probabilistie programming framework in python. PeerJ Computer Science, 9te1516, 2023.
Jonas Adler and Ozan Öktem. Deep Bayesian inversion. arXiv preprint 181 1,00090 2018.
Michael S Albergo, Nicholas M Boffi, and Eric Vanden- Yanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.088 087977 2023a.
Michael Albergo, Mark Goldstein, Nicholas M Boffi, Rajesh Ranganath, and Eric Vanden- Eijuden.. 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 arXivv2307.041000 2023.
Martin Alnæs, Jan Blechta, Johan Hake, August Johansson, Benjamin Kehlet, Anders Logg, Chris Richardson, Johannes Ring, Marie Rognes, and Garth Wells. The FFniCS project version 1.5. Archive of nunerical software, 3( (100), 2015.
Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: In Metric Spaces and in the Space of Probahility Measures. Springer Science & Business Media, 2005.
Luigi Ambrosio, Alberto Bressan, Dirk Helbing, Axel Klar, Enrique Zuazua, Luigi Ambrosio, and Nicola Gigli. user's guide to optimal transport. Modelling and Optimisation of Flows on Networks: Cetraro, Italy 2009,
Editors: Benedetto Piccoli, Michel Rascle, pages 1 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 empiries. 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-fret inference. arXiv preprint arXiv:2006.0677 067555 2020.
Ricardo Baptista, Youssef Marzouk, and Olivier Zahm. On the representation and learning of monotonc triangular transport maps. Foundations of Computational Mathematics, pages 1-46, 2023.
Raphaël Barboni, Gabriel Peyré, and François-Xaviss Vialard. Understanding the training of infinitely deep and wide ResNets with conditional optimal transport. arXiv preprint arXiv:2403.12882 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 transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
David M Blei, Alp Kucukelbir, and Jon D McAuliffc. Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):859-877 877, 2017.
Vladimir Igorevich Bogachev and Maria Aparecida Soares Ruas. Measure Theory, volume 2. Springer, 2007.
Cajo JF Ter Braak. 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 's transport to Brenier's map and a continuation method for optimal transport. SIAM Journal 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-AOS1400 URL https:
//doi. //doi.org/10 10. 1214/15 15 A0S1401

---

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.18700 2024.
L. Cotter, G. O. Roberts, A. M. Stuart, and D. Whitc. 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 iaaulation-based inference. Proceedings of the National Academy of Sciences, (17744):30555-30000 2020.
Giuseppe Da Prato and Jerzy Zabezyk. 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.6995: 2013.
Aram Daviyan, Sepehr Sameni, and Paolo Favaro. Efficient video prediction via sparsely conditioned flow matching. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 23263- 23274, 2023.
Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar 7. Alaya, Aurélic Boisbunon, Stanislas Chambon, Lactitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, Léo Gautheron, Nathalie T.H. Gayraud, Hicham Janati, Alain Rakotomamoniy. levgen 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 Reseurch, 22(78):1-8, 2021. URL http: //jmrr . org/ pttpps/pppss/// 2/////222222//22 html.
Giulio Franzese, Giulio Corallo, Simone Rossi, Markus Heinonen, Maurizio Filippone, and Pictro 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. 082295, 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):55--570, 2000.
Bamdad Hosseini, Alexander W Hsu, and Amirhossein Taghvaei. Conditional optimal transport function spaces. arXiv preprint ar3: I. /00662 2023.
Noboru Isobe, Masanori Koyama, Kohci 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 Transuc- tiony on Pattern Analysis and Machine Intelligence, 2022.
Young-geum Kim, Kyungbok Lee, Youngwon Choi, Joong-Hu Ho Won, and Myunghee Cho Paik. Wasserstein geodesic generator for conditional distributions. arXiv preprint arXiv:2308. 10744, 2023.
Günter Klambauer. Thomas Unterthiner, Andreas Mayr. and Sepp Hochreiter. Self-normalizing neural networks. Advances in neural information processing systems, 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. 00488,,881111111 2021.
12

---

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.0888) 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 arXi::2302.07400 2023.
Yaron L Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Fleventh International Conference 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.03000. 2022.
Robert McCann. A convexity principle for interacting gases. Advances in Mathematics, 128(1):153-17, 1997.
Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv preprint wrXiv:1411.1700 2014.
George Papamakarios, David Sterratt, and lain 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.
F. Pedregosa, G. Varoquaux, Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825 2830, 2011.
-raa-Alexandre Pooladian, Heli
Carles Domingo- Enrich, Brandon Amos, Yaron Lipman, and Ricky Chen. Multisample flow matching: Straightening flows with minibatch couplings. arXiv preprint arXiv:2304. 14772, 2023.
Poomnim 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.0648) 2022.
Deep Ray, Harisankar Ramaswany, Dhruv V Patel, and Assad A Oberai. The efficacy and generalizability of conditional gans for posterior inference physics-base bresics-based inverse problems. arXiv preprint arXiv:2202.07777, 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 computer vision, pages 145000, 2017.
Filippo Santambrogio. Optimal transport for applicd mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Louis Sharrock, Jack Simons, Song Liu, and Mark Beaumont. Sequential neural score estimation: Likelihood- free 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, 6944666133648 2016.
Cédric Villani al. Optimal Transport: Old and New, volume 338. Springer, 2009.
Zheyu Oliver Wang, Ricardo Baptista, Youssef Marzouk, I Lars Ruthotto, and Deepanshu Verma. Efficient neural network approaches for conditional optimal transport with applications in Bayesian inference. arXiv preprint arXiv:2310.16675. 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 Sysiems. 36, 2024.
13

---

Appendix: Table of Contents
Optimal Transport
14
Conditional Optimal Transport
15
Closed-Form Conditional Wasserstein Distance for Gaussian Measures
16
Proofs: Section 4
17
D.I Metric Properties
17
D.2 Geodesics
19
Proofs: Section 5
E.1 Continuity Equation
21
E.2 Absolutely Continuous Curves
22
F
Experiment Details
25
F.1
2D Synthetic Data
26
F.2
Lotka-Volterne Dynamical System
26
F33 Inverse Darcy Flow
33
Minibatch COT
34
A
Optimal Transport
We provide here brief and informal overview of optimal transport in the standard unconditional setting. For details, refer the standard references of Villani et al. [2009], Santambrogio [2015], Ambrosio et al. [2005]. Let X be separable metric space and fix cost function X X + R [+o0). Suppose have Borel measures 7), v E P(X). The Monge problem seeks to find a measurable transport map I' X + X minimizing the expected cost of transport, i.c. corresponding to the optimization problem
{ c(x, ((), dn(z) Twn v
inf T
(14)
This optimization problem is challenging, though, it involves nonlinear constraint and the set of feasible maps may be empty. In contrast, the Kantorovich problem a relaxation which seeks to find optimal coupling y E I[(n, v), i.e. probability distribution over X X with marginals 17, v, which solves
inf ( clzo, 21) dy(20, m1) 7
I y E M(n
(15)
Under fairly weak conditions (e.g.. the cost is lower semicontinuous and bounded from below [Ambrosio et al., 2013, Theorem 2.51). minimizers to the Kantorovich problem are guaranteed to exist. the cost function c(2n, Ji ) = Jg 1 for some < do, under sufficient regularity conditions on 1) solution 7* to the Monge problem is guaranteed to exist and, moreover, the coupling " = (I, T*) # optimal for the Kantorovich problem. See Ambrosio al. [2013, Chapter and Ambrosio et al. [2005, Theorem 60].
14

---

Wasserstein Space. In the special case that c(zo. 21] 120 z11² P,(X), the Kantorovich problem admits finite-cost solution. The cost of such an optimal coupling is the w--asserstein distance
I for 1 S oo, and e
Wz,n, min U
Xxx 100 xn 2nl d' 01) 1 E M(n,
(16)
which, the name suggests, metric the space (x() [Ambrosio 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- p-Wasserstein distance can be obtained by finding a time-dependent vector field transforming 7y to v across time / E [0, 1] with minimal energy:
U (xxx)) dt to 77, 71 v4 0 div(er))) = =0 =0). (17) (e))
Wp Wy(n min
Here, we constrain our minimization problem over the set of measures and vector fields (7t, v)) interpolating between 1j and v. satisfying continuity equation (see Section 5). In Section 4, we study generalization of the Wasserstein distances for conditional optimal transport problems. In particular, Theorem provides a generalization of the Benamou-Brenier theorem to the conditional setting which recovers conditional Wasserstein distance.
B Conditional Optimal Transport
This section contains additional discussion regarding the static COT problem, supplementing Section 3.1. We refer Hosseini al. [2023], Baptista et al. [2020], and Chemseddine al. [2023] for further results and details.
Given source and target measures 17, E PY U), the conditional Monge problem seeks to find a triangular mapping solving
inf U dy, u, T(,, u)) dn(y,,) TTn u, (v, (y, Tv,, w))
(18)
The conditional Monge problem also admits a relaxation under which one only considers couplings whose Y T-components are almost surely equal. To that end, we consider the subset & C (Y U)² whose Y components are identical, i.e.,
6 == (30, uo UI, u1) E (Y x
= =
(19)
and we define the set of rrsttttti probability measures Ry CP((Y U)) such that every E Ry is concentrated on 66. In other words, if E Ry, then samples (yo, uo 31, u1) have yo y1 almost surely. In addition. for any 17, P(Y x U), we define the triangular couplings
lly (17, v) to be the probability measures in Ry whose marginals are 1) and v. i.e.
my, v) = {o Ry
1.2 #
= 1/, F 3.4
v)
(20)
The conditional Kantorovich problem seeks a triangular coupling >* solving
inf (Jrxey cyyo, uo U1, u1) dy uo, Vı, w1) ) (1,
(21)
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 Kan- torovich problem induce optimal couplings for j-almostt every conditional measure. Assuming sufficient regularity assumptions on the conditional measures, unique solutions the conditional Monge problem exist. We restate these results here for the sake of completeness.
15

---

Proposition 3 (Prop 3.3 [Hosseini et al., 2023])
Fix T,V PV x U.. Suppose the cost function с continuous, inf -Oc, and there exists finite cost coupling 7 E Hl (1), V. Then, the conditional Kantorovich problem admits minimizer Y. Moreover, 7 (y1, Uo; U1) = 30 (uo; un Jo(n) - (o) where for p-almost every y the measure Y is an optimal coupling for r)", D² under the cost ca ("o, "1) c(y. uo U, u1)
Proposition 4 (Prop 3.8 [Hosscini et al., 2023])
Fix 1 < < Do and VV. E (Y X U). Suppose c(Jo: ug y1, w1) Uo -o-- unp. . If T7" assign zero measure to Gaussian null sets for p-almost every y, then there is a unique solution T to the conditional Monge problem, and 7 (I, +1 is the unique solution to the conditional Kantorovich problem. If v" also assign zero measure to Gaussian null seis for p-almost every y, then T is injective n-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 in the main paper.
Suppose Y Rd and U Rd are Euclidean spaces (of possibly different dimensions), and that n. E (Y x U) are Gaussians of the form
N ((
A" Σ
(
m m
A Σ
m
A"
== M
AuD
(22)
where Rd, my my Rd", and the (block) covariance matrices are Σ E Rdxd A" E Raxd' and 2. 2 d'
This form chosen to ensure that 1) and v have equal Y-marginals. follows that μ = 2
(m, E). Let
Q" We have that the conditionals 17", 1r are available in closed-form.
R= RA" A
(23)
n N m + A3 Σ (y Q) D N (m" A - 5- (-y))- m).
(24)
Thus, for any fixed y, we use the known closed- form unconditional Wasserstein distance to obtain
W v²) m m R(y R(y-m)}² Tr Q" Q" -((10)180)00)))))) (25)
We now take expectation over μ N(m, 2) to compute W Observe that R(y m)
No, RER") and that En [[ m) = Tr( )) Thus,
2 (,, v) = Eynn (1"""")))
(26)
Eynw [m mx! 2(m² m R(y m)) RRy m]]]
(27)
+ Tr Q" + Q" 2 (2((00)) (109))))())))))))
2/2
= |m m2/ Tr Q" Q" -2 -(((00)120)000))11 1/2 RER")
(28)
This form, perhaps unsurprisingly, closely resembles the unconditional Wasserstein distance between two Gaussians, except for the presence of an additional Tr (REERT) term. Note that when 77, have uncorrelated Y, components, we precisely recover W3 ( 77 =T as one may expect. As a special case of interest, Y U R and
w -= N (0,
1 1)
= 100
lel <1
(29)
16

---

then we obtain as a special case of Equation (26) that W 2(1 - VI- pI- This is zero if
and only p = 0, i.e. 17 =
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 WH vr) may be viewed as the minimal value of the constrained Kantorovich
problem in Equation (3) when one takes the cost be the metric on the space Y U. Similar results, relating the conditional Wasserstein distance to triangular couplings, have appeared previously, but our proof independent of these prior works [Chemseddine et al., 2023, Gigli, 2008].
Proposition 5 (Equivalent Formulation of the Conditional Wasserstein Distance) Fix q) v E (Y x U) and $ < OC. Then, W (7,, v) is well-defined, finite, and
(n, v) min (oowor d un Un, u1) dr 1 (1.))
(30)
where Wh' (7,, v) represents the p-th power of the conditional p- Wasserstein distance.
Proof. The cost function d' is clearly continuous and non-negative, and hence by Proposition 3 it suffices to exhibit a finite-cost coupling 7 E ITy v) between 1) and v. Indeed, take the conditionally independent coupling
(Bu. up V1, u1) = nuw (1)x((1 11)) (0)
(31)
which is clearly in IIy [1,(n,.), v). We then have that
d' dyo, uo y1, u1) dry un U1, u1) = 1 Srreoy (1 uo) - (u1, un)llxx dr(u0, uo. u2)
2 [ Jaxxy (1130. (II(8, I + (1,,, un xx)) dr(30. uo. v1.
-2 (Fr [1 uo)l dn(ya, uo) frw u1)lix du(en, vi) to0.
Hence, Equation (30) admits minimizer x* E [ly(n, v). By Proposition 3, this minimizer may be taken have the form y = -m (ug, un)ō(y1 - (o) where * (uo, uz) is (1y))-almost surely optimal coupling between 17 wa for the cost lug uul". Thus,
J d dr d
f |un - -ol d u1) dy(y)
(32)
[ v") = (F,(n,
(33)
Here, emphasize that the μ-almost sure uniqueness of the disintegrations of 17, along Y result in a well-defined expression.
Moreover, if 17 E (Y U) it follows that 17" E (U) for p-a.e because
[ Julp (1) du(v) = [. / (v,, u)I dry du(v)
(34)
= fr uP u) < toc.
(35)
Thus all considered Wasserstein distances on U are finite.
17

---

We now proceed to prove several metric properties of our distance. Proposition 2 (Some Properties of w
1 p < oc.
(a) W is well-defined, finite, and equals the minimal conditional Kantorovich cost.
(b) Wh is metric on the spuce PH (Y x U).
(c) There does not exist C > Such that WH, v) & (W for all v PbYY x
(d) For all E PY x U). w (TH (T # S WW((, v) and W W,(,, v) S Wz
Proof. Part (a). This is simply a restatement of Proposition 5.
Part (b). Fix 77, P E PY x U). Since Wp is a metric on (U)), we immediately obtain the symmetry of WP. Moreover, we have that W (1, v) = 0 if and only if 17 = vid for p-aamost every y. Thus, if W (n, v) = and E C Y x U is Borel measurable,
n nEE) 1 (1((E")d(()) J, ("(E" v(E).
(36)
which shows that 17 v. Here, E {u (y, u) E3 is the y-slice of E. Conversely, if 1) V. then ry² = D up to f--ulll set by the essential uniqueness of disintegrations. Thus, (1 v) = 0 if and only if n- = v.
By Minkowski's 's inequality and the triangle inequality for Wp on PU)), we see
Wh (n, v) S (E (W, ,²) + Wo'o", 2²) P)) /p + Enyny Wp(0' , w2)) (w, + (o,
(37)
S Engn WW3 /")]
(38)
(39)
Part (c). We provide counterexample. Fix any uo * 0 E U and yo: 1/1 E Y such that yo * V/-
Define μ = 2 (oyo + tg,)) Set up (k + 1)zu for k = 1,2.... and for each k define two measures
on Y x U by
1
1
T 2 (6gons + Sgrna)
uk 2 (ōgrno + dunvo)
(40)
It is clear that
WHl (1k, vk) = /uol WPink 2k) = min ( (uol"lll",, "ol").
(41)
Moreover, as k + we have W7' 2k) but Wr'kk, 7k) remains bounded. See Figure 4.
Part (d). First, the unconditional distance Wp (17, v) may be obtained via an unrestricted coupling in [[ v), i.e. the set of all joint measures on Y x U having marginals 17, v. Since II v) 2 (y,n) by part (a) we see that W,,n,) v) S Wn,
Let (3o, uo. y1, u1) = 7 (uo. un u1)6(yn - yo)u(()) be an optimal 7* E Ty v). We claim that 7(un, u1) : Jy 7 (wo, u1) du(()) couples u and ny.. Let T" (wo, u1) - uo be the projection onto the first coordinate of U U. Observe that for p-almost every y. we have that yy E [[('",, 2²) is optimal, and, in particular, = = sp. Fix an arbitrary & E Ca(U). We then have
1 p(uo) do ( = L (p o (po ( u1)
(42)
/ /
I (po x dr x)x "mo. du(s) [ pluo) d " dp(y) / TxU
(43)
= 1, / (l
pluo) dn(uu.v)
(44)
= - (pox")dn(uo) v) = / ş co cdrffi(uo))
(45)
18

---

91
3/0
Uo
Figure 4: The counterexample in Proposition The measure yi is shown black and the measure V4 is shown in white.
Thus 7 " 7- A similar argument shows that for the map n² (wo: U1) u1 have # = T U so that y E I1 π /
Now, as nr (wo: u1) E M( , is p-almost surely optimal the usual Wasserstein sense,
WP v) = / / uo - =1l dry (uo. u1) du(y))
(46)
U22 nuo ui dy(wu, u1)
(47)
2 WP WT (7 v
(48)
U. "1 u is coupling but potentially sub-optimal. since 7 E II[( 11((=
D.2 Geodesics
We now study the geodesics in the space (Y x U).
Theorem (IPH (Y x U) is Geodesic Space)
For any 4, v E Pu x U), there exists a constant speed geodesic between and v.
Proof. Write d (Y x U) + Y x U for the linear interpolant
(1 No, //, u1) (tyo (1 t)y1, two (1 t)u1)
0 0
(49)
Let E ITy vn) be an optimal restricted coupling, and consider the path of measures in Pp x U) given by
7t
7
0S15L.
(50)
Step We check that for each 0 5 / 1, have "t E (Y U). That is, need to check that for all Borel C Y, have 7A U) p(A). Indeed, recall that restricted measures concentrated on the set & (see Equation (19)). Thus,
7A x U) 7" (Ax
7 {VV, uo y, u1) y A}
(A) (='om119)
= ( = p(A)
i.e. 7A(A x Y) = ((4) as claimed.
Step two: We show that W/ (7t, 7s) - 1 (17, v). Set of (De,>a)#'** for 0 <1
We claim >f E (y) Indeed, we have T 7/ = because for all Borel A CC XU,
(Nr, (Ax U) x" (2 (A)) (A))
(51)
19

---

An analogous calculation shows that " = Yo, so that 32 E II( Ya). We now check that Ii E Ry(Y x U). Indeed, suppose E' E Y x U is a Borel set such that Eng Ø. In other words, for every (yo, uo, M1, u1) E E we have yo # yı. Set D = (A,, (B). We claim Dn = Ø, so that
(FE) (dc: )2) #7 7 (B) = (A,, (F))
(52)
= =7 86) =0.
(53)
Indeed, if c (y, no y, u1) E 8 then
(x..))(()) = (y, tuo + (1 t) y. suo + (1 =s)u) ) $ E
(54)
= c $ (Tee TT) (E).
(55)
Thus of E I1 (7. as claimed. Now, have
W 76) -
(xxery d (1o. uo, 21, u1) dAA (yo, uo, V1, u1)
J Juyo dp (Aelyo: uo, y1, u1, (vo. uo: y1, u1)) dr (yo: un, u0,31, u1)
J JoYx (I(t- (I(t-s) ())00 un)1 + 1(t suo 11)1 dy uo, V1, U1)
= || - s J(yxup d' uo V1, u1) dy (vo= "0, v1, u1)
|t s sl"WW
Conversely, an application of the previous inequality and the triangle inequality show that for 058S151,
v) < W 72) + W7( + 71) + W)
(56)
I sW7 2) + W7 76) + (1 (1- (t)
(57)
Rearranging the previous inequality implies |t - s.W - Wfi(te 76) for all s, E [0, 1]. and hence W (16, 7i)) = |t- sjPWH(n, v).
Theorem (Conditional McCann Interpolants) Fix n.v (Y U). Suppose T (y, u) (y. (u u)) is an injective triangular map solving
the conditional Monge problem (1). Define the maps T, Y x - x U for 0 < t S via T, (1 t)1 + ET" and define the curve of measures "/- = [Ti] E P7 U). Then,
(a) (rx) is absolutely continuous and a constant speed geodesic between
(b) The vector field ve (T; (y. u)) = (0, Ti u) u) generates the path "tt in the sense that (Tt, 04) solve the continuity equation (9).
Proof. Consider the function wt Y x U + U given by
We (y, u) (0, T7 (y, u) u. (0, we u)) and note this is precisely wx u) a,Ti (y, u). Define the vector field
(58)
(e u) = We a T (y. u) = (1, (we, o Thi T u))
(59)
For any & E Cyl(Y U). we have
d dt Yxu
/ p(v, 1) d(il,.)))
d dt yxu ply, (T
/
(60)
dl. Jyx plu, Tū, (u u)) dn(y.u)
d
(61)
1 (V Ti (v. w), (( du(w,u)
(62)
= -r (V u), v2 u)) dy
(63)
20

---

which shows that (7, 0x) solve the continuity equation.
Now, note that for S bs have
[.a | dt= S. (ucelo1007711 (v, drdeen 1/r
dt
(64)
J." Urre (eee lwoe u) du(e. 1/D
dt
(65)
S. (f -u u) an(v. (b- a),,((n,
1/8
de (66)
(67)
In particular, fo || dt < 00 and so by Theorem 4 (7x) is absolutely continuous.
A similar calculation shows that (6 - a)W7((, = W7 7a) = f In((t),, where the last line follows from the absolute continuity of % Thus, ||0:lle7(m.xxx0) l/'((t) for almost every E [0, 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 Theorem Informally, solution to the continuity equation with triangular vector field will result in the conditional measures almost surely satisfying the continuity equation as well.
Lemma (Triangular Vector Fields Preserve Conditionals)
Suppose ve u) (0, of (y, u)) is triangular and that (7x) C (Y U) a path of measures such that (we, 74) satisfy the continuity equation in the sense of distributions. Then, follows that for p--almost every E Y, we have dzy + V. (u (y. )n')) =.
Proof. Fix any & E Cyl(U I). Suppose vi Cyl(Y) is given, and note that (y))p(u, t) Cy x x I). As (Ve.7) solve the continuity equation, follows from the triangular structure of 2e that upon testing against thip we have
S.J, t(v) S, (Oup(u, (o (of(),,)) u), Vus(l, du)(u) dply) di 0.
(68)
Because (y)) E Cyl(Y), it is of the form ((T())) where क Y + R for some k 2 and E C22((Tt) Taking p to be a sequence of smooth approximations to the indicator function of an arbitrary rectangle E = E1 x Eg x x Es C RR we see
J. J (Diplu, t) + (0/(1) (y, u), vup(u,t)) dr dt du(y) =
(69)
As Y is separable, the Borel o-algebra on Y is generated by the cylinder sets, i.e. those which are precisely of the form л -FF) for some finite dimee-iimesiionll rectangle E.. We have thus shown that for an arbitrary Borel measurable set E Y.
[. (Bup(u, + (u. (u)p(,))) dn) di de = 0.
(70)
From this, it follows that J.I. (Dop(u, t) + (of (v.u), u), Vap(u, 6) (u) du(y) dt = 0 p-almost every
(71)
21

---

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 ig (9 Y x U) + (" Y x U) for 1/p + 1/q = 1 via
1 11019-2u w 10
w + 0
(((w)
(72)
=
which is the Fréchet differential of the convex functional calculation shows that this map satisfies
| A straightforward
(w)) (P 110 fyxe (g(w),, w) d7.
(73)
See also Ambrosio et al. [2005, Chapter 8].
Theorem (Absolutely Continuous Curves in PH x U)) Let I C R be an open interval, and suppose " - (Y x U) is an absolutely continuous in the Wh metric with (t(((t) E (I). Then, there exists Borel vector field vely. u) such that
(a) vr is triangular
(b) Ue E ("(,, Y U) and [ S (ll(t)) for a.. t
(c) (ve, 7t) solve the continuity equation the sense of distributions.
Proof. Assume without loss of generality that (t((t)) (f) and that I (0. 1) [Ambrosio et al., 2005, Lemma 1.1.4, Lemma 8. 8.1.3]. Fix any yp E Cyl(Y x U). For s, t / there exists an optimal triangular coupling "s E 1ly (1 7e). By Hölder's inequality,
hx(p)) - 7/s (p) s (ip())
(74)
follows that t + 7e(()) is absolutely continuous. We can introduce the upper semicontinuous and bounded map
V uo)l
(yo, uo) = (v1, ui) (yo, uo) # (y/1, u1)
(yo uo y1, (1)
[e(gn.uo)= - wn)l [(00. - (Mn
(75)
For [h| sufficiently small, choose any optimal coupling Yc+A)A E (7s+ 7/s) and note that
(yp)
(x(p)
/ Oxxuya
s hhl
(1oo. un) (y11 (yo, un Vı, U1) dy(s+))s
(76)
|/|
Wh Ris) hn
Vrrzy (90, uo y1. u1) d'(o+))a,
1/9
(77)
If t is a point of metric differentiability for t + note that ((++)e + (I,T) 4 narrowly, where is the identity map on x U. Moreover, since 7/. E (Y x U), it follows that on the diagonal we have that almost surely Hyyo, un, Yo, 41) = <(I Va un)). Thus,
h444(4)) - re(p))
< <l'((t) ( [ uo, yo, uo) dr uo)
1/9
lim sup
(78)
h
1h
-tt((t) [| ((VuP))ll
(79)
Taking Q = Y x U x y f ded fix any E Cyl(Q). We have that
I 0, u, 9) d)
(80)
J (bre p(u. a. u.
=lim lim
=
plu,,,
dro
1(xx0) ply, u, s) (y. ds.
22

---

An application of Fatou's Lemma, Equation (78), and Hölder's inequality gives us
1. u, dx(,,, for any interval J C with supp cJxYXU.
((1H/())a)) 1/P (u., s u,
2/9
(81)
Fix the subspace
V = [((V u, 8)) : v E Cyl(Q)) CYxU and denote by V its (r,, Y U I) closure. Define the linear functional L V R via
(82)
L(V.+) - 1 0,ply. u. dry,, s)
(83)
and note that Equation (81) implics that L is bounded linear functional on V. Thus (by Hahn- Banach and the fact that V V is dense) we may uniquely extend L to V. We thus have a convex minimization problem
min w - S |wly u. s)18 dr u, s) L(w)
(84)
which admits the unique solution w such that (w) T. that the above functional coercive and hence admits a minimizer which we may obtain via its differential as consequence of convexity. Thus, obtain triangular vector field v (w()) such that for all E Cyl(Q).
In particular, the estimate (81) shows
(v, vpp) / (u(y. (u(y, s),,, Villy, s)) dy(y, u, (L, Vp) = =- 0xply, dy u, s,
(85)
This precisely shows that (we, 7<) is triangular distributional solution the continuity equation.
Now, choose any interval / C / and choose a sequence =f* E (J), with 0 < T/* < and n + L, - oc. Moreover choose a sequence (V.Pa) V converging w (p(v) L7 Q). Our previous calculations give
I n (²)))) u, 5)1 dr(u, / ()), n Jo
lim / (t, vw d7
(86)
lim (L, Vn(T pm)) ( ((PP())) n-000
1/8 (nrr lulp dy
1/8
(87)
Taking & + xo we see that
I. (ur u) dyll, dt s (s)(
(88)
and since J C was arbitrary, we conclude
[1ull st1((t)
a.e.-t.
(89)
We now prove, in some sense, converse of the previous theorem.
Theorem 4 (Continuous Curves Generated by Triangular Vector Fields) Suppose that it I + (Y U) is narrowly continuous and (ve) is a triangular vector field such that (7e, v.) solve the continuity equation with [ivell E (I). Then, ?t. - P4 x U) is absolutely continuous in the WHth metric and hn((t) Le x for almost every t.
23

---

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 u-almost every y, followed by an application of this theorem. By Lemma 1, for ja-almost every y we have that (7h ot (v, -)) solve the continuity equation distributionally on I U.
By Jensen's s inequality (and the assumption 2 1) we see
1/p dt J /ellll dt = E Byny [|o] (v, -) -)
(90)
2 / Eynn [118 (v. - -)1 de
(91)
1 11 (v. -) -) di at)
= Eager
(92)
Since the first term is finite, it follows that
| (v. -11 E (T)
p-almost every y.
(93)
Now Ambrosio et al. [2005, Lemma 8.1.2] shows that for p-almost every y we have that (7 admits
narrowly continuous representative (5%) with 7/ = x for almost every t. It follows from Ambrosio et al. [2005, Theorem 8.3.1] IJ that for any t S t2 in L, we have
W7 72) $ (t2 t1)=-1
1/ (y, u)) di (u) dt
(94)
[
= =(t - 11)
lof (v. ") (w) dt
(95)
where the second line follows as -. 7/ for almost every t.
Let 5: fr 5% du(y) be the measure obtained via marginalizing over the Y-variables. Taking an expectation over ~ μu, the previous inequality shows us that
WH (fte n2a)
- t2 12
J
12
|10
|10 dt.
(96)
(t2 t1)
Now, note that t is almost surely Lebesgue point of the right- side and fe, %t. Taking l2 L along sequence where it Mh shows us that
11'(()) s [wellep(ao).x
(97)
for almost every f E J.
In the case that U is infinite dimensional, fix any y E Y such that Lemma 1 holds (which is of full measure) and fix a countable orthonormal basis (ex) for U. Set n" - R2d be the projection operator for this basis, i.e. u ((u, ei),... ea)). We consider the collection of finite d.y nd ?." By the same argument in Ambrosio al. [2005, day a solve the continuity Theorem 8.3.1] there exists vector field d.y on Rd such that (re. equation and
dimensional conditional measures 7:
d.y
1o (v. -
(98)
LP Rd)
follows from the dinteeddmeesional case above that for almost every t S 12, we have
WP d.y d.N) s (t2 -1)p=1
/ (v.-))l dt.
(99)
Let a.a.d 2z v.d where (z²))= d U maps + E =kek- As d Ad.y + 7/0 narrowly for all 1 E T. Since d= is an isometry, Ambrosio et al. [2005, Lemma 7.1.4] shows that
have
WP( Nh S inm Wp(n d.y d.y) S (tz t1)p-1 1/ a1l| (v, -11- dt. (100)
d
24

---

Now, integration with respect to du(y) yields
W nn)) I (t2 - t2 tz)0-1 S.x l leellloptneree dt.
(101)
Taking t2 + 11 shows that for almost every t we have
1r/(()) 1| ||0
(102)
Together, Theorem 3 and Theorem give dynamical interpretation of the conditional Wasserstein distance. The following result a conditional analogue of the well-known Benamou- 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-Brent
Let 1 < p < co. For any T), v LY U), we have
{o [Iolllee (ve, 7) solve (9), 70 = 71 = v, and Vt is triangular
p (1, (n, = min (Te.
Proof. Write M for the infimum on the right-hand side.
First, suppose that (Ut, pu) are admissible and So leelll (o()) < oo. It follows from Theorem 4 that (7x) is an absolutely continuous curve in (V x U) and |ieell 2 h'((() Thus,
( (1(0))a)))
s $ Ie [eellerpuu di M.
WH (T, (1, S
(103)
Conversely, by Theorem there exists a constant speed geodesic (7:) c PY x U) connecting and v. Recall that constant speed geodesics are absolutely continuous. By Theorem 3, there exists a Borel triangular vector field 1 such that (we, 7:)) solve the continuity equation, and moreover || [Iuellep(.. YxU) S (t)(t.. In fact, because (ve, v) solve the continuity equation, Theorem 4 yields that [leell XU) h'(().
WF(1),) (n, 2) for almost Since 74 is a constant speed geodesic in (Y U), it follows that [p' W every & (0, 1). Hence,
(1) 111(t)9 dt= 1 [eellBeeteexxx M.
(104)
Thus, WH.P(n. 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 be trained 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 setting for the 2D and Lotka-Volterra experiments, these share 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= o²1 where a a hyperparameter.
Our implementation of FM is adapted from the torchefm 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.
25

---

Table 4: Hyperparameter grid used for random search of the FM and COT-FM models on the 2D and Lotka- Lotka-Volterra datasets.
Hyperparameter
Description
Values
COT coupling strength
[1e-6, le-4, le-2, le-1]
σ
Variance for C = I in (10)
[le-3, 1e-2, le-1, 5c- I]
Batch Size
Training batch size Layer width in MIP
[256, 512, 1024]
Width
[256, 512, 1024, 2048]
LR
Learning rate Number of MLP layers
[1e-4, 3e-4, 7e-4, 1e-3]
Layers
[4. 6, 8]
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 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 W2 error the joint target distribution on 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
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., 201 1]. The moons dataset is generated with noise=0.05 followed by standard scaling with mean of m = (0.5, 0.25) and standard deviation of a = (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.72 75, followed by projection to the first two coordinates and re-scaling by 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.
Data Generation.
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 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 prohibitively expensive in terms of memory usage. However, note that use large batch sizes, meaning that the COT plan we find in this way should not be too far from optimal. All couplings 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 de- scribed in the main paper, we assume p(0) = (30, 1) and that log(u) ((m, 0.52) with m = (-0.125, -3, -0.125,-533 -3). Given parameters u E Ryoo we simulate Equation (13) for t E (0.2...., 20] to obtain solution 2(u) 2222 An observation y R23 is obtained by the addition of log-normal noise, i.e. log(y) ~ ((log(z(u), 0.11). We thus may simulate many (v, u) pairs from the target measure for training.
We generate 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
26

---

True
COT-FM (Our)
PCP-Map
COT-Flow
FM
Checkerboard
Moons
Ciroles
Swissroll
U
U
U
U
U
Figure 5: Samples from the ground-truth joint target distribution and the various models for the 2D datasets. Samples from COT-FM more closely match the ground-truth distribution than the baselines. common failure mode for the baselines is to generate samples from regions with zero support under the true data distributions. Table contains quantitative evaluation.
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 used in this experiment. The COT couplings computed in the same way as the previous section, and as described in Section 7.
27

---

COT-FM (Ours)
PCP-Map
COT-Flow
FM
Checkerboard
Moons
Circles
Swissroll
Figure 6: Conditional KDEs shown for each of the methods on the 2D datasets. The conditioning variable y is fixed the horizontal dashed line shown in Figure 5. In all plots, the orange solid line indicates the CKDE of the ground-truth joint samples. In each column, the dashed blue line indicates the CKDE of samples generated from the respective method.
28

---

Lotka- Volterra Samples: COT-FM (Ours)
0.50
0.73
1.00
1.25
B
1.26
0.50
0.0
0.1
0.05
0.10
0.0
5
1.25
0.1
1.3
0.30
0.7
0.02
0.06
0.02
0.06
0.02
0.06
0.02
0.04
0.06
Figure 7: KDE plots of the samples on the Lotka- Volterra system, using the settings described in Section 7. Plots include one-dimensional KDEs the diagonal, as well all two-dimensional pairs. In all plots, samples from MCMC are drawn in orange, and samples from our method (COT-FM) are indicated in blue. The true unknown parameters are indicated by the red vertical line the diagonal plots, or the black x in the off- off-diagonal plots.
29

---

Volte-Volterr Samples: PCP-Map
0.30
0.73
1.00
1:25
0.50
1.25
8
1.25
0.1
1.3
0.30
0.7
0.02
0.06
0.02
0.06
0.02
0.06
0.02
0.04
0.06
Figure 8: KDE plots of the samples on the Lotka Volterra system, using the settings described in Section 7. Plots include one-dimensional KDEs the diagonal, as well all two-dimensional pairs. In all plots, samples from MCMC are drawn in orange, and samples from PCP-Map are indicated in blue. The true unknown parameters are indicated by the red vertical line the diagonal plots, or the black x in the off-diagonal plots.
30

---

Lotka- Volterra Samples: COT-Flow
0.30
0.73
1.00
1:25
0.50
1.25
8
1.25
0.1
1.3
0.30
0.7
0.02
0.06
0.02
0.06
0.02
0.06
0.02
0.04
0.06
Figure 9: KDE plots of the samples on the Lotka Volterra system, using the settings described in Section 7. Plots include one-dimensional KDEs the diagonal, as well all two-dimensional pairs. In all plots, samples from MCMC are drawn in orange, and samples from COT-Flow are indicated in blue. The true unknown parameters are indicated by the red vertical line in the diagonal plots, or the black x in the off-diagonal plots.
31

---

Lotka- Voller- Samples: FM
0.50
0.73
1.00
1/25
0.50
1.25
8
1.25
0.1
1.3
0.30
0.02
0.06 0.02
0.06
0.02
0.06
0.02
0.04
0.06
Figure 10: KDE plots of the samples on the Lotka- Lotka-Volterra system, using the settings described in Section 7. Plots include one-dimensional KDEs the diagonal, as well all two-dimensional pairs. In all plots, samples from MCMC are drawn in orange, and samples from flow matching (FM) are indicated in blue. The true unknown parameters indicated by the red vertical line the diagonal plots, or the black x in the off- off-diagonal plots.
32

---

F33 Inverse Darey 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 regular 40 x 40 grid. The parameters are then exponentiated and used to simulate the permeability fields from the forward model 3 solving the Darcy flow PDE, using FEniCS [Alnæs al., 2015]. Stochasticity arises from adding Gaussian noise to the permeability fields, obtaining
(u)) +, 6 ~ (0, x² T). For our experiments we observe 100 x 100 grid, and we use a 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 data point for reference, Our source and target training sets contain 1 x 10 samples each, and our test set comprises 5 x 10ss samples. We remark that although y and are observed on grid their resolution does not need to be fixed, allowing for training at different resolutions.
Permeability field
Pressure field p
Observed
Figure 11: Example of one random data point from the Darcy flow dataset.
Models. In order make learning feasible infiiitediimensional Hilbert spaces, we adapt the architecture of a Fourier Neural Operator (FNO) [Li et al., 2020] from the neuraloperator package [Kovachki al., 2021] accommodate for conditioning information observed 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 x besddn_chhanneee original architecture. For all of the models in consideration, we fix the architecture to be have hidden hidden_channele channels = 64, projection_cee shheetiioncchamelll channels 256, and 32 Fourier modes. We train each model for 1500 epochs, and hyperparameters for each architecture are selected as follows:
input_diie
hidden_ channels x input_ _im, before following the hidden__channeeeenn
WaMGAN [Hosseini al., 2023]: using adaptation to the FNO architecture of the original code", we perform grid search as detailed in Table 5. We found the training procedure to be rather unstable, and for this reason 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 4, full critic iterations. and monotone penalty of x 10 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 x 10-1, and the covariance operator C' is set to match that of the prior, but rescaled by a factor of = 10-3 33 We use the code from the original repository".
COT-FFM: we set ε = x 10 in the cost function used to build the COT plan. The learning rate and C chosen to be the same FFM. In order build COT couplings, take the source measure to be the product measure π T 17 x (0, C). Approximate couplings are obtained on minibatches size 256.
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 the optimal vector field minimizing (12). Hence, the OT-CFM model [TTnn et al., 2023] is inapplicable here.
http:///sstiiioooo//o/Sooooooooo http::/giiiio.ocmmmmmaaiieinnnee flow_matching
33

---

Table 5: Hyperparameter search space for WaMGAN
Parameter
Search Space
Learning rate
[1 10 5 10 1 x 10 10-4)
Full critic iter.
5, 10]
Monotone penalty
[1x 10
x,5x 10 10 1 10
x,5x (1,5, 10)
Gradient penalty
Sampling, The resulting amortized sampler, denoted for simplification by the mapping (v, un) H M1 (y, un), will parameterize an approximate posterior measure. Notice that, contrast 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
(f) = ( dN(O,C)
Es(to (yk, uo,x)), {uo lid. NNO, (105)
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 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 's conditional 2- Wasserstein distance by sampling 10, 000 points from the source distribution, flowing each source point along the model's '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. as a function of batch size. We see that even a relatively small batch size of 16, the resulting error in the (squared) distances is less than 10% of the true value (&s 0.678). While this experiment is only feasible on synthetic data (as we must compute the true distance), nonetheless demonstrates that even with small batches we may recover the true distance.
Transport Distance Crror
0.15
0.1
True
0.00
-0.5
166
32
64
118
256
512
1024
Baicch 5ime
Figure 12: Error in the squared transportation distance function of batch size. Values closer to zero indicate a better approximation of the COT cost. For fairly small batch sizes (> 16) the magnitude of the error is stable and relatively small. Shading indicates standard deviation computed over ten random evaluation sets.
34

---

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 work 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 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 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.
is fine to include aspirational goals as motivation as long as is clear that these goals 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 Section 8.
Guidelines:
The answer NA means that the paper has 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 (c.g.. independence assumptions, noiseless settings, model sell-specfiiattinn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
35

---

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 cross- referenced.
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
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Our experimental settings are clearly described in Section 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, No answer 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 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 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 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 replicate the results, access a hosted model (e.g.. in the case of 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 submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or 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 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 possible for other researchers to have some path to reproducing verifying the results.
36

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
We will release code open-source upon acceptance of the paper for publication.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nis. wwwpccccppcc.aeeecc.ccccc.ccetle http:/ppplpgggpsss/e//pppssa httpg//ee//oeeuueiiieeeeeeeeeeee
for more details.
The authors should provide instructions on data access and preparation, including how 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. only 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
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 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, 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 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 emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
37

---

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.
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 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 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
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 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 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 failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /ttpa://///eip.....e http:ccccpbbccceeeeeeeee
ines?
Answer:
[Yes]
Justification:
We have reviewed the ethics guidelines and conform 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws 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
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
38

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess 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 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 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
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 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 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 best faith effort.
12. Licenses for existing assets
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?
Question:
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 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.
39

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of 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 not available online, the authors are encouraged to reach out to the asset's 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/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion 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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending 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.
40

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
41


Reviews:
[{"review_id": "Hf4PQVIvdy", "reviewer": "Reviewer_QX6V", "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.", "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. \n\nComment: I would suggest citing work on rectified flows as a concurrent invention of flow matching. \nhttps://arxiv.org/abs/2209.03003", "limitations": "Adequately discussed.", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n- The results effectively demonstrate how this approach can be used in a variety of settings.\n- The work is well written and fairly clear", "weaknesses": "- Only applied to relatively low dimensional problems.\n- 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\n    \n    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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their valuable feedback!\n\n> It would be good to know empirically how far from optimal the learned transport maps are, as this is a known limitation of mi[n]ibatch-based approaches.\n\nWe thank the reviewer for pointing this out. Please see our global response for a discussion regarding minibatches and COT. Overall, we find that our method COT-FM can recover nearly optimal transport distances even when trained at small batch sizes.\n\n> Only applied to relatively low dimensional problems.\n\nThis is a fair point. However, we would like to emphasize that our most significant contributions are to the theory of conditional optimal transport, laying the foundations for future applications. \n\nMoreover, many of the applications we have in mind are inverse problems, which are often fairly low-dimensional, e.g. recovering the parameters of an ODE as in our Lotka-Volterra experiment in Section 7. We also demonstrate our method on image data for the Darcy Flow inverse problem. While this is still relatively low-dimensional (at a grid size of 40x40), this was largely due to the computational cost of solving PDEs at high resolutions, which is necessary for both producing the training data as well as running MCMC as a point of comparison. \n\nThese results suggest that the method should scale to higher dimensional problems, which is an interesting avenue for future work.\n\n> 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.\n\nThank you for this suggestion -- we agree that this could be more clearly explained. We elaborate here on the key differences. \n\nWhen the conditioning variable takes values in a finite set, like a class label, the conditional optimal transport problem becomes much simpler. This is because, in this scenario, we typically have many observations (in the $U$ space) for any single, given conditioning variable $y$ -- e.g., many images all coming from the same class. In this case, one can simply solve the usual optimal transport problem for each class independently. \n\nHowever, we are largely interested in problems where we only have a single observation $u$ for each given $y$ -- this is the case, for instance, in inverse problems. Here, it is impossible to solve the optimal transport problem directly for each $y$ independently as there is simply not enough data. \n\nWe describe this briefly in lines 98-106, but we will make this point more clear. We will include pseudocode in our camera-ready submission to highlight the differences between our method and existing work.\n\n\n> I would suggest citing work on rectified flows as a concurrent invention of flow matching.\n\nWe thank the reviewer for the reference -- we are aware of rectified flows, and we have updated our paper with a short discussion of this work and how it relates to our submission and flow matching more broadly.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jhG9sDwLrm", "reviewer": "Reviewer_hqSP", "summary": "This 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.", "questions": "- 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?", "limitations": "The limitation is discussed in Weakness section.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper successfully extends the Benamou-Brenier theorem to the context of dynamical conditional optimal transport. \n- The paper is easy to follow.", "weaknesses": "- 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback!\n\n> The paper appears to be extremely similar to [1] [...] I would like to ask authors to discuss the difference with result in [1].\n\nWe first would like to remind the reviewer of the [NeurIPS concurrent work policy](https://neurips.cc/Conferences/2024/PaperInformation/NeurIPS-FAQ). Given that [Chemseddine 24] first appeared on 27 March 2024 and the NeurIPS submission deadline was 22 May 2024, this work counts as concurrent. This work was developed independently and concurrently with our submission, which we briefly mention in lines 78-80. However, there are several differences. An updated version of our paper will elaborate on these.\n\n1. [Chemseddine 24] work in finite-dimensional Euclidean spaces. Our theory is applicable in more general spaces -- including many infinite-dimensional function spaces. This necessitates additional techniques in our proofs (particularly Lemma 1 Theorem 4). Theorem 18 of [Chemseddine 24] can be viewed as a special case of our Theorem 5 under a finite-dimensional assumption. \n\n2. Prop 6 [Chemseddine 24] shows existence of vector fields generating a path of measures when the path is induced by an optimal plan. We show a stronger result in Theorem 3, where we characterize **all** absolutely continuous curves, rather than only those induced by an optimal plan. We also prove a converse of this statement in Theorem 4 which does not appear in [Chemseddine 2024]. \n\n> The paper discusses dynamical COT in Sections 4-5. [...] the experiments seem to address conditional generation rather than conditional optimal transport.\n\nAs discussed in our abstract and introduction, we are indeed primarily interested in applications in conditional generative modeling. Inspired by the success of OT couplings in unconditional generative modeling [Tong 23, Pooladian 23], we sought to employ similar techniques. However, prior to our work, the theory of conditional optimal transport was not sufficiently developed to justify this approach. Thus, a major contribution of our work is the development of such a theory. We believe that our theoretical contributions provide a foundation for COT in general, beyond only our proposed COT-FM method.\n\nWe agree with the reviewer that this point could have been made more clear, and we will update our paper to discuss the relationship between Sections 4-5 and Section 6. \n\n> This method solves dynamic COT only when the given joint coupling is the solution of COT. [...] the application as an OT method is extremely limited.\n\nYour understanding is correct -- however, we respectfully disagree that our method has limited applicability.\n\nFirst, approximating a static COT coupling is relatively straightforward [Carlier 10]. We discuss this in lines 257-265. However, **such a coupling is necessarily empirical** -- that is, the coupling is only between the observed data, and it is not straightforward to generalize to unseen data. Although our method requires these empirical static couplings at training time, we may apply our method to simulate the transport for unseen data. This is particularly useful when one is interested in modeling the transport for large datasets, for which the standard approaches to COT would be prohibitively expensive in terms of memory.\n\nSecond, as mentioned previously, we are primarily interested in applications in generative modeling. As demonstrated by our experiments, computing the static COT couplings is not a prohibitive step in solving conditional generation tasks. This is further emphasized by the success of OT based flow-matching algorithms for unconditional generation [Tong 23, Pooladian 23].\n\n> It is unclear if the experimental settings are appropriate\n\nAs discussed above, our methodological contributions are focused on conditional generative modeling. Our experiments are aligned with this aim -- in Section 7, we demonstrate that our proposed method, COT-FM, obtains strong empirical performance across a range of conditional generation tasks. We would be happy to hear suggestions for additional experiments that would further improve the paper. \n\n> the experiments were conducted on very small datasets.\n\nIn Appendix F, we discuss the size of our datasets -- namely, we train on 20,000 datapoints for our 2D experiments and 10,000 datapoints for our Lotka-Volterra and Darcy flow experiments. We respectfully disagree that these datasets are too small, as many real-world applications of inverse problems have datasets of roughly this magnitude. Moreover, we believe our method will scale to larger datasets, as the additional overhead as compared to e.g. standard flow matching is not prohibitively expensive.\n\n> Does it mean that FM model learn geodesics?\n\nYes -- the path of distributions modeled by COT-FM is a geodesic in the conditional Wasserstein space, in the sense of Theorem 1. \n\nInformally, Theorem 5 shows these geodesics are induced by triangular vector fields. Moreover, we show in Theorem 2(a) an optimal COT mapping induces a geodesic in this space, where the vector field producing this geodesic is given in Theorem 2(b). The vector fields in Theorem 2(b) are precisely those we use to learn COT-FM in Equation (7). While there is some necessary additional smoothing, in the limit of zero covariance we recover the true geodesics in theory. \n\nWe briefly discuss this in Lines 254-256 of our submission, but we agree that this could have been explained more clearly. If accepted, the camera-ready version of our paper will include a more extensive discussion.\n\n\n1. Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching. Chemseddine et al., 2024\n2. Improving and generalizing flow-based generative models with minibatch optimal transport. Tong et al., 2023.\n3. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. Pooladian et al., 2023.\n4. From Knothe's transport to Brenier's map and a continuation method for optimal transport. Carlier et al., 2010", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "i7ZrPWZL9C", "reviewer": "Reviewer_yVF4", "summary": "This 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.", "questions": "I 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}", "limitations": "N/A", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "This 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.", "weaknesses": "N/A :)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback and positive review! We have updated our paper to incorporate the various suggested references and stylistic edits.\n\n> Is there a clear way to choose the $\\epsilon$ parameter for the COT Flow Matching? Any heuristics whatsoever?\n\nThis is an interesting question -- while we do not have any rigorous results in this direction, we can perhaps give some intuition. Ideally, one would choose a small value of $\\epsilon$ -- as suggested by the results of Carlier [2010], with unlimited data this would indeed recover a triangular map. However, we expect there to be dependence between $\\epsilon$ and the size of the dataset (or batch size, if using minibatch couplings). Loosely speaking, $\\epsilon$ controls the tradeoff between triangularity and coupling nearby points; small values of $\\epsilon$ encourage the maps to be triangular at the cost of a potentially large distance in the $U$ component, whereas large values of $\\epsilon$ allow greater violations of triangularity while coupling points that are nearby in $U$ distance.\n\nMoreover, we would expect this $U$ distance to be larger for small batches of data, since there may not be a source point with the same (or nearby) $Y$ values as a target point. Thus we would expect that $\\epsilon$ and the batch size should be negatively correlated.\n\nHowever, in practice, $\\epsilon$ is a nuisance parameter which can be tuned just as any other parameter in the learning process. In practice, we tuned $\\epsilon$ through a grid search, measuring the loss on held-out validation data.\n\nWe will include a discussion of these notions in an updated camera-ready version of our paper.\n\n> Is equation (9) not due to the original flow matching papers?\n\nYes, essentially -- a special case of Equation (9) appears in the early work of [Theorem 2, Lipman 23], where the coupling depends only on the target variable $x_1$. An equivalent loss appears in Proposition 1 of [Albergo 23A], but still using an independent coupling. Theorem 3.2 of [Tong et al., 2023] extends this loss to general couplings, and [Albergo 23B, Pooladian 23] propose a similar use of general couplings. We will update the citation in our paper to better reflect this.\n\n1. Flow Matching for Generative Modeling. Lipman et al., 2023.\n2. Building Normalizing Flows with Stochastic Interpolants. Albergo et al., 2023(A).\n3. Stochastic Interpolants with Data-Dependent Couplings. Albergo et al., 2023(B).\n4. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. Tong et al., 2023.\n5. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. Pooladian et al., 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hcHuGwvn95", "reviewer": "Reviewer_i1c7", "summary": "This 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$.", "questions": "- 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?", "limitations": "-\tThe authors addressed the potential negative societal impact of their work.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- Whether the COT-FM can recover the dynamic optimal transport requires further clarification.\n- Please see the Questions Section below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for their constructive feedback and suggestions.\n\n> I would like to clarify the connection between Section (4, 5) and Section 6 [...] Section (4,5) justify the triangular parametrization of the standard Flow Matching model within the Conditional Wasserstein Space. Is this correct?\n\nYes, your understanding is correct. We agree that the relationship between our theoretical and methodological contributions could be better clarified, and we will update our paper to make this more clear.\n\nHere, we provide some additional context. Inspired by the empirical success of OT couplings in unconditional generative modeling [Tong 23, Pooladian 23], we sought to employ similar techniques for conditional generative modeling. However, prior to our work, the theory of conditional optimal transport was not sufficiently developed to justify this approach. Thus, a major contribution of our work (Sections 4, 5) is the development of such a theory which provides a principled foundation for COT-FM (Section 6). However, we believe that our theoretical contributions provide a foundation for COT in general, beyond only our proposed COT-FM method.\n\nWe additionally would like to clarify that COT-FM requires more than simply using triangular vector fields in the standard flow matching model. In particular, COT-FM requires one to additionally compute the static COT couplings, as discussed in lines 267-265.\n\n\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\nOur COT-FM indeed recovers the dynamic COT paths in the limit of zero smoothing. We write in Line 256 that this follows from a “pointwise application of [Tong et al., 2023, Prop. 3.4]”. Here, by pointwise, we mean that we may apply the result of [Tong 23] for a single, fixed $y$ to recover the optimal path conditioned on $y$ -- and thus, over all possible values of $y$.\n\nSomewhat more precisely, Prop 3.4 of [Tong 23] shows in the unconditional setting that $v_t^\\sigma(u) \\to v_t(u)$ as $\\sigma \\to 0$, where $v_t^\\sigma(u)$ is the smoothed flow matching vector field, and $v_t(u)$ is the dynamic OT vector field. In the *conditional* setting, our vector fields are of the form $v_t^\\sigma(y, u)$. However, if we fix $y$, this vector field can be viewed as a vector field purely over the $U$ component, thanks to the triangularity assumption. This allows us to directly apply Prop 3.4 of [Tong 23] -- in particular, because the COT problem is essentially a collection of unconditional OT problems, one for each fixed $y$. See e.g. [Hosseini 23].\n\nWe agree that this point could have been explained more clearly, and we will update our paper with a complete, formal proof of this result.\n\nIn our global response, we additionally conduct an experiment to measure the degree to which the COT paths are optimal. Overall, we find that the paths learned by COT-FM are close to optimal, even when using the minibatch approximation.\n\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\nWe would like to clarify that we measure the joint metrics for practical reasons. Namely, to measure the conditional MMD or Wasserstein distances, we would require the ability to sample many values of $u \\sim \\nu (u \\mid y)$ from the data distribution, conditioned on a given $y$. However, in many practical instances (including the Lotka-Volterra, Darcy Flow, and many of the 2D datasets), this is not possible, as the data generation mechanism either directly produces sample $(y, u) \\sim \\nu(y, u)$ from the joint distribution (as in the 2D data) or produces samples $(y, u) \\sim \\nu(u) \\nu(y \\mid u)$ from a prior over $u$ and a forward model $\\nu(y \\mid u)$ (as in the Lotka-Volterra and Darcy Flow problems).\n\nMoreover, even in cases where we can indeed sample $u \\sim \\nu(u \\mid y)$, we found that estimators of the conditional distances have prohibitively high variance, as it is in essence a nested sampling problem. \n\nThus, some approximate measure of model fit is necessary. We believe the joint metrics are a reasonable proxy since, by Proposition 1, any triangular mapping which couples the joint distributions necessarily couples the conditional distributions.\n\nWe thank the reviewer for pointing out this difference, and we will use the additional space afforded by a camera-ready version of our paper to discuss these challenges and justifications.\n\n\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\nPlease see our global response for a discussion regarding the role of minibatches in our method. We would additionally like to point out that we do indeed compare our method against flow matching with an independent coupling. This is “FM” in Table 1 and Table 2, and “FFM” in Table 3. Overall, there is a clear gap between COT-FM and FM with independent couplings across all experiments considered in our work.\n\nWe agree, though, that this could have been better explained in our submission. We will update our paper with a more clear explanation of the baseline flow matching model. \n\n1. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. Tong et al., 2023.\n2. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. Pooladian et al., 2023.\n3. Conditional Optimal Transport on Function Spaces. Hosseini et al., 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Summary
We would like to thank the reviewers for their detailed and valuable feedback. We are encouraged to hear that the reviewers found our theoretical contributions to be a strength (QX6V hqSP, i1c7, yVF4) and that the submission is well-written (QX6v, hqSP, yVF4, i1c7). We are also glad the reviewers found our proposed method “quite elegant” (yVF4), with experimental results that clearly demonstrate the strengths of the approach (QX6v, yVF4).  

# Minibatch COT
Several reviewers had questions regarding the role of minibatches in our setup. While the use of minibatches is needed to scale the method to large datasets, we believe that this limitation is not prohibitive in practical adaptations of these methods. Our experiments (Tables 1, 2, 3) demonstrate that even with a minibatch approximation, COT-FM outperforms standard flow matching (without the use of COT).

The precise relationship between minibatch OT and the standard OT problem is an area of active research. However, there have been some promising results, at least in the unconditional setting. For instance, [Bernton 19] show that minimizers of the batched Wasserstein distance converge to minimizers of unbatched Wasserstein distance (as the batch size grows), and [Sommerfeld 19] establish non-asymptotic bounds between the Wasserstein distance and its minibtach approximation. Studying these relationships for the conditional OT problem is an open, challenging problem and an exciting direction for future work. 

To further investigate the role of minibatches, we conducted an additional experiment. In this experiment, we use a standard Gaussian as our source distribution, and a Gaussian with covariance $\rho = 0.75$ as the target. We choose these distributions as the COT distance is available in closed form, as we prove in Section 4 / Appendix C. We then train our proposed COT-FM model on a range of different batch sizes (leaving all other hyperparameters fixed except for the learning rate). We then measure the degree to which the transport learned by COT-FM is optimal. Figures for our preliminary results are contained in the attached response. 

In Figure 1, we plot the difference between the true value of the squared COT distance $W_2^{\mu, 2}$, and the squared transport cost resulting from the trained COT-FM model. This learned cost is estimated by sampling 10,000 points from the source distribution, flowing each point along the model’s learned vector field, and computing the Euclidean distance between the initial and terminal location of each point. We see that even for relatively small batch sizes, the transport costs from our learned COT-FM model closely match the ground-truth cost.

In Figure 2, we additionally measure the degree to which the learned dynamic COT paths are optimal. To do so, we measure the difference between the path energy of the model $\int_0^1 \int |v_t^\theta| d p_t d t$ and the known ground-truth distance $W_2^{\mu, 2}$. By Theorem 5 of our work, these two quantities should be equal if the learned dynamic paths are optimal. We find that the deviation decreases for larger batch sizes, but that the magnitude of the deviation is fairly small relative to the distance.

Overall, these two results indicate that COT-FM is able to obtain good approximations of both the dynamic and static COT solution, even when using minibatches. These results are in agreement with similar results in the unconditional case [Tong 23, Fatras 21]. We will conduct additional experiments with the datasets considered in the main submission to assess the degree of optimality in an updated version of our paper.

1. On parameter estimation with the Wasserstein distance. Bernton et al., 2019.
2. Improving Mini-batch Optimal Transport via Partial Transportation. Nguyen et al., 2021.
3. Optimal Transport: Fast Probabilistic Approximation with Exact Solvers. Sommerfeld et al., 2019.
4. Unbalanced minibatch Optimal Transport; applications to Domain Adaptation. Fatras et al., 2021.
5. Improving and generalizing flow-based generative models with minibatch optimal transport. Tong et al., 2023.

Meta-review: The reviewers were quite split on this submission. Taking into account my own evaluation, as well as the expertise of the reviewer, I recommend to accept this paper.
Decision: Accept (poster)