Optimal Flow Matching: Learning Straight Trajectories in Just One Step

Dear reviewer, thank you for your feedback and kind words. Please find below the answers to your questions.

(1) Convergence time. Plots of convergence.

Following your request, we provide the plots of convergence (in $L^2$-UVP) depending on training time in the experiment on the W2 benchmark (Section 4.2) in dimension $D=32$. Please see Figure 3 in the PDF file. OFM iteration takes more time than iteration of OT-CFM, RF or $c-RF$ during the first FM round. After the first FM round, RF and $c-$RF slow down considerably and lose to OFM. Fortunately, OFM requires much fewer iterations to achieve the best metrics and outperforms all other methods in the first hours of the training.

We will include this requested plot to the appendix of the final version of the paper.

(2) Comparing OFM to other baselines on standard generation benchmarks.

In theory, OFM can work directly in pixel-space (e.g., of CIFAR10) and has no limitation for dimensionality. However, in practice, the ML community has not yet designed input convex neural networks (ICNN) architectures that work well in the pixel-space (see, e.g., [1] for reference). Hence, for now image generation probably remains infeasible, and we encourage the ML community to invest into developing ICNN architectures.

(3) On the image-to-image transfer task, do you use unpaired training data? Is it costly to compute the minibatch OT? Still, in this case, what is the tensor shape of the 512-dimensional latent space? I think 512 is much smaller than other latent space models.

Yes, we use unpaired data. The latent space of ALAE autoencoder has dimension $D=512$ (e.g., the full encoded dataset has size $N\times 512$, where $N$ is the number of samples). Discrete OT computation in $512$-dimensional space actually is not very costly, since this task is much easier than the conjugation or hessian inversion used in training. OT computation requires $O(D\cdot B^3)$ time, where $B$ is the batch size (it is usually considerably smaller than dimension $D$). For comparison, the hessian inversion itself takes $O(BD^3)$.

Concluding remarks for all questions. We thank the reviewer one more time for high evaluation of our work. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase.

References.

[1] Korotin, Alexander, et al. "Do neural optimal transport solvers work? a continuous Wasserstein-2 benchmark." Advances in neural information processing systems 34 (2021): 14593-14605.

Dear reviewer, thank you for your detailed feedback. Please find below the answers to your questions.

(1) The reduction of our OFM to sampling of $t = 1$ only.

Our method requires considering all times $t\in [0,1]$ by its design. Considering $t=1$ breaks the theory and does not work. Indeed, this can be seen both theoretically and practically. In case $t=1$, loss (17) in OFM becomes $\min\limits_{\Psi \text{ convex}}\int\limits_{\mathbb{R}^D \times \mathbb{R}^D} \left|\left|\nabla \overline{\Psi} (x_1) - x_0 \right|\right|^2 \pi(x_0, x_1) dx_0 dx_1.$

Without $t\sim U[0,1]$, loss becomes plan-dependent, and OT recovery is not guaranteed. Hence, we lose both the desired connections with FM and OT. To also confirm this empirically, we run this optimization on a 2D Gaussian to Mixture of Gaussians setup for two plans: minibatch with 4 and 1024 samples in batch. The results are depicted in Figure 1 in the PDF file. One may clearly see that the learned map does not even match the target distribution. We also highlight that the objective above is not equivalent to the dual OT objective $\mathcal{L}_{OT}$ (equation (4) of our paper), which involves the potential $\Psi$ and its conjugate $\overline{\Psi}$.

(2) Benefits of the dynamic formulation ($t\in [0,1]$) over the static one (in addition to the strong convexity).

The main reason to consider $t \sim U[0,1]$ is to pave a new bridge between two emerging subfields in deep learning: optimal transport and flow matching. We introduce novel methodology relating both of these problems. This development potentially allows to use the advantage of one field to improve the methods in the other one. And, indeed, stronger convexity is a good benefit of the dynamic loss, but it is a bonus, not an initial goal.

(3) Comparison/relation to the dual OT optimization approach [1].

We primarily aimed to compare flow matching-based methods with each other. Metrics from MMv1 are included for completeness, and they already give a general understanding about the quality and baselines of modern static OT methods. Method [1] surely demonstrates the better result, as it used advanced amortization techniques to optimize $\mathcal{L}_{OT}$. We will include the results of the method in the table.

