Optimal Flow Matching: Learning Straight Trajectories in Just One Step

Authors: nikita  kornilov
petr  mokrov
alexander  gasnikov
alexander  korotin
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 15:06:10

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

Main paper: 
 introduction: 
 Recent success in generative modeling [41, 17, 9] is mostly driven by Flow Matching (FM) [38] models. These models move a known distribution to a 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, it 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 is 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 a heuristic does not guarantee straight paths because of minibatch OT biases, see, e.g., [55, Figure 1, right] for the practical illustration.Contributions. In this paper, we fix the above-mentioned problems of the straightening methods. We propose a novel Optimal Flow Matching (OFM) approach ( §3) that after a 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., it solves the Benamou-Brenier problem (Figure 1 ). We demonstrate the potential of OFM in the series of experiments and benchmarks ( §4).Optimal Flow Matching (just one step, ours)Figure 1 : Our Optimal Flow Matching (OFM). For any initial transport plan π between p 0 and p 1 , OFM obtains exactly straight trajectories (in just a single FM loss minimization) which carry out the OT displacement for the quadratic cost function.The main idea of 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, one can optionally use minibatch OT or any other transport plan as the input, and this is completely theoretically justified. 
background and related works: 
 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).Notations. For vectors x, y ∈ R D , we denote the inner product by ⟨x, y⟩ and the corresponding ℓ 2 norm by ∥x∥ := ⟨x, x⟩. We use P 2,ac (R D ) to refer to the set of absolute continuous probability distributions with the finite second moment. For vector x ∈ R D and distribution p ∈ P 2,ac (R D ), notation x ∼ p means that x is sampled from p. For the push-forward operator, we use symbol #.Monge's and Kantorovich's formulations. Consider two distributions p 0 , p 1 ∈ P 2,ac (R D ) and a cost function c : R D × R D → R. Monge's Optimal Transport formulation is given by inf T #p0=p1 R D c(x 0 , T (x 0 ))p 0 (x 0 )dx 0 ,where the infimum is taken over measurable functions T : R D → R D which satisfy the masspreserving constraint T #p 0 = p 1 . Such functions are called transport maps. If there exists a transport map T * that achieves the infimum, then it is called the optimal transport map.Since the optimal transport map T * in Monge's formulation may not exist, there is Kantorovich's relaxation for problem (1) which addresses this issue. Consider the set of transport plans Π(p 0 , p 1 ), i.e., the set of joint distributions on R D × R D which marginals are equal to p 0 and p 1 , respectively. Kantorovich's Optimal Transport formulation isEQUATIONWith mild assumptions on p 0 , p 1 , the infimum is always achieved (possibly not uniquely). An optimal plan π * ∈ Π(p 0 , p 1 ) is called an optimal transport plan. If optimal π * has the form [id, T * ]#p 0 , then T * is the solution of Monge's formulation (1) .In our paper, we mostly consider the quadratic cost function c(x 0 , x 1 ) = ∥x0-x1∥ 2 2Dual formulation. For the quadratic cost, problem (3) has the equivalent dual form [60] :EQUATIONwhere the minimum is taken over convex functions Ψ(x) : R D → R. Here Ψ(x1 ) := sup x0∈R D [⟨x 0 , x 1 ⟩ -Ψ(x 0 )] is the convex (Fenchel) conjugate function of Ψ. It is also convex.The term CONST(p 0 , p 1 ) does not depend on Ψ. Therefore, the minimization (3) over transport plans π is equivalent to the minimization of L OT (Ψ) from (4) over convex functions Ψ. Moreover, the optimal transport map T * can be expressed via an optimal Ψ * (the Brenier potential [60] ), namely,EQUATION)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(t, •) ≡ u t (•) : [0, 1] × R D → R Dbe a vector field and {{z t } t∈[0,1] } be the set of random trajectories such that for each trajectory {z t } t∈[0,1] the starting point z 0 is sampled from p 0 and z t satisfies the differential equation:EQUATIONIn other words, the trajectory {z t } t∈[0,1] is defined by its initial point z 0 ∼ p 0 and goes along the speed vector u t (z t ). Under mild assumptions on u, for each initial z 0 , the trajectory is unique.Let ϕ u (t, •) ≡ ϕ u t (•) : [0, 1] × R D → R D denote the flow map, i.e., it is the function that maps the initial z 0 to its position at moment of time t according to the ODE (6), i.e.,EQUATIONIf initial points z 0 of trajectories are distributed according to p 0 , then (6) defines a distribution p t of z t at time t, which can be expressed via with the push-forward operator, i.e., p u t := ϕ u t #p 0 . Benamou-Brenier problem. Dynamic OT is the following minimization problem:EQUATIONs.t. ϕ u 1 #p 0 = p 1 .In (8) , we look for the vector fields u that define the flows which start at p 0 and end at p 1 . 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 * = ∇Ψ * and the dynamic OT solution u * . Namely, for every initial point z 0 , the vector field u * defines a linear trajectory {z t } t∈[0,1] :EQUATIONThere exist a variety of continuous OT solvers [21, 52, 54, 43, 19, 14, 59, 15, 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, we focus only on the most relevant ones, called the ICNN-based solvers [54, 33, 43, 2] . These solvers directly minimize objective L OT from (4) parametrizing a class of convex functions with convex in input neural networks called ICNNs [3] (for more details, see "Parametrization of Ψ" in §3.2). Solvers details may differ, but the main idea remains the same. To calculate the conjugate function Ψ(x 1 ) at the point x 1 , they solve the convex optimization problem from conjugate definition. Envelope Theorem [1] allows obtaining closed-form formula for the gradient of the loss.In this section, we recall popular approaches [40, 39, 48] to find fields u which transport a given probability distribution p 0 to a target p 1 and their relation to OT.To find such a field, one samples points x 0 , x 1 from a transport plan π ∈ Π(p 0 , p 1 ), e.g., the independent plan p 0 × p 1 . The vector field u is encouraged to follow the direction x 1 -x 0 of the linear interpolation x t = (1 -t)x 0 + tx 1 at any moment t ∈ [0, 1], i.e., one solves:EQUATIONFigure 2 : Flow Matching (FM) obtains a vector field u moving p 0 to p 1 . FM typically operates with the independent transport plan π = p 0 × p 1 .We denote the solution of (10) and the flow map (7) by u π and ϕ π , respectively. The concept of FM is depicted in Figure 2 .The intuition of this procedure is as follows: linear interpolation x t = (1 -t)x 0 + tx 1 is an intuitive way to move p 0 to p 1 , but it requires knowing x 1 . By fitting u with the direction x 1 -x 0 , one yields the vector field that can construct this interpolation without any information about x 1 .The set of trajectories {{z t } t∈[0,1] } generated by u π t (with z 0 ∼ p 0 ) has a useful property: the flow map ϕ π 1 transforms distribution p 0 to distribution p 1 for any initial transport plan π. Moreover, marginal distribution p t = ϕ π t #p 0 is equal to the distribution of linear interpolation x t = (1 -t)x 0 + tx 1 for any t and x 0 , x 1 ∼ π. This feature is called the marginal preserving property.To push point x 0 according to learned u, one needs to integrate ODE (6) via numerical solvers. The vector fields with straight (or nearly straight) paths incur much smaller time-discretization 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 p 0 × p 1 . The two most popular approaches to choose better plans are Optimal Transport Conditional Flow Matching [48, 55] and Rectified Flow [40] . If one uses the OT plan π * as the initial plan for FM, then it returns the Brenier's vector field u * , which generates exactly straight trajectories (9) . However, typically, the true OT plan π * is not available. In such a case, in order to achieve some level of straightness in the learned trajectories, a natural idea is to take the initial plan π to be close to the optimal π * . 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 p 0 and p 1 . 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 .Flow MatchingThe main drawback of OT-CFM is that it recovers only biased dynamic OT solution. In order to converge to the true transport plan the batch size should be large [6] , while with a 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.In [40] , the authors propose an iterative approach to refine the plan π, straightening the trajectories more and more with each iteration. Formally, Flow Matching procedure denoted by FM takes the transport plan π as input and returns an optimal flow map via solving (10) :EQUATION)Flow Matching (step 1)Flow Matching (step 2)...Figure 4 : Rectified Flow iteratively applies FM to 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 on K-th iteration has the following update ruleEQUATIONwhere ϕ K , π K denote flow map and transport plan on K-th iteration, respectively.With each new FM iteration, the generated trajectories {{z t } t∈[0,1] } K provably become more and more straight, i.e., error in approximationz K t ≈ (1 -t)z K 0 + tz K 1 , ∀t ∈ [0, 1]decreases as 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 c the flow map ϕ π 1 from Flow Matching yields lower or equal transport cost than initial transport plan π:EQUATIONIntuitively, 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 x 0 and x 1 can. With each iteration of RF (12) , transport costs for all convex cost functions do not increase, but, for a given cost function, convergence to its own 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 c. In this modification, called c-Rectified Flow (c-RF), the authors slightly change the FM training objective and restrict the optimization domain only to potential vector fields u t (•) = ∇c(∇f t (•)), where f t (•) : R D → R is an arbitrary time-dependent scalar valued function and c is the convex conjugate of the cost function c. In case of the quadratic cost function, the training objective remains the same, and the vector field u t is set as the simple gradient ∇f t (•) of the scalar valued function f t .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., ϕ K 1 #p 0 ̸ = p 1 ), and this inexactness only grows with iterations. In addition, training of (c-)RF becomes non-simulation free after the first iteration, since to calculate the plan π K+1 = [id, ϕ K+1 ]#p 0 it has to integrate ODE. 
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 ( §3.3). All our proofs are located in Appendix A.We want to design a method of moving distribution p 0 to p 1 via exactly straight trajectories. Namely, we aim to obtain straight paths from the solution of the dynamic OT (8) . Moreover, we 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 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 if it generates linear trajectories {{z t } t∈[0,1] } such that there exist a convex function Ψ : R D → R, which for any path {z t } t∈[0,1] pushes the initial point z 0 to the final one as z 1 = ∇Ψ(z 0 ), i.e.,EQUATIONEquation ( 14) does not provide a closed formula for u Ψ as it depends on z 0 . The explicit formula is constructed as follows: for a time t ∈ [0, 1] and point x t , we can find a trajectory {z t } t∈[0,1] s.t.EQUATIONFigure 5 : An Optimal Vector Field: a vector field u Ψ with straight paths is parametrized by a gradient of a convex function Ψ.and recover the initial point z 0 . We postpone the solution of this problem to §3.2. For now, we define the inverse of flow map (7) as (ϕ Ψ t ) -1 (x t ) := z 0 and the vector fieldu Ψ t (x t ) := ∇Ψ(z 0 ) -z 0 = ∇Ψ((ϕ Ψ t ) -1 (x t )) -(ϕ Ψ t ) -1 (x t ), which generates ODE (14) , i.e., dz t = u Ψ t (z t )dt. The concept of optimal vector fields is depicted on Figure 5 .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 Ψ * (5).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 u Ψ into FM loss from (10) and define our Optimal Flow Matching loss:EQUATIONx t = (1 -t)x 0 + tx 1 .Our Theorem 1 states that OFM solves the dynamic OT via single FM minimization for any initial π. Theorem 1 (OFM and OT connection). Consider two distributions p 0 , p 1 ∈ P ac,2 (R D ) and any transport plan π ∈ Π(p 0 , p 1 ) between them. Then, the dual Optimal Transport loss L OT (4) and Optimal Flow Matching loss L π OF M (16) have the same minimizers, i.e., arg minconvex Ψ L π OF M (Ψ) = arg min convex Ψ L OT (Ψ).In this subsection, we explain the details of optimization of our Optimal Flow Matching loss (16) .Parametrization of Ψ. In practice, we parametrize the class of convex functions with Input Convex Neural Networks (ICNNs) [3] Ψ θ and parameters θ. These are scalar-valued neural networks built in such a way that the network is convex in its input. They consist of fully-connected or convolution blocks, some weights of which are set to be non-negative in 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 of OFM loss (16) . Proposition 1 (Explicit Loss Gradient Formula). The gradient of L π OF M can be calculated asz 0 = NO-GRAD (ϕ Ψ θ t ) -1 (x t ) , dL π OF M dθ := d dθ E t;x0,x1∼π NO-GRAD 2 t∇ 2 Ψ θ (z 0 ) + (1 -t)I -1 (x 0 -z 0 ) t , ∇Ψ θ (z 0 ) ,where variables under NO-GRAD remain constants during differentiation.Flow map inversion. In order to find the initial point z 0 = (ϕ Ψ t ) -1 (x t ), we note that (15)x t = (1 -t)z 0 + t∇Ψ(z 0 ) is equivalent to ∇ (1 -t) 2 ∥ • ∥ 2 + tΨ(•) -⟨x t , •⟩ (z 0 ) = 0.The function under gradient operator ∇ has minimum at the required point z 0 , since at z 0 the gradient of it equals 0. If t < 1 the function is at least (1 -t)-strongly convex, and the minimum is unique. The case t = 1 is negligible in practice, since it has zero probability to appear during training.We can reduce the problem of inversion to the following minimization subproblemEQUATION)Optimization subproblem ( 17) is at least (1 -t)-strongly convex and can be effectively solved for any given point x t (in comparison with typical non-convex optimization tasks).Algorithm. The Optimal Flow Matching pseudocode is presented in listing 1. We estimate math expectation over plan π and time t with uniform distribution on [0, 1] via unbiased Monte Carlo.Input: Initial transport plan π ∈ Π(p 0 , p 1 ), number of iterations K, batch size B, optimizer Opt, sub-problem optimizer SubOpt, ICNN Ψ θ 1: for k = 0, . . . , K -1 doEQUATIONCalculate linear interpolationx i t i = (1 -t i )x i 0 + t i x i 1 for all i ∈ 1, B; 5:Find the initial points z i 0 via solving the convex problem with SubOpt:z i 0 = NO-GRAD arg min z i 0 (1 -t i ) 2 ∥z i 0 ∥ 2 + t i Ψ θ (z i 0 ) -⟨x i t i , z i 0 ⟩ ; 6: Calculate loss LOF M LOF M = 1 B B i=1 NO-GRAD 2 t i ∇ 2 Ψ θ (z i 0 ) + (1 -t i )I -1 (x i 0 -z i 0 ) t i , ∇Ψ θ (z i 0 ) ; 7:Update parameters θ via optimizer Opt step with d LOF M dθ ; 8: end forIn 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 a novel theoretical bridge between OT and FM. Such a direct connection can lead to the adoption of the strengths of both methods and a 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 as 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 1.Rectified Flow [40, 39] . In Rectified Flows [40] , the authors iteratively apply Flow Matching to refine the obtained trajectories. However, in each iteration, RF accumulates error since one may not 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 to the OT plan for any cost function c, 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. In EOT and BM, instead of deterministic ODE and flows, 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 KL 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).In this subsection, we provide the OFM's theoretical properties, which give an intuition for understanding of its main working principles and behavior. Proposition 2 (Simplified OFM Loss). We can simplify (16) to a more suitable form:EQUATIONThe simplified form (18) shows that OFM loss actually measures how well Ψ restores initial points x 0 of linear interpolations depending on future point x t and time t.Generative properties of OFM. In this paragraph, we provide another view on our OFM approach.In our OFM, we aim to construct a vector field u which is as close to the dynamic OT field u * as possible. We can use the least square regression to measure the distance between them:DIST(u, u * ) :=EQUATION)Proposition 3 (Intractable Distance). The distance DIST(u, u * ) between an arbitrary vector field u and OT field u * equals to the FM loss from (10) with the optimal plan π * , i.e.,DIST(u, u * ) = L π * F M (u) -L π * F M (u * ) =0 .We can not minimize intractable DIST(u, u * ) since the optimal plan π * is unknown. In OT-CFM [55] , authors heuristically approximate π * in L π * F M (u), but obtain biased solution. Surprisingly, for the optimal vector fields, the distance can be calculated explicitly via any known plan π. Proposition 4 (Tractable Distance For OFM). The distance DIST(u Ψ , u Ψ * ) between an optimal vector field u Ψ generated by a convex function Ψ and the vector field u Ψ * with the Brenier potential Ψ * can be evaluated directly via OFM loss (16) and any plan π:EQUATIONIn (20) , the first term is our tractable OFM loss, and the second term does not depend on Ψ. 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. 
experimental illustrations: 
 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, we apply our approach for solving high-dimensional unpaired image-to-image translation in the latent space of pretrained ALAE autoencoder ( §4.3). The PyTorch implementation of our method is publicly available at https://github.com/Jhomanik/Optimal-Flow-MatchingThe technical details of our experiments (architectures, hyperparameters) are in the Appendix B.In this subsection, we illustrate the proof-of-concept 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 a standard Gaussian p 0 = N (0, I) and a Mixture of eight Gaussians p 1 depicted in the Figure 6a . We consider different stochastic plans π: independent plan p 0 × p 1 (Figure 6b ), minibatch and antiminibatch (Figures 6c, 6d ) discrete OT (quadratic cost) with batch size B mb = 64. In the antiminibatch case, we compose the pairs of source and target points by solving discrete OT with minus quadratic cost -∥x -y∥ 2 2 . 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 to the original FM (10) , and show that, in comparison with our OFM, the resulting paths obtained by FM considerably depend on the plan. In this subsection, we quantitatively compare our OFM and other methods testing their ability to solve OT. We run our OFM, FM based methods and OT solvers on OT Benchmark [33] . The authors provide high-dimensional continuous distributions p 0 , p 1 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 L 2 -UVP(T ) := 100 • ∥T -T * ∥ 2 L 2 (p0) /Var(p 1 )% [33] . It directly computes the normalized squared error between OT map T * and learned map T .Competitors. We evaluate Conditional Flow Matching (OT-CFM), Rectified Flow (RF), c-Rectified 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 L OT (4) by parametrizing Ψ with ICNNs and calculating Ψ(x 1 ) via solving a convex optimization subproblem. The latter is similar to our inversion (17) . Additionally, in [2] , the authors use MLPs to parametrize Ψ, and we include these results as well. Following [33] , we also provide results for a linear OT map (baseline) 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 B mb = 64. Results. Among FM-based methods, our OFM with any plan demonstrates the best results. For all plans, OFM convergences to close final solutions and metrics. Minibatch plan provides a little bit better results, especially in high dimensions. In theory, the OFM results for any plan π must be similar. However, in stochastic optimization, plans with large variance yield convergence to slightly worse solutions.MLP-based OT solver usually beats our OFM, since MLPs do not have ICNNs' limitations in practice. However, usage of MLP is an empirical trick and is not completely justified. We also run OFM with MLP instead ICNN, and, unfortunately, the method fails to converge.RF demonstrates worse performance than even linear baseline, but it is ok since it is not designed to solve W 2 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.Another task that involves learning a 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-trained ALAE autoencoder [46] on 1024 × 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 as the source and target distributions p 0 and p 1 . At the inference stage, we take a new (unseen) adult 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 The batch size for minibatch OT methods (⌊OFM, MB⌉, ⌊OT-CFM⌉) is B mb = 128. Our OFM converges to 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-based 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 000000D730324P540002).In this section, we place the proofs of all our results from the main manuscript and some auxiliary results. We reorder the proofs of Prop. 1 and Prop. 2 since the former is based on the latter.Note that in all of our theoretical derivations, if not stated explicitly, we assume the differentiability of convex potential Ψ at given points z 0 , x t 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.t Lebesgue measure [50] . Therefore, since we consider absolutely continuous reference distributions p 0 , p 1 (see §2.1), the differentiability of Ψ at the considered points also holds almost surely.Proof of Proposition 2 (Simplified OFM Loss)Proof. By definition L π OF M (Ψ) equals toEQUATION)For fixed points x 0 , x 1 and time t in integrand, we find a point z 0 = (ϕ Ψ t ) -1 (x t ) such that in moment t ∈ [0, 1] it is transported to point x t = (1 -t)x 0 + tx 1 . This point z 0 satisfies equalityx t = t∇Ψ(z 0 ) + (1 -t)z 0 .We define the vector field u Ψ t asu Ψ t (x t ) = ∇Ψ(z 0 ) -z 0 = x t -z 0 t .Putting u Ψ t (x t ) in the integrand of ( 21), we obtain simplified integrand∥x 1 -x 0 -u Ψ t (x t )∥ 2 = x 1 -x 0 - x t -z 0 t 2 = 1 t 2 ∥tx 1 -tx 0 -((1 -t)x 0 + tx 1 ) + z 0 ∥ 2 = 1 t 2 ∥z 0 -x 0 ∥ 2 = (ϕ Ψ t ) -1 (x t ) -x 0 t 2 .Proof of Proposition 1 (Explicit Loss Gradient Formula)Proof. Point z 0 = (ϕ Ψ θ t ) -1 (x t ) now depends on parameters θ. We differentiate the integrand from the simplified OFM loss (18) for fixed points x 0 , x 1 and time t, i.e.,EQUATIONFor point z 0 , the equation ( 27) holds true:EQUATIONWe differentiate (23) w.r.t. θ and obtain0 = (1 -t) dz 0 dθ + t∇ 2 Ψ θ (z 0 ) dz 0 dθ + t ∂∇Ψ θ ∂θ (z 0 ) ⇒ dz 0 dθ = -t∇ 2 Ψ θ (z 0 ) + (1 -t)I -1 • t ∂∇Ψ θ ∂θ (z 0 ).Therefore, we haveEQUATION)Now the differentiation over θ is located only in the right part of (24) in the term ∂∇Ψ θ ∂θ . Hence, point z 0 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 t.The following two Lemmas are used to prove our main theoretical result, Theorem 1. Lemma 1 (Properties of convex functions and their conjugates). Let Ψ : R D → R be a convex function; x 0 , x 1 ∈ R D . Let Ψ and Ψ be differentiable at x 0 and x 1 correspondingly. Then the following statements are equivalent:(i) x 1 = ∇Ψ(x 0 ); (ii) x 0 = arg max z∈R D ⟨x 1 , z⟩ -Ψ(z)};(iii) Fenchel-Young's equality: Ψ(x 0 ) + Ψ(x 1 ) = ⟨x 1 , x 0 ⟩;(iv) x 0 = ∇Ψ(x 1 ) ; (v) x 1 = arg max z∈R D ⟨z, x 0 ⟩ -Ψ(z)} ;Proof. The lemma is a 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 x 0 , x 1 ∈ R D and a convex function Ψ, the following equality holds true:EQUATIONwherex t = tx 0 + (1 -t)x 1 .Proof. Following Proposition 2, we use the simplified loss form, i.e.,EQUATION)where z 0 = z 0 (t) = (ϕ Ψ t ) -1 (x t ) satisfies the equality:EQUATIONNext, we substitute (26) into rhs of (25) integrate w.r.t. time t from 0 excluding to 1 excluding (This exclusion does not change the integral):EQUATIONTo further simplify (28) we need some preliminary work. Following (27) we note:EQUATION)Changing in (28) time variable t to s = t 1-t , ds = dt (1-t) 2 and substitution of (30) yield:EQUATIONNow we use properties of conjugate functions (Lemma (1)):Ψ(∇Ψ(x 1 )) (iv) + (iii) = ⟨∇Ψ(x 1 ), x 1 ⟩ -Ψ(x 1 ), ∇Ψ(∇Ψ(x 1 )) (iv) + (i) = x 1 .This allows us to simplify (36):(36) = ⟨x 1 , ∇Ψ(x 1 )⟩ -⟨x 1 , x 0 ⟩ + 2(Ψ(x 0 ) + Ψ(x 1 ) -⟨∇Ψ(x 1 ), x 1 ⟩) + ⟨(∇Ψ(x 1 ), x 1 ⟩ -⟨x 0 , ∇Ψ(x 0 )⟩ + ⟨x 0 , ∇Ψ(x 0 )⟩ -⟨x 0 , x 1 ⟩ = 2[Ψ(x 0 ) + Ψ(x 1 ) -⟨x 0 , x 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 1 (OFM and OT connection)Proof. Main Integration Lemma 2 states that for any fixed points x 0 , x 1 we have1 0 ∥x 1 -x 0 -u Ψ t (x t )∥ 2 dt = 2[Ψ(x 0 ) + Ψ(x 1 ) -⟨x 0 , x 1 ⟩].Taking math expectation over any plan π (integration w.r.t. points x 0 , x 1 ∼ π ) givesEQUATION0 ∥u Ψ t (x t ) -(x 1 -x 0 )∥ 2 dt =L π OF M (Ψ) = 2 • E x0,x1∼π [Ψ(x 0 ) + Ψ(x 1 )] =L OT (Ψ) -2 • E x0,x1∼π [⟨x 0 , x 1 ⟩] =:CONST ′ (π), where CONST ′ (π) does not depend on Ψ. Hence, both minimums of OFM loss L π OF M (Ψ) and of OT dual form loss L OT (Ψ) are achieved at the same functions.Proof. We recall the definitions of DIST(u, u * ) (19) and FM loss L π * F M (u) (10) :DIST(u, u * ) = 1 0 R D ∥u t (x t ) -u * t (x t )∥ 2 ϕ * t #p 0 (x t ) =p * t (xt) dx t dt, L π * F M (u) = 1 0      R D ×R D ∥u t (x t ) -(x 1 -x 0 )∥ 2 π * (x 0 , x 1 )dx 0 dx 1      , x t = (1 -t)x 0 + tx 1 .In the optimal plan π * , each point x 0 almost surely goes to the single point ∇Ψ * (x 0 ). Hence, in FM loss, we can leave only integration over initial points x 0 substituting x 1 = ∇Ψ * (x 0 ) for fixed time t:EQUATIONWe notice that dynamic OT vector field u * = u Ψ * is the optimal one with potential Ψ * . Moreover, for any point x t = (1 -t)x 0 + t∇Ψ * (x 0 ) generated by u * , we can calculate u * t (x t ) = u Ψ * t (x t ) = ∇Ψ * (x 0 ) -x 0 . It is the same expression as from (37), i.e., (37) In what follows, we provide approximate running times for training our OFM and other FM-based method in different experiments with hyper-parameters provided in Table 3 .= R D ∥u t (x t ) -(∇Ψ * (x 0 ) -x 0 )∥ 2 p 0 (x 0 )dx 0 = R D ∥u t (x t ) -u * t (x t )∥ 2 p 0 (x 0 )dx 0 , x t = (1 -t)x 0 + t∇Ψ * (x 0 ).In the Illustrative 2D experiment, the training takes ≈ 1.5 hours on a single 1080 ti GPU. In the Wasserstein-2 benchmark, the computation time depends on the dimensionality D = 2, 4, . . . , 256. Totally, all the benchmark experiments (both with Ind and MB plan π) take ≈ 3 days on three A100 GPUs. In the ALAE experiment, the training stage lasts for ≈ 5 hours on a single 1080 ti GPU.For better understanding of methods' behaviour over time, we depict achieved L 2 -UVP metric on the benchmark (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.In order to train our OFM, we need to 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 an approximate solution of (17) at point x t and time t with an extra MLP A ϕ (•, •) :EQUATIONand then run sub-problem solver (LBFGS) initialized with A ϕ (x t , t) until convergence. We modify the training pipeline and include learning of parameters ϕ of A ϕ in Algorithm 2.Input: Initial transport plan π ∈ Π(p 0 , p 1 ), number of iterations K, batch size B, optimizer Opt, amortization optimizer AmorOpt, sub-problem optimizer SubOpt, ICNN Ψ θ , MLP A ϕ 1: for k = 0, . . . , K -1 do Calculate linear interpolation x i t i = (1 -t i )x i 0 + t i x i 1 for all i ∈ 1, B;Compute initialization z i init = A ϕ (x i t i , t i ) for all i ∈ 1, B;6:Find detached solution z i 0 of (17) via SubOpt initialized with z i init for all i ∈ 1, B;7: Calculate OFM loss LOF M LOF M = 1 B B i=1 NO-GRAD 2 t i ∇ 2 Ψ θ (z i 0 ) + (1 -t i )I -1 (x i 0 -z i 0 ) t i, ∇Ψ θ (z i 0 ) ; 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 (ϕ Ψ θ t ) -1 (•) via solving strongly convex subproblem (17) . In practice, we 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 a dozen of opportunities for improvement, but leave such advanced methods for future research.(b) ICNNs. It is known that ICNNs may underperform compared to 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 ∇ 2 Ψ(•), which is expensive. We point to addressing this limitation as a promising avenue for future studies.Broader impact. This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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?
