Dynamic Conditional Optimal Transport through Simulation-Free Flows

Authors: gavin  kerrigan
giosue  migliorini
padhraic  smyth
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 18:24:52

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.

Main paper: 
 introduction: 
 Many fundamental tasks in machine learning and statistics may be posed as modeling a conditional distribution ν(u | y) where obtaining an analytical representation of ν(u | y) is impractical. While sampling-based approaches such as Markov Chain Monte Carlo (MCMC) methods are useful, they suffer from several limitations. First, MCMC requires numerous likelihood evaluations, rendering it prohibitively expensive in scientific and engineering applications where the likelihood is determined by an expensive numerical simulator. Second, MCMC must be run anew for every observation y, which is impractical in applications such as Bayesian inverse problems [Dashti and Stuart, 2013] and generative modeling [Mirza and Osindero, 2014] . These limitations motivate the need for a likelihood-free [Cranmer et al., 2020] and amortized [Amos et al., 2023] approach. While methods like ABC [Beaumont, 2010] and variational inference [Blei et al., 2017] partially address these challenges, they are difficult to scale to high dimensions or have limited flexibility.Recently, generative models such as normalizing flows [Papamakarios et al., 2019 [Papamakarios et al., , 2021]] , GANs [Ramesh et al., 2022] , and diffusion models [Sharrock et al., 2022] have shown promise in amortized and likelihood-free inference. These models may be viewed in the framework of measure transport [Baptista et al., 2020] , where samples u ∼ η(u) from a tractable source distribution are transformed by a mapping T (y, u) such that the transformed samples are approximately distributed as ν(u | y). One way to achieve this is through triangular mappings [Baptista et al., 2020 , Spantini et al., 2022] , where a joint source distribution η(y, u) is transformed by a mapping of the form T : (y, u) → (T Y (y), T U (y, u)). Under suitable assumptions, if T transforms the source η(y, u) into the target ν(y, u), then T U (y, -) couples the conditionals η(u | y) and ν(u | y).Typically, such a map T is not unique [Wang et al., 2023] , and a natural idea is thus to regularize the transport and search for an admissible mapping that is in some sense optimal. In other words, learning a conditional sampler may be phrased as finding a conditional optimal transport (COT) map. While there exists some work on learning COT maps, these approaches often rely on a difficult adversarial optimization problem [Baptista et al., 2020 , Hosseini et al., 2023 , Bunne et al., 2022 , Ray et al., 2022] or simulating from the model during training [Baptista et al., 2023 , Wang et al., 2023] . In this work, we propose a conditional generative model for likelihood-free inference based on a dynamic formulation of conditional optimal transport. Specifically, our contributions are as follows:1. We develop a general theoretical framework for dynamic conditional optimal transport in separable Hilbert spaces. Our framework is applicable in infinite-dimensional spaces, enabling applications in function space Bayesian inference. In Section 4, we study the conditional Wasserstein space P µ p (Y × U ) and show that this space admits constant speed geodesics between any two measures. In Section 5, we characterize the absolutely continuous curves of measures in P µ p (Y ×U ) via the continuity equation and triangular vector fields. As a consequence, we obtain conditional generalizations of the McCann interpolants [McCann, 1997] and the Benamou-Brenier Theorem [Benamou and Brenier, 2000] .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 a path of measures interpolating between an arbitrary source and target distribution via a 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-dimensional problem arising from the Darcy Flow PDE. Our method shows competitive performance against recent COT methods. 
related work: 
 Conditional Optimal Transport. Conditional Optimal Transport (COT) remains relatively underexplored in both machine learning and related fields. Recent approaches learn static COT maps via input convex networks [Bunne et al., 2022 , Wang et al., 2023] or normalizing flows [Wang et al., 2023] . In addition, there have been a number of heuristic approaches to conditional simulation through W-GANs [Sajjadi et al., 2017 , Adler and Öktem, 2018 , Kim et al., 2022 , 2023] , for which Chemseddine et al. [2023] provide a rigorous basis. Closely related to our work are those which employ triangular plans [Carlier et al., 2016, Trigila and Tabak, 2016] , which have been modeled through GANs in Euclidean spaces [Baptista et al., 2020] and function spaces [Hosseini et al., 2023] .In contrast, we use a novel dynamic formulation of COT, which we model through a generalization of flow matching [Lipman et al., 2022 , Albergo et al., 2023b , Liu et al., 2022] . This allows us to use flexible architectures while avoiding the difficulties of training GANs [Arora et al., 2018] .Simulation-Free Continuous Normalizing Flows. Flow matching [Lipman et al., 2022] (and the closely related stochastic interpolants [Albergo et al., 2023a] and rectified flows [Liu et al., 2022] ) are a class of methods for building continuous-time normalizing flows in a simulation-free manner. Notably, these works do not approximate an optimal transport between the source and target measures. Pooladian et al. [2023] and Tong et al. [2023] propose instead to couple the source and target distributions via optimal transport, leading to marginally optimal paths. In this work, we study an extension of these techniques for conditional generation.While some works [Davtyan et al., 2023 , Gebhard et al., 2023 , Isobe et al., 2024 , Wildberger et al., 2024] have applied flow matching for conditional generation, these approaches do not employ COT. Notably, these approaches are limited to the finite-dimensional setting, whereas our method adds to the growing literature on function-space generative models [Hosseini et al., 2023 , Kerrigan et al., 2023 , 2024 , Lim et al., 2023 , Franzese et al., 2024] . Concurrent work by Chemseddine et al. [2024] develops the foundation of dynamic COT with applications to flow matching, and concurrent work by Barboni et al. [2024] develop the theory of dynamic COT for the purposes of studying infinitely deep ResNets. However, these works focus on the Euclidean setting, whereas our methods are applicable in general separable Hilbert spaces.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The limitations of the work are discussed in Section 8.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