Overall, advanced amortization techniques such as [1] are orthogonal to our theoretical and methodological study and can be used on top of our method to further improve it. To provide the empirical evidence, we adapted the amortization scheme from [1] to our OFM. Namely, we find an approximate conjugation solution with an extra MLP $A_\phi(x_t,t):\mathbb{R}^D \times [0,1] \to \mathbb{R}^D$: \begin{eqnarray} A_\phi(x_t, t) \approx \arg\min_{z_0 \in \mathbb{R}^D} \left[ \frac{(1-t)}{2} | z_0|^2 + t \Psi (z_0) - \langle x_t , z_0 \rangle \right], \end{eqnarray} and repeat the training pipeline from [1] but with an extra condition on $t$. We run an additional OFM experiment on $D = 32$ with amortization and regression loss. This augmentation gives a boost for metrics in comparison with basic minimization techniques, see Figure 4 in the PDF file. In the final version of the paper, we will add an appendix about the amortization technique for OFM.

(4) 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?

We are sorry for this inconsistency in our text. In all our experiments, we used the LBFGS optimizer. We kindly propose you to take a look at our code. We will fix the misprint in Appendix B in the final version. Regarding Adam, in our early experiments we indeed experimented with this optimizer, but eventually found that LBFGS works better (it is faster when achieving the same tolerance). We humbly think that the detailed comparison between the optimizers is out of scope of our work.

(5) ICNNs are considered. Does OFM also work with standard MLPs/CNNs, etc.?

Following your request, we conducted several additional experiments with MLP parameterization of $\Psi$ instead of ICNN. Unfortunately, they failed to converge. Probably, this happens due to violation of conjugated functions' properties for non-convex functions required in OFM. We leave studying this phenomenon for future works.

(6) Quantitative results are provided for the unpaired image-to-image translation experiments.

Following your request, we calculated FID for I2I experiments. See our General response section and Table 1 from the attached pdf. According to the results, our OFM method beats all the baselines.

(7) How do the illustrative 2D examples look like for standard FM training with ID, MB, anmiMB plans? Do the endpoints/trajectories differ?

Following your request, we run experiments with FM under independent, minibatch and anti-minibatch plans on 2D example, see Figure 2 in the PDF file. Unlike OFM with straight paths, FM trajectories differ a lot depending on plan, curvature increases in order "mb", "idp" and "anti-mb". We, in particular, highlight the high curvature of antimibatch's trajectories. The generated distributions of end points and the target distribution are close in all the cases.

Concluding remarks for all questions. We thank the reviewer one more time for the interesting and important questions. Please respond to our post to let us know if the clarifications and additional experiments above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.

Dear reviewer nvT5.

Thank you for your time efforts to review our paper.

The discussion period is coming to the end. We would appreciate your feedback on our responses. We are happy to answer any further questions during the remaining discussion period.

Thank you in advance, Authors.

Dear reviewer, we thank you for your feedback. Please find below the answers to your questions.

(1) In essence, does the author claim that the loss presented in equation (16) is minimized by the velocity field of the optimal transport map for any plan ? As stated by the author, it seems this loss is by definition Flow Matching (FM) loss constrained to convex functions.

Yes, we claim this. This is our main theoretical result presented in Section 3.1.

(2) For the independent plan, i.e., the minimizer for the FM loss is explicitly given in [1], where for the linear interpolant. This minimizer is dependent solely on the source and target distribution (i.e., the plan for this case) and does not yield the optimal transport. If you could please explain how you overcame this.

If we consider unconstrained FM loss minimization (w.r.t. the optimized vector field) as you suggest, the analytical solution is indeed given by the formula which you provided above. However, in our OFM, we propose to optimize FM loss only over specific optimal vector fields. In this case, minimizers of constrained and unconstrained problems are not the same. We theoretically prove that the minimizer in this case is the true OT vector field.

(3) The experiment section is rather minimal. In particular, it is not clear what Figure 3 in the Image-to-Image task presents.

The primal goal of the experimental section is to empirically confirm our theoretical findings. Namely, we check that our OFM always recovers the true OT map regardless of the input transport plan. We humbly believe that the section clearly demonstrates this:

1. In qualitative 2D-toy example of Section $4.1$ (Figure $2$), our OFM recovers the same trajectories for all input plan, even for misleading anti-minibatch plan.