- The assumption that the optimal transport map can be recovered in one iteration might not hold in all scenarios, particularly in more complex distributions.
- How do you calculate the initial transport plan? Does OFM suffer from distorted or biased dynamic OT solutions like OT-CFM?
- 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?
- 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]?
questions: questions: - Does OFM still work in the pixel space or a higher-resolution dataset?
- Could you provide empirical running time and OFM time complexity?
- Does the method work for the conditional generation, which is feasible for Rectified flow?
- 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: limitations: **References**

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

[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.
paper_summary: paper_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.
strength_weakness: strength_weakness: The paper is written in an unprofessional form:
1. Theorems and propositions are structured in a way which makes it very hard to follow.

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

3. The proof of Theorem 1 in the appendix is unclear and unsatisfying.

4. The experiment section is rather minimal. In particular, it is not clear what Figure 3 in the Image-to-Image task presents.
questions: 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], 
$$ 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, $$
where 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.

[1] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).
limitations: limitations: -
paper_summary: paper_summary: This work proposes Optimal Flow Matching (OFM), which restricts standard Flow Matching to specific vector fields that yield straight trajectories by design. This is implemented by considering vector fields such that a convex function exists whose gradient pushes the initial to the final point.

The authors propose to solve this optimization problem with training that involves solving a convex optimization problem in each iteration to obtain the inverse of the flow map. Then, it is shown that this procedure has the same minimizer as the one minimizing the dual OT loss. The authors empirically validate OFM quantitatively on the Wasserstein-2 benchmark and qualitatively on FFHQ.
strength_weakness: strength_weakness: Missing comparison and contextualization of related work:
- In [1], it was proposed to approximate the computation of the conjugate with amortized optimization and solve this with Adam or L-BFGS. The proposed *SubOpt* in this work (solved with Adam/L-BFGS) is equivalent to the convex conjugate optimization when $t=1$.  A discussion between these two optimization schemes should be included. Specifically, how does OFM compare to [1] when only sampling $t=1$? What is gained by considering $t\sim \mathcal{U}[0,1]$? Additionally, the Wasserstein-2 benchmark results from [1] are not considered in the paper. This is the most related method and outperforms the method proposed in the paper on all dimensions. These results should be included as well as a comparative discussion.

