Summary
This work proposes Optimal Flow Matching (OFM), which restricts standard Flow Matching to specific vector fields that yield straight trajectories by design. This is implemented by considering vector fields such that a convex function exists whose gradient pushes the initial to the final point.
The authors propose to solve this optimization problem with training that involves solving a convex optimization problem in each iteration to obtain the inverse of the flow map. Then, it is shown that this procedure has the same minimizer as the one minimizing the dual OT loss. The authors empirically validate OFM quantitatively on the Wasserstein-2 benchmark and qualitatively on FFHQ.
Strengths
- The authors derive a new training procedure based on Flow Matching restricted to fully straight vector fields implemented through the gradient of a convex function.
- The authors provide a new insight into the connection between minimizing an OT loss and Flow Matching. Specifically, they show that the minimizer of these two losses is equivalent when restricted to straight vector fields.
Weaknesses
Missing comparison and contextualization of related work:
- In [1], it was proposed to approximate the computation of the conjugate with amortized optimization and solve this with Adam or L-BFGS. The proposed *SubOpt* in this work (solved with Adam/L-BFGS) is equivalent to the convex conjugate optimization when $t=1$. A discussion between these two optimization schemes should be included. Specifically, how does OFM compare to [1] when only sampling $t=1$? What is gained by considering $t\sim \mathcal{U}[0,1]$? Additionally, the Wasserstein-2 benchmark results from [1] are not considered in the paper. This is the most related method and outperforms the method proposed in the paper on all dimensions. These results should be included as well as a comparative discussion.
Experiments:
- No quantitative results are provided for the unpaired image-to-image translation experiments. It is hard to generalize the results over a few qualitative results. Thus, including some commonly used quantitative metrics like FID or the ones used in the Wasserstein-2 benchmark on CelebA64, as done in [1], would be very beneficial.
Minor Weaknesses:
- Inconsistency on the optimizer used to solve *SubOpt*. In Appendix C, it is stated that "in all our experiments as the subOpt optimizer we use LBFGS"[line 557] while in Appendix B, it says "(with Adam optimizer)"[line 534]. Which one is actually used? How do Adam and LBFGS compare?
- Some of the wording can be improved, e.g. "... via fancy formula (18) is not trivial and requires tricky integration techniques."[line 194]
[1] Brandon Amos. "On amortizing convex conjugates for optimal transport". In ICLR 2023.
Questions
- The goal of OFM is to learn a time-independent one-step model that translates $x_0$ to $x_1$. Then, during inference, no ODE is solved. As mentioned above, some more motivation could improve clarity here. Specifically, explaining why it is beneficial to formulate this in the FM framework considering $t\sim \mathcal{U}[0,1]$. Would the proposed method also work when only sampling $t=1$? The authors only mention that the subproblem (20) is $(1-t)$-strongly convex, which gives a benefit compared to $t=1$. Are there other benefits of the proposed formulation?
- In Figure 2, it is shown that OFM learns the same solution for different choices of couplings. In [2], it was shown that only the OT coupling is preserved when using standard Flow Matching. What do these illustrative examples look like for standard FM training? I assume they would learn different trajectories, but do the learned endpoints differ?
- In this work, only ICNNs are considered. As shown in previous works, even when wanting to learn a convex function, MLPs can also work well, i.e., in [1], they work better compared to ICNNs. Does OFM also work with standard MLPs/CNNs/etc? How is the comparative performance?
[1] Brandon Amos. "On amortizing convex conjugates for optimal transport". In ICLR 2023.
[2] Valentin De Bortoli and Guan-Horng Liu and Tianrong Chen and Evangelos A. Theodorou and Weilie Nie. "Augmented Bridge Matching". In Arxiv 2023.
Limitations
- The limitations of the flow map inversion, the Hessian computation, and ICNNs are laid out in the paper, and details for the experiments' computational time are provided.
- While these computational times give some insight, detailed information on time and space complexity with respect to the input dimension would be even more beneficial. Additionally, comparing the runtime of OFM and standard FM training would provide a more illustrative comparison of raw computational time.