2. In Section $4.2$, we quantitatively evaluate the ability to solve OT problems. We run our OFM, other FM-based methods and popular OT solver on OT Benchmark which knows ground truth OT solution. Then we directly compare obtained solutions with ground truth via L2-UVP and cosine metrics. We can see that our OFM outperforms other FM-based methods in recovering OT and demonstrates metrics close to OT solvers.

Additionally, we conduct an illustrative experiment with image-to-image translation. Here the usage of OT is motivated by the necessity to keep the image content during the translation. In Figure $3$, we can see that OFM's results are the same for the both considered input plans and also the most plausible in comparison with adult images.

As per request of the other reviewers, we also calculated FID on the I2I problem to quantitatively estimate the performance, see Table 1 from the attached pdf. Our OFM method outperforms all the baselines on this task.

(3) Theorems and propositions are structured in a way which makes it very hard to follow.

We are sorry that you found our paper hard to follow, but we understand that opinions on this topic differ, e.g., Reviewer qt8g, in contrast, found our paper well-written and easy to follow. To address your comment, we will restructure Section 3.1 in the final version. Specifically, we will split in to 2 subsections:

• Subsection 3.1.1. (major) will contain optimal vector fields, our proposed loss and direct statement of the main Theorem 1. In turn, technical Lemma 1 will be moved to Appendix. This will help the reader to directly find the main result without any disctractions on technical things.

• Subsection 3.1.2 (auxiliary) will contain various supportive results: Propositions 1-3 (loss reformulation, properties of the loss). Hence, the reader which is interested in the main result only will be able to easily skip this subsection and continue to the next (practical) section.

Could you please reply, does such a change suit you well?

(4) The proof of Theorem 1 in the appendix is unclear and unsatisfying.

Please note that our main technical result is Lemma 1. It allows to take integral in FM loss over time $t$ and obtain analytical formula (18). The proof is based on rather non-trivial techniques from calculus and convex analysis: change of variables, line integrals, differentials and convex conjugation. We tried to do our best to make the proof accessible to the general ML community. Could you please point to particular aspects of the proof to clarify?

In turn, our main Theorem 1 follows from Lemma 1. We take math expectation over plan $\pi$ from both sides of (18) and notice that the left part of (18) corresponds to OFM loss and the right part matches OT dual loss (4) up to constant.

(5) Accessible main proof idea/proof sketch in the main paper.

We do not include any sketch of the proofs because it is purely technical (but still very non-trivial) and heavily relies on doing calculus tricks such as the change of variables.

In Lemma's proof (line 480), we firstly simplify OFM loss using Proposition 1. Then we make two changes of time variable $t$: $s = \frac{1}{t}, s' = \frac{1}{s-1}$, use eqn (15) and obtain integral in line 488. In line 489, we consider the curve $z_0(s')$ and take its differential $dz_0$ using implicit definition of $z_0(s')$ from eqn (15). Starting from here, we take line integrals over line $z_0$ to get the sum from (27). Next we use closed forms of differentials, integrate it from $t = 0$ to $t = 1$ and obtain analytical sum (28). Finally, we use standard inversion and Fenchel-Young’s equalities to simplify the final sum.

If you think a sketch of this type could be useful in the main text, we are happy to add it to the final version of the paper.

Concluding remarks for all questions. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.

I want to thank the authors for their response. Following the reviewers answers, I have invested more time in better following the authors claim and proofs. It is now my understanding that the key results of the paper is that the minimizer of FM loss constrained to velocity field which are a gradient of convex functions is the OT map. I acknowledged this is an important results and I will raise my score.

Dear reviewer, thank you for your feedback. Please find the answers to your questions and comments below.

(1) Pretrained autoencoders. Does OFM still work in the pixel space or a higher-resolution dataset?

In theory, OFM can work directly in pixel space and has no limitation for dimensionality. However, in practice, there are yet no input convex neural networks (ICNN) architectures that work well in the pixel-space (see, e.g., [1] for discussion). Hence, we use pretrained encoders to process images.

We emphasize that this aspect is mostly related to ICNN rather than to our developed OFM Framework. We believe that our novel theoretical findings will inspire the ML community to develop new well-performing ICNNs that may work well in the pixel space. 