Experiments:
- No quantitative results are provided for the unpaired image-to-image translation experiments. It is hard to generalize the results over a few qualitative results. Thus, including some commonly used quantitative metrics like FID or the ones used in the Wasserstein-2 benchmark on CelebA64, as done in [1], would be very beneficial.

Minor Weaknesses:
- Inconsistency on the optimizer used to solve *SubOpt*. In Appendix C, it is stated that "in all our experiments as the subOpt optimizer we use LBFGS"[line 557] while in Appendix B, it says "(with Adam optimizer)"[line 534]. Which one is actually used? How do Adam and LBFGS compare?
- Some of the wording can be improved, e.g. "... via fancy formula (18) is not trivial and requires tricky integration techniques."[line 194]

[1] Brandon Amos. "On amortizing convex conjugates for optimal transport". In ICLR 2023.
questions: questions: - The goal of OFM is to learn a time-independent one-step model that translates $x_0$ to $x_1$. Then, during inference, no ODE is solved. As mentioned above, some more motivation could improve clarity here. Specifically, explaining why it is beneficial to formulate this in the FM framework considering $t\sim \mathcal{U}[0,1]$. Would the proposed method also work when only sampling $t=1$? The authors only mention that the subproblem (20) is $(1-t)$-strongly convex, which gives a benefit compared to $t=1$. Are there other benefits of the proposed formulation?
- In Figure 2, it is shown that OFM learns the same solution for different choices of couplings. In [2], it was shown that only the OT coupling is preserved when using standard Flow Matching. What do these illustrative examples look like for standard FM training? I assume they would learn different trajectories, but do the learned endpoints differ?
- In this work, only ICNNs are considered. As shown in previous works, even when wanting to learn a convex function, MLPs can also work well, i.e., in [1], they work better compared to ICNNs. Does OFM also work with standard MLPs/CNNs/etc? How is the comparative performance?