background and notation: 
 Let X, X ′ represent arbitrary separable Hilbert spaces, equipped with the Borel σ-algebra. We use P(X) to represent the space of Borel probability measures on X, and P p (X) ⊆ P(X) to represent the subspace of measures having finite pth moment. If η ∈ P(X) is a probability measure on X and T : X → X ′ is measurable, then the pushforward measure T # η(-) = η(T -1 (-)) is a probability measure on X ′ . Maps of the form e.g. π X : X × X ′ → X represent the canonical projection.We assume that we have two separable Hilbert spaces of interest. The first, Y , is a space of observations, and the second, U , is a space of unknowns. These spaces may be of infinite dimensions, but a case of practical interest is when Y and U are finite dimensional Euclidean spaces. We will consider the product space Y × U , equipped with the canonical inner product obtained via the sum of the inner products on Y and U , under which the space Y × U is also a separable Hilbert space. Let η ∈ P(Y × U ) be a joint probability measure. The measures π Y # η ∈ P(Y ) and π U # η ∈ P(U ) obtained via projection are the marginals of η. We use η y ∈ P(U ) to represent the measure obtained by conditioning η on the value y ∈ Y . By the disintegration theorem [Bogachev and Ruas, 2007, Chapter 10] , such conditional measures exist and are essentially unique, in the sense that there exists a Borel set E ⊆ Y with π Y # η(E) = 0, and the η y are unique for y / ∈ E.In conditional optimal transport, we are given a target measure ν ∈ P(Y × U ) and some source measure η ∈ P(U ). We seek a transport map T : Y × U → U such that, for any given y ∈ Y , the mapping T (y, -) : U → U transforms the source distribution η into the conditional distribution ν y , i.e., T (y, -) # η = ν y . In a sense, T can be thought of as a collection of transport maps indexed by y ∈ Y . If such a T were available, by drawing samples u 0 ∼ η and transforming them, one would obtain samples T (y, u 0 ) ∼ ν y . Solving this transport problem for each fixed y is expensive at best, or impossible when only has a single (or no) samples (y, u) ∼ ν for any given y. Thus, one must leverage information across different observations y. To that end, recent work has focused on the notion of triangular mappings T : Hosseini et al., 2023 , Baptista et al., 2020] of the form T (y, u) = (T Y (y), T U (T Y (y), u)) for some T Y : Y → Y and T U : Y × U → U . Triangular mappings are of interest as they allow us to obtain conditional couplings from joint couplings. Proposition 1 (Theorem 2.4 [Baptista et al., 2020] , Prop. 2.3 [Hosseini et al., 2023] ) y) for π Y # η-almost every y. In many scenarios of practical interest, the source measure η and the target measure ν have the same Y -marginals. We will henceforth make this assumption, and use µ = π Y # η = π Y # ν to represent this marginal. In this case, we may take T Y to be the identity mapping, so that the conclusion of Proposition 1 simplifies to T U (y, -) # η y = ν y for µ-almost every y. We note that in situations where such an assumption does not hold, one may simply preprocess the source measure η via an invertible mapping Hosseini et al., 2023, Prop 3.2] . Given a source and target measures η, ν ∈ P µ (Y × U ) and a cost function c : (Y × U ) 2 → R, the conditional Monge problem seeks to find a triangular mapping solvingY × U → Y × U [Suppose η, ν ∈ P(Y × U ) and T : Y × U → Y × U is triangular. If T # η = ν, then T U (T Y (y), -) # η y = ν T Y (T Y satisfying [T Y ] # [π Y # η] = π Y # ν [EQUATIONThe conditional Monge problem also admits a relaxation under which one only considers couplings whose Y -components are almost surely equal. To that end, for η, ν ∈ P µ p (Y × U ) we define the set of triangular couplings Π Y (η, ν) to be the couplings of η and ν that almost surely fix the Y -components,EQUATIONIn other words, a triangular coupling γ ∈ Π Y (η, ν) has samples (y 0 , u 0 , y 1 , u 1 ) ∼ γ such that y 0 = y 1 almost surely. The conditional Kantorovich problem seeks a triangular coupling solvinginf γ (Y ×U ) 2 c(y 0 , u 0 , y 1 , u 1 ) dγ(y 0 , u 0 , y 1 , u 1 ) | γ ∈ Π Y (η, ν) .(3)Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• 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. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: The potential societal impacts are discussed in Section 8. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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 of constructing geometric tangent spaces in the usual Wasserstein space. Definition 1 (Conditional Wasserstein Space) Suppose µ ∈ P(Y ) is given and 1 ≤ p < ∞. The conditional p-Wasserstein space isEQUATIONWe now equip P µ p (Y × U ) with a metric W µ p , the conditional Wasserstein distance. Intuitively, the conditional Wasserstein distance measures the usual Wasserstein distance between all of the conditional distributions in expectation under the fixed Y -marginal µ.Definition 2 (Conditional p-Wasserstein Distance) Suppose η, ν ∈ P µ p (Y × U ) and 1 ≤ p < ∞. The function W µ p : P µ p (Y × U ) × P µ p (Y × U ) → R, W µ p (η, ν) = E y∼µ W p p (η y , ν y ) 1/p = Y W p p (η y , ν y ) dµ(y) 1/p (5)is the conditional p-Wasserstein distance. W p is the usual Wasserstein distance for measures on U .By Jensen's inequality we haveW µ p (η, ν) ≥ E y∼µ [W p (η y , ν y )].In the following, we show that the conditional Wasserstein distance is a well-defined metric as well as a few other metric properties. Proposition 2 (Some Properties ofW µ p ) Let 1 ≤ p < ∞.(a) W µ p is well-defined, finite, and equals the minimal conditional Kantorovich cost. (b) W µ p is a metric on the space P µ p (Y × U ). (c) There does not exist C > 0 such that W µ p (η, ν) ≤ CW p (η, ν) for all η, ν ∈ P µ p (Y × U ).(d) For all η, ν ∈ P µ p (Y × U ), W p π U # η, π U # ν ≤ W µ p (η, ν) and W p (η, ν) ≤ W µ p (η, ν).Proposition 2(c, d) together shows that one should expect the topology generated by W µ p to be stronger than the unconditional distance W p . Here, we note that Gigli [2008] and Chemseddine et al. [2023] previously showed that W µ p is a metric through an equivalence with restricted couplings. Our approach builds on the results of Hosseini et al. [2023] and is somewhat more direct, and hence our proofs may be of independent interest.For the sake of concreteness, we include an example where the conditional 2-Wasserstein distance may be explicitly computed. Here, the necessary calculations follows from the fact that the conditional distributions of a multivariate are again Gaussian, and Gaussian distributions admit a closed-form expression for the usual unconditional 2-Wasserstein distance.Example: Gaussian Measures. Suppose Y = U = R and that η, ν ∈ P µ p (Y × U ) are Gaussians η = N (0, I) ν = N 0, 1 ρ ρ 1 |ρ| < 1. (6) It follows that µ = π Y # η = π Y # ν = N (0, 1). As η y , ν y are Gaussians, their W 2 distance admits a closed form and we can directly compute the expectation in Equation ( 5) to obtainW µ,2 2 (η, ν) = 2(1 -1 -ρ 2 ). This is zero if and only if ρ = 0, i.e. η = ν. However,π U # η = π U # ν = N (0, 1) and W 2 (π U # η, π U # ν) = 0 regardless of ρ. Moreover, the uncondi- tional distance is W 2 2 (η, ν) = 2 2 - √ 1 -ρ - √ 1 + ρ , from which it is easy to verify that W 2 (η, ν) ≤ W µ 2 (η, ν).See Appendix C for a similar derivation which applies to arbitrary Gaussians.Conditional Wasserstein Space as a Geodesic Space. We now turn our attention to the geodesics in P µ p (Y × U ). In particular, we show that there exists a constant speed geodesic between any two measures in P µ p (Y × U ), generalizing a similar result in the unconditional setting [Santambrogio, 2015, Theorem 5.27] . Moreover, we show that under suitable regularity assumptions, solutions to the conditional Monge problem (1) induce constant speed geodesics. Our motivation for studying geodesics in P µ p (Y × U ) is practical -in Section 6, we show how one can model geodesics in P µ p (Y × U ) in order to obtain a conditional flow-based model whose paths are easy to integrate. A curve is a continuous function γ • :EQUATIONIf (γ t ) is an absolutely continuous curve, then its metric derivativeEQUATIONexists for almost every t ∈ (a, b), and, moreover, we almost surely have |γ ′ |(t) ≤ m(t) pointwise for any m satisfying Equation ( 7) [Ambrosio et al., 2005 , Theorem 1.1.2]. A curve (γ t ) is called a constant speed geodesic if for all a < s ≤ t < b, we have W µ p (γ s , γ t ) = |t -s|W µ p (γ a , γ b ).It is straightforward to show that every constant speed geodesic is absolutely continuous. Theorem 1 (P µ p (Y × U ) is a Geodesic Space) For any η, ν ∈ P µ p (Y × U ), there exists a constant speed geodesic between η and ν.When an optimal triangular coupling γ ⋆ ∈ Π Y (η, ν) is induced by an injective triangular map T ⋆ , we may recover a constant speed geodesic in P µ p (Y × U ), generalizing the McCann interpolant [McCann, 1997] to the conditional setting. We refer to Proposition 4 for sufficient conditions on η, ν under which such a T ⋆ exists. Informally, samples from (y 0 , u 0 ) ∼ η flow in a straight path at a constant speed to their destination T ⋆ (y 0 , u 0 ). Theorem 2 (Conditional McCann Interpolants) Fix η, ν ∈ P µ p (Y × U ). Suppose T ⋆ (y, u) = (y, T ⋆ U (y, u)) is an injective triangular map solving the conditional Monge problem (1). Define the maps T t : Y × U → Y × U for 0 ≤ t ≤ 1 via T t = (1 -t)I + tT ⋆ , and define the curve of measures γ t = [T t ] # η ∈ P γ p (Y × U ). Then, (a) (γ t ) is absolutely continuous and a constant speed geodesic between η, ν (b) The vector field v t (T ⋆ t (y, u)) = (0, T ⋆ U (y, u) -u) generates the path γ t , in the sense that (γ t , v t ) solve the continuity equation (9). 