(2) The proposal is compared with OT-CFM and Rectified Flow. However, the main task of the baselines in their original paper was image generation in pixel space. Is it fair to exclude this task from the experiments?

Please see the previous answer for the reason of excluding this type of experiments.

(3) The assumption that the optimal transport map can be recovered in one iteration might not hold in all scenarios, particularly in more complex distributions.

Dear reviewer, this is not an assumption. We provide and prove a rigorous theoretical statement (Theorem $1$) that our Optimal Flow Matching retrieves OT between any two absolutely continuous distributions in just one flow matching iteration. Moreover, our experiments on high-dimensional standard benchmark [1] in the field of OT (Table $1$) demonstrate that OFM actually finds a solution whose metrics are close to or better than those of the state-of-the-art OT approaches.

(4) Does OFM suffer from distorted or biased dynamic OT solutions like OT-CFM?

No, it does not. In our OFM, we run only one FM iteration with any transport plan as the initial plan (Theorem $1$) and provably recover the true OT map. That is, unlike OT-CFM, the initial plan theoretically does not affect the final solution (again, which is the OT map between distributions), i.e., there is no bias or distortion effect.

(4) How do you calculate the initial transport plan?

In 2D toy example (Section $4.1$), we practically demonstrate that OFM's final solutions do not depend on initial plans. In this setup, the considered plans are:

(a) Independent, where we independently sample batches from both distributions,

(b) Minibatch, where we independently sample batches and then rearrange them according to Discrete Optimal Transport (with $\ell^2$ cost) between them,

(c) Anti-Minibatch, where we independently sample batches and then rearrange them according to Discrete Optimal Transport based on negative Euclidean distance.

The same approach with plans is used in the other experiments.

(4) It is essential to provide more justification for using metrics in the experimental section. What are the advantages of these metrics over others in works on Unpaired I2I Translation tasks?

The metrics which we use in Section 4.2 are specially designed for evaluation of the quality of the recovered OT solution in the benchmark experiment. These metrics are generally not related to unpaired I2I problems. Our goal there is to confirm that our OFM indeed learns OT. Specifically, we use the unexplained variance percentage (UVP) as the main metric: $\mathcal{L}^2$-UVP$(T) := 100 \cdot \| T - T^*\|^2_{\mathcal{L}^2(p_0)} / \text{Var}(p_1) \%$. It directly computes the squared error between ground truth OT map $T^{*}$ (which is known by the construction of the benchmark [1, Sec. 4.2]) and learned map $T$ and then normalizes the answer. Following the benchmark, we additionally compute cosine similarity between ground truth map directions and directions of $T.$

Following your request, we additionally calculated FID for Unpaired Image-to-Image experiment from Section $4.3$. See our General response section and Table 1 from the attached PDF. According to the table, our method beats all the baselines.

(8) Does the method work for the conditional generation?

Unfortunately, ICNNs which we use in our method are not straightforward to adapt to conditional generation. In particular, to our knowledge, there are yet no research papers developing scalable ICNNs (especially, conditional) to work with pixel spaces. This is why we leave this interesting question for future research.

(9) Could using an ODE solver for sampling improve the performance of OFM?

If one uses ODE solver for the vector field generated by OFM, one gets the same solution as the one obtained by simply calculating the final point. It happens because OFM trajectories between the initial and final points are perfectly straight.

Concluding remarks. We thank the reviewer one more time for the interesting and important questions. Please respond to our post to let us know if the clarifications above suitably address your concerns about our work. We are happy to address any remaining points during the discussion phase; if the responses above are sufficient, we kindly ask that you consider raising your score.

References:

[1] Korotin, Alexander, et al. "Do neural optimal transport solvers work? a continuous Wasserstein-2 benchmark." Advances in neural information processing systems 34 (2021): 14593-14605.

We are very grateful to the Reviewers for valuable questions, fruitful discussions and raising the scores. In the final version, we will make the following edits:

1. include time convergence graph (Appendix), extra OT baselines and OFM improved metrics (few rows in Table 1), I2I FID (Section $4.3$), FM examples and amortization (Appendix), (nvT5, qt8g);

2. restructure Section 3, (following the reply to YPMs);

3. discuss the intricacy of high-dimensional (i.e., image data) generation in Limitations section, (nvT5, Bhuh).