[1] Brandon Amos. "On amortizing convex conjugates for optimal transport". In ICLR 2023.

[2] Valentin De Bortoli and Guan-Horng Liu and Tianrong Chen and Evangelos A. Theodorou and Weilie Nie. "Augmented Bridge Matching". In Arxiv 2023.
limitations: limitations: - The limitations of the flow map inversion, the Hessian computation, and ICNNs are laid out in the paper, and details for the experiments' computational time are provided.
- While these computational times give some insight, detailed information on time and space complexity with respect to the input dimension would be even more beneficial. Additionally, comparing the runtime of OFM and standard FM training would provide a more illustrative comparison of raw computational time.
paper_summary: paper_summary: This paper aims to solve the optimal transport problem between two distributions within a flow matching framework.
This 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.
The 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.
However, optimizing the optimal flow matching loss can be computationally expensive, as the velocity model is only implicitly defined.
Despite this, the study of the relationship between the optimal flow matching problem and optimal transport can be a strong contribution.
strength_weakness: strength_weakness: - 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.

- As the authors claim that OFM learns straighter trajectories, comparing OFM to other baselines on standard generation benchmarks, e.g., CIFAR10, would be better.
questions: 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: limitations: the authors adequately addressed the limitations

Comments: title: Thank you! and concluding comments
The discussion period deadline is approaching...
Reviewer response
comment: We are very grateful to the Reviewers for valuable questions, fruitful discussions and raising the scores. In the final version, we will make the following edits: 

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

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