conditional benamou-brenier theorem: 
 In this section, we prove a characterization of the absolutely continuous curves in P µ p (Y × U ). As a corollary, we obtain a conditional generalization of the Benamou-Brenier Theorem [Benamou and Brenier, 2000] , giving us a dynamic characterization of the conditional Wasserstein distance. Roughly speaking, all such curves are generated by a vector field on Y × U which has zero velocity in the Y component. This is natural, as all measures in P µ p (Y × U ) have a fixed Y -marginal µ. Such a vector field can be informally seen as tangent to a curve of measures, and is the dynamic analogue of the triangular maps discussed in Section 3. More formally, given an open interval I ⊆ R, a time-dependent Borel vector field v :I × Y × U → Y × U is said to be triangular if there exists a Borel vector field v U : I × Y × U → U such that v t (y, u) = 0, v U t (y, u) .Continuity Equation. We introduce some necessary background which allows us to link vector fields to curves of measures. The continuity equation ∂ t γ t + div(v t γ t ) = 0 describes the evolution of a measure γ t which flows along a given vector field v t [Ambrosio et al., 2005, Chapter 8 ]. This equation must be understood in the sense of distributions, i.e. for every φ in a space of test functions,EQUATIONWe consider cylindrical test functions φ ∈ Cyl(Y ×U ×I), i.e. of the form φ(y, u, t) = ψ(π d (y, u), t) where π d : Y × U → R d maps (y, u) → (⟨(y, u), e 1 ⟩, . . . , ⟨(y, u), e d ⟩) and {e 1 , e 2 , . . . , e d } is any orthonormal family in Y × U . In the finite dimensional setting, one may take φ ∈ C ∞ c (Y × U ) to be smooth and compactly supported [Ambrosio et al., 2005, Remark 8.1.1] .In Appendix E, we prove Lemma 1, which is key in proving Theorem 4 below. Informally, Lemma 1 states that if the weak continuity equation ( 9) is satisfied for a joint distribution and triangular vector field, then the continuity equation is also satisfied for the corresponding conditional distributions and U components of the vector field.Lemma 1 (Triangular Vector Fields Preserve Conditionals) Suppose v t (y, u) = (0, v U t (y, u)) is triangular and that (γ t ) ⊂ P µ p (Y × U ) is a path of measures such that (v t , γ t ) satisfy the continuity equation in the sense of distributions. Then, it follows that for µ-almost every y ∈ Y , we have∂ t γ y t + ∇ • (v U t (y, -)γ y t ) = 0.We note that having v t be triangular is sufficient, but certainly not necessary, for the conditional continuity equation to almost surely hold. For instance, the vector field in R d that rotates N (0, I) about the origin is not triangular yet preserves all conditional distributions.Absolutely Continuous Curves. In this section, we state our characterization of absolutely continuous curves in P µ p (Y × U ). Informally, given such a curve, Theorem 3 provides us with a triangular vector field which generates the curve, in the sense that the pair solve the continuity equation. As a corollary of Theorem 3 and Theorem 4, we obtain a conditional version of the Benamou-Brenier theorem [Benamou and Brenier, 2000 ]. Once we have our characterization of absolutely continuous curves provided by these theorems, the proof of Theorem 5 largely follows the unconditional case (see e.g. Ambrosio et al. [2005, Chapter 8] ), but we include it for the sake of completeness.Theorem 5 (Conditional Benamou-Brenier) Let 1 < p < ∞. For any η, ν ∈ P µ p (Y × U ), we haveW p,µ p (η, ν) = min (γt,vt) 1 0 ∥v t ∥ p L p (µt) dt | (v t , γ t ) solve (9), γ 0 = η, γ 1 = ν,and v t is triangular . 
cot flow matching: 
 We have thus far seen that the COT problem (3) admits a dynamic formulation by Theorem 5, where one may take the underlying vector fields to be triangular. We use these results to design a principled model for conditional generation based on flow matching [Lipman et al., 2022 , Albergo et al., 2023b , Liu et al., 2022 , Tong et al., 2023 , Pooladian et al., 2023] . We hereafter use the squared-distance cost (i.e. p = 2). Flow Matching. We assume that we have access to samples z 0 = (y 0 , u 0 ) ∼ η(y 0 , u 0 ) ∈ P µ p (Y × U ) from a source measure, and samples z 1 = (y 1 , u 1 ) ∼ ν(y 1 , u 1 ) ∈ P µ p (Y × U ) from a target measure. Let z = (z 0 , z 1 ) ∼ ρ(z 0 , z 1 ) ∈ Π(η, ν) be any coupling of the source and target measure. We specify a collection of measures and vector fields on Y × U viaEQUATION)where C is any trace-class covariance operator [Da Prato and Zabczyk, 2014] . As is standard in flow matching [Lipman et al., 2022 , Kerrigan et al., 2024] , we obtain from Equations (10) a marginal measure γ t (y, u) and vector field v t (y, u) satisfying the continuity equation viaEQUATIONThis marginal path (γ t )1 t=0 interpolates between the source measure (t = 0) and a smoothed version of the target measure (t = 1). To transform source samples from η into target samples from ν, we seek to learn the intractable vector field v t (y, u) with a model v θ (t, y, u) by minimizing the loss 1L(θ) = E t,ρ(z),γt(y,u|z) v θ (t, y, u) -v t (y, u | z) 2 (12)which has the same θ-gradient as the MSE loss to the true vector field u t (y, u) [Tong et al., 2023] .COT Flow Matching. In the preceding section, ρ(z) may be an arbitrary coupling between η and ν. Motivated by Proposition 1, we will choose ρ to be a COT coupling. Under sufficient regularity conditions (see Appendix B), this COT plan will be induced by a triangular map. In turn, Theorem 2 gives us that this triangular map is generated by a triangular vector field of the form (10). Thus, we parametrize our model u θ to also be triangular. Moreover, we recover the optimal dynamic transport given in Theorem 5 as Tr(C) → 0 by a pointwise application of [Tong et al., 2023, Proposition 3.4 ].Given a collection of samples {z i 0 , z i 1 } n i=1 drawn from η and ν, we approximate a conditional optimal coupling ρ using standard numerical techniques with the cost function c ϵ (y 0 , u 0 , y 1 , u 1 ) = |y 1 -y 0 | 2 + ϵ|u 1 -u 0 | 2 for some 0 < ϵ ≪ 1. Intuitively, such a cost penalizes mass transfer along the Y dimension, which is precisely the constraint sought in the COT problem (3). As ϵ ↓ 0, we recover the true optimal triangular map [Carlier et al., 2010 , Hosseini et al., 2023] . The COT coupling 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. can either be precomputed for small datasets or computed on each minibatch drawn during training.Checkerboard Moons Circles Swissroll W 2 (1e-2) MMD (1e-3) W 2 (1e-2) MMD (1e-3) W 2 (1e-2) MMD (1e-3) W 2 (1e-2) MMD (1e-3) PCP-While the use of minibatches is a computational necessity, we find that surprisingly small batch sizes still yields accurate approximations of the true COT mapping using our COT-FM method. See Appendix G.After training, we obtain a learned triangular vector field v θ (t, y, u). Given an arbitrary fixed y ∈ Y , we may approximately sample from the target ν(u | y) by sampling u 0 ∼ η(u 0 | y) and numerically solving the corresponding flow equation ∂ t (y, u t ) = v θ (t, y, u t ) with initial condition (y, u 0 ).Source Measure. Our framework is agnostic to the choice of source measure η, allowing for flexibility in the modeling process. The main requirement is that the Y -marginals of the source η and target η must match. In some scenarios, this is trivially satisfied. If one is interested in using a source distribution which is simply random noise, one may take η(y 0 , u 0 ) = π Y # ν(y 0 ) ⊗ η U (u 0 ) to be the product of two independent distributions where η U is arbitrary, e.g. Gaussian noise. We now illustrate our methodology (COT-FM) on a variety of conditional simulation tasks. We compare our method against several competitive baselines, namely PCP-Map [Wang et al., 2023] , COT-Flow [Wang et al., 2023] , and WaMGAN [Hosseini et al., 2023] . These baselines are chosen as they reflect current state-of-the-art approaches to learning COT maps. We additionally compare against flow matching [Lipman et al., 2022 , Wildberger et al., 2024] without COT, i.e. where the coupling between the source and target measures is the independent coupling ρ(z 0 , z 1 ) = η ⊗ ν. This baseline serves as an ablation for the COT component of our model. 
