Optimal Flow Matching: Learning Straight Trajectories in Just One Step
Authors: Nikita Maksimovich Kornilov
Petr Mokrov
Alexander Gasnikov
Alexander Korotin
Venue: NeurIPS 2024 poster (2024)
Keywords: Flow Matching
Optimal Transport
Rectified Flow
straight trajectories
Primary area: generative_models

Abstract: Over the several recent years, there has been a boom in development of Flow Matching (FM) methods for generative modeling. One intriguing property pursued by the community is the ability to learn flows with straight trajectories which realize the Optimal Transport (OT) displacements. Straightness is crucial for the fast integration (inference) of the learned flow's paths. Unfortunately, most existing flow straightening methods are based on non-trivial iterative FM procedures which accumulate the error during training or exploit heuristics based on minibatch OT. To address these issues, we develop and theoretically justify the novel Optimal Flow Matching approach which allows recovering the straight OT displacement for the quadratic transport in just one FM step. The main idea of our approach is the employment of vector field for FM which are parameterized by convex functions. The code of our OFM implementation and the conducted experiments is available at https://github.com/Jhomanik/Optimal-Flow-Matching

Optimal Flow Matching: Learning Straight Trajectories in Just One Step
Nikita Kornilov
Petr Mokrov
Skolkovo Institute of Science and Technology R.Center for AI, Innopolis University
Skolkovo Institute of Science and Technology petr. www.kooo@@kkkttch.h.o.
Moscow Institute of Physics and Technology kornilov nm@phystech. edu
Alexander Gasnikov
Alexander Korotin
Innopolis University
Skolkovo Institute of Science and Technology Artificial Intelligence Research Institute
Moscow Institute of Physics and Technology Steklov Mathematical Institute of RAS' gasnikov@yandex.ag
wwwotin@skoltech.......o.. ru
Abstract
Over the several recent years, there has been boom in development of Flow Matching (FM) methods for generative modeling. One intriguing property pursued by the community is the ability to learn flows with straight trajectories which realize the Optimal Transport (OT) displacements. Straightness is crucial for the fast integration (inference) the learned flow's ' paths. Unfortunately, most existing flow straightening methods are based non-trivial iterative FM procedures which accumulate the error during training or exploit heuristics based on minibatch OT. To address these issues, develop and theoretically justify the novel Optimal Flow Matching (OFM) approach which allows recovering the straight OT displacement for the quadratic transport in just one FM step. The main idea of our approach is the employment of vector field for FM which are parameterized by convex functions. The code of our OFM implementation and the conducted experiments available at https:: / /tt:://gtttttt com/ Cttommaik/Oppiimii Flow -ato--aaccnnn....
Introduction
Recent success generative modeling [41, 17, 9] mostly driven by Flow Matching (FM) [38] models. These models move a known distribution to target one via ordinary differential equations (ODE) describing the mass movement. However, such processes usually have curved trajectories, resulting in time-consuming ODE integration for sampling. To overcome this issue, researches developed several improvements of the FM [39, 40, 48), which aim to recover more straight paths.
Rectified Flow (RF) method [39, 40] iteratively solves FM and gradually rectifies trajectories. Unfortunately, in each FM iteration, accumulates the error, see [40, $2.2] and [39, $6). This may spoil the performance of the method. The other popular branch of approaches to straighten trajectories based on the connection between straight paths and Optimal Transport (OT) [60]. The main goal of OT is to find the way to move one probability distribution to another with the minimal effort. Such OT maps are usually described by ODEs with straight trajectories. In OT Conditional Flow Matching (OT-CFM) [48, 55], the authors propose to apply FM on top of OT solution between batches from considered distributions. Unfortunately, such heuristic does not guarantee straight paths because of minibatch OT biases, see, e.g., [55, Figure 1, right] for the practical illustration.
Russian Academy Sciences
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Contributions. In this paper, we fix the menteomme problems of the straightening meth- ods. We propose novel Optimal Flow Matching (OFM) approach (33) that after single FM iteration obtains straight trajectories which can be simulated without ODE solving. It recovers OT flow for the quadratic transport cost function, i.e., solves the Benamou- problem (Figure 1). We demonstrate the potential of OFM in the series of experiments and benchmarks (g4). The main idea our OFM is to consider during FM only specific vector fields which yield straight paths by design. These vector fields are the gradients of convex functions, which in practice are parametrized by Input Convex Neural Networks [3]. In OFM, can optionally use minibatch OT or any other transport plan as the input, and this is completely theoretically justified.
(aila)))
Pl
Optimal Flow Matching (just step, ours)
2 Background and Related Works
20
P1
In this section, we provide all necessary backgrounds for the theory. First, we recall static ($2.1) and dynamic ($2.2) formulations of Optimal Transport and solvers ($2.3) for them. Then, we recall Flow Matching ($2.4.1) and flow straightening approaches: OT-CFM ($2.4.2) and Rectified Flow ($2.4.3).
PD
Figure 1: Our Optimal Flow Matching (OFM), For any ini- tial transport plan between Po and pt, OFM obtains ex- actly straight trajectories (in just a single FM loss min- imization) which carry out the OT displacement for the quadratic cost function.
Notations. For vectors 2. y E TR., we denote the inner product by (ar, v) and the corresponding & norm by ||^||| == v zr). We use P2, (Tk D) refer to the set of absolute continuous probability distributions with the finite second moment. For vector distribution E (2)) notation x ~ p means that æ is sampled from p. For the push-forward operator, we use symbol #.
R and
2.1 Static Optimal Transport
Monge's and Kantorovich's formulations. Consider two distributions Po, Pr E P2,ac((RD) and a cost function e: IR.. x RDD + IR. Monge's 's Optimal Transport formulation is given by
inf ##pe=p1 2 c( T(zn))Po(zo)dl
(1)
where the infimum is taken over measurable functions T' IRD + RD which satisfy the mass- preserving constraint Typo Pl. Such functions are called transport maps. If there exists transport map 7 that achieves the infimum, then it is called the optimal transport map.
Since the optimal transport map I in Monge's 's formulation may not exist, there is Kantorovich's relaxation for problem (I) which addresses this issue. Consider the set transport plans [1 Pt). i.., the set of joint distributions on R x TR' which marginals are equal pn and pr, respectively. Kantorovich's 's Optimal Transport formulation is
inf / rETT(po.P))
c(ro, xi (xo, æ )dægdz
With mild assumptions on Po. Pl, the infimum is always achieved (possibly not uniquely). An optimal plan x² E 1I(po. P1) is called an optimal transport plan. If optimal x has the form [id, 1" I #oo. then T is the solution of Monge's 's formulation (I).
Quadratic cost function. In our paper, we mostly consider the quadratic cost function c(x,,
[ In this case, infimums in both Monge's and Kantorovich's OT are always uniquely attained [60, Brenier's Theorem 2.] They related by л² [id, #' Moreover, the optimal values of (1) and (2) are equal to each other. The square root of the optimal value is called the Wasserstein-2 distance (P P1) between distributions Pa and p1, i.e.,
| zoll
(Po. PI)
min TEII(po.P1)
x x )drodr
2
RP xPRR
||00 T(x0)112
min
Po po
(3)
T T# #PPP RD
2

---

W2(po,P1) pI) = inf Dual formulation. For the quadratic cost, problem (3) has the equivalent dual form [60]:
/ (er))r )dz (4)
(2()o. CONST(Po, PI)
min
((zo))o(rra)ddg +
RD
(V)
where the minimum is taken over convex functions (x) RD + R. Here ((ri) =
S [ 01) (20)] the convex (Fenchel) conjugate function of v. Is also convex.
The term CONST(Po, p1) does not depend on V. Therefore, the minimization (3) over transport plans is equivalent to the minimization of Cor (V) from (4) over convex functions V. Moreover, the optimal transport map T can be expressed via an optimal u (the Brenier potential [60]), namely,
T VV
(5)
2.2 Dynamic Optimal Transport
In [5], the authors show that the calculation of Optimal Transport map in (3) for the quadratic cost can be equivalently reformulated in a dynamic form. This form operates with a vector fields defining time-dependent mass transport instead of just static transport maps.
Preliminaries. We consider the fixed time interval [0, 1]. Let u((,-) = v((+)= [0. x RD RD
be vector field and {{ (0, } be the set of random trajectories such that for each trajectory { (=)0) the starting point 2p is sampled from Po and 2 satisfies the differential equation:
dat
(=(=x)dt, 20 ~-Po-
(6)
In other words, the trajectory fz (e is defined by its initial point speed vector uz(2)). Under mild assumptions on u, for each initial žn, the trajectory is unique.
Po and goes along the
Let (t, -) of [0, 1] RD > R² denote the flow map, i.c., it is the function that maps the initial 2 to its position at moment of time t according to the ÓDE (6), i.c.
døy
2 (20) 20-
(7)
If initial points 2p of trajectories distributed according po. then (6) defines distribution pt of 2, at time 1, which can be expressed via with the push-forward operator, i.e., P == o8 #Po
Benamou- Brenier problem. Dynamic OT is the following minimization problem:
S.¹ J
Iue (2x))
of (1
(8)
a
2
(x))
s.t. Φ P1-
In (8), we look for the vector fields u that define the flows which start Po and end P1. Among such flows, we seek for the one which has the minimal kinetic energy over the entire time interval.
There is a connection between the static OT map T = V and the dynamic OT solution u". Namely, for every initial point 2g, the vector field u defines linear trajectory (2
2 = IVV + (1 - 1) 11 E [0.1]
(9)
2.3 Continuous Optimal Transport Solvers
There exist variety of continuous OT solvers [21, 52, 54, 43, 19, 14, 59, 34, 51, 40, 36, 35, 13, 18, 57, 2, 55, 23, 44, 4, 20]. For a survey of solvers designed for OT with quadratic cost, see [33]. In this paper, focus only the most relevant ones, called the ICNN-base These solvers directly minimize objective Lor from (4) parametrizing a class of convex functions with convex in input neural networks called ICNNs [3] (for more details, see "Parametrization of up" in $3.2). Solvers details may differ, but the main idea remains the same. To calculate the conjugate function (xx) at the point æ1, they solve the convex optimization problem from conjugate definition. Envelope Theorem [1] allows obtaining closed-form formula for the gradient of the loss.
solvers [54, 33, 43, 2].
2.4 Flow Matching Framework
In this section, we recall popular approaches [40, 39, 48] to find fields u which transport given probability distribution po to a target P1 and their relation OT.
3

---

2.4.1 Flow Matching (FM)
To find such a field, one samples points 2g, #1 from transport plan я e [[p,, Pi), e.g. the independent plan Po x Pr. The vector field u is encouraged to follow the direction 21 2o of the linear interpolation 2, = (1 t) tæ at any moment E [o, 1], i.e., one solves:
min C7 (u): (w)=
[xx(( (x1 x0))1>>(<0 )drodx yll, (1 ()aro + tz
R XR
(10)
We denote the solution of (10) and the flow map (7) by u" and d* respectively. The concept of FM is depicted in Figure 2.
3ig
The intuition of this procedure is as follows: linear interpolation I = t) + fx an intuitive way to move Po to p1. but it requires knowing 21. By fitting u with the direction I2 - Zo. one yields the vector field that can construct this interpolation without any information about 21-
P1
p0
Flow Matching
The set of trajectories ((01) generated by oz (with 2n po) has a useful property: the flow map of transforms distribution Po to distribution P1 for any initial transport plan и. Moreover, marginal distribution
Me (are)
of #po is equal to the distribution of linear interpolation x; = (1 - t) + tx for any and zo, I1 n. This feature is called the marginal preserving property.
Figure 2: Flow Matching (FM) obtains a vector field u mov- ing Po to P1. FM typically operates with the independent transport plan po x p1-
To push point 20 according to learned u, needs to integrate ODE (6) via numerical solvers. The vector fields with straight (or nearly straight) paths incur much smaller time-dicreetizatio error and increase effectiveness of computations, which is in high demand for applications.
Researchers noticed that some initial plans æ can result in more straight paths after FM rather than the standard independent plan po P1. The two most popular approaches to choose better plans Optimal Transport Conditional Flow Matching [48, 55] and Rectified Flow [40].
2.4.2 Optimal Transport Conditional Flow Matching (OT-CFM)
one uses the OT plan n" as the initial plan for FM, then it returns the Brenier's vector field u", which generates exactly straight trajec- tories (9). However, typically, the true OT plan x* is not available. In such case, in order achieve some level of straightness in the learned trajectorics, a natural idea to take the initial plan π to be close to the optimal x². Inspired by this, the authors of OT-CFM [48, 55] take the advantage of minibatch OT plan approximation. Firstly, they independently sample batches of points from Po and p1. Secondly, they join the batches together according to the discrete OT plan between them. The resulting joined batch is then used in FM. The concept of OT-CFM is depicted in Figure 3.
P0
P)
Flow Matching
P1
The main drawback of OT-CFM is that it recovers only biased dy- namic OT solution. In order to converge to the true transport plan the batch size should be large [6], while with growth of batch size computational time increases drastically [56]. In practice, batch sizes that ensure approximation good enough for applications are nearly infeasible to work with.
Figure 3: OT-CFM uses mini- batch OT plan to obtain more straight trajectories.
2.4.3 Rectified Flow (RF)
In [40], the authors propose an iterative approach to refine the plan w, straightening the trajectories more and more with each iteration. Formally, Flow Matching procedure denoted by FM takes the transport plan N as input and returns an optimal flow map via solving (10):
(M
(1I)

---

50"
Pl
P1
Pr
Flow Matching (step I)
Flow Matching (step 2)
Flow Matching (step Kj
20
P1
P1
Pd
P0
P9
Figure 4: Rectified Flow iteratively applies FM straighten the trajectories after each step.
One can iteratively apply FM to the initial transport plan (e.g.. the independent plan), gradually rectifying it. Namely, Rectified Flow Algorithm K-th iteration has the following update rule
K+1 FM(*"),
π K+I fid, ø K #po.
(12)
where oK denote flow map and transport plan on K-th iteration, respectively.
With each new FM iteration, the generated trajectories {(2)) } provably become more and more straight, i.e., error in approximation 2, (1-t)af t22 Vt E [0, 1] decreases the number of iterations K grows. The concept of RF is depicted on Figure 4.
The authors also notice that for any convex cost function the flow map of from Flow Matching yields lower or equal transport cost than initial transport plan ı
[ c( of (olo)))o(()))dr S RD
f c( (T 2n )dzdd11-
(13)
JRRRRD
Intuitively, the transport costs are guaranteed to decrease because the trajectories of FM as solutions of well-defined ODE do not intersect each other, even if the initial lines connecting zo and 2ı can. With each iteration RF (12), transport costs for convex cost functions do not increase, but, for given cost function, convergence to its OT plan is not guaranteed. In [39], the authors address this issue and, for any particular convex cost function c, modify Rectified Flow to converge to OT map for In this modification, called c- c-Rectified Flow (c-RF), the authors slightly change the FM training objective and restrict the optimization domain only to potential vector fields us
Ve(V((.)) where f+()) RDD - TR is arbitrary time-dependent scalar valued function and is the convex conjugate of the cost function c. case of the quadratic cost function, the training objective remains the same, and the vector field u, is set the simple gradient f-((-) of the scalar valued function f(-
Unfortunately, in practice, with each iteration (c-)RF accumulates error caused by inexactness from previous iterations, the issue mentioned in [39, $6, point 3]. Due to neural approximations, we can not get exact solution of FM (e.g. of #Pu + P1), ), and this inexactness only grows with iterations. In addition, training of (c-)RF becomes non-simulation free after the first iteration, since calculate the plan x K+1 id,, o +1 #o it has to integrate ODE.
3 Optimal Flow Matching (OFM)
In this section, we provide the design of our novel Optimal Flow Matching algorithm (1) that fixes main problems of Rectified Flow and OT-CFM approaches described above. In theory, it obtains exactly straight trajectories and recovers the unbiased optimal transport map for the quadratic cost just in one FM iteration with any initial transport plan. Moreover, during inference, our OFM does not require solving ODE to transport points.
We discuss the theory behind our approach ($3.1), its practical implementation aspects ($3.2) and the relation to prior works (53.3). All our proofs are located in Appendix A.
3.1 Theory: Deriving the Optimization Loss
We want to design a method moving distribution po to pı via exactly straight trajectories. Namely, we aim to obtain straight paths from the solution of the dynamic OT (8). Moreover, want to limit ourselves to just one minimization iteration. Hence, we propose our novel Optimal Flow Matching (OFM) procedure satisfying the above-mentioned conditions. The main idea of our OFM is to
5

---

minimize the Flow Matching loss (10) not over all possible vector fields u, but only over specific optimal ones, which yield straight paths by construction and include the desired dynamic OT field u².
Optimal vector fields. We say that a vector field u. is optimal it generates linear trajectories ((=)) such that there exist a convex function Vp IR_ IR, which for any path (21)1 pushes the initial point 2g the final one 21
(z)) ie., (1 t) + (2)) t te [0,
2t
The function V defines the ODE
dzt
(V)( - =0) == =0
(14)
Equation (14) does not provide a closed formula for u it depends on zo. The explicit formula is constructed as follows: for a time & E [0, 1] and point Lt, we can find trajectory (zrheeoo.1) s.t.
2t = = t)20 + (2))
(15)
and recover the initial point 20. We postpone the so- lution of this problem to $3.2. For now, we define 20
(V((2))
the inverse of flow map (7) (xx))
vir(e)) Pl
and the vector field u, (x+) = VV)(2)) 20 =
(V((o))) (zx))
(o) (I2), which generates
P0
ODE (14), i... dæp = 4 (2) dt.. The concept of op- timal vector fields is depicted on Figure 5.
20
(m)) (1) Fp
We highlight that the solution of dynamic OT lies in the class of optimal vector fields, since it generates linear trajectories (9) with the Brenier potential T (5).
Figure 5: An Optimal Vector Field: vector field u with straight paths is parametrized by gradient a convex function V.
Training objective. Our Optimal Flow Matching (OFM) approach is as follows: we restrict the optimization domain of FM (10) with fixed plan only to the optimal vector fields. We put the formula for the vector field uy into FM loss from (10) and define our Optimal Flow Matching loss:
COFM( (T) = CF )
I| (#+) - (x1 =o)ll2 vn ddroddydd (16)
RRPXR
(1-0)20 tx
Our Theorem states that OFM solves the dynamic OT via single FM minimization for any initial Theorem 1 (OFM and OT connection). Consider two distributions po: Pi E (TR and any transport plan æ II pi) between them. Then, the dual Optimal Transport loss Lor (4) and Optimal Flow Matching loss ForM (16) have the same minimizers, i.e.,
arg min LO (Vr) arg min LoT (V).
convex
V
3.2 Practical implementation aspects
In this subsection, we explain the details of optimization of our Optimal Flow Matching loss (16).
Parametrization of V. In practice, parametrize the class of convex functions with Input Convex Neural Networks (ICNNs) [3] Ve and parameters These are scalar- scalar-valued neural networks built in such a way that the network is convex in its input. They consist of fully-connected convolution blocks, some weights of which are set to be non-negative order to keep the convexity. In addition, activation functions are considered to be only non-decreasing and convex in each input coordinate. These networks are able to support most of the popular training techniques (e.g., gradient descent optimization, dropout, skip connection, etc.). In Appendix B, we discuss the used architectures.
OFM loss calculation. We provide an explicit formula for gradient OFM loss (16).
Proposition 1 (Explicit Loss Gradient Formula). The gradient of FT can be calculated as
NO-GRAD (0
d. = Ei.ro,ar mmm
NO-GRAD { (tV2)0(20) 2 + +11 - -)) -1 (ro 20) }. VVo(2o))
d0 de where variables under NO-GRAD remain constants during differentiation.
t

---

Flow map inversion. In order to find the initial point 20 = -0,,) - (xx)) we note that (15)
3t = 1)z0 (20)
equivalent to
(1 (1-1) 1111 + (m2, (20) = O.
V
2
The function under gradient operator V has minimum at the required point zo, since at 2o the gradient of it equals 0. I t < the function is at least (1 -)-strongly convex, and the minimum is unique. The case i = 1 negligible in practice, since has zero probability appear during training.
We can reduce the problem of inversion to the following minimization subproblem
(1 (1-t), =01 3 (20) (aze =0)
(0 M/ (xx) arg min S
(17)
2
Optimization subproblem (17) is at least (1 - -)-stonngly convex and can be effectively solved for any given point (in comparison with typical non-convex optimization tasks).
Algorithm. The Optimal Flow Matching pseudocode is presented in listing We estimate math expectation over plan π and time t with uniform distribution on [0, 1] via unbiased Monte Carlo.
Algorithm Optimal Flow Matching
Input: Initial transport plan E [P (i)) number of iterations K, batch size B, optimizer Opt, sub-problem optimizer SubOpt, ICNN V
1: for k = O..... K K- do
) 1= 2: Sample batch {(r) x
of size B from plan æ.; 3: Sample times batch (t*) of size B from 0 1];
4: Calculate linear interpolation x (1 Da t'xj for all i B:
5: Find the initial points zi via solving the convex problem with SubOpt:
{
(1 (1-(") (*)
|= ( - (5t z6) }
20 = NO-GRAD
arg min =6
6: Calculate loss LOFM
1 - Σ (NO-GRAD (2 (1--))+11+))) (1 (2) ŁOFM
2b) ) vWo(z))
7: Update parameters 0 via optimizer Opt step with dCoFM. db
8: end for
3.3 Relation to Prior Works
In this subsection, we compare our Optimal Flow Matching and previous straightening approaches. One unique feature of OFM is that it works only with flows which have straight paths by design and does not require ODE integration to transport points. Other methods may result in non-straight paths during training, and they still have to solve ODE even with near-straight paths.
OT Solvers [54, 43, 2]. According to Theorem 1, our OFM and dual OT solvers basically minimize the same OT loss (4). However, our OFM actively utilizes the temporal component of the dynamic process. It allows us to pave novel theoretical bridge between OT and FM. Such a direct connection can lead to the adoption of the strengths of both methods and deeper understanding of them.
OT-CFM [48, 55]. Unlike our OFM approach, OT-CFM method retrieves biased OT solution, and the recovery of straight paths is not guaranteed. In OT-CFM, minibatch OT plan appears a heuristic that helps to get better trajectories in practice. In contrast, usage of any initial transport plan и in our OFM is completely justified in Theorem
Rectified Flow [40, 39]. Rectified Flows [40], the authors iteratively apply Flow Matching to refine the obtained trajectories. However, in each iteration, RF accumulates enror since one may not
7

---

learn the exact flow due to neural approximations. In addition, RF does not guarantee convergence to the OT plan for the quadratic cost. The c-Rectified Flow [39] modification can converge the OT plan for any cost function but still remains iterative. In addition, RF and c-RF both requires ODE simulation after the first iteration to continue training. In OFM, we work only with the quadratic cost function, but retrieve its OT solution in just one FM iteration without simulation of the trajectories.
Light and Optimal Schrödinger Bridge. In [22], the authors observe the relation between Entropic Optimal Transport (EOT) [42, 12] and Bridge Matching (BM) [53] problems. These are stochastic analogs of OT and FM, respectively. one considers stochastic processes with non-zero stochasticity. The authors prove that, during BM, one can restrict considered processes only to the specific ones and retrieve the solution of EOT. Hypothetically. our OT/FM case is a limit of their EOT/BM case when the stochasticity tends to zero. Proofs in [22] for EOT are based on sophisticated KI. divergence properties. We do not know whether our results for OFM can be derived by taking the limit of their stochastic case. To derive the properties of our OFM, we use completely different proof techniques based on computing integrals over curves rather than KL-based techniques. Besides, in practice, the authors of [22] mostly focus on Gaussian mixture parametrization while our method allows using neural networks (ICNNs).
EOT and BM, instead of deterministic ODE and flows,
3.4 Theory: properties of OFM
In this subsection, we provide the OFM's theoretical properties, which give an intuition for under- standing of its main working principles and behavior. Proposition 2 (Simplified OFM Loss),
We can simplify (16) to more suitable form:
(6." ( (xaa 3o
COFM (-)=
(Fo, x1 droddr dt, xz = (1-t)x0 + tzr (18)
t
RPP XPD
The simplified form (18) shows that OFM loss actually measures how well yi restores initial points To of linear interpolations depending on future point x; and time
Generative properties of OFM. In this paragraph, we provide another view on our OFM approach. In our OFM, we aim to construct vector field u which is as close the dynamic OT field u as possible. We can use the least square regression to measure the distance between them:
DIST(U, u²) = S. f [lue(z)) - uz(xx) (xx)) dzo
(19)
: (2.)
Proposition 3 (Intractable Distance). The distance DIST(U, u²) between an arbilrary vector field u and OT field u² equals to the FM loss from (10) with the optimal plan n², iee,
DIST(U, u²) M) (Fr
We can not minimize intractable DIST(u, u²) since the optimal plan x² is unknown. In OT-CFM [55], authors heuristically approximate x²" in LT (u). but obtain biased solution. Surprisingly. for the optimal vector fields, the distance can be calculated explicitly via any known plan I.
Proposition (Tractable Distance For OFM). The distance DIST(u v u²²) between an optimal vector field R generated by a convex function V and the vector field " with the Brenier potential 1 can be evaluated directly via OFM loss (16) and any plan n:
DIST(u u "u') L (u²) CF = -
(20)
In (20), the first term is our tractable OFM loss, and the second term does not depend on V. Hence, during the whole minimization process in our OFM, we gradually lower the distance (19) between the current vector field and the dynamic OT field up to the complete match.
4 Experimental Illustaaioon
In this section, we showcase the performance of our Optimal Flow Matching method on illustrative 2D scenario ($4.1) and Wasserstein-2 benchmark [33] ($4.2). Finally, apply our approach for solving high-dimensional unpaired image-to-image translation in the latent space pretrained ALAE autoencoder ($4.3). The PyTorch implementation of our method is publicly available at
8

---

sttps://githnn ctw/
The technical details of our experiments (architectures, hyperparameters) are in the Appendix B.
4.1 Illustrative 2D Example
In this subsection, we illustrate the proof-of-concep of our OFM on 2D setup and demonstrate that OFM's solutions do not depend on the initial transport plan т. We run our Algorithm 1 between standard Gaussian Pa (O, I) and a Mixture of eight Gaussians P1 depicted in the Figure 6a. We consider different stochastic plans т: independent plan Pa X P1 (Figure 6b), minibatch and antiminibatch (Figures 6c, 6d) discrete OT (quadratic cost) with batch size Bmb 64. In the antiminibatch case, we compose the pairs of source and target points by solving discrete OT with minus quadratic cost || - y/ا.. The fitted OFM maps and trajectories are presented in Figure 6. We empirically see that our OFM finds the same solution for all initial plans π.
For completeness, in Appendix B.2, we apply these plans the original FM (10), and show that, in comparison with our OFM, the resulting paths obtained by FM considerably depend on the plan.
Ritted distibution
"ajechory end (1ihed)
Trajectory (tited]
To
(a) Input and target distributions pa and pi-
(b) Our fitted OFM:
(c) Our fitted OFM: minibatch π.
(d) Our fitted OFM; antiminibatch
independent π Po x P1-
Figure 6: Performance of our Optimal Flow Matching on Gaussian Eight Gaussians 2D setup.
4.2 High-dimensional OT Benchmarks
In this subsection, we quantitatively compare our OFM and other methods testing their ability to solve OT. We run OFM, FM based methods and OT solvers OT Benchmark [33]. The authors provide high- nigh-dimenssooal continuous distributions po: P1 for which the ground truth OT map T for the quadratic cost is known by the construction. To assess the quality of retrieved transport maps. we use standard unexplained variance percentage C2-UVP(T) := 100 | - T l22=(pp) Var(p,) 5% [33]. It directly computes the normalized squared error between OT map 1" and learned map I'.
Competitors. We evaluate Conditional Flow Matching (OT-CFM), Rectified Flow (RF), ceRetiified Flow (c-RF), the most relevant OT solver MMv-1 [54] and its amortized version from [2]. In [54] and [2], the authors directly minimize the dual formulation loss Lor (4) by parametrizing V/ with ICNNs and calculating ((11) via solving convex optimization subproblem. The latter is similar to our inversion (17). Additionally, in [2], the authors use MLPs to parametrize V, and we include these results well. Following [33]. we also provide results for linear OT map (bascline) which translates means and variances of distributions to each other. For our OFM, we consider two initial plans: independent plan (Ind) and minibatch OT (MB), the batch size for the latter is Bmb 64. The overall results are presented in Table
More details are given in Appendix B.3.
Solver
Solver type
D=2
D
D=8
D 16
D 32
D 64
D = 128
256
MM-1 [M4]
0.2
1.U
1.8
1.4
6.9
8.1
2.2
2.8
Amortization, ICNN" [2]
Dual OT solver
0.26
0.78
1.6
1.1
1.9
4.2
1.8
2.0
Amurtzation. MLP** [2]
0.03
0.22
0.6
0.8
2.0
0.67
0.59
Lineur" [33]
Riseline
14.1
14.9
27.3
41.6
55.3
63.9
63.6
67.4
OTC [55]
0.16
0.13
2:27
4.33
7.9
11.4
12.1
27.5
40)
8.58
46
63.33
63.
(39]
Flow Matching
1.56
13.11
17.87
35.39
76.48
OFM Ind (Ours)
0.19
0.61
1.4
1.1
1.47
8.35
1.96
3.96
OHM MB (Ours)
0.15
0.52
1.2
1.0
1.2
7.2
1.5
2.9
Table I: L² - UVP values of salvers litted high dighndinnnsional benchmarks dimensions D 2, 16, 32, 64, 128, 256.
The hest metric over Flow Matching based methods is bolded. Metrics are taken from [33]. Metrics are taken from [2].
9

---

Results. Among FM-based methods, our OFM with any plan demonstrates the best results. For all plans, OFM convergences to close final so- lutions and metries. Minibatch plan provides a little bit better results, especially in high di- mensions. theory, the OFM results for any plan π must be similar. However, in stochas- tic optimization, plans with large variance yield convergence slightly worse solutions.
(OURS) (OURS)
OFM Ma OFM
MLP-based OT solver usually beats our OFM, since MLPs do not have ICNNs' limitations in practice. However, usage of MI.P is empir- ical trick and is not completely justified. We also run OFM with MLP instead ICNN, and,
RF
OTCHM CRF
unfortunately, the method fails to converge.
RF demonstrates worse performance than even linear baseline, but it is ok since it is not de- signed solve W2 OT. In turn, c-RF works better, but rapidly deteriorates with increasing dimensions. OT-CFM demonstrates the best results among baseline FM-based methods, but still underperforms compared to our OFM solver in high dimensions.
Figure 7: Unpaired Adult by FM solvers, ALAE 1024 x 1024 FFHQ latent space.
4.3 Unpaired Image-to-imet to-agg Transfer
Another task that involves learning translation between two distributions is unpaired image-to-image translation [63]. We follow the setup of [32] where translation is computed in the 512 dimensional latent space of the pre- rre-trained ALAE autoencoder [46] on 1024 x 1024 FFHQ dataset [29]. In particular, we split the train FFHQ sample (60K faces) into children and adults subsets and consider the corresponding ALAE latent codes the source and target distributions Po and p1. At the inference stage, we take a new (unseen) adull face from a test FFHQ sample, extract its latent code, process with our learned model and then decode back to the image space. The qualitative translation results and FID metric [25] are presented in Figure 7 and Table 2, respectively.
OFM, Ind
OFM. MB
Method
(Ours)
(Ours)
RF
c-RF
RF
OT-CFM
FID
11.8
11.0
21.0
13.5
12.9
Table 2: FID metric on Adult translation task for the Flow Matching based methods.
The batch size for minibatch OT methods ((OFM, MB], LOT-CFM) is Bub = 128. Our OFM con- verges nearly the same solution for both independent and MB plans, and demonstrates qualitatively plausible translations. The most similar results to our method are demonstrated by [c-RF]. Similar to OFM, this method (in the limit of RF steps) also recovers the quadratic OT mapping.
Discussion
Potential impact. We believe that our novel theoretical results have a huge potential for improving modern flow matching-bass methods and inspiring the community for further studies. We think this is of high importance especially taking into account that modern generative models start to extensively use flow matching methods [61, 41, 17].
Limitations and broader impact are discussed in Appendix C.
Acknowledgement
The work of N. Kornilov has been financially supported by The Analytical Center for the Government of the Russian Federation (Agreement No. 70-2021-00143 01.11.2021, IGK 000000D730334554000)..
10

---

References
[1] SN Afriat. Theory of maxima and the method of lagrange. SIAM Journal on Applied Mathe- matics, (M33):343-57, 1971.
[2] Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
[3] Brandon Amos, Lei Xu, and Zico Kolter. Input convex neural networks. In International Conference Machine Learning, pages 146-155. PMLR, 2017.
[4] Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev. Neural optimal transport with general cost functionals. The Twelfth International Conference on Learning Representations, 2024.
[5] Jean-David Benamou and Yann Brenier. A computational fluid mechanics solution to the monge-kantorovich mass transfer problem. Numerische Mathematik, 84(3):375-393, 2000.
[6] Espen Bernton, Pierre E Jacob, Mathieu Gerber, and Christian Robert. On parameter estimation with the wasserstein distance. Information and Inference: A Journal of the IMA, 8(4):657-676, 2019.
[7] Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
[8] Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised training of conditional monge maps. Advances in Neural Information Processing Systems, 35:6859-688 6872, 2022.
(9] Hanqun Cao, Cheng Tan, Zhangyang Gao, Yilun Xu, Guangyong Chen, Pheng-Ann Heng, and Stan Z. Li. A survey generative diffusion models. IEEE Transactions Knowledge and Data Engineering, 2024.
[10] Shreyas Chaudhari, Srinivasa Pranav, and José MF Moura. Gradient networks. arXiv preprint 2rXiv:2404.03361, 2024.
[11] Yize Chen, Yuanyuan Shi, and Baosen Zhang. Optimal control via neural networks: convex approach. arXiv preprint arXiv: 1835, 2018.
[12] Yongxin Chen, Tryphon Georgiou, and Michele Pavon. On the relation between optimal transport and schrödinger bridges: A stochastic control viewpoint. Journal of Optimization Theory and Applications, 169:67 1-61,, 2016.
[13] Jaemoo Choi, Jaewoong Choi, and Myungjoo Kang. Generative modeling through the semi- dual formulation of unbalanced optimal transport. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[14] Max Daniels, Tyler Maunu, and Paul Hand. Score-based generative neural networks for large- scale optimal transport. Advances in neural information processing systems, 34:12955-12966 2021.
[15] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications score-bar based generative modeling. Advances in Neural Information Processing Systems, 34:17695-17709, 2021.
[16] John R Dormand and Peter J Prince. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-25 29119992261111 1980.
[17] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, 2024.
[18] Jiaojiao Fan, Shu Liu, Shaojun Ma, Hao-Min Zhou, and Yongxin Chen. Neural monge map estimation and its applications. Transactions on Machine Learning Research, 2023. Featured Certification.

---

[19] Jiaojiao Fan, Amirhossein Taghvaci, and Yongxin Chen. Scalable computations of wasserstein barycenter via ipput convex neural networks. In Marina Meila and Tong Zhang, editors, Proceed- ings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 1571-1581. PMLR, 18-24 Jul 2021.
[20] Milena Gazdieva, Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Extremal domain translation with neural optimal transport. Advances in Neural Information Processing Systems, 06:40331-40413 2023.
[21] Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large-scale optimal transport. Advances in neural information processing systems, 29, 2016.
[22] Nikita Gushchin, Sergei Kholkin, Evgeny Burnaev, and Alexander Korotin. Light and optimal schrödinger bridge matching. In Forty-first International Conference on Machine Learning. 2024.
[23] Nikita Gushchin, Alexander Kolesov, Alexander Korotin, Dmitry P Vetrov, and Evgeny Burnaev. Entropic neural optimal transport via diffusion processes. Advances in Neural Information Processing Systems, 36, 2024.
[24] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition (CVPR), June 2020.
[25] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[26] J Hiriart- Urruty and Yves Lucet. Parametric computation of the legendre- conjugate with application to the computation of the moreau envelope. Journal of Convex Analysis, 14(3):657, 2007.
[27] Pieter-Jan Hoedt and Günter Klambauer. Principled weight initialisation for input-convex neural networks. Advances Neural Information Processing Systems, 36, 2024.
[28] Chin-Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex potential flows: Universal probability distributions with optimal transport and convex optimization. In International Conference on Learning Representations, 2021.
[29] Tero Karras, Samuli Laine, and Timo Aila. style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 4401 4401-4410, 2019.
[30] Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations, 2021.
[31] Alexander Korotin, Vage Egiazarian, Lingxiao Li, and Evgeny Burnaev. Wasserstein iterative networks for barycenter estimation. Advances in Neural Information Processing Systems, 35:15672-15686, 2022.
[32] Alexander Korotin, Nikita Gushchin, and Evgeny Burnaev. Light schrödinger bridge. In The Twelfth International Conference on Learning Representations, 2023.
[33] Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? a continuous wasserstein-2 benchmark. Advances in neural information processing systems, 34: 14:14553-, 2021.
[34] Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wasserstein-
2 barycenter estimation without minimax optimization. In International Conference on Learning Representations, 2021.
[35] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Kernel neural optimal transport.
In The Eleventh International Conference Learning Representations, 2023.
12

---

[36] Alexander Korotin, Daniil Selikhanovych, and Evgeny Burnaev. Neural optimal transport. In The Eleventh International Conference on Learning Representations. 2023.
[37] Maciej Ławryńezuk. Input convex neural networks in nonlinear predictive control: A multi- model approach. Neurocompuring. 513:273-22 293, 2022.
[38] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. The Eleventh International Conference on Learning Representations, 2023.
[39] Qiang Liu. Rectified flow: A marginal preserving approach to optimal transport. arXiv preprint arXiv:2209. /4577, 2022.
[40] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
[41] Xingchao Liu, Xiwen Zhang, Jianzhu Ma, Jian Peng, and qiang liu. Instaflow: One step is enough for high-quality diffusion- based text-to-image generation. In The Twelfth International Conference on Learning Representations, 2024.
[42] Christian Léonard. A survey of the schrödinger problem and some of its connections with optimal transport. Discrete and Continuous Dynamical Systems, 34(4):1533-1574 2014.
[43] Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pages 1672-6681. PMLR, 2020.
[44] Petr Mokrov, Alexander Korotin, Alexander Kolesov, Nikita Gushchin, and Evgeny Burnaev. Energy-guided entropic neural optimal transport. The Twelfth International Conference on Learning Representations, 2024.
[45] Daniel Morales Morales-Brotons, Thijs Vogels, and Hadrien Hendrikx. Exponential moving average of weights in deep learning: Dynamics and benefits. Transactions on Machine Learning Research, 2024.
[46] Stanislav Pidhorskyi, Donald A Adjeroh, and Gianfranco Doretto. Adversarial latent au- toencoders. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14104-14113, 2020.
[47] ES Polovinkin and MV Balashov. Elements of convex and strongly convex analysis, 2007.
[48] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Ennich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. In International Conference on Machine Learning, pages 28100-28127. PMLR, 2023.
[49] Jack Richter-Powell, Jonathan Lorraine, and Brandon Amos. Input convex gradient networks. arXiv preprint arXiv:2111.12187, 2021.
[50] Ralph Tyrell Rockafellar. Convex analysis. 2015.
[51] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps. International Conference on Learning Representations, 2022.
[52] Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Mathieu Blondel. Large scale optimal transport and mapping estimation. In International Conference on Learning Representations, 2018
[53] Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
[54] Amirhossein Taghvaei and Amin Jalali. 2-wasserstein approximation via restricted convex potentials with application to improved training for gans. arXiv preprint arXiv: 1902.07 2019.
13

---

[55] Alexander Tong, Kilian FATRAS, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing flow-based genera- tive models with minibatch optimal transport. Transactions on Machine Learning Research, 2024. Expert Certification.
[56] Nazarii Tupitsa, Pavel Dvurechensky, Darina Dvinskikh, and Alexander Gasnikov. Numerical methods for large-scale optimal transport. arXiv preprint 2r 2022.
[57] Théo Uscidda and Marco Cuturi. The monge gap: A regularizer learn all transport maps. In International Conference on Machine Learning, pages 34709-34733, PMLR, 2023.
[58] Bryan Van Scoy, Randy A Freeman, and Kevin M Lynch. The fastest known globally convergent first-order method for minimizing strongly convex functions. IEEE Control Systems Letters, 2(1):49 54, 2017.
[59] Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9):1 134, 2021.
[60] Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
[61] Hanshu Yan, Xingchao Liu, Jiachun Pan, Jun Hao Liew, Qiang Liu, and Jiashi Feng. Perflow: Piecewise rectified flow as universal plug-and-play accelerator, 2024.
[62] Shu Yang and B Wayne Bequette. Appimization-based control using input convex neural networks. Computers Chemical Engineering, 144:107143, 2021.
[63] Jun-Yar Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using sece-cconsistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pages 2223-2232, 2017.
14

---

A Proofs and auxiliary statements
In this section, place the proofs of all results from the main manuseript and some auxiliary results. We reorder the proofs of Prop. and Prop. since the former is based on the latter.
Note that in all of our theoretical derivations, if not stated explicitly, assume the differentiabillty of convex potential V/ at given points 20, 4t etc. This assumption is done for simplicity and does not spoil our theory. The convex functions are known to be differentiable almost surely w.r. Lebesgue measure [50]. Therefore, since we consider absolutely continuous reference distributions po. P1 (see $2.1), the differentiability of Y at the considered points also holds almost surely.
Proof of Proposition 2 (Simplified OFM Loss)
Proof. By definition ( equals to
COFM (V):=
|u (xx) (01
zu)l ( x1) dodddd dt, 22 (1 t) tz (21)
RDXRD
For fixed points To, 21 and time t integrand, we find a point 20 (z)) such that in moment t E [0, 1] it is transported to point 2₁ = (1 t) + tzj, This point 20 satisfies equality
= +V (20) + (1 - t)
We define the vector field i v
as
uy (xt)
(2)(2)) -0
St. 20
t
Putting u. (1x) in the integrand of (21), we obtain simplified integrand
2
| xo uz (=1)
11 x7u
t
|t,, - txo
((1 t)æo + txx)) ) + 2o1l
(0t v ) (xre) - 20
20
xo
t
Proof of Proposition 1 (Explicit Loss Gradient Formula)
Proof. Point 2u (0, > now depends parameters 6. We differentiate the integrand from the simplified OFM loss (18) for fixed points zo 2i and time t, i.c..
1 gz
|z - 2oll2)
= 2
20
dz
d0
(22)
12
d0
For point 20, the equation (27) holds true:
It = - (1-t)zo + (VVe(20).
(23)
We differentiate (23) w.r.t. 0 and obtain
0 (1 t) dzp
dz0
0V
+ 2 (20)
(50)
d0
d0
30
dzg
oV V
- (IV2 + (1 ())
(20)-
db
a0
Therefore, we have
2 Eo 20 (FV2 + (1 - (1) 00
-1 8V
(22) =
(20 d0
t
(2(122 + (1 t)1) 1 (xo 20)
0V V
(2o)
(24)
80
15

---

Now the differentiation over & is located only in the right part of (24) in the term 20 and the left part of (24) can be considered as constants during differentiation. To get the gradient of OFM loss we also need to take math expectation over plan æ and time l.
@ Hence, point 26
The following two Lemmas are used to prove our main theoretical result, Theorem 1.
Lemma 1 (Properties of convex functions and their conjugates). Let V TRD R be convex function; I0 E IRP Let & and y be differentiable at zo and ø1 correspondingly. Then the following statements are equivalent:
(i) (zo):
(ii) arg max { (x1 2) - (=))];
ZERD
(iii) Fenchel- Young's equality: (Fro) + Vxx1) = #o);
(iv) Zo (11):
(v) I1 arg max (2,,00) +=)=));
ZER"
Proof. The lemma is simplified version of [47, Theorem 1.16.4). Also, the proof can be constructed by combining facts from [7, $3.3).
Lemma 2 (Main Integration Lemma). For any two points 2g. E R and a convex function v the following equality holds true:
(Z)) (21) (xo, 21)1, (lu@( (Ix) (21 =0)1 -01--2
(25)
where = two + (1 -1-18 1)x1
Proof. Following Proposition 2, we use the simplified loss form, i.e.,
|| (xx)) - (a1 - zo)11 マ 10 xoll?
(26)
where = (t) = (0, > (xx) satisfies the equality:
Tt v)
(20) + (1 - t)
(27)
Next, we substitute (26) into rhs of (25) integrate w.r.t. time from 0 excluding to 1 excluding (This exclusion does not change the integral):
S. [x+)
(xx) - (x1 -roo)l ('' 82 1 2dt.
(28)
To further simplify (28) we need some preliminary work. Following (27) we note:
21 =V + t) t)ao tzı *
((VV(=)) x11)
(1 - () =0) +
(V - -x1)
(Ii - 20) =
(29)
t
(1 1)2
[V)(())) - =211
|| - zoll2
(30)
t
Changing in (28) time variable to s = -
1
ds
dt 7T-
and substitution of (30) yield:
m
(1 2)2
dt
2 (2)()) - xool 1dt
1 (t) - zoll
[V((2(()) - 1ds.
(31)
(1 - 1)2
16

---

We notice that set of points zu(s(t)) (o,)) (0"0-((am) (xx), (0, )) forms curve in R with parameter t (or s(t)). Now we are to process formula (31) by switching from the integration w.r.t. parameter to the integration along this curve. To do it properly we need two things:
1. Limits of integration. The limits of integration along the curve 20(
==(t)l Io
(32)
zo(t)l(=1 (VV) (z1)
(z1)
Lemmn (iv) (V(Z)).
2. Expression under integral sign w.r.t. differential dag. Starting with (29), we derive:
(29) 8 (ZV()))) - x1) = (xo - 20) -
( - =1)
d/z0 zol =
s V + (V)((2)) x1)ds -d20
(V xi)ds = =((sV2((z)) + T)dzo.
(33)
Now we proceed with (31):
(31) =
(VV(2)) x1, VV(=z)) - xxi) ds
(33
(0) - VV((2)) (8 (2)) + T)dzo)
20
(21 (V((2)) dzu) (s(01 V-V(22)), 22 (2))dd))
20
29)
(0) (V(22), dog) + / (z0 - vo. V2 (Z)))
(34)
We notice that
d(20, VI((2)) = (20. (2)(22))2z)) + (dzo, (20(22)) =
(zo, ( = d(zo. vV((2)) - (=)) dzo)-
As a consequence, we further proceed with (34):
(34) = / (01 - V.V(2)) + (20- xu. (2)
(2), d2o) - / (71((2)).22))
20
dl.00 v-V(22)) / (VV(2o), d2o) (Ro: (=((22))20)
20
(x1, d2u) -2 2 (z)) dzo) + d(zo, V.L(2o)) - -f (2o. (2)(2)) (35)
20
Under all integrals we have closed form differentials
(2), dzo) = d(zn, 20):
(V dzo)
dVV(2o),
(vo. V VZ)(zz)) = d(co. VV)(=)))
We integrate them from initial point zo to the final VV((I) according to limits (32) and get
(35) = / (((11 z0) -2 2 f [dV(2))) + d(20, VV((2)) / d(au, V
20
20
20
(Z1, (V((2)) - (211 2o) + (x) V(VV(zz))) (WVVE))
(Z0, (V((Z)) + (2o. v(((2))) - (2o. VV((V(Zz))).
(36)
17

---

Now we use properties of conjugate functions (Lemma (I)):
((V((z1))
(i))
x)) -
VU(VV(zi)) (++
21.
This allows us to simplify (36):
(36) (1i, (2()))) (2) mo) + (20 + +(()) ((()))((V(())))) + (VV((z))
((VV)
(ro, VV (2o)) + (no, Vv(xx)) - (10,
= (20) - T ((1) (1).
Integrating equality (25) over the given transport plan æ and considering the formulas for the losses (4) and (16), we derive our Theorem 1.
Proof of Theorem (OFM and OT connection)
Proof. Main Integration Lemma 2 states that for any fixed points zo, 2i we have
| xo u (xx) (x)) + (x1) - (100, 21)]-
Taking math expectation any plan (integration w.r.f. points To, 21 gives
Ezg.z1 ~5 h," ||u2 (2xa) (x) so)ll 2. Ezo. (V ((x1)
2 Ezo ( mn),
=Lor(V)
=CoNsT((x))
=fa (V)
where (T) does not depend on V. Hence, both minimums of OFM loss CorM (4) and of OT dual form loss Lor(Y) are achieved at the same functions.
Proof of Proposition 3 (Intractable Distance)
Proof. We recall the definitions of DIST(U, u²) (19) and FM loss CF (w) (10);
DIST(U, u²) S/' f/0 [lu((ss) 2x (2:) (poo(r))dddd
mpt (m))
(u) S.'
[(u((x)) - (x) xo) 7T "xo, xn ddrgday
.t = ()ar tan-
In the optimal plan πτ, each point zu almost surely goes to the single point V/' V Hence, in FM loss, can leave only integration over initial points zo substituting æ1 V V (2a) for fixed time t:
/ RPXR
[lue(x)) - (x1 - =00)] 2 (wo, æ dzgdr
f [lue(zz) (xo) - (o)))))(( RD
f)xo + (vo) (37)
We notice that dynamic OT vector field u² " the optimal one with potential Vr* Moreover, for any point x, (1 t)xo + +VV'((z)) generated by u", can calculate ug (æx) = u (Ix) V V (20) Zo. It is the same expression as from (37),
f RD
(37)
[lv)(x)) (VV) zo)ll p)l2po(xa)ddrg
/
[lue(1r)) uz (oo))))) xe 1)20 +/VV'(s0) (200.
R
18