3) discuss the intricacy of high-dimensional (i.e., image data) generation in Limitations section, (nvT5, Bhuh).
Dear Reviewer, we thank you for your answer and for highlighting the theoretical contribution of our work. We are very grateful for the score increase. Let us leave some concluding comments.

1. We agree that the reliance on the ICNNs makes the practical adaptation of our methodology to high-dim generative image setups to be tricky. Following your suggestions, we will expand the corresponding limitation subsection (**ICNNs**) accordingly.

2. The work [1] which is based on the amortization methodology will receive due attention in our revised manuscript. In particular, we acknowledge that we will include the benchmark results of [1] in our Table. With regard to the superior performance of [1], we humbly believe that achieving the SOTA metrics on the benchmark lies primarily on the surface of technical solutions. I.e., particular architectures, optimizers, schedulers, regularizers etc. To support our point, we provide the additional results for our method fitted on the benchmark, but with additional exponential moving average (EMA) of the trained model weights. Note that EMA does not change the training process itself - it just creates a smoothed copy of the model whose weights are updated (at each new training iteration $t + 1$) as $W^{\text{ema}}\_{t + 1} = \alpha W^{\text{ema}}\_{t} + (1 - \alpha) W\_{t + 1}$, where $W_{t + 1}$ are the newly updated original trained weights. Our results for $\mathcal{L}^2$-UVP are in the table below ($\alpha = 0.999$). For completeness, we also recall the metrics for our models without EMA and additionally grab the benchmark results from [1].