experiments: 
 Overall, our method (COT-FM) typically outperforms these baselines across the diverse and challenging set of tasks we consider. We find that PCP-Map [Wang et al., 2023 ] is a strong baseline, but we emphasize that this model relies on the use of an input-convex neural network [Amos et al., 2017] and it is hence unclear how to adapt this method to e.g. images. Appendix F contains further details and results for all of our experiments.2 2D Synthetic Data. We first consider synthetic distributions where Y = U = R. Our source measure is taken to be the independent product η(y, u) = π Y # ν ⊗ N (0, 1). We plot ground-truth joint distributions and samples for two datasets in Figure 1 . See Appendix F for additional results. Samples from our method (COT-FM) closely match those from the ground-truth distribution, whereas samples from PCP-Map and COT-Flow [Wang et al., 2023] can produce samples in regions of zero support under the ground-truth distribution. In Table 1 , we provide a quantitative analysis, where we measure the W 2 and MMD distances between the generated and ground-truth joint distributions. This is motivated by Proposition 1, as triangular maps which couple the joint distributions necessarily couple the conditional distributions. Our method outperforms the baselines across all metrics.Lotka-Volterra (LV) Dynamical System. Here we estimate parameters of the Lotka-Volterra (LV) model given only noisy observations of its solution. The LV model has parameters u = (α, β, γ, δ) ∈ R 4≥0 and a pair of coupled nonlinear ODEs of the form whose solution p(t) = (p 1 (t), p 2 (t)) ∈ R 2 ≥0 represents the number of prey and predator species at time t ∈ [0, T ]. Following Alfonso et al. [2023] , we assume p(0) = (30, 1) and that log(u) ∼ N (m, 0.5I) with m = (-0.125, -3, -0.125, -3). Given parameters u ∈ R 4 ≥0 , we simulate Equation ( 13) for t ∈ {0, 2, . . . , 20} to obtain a solution z(u) ∈ R 22EQUATION≥0 . An observation y ∈ R 22 ≥0 is obtained by the addition of log-normal noise, i.e. log(y) ∼ N (log(z(u), 0.1I). We thus may simulate many (y, u) pairs from the target measure for training.As a benchmark, we follow the settings of Alfonso et al. [2023] and choose parameters u = (0.83, 0.041, 1.08, 0.04) to generate a single observation y as described above. In Figure 2 , we plot a histogram of 10, 000 samples from the posterior ν(u | y) of COT-FM.Since the ground-truth posterior is intractable, we compare against differential evolution Metropolis MCMC [Braak, 2006] . Samples from our method qualitatively resemble those from MCMC, and the posterior mode is typically close to the true unknown u (shown in red). Our method is quantitatively closest to the MCMC samples in the W 2 metric, and competitive in the MMD metric (Table 2 ).Darcy Flow Inverse Problem. Here we consider an infinitedimensional Bayesian inverse problem from the 2D Darcy flow PDE. The setting is adapted from Hosseini et al. [2023] . We opt to compare against WaMGAN [Hosseini et al., 2023] , as this is currently the only other extant amortized function-space COT method, and FFM [Kerrigan et al., 2023] as a function-space flow matching ablation. The Darcy flow PDE is an elliptic equation on a smooth domain Ω ⊆ R d which relates a permeability field exp(u), a pressure field ρ, and a source term f via -div exp(u)∇ρ = f on Ω subject to ρ = 0 on ∂Ω. Our goal is to recover the permeability u from noisy measurements y of the pressure ρ. Both the unknown u and observations y are functions and thus infinite-dimensional. To define our target measure, we specify a prior ν(u) = N (0, C) with a Matérn kernel C of lengthscale ℓ = 1/2 and ν = 3/2. Given u ∼ η(u), the Darcy flow PDE is solved numerically [Alnaes et al., 2015] to obtain a solution ρ(u) observed at some finite but arbitrary number of points {x 1 , . . . , x n } ⊂ R 2 . An observation y(u) is obtained by adding Gaussian noise to each observation, i.e. y(u) ∼ N (ρ(u), σ 2 I) where σ = 2.5 × 10 -2 . We implement all models via a Fourier Neural Operator [Li et al., 2020] , allowing us to work with arbitrary discretizations, as required by the functional nature of this problem.We provide an illustration in Figure 3 . As the true posterior is intractable, we compare against preconditioned Crank-Nicolson (pCN) [Cotter et al., 2013] , a function-space MCMC method. In Figure 3 , we plot the mean posteriors obtained from the various methods. Qualitatively, both COT-FM and FFM are good approximations to pCN, while WaMGAN has visual artifacts. However, the MSE 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 a quantitative comparison between the methods on a test set of 5,000 samples, where we measure MSE and CRPS [Hersbach, 2000] . We compare the ensemble mean of 10 samples against the true u value as running pCN for each observation is prohibitively expensive. COT-FM outperforms FFM and WaMGAN in terms of MSE and is on-par with FFM in terms of CRPS. See Appendix F for further details. 