---

Finally, we change the variable zo to 2t o (ro). and probability changes as po = @ Po (d, = Pp (d,), In new variables, we obtain the result
(1x)( - (2r xo)) " )drodx = |lue - (xx))) ddx.
RP XR"
Hence, the integration over time gives the desired equality
and (Far(")) = (uST(u", u²) DIST(u, u²) = (w),
Proof of Proposition 4 (Tractable Distance For OFM)
Proof. For the vector field u", we apply the formula for intractable distance from Proposition 3, i.e,
DIST(W .u") (u") L (u**) (16) CorM (V) C (V**)
According to Main Integration Lemma 2, for any plan and convex function V, have equality
Exg.=1 h,' |ux (ara) (x) =t) 2 B (V(o)) + ((x1))
2 Ezzo,a ( m1)
-oor(v)
(CoNsT((F)
Since CONST' (ir) does not depend on V, we have the same constant with v y/* and can eliminate it, i.c.,
[OFM (V) Lor(T) CONST (OFM((") = Lor(V*) CONST'(T)
$
OOFM (ern(M)) (V) COFM (V*) = 2- Lor(w) (V) 2 Lor (V²).
(38)
The right part of (38) does not depend on plan æt, thus, the left part is invariant for any plan including optimal plan n², i.e.
for (DrM((")) (TFM(4) (oFM("") DIST(U ,
Experiments details
B.1 OFM implementation
To implement our proposed approach in practice, we adopt fully-connected ICNN architectures proposed in [30, Appendix B2] (W2GN_ICNN) and [28, Appendix E1] (CPF_ICHN). To ensure the convexity, both architectures place some restrictions on the NN's weights and utilized activation functions, see the particular details in the corresponding papers. We take the advantage of their official repositories:
http:://thh..cccoooccaacctttttts //g:::///ittbb com/Cw Ht///CCWHHHuaae Huang/CP-Flox
We aggregate the hyper-parameters of our Algorithm and utilized ICNNs for different experiments in Table 3. all our experiments as the SubOpt optimizer we use LBFGS (torch, optim. LBFGS) with Kgub optimization steps and early stopping criteria based gradient norm. To find the initial point 2 (Step 5 of our Algorithm 1), we initalize SubOpt with z/.- As the Opt optimizer we adopt Adam with learning rate ir and other hyperparameters set to be default.
Similarly to OT-CFM, in some of our experiments we use con-independent initial plans to improve convergence. We construct π as follows: for independently sampled minibatches Xo, X1 of the same size B, we build the optimal discrete map and apply We stress that considering minibatch OT for our method is done exclusively speed up the training process. Theoretically, our method agnostic to initial plan π and guaranteed have an optimum in dynamic OT solution.
Minibatch.
reorder the pairs samples.
19

---

Experiment
ICNN architecture V/
K
B
Ir
K
Illustrative 2D
ICN_, R- - R. Softplus, [1024, 1024]
30K
1024
10"
5
W2 bench., dim. D
W2GN ICNN, - R. CELU, [128, 128, 64]
30K
1024
10-3
50
ALAE
W2GN_ICNN, 12 + R. CELU, [1024, 1024]
10K
128
10
10
Table 3: Hyper-parameters of our OFM solvers in different experiments
Figure 8: Performance of Flow Matching on Caussian Gaussians 2D setup.
B.2 2D Example: Comparison with Flow Matching
In this subsection, we illustrate that the restriction of the optimization domain only to optimal vector fields in FM loss is crucial for the plan independency and straightmess the obtained trajectories.
For that, we run the same setup from Section 4.1 but with vanila Flow Matching instead of OFM. The obtained trajectories and learned distributions for different initial plans are depicted in Figure
In comparison with our OFM (Figure 6), basic FM yields more curved trajectories, especially with the misleading anti-minibatch plan. The learned distributions for all plans are similar to the target.
B.3 Benchmark details
Metries. Following the authors of the benchmark [33], assess the quality of retrieved transport map T between Po and pr, we use unexplained variance percentage (UVP): C²-UVP(TT := 100 T T (po) /Var() (Vr((P)))))) (2.% For values L UVP((T = 0%, T approximates T", while for values 2 100% T's far from optimal. We also calculate the cosine similarity between ground truth directions T id and obtained directions T id, i.e.,
(T id, T id) (P()))
cos(T T id)
[1-1,
T id (po) 1|7 -idl/c<(pp)
For good approximations the cosine metric is approaching 1. We estimate L -VVP and cos metrics with 2 samples from Pu-
In the experiments, we use the exponential moving average (EMA) [45, 24] of the trained model weights. EMA creates a smoothed copy of the model whose weights are updated at each new training iteration L. + as 0eme a0.ma + (1 or 0,+1, where 0,+1 1 are the newly updated original trained weights. We calculate final metrics with 0.999.
Solvers' results for cos metric are presented in Table 4.
Solver
Solver
D=2
D
D=
D=
D=
D=
D=
D56
[S41
Dual OT solver
0.99
0.89
0.99
0.99
0.98
0.99
0.99
0.99
Reseline
0.75
0.80
0.73
0.73
0.76
0.75
0.77
0.77
OTFCPM MB [55]
0.999
0.985
0.078
0.068
0.975
0.96
0.949
0.915
RE [40]
0.87
0.75
0.65
0.67
0.72
0.70
0.70
0.70
CRR (39)
OFM Ind (Ours)
0.999
0.993
0.993
0.993
0.999
0.966
0. 0.992
0.98
OFM MB (Ours)
0.999
0.994
0.995
0.994
0.999
0,970
0.994
0.986
Table 4: cos values of solvers fitted on high digh-diimessonal benchmarks in dimensions D 2, 4, 8, 16, 32, 64, 128, 256.
best
Mutching based bolded. Metrics and buseline are [33].
20

---

Details of Solvers. Neural networks' architectures of competing Flow Matching methods and their parameters used in benchmark experiments are presented in Table 5. In this Table, "FC" stands for "fully-connected".
Selver
Anchitecior
Activattion
Hidden layers
Opinnizer
Batch
Tearning rate
Ther. per round in rounds
CFM
PCNN
RD x [0, - RDD
[128, 128,
RMSpopp
1024
10-3
[FF(40]
RD x [u. RD
PONN
[128, 128,
RMSProp
10"4
cRF
PONN
RD 1u.1]
[128, 128,
RMSProp
10-6
Table Parameters models fitted on benchmark dimensions 2, 2,4,8 16, 32, 64, 128, 256.
Time variable in (c-)RF and OT-CFM's architectures is added as one more dimensionality in input without special preprocessing. In RF and c-RF, ODE are solved via Explicit Bunge-Kutta method of order 5(4) [16] with absolute tolerance 10-4 calculated via autograd of PyTorch.
10-6 6 In OFM and c-RF, gradients over input are
Following the authors of RF [40], we run only rounds in RF. further rounds, straightness and metrics change insignificantly, while the eπror of target distribution learning still accumulates.
Our implementations of OT-CFM [55] and RF [40] are based the official repositories:
https: /ttt:://gtthb...t https: http:////gooooatt//eeaaceeeaaaae
com
Implementation of c-RF follows the RF framework with the modification of optimized NN's archi- tecture. Instead of RD x [0, 1] + RD net, we parametrize time-dependent scalar valued model RD x [0, 1] R which gradients
to be the vector field.
B.4
Unpaired Image-to-image transfer details
To conduct the experiments with high-dimensional 121 translation empowered with pretrained ALAE autoencoder, we adopt the publicly available code:
https: / /tops://githinn cot./SStollklt//iggghht-Maaa.a.a
Additional qualitative results for our method are provided in Figure 9.
COURS OPM 10
(OURS) MB
Figure 9: Unpaired 121
by OFM solver, ALAE x 1024 FFHQ latent space. The samples uncurated.
21

---

B.5 Computation time
In what follows, we provide approximate run- ning times for training our OFM and other FM-based method different experiments with hyper-parameters provided in Table 3.
ove
In the Illustrative 2D experiment, the training takes R 1.5 hours on a single 1080 ti GPU. In the Wasserstein-2 benchmark, the computa- tion time depends on the dimensionality D 2, .,, 256. Totally, all benchmark exper- iments (both with Ind and MB plan πT) take ~ 3 ALAE exper-
10-
CTM
OFN (Curs)
days on three 100 GPUs. iment, the training stage lasts for x 5 hours a single 1080 GPU.
Figure, 10: L -VVP metric depending the
elapsed training time in dimension D 32.
For better understanding of methods' behaviour over time, we depict achieved -V- metric on the benehmark (D 32) depending on elapsed training time in Figure 10. We note that the training iteration of OFM is computationally expensive, but it requires less steps to achieve the best results.
B.6 Amortization technique
In order to train our OFM, we need efficiently solve subproblem (17). As an example of more advanced technique rather than LBFGS solver, we discuss amortization trick proposed in [2].
Namely, we find approximate solution of (17) at point 2; and time I with extra MLP Ag (+))):
RD x [
(1 1
||0012 2 + (20 - (D 20)
Ag t) so arg min
(39)
2
and then sub-problem solver (LBFGS) initialized with Ag(2,, t) until convergence. We modify the training pipeline and include learning of parameters & of Ag in Algorithm
Algorithm 2 Optimal Flow Matching with Amortization
Input: Initial transport plan [I Pi), number of iterations K, batch size B, optimizer Opt, amortization optimizer AmorOpt, sub-problem optimizer SubOpt, ICNN Va MLP Ag.
1: for k= = ..... K - do
2: Sample batch {(r). xii ) of size B from plan n;;
3: Sample times batch [ 1 of size from U10, 1);
4: Calculate linear interpolation x². = = t5 )²) + t'a for all E I,
5: Compute initialization 2//ixr = (M( (t) for all E 1, B;
6: Find detached solution 2 of (17) via SubOpt initialized with aimat for all i I.B:
7:
Calculate OFM loss LOFM
ŁOFM
1 B
Σ
NN-GRAD (2
(1- -1- (z) 2b) }
} vWo(st))
Update parameters 8 via optimizer Opt step with ddoFM, db
8:
9: Calculate Amortization loss Aamor
B
Amnor
2 11:
B
1=1
10: Update parameters & via optimizer AmorOpl step with Amor.
db
11: end for
22

---

During the experiments, we did not find any improvements of the final metrics, in comparison with the original OFM with the same hyperparameters. However, this augmentation potentially can cause a shrinking of the overall training time. During training, A&(;;) :) learns to predict more and more accurate initial solution zim and, thus, reduces the required number of the expensive SubOpt steps.
C Limitations and Broader Impact
We find 3 limitations of our OFM which are to be addressed in the future research.
(a) Flow map inversion. During training, we need to compute (-))) via solving strongly convex subproblem (17). In practice, approach it by the standard gradient descent (with LBFGS optimizer), but actually there exist many improved methods to solve such conjugation problems more effectively in both the optimization [58, 26] and OT [2, 43]. This provides dozen of opportunities for improvement, but leave such advanced methods for future research.
(b) ICNNs. It is known that ICNNs may underperform compared regular neural networks [33, 31]. Thus, ICNN parametrization may limit the performance of our OFM. Fortunately, deep learning community actively study ways to improve ICNNS [10, 8, 49, 27] due to their growing popularity in various tasks [62, 37, 11]. We believe that the really expressive ICNN architectures are yet to come.
(c) Hessian inversion. We get the gradient of our OFM loss via formula from Proposition 1. There we have to invert the hessian (-), which is expensive. We point to addressing this limitation as promising avenue for future studies.
Broader impact. This paper presents work whose goal is to advance the field of Machine Learn- ing. There are many potential societal consequences of our work, none of which feel must be specifically highlighted here.
23

---

D
Check List
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In Abstract and Introduction, completely describe our contributions. For every contribution, we provide a link to the section about
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in Appendix C.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Proofs provided in Appendix A.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Experimental details are discussed in Appendix B. Code for the experiments is provided in supplementary materials. All the datasets available in public.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [YES]
Justification:Codee will be made public after the paper acceptance (now we provide it in the supplementary). Experimental details are provided in Appendix B. All the datasets are publicly available.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Experimental details are discussed in Appendix B.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Unfortunetely, running each experiment several times for statistics is too computentionaly expensive.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
24

---

Justification: In Appendix B, we mention time and resources.
9. Code Of Ethies
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification: Research conforms with NeurlPS Code of Ethics.
Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer:
[Yes]
Justification:
We discuss broader impact in Appendix C.
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The research does not need safeguards.
Licenses 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: We cite each used assets.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: New code is attached in the supplementary materials. The license will be provided after acceptance.
14. 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: The research does not engage with Crowdsourcing or Human subjects.
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: The research does not engage with Crowdsourcing or Human subjects.
25


Reviews:
[{"review_id": "jEKyai2EC1", "reviewer": "Reviewer_Bhuh", "summary": "This paper introduces the Optimal Flow Matching (OFM) algorithm, which improves upon Rectified Flow and OT-CFM by generating exact straight trajectories and recovering the optimal transport map in one iteration. OFM optimizes Flow Matching loss using vector fields and a convex function. The algorithm is implemented with Input Convex Neural Networks (ICNNs) and includes an explicit gradient formula for loss calculation. Experiments demonstrate OFM's better performance on 2D, high-dimensional benchmarks,  and in unpaired image-to-image translation. OFM achieves good results in the latent space of a pretrained autoencoder without needing ODE integration.", "questions": "- Does OFM still work in the pixel space or a higher-resolution dataset?\n- Could you provide empirical running time and OFM time complexity?\n- Does the method work for the conditional generation, which is feasible for Rectified flow?\n- OFM does not need ODE for the sampling process; however, could using an ODE solver for sampling improve the performance of OFM, which is some kind of quality-time tradeoff?", "limitations": "**References**\n\n[1] Zhu, Jun-Yan, et al. \"Unpaired image-to-image translation using cycle-consistent adversarial networks.\" Proceedings of the IEEE international conference on computer vision. 2017.\n\n[2] Xie, Shaoan, et al. \"Unpaired image-to-image translation with shortest path regularization.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "Strengths: \n- The problem is sound and essential. \n- The proposal is supported by both theoretical and empirical evidence.", "weaknesses": "- The success of the algorithm in high-dimensional tasks relies on pretrained autoencoders. To what extent does the performance of the pretrained autoencoders affect the performance of the proposal?\n- The assumption that the optimal transport map can be recovered in one iteration might not hold in all scenarios, particularly in more complex distributions.\n- How do you calculate the initial transport plan? Does OFM suffer from distorted or biased dynamic OT solutions like OT-CFM?\n- 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?\n- 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 Image-to-Image Translation tasks, such as [1] and [2]?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your feedback. Please find the answers to your questions and comments below.\n\n**(1) Pretrained autoencoders. Does OFM still work in the pixel space or a higher-resolution dataset?**\n\nIn 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.\n\nWe 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. \n\n**(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?**\n\nPlease see the previous answer for the reason of excluding this type of experiments.\n\n\n**(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.**\n\nDear 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.\n\n**(4) Does OFM suffer from distorted or biased dynamic OT solutions like OT-CFM?**\n\nNo, 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.\n\n\n**(4) How do you calculate the initial transport plan?**\n\nIn 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:\n\n (a) Independent, where we independently sample batches from both distributions,\n\n (b) Minibatch, where we independently sample batches and then rearrange them according to Discrete Optimal Transport (with $\\ell^2$ cost) between them,\n\n (c) Anti-Minibatch, where we independently sample batches and then rearrange them according to Discrete Optimal Transport based on **negative** Euclidean distance.\n\nThe same approach with plans is used in the other experiments.\n\n\n**(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?**\n\nThe 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.$\n\nFollowing 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.\n\n\n**(8) Does the method work for the conditional generation?**\n\nUnfortunately, 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.\n\n**(9) Could using an ODE solver for sampling improve the performance of OFM?**\n\nIf 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.\n\n**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.\n\n**References**:\n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qRifoQIjmP", "reviewer": "Reviewer_YPMs", "summary": "The paper goal is to learn the optimal transport map between any two distributions. The approach  derives a flow matching loss which is minimized by the velocity field of the optimal transport map.", "questions": "In essence, does the author claim that the $\\mathcal{L}^{\\pi}\\_{OFM}(\\Psi)$ loss presented in equation (16) is minimized by the velocity field of the optimal transport map for any plan $\\pi$? As stated by the author, it seems this loss is by definition Flow Matching (FM) loss constrained to convex functions. For the independent plan, i.e., $\\pi=p_0\\times p_1$ the minimizer for the FM loss is explicitly given in [1], \n$$ u_t(x) = \\frac{1}{p_t(x)}\\int u_t(x|x_1)\\frac{p_t(x|x_1)p_1(x_1)}{p_t(x)}dx_1, $$\nwhere for the linear interpolant $u_t(x|x_1) = \\frac{x_1-x}{1-t}$. 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.\n\n[1] Lipman, Yaron, et al. \"Flow matching for generative modeling.\" arXiv preprint arXiv:2210.02747 (2022).", "limitations": "-", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The paper investigate the interesting subject of the optimal transport problem.", "weaknesses": "The paper is written in an unprofessional form:\n1. Theorems and propositions are structured in a way which makes it very hard to follow.\n\n2. The author does not make any attempt to give an accessible main proof idea/proof sketch in the main paper. Specifically, it is expected to provide one for Theorem 1 which seems to be the key claim of the paper.\n\n3. The proof of Theorem 1 in the appendix is unclear and unsatisfying.\n\n4. The experiment section is rather minimal. In particular, it is not clear what Figure 3 in the Image-to-Image task presents.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, we thank you for your feedback. Please find below the answers to your questions.\n\n**(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.**\n\nYes, we claim this. This is our main theoretical result presented in Section 3.1.\n\n**(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.**\n\nIf 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.\n\n**(3) The experiment section is rather minimal. In particular, it is not clear what Figure 3 in the Image-to-Image task presents.**\n\nThe 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:\n\n1. 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*.\n\n2. 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*.  \n\nAdditionally, 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.\n\nAs 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.\n\n**(3) Theorems and propositions are structured in a way which makes it very hard to follow.**\n\nWe 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:\n\n- 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.\n\n- 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.\n\n*Could you please reply, does such a change suit you well?*\n\n**(4) The proof of Theorem 1 in the appendix is unclear and unsatisfying.**\n\nPlease 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?\n\nIn 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. \n\n**(5) Accessible main proof idea/proof sketch in the main paper.**\n\nWe 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.\n\nIn 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.\n\nIf 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.\n\n**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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ynJSkFd1Tx", "reviewer": "Reviewer_nvT5", "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.\n\nThe 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.", "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?\n- 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?\n- 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?\n\n[1] Brandon Amos. \"On amortizing convex conjugates for optimal transport\". In ICLR 2023.\n\n[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.\n- 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.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n- 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:\n- 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.\n\nExperiments:\n- 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.\n\nMinor Weaknesses:\n- 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?\n- Some of the wording can be improved, e.g. \"... via fancy formula (18) is not trivial and requires tricky integration techniques.\"[line 194]\n\n[1] Brandon Amos. \"On amortizing convex conjugates for optimal transport\". In ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your detailed feedback. Please find below the answers to your questions.\n\n**(1)  The reduction of our OFM to sampling of $t = 1$ only.**\n\nOur 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\n$$\\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.$$\n\nWithout $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.\nTo 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.\nWe 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}$.\n\n**(2) Benefits of the dynamic formulation ($t\\in [0,1]$) over the static one (in addition to the strong convexity).**\n\n\nThe 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.\n\n**(3) Comparison/relation to the dual OT optimization approach [1].**\n\n\nWe 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.\n\nOverall, 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 $:\n\\begin{eqnarray}\nA_\\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],\n\\end{eqnarray}\nand 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.\n\n**(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?**\n\n\nWe 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.\n\n\n**(5) ICNNs are considered. Does OFM also work with standard MLPs/CNNs, etc.?**\n\nFollowing 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.\n\n\n\n**(6) Quantitative results are provided for the unpaired image-to-image translation experiments.**\n\n\nFollowing 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.\n\n\n**(7) How do the illustrative 2D examples look like for standard FM training with ID, MB, anmiMB plans? Do the endpoints/trajectories differ?**\n\n\nFollowing 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.\n\n**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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0zcK00aTGm", "reviewer": "Reviewer_qt8g", "summary": "This paper aims to solve the optimal transport problem between two distributions within a flow matching framework.\nThis can be achieved by finding the best velocity model $u_t$ in a class of optimal vector fields, i.e., $u_t$ be implicitly induced by defining the $x_1=\\nabla \\Phi(x_0)$ where $\\Phi$ is a convex function.\nThe authors prove that the resulting (optimal) flow matching loss is equivalent to the dual formulation of the OT problem and regression towards the dynamic OT field.\nHowever, optimizing the optimal flow matching loss can be computationally expensive, as the velocity model is only implicitly defined.\nDespite this, the study of the relationship between the optimal flow matching problem and optimal transport can be a strong contribution.", "questions": "On the image-to-image transfer task, do you use unpaired training data? If the answer is yes, 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.", "limitations": "the authors adequately addressed the limitations", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- In general, this paper is well-written and easy to follow.\n\n- Lemma 1 and its corollaries make a strong contribution to the community. It proves that the OT problem can be equivalently solved by simply optimizing the flow matching loss in the ``optimal vector fields'' family. In fact, this is a rather surprising result for me.\n\n- In theory, the proposed optimal flow matching method can learn the OT map after one round of training, regardless of the reference distribution.\n\n- The experiments are well-designed and convincing. The 2D experiment with different reference distributions (especially the anti-minibatch) is illustrative.", "weaknesses": "- The practical implementation of OFM is rather computationally expensive. It would be fine to add a table to directly compare the training times of different methods, or a figure showing the convergence as a function of computation time. I don't think this will lessen the contribution of this paper.\n\n- As the authors claim that OFM learns straighter trajectories, comparing OFM to other baselines on standard generation benchmarks, e.g., CIFAR10, would be better.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear reviewer, thank you for your feedback and kind words. Please find below the answers to your questions.\n\n**(1) Convergence time. Plots of convergence.**\n\nFollowing 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.\n\nWe will include this requested plot to the appendix of the final version of the paper.\n\n\n**(2) Comparing OFM to other baselines on standard generation benchmarks.**\n\nIn 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.\n\n\n\n**(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.**\n\n\n\nYes, 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)$. \n\n\n\n**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.\n\n**References.**\n\n[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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their feedback and time. We appreciate that reviewers acknowledge: importance and soundness of the chosen problem (Bhuh, YPMs), our theoretical contribution connecting FM and OT methods (Bhuh, nvT5, qt8g), and experimental verification of the theoretical part (qt8g, Bhuh). Please find the answers to your questions in the respective posts below your reviews.

We additionally calculated FID for our considered Unpaired Image-to-Image problem to expand the experiments and make them more representative. The obtained results exhibit that our OFM approach quantitatively outperforms the considered flow-based baselines.

We attach the **PDF file** to this post with requested experiments containing:

- **(All Reviewers)** Table 1, FID metric for unpaired Image-to-Image real data problem. The considered methods are OFM, RF, c-RF and OT-CFM.

- **(Reviewer nvT5)** Figure 1, OFM setting $t \equiv 1$. The title "gt" stands for the source and target distributions. Titles "mb 64" and "mb 16" stand for minibatch plan with batchsize $64$ and $16$, respectively.

- **(Reviewer nvT5)** Figure 2, default FM with independent ("idp"), anti-minibatch ("anti mb") and minibatch ("mb") plans. The title "gt" stands for the source and target distributions.

- **(Reviewer qt8g)** Figure 3, convergence of L2-UVP loss depending on training time for OFM, FM, c-RF and OT-CFM methods.

- **(Reviewer nvT5)** Figure 4, OFM with and without amortization. The graph depicts L2-UVP metric during training.

Meta-review: This paper shows that the minimizer of the flow matching loss, where the velocity field is constrained to the form of W2-optimal velocity field (i.e., the velocity of the displacement map), is the actual Optimal Transport map. Based on this observation the author model the learnable velocity via the ICNN (convex functions) architecture and demonstrate the efficacy of their approach on a few 2D and image datasets, and comparing to other flow matching alternatives looking to find the Optimal Transport map. 

Pros:
+ Reviewers found the main theoretical result in the paper an important one. In particular the use of Lemma 1 to prove main theorem. 
+ New and interesting connection between Flow Matching and Optimal Transport. 
+ A sufficient set of experiments supporting the theoretical findings. 

 
Main reviewers' concerns included:
- Readability and lacking exposition of the main theoretical results > authors have addressed this and suggested a fix in final version. 
- Unclear relation to SubOpt > authors clarified their method requires optimizing over all times $t\in [0,1]$.
- Scalability issues of ICNN, computational overhead of the method, convergence times > authors will address this in limitation and/or experiments and/or appendix section in the revised version. 

Overall this paper passes the bar for a NeurIPS paper, and we ask the authors to incorporate the promised revisions in their camera ready version.
Decision: Accept (poster)