We are grateful to the reviewer for their constructive feedback and suggestions.

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?

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

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

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

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?

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

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

We agree that this point could have been explained more clearly, and we will update our paper with a complete, formal proof of this result.

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

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?

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

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

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

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

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?

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

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

1. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. Tong et al., 2023.
2. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. Pooladian et al., 2023.
3. Conditional Optimal Transport on Function Spaces. Hosseini et al., 2023.

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.

Is there a clear way to choose the $\epsilon$ parameter for the COT Flow Matching? Any heuristics whatsoever?

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

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

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

We will include a discussion of these notions in an updated camera-ready version of our paper.

Is equation (9) not due to the original flow matching papers?

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

1. Flow Matching for Generative Modeling. Lipman et al., 2023.
2. Building Normalizing Flows with Stochastic Interpolants. Albergo et al., 2023(A).
3. Stochastic Interpolants with Data-Dependent Couplings. Albergo et al., 2023(B).
4. Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport. Tong et al., 2023.
5. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. Pooladian et al., 2023.

We thank the reviewer for their feedback!

The paper appears to be extremely similar to [1] [...] I would like to ask authors to discuss the difference with result in [1].

We first would like to remind the reviewer of the NeurIPS concurrent work policy. 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.

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

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

The paper discusses dynamical COT in Sections 4-5. [...] the experiments seem to address conditional generation rather than conditional optimal transport.

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

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

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.

Your understanding is correct -- however, we respectfully disagree that our method has limited applicability.

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

Second, 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].

It is unclear if the experimental settings are appropriate

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

the experiments were conducted on very small datasets.

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

Does it mean that FM model learn geodesics?

Yes -- the path of distributions modeled by COT-FM is a geodesic in the conditional Wasserstein space, in the sense of Theorem 1.

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

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

1. Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching. Chemseddine et al., 2024
2. Improving and generalizing flow-based generative models with minibatch optimal transport. Tong et al., 2023.
3. Multisample Flow Matching: Straightening Flows with Minibatch Couplings. Pooladian et al., 2023.
4. From Knothe's transport to Brenier's map and a continuation method for optimal transport. Carlier et al., 2010

Dear reviewer hqSP,

Thank you for your work evaluating this submission.

The internal policy (link) 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.

We thank the reviewer for their valuable feedback!

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.

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

Only applied to relatively low dimensional problems.

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

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

These results suggest that the method should scale to higher dimensional problems, which is an interesting avenue for future work.

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.

Thank you for this suggestion -- we agree that this could be more clearly explained. We elaborate here on the key differences.

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

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

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

I would suggest citing work on rectified flows as a concurrent invention of flow matching.

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