conclusion: 
 We analyze conditional optimal transport from a geometric and dynamical point of view. Our analysis culminates in the characterization of absolutely continuous curves of measures in a conditional Wasserstein space, resulting in a conditional analog of the Benamou-Brenier Theorem.We use these result to build on the framework of triangular transport and flow matching to develop simulation-free methods for conditional generative models. Our methods are applicable across a wide class of problems, and we demonstrate our methodology on several challenging inverse problems.Limitations and Broader Impacts. A limitation COT-FM is that computing the full COT plan can be expensive for large datasets, necessitating the use of minibatch approximations potentially resulting in sub-optimal plans. While this approximation does not limit the practical applicability of our method, an interesting challenge is to characterize the precise relationship between this minibatch approximation and the full COT plan. Moreover, computing the COT plan incurs a small additional computational cost compared to standard flow matching. As with all generative models, a potential negative impact is the potential for disinformation through generated samples being purported as real.Appendix:  
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We believe that our paper does not pose any such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Any assets used are properly cited throughout the paper. We provide addiitonal details in Appendix F.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: Our paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: Our paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We provide here a brief and informal overview of optimal transport in the standard unconditional setting. For more details, we refer to the standard references of Villani et al. [2009] , Santambrogio [2015] , Ambrosio et al. [2005] . Let X be a separable metric space and fix a cost function c : X × X → R ∪ {+∞}. Suppose we have two Borel measures η, ν ∈ P(X). The Monge problem seeks to find a measurable transport map T : X → X minimizing the expected cost of transport, i.e. corresponding to the optimization problemEQUATION)This optimization problem is challenging, though, as it involves a nonlinear constraint and the set of feasible maps may be empty. In contrast, the Kantorovich problem is a relaxation which seeks to find an optimal coupling γ ∈ Π(η, ν), i.e. a probability distribution over X × X with marginals η, ν, which solvesEQUATION)Under fairly weak conditions (e.g., the cost is lower semicontinuous and bounded from below [Ambrosio et al., 2013, Theorem 2.5] ), minimizers to the Kantorovich problem are guaranteed to exist. If the cost function is c(x 0 , x 1 ) = |x 0 -x 1 | p for some 1 < p < ∞, under sufficient regularity conditions on η a solution T ⋆ to the Monge problem is guaranteed to exist and, moreover, the coupling γ ⋆ = (I, T ⋆ ) # η is optimal for the Kantorovich problem. See Ambrosio et al. [2013, Chapter 2] and Ambrosio et al. [2005, Theorem 6.2.10 ].Wasserstein Space. In the special case that c(x 0 , x 1 ) = |x 0 -x 1 | p for 1 ≤ p < ∞, and η, ν ∈ P p (X), the Kantorovich problem admits a finite-cost solution. The cost of such an optimal coupling is the p-Wasserstein distanceW p p (η, ν) = min γ X×X |x 0 -x 1 | p dγ(x 0 , x 1 ) | γ ∈ Π(η, ν) (16)which, as the name suggests, is a metric on the space P p (X) [Ambrosio et al., 2005, Section 7 .1] [Santambrogio, 2015, Section 5.1] . The Wasserstein distance admits a dynamical formulation via the Benamou-Brenier theorem [Benamou and Brenier, 2000] . Namely, the p-Wasserstein distance can be obtained by finding a time-dependent vector field transforming η to ν across time t ∈ [0, 1] with minimal energy:W p (η, ν) = min (γt,vt) 1 0 X |v t (x)| p dγ t (x) dt | γ 0 = η, γ 1 = ν, ∂ t γ t + div(v t γ t ) = 0 . (17)Here, we constrain our minimization problem over the set of measures and vector fields (γ t , v t ) interpolating between η and ν, satisfying a continuity equation (see Section 5). In Section 4, we study a generalization of the Wasserstein distances for conditional optimal transport problems. In particular, Theorem 5 provides a generalization of the Benamou-Brenier theorem to the conditional setting which recovers a conditional Wasserstein distance.This section contains additional discussion regarding the static COT problem, supplementing Section 3.1. We refer to Hosseini et al. The conditional Monge problem also admits a relaxation under which one only considers couplings whose Y -components are almost surely equal. To that end, we consider the subset C ⊂ (Y × U ) 2 whose Y components are identical, i.e.,C := (y 0 , u 0 , y 1 , u 1 ) ∈ (Y × U ) 2 | y 0 = y 1 (19)and we define the set of (Y )-restricted probability measures R Y ⊂ P (Y × U ) 2 such that every γ ∈ R Y is concentrated on C . In other words, if γ ∈ R Y , then samples (y 0 , u 0 , y 1 , u 1 ) ∼ γ have y 0 = y 1 almost surely. In addition, for any η, ν ∈ P(Y ×U ), we define the set of triangular couplings Π Y (η, ν) to be the probability measures in R Y whose marginals are η and ν, i.e.EQUATION)The conditional Kantorovich problem seeks a triangular couplingEQUATIONHosseini et al. [2023] prove the existence of minimizers to the conditional Kantorovich and Monge problems under very general assumptions. Moreover, optimal couplings to the conditional Kantorovich problem induce optimal couplings for µ-almost every conditional measure. Assuming sufficient regularity assumptions on the conditional measures, unique solutions to the conditional Monge problem exist. We restate these results here for the sake of completeness.Proposition 3 (Prop 3.3 [Hosseini et al., 2023] ) Fix η, ν ∈ P µ (Y × U ). Suppose the cost function c is continuous, inf c > -∞, and there exists a finite cost coupling γ ∈ Π Y (η, ν). Then, the conditional Kantorovich problem admits a minimizer γ ⋆ . Moreover, γ ⋆,y0 (y 1 , u 0 , u 1 ) = γ⋆,y0 (u 0 , u 1 )δ(y 1 -y 0 ) where for µ-almost every y the measure γ ⋆,y is an optimal coupling for η y , ν y under the cost c y (u 0 , u 1 ) = c(y, u 0 , y, u 1 ) Proposition 4 (Prop 3.8 [Hosseini et al., 2023] ) Fix 1 < p < ∞ and η, ν ∈ P µ p (Y × U ). Suppose c(y 0 , u 0 , y 1 , u 1 ) = |u 0 -u 1 | p . If η y assign zero measure to Gaussian null sets for µ-almost every y, then there is a unique solution T ⋆ to the conditional Monge problem, and γ ⋆ = (I, T ⋆ ) # η is the unique solution to the conditional Kantorovich problem. If ν y also assign zero measure to Gaussian null sets for µ-almost every y, then T ⋆ is injective η-almost everywhere.In this section, we provide additional details and results regarding the closed-form conditional Wasserstein distance for Gaussian distributions. See Section 4 in the main paper.Suppose Y = R d and U = R d ′are Euclidean spaces (of possibly different dimensions), and that η, ν ∈ P µ p (Y × U ) are Gaussians of the formη = N m m η u , Σ Λ η Λ η T Σ η u ν = N m m ν u , Σ Λ ν Λ ν T Σ ν u (22) where m ∈ R d , m η u , m ν u ∈ R d ′, and the (block) covariance matrices areΣ ∈ R d×d , Λ η , Λ ν ∈ R d×d ′ , and Σ η u , Σ ν u ∈ R d ′ × d ′ .This form is chosen to ensure that η and ν have equalEQUATION)We have that the conditionals η y , ν y are available in closed-form:EQUATION)Thus, for any fixed y, we use the known closed-form unconditional Wasserstein distance to obtainEQUATION)We now take an expectation overEQUATION)This form, perhaps unsurprisingly, closely resembles the unconditional Wasserstein distance between two Gaussians, except for the presence of an additional Tr(RΣR T ) term. Note that when η, ν have uncorrelated Y, U components, we precisely recoverW 2 2 (π U # η, π U # ν) as one may expect. As a special case of interest, if Y = U = R and η = N (0, I) ν = N 0, 1 ρ ρ 1 |ρ| < 1 (29)then we obtain as a special case of Equation ( 26) that W µ,2 2 (η, ν) = 2(1 -1 -ρ 2 ). This is zero if and only if ρ = 0, i.e. η = ν.In this section, we provide detailed proofs of our claims in Section 4, regarding the metric properties of the conditional Wasserstein space.We first note that W µ p (η, ν) may be viewed as the minimal value of the constrained Kantorovich problem in Equation (3) when one takes the cost to be the metric on the space Y × U . Similar results, relating the conditional Wasserstein distance to triangular couplings, have appeared previously, but our proof is independent of these prior works [Chemseddine et al., 2023 , Gigli, 2008] . Proposition 5 (Equivalent Formulation of the Conditional Wasserstein Distance) Fix η, ν ∈ P µ p (Y × U ) and 1 ≤ p < ∞. Then, W µ p (η, ν) is well-defined, finite, andEQUATIONwhere W µ,p p (η, ν) represents the p-th power of the conditional p-Wasserstein distance.Proof. The cost function d p is clearly continuous and non-negative, and hence by Proposition 3 it suffices to exhibit a finite-cost coupling γ ∈ Π Y (η, ν) between η and ν. Indeed, take the conditionally independent couplingEQUATIONwhich is clearly in Π Y (η, ν). We then have that(Y ×U ) 2d p (y 0 , u 0 , y 1 , u 1 ) dγ(y 0 , u 0 , y 1 , u 1 ) =(Y ×U ) 2 ∥(y 0 , u 0 ) -(y 1 , u 1 )∥ p Y ×U dγ(y 0 , u 0 , y 1 , u 1 ) ≤ 2 p (Y ×U ) 2 ∥(y 0 , u 0 )∥ p Y ×U + ∥(y 1 , u 1 )∥ p Y ×U dγ(y 0 , u 0 , y 1 , u 1 ) = 2 p Y ×U ∥(y 0 , u 0 )∥ p Y ×U dη(y 0 , u 0 ) + Y ×U ∥(y 1 , u 1 )∥ p Y ×U dν(y 1 , u 1 ) < +∞.Hence, Equation (30) admits a minimizer γ ⋆ ∈ Π Y (η, ν). By Proposition 3, this minimizer may be taken to have the form γ ⋆ = γ ⋆,y1 (u 0 , u 1 )δ(y 1 -y 0 )µ(y 1 ) where γ ⋆,y1 (u 0 , u 1 ) is µ(y 1 )-almost surely an optimal coupling between η y1 , ν y1 for the cost |u 1 -u 0 | p . Thus,EQUATIONHere, we emphasize that the µ-almost sure uniqueness of the disintegrations of η, ν along Y result in a well-defined expression.Moreover, if η ∈ P µ p (Y × U ) it follows that η y ∈ P p (U ) for µ-a.e. y, becauseEQUATION)Thus all considered p-Wasserstein distances on U are finite.We now proceed to prove several metric properties of our distance.Proposition 2 (Some Properties ofW µ p ) Let 1 ≤ p < ∞.(a) W µ p is well-defined, finite, and equals the minimal conditional Kantorovich cost. (b) W µ p is a metric on the space P µ p (Y × U ). (c) There does not exist C > 0 such that W µ p (η, ν) ≤ CW p (η, ν) for all η, ν ∈ P µ p (Y × U ).(d) For all η, ν ∈ P µ p (Y × U ), W p π U # η, π U # ν ≤ W µ p (η, ν) and W p (η, ν) ≤ W µ p (η, ν).Proof. Part (a). This is simply a restatement of Proposition 5.Part (b). Fix η, ν, ρ ∈ P µ p (Y × U ). Since W p is a metric on P p (U ), we immediately obtain the symmetry of W µ p . Moreover, we have that W µ p (η, ν) = 0 if and only if η y = ν y for µ-almost every y. Thus, ifEQUATION)which shows that η = ν. Here, E y = {u | (y, u) ∈ E} is the y-slice of E. Conversely, if η = ν, then η y = ν y up to a µ-null set by the essential uniqueness of disintegrations. Thus, W µ p (η, ν) = 0 if and only if η = ν. By Minkowski's inequality and the triangle inequality for W p on P p (U ), we seeEQUATIONPart (c). We provide a counterexample. Fix any u 0 ̸ = 0 ∈ U and y 0 , y 1 ∈ Y such that y 0 ̸ = y 1 . Define µ = 1 2 (δ y0 + δ y1 ). Set u k = (k + 1)u 0 for k = 1, 2, . . . and for each k, define two measures on Y × U byEQUATION)It is clear thatEQUATIONMoreover, as k → ∞ we have W µ p (µ k , ν k ) → ∞ but W p p (ν k , η k ) remains bounded. See Figure 4 . Part (d). First, the unconditional distance W p (η, ν) may be obtained via an unrestricted coupling in Π(η, ν), i.e. the set of all joint measures on Y × U having marginals η, ν. Since Π(η, ν) ⊇ Π Y (η, ν), by part (a) we see that W p (η, ν) ≤ W µ p (η, ν). Let γ ⋆ (y 0 , u 0 , y 1 , u 1 ) = γ ⋆,y1 (u 0 , u 1 )δ(y 1 -y 0 )µ(y 1 ) be an optimal γ ⋆ ∈ Π Y (η, ν). We claim that γ(u 0 , u 1 ) := Y γ ⋆,y (u 0 , u 1 ) dµ(y) couples π U # η and π U # ν. Let π 0 : (u 0 , u 1 ) → u 0 be the projection onto the first coordinate of U × U . Observe that for µ-almost every y, we have that γ ⋆,y ∈ Π(η y , ν y ) is optimal, and, in particular, π 0 # γ ⋆,y = η y . Fix an arbitrary φ ∈ C b (U ). We then have ThusEQUATIONπ U # γ = π U # η.A similar argument shows that for the map π 1 : (u 0 , u 1 ) → u 1 we haveπ 1 # γ = π U # ν, so that γ ∈ Π(π U # η, π U # ν). Now, as γ ⋆,y1 (u 0 , u 1 ) ∈ Π(η y1 , ν y1) is µ-almost surely optimal in the usual Wasserstein sense,W p,µ p (η, ν) = Y U 2 |u 0 -u 1 | p dγ ⋆,y (u 0 , u 1 ) dµ(y) (46) = U 2 |u 0 -u 1 | p dγ(u 0 , u 1 ) (47) ≥ W p p (π U # η, π U # ν) (48) since γ ∈ Π(π U # η, π U # ν) is a coupling but potentially sub-optimal.We now study the geodesics in the space P µ p (Y × U ). Theorem 1 (P µ p (Y × U ) is a Geodesic Space) For any η, ν ∈ P µ p (Y × U ), there exists a constant speed geodesic between η and ν.EQUATIONLet γ ⋆ ∈ Π Y (η, ν) be an optimal restricted coupling, and consider the path of measures in P p (Y × U ) given byEQUATION)Step one: We check that for each 0 ≤ t ≤ 1, we have γ t ∈ P µ p (Y × U ). That is, we need to check that for all Borel A ⊆ Y , we have γ t (A × U ) = µ(A). Indeed, recall that restricted measures are concentrated on the set C (see Equation ( 19)). Thus,γ t (A × U ) = γ ⋆ λ -1 t (A × U ) = γ ⋆ {(y, u 0 , y, u 1 ) | y ∈ A} = π 1 # γ ⋆ (A) = (π 1 • π 1,2 ) # γ ⋆ (A) = π 1 # η(A) = µ(A)i.e. γ t (A × Y ) = µ(A) as claimed.Step two: We show thatEQUATIONAn analogous calculation shows that π 3,4 # γ s t = γ s , so that γ s t ∈ Π(γ t , γ s ). We now check thatγ s t ∈ R Y (Y × U ). Indeed, suppose E ⊆ Y × U is a Borel set such that E ∩ C = ∅.In other words, for every (y 0 , u 0 , y 1 , u 1 ) ∈ E we have yEQUATIONIndeed, if c = (y, u 0 , y, u 1 ) ∈ C , thenEQUATIONThusγ s t ∈ Π Y (η, ν) as claimed. Now, we have W µ,p p (γ t , γ s ) ≤ (Y ×U ) 2 d p (y 0 , u 0 , y 1 , u 1 ) dλ s t (y 0 , u 0 , y 1 , u 1 ) = (Y ×U ) 2 d p (λ t (y 0 , u 0 , y 1 , u 1 ), λ s (y 0 , u 0 , y 1 , u 1 )) dγ ⋆ (y 0 , u 0 , y 1 , u 1 ) = (Y ×U ) 2 |(t -s)(y 0 -y 1 )| 2 + |(t -s)(u 0 -u 1 )| 2 p/2 dγ ⋆ (y 0 , u 0 , y 1 , u 1 ) = |t -s| p (Y ×U ) 2 d p (y 0 , u 0 , y 1 , u 1 ) dγ ⋆ (y 0 , u 0 , y 1 , u 1 ) = |t -s| p W µ,p p (η, ν).EQUATIONEQUATIONRearranging the previous inequality implies |t -s|W µ p (η, ν) ≤ W µ p (γ s , γ t ) for all s, t ∈ [0, 1], and hence W µ p (γ t , γ s ) = |t -s| p W µ p (η, ν). Theorem 2 (Conditional McCann Interpolants) Fix η, ν ∈ P µ p (Y × U ). Suppose T ⋆ (y, u) = (y, T ⋆ U (y, u)) is an injective triangular map solving the conditional Monge problem (1). Define the maps T t : Y × U → Y × U for 0 ≤ t ≤ 1 via T t = (1 -t)I + tT ⋆ , and define the curve of measures γ t = [T t ] # η ∈ P γ p (Y × U ). Then, (a) (γ t ) is absolutely continuous and a constant speed geodesic between η, ν (b) The vector field v t (T ⋆ t (y, u)) = (0, T ⋆ U (y, u) -u) generates the path γ t , in the sense that (γ t , v t ) solve the continuity equation (9).EQUATIONEQUATIONwhich shows that (γ t , v t ) solve the continuity equation.EQUATIONIn particular, 1 0 ∥v t ∥ L p (γt,Y ×U ) dt < ∞In this section, we provide proofs of all claims made in Section 5.We begin with a lemma that is used in the proof of Theorem 4. Informally, a solution to the continuity equation with a triangular vector field will result in the conditional measures almost surely satisfying the continuity equation as well. Lemma 1 (Triangular Vector Fields Preserve Conditionals) Suppose v t (y, u) = (0, v U t (y, u)) is triangular and that (γ t ) ⊂ P µ p (Y × U ) is a path of measures such that (v t , γ t ) satisfy the continuity equation in the sense of distributions. Then, it follows that for µ-almost every y ∈ Y , we have ∂ t γ y t + ∇ • (v U t (y, -)γ y t ) = 0.Proof. Fix any φ ∈ Cyl(U × I). Suppose ψ ∈ Cyl(Y ) is given, and note that ψ(y)φ(u, t) ∈ Cyl(Y × U × I). As (v t , γ t ) solve the continuity equation, it follows from the triangular structure of v t that upon testing against ψφ we haveEQUATION) Because ψ(y) ∈ Cyl(Y ), it is of the form ρ(π(y)) where π : Y → R k for some k ≥ 1 and ρ ∈ C ∞ c (R k ).Taking ρ to be a sequence of smooth approximations to the indicator function of an arbitrary rectangleEQUATION)As Y is separable, the Borel σ-algebra on Y is generated by the cylinder sets, i.e. those which are precisely of the form π -1 (E) for some finite-dimensional rectangle E. We have thus shown that for an arbitrary Borel measurable set E ⊆ Y ,EQUATION)From this, it follows thatEQUATIONWe now proceed to prove the main results of this section. First, we introduce some preliminary notions. We define the map j q : L q (γ, Y × U ) → L p (γ, Y × U ) for 1/p + 1/q = 1 viaj q (w) = |w| q-2 w w ̸ = 0 0 w = 0 (72)which is the Fréchet differential of the convex functional 1 q ∥w∥ q L q (γ,Y ×U ) . A straightforward calculation shows that this map satisfiesEQUATION)See also Ambrosio et al. [2005, Chapter 8 ]. Proof. Assume without loss of generality that |γ ′ |(t) ∈ L ∞ (I) and that I = (0, 1) [Ambrosio et al., 2005, Lemma 1.1.4, Lemma 8.1.3] . Fix any φ ∈ Cyl(Y × U ). For s, t ∈ I there exists an optimal triangular coupling γ st ∈ Π Y (γ s , γ t ). By Hölder's inequality,EQUATIONIt follows that t → γ t (φ) is absolutely continuous. We can introduce the upper semicontinuous and bounded mapEQUATION)For |h| sufficiently small, choose any optimal coupling γ (s+h)h ∈ Π Y (γ s+h , γ s ) and note thatEQUATION)If t is a point of metric differentiability for t → γ t , note that γ (t+h)t → (I, I) # γ t narrowly, where I is the identity map on Y × U . Moreover, since γ t ∈ P µ p (Y × U ), it follows that on the diagonal we have that almost surely H(y 0 , u 0 , y 0 , u 1 ) = ι(|∇ u φ(y 0 , u 0 )|. Thus,lim sup h→0 |γ t+h (φ) -γ t (φ)| |h| ≤ |γ ′ |(t) Y ×U|H| q (y 0 , u 0 , y 0 , u 0 ) dγ t (y 0 , u 0 )EQUATION)Taking Q = Y × U × I and γ = γ t dt, fix any φ ∈ Cyl(Q). We have thatQ ∂ s φ(y, u, s) dγ(y, u, s) = lim h↓0 I 1 h Y ×U φ(y, u, s) dγ s (y, u) - (Y ×U )φ(y, u, s) dγ s+h (y, u) ds.(80)An application of Fatou's Lemma, Equation ( 78), and Hölder's inequality gives usQ ∂ s φ(y, u, s) dγ(y, u, s) ≤ J |γ ′ |(s) ds 1/p Q |∇ u φ(y, u, s)| q dµ(y, u, s) 1/q (81)for any interval J ⊂ I with supp φ ⊂ J × Y × U .Fix the subspaceV = {ι(∇ u φ(y, u, s)) : φ ∈ Cyl(Q)} ⊆ Y × U (82)and denote by V its L q (γ, Y × U × I) closure. Define the linear functional L :EQUATIONand note that Equation ( 81) implies that L is a 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 problemEQUATIONwhich admits the unique solution w such that j q (w) -L = 0. In particular, the estimate (81) shows that the above functional is coercive and hence admits a minimizer which we may obtain via its differential as a consequence of convexity. Thus, we obtain a triangular vector field v = j q (w) such that for all φ ∈ Cyl(Q),EQUATIONThis precisely shows that (v t , γ t ) is a triangular distributional solution to the continuity equation. Now, choose any interval J ⊂ I and choose a sequenceEQUATION)Taking k → ∞ we see thatEQUATIONand since J ⊂ I was arbitrary, we concludeEQUATION)We now prove, in some sense, a converse of the previous theorem.Theorem 4 (Continuous Curves Generated by Triangular Vector Fields) Suppose that γ t : I → P µ p (Y × U ) is narrowly continuous and (v t ) is a triangular vector field such that (γ t , v t ) solve the continuity equation with ∥v t ∥ L p (γt,Y ×U ) ∈ L 1 (I). Then, γ t :I → P µ p (Y × U ) is absolutely continuous in the W µ p metric and |γ ′ |(t) ≤ ∥v t ∥ L p (µ,Y ×U ) for almost every t.Proof. We first assume that U is finite dimensional. Our strategy is to check the hypotheses necessary for Ambrosio et al. [2005, Theorem 8.3.1] to hold for µ-almost every y, followed by an application of this theorem. By Lemma 1, for µ-almost every y we have that (γ y t , v U t (y, -)) solve the continuity equation distributionally on I × U . By Jensen's inequality (and the assumption p ≥ 1) we seeEQUATION)Since the first term is finite, it follows thatEQUATION)Now Ambrosio et al. [2005, Lemma 8.1.2] shows that for µ-almost every y we have that (γ y t ) admits a narrowly continuous representative (γ y t ) with γy t = γ y t for almost every t. It follows from Ambrosio et al. [2005, Theorem 8.3 .1] that for any t 1 ≤ t 2 in I, we haveEQUATIONwhere the second line follows as γy t = γ y t for almost every t. Let γt = Y γy t dµ(y) be the measure obtained via marginalizing over the Y -variables. Taking an expectation over y ∼ µ, the previous inequality shows us thatEQUATIONNow, note that t 1 is almost surely a Lebesgue point of the right-hand side and γt1 = γ t1 . Taking t 2 → t 1 along a sequence where γt2 = γ t2 shows us thatEQUATIONfor almost every t ∈ I.In the case that U is infinite dimensional, fix any y ∈ Y such that Lemma 1 holds (which is of full measure) and fix a countable orthonormal basis (e k ) for U . Set π d : U → R d to be the projection operator for this basis, i.e. u → (⟨u, e 1 ⟩, . . . , ⟨u, e d ⟩). We consider the collection of finite dimensional conditional measures γ d,y t = π d # γ y t . By the same argument in Ambrosio et al. [2005, Theorem 8.3 .1], there exists a vector field v d,y t on R d such that (γ d,y t , v d,y t ) solve the continuity equation andEQUATIONIt follows from the finite-dimensional case above that for almost every t 1 ≤ t 2 , we haveEQUATIONLet γy,d t = (π d ) ⋆ # γ y,d t where (π d ) ⋆ : R d → U maps z → d k=1 z k e k . As d → ∞ we have γd,y t → γ yt narrowly for all t ∈ I. Since (π d ) ⋆ is an isometry, Ambrosio et al. [2005, Lemma 7.1.4 ] shows thatEQUATIONNow, integration with respect to dµ(y) yieldsEQUATIONTaking t 2 → t 1 shows that for almost every t we haveEQUATIONTogether, Theorem 3 and Theorem 4 give us a dynamical interpretation of the conditional Wasserstein distance. The following result is a conditional analogue of the well-known Benamou-Brenier Theorem [Benamou and Brenier, 2000] . Here, we note that the following proof follows the standard proof closely -the main legwork in obtaining this conditional generalization is through the previous two theorems.Theorem 5 (Conditional Benamou-Brenier) Let 1 < p < ∞. For any η, ν ∈ P µ p (Y × U ), we haveW p,µ p (η, ν) = min (γt,vt) 1 0 ∥v t ∥ p L p (µt) dt | (v t , γ t ) solve (9), γ 0 = η, γ 1 = ν,and v t is triangular .Proof. Write M for the infimum on the right-hand side.First, suppose that (v t , µ t ) are admissible andEQUATIONConversely, by Theorem 1 there exists a constant speed geodesic (γ t ) ⊂ P µ p (Y × U ) connecting η and ν. Recall that constant speed geodesics are absolutely continuous. By Theorem 3, there exists a Borel triangular vector field v t such that (v t , γ t ) solve the continuity equation, and moreover ∥v t ∥ L p (µt,Y ×U ) ≤ |γ ′ |(t). In fact, because (v t , γ t ) solve the continuity equation, Theorem 4 yields that∥v t ∥ L p (µt,Y ×U ) = |γ ′ |(t).Since γ t is a constant speed geodesic in P µ p (Y × U ), it follows that |µ ′ |(t) = W µ p (η, ν) for almost every t ∈ (0, 1). Hence,EQUATIONThus, W µ,p p (η, ν) = M as desired.In this section, we provide additional details regarding all of our experiments, as well as additional results not contained within the main paper. All models can be trained on a single GPU with less than 24 GB of memory, and our experiments were parallelized over 8 such GPUs on a local server. We first describe our setting for the 2D and Lotka-Volterra experiments, as these share a similar setup. Details for the Darcy flow inverse problem are described in the corresponding section.Models. For FM and COT-FM, our model architecture is an MLP with SeLU activations [Klambauer et al., 2017] . Time conditioning is achieved by concatenating the time variable as an input to the network. The covariance operator C chosen in the path of measures in Equation ( 10) is taken to be C = σ 2 I where σ is a hyperparameter.Our implementation of FM is adapted from the torchcfm package Tong et al. [2023] , available under the MIT License. For PCP-Map and COT-Flow, we adapt the open-source implementations from Wang et al. [2023] , available under the MIT License. [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 a random grid search over 100 hyperparameter settings using the grid described in Table 4 . For all model types, we select the best model used to generate the results in the paper as the training checkpoint that resulted in the lowest W 2 error to the joint target distribution on a held-out validation set. For training, we use the Adam optimizer where we only tune the learning rate, leaving all other settings as their defaults in pytorch.Data Generation. This experiment consists of four 2D synthetic datasets, where Y = U = R.The datasets moons, circles, swissroll are available through scikit-learn [Pedregosa et al., 2011] .The moons dataset is generated with noise=0.05 followed by standard scaling with a mean of m = (0.5, 0.25) and standard deviation of σ = (0.75, 0.25). The circles dataset is generated with factor=0.5 and noise=0.05. The swissroll dataset is generated with noise=0.75, followed by projection to the first two coordinates and re-scaling by a factor of 12. All other unstated parameters are left as their default values. We use the code available from Hosseini et al. [2023] to generate the checkerboard dataset. For all datasets, we generate a training set (i.e., samples from the target distribution) of 20, 000 samples and 1, 000 held-out validation samples for model selection. Means and standard deviations in Table 1 are reported across five independent testing sets of 5,000 samples for the best representative of each model type.In COT-FM, to generate samples from the source distribution, we sample an additional 20, 000 points from the target distribution and keep only the Y coordinates. This ensures that the source and target have equal Y marginals. During training, standard Gaussian noise N (0, 1) is sampled for the U coordinate of these source points at each minibatch.We use minibatch COT couplings [Tong et al., 2023] in this experiment as computing the full COT plan was prohibitively expensive in terms of memory usage. However, we note that we use large batch sizes, meaning that the COT plan we find in this way should not be too far from optimal. All couplings are computed using the POT Python package [Flamary et al., 2021] .Data Generation. We adopt the settings of Alfonso et al. [2023] for this experiment. As described in the main paper, we assume p(0) = (30, 1) and that log(u) ∼ N (m, 0.5I) with m = (-0.125, -3, -0.125, -3). Given parameters u ∈ R 4 ≥0 , we simulate Equation (13) for t ∈ {0, 2, . . . , 20} to obtain a solution z(u) ∈ R 22≥0 . An observation y ∈ R 22 ≥0 is obtained by the addition of log-normal noise, i.e. log(y) ∼ N (log(z(u), 0.1I). We thus may simulate many (y, u) pairs from the target measure for training.We generate a training set of 10, 000 (y, u) pairs using the procedure described above and a held-out validation set of 10, 000 (y, u) pairs for model selection. Means and standard deviations in Table 2 are reported across five independent testing sets of 5,000 samples for the best representative of each model type. Figure 2 and Figures 9, 8 , 7, 10 show 10, 000 samples from each model, as well as 10, 000 samples from the differential evolution Metropolis MCMC sampler [Braak, 2006] A common failure mode for the baselines is to generate samples from regions with zero support under the true data distributions. Table 1 contains a 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 set used in this experiment. The COT couplings are computed in the same way as the previous section, and as described in Section 7. Figure 7 : KDE plots of the samples on the Lotka-Volterra system, using the settings described in Section 7. Plots include one-dimensional KDEs on the diagonal, as well as 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 in the diagonal plots, or the black x in the off-diagonal plots. For our experiments we observe y on a 100 × 100 grid, and we use σ = 2.5 × 10 -2 . We note that this level of noise is quite considerable, as it accounts for roughly 60% of the variability in the y. Figure 11 showcases a data point for reference.Our source and target training sets contain 1 × 104 samples each, and our test set comprises 5 × 103 samples. We remark that although y and u are observed on a grid their resolution does not need to be fixed, allowing for training at different resolutions. Models. In order to make learning feasible in infinite-dimensional Hilbert spaces, we adapt the architecture of a Fourier Neural Operator (FNO) [Li et al., 2020] from the neuraloperator package [Kovachki et al., 2021] to accommodate for conditioning information observed at an arbitrary resolution. We do so by introducing a projection layer mapping the conditioning information to match the hidden channels of the input lifting block, and a pooling operation to project to the input dimensions. The two are then concatenated and passed through an FNOBlock mapping from (2 × hidden_channels) × input_dim to hidden_channels × input_dim, before following the original architecture. For all of the models in consideration, we fix the architecture to be have hidden_channels = 64, projection_channels = 256, and 32 Fourier modes. We train each model for 1500 epochs, and hyperparameters for each architecture are selected as follows:• WaMGAN [Hosseini et al., 2023] : using an adaptation to the FNO architecture of the original code 3 , we perform a grid search as detailed in Table 5 . We found the training procedure to be rather unstable, and for this reason we checkpoint the model every 100 epochs and report the results for the best performing model at its best checkpoint. We found this to be a model with learning rate 1 × 10 -4 , 2 full critic iterations, and monotone penalty of 1 × 10 -3 . The gradient penalty parameter did not seem to significantly affect performance on the test set, and was set to 5. • FFM [Kerrigan et al., 2024] : the learning rate is fixed to 5 × 10 -4 , and the covariance operator C is set to match that of the prior, but rescaled by a factor of σ = 1 × 10 -3 . We use the code from the original repository 4 . • COT-FFM: we set ϵ = 1 × 10 -5 in the cost function used to build the COT plan. The learning rate and C are chosen to be the same as FFM. In order to build COT couplings, we take the source measure to be the product measure π Y # η × N (0, C). Approximate couplings are obtained on minibatches of size 256.It should be noted that in any scenario where the source and the target U -marginals are identical, using the OT coupling would yield the identity mapping as the optimal vector field minimizing (12). Hence, the OT-CFM model [Tong et al., 2023] is inapplicable here. Learning rate {1 × 10 -3 , 5 × 10 -4 , 1 × 10 -4 } Full critic iter.{2, 5, 10} Monotone penalty {1 × 10 -3 , 5 × 10 -2 , 1 × 10 -1 } Gradient penalty {1, 5, 10}Sampling. The resulting amortized sampler, denoted for simplification by the mapping (y, u 0 ) → u 1 = TU (y, u 0 ), will parameterize an approximate posterior measure. Notice that, in contrast to classical variational inference techniques, no distributional assumptions are made about the approximate posterior. In turn, integrals are obtained numerically by Monte Carlo sampling K samples from the prior, resulting in the approximationEQUATIONIn this section, we perform an additional experiment demonstrating that our COT-FM is able to obtain good approximations of the true COT map even when trained on relatively small minibatches.We work on Y = U = R and use a standard Gaussian source and a Gaussian target distribution with covariance ρ = 0.75. These are chosen so that we may evaluate, in closed-form, the true conditional 2-Wasserstein distance, as detailed in Section C. We then train our COT-FM method using various batch sizes, and measure the resulting model's conditional 2-Wasserstein distance by sampling 10, 000 points from the source distribution, flowing each source point along the model's learned vector field, and computing the resulting squared distance to the corresponding terminal point.In Figure 12 , we plot the resulting deviation from the true value of W µ,2 2 as a function of batch size. We see that even at a relatively small batch size of 16, the resulting error in the (squared) distances is less than 10% of the true value (≈ 0.678). While this experiment is only feasible on synthetic data (as we must compute the true distance), it nonetheless demonstrates that even with small batches we may recover the true distance. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.38th Conference on Neural Information Processing Systems (NeurIPS 2024).Previous work has referred to this as the conditional flow matching loss[Tong et al.,2023], which is not to be confused with the notion of conditioning that we focus on in this work.Code for all of our experiments is available at https://github.com/GavinKerrigan/cot_fmhttps://github.com/TADSGroup/ConditionalOT2023https://github.com/GavinKerrigan/functional_flow_matching 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - Only applied to relatively low dimensional problems.
- It might be useful to clearly highlight the power and generalization of this work over a “simple” conditional optimal transport formulation which simply conditions on a single-class variable. e.g. https://github.com/atong01/conditional-flow-matching/blob/main/examples/images/conditional_mnist.ipynb
    
    While it is clear from a deep enough reading I suggest the authors might want to highlight these differences for potentially broader appeal. This might be done through an algorithm box or other presentation.