| Solver/dim        | 2        | 4        | 8       | 16      | 32      | 64      | 128      | 256      |
|-------------------|----------|----------|---------|---------|---------|---------|----------|----------|
| OFM MB (ema)      |  $\textit{0.15}$   | $\textit{0.52}$   | $\textit{1.2}$   | $\textit{1.0}$   |  $\textbf{1.2}$ | 7.2     | $\textit{1.5}$    | 2.9      |
| OFM MB            | 0.33     | 0.81     | 1.5     | 2.0     | 3.4     | 10      | 6.7      | 11.7     |
| [1] (ICNNs) | 0.26     | 0.78     | 1.6     | 1.1     | $\textit{1.9}$     | $\textit{4.2}$   | 1.6      | $\textit{2.0}$    |
| [1] (MLP)   | $\textbf{0.03}$ | $\textbf{0.22}$ |  $\textbf{0.6}$ |  $\textbf{0.8}$ | 2.0     |  $\textbf{2.1}$ |  $\textbf{0.67}$ |  $\textbf{0.59}$ |

As we can see, the utilization of a rather cheap technique allows obtaining much better results, compared to those reported in our article. Regarding [1], it seems they did not utilize EMA in their experiments. Probably, the usage of EMA in their experiments also will lead to better results. Our newly provided table is not about competing with [1]. We just want to stress that the point of our benchmark experiments is to demonstrate the applicability and success of our proposed methodology - and we humbly think that we managed to achieve this goal - by demonstrating rather good and close-to-the-best metrics.

