Dynamic Conditional Optimal Transport through Simulation-Free Flows

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

Strengths

- This paper presents theoretical analysis of the Conditional Wasserstein Space, such as the characterization of the absolutely continuous path and the conditional generalization of the Benamou-Brenier Theorem. - This paper proposes a Flow Matching model for conditional generation. - This paper is easy to follow.

Weaknesses

- Whether the COT-FM can recover the dynamic optimal transport requires further clarification. - Please see the Questions Section below.

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?

Rating

4

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

- The authors addressed the potential negative societal impact of their work.

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.

Strengths

This paper has many strengths! It is well-written, fits nicely in the conference format without omitting many details, and has appropriate experiments. The proposed methodology is also quite elegant, and circumvents many issues other methods face.

Weaknesses

N/A :)

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

Rating

8

Confidence

5

Soundness

4

Presentation

4

Contribution

3

Limitations

N/A

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.

Strengths

- The paper successfully extends the Benamou-Brenier theorem to the context of dynamical conditional optimal transport. - The paper is easy to follow.

Weaknesses

- The paper appears to be extremely similar to [1] both in theoretically and empirically. Particularly, Theorem 18, which discusses a Benamou-Brenier-like formula and it applies COT with flow matching as like this work. I would like to ask authors to discuss the difference with result in [1]. - 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

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

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

The limitation is discussed in Weakness section.

Summary

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

Strengths

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

Weaknesses

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

Questions

It would be good to know empirically how far from optimal the learned transport maps are, as this is a known limitation of miqibatch-based approaches. Comment: I would suggest citing work on rectified flows as a concurrent invention of flow matching. https://arxiv.org/abs/2209.03003

Rating

6

Confidence

2

Soundness

4

Presentation

3

Contribution

3

Limitations

Adequately discussed.