questions: questions: It would be good to know empirically how far from optimal the learned transport maps are, as this is a known limitation of miqibatch-based approaches. 

Comment: I would suggest citing work on rectified flows as a concurrent invention of flow matching. 
https://arxiv.org/abs/2209.03003
limitations: limitations: Adequately discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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].

- 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.

- 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.

[1] Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching (arxiv, v1 released in March, 2024)
questions: 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: limitations: The limitation is discussed in Weakness section.
paper_summary: paper_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.
strength_weakness: strength_weakness: N/A :)
questions: questions: I am using this space for comments and suggestions as well as questions.

- 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
- 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.
- 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)
- Is equation (9) not due to the original flow matching papers?
- For equation (6): I have never heard anyone say the equation should be "understood distributionally". Maybe consider "in the sense of distributions" 
- Stylistic comment: Maybe omit "unconditional" from the title of Appendix A? This is not really used  

@article{carlier2010knothe,
  title={From Knothe's transport to Brenier's map and a continuation method for optimal transport},
  author={Carlier, Guillaume and Galichon, Alfred and Santambrogio, Filippo},
  journal={SIAM Journal on Mathematical Analysis},
  volume={41},
  number={6},
  pages={2554--2576},
  year={2010},
  publisher={SIAM}
}

@article{liu2022flow,
  title={Flow straight and fast: Learning to generate and transfer data with rectified flow},
  author={Liu, Xingchao and Gong, Chengyue and Liu, Qiang},
  journal={arXiv preprint arXiv:2209.03003},
  year={2022}
}
limitations: limitations: N/A
paper_summary: paper_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$.
strength_weakness: strength_weakness: - Whether the COT-FM can recover the dynamic optimal transport requires further clarification.
- Please see the Questions Section below.
questions: 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?
- 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?
- 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? 
- 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: limitations: -	The authors addressed the potential negative societal impact of their work.

Comments: title: Comparison to previous work
comment: I thank the authors for their clarifications and additional experiments. I still believe this makes a potentially useful theoretical contribution but with limited empirical evaluation and therefore maintain my score.
I appreciate the author for the clarification. I agree that the main contribution of this work is the development of the dynamical conditional optimal transport theory, and it is quite novel. I also agree that the COT map can be obtained as the mini-batch size approaches infinity (as discussed in [Tong, 2023]). The additional experiments, which demonstrate that it is possible to approximate the actual COT with large batch sizes, enhance the soundness of the approach. Although there are some aspects of the methodology that might be open to discussion, considering the theoretical contributions and the additional experiments, I would like to raise my score to 5.
I appreciate the author for their clarifications. However, I still believe that additional quantitative evaluation as a conditional generative model would provide more solid support for this work. Hence, I will maintain my current score.
Dear reviewer hqSP,

Thank you for your work evaluating this submission.

The internal policy ([link](https://neurips.cc/Conferences/2024/PaperInformation/NeurIPS-FAQ)) states that work appearing two months or less before the deadline - which is the case of the work you refer to - is considered concurrent, and that authors should not be expected to compare to such work.

If part of your review was based on this, it will be possible to update it during the discussion phase.