One more time, we are grateful to the reviewer for the valuable comments, and we will incorporate the necessary edits in the revised manuscript.
Thank you for the detailed comments and added experiments. This cleared up a lot of my questions, and I've decided to raise my score 4->5 as I think the paper has an interesting theoretical contribution.

I still have concerns regarding the scalability of the method because of its reliance on ICNNs making it very hard to scale up to real word examples like pixel-space image generative modeling. The authors mention that the main competitive methods of their proposed method are different Flow Matching methods but the work does not include results on any common generative modeling benchmarks as also mentioned by Reviewer Bhuh. These other FM methods don't have any problems scaling up to arbitrary dimensions. This should be sufficiently discussed and contextualized in the revised manuscript. Also as mentioned in my rebuttal, I believe [1] is a more closely related work to the proposed method with respect to the Optimal Transport benchmark and outperforms it across all dimensions.
Dear reviewer nvT5.

Thank you for your time efforts to review our paper.

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

Thank you in advance, Authors.
Thank you the the authors for their responses. Your responses partially address my concerns about the paper. I have raised the score by one point to reflect my positive attitude toward the additional explanation and empirical results.
I want to thank the authors for their response. Following the reviewers answers,  I have invested more time in better following the authors claim and proofs. It is now my understanding that the key results of the paper is that the minimizer of  FM loss constrained to velocity field which are a gradient of convex functions is the OT map. I acknowledged this is an important results and I will raise my score.
I thank the authors for their clarifications and additional experiments. I will keep my score.