Explicit Flow Matching: On The Theory of Flow Matching Algorithms with Applications

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:13:30

Abstract: This paper proposes a novel method, Explicit Flow Matching (ExFM), for training and analyzing flow-based generative models. ExFM leverages a theoretically grounded loss function, ExFM loss (a tractable form of Flow Matching (FM) loss), to demonstrably reduce variance during training, leading to faster convergence and more stable learning. Based on theoretical analysis of these formulas, we derived exact expressions for the vector field (and score in stochastic cases) for model examples (in particular, for separating multiple exponents), and in some simple cases, exact solutions for trajectories. In addition, we also investigated simple cases of Diffusion Generative Models by adding a stochastic term and obtained an explicit form of the expression for score. While the paper emphasizes the theoretical underpinnings of ExFM, it also showcases its effectiveness through numerical experiments on various datasets, including high-dimensional ones. Compared to traditional FM methods, ExFM achieves superior performance in terms of both learning speed and final outcomes.

Main paper: 
 introduction: 
 In recent years, there has been a remarkable surge in Deep Learning, wherein the advancements have transitioned from purely neural networks to tackling differential equations. Notably, Diffusion Models [16] have emerged as key players in this field. This models transform a simple initial distribution, usually a standard Gaussian distribution, into a target distribution via a solution of Stochastic Differentiable Equation (SDE) [1] or Ordinary Differentiable Equation (ODE) [2] with right-hand side representing a trained neural network. The Conditional Flow Matching (CFM) [9] technique, which we focus on in our research, is a promising approach for constructing probability distributions using conditional probability paths, which is notably a robust and stable alternative for training Diffusion Models. The development of the CFM-based approach includes various techniques and heuristics [4, 7, 13] aimed at improving convergence or quality of learning or inference. For example, in the works [19, 20, 10] it was proposed to straighten the trajectories between points by different methods, which led to serious modifications of the learning process. We refer the reader for, example, to the paper [20] where different FM-based approaches are summarised, and to the paper [9] for the connection between Diffusion Models and CFM.In our work, we introduced an approach which we called Explicit Flow Matching (ExFM), to consider the Flow Matching framework theoretically by modifying the loss and writing the explicit value of the vector field. Strictly speaking, the presented loss is a tractable form of the FM loss, see Eq. (5) of [9] . Base on this methods we can improve the convergence of the method in practical examples reducing the variance of the loss, but the main focus of our paper is on theoretical derivations.Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.Our method allows us to write an expression for the vector field in closed form for quite simple cases (Gaussian distributions), however, we note that Diffusion Models framework in the case of a Gaussian Mixture of two Gaussian as a target distribution is still under investigation, see recent publications [15, 8] .Our main contributions are:1. A tractable form of the FM loss is presented, which reaches a minimum on the same function as the loss used in Conditional Flow Matching, but has a smaller variance; 2. The explicit expression in integral form for the vector field delivering the minimum to this loss (therefore for Flow Matching loss) is presented.3. As a consequence, we derive expressions for the flow matching vector field and score in several particular cases (when linear conditional mapping is used, normal distribution, etc.); 4. Analytical analysis of SGD convergence showed that our formula have better training variance on several cases; 5. Numerical experiments show that we can achieve better learning results in fewer steps.Flow matching is well known method for finding a flow to connect samples from two distribution with densities ρ 0 and ρ 1 . It is done by solving continuity equation with respect to the time dependent vector field v(x, t) and time-dependent density ρ(x, t) with boundary conditions:EQUATIONFunction ρ(x, t) is called probability density path. Typically, the distribution ρ 0 is known and it is chosen for convenience reasons, for example, as standard normal distribution ρ(x) = N (x | 0, I).The distribution ρ 1 is unknown and we only know the set of samples from it, so the problem is to approximate the vector field v(x, t) ≈ v(x, t) using these samples. To make problem (1) well defined, one usually imposes additional regularity conditions on the densities, such as smoothness. The rigorous justification of the obtained results we put in the Appendix, leaving the general formulations of theorems and ideas in the main text.From a given vector field, we can construct a flow ϕ t , i. e., a time-dependent map, satisfying the ODE ∂ϕt(x) ∂t = v(ϕ t (x), t) with initial condition ϕ 0 (x) = x. Thus, one can sample a point x 0 from the distribution ρ 0 and then using this ODE obtain a point x 1 = ϕ 1 (x 0 ) which have a distribution approximately equal to ρ 1 . For given boundary ρ 0 and ρ 1 , the vector field or path solutions are not the only solutions, but if we have found any solution, it will already allow us to sample from the unknown density rho 1 . However, if the problem is more narrowly defined, e. g., one needs to have a map that is close to the Optimal Transport (OT) map, we have to impose additional constraints.The problem of finding any vector field v is solved in conditional manner in the paper [9] , where so-called Conditional Flow Matching (CFM) is present. Namely, the following loss function was introduced for the training a model v θ which depends on parameters θEQUATIONwhere ϕ t,x1 (x 0 ) is some flow, conditioned on x 1 (one can take ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 + σ s tx 0 in the simplest case, where σ s > 0 is a small parameter need for this map to be invertable at any 0 ≤ t ≤ 1). Hereinafter the dash indicates the time derivative. Time variable t is uniformly distributed: t ∼ U[0, 1] and random variables x 0 and x 1 are distributed according to the initial and final distributions, respectively: x 0 ∼ ρ 0 , x 1 ∼ ρ 1 . Below we omit specifying of the symbol E the distribution by which the expectation is taken where it does not lead to ambiguity.Model training using loss (2) argument values (x t , t). Indeed, a fixed point x t = ϕ t,x1 (x 0 ) can be obtained by an infinite set of x 0 and x 1 pairs, some of which are directly opposite, and at least for small times t the probability of these different directions may not be significantly different. At the same time, data ϕ ′ t,x1 (x 0 ) on which the model learns significantly different for such different positions of pairs x 0 and x 1 . Thus, the model is forced to do two functions during training: generalize and take the mathematical expectation (clean the data from noise).In our approach, see Fig. 1 (a), we feed the model input with cleaned data with small variance. Thus, the model only needs to generalize the data, which happens much faster (in fewer training steps).Moreover, in the process of constructing the modified loss, we have developed the exact formula for the vector field, see Eq. ( 11), (34). The existence of an explicit formula for the vector field is of great importance not only from a theoretical but also from a practical point of view.2 Main idea 
modified objective: 
 Lets expand the last two mathematical expectations in the loss (2) and substitute variables using map ϕ t,x1 , passing from the point x 0 to its position x t = ϕ t,x1 (x 0 ) at time t:EQUATION)We assume, that the map ϕ t,x1 is invertible at each 0 < t < 1, i. e. that ϕ -1 t,x1 (x t ) exits on this time interval and for all x t = {ϕ t (x 0 ) | ∀x 0 : ρ(x 0 ) > 0}. Eq. ( 3) can be seen as a transition from expectation on the variable x 0 ∼ ρ 0 to expectation on the variable x t ∼ ρ x1 (•, t), whereEQUATION) ∂x . See paper [5] for details about the push-forward operator "*". Our representation (3) is very similar to expression (9) of the cited paper [9] , only we write it in terms of the conditional flow rather than the conditional vector field.To obtain the modified loss, we return to end of the standard CFM loss representation in (3) . It is written as the expectation over two random variables x 1 and x t having a common distribution densityEQUATIONwhich, generally speaking, is not factorizable. Let us rewrite this expectations in terms of two independent random variables, each of which have its marginal distribution. The marginal distribution ρ m of x t can be obtained via integration:EQUATIONwhile the marginal distribution of x 1 is just (unknown) function ρ 1 . Let for convenience w(t, x 1 , x) =ϕ ′ t,x1 ϕ -1 t,x1 (x) 1 .We haveEQUATIONwhere we introduce a conditional distributionEQUATIONThe key feature of the representation ( 6) is that the integration variables x 1 and x are independent.Thus, we can evaluate them using Monte Carlo-like schemes in different ways. However, we go further and make a modification to this loss to reduce the variance of Monte Carlo methods.Note that so far the expression for L CFM have not changed, it has just been rewritten in different forms. Now we change this expression so that its numerical value, generally speaking, may be different, but the derivative of the model parameters will be the same. We introduce the following lossEQUATION)Theorem 2.1. Losses L CFM in Eq. (2) and L ExFM in Eq. ( 8) have the same derivative with respect to model parameters:EQUATIONProof is in the Appendix A.1.In the presented loss L ExFM , the integration (outside the norm operator) proceeds on those variables on which the model depends, while inside this operator there are no other free variables. Thus, using this kind of loss, it is possible to find an exact analytical expression for the vector field for which the minimum of this loss is zero (unlike the loss L CFM ). Namely, we haveEQUATIONWe can obtain the exact form of this vector field given the particular map ϕ t,x1 . For example, the following statement holds:Corollary 2.2. Consider the linear conditioned flow ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 which is inevitable as 0 ≤ t < 1. Then w(t, x 1 , x) = x1-x 1-t , ρ x1 (x, t) = ρ 0 x-x1t 1-t 1(1-t) d and the loss L ExFM in Eq. (8) reaches zero value when the model of the vector field have the following analytical formv(x, t) = (x 1 -x)ρ 0 x -x 1 t 1 -t ρ 1 (x 1 ) dx 1 (1 -t) ρ 0 x -x 1 t 1 -t ρ 1 (x 1 ) dx 1 . (11)This is the exact value of the vector field whose flow translates the given distribution ρ 0 to ρ 1 .Complete proofs are in the Appendix A.3.1. Note that the result (11) is not totally new, for example, a similar result (though in the form of a general expression rather than an explicit formula), was given in [19] , Eq. ( 9). However, our contribution consists of both the general form (10) and practical and theoretical conclusions from it (see below).Remark 2.3. In the case of the initial and final times t = 0, 1, Eq. ( 11) is noticeably simplerEQUATION)This expression for the initial velocity means that each point first tends to the center of mass of the unknown distribution ρ 1 regardless of its initial position.Extensions to SDE Now let the conditional map be stochastic:ϕ t,x1 = (1 -t)x 0 + tx 1 + σ e (t)ϵ,where ϵ ∼ N (0, 1). Typically, σ e (0) = σ e (1) = 0, for example, σ e (t) = t(1 -t)σ e .Note that this formulation covers (with appropriate selection of the σ e (t) parameter) the case of diffusion models [20] .Then, we can write the exact solution for a so-called score and flow matching objective (see [20] for details)L [SF] 2 M (θ) = E ∥v θ (x, t) -u • t (x)∥ 2 flow matching loss +λ(t) 2 ∥s θ (x, t) -∇ log p t (x)∥ 2 score matching loss.that corresponds to this map. In the last expression, the following explicit conditional expressions are considered in the cited paper for the case σ e (t) = t(1 -t)σ e u • t (x) = 1 -2t t(1 -t) (x -(tx 1 + (1 -t)x 0 )) + (x 1 -x 0 ), ∇ log p t (x) = tx 1 + (1 -t)x 0 -x σ 2 e t(1 -t).The exact solution (our result, explicit analog of the Eq. ( 10) from [20] ) under consideration has the form (44) and (46) and, for example for the for the Gaussian ρ 0 this expressions reduced to the Eq. ( 49) and (50), correspondingly. See Appendix E for the details on this case. Let us consider the difference between our new scheme based on loss L ExFM and the classical CFM learning scheme. As a basis for the implementation of the learning scheme, we take the open-source code2 from the works [20, 19] .Consider a general framework of numerical schemes in classical CFM. We first sample m random time variables t ∼ U[0, 1]. Then we sample several values of x. To do this, we sample a certain number n samples {x i 0 } n i=1 from the "noisy" distribution ρ 0 , and the same number n of samples {x i 1 } n i=1 from the unknown distribution ρ 1 . Then we pair them (according to some scheme), and get n samples asx j,i = ϕ t j ,x i 1 (x i 0 ) (e.g. a linear combination in the simple case of linear map: x j,i = (1-t j )x i 0 +t j x i 1 ), ∀i = 1, 2, . . . , n; ∀j = 1, 2, . . . , m. Note, than one of the variable n or m (or both) can be equal to 1.At the step 2, the following discrete loss is constructed from the obtained samplesEQUATION)Finally, we do a standard gradient descent step to update model parameters θ using this loss.The first and last step in our algorithm is the same as in the standard algorithm, but the second step is significantly different. Namely, we additionally generate a sufficiently large number N ≫ n • m of samples x 1 from the unknown distribution ρ 1 , sampling (N -n) new samples and adding to it the samples {x i 1 } n 1 that are already obtained on the previous step.Then we form the following discrete loss which replaces the integral on x 1 in L ExFM by its evaluation v d by self-normalized importance sampling or rejection sampling (see Appendix B for details)EQUATION)For example, if we use self-normalized importance sampling and assume that the Jacobian det ∂ϕ -1 t,x1 (x) ∂x do not depend on x 1 , we can writeEQUATIONTheorem 2.4. Under mild conditions, the error variance of the integral gradient (9) using the Monte Carlo method (14) is lower than using formula (13) with the same number n • m of samples for {x}.Sketch of the proof is in the Appendix A.2. The steps of our scheme are formally summarized in Algorithm 1.Particular case of linear map and Gaussian noise Let ϕ t,x1 be the linear flow: ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 . and consider the case of standard normal distribution for the initial density ρ 0 :ρ 0 (x) ∼ N (x | 0, I).Then in the case of using self-normalized importance sampling, we haveEQUATIONHere, the lower index k in SoftMax stands for the k-th component, and the SoftMax operation itself came about due to exponents in the Gaussian density as a more stable substitute for computing than directly through exponents.Extension of other maps and initial densities ρ 0 Common expression (10) can be reduced to closed form for the particular choices of density ρ 0 and map ϕ (consequently, expression for w). We summarise several known approaches for which FM-based techniques can be applied in Table 1 3 .See Appendix C and D for derivations of formulas and for more extensions.Table 1 : Correspondence between some methods which can reduced to FM framework and our theoretical descriptions of them.Probability Path q(z) µ t (z) σ t Explicit expressions: vector field (VF) and score (S) Var. Exploding [17] ρ 1 (x 1 )x 1 σ 1-t VF: (32)Var. Preserving [6] ρ 1 (x 1 ) α 1-t x 1 1 -α 2 1-tVF: (31) Flow Matching [9] ρ 1 (x 1 ) tx 1 tσ s -t + 1 VF: (11) if σ = 0; and (26) Independent CFM ρ 0 (x 0 )ρ 1 (x 1 ) tx 1 + (1 -t)x 0 σ VF: (10) Schrödinger Bridge CFM [20] ρ 0 (x 0 )ρ 1 (x 1 ) tx 1 + (1 -t)x 0 σ t(1 -t) Can be obtained by SDE using VF: (49), S:(50)Complexity We assume that the main running time of the algorithm is spent on training the model, especially if it is quite complex. Thus, the running time of one training step depends crucially on the number n • m of samples {x} and it is approximately the same for both algorithms: the addition of points x 1 entails only an additional calculation using formula (16) , which can be done quickly and, moreover, can be simple parallelized.Ensuring the stability of optimization is vital. Let ∆θ be changes in parameters, obtained by SGD with step size γ/2 applied to the functional from Eq. ( 13):EQUATIONFor simplification, we consider a function, v θ (x, t), capable of perfectly fitting the CFM problem and providing an optimal solution for any point x and time t. For a linear conditional flow at a specific point x j,i ∼ ρ x i 1 (•, t j ) at time t j ∼ U (0, 1), the update ∆v(x j,i , t j ) can be represented as follows:EQUATIONwhere xi 0 =x j,i -t j x i 1 1-t j. We define the dispersion D x,x1 f (x, x 1 ) for x ∼ ρ x1 (•, t) and x 1 ∼ ρ 1 as:EQUATION)Proposition 2.5. At the time t = 0, the dispersion of update in the form (18) have the following element-wise lower bound:D x j,i ,x i 1 ∆v(x j,i , 0) = γ 2 D x i 1 x i 1 + γ 2 D x j,i ,x i 1 (x j,i + v(x j,i , 0)) ≥ γ 2 D x i 1 x i 1 .Equality is reached when the model v(x j,i , 0) has exact values equal to (12) .Given that the dispersion cannot be reduced with an increase in batch size, the only available option is to decrease the step size of the optimization method, i. e., reduce the learning rate slowing down the convergence. The situation is much better for the proposed loss in (14) . We can express the update ∆v(x j,i , t j ) in the case of ExFM objective as:EQUATIONwherex j,i ∼ ρ x i 1 (•, t j ), x k 1 ∼ ρ 1 and ρ x j,i |x k 1 , t j = ρ 0 x j,i -t j x k 1 1-t j / N k=1 ρ 0 x j,i -t j x k 1 1-t j. Similar to the derivations in the previous part, we can found simplified form for the dispersion of update at t = 0.Proposition 2.6. At the time t = 0, the dispersion of update from (20) have the following element-wise lower bound:D x j,i ,x k 1 ∆v(x j,i , 0) = γ 2 N D x k 1 x k 1 + γ 2 D x j,i ,x k 1 (x j,i + v(x j,i , 0)) ≥ γ 2 N D x k 1 x k 1 .Equality is reached when the model v(x j,i , 0) has exact values equal to (12) .In comparison to CFM, the dispersion of the update is N times smaller than the dispersion of the target distribution and could be controlled without impeding convergence by adjusting the number of samples N . In Figure 1 -8.51e-02 ± 4.85e-02 1.64E-01 ± 4.18E-02 5.22E-02 ± 3.92E-02 GAS -5.53e+00 ± 3.66e-02 -5.00E+00 ± 2.56E-02 -5.48E+00 ± 2.90E-02 HEPMASS 2.16E+01 ± 6.31E-02 2.21e+01 ± 6.13e-02 2.16E+01 ± 4.32E-02 BSDS300 -1.29E+02 ± 8.40E-01 -1.29E+02 ± 8.97E-01 -1.32e+02 ± 6.39e-01 MINIBOONE 1.34e+01 ± 1.95e-04 1.42E+01 ± 1.29E-04 1.43E+01 ± 9.22E-053 Numerical Experiments Toy 2D data We conducted unconditional density estimation among eight distributions. Additional details of the experiments see in the Appendix H. We commence the exposition of our findings by showcasing a series of classical 2-dimensional examples, as depicted in Fig. 3 and Table 2 .Our observations indicate that ExFM adeptly handles complex distribution shapes is particularly noteworthy, especially considering its ability to do so within a small number of epochs. Additionally, the visual comparison underscores the evident superiority of ExFM over the CFM approach. The primary experimental setup follows that used in [19] . Additional details on the hyperparameters used are available in Appendix H. Based on the findings presented in Table 4 , we determine that ExFM-S surpasses I-CFM on all four datasets in terms of generative performance (W 2 ) and also outperforms in terms of OT optimality (NPE) on two of them, exhibiting similar results on the remaining datasets. It also demonstrates performance similar to OT-CFM. While ExFM-S is not as robust as the basic ExFM, it enables the matching of one dataset to another (moons → 8gaussians) as it does not necessitate the presence of an explicit formula for ρ 0 . Among other things, this experiment demonstrates the feasibility of our methods when both distributions ρ 0 and ρ 1 are unknown.Table 4 : ExFM-S evaluation on four toy datasets (µ ± σ over three seeds). For comparison we take I-CFM, OT-CFM, and ExFM (no values for moons → 8gaussians due to the absence of explicit formula for ρ 0 ). Performance in generative modeling (W 2 ) and dynamic OT optimality (NPE) is assessed. The best result for each metric is highlighted in bold. Instances where we outperform CFM are underscored. Metric → W2 ↓ NPE ↓ Algorithm ↓ Dataset → N → moons N → 8gaussians moons → 8gaussians N → 2spirals N → moons N → 8gaussians moons → 8gaussians N → 2spirals I-CFM 0.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification:Guidelines: Limitation is discussed both in the theoretical part (in particular, in the formulation of theorems) and in the practical part, where we can see at which cases our method works better or worse• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Consider the case where we only have access to the samples {x i 0 } N0 i=1 and {x i 1 } N1 i=1 from both distributions, ρ 0 and ρ 1 , but do not know their explicit expressions. In this case, we can estimate the vector field using by a method similar to the one we used to estimate the vector field in (15):EQUATION.Similar for the scoreEQUATION.In addition, we can also use the importance sampling method in this case. Namely we can use both approaches: self-normalized importance sampling and rejection sampling, similar to what is described in Sec. B F Consistency of Eq. (24) in the case of optimal transportLet us analyze what happens if in formula (24) the joint density π represents the following Dirac delta-function 4 : π(x 0 , x 1 ) = δ x 0 -F (x 1 ) ,i. e. we have a deterministic mapping F from x 1 to x 0 . Then, the Eq. (34) come tov(x, t) = (x 1 -x) δ ϕ -1 t,x1 (x) -F (x 1 ) dx 1 (1 -t) δ ϕ -1 t,x1 (x) -F (x 1 ) dx 1 .Let y(x, t) be the unique solution of the equationEQUATIONconsidered as an equation on y. Thenv(x, t) = x -y(x, t) 1 -t .Now, let us use linear mapping ϕ t,x1 (x) = x 1 t + x(1 -t), with inverse ϕ -1 t,x1 (x) = x-tx1 1-t , and consider the simplest case when the original distribution is a d-dimensional standard Gaussian and ρ 1 is a d-dimensional Gaussian with mean µ and diagonal variance Σ = diag(σ). We know the OT correspondence between Gaussians, namelyF (x 1 ) i = (x 1 -µ) i Σ ii , ∀1 ≥ i ≥ d.Here and further by index i we denote ith component of the corresponding vector. Then, the Eq. (53) reads as(x -yt) i 1 -t = (y -µ) i Σ ii , with the solution y(x, t) i = µ i (1 -t) + x i Σ ii 1 + (Σ ii -1)t .Then the expression for the vector field isv(x, t) i = µ i + x i (Σ ii -1) 1 + (Σ ii -1)t .Now, knowing the expression for velocity, we can write the equations for the trajectories x(t):   x ′ (t) i = µ i + (x(t)) i (Σ ii -1) 1 + (Σ ii -1)t , x(0) i = (x 0 ) i .This equation have closed-form solution:x(t) = µt + x 0 -(1 -σ) tx 0 .Analyzing the obtained solution, we conclude that, first, the trajectories obey the given mapping F :F (x(1)) i = (x 0 ) i = (x(1) -µ) i Σ ii ,And, second, the trajectories are straight lines (in space), as they should be when the flow carries points along the optimal transport.As a final conclusion, note that, of course, if we are mapping optimal transport F , then it is meaningless to use numerical formula (16) . However, usually the exact value of the mapping F is not known, and our theoretical formula (34) can help to rigorously establish the error that is committed when an approximate mapping is used instead of the optimal one. G Analytical derivations for example in Fig. 1(b) G 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Justification: All theorems are formulated according to strict mathematical rules. There are proofs or proof sketches in the text or Appendix.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
conclusions: 
 The presented method introduces a new loss function in tracrable form (in terms of integrals) that improves upon the existing Conditional Flow Matching approach. New loss as a function of the model parameters, reaches zero at its minimum. Thanks to this, we can: a) write an explicit expression for the vector field on which the loss minimum is achieved; b) get a smaller variance when training on the discrete version of the loss, therefore, we can learn the model faster and more accurately.Numerical experiments conducted on toy 2D data show reliable outcomes under uniform conditions and parameters. Comparison of the absolute values of loss for the proposed method and for CFM for the same distributions show that the absolute values of loss for these models differ strikingly, by a factor of 10 2 -10 3 . Experiments on high-dimensional datasets also confirm the theoretical deductions about the variance reduction of our method. However, we emphasize that we do not expect to use the proposed method in its pure form. On the contrary, we expect that the theoretical implications of our formulas will contribute to the construction of better learning or inference algorithms in conjunction with other heuristics or methods.Algebraic analysis of variance for some cases (in particular, for the case t = 0 or for the case of two Gaussians as initial and final distributions) show an improvement in variance when using the new loss. However, it is rather difficult to analyze in the general case, for all times t and general distributions ρ 0 and ρ 1 .Having the expression for the vector field and score in the form of integrals, we can explicitly write out their expressions for some simple cases; in the case of Gaussian distributions we can also write out the exact solution for the trajectories. Thus, our approach allows one to advance the theoretical study of FM-based and Diffusion Model-based frameworks.NeurIPS Paper Checklist Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes]Justification: Yes, all parameters and hyperparameters required for reproducibility are described in the Appendix Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
open access to data and code: 
 Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [NA] Justification: We do not provide a code, paper is mostly theoretical. We use open datasets.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Yes, all parameters and hyperparameters required for reproducibility are described in the Appendix Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
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: [Yes] Justification: For several experiments variances are given.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: The experiments are simple enough to be reproduced on publicly available resources Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Paper is mostly theoretical, there's little chance there could be a negative impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our work has a very low risk of misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 
licenses for 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: Links are provided to articles with open repositories whose code was used in the workGuidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: We do not provide a code. All Theorems and Satetements are well formulated.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
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: we have no crowdsourcing Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
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? 
 
 Appendix: Justification: The article and experiments meet all the requirements of the code of ethics, all citations for materials used are given.Guidelines: The research conducted in the article and its text meet the ethics codex.• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There are no experiments with subjects in the paper, only numerical experiments.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.A.1 Proof of the Theorem 2.1Proof. We need to proof, that dLCFM(θ) dθ = dLExFM(θ) dθ .To establish the equivalence of L CFM and L ExFM up to a constant term, we begin by expressing L CFM in the format specified by equation ( 6):L CFM = E t,x1,x∼ρm(•,t) ∥v θ (x, t) -w(t, x 1 , x)∥ 2 × ρ c (x|x 1 , t)/ρ 1 (x 1 ).Utilizing the bilinearity of the 2-norm, we can rewrite L CFM as:L CFM = E t,x1,x∼ρm(•,t) ∥v θ (x, t)∥ 2 ρ c (x|x 1 , t) ρ 1 (x 1 ) - 2E t,x1,x∼ρm(•,t) v θ (x, t) T • w(t, x 1 , x)ρ c (x|x 1 , t) ρ 1 (x 1 ) + C. (21)Here, T denotes transposed vector, dot denotes scalar product, C represents a constant independent of θ.Noting that E x1 ρ c (x|x 1 , t)/ρ 1 (x 1 ) = 1:E x1 ρ c (x|x 1 , t) ρ 1 (x 1 ) = ρ x1 (x, t)ρ 1 (x 1 ) dx 1 ρ x1 (x, t)ρ 1 (x 1 ) dx 1 = 1,we can simplify the first term in the expansion (21):EQUATION)For our loss L ExFM in the form (8) we also use the bilinearity of the norm:EQUATION)Comparing the last expression and the Eq. ( 21) with the modification (22) and also taking into account the independence of random variables x and x 1 , we come to the conclusion that L ExFM is equal to L CFM up to some constant independent of the model parameters.A.2 Sketch of the proof of the Theorem 2.4Proof. We need to prove that DdL d ExFM (θ) dθ ≤ D dL d CFM (θ) dθ, where L d ExFM (θ) and L d CFM (θ) discrete loss functions presented in ( 14) and ( 13). Firstly, let us rewrite the derivative of loss functions using the bilinearity:dL d ExFM (θ) dθ = 2 i,j dv θ (x j,i , t j ) dθ T • v θ (x j,i , t j ) -v d (x j,i , t j ) .Note that in this expression, values x j,i as well as t j , which are included in the argument of the function v, are fixed (our goal to calculate the variance with fixed model arguments). Thus, we need to consider the variance of the remaining expression arising from the randomness of x k 1 .Recall (below we will omit the indices at variables x and t),v d (x, t) = N k=1 w(t, x k 1 , x) • ρ 0 ϕ -1 t,x k 1 (x) N k=1 ρ 0 ϕ -1 t,x k 1 (x) .Note, that if N = 1, (i. e. we do not sample any additional points other than the ones we have already sampled) this expression is exactly the same as the derivative of the common discretized CFM lossdL d CFM (θ) dθ .Moreover, recall that one of the points (without loss of generality, we can assume that its index is 1)x 1 1 is added from the set from which point x was derived: x = ϕ t,x 1 1 (x 0 ). (Here x 0 is the paired point to x 1 1 )Thus, we can rewrite expression for v d :EQUATION)Thus, our task was reduced to evaluating how well the additional terms (for k starting from 2) improve approximate of the original integrals that are in loss (8) .So, we need to estimate the following dispersion ratio, where in the numerator is the variance of discrete loss CFM, and in the denominator -the variance of loss ExFM:k D = D v θ (x, t) -w(t, x 1 1 , x) D v θ (x, t) - N k=1 w(t,x k 1 ,x)•ρ0 ϕ -1 t,x k 1 (x) N k=1 ρ0 ϕ -1 t,x k 1 (x)The smaller coefficient k D is, the better the proposed loss ExFM works.Formally, we can write our problem as an importance sampling problem for the following integral:I = f (x)p(x) dx .This integral we estimate by sample mean of the following expectation over some random variable with density function q(x):I = E x∼q w(x)f (x) with w(x) = p(x) q(x) .We replace the exact value of I with the valueI = N k=1 w(x i 1 )f (x k 1 ) N i=k w(x k 1 ).It follows from the strong law of large numbers that in the limit N → ∞, I → I almost surely. From the central limit theorem we can find the asymptotic variance:EQUATION)In our case (loss L ExFM ), we have q(x 1 ) = ρ 1 (x 1 ), f (x 1 ) = w(t, x 1 , x) and w(x1 ) = ρ 0 ϕ -1 t,x1 (x) .Despite the fact that the equation ( 25) for the variance contains N in the denominator, it is rather difficult to give an estimate of its behavior in general. The point is that this formula is well suited for the case when w in it is of approximately the same order. In the considered case, this is achieved at times t noticeably less than 1.But in the case, when t is closed to 1 we have, for example, for the linear map, thatw(x 1 ) = ρ 0 ϕ -1 t,x1 (x) = ρ 0x -x 1 t 1 -t and this function has a sharp peak near the point x/t if it is considered as a function of x 1 . Thus, at such values of t, only a small number of summands will give a sufficient contribution to the sum compared to the first term.Finally, inequality k D < 1 is formally fulfilled, but how much k D is less than one depends on many factors.To justify the expression (11) , we use a invertable transformation and then strictly take the limit σ s → 0.Expression Eq. ( 11), ( 16) are obtained for the simple map ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 which is not invertable at t = 1. For the map with small regaluraziting parameter σ s > 0 ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 + σ s x 0 , which is invertable at all time values 0 ≤ t ≤ 1, Eq. ( 11), ( 16) needs modifications. Namely, for this map the following exact formulas holds truev(x, t) = w(t, x 1 , x)ρ c (x|x 1 , t)ρ 1 (x 1 ) dx 1 = x 1 -x(1 -σ s ) ρ 0 x-x1t 1+σst-t ρ 1 (x 1 ) dx 1 (1 + σ s t -t) ρ 0 x-x1t 1+σst-t ρ 1 (x 1 ) dx 1 .(26) By direct substitution we make sure that for this vector fieldEQUATIONandEQUATION)where we perform change of the variables y ← x1-x σst .A.3.1 Prof of the explicit formula (11) for the vector field Assumption A.1. Density ρ 1 is continuous at any point x ∈ (-∞, ∞).Theorem A.2. In equations (26), ( 27) and (28) we can take the limit σ s → 0 under integrals to get Eq. ( 11) and (12) .Proof. Assuming that the distribution ρ 1 has a finite first moment: | ξρ 1 (ξ) dξ | < C 1 and that the density of ρ 0 is bounded: ρ 0 (x) < C 2 , ∀x ∈ (-∞, ∞), we obtain that the integrand functions in the numerator and denominator in the Eq. ( 26) can be bounded by the following integrable functions independent of σ s and t:ρ 0 x -x 1 t 1 + σ s -t ρ 1 (x 1 ) < C 1 ρ 1 (x 1 )and0 ≤ x 1 ρ 0 x -x 1 t 1 + σ s t -t ρ 1 (x 1 ) < x 1 C 1 ρ 1 (x 1 ), x ≥ 0, 0 > x 1 ρ 0 x -x 1 t 1 + σ s t -t ρ 1 (x 1 ) > x 1 C 1 ρ 1 (x 1 ), x < 0.It follows that both integrals in expression (26) converge absolutely and uniformly. So, we can swap the operations of taking the limit and integration, and we can take the limit σ s → 0 in the integrand for any time t ∈ [0, t 0 ] for arbitrary t 0 < 1.Now, let us consider the case t = 1. From Assumption A.1 the boundedness of the density ρ 1 follows: ∞, ∞ ). Thus, integrand functions in the numerator and denominator in the Eq. ( 28) can be bounded by the following integrable functions independent of σ s :ρ 1 (x) < C 2 , ∀x ∈ (-ρ 0 (y)ρ 1 (x -yσ s ) < ρ 0 (y)C 2and 0 ≤ yρ 0 (y)ρ 1 (x -yσ s ) < yC 2 ρ 0 (y), y ≥ 0, 0 > yρ 0 (y)ρ 1 (x -yσ s ) > yC 2 ρ 0 (y), y < 0. The existence of the limit limσs→0 ρ 1 (x -yσ s ) = ρ 1 (x),follows from Assumption A.1.Finally, we conclude that formula (11) , regarded as the limit σ s → 0 of the (26) at any t ∈ [0, 1], is true.Theorem A.3. The vector field in Eq. ( 11) delivers minimum to the Flow Matching objective (see the work [9] ), E t E x∼ρ(x,t) ∥v(x, t) -v(x, t)∥, where ρ(x, t) and v(x, t) satisfy the equation (1) with the given densities ρ 0 and ρ 1 .Proof. The proof is based on the previous statements and on a Theorem 1 from [9] (that the marginal vector field based on conditional vector fields generates the marginal probability path based on conditional probability paths.To complete the proof, we must justify that, with σ s tending to zero, the marginal path at t = 1 coincides with a given probability ρ 1 .Consider the marginal probability path p t (x, t)EQUATION)where p t (x|x 1 , σ s ) is conditional probability paths obtained by regularized linear conditional map.Distribution p t in the time t = 0 is equal to standard normal distribution p 0 (x|x 1 , σ s ) = N (x | 0, 1)and at the time t = 1 it is a stretched Gaussian centered atx 1 : p 1 (x|x 1 , σ s ) = N (x | x 1 , σ s I).Substituting p 1 into the Eq. ( 29) and considering that there exists a limit σ s → 0 due to Assumption A.1, we obtainp 1 (x) = lim σs→0 p t (x|x 1 , σ s )ρ 1 (x 1 )dx 1 = ρ 1 (x 1 ).This finish the proof.A.3.2 Learning procedure for σ s > 0Using standard normal distribution as initial density ρ 0 , and the regularized map ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 + σ s tx 0 we obtain the following approximation formulav d (x, t) = N k=1x k 1 -x(1-σs)1-t(1-σs) exp Y k N k=1 exp(Y k ), whereY k = - 1 2 x -t • x k 1 2 R d 1 -t(1 -σ s ) .In practical applications, the exponent calculation is replaced by the SoftMax function calculation, which is more stable.In general, we need to estimate the following expressionI(η) = w(x 1 , η)f (x 1 , η)ρ 1 (x 1 ) dx 1 f (x 1 , η)ρ 1 (x 1 ) dx 1 .In particular, substituting η → {x, t}, w(x, η) → (x 1 -x)/(1 -t) we obtain formula (11) and similar ones with similar substitutions.If we can sample from the ρ 1 distribution, we can estimate this integral in two ways: self-normalized importance sampling and rejection sampling.Let X = {x k 1 } N k=1 be N samples from the distribution ρ 1 .Self-normalized Importance Sampling In this caseEQUATION)This estimate is biased in theory, but there several methods to reduce this bias and improve this estimate, see, for example, [3] . Our numerical experiments generally show that the estimation (30) in the form is already sufficient for stable results; we don not observe any bias.Rejection sampling Let Y = {y k } M k=1 ⊂ X be a subset of the the initially given set of samples, which is formed according to the following rule. Let C = sup x ρ 1 (x). For a given sample x j 1 we generate a random uniformly distributed variable ξ j ∼ U(0, 1) and if f (x j 1 ) ≥ Cξ j , then we put the point x j k to the set Y; otherwise we reject it.Having formed the set Y, we evaluate the integral asI(η) ≈ 1 M M k=1w(y k , η).To justify the last estimation, we note, that the points from the set Y are distributed according to (non-normalized) density ρ(x)f (x, η)ρ 1 (x). One can show it using the proof of the rejection sampling method. This is the same density as in Eq. ( 7) and thus we estimate the expression (10) using Important Sampling without any additional denominator.Comparison When we apply these techniques to evaluating the expression for the vector field, we know that when the time parameter t is close to 1, the function f (x 1 , η) (which is a scaled ρ 0 ) has a peak at the point x = x 1 . This means that only a small number of points from the original set will end up in the set Y. Moreover, in the case when the time t is very close to one and the data are well separated, only one point x 1 will end up in Y. This explains why we initially put this point in the set X , because otherwise it would be possible that the set Y is empty and M = 0.As a future work, we indicate a theoretical finding of the probability of hitting a particular point x 1 in the set Y and, thus, a modification of our algorithm, when the sample x 1 will not always go to the set X , but with some probability -the greater the t the closer this probability to 1.Algorithm 1 Vector field model training algorithm Require: Sampler from distribution ρ 1 (or a set of samples); parameters n and m (number of spatial and time points, correspondingly); parameter N (number of averaging point); model v θ (x, t); algorithm with parameters for SGD Ensure: quasi-optimal parameters θ for the trained model 1: Initialize θ (maybe random) 2: while exit condition is not met do3: Sample m points {t j } from U[0, 1] 4: Sample n points pairs {x i 0 , x i 1 } n i=1 from joint distribution π (π(x 0 , x 1 ) = ρ 0 (x 0 )ρ 1 (x 1 ) if variables are independent) 5: Sample N -n points {x l 1 } from ρ 1 and form {x k 1 } = {x i 1 } ∪ {x l 1 } //We can take all available samples as {x k 1 } if we don't have access to a sampler, but only ready-made samples.For all i and j calculate the sum at the right side of (14) (using ( 16) if ρ 0 is standard Gaussian or (24) in general)7:Calculate the sum on i and j in discrete loss (14) , and take backward derivative, obtaining approximate grad G ≈ ∇ θ L ExFM of loss L ExFM on model parameters θ.Update model parameters θ ← SGD(θ, G) 9: end while General form of the proposed Algorithm is given in Alg 1.When using other maps, formula (11) is modified accordingly. For example, if we use the regularized map ϕ t,x1 (x 0 ) = (1 -t)x 0 + tx 1 + σ s tx 0 , we get the formula (26).Note, that in this case the final density ρ(x, 1), obtained from the continuity equation is not equal to ρ 1 , but is its smoothed modification.When using a different initial density ρ 0 (not the normal distribution), an obvious modification will be made to formula (16) .Diffusion-like models We can treat so-called Variance Preserving [6] model as CFM with the mapϕ t,x1 (x) = α 1-t x + 1 -α 2 1-t x 1 .and ρ 0 as standard normal distribution: ρ 0 = N • 0, 1 2 In this case, the common expression (10) for vector filed transforms toEQUATIONwhere α ′ s = dαs ds .Similarity we can treat so-called Variance Exploding [17] model as CFM with the mapϕ t,x1 (x) = σ 1-t x + x 1 .and ρ 0 also as standard normal distribution: ρ 0 = N • 0, 1 2 In this case, the common expression (10) for vector filed transforms toEQUATIONwhere σ ′ s = dσs ds .Joint Distribution Moreover, in addition to the independent densities x 0 ∼ ρ 0 and x 1 ∼ ρ 1 , we can use the joint density {x 0 , x 1 } ∼ π(x 0 , x 1 ). In the papers [20, 19] , optimal transport (OT) and Schrödinger's bridge are taken as π. In this case the expression for the vector field changes insignificantly: the conditional probability ρ c from Eq. ( 7) is subject to change:EQUATIONThen, Eq. ( 10) remains the same in general case. In the case of linear ϕ, the extension of Eq. ( 11)EQUATIONIn all of the above cases, the essence of Algorithm 1 does not change (except that in the case of dependent x 0 and x 1 we should be able either to calculate the value of π ϕ -1 t,x1 (x), x 1 /ρ 1 (x 1 ) or to estimate it).In this section, we present several analytical results that directly follow from our exact formulas for the vector field, which, to the best of our knowledge, have not been published before.Consider the flow from a one-dimensional Gaussian distribution ρ 0 ∼ N • µ 0 , σ 2 0 into another (with other parameters) Gaussian distribution ρ 1 ∼ N • µ 1 , σ 2 1 . Note that in this case the generalization to the multivariate case is done directly, so the spatial variables are separated. EQUATIONbut the expression for the path x(t) is unknown. Using Eq. ( 44)-( 46) we can explicitly calculate vector field v and score s with the setup as in Sec. D.1 but with additional noise, i. e. in the stochastic case.Consider like in the Sec. D.1 the flow from a one-dimensional standard Gaussian distribution ρ 0 ∼ N • 0, 0 2 into another (with other parameters) Gaussian distribution ρ 1 ∼ N • µ 1 , σ 2 1 but with additional noise as described above.In this case we have for the field.EQUATIONWe can solve ODE with this field and get the expression for the trajectories, starting from the given point x 0 :EQUATION)These trajectories, for different x 0 are depicted in Fig. 6 . For the score s in the considered case we haves(x, t) = tµ 1 -x (1 -t) 2 + t(1 -t)σ 2 e + t 2 σ 2 1Thus, we can explicitly write expressions for the stochastic process for the evolution from the initial distribution rho 0 (standard Gaussian) to the final distribution ρ 1 :dx(x) = x tσ 2 1 + (1 -t)σ 2 e /2 -(x -µ 1 ) (1 -t) + tσ 2 e /2 t(1 -t)σ 2 e + σ 2 1 t 2 + (1 -t) 2 + + g 2 (t) 2 tµ 1 -x (1 -t) 2 + t(1 -t)σ 2 e + t 2 σ 2 1 dt + g(t) dW (t) .Here g(t) is arbitrary smooth function. In the case of Shrödinger Bridge we take g(t) = σ e t(1 -t).E Detail on the SDE case E.1 Optimal vector field and score for stochastic mapFollowing [20] we consider a so-called Brownian bridge B(t) from x 0 to x 1 with constant diffusion rate σ e . This stochastic process can be expressed through a multidimensional standard Winner process W (t) asEQUATIONThus, the conditional distribution p(t, x | x 0 , x 1 ) conditioned on the starting x 0 and end point x 1 is Gaussian:p(x, t | x 0 , x 1 ) = N x (1 -t)x 0 + tx 1 , σ 2 e t(1 -t) .We can not directly use the results Theorem 3 from [9] (or similar Theorem 2.1 from [19] ) for the Gaussian paths, as in this case σ(0) = 0. To circumvent this obstacle and to be able to write an expression for the conditional velocity, we assume that we have a Gaussian distribution with a very narrow peak at the initial (t = 0) and final (t = 1) points. In other words, we will consider conditional probabilities of the formEQUATIONwhere parameter η is small enough. Then we can use the above Theorems and immediately writeEQUATIONAfter integrating over x 0 and x 1 , we can take the limit η → 0. Thus, now for fixed x 0 and x 1 we do not have a fixed value of x t in which to train the model, but a random one. In general case, we end up to the loss:EQUATIONwhere π(x 1 , x 0 ) is the density of the joint distributions with the marginal equal to the two given probabilities:π(x 1 , x 0 ) dx 1 = ρ 0 (x 0 ), π(x 1 , x 0 ) dx 0 = ρ 1 (x 1 ).In the simple case, π(x 1 , x 0 ) = ρ 0 (x 0 )ρ 1 (x 1 ). Vector field in Eq. ( 43) if taken in the form of Eq. ( 42). Now, we can obtain an explicit form for the vector field v at which the written loss is reached its minimum by performing the same calculations as in the derivation of formula (10) :EQUATION)As in the work [20] we can also train score network. Namely, as marginals for Brownian bridge are Gaussian, we can write explicit conditional score for conditional probabilistic path∇ log p(x, t | x 0 , x 1 ) = µ(t) -x σ 2 e (t) = x 0 (1 -t) + x 1 t -x σ 2 e t(1 -t).In the work [20] the following loss is introduced to train a model for this scoreEQUATION)Similar to (44), for the optimal score s we have:EQUATION)where p is given in (41).Note that the obtained vector field gives marginal distributions p(x, t), which (in the limit η → 0) at t = 1 leads to the distribution we need: p(x, t = 1) = ρ 1 (x). However, the addition of the stochastic term allows us to extend the scope of application of the explicit formula for the vector field. In particular, it can be applied to the situation when we have two sets of samples and both distributions are unknown, as well as the possibility of constructing SDE and solving it using, for example, the Euler-Maruyama method (see examples below).As consequence of Theorem 3.1 from [20] we have that, if v is given by Eq. ( 44) then ODEEQUATIONrecovers the marginal ρ(x, t) (with the given initial conditions) of the stochastic process P (t) which is obtained by marginalization conditional Brownian bridge (40) over initial and target distributionP (t) = B(t | x 0 , x 1 )π(x 0 , x 1 ) dx 0 dx 1 .As the second consequence of this Theorem, the SDEEQUATIONgenerates so-called Markovization of the process P (t). Indeed, we can rewrite PDE Eq. ( 47) in the form∂ρ(x, t) ∂t = -div ρ(x, t)v(x, t) + g 2 (t) 2 ∇ρ(x, t) + g 2 (t) 2 ∆ρ(x, t),where nabla operator is defined as ∆ = div ∇. Thus, we get the Fokker-Planck equation for the density of the stochastic process (48).In particular case of Brownian bridge when σ e (t) = σ ϵ t(1 -t), then σ ′ e (t) = σ ϵ (1 -2t)/ 2 t(1 -t) . In this section we consider simple case of separable variables π(x 0 , x 1 ) = ρ 0 (x 0 )ρ 1 (x 1 ).In the case, when ρ 0 is standard Gaussian distribution: ρ 0 = N • 0, 1 2 , we can take integral on x 0 and then take the limit η → 0 in the expressions for v and s. First, consider the expression for v: where we use explicit expression (41) for conditional density path and Eq. (42) for conditional velocity:EQUATIONwhere w(x, t | x 1 ) is the conditional velocity, generated by the conditional map ϕt,x1 (x) = σ 2 e t(1 -t) + (1 -t) 2 + tx 1 : w(x, t | x 1 ) = x 1 -x 1 -t + tσ 2 e + σ 2 e (1 -2t)x + tx 1 2 (1 -t) 2 + (1 -t)tσ 2 e .Thus, note that in the case of Gaussian distributions, all the difference between this expression and the expression without the stochastic part is the appearance of additional (time-dependent, in general) variance. Marginal distributions are still Gaussian's. Similar, using Eq. ( 46) we have for the score s:s(x, t) = (tx 1 -x)N x x 1 t, σ 2 e t(1 -t) + (1 -t) 2 ρ 1 (x 1 )dx 1 (1 -t) 2 + (1 -t)tσ 2 e N (x| x 1 t, σ 2 e t(1 -t) + (1 -t) 2 ) ρ 1 (x 1 )dx 1 = = (tx 1 -x)ρ 0 x-x1t √ σ 2 e t(1-t)+(1-t) 2 ρ 1 (x 1 )dx 1 (1 -t) 2 + (1 -t)tσ 2 e ρ 0 x-x1t √ σ 2 e t(1-t)+(1-t) 2 ρ 1 (x 1 )dx 1 . (50) E.3.To derive the analytical expression for the optimal flow velocity in the case of two normal distributions ρ 0 ∼ N (0, I) and ρ 1 ∼ N (µ, σ 2 I), we start by substituting µ 0 = 0,EQUATIONwherew(t) = tσ 2 + t -1 (1 -t) 2 + t 2 σ 2 ,and C is constant independent of x. We then redefine the dispersion based on Eq. ( 19) using x = (1 -t)x 0 + tx 1 with x 0 ∼ ρ 0 and x 1 ∼ ρ 1 :EQUATIONThis leads us to the final expression:D x,x1 ∆v(x, t) = D x0,x1 ((1 -w(t))x 1 -(1 + w(t)(1 -t))x 0 ) = = (1 + w(t)(1 -t)) 2 D x0 x 0 + (1 -w(t)) 2 D x1 x 1 .This provides a comprehensive representation of the updated dispersion for the CFM objective at any given time t.Require: Density function for initial distribution ρ 0 ; sampler for target distribution ρ 1 ; parameter M (number of samples for evaluation); parameter N (number of samples from ρ 1 for certain samples x ∼ ρ m (x, t)); optimal model v(x, t); time for evaluation t. Ensure: numerical evaluation of dispersion update for ExFM objective1: Sample (M • N ) samples x i,j 1 from ρ 1 , where i ∈ [1, M ] and j ∈ [1, N ] 2: Sample (M ) samples x i 0 from ρ 0 , where i ∈ [1, M ] 3: Compute points x i as (1 -t)x i 0 + tx i,0 1 4: Compute v d (x i , t) = N j=1ρi,j (t)x i,j 1 -x i 1-t , where ρi,j (t) = ρ 0x i -tx i,j 1 1-t / N j=1 ρ 0 x i -tx i,j 1 1-t 5: Compute and return dispersion D i (v(x i , t) -v d (x i , t))The analytical derivation of the updated dispersion for the ExFM objective proves to be complex in practice. Therefore, for the example at hand, a numerical scheme was employed for evaluation. The procedure outlined in Alg. 2 was utilized for this task. The experiment's parameters for the algorithm were as follows: M = 200k, N = 128, ρ 0 = N (0, I), ρ 1 = N (µ, σ 2 I), and the optimal model v(x, t) was derived from equation ( 54).To ensure the reliability and impartiality of the outcomes, we carried out the experiment under uniform conditions and parameters. Initially, we generated a training set of batch size N = 10,000points. The employed model was a simple Multilayer Perceptron with ReLu activations and 2 hidden layers of 512 neurons, Adam optimizer with a learning rate of 10 -3 , and no learning rate scheduler.We determined the number of iteration steps equal to 10000. Subsequently, we configured the mini batch size n = 256 during the training procedure, with the primary objective of minimizing the Mean Squared Error (MSE) loss. The full training algorithm and notations can be seen in Algorithm 1. To perform sampling, we employed the function odeint with dopri5 method from the python package torchdiffeq import odeint with atol and rtol equal 1e -5.The power dataset (dimension = 6, train size = 1659917, test size = 204928) consisted of electric power consumption data from households over a period of 47 months. The gas dataset (dimension datasets of natural images [11] .These diverse multivariate datasets are selected to provide a comprehensive evaluation of performance across various domains. To maintain consistency, we followed the code available at the given GitHub link 5 to ensure that the same instances and covariates were used for all the datasets.To ensure the correctness of the experiments we conduct them with the same parameters. The model was trained over 2,000 iterations with a batch size of 128. Inference was conducted using the Euler solver for Ordinary Differential Equations (ODE) with 100 steps. To validate the models, the POT library was employed to compute the Wasserstein distance based on 4,000 samples. The experiments were performed on a single Nvidia H100 GPU with 80gb memory.We conducted experiments related to high dimensional data, the parameters for training were taken from the open-source code 6 from the works [20, 19] . We saved the leverage of additional heuristics(EMA, lr scheduler). We use the generalized Energy Distance [18] (or E-metrics) to the metric space.Consider the null hypothesis that two random variables, X and Y , have the same probability distributions: µ = ν . The 2-Wasserstein distance [14] , also called the Earth mover's distance or the optimal transport distance W is a metric to describe the distance between two distributions, representing two different subsets A and B. For continuous distributions, it is:W := W (F A , F B ) = 1 0 F -1 A (u) -F -1 B (u) 2 du 1 2, where F A and F B are the corresponding cumulative distribution functions and F -1 A and F -1 B the respective quantile functions.To compute the NLL, we first sampled N = 5000 samples {x s i } N i=1 from the target distribution. Then we solved the following inverse flow ODE: For images evaluation we take Fréchet inception distance (FID) metrics, in particular the implementation from [12] . The main idea of FID metrics is to measure the gap between two data distributions, such as between a training set and samples from a trained model. After resizing the images, and feature extraction, the mean (µ, μ) and covariance matrix (Σ, Σ) of the corresponding features are used to compute FID:FID = ||µ -μ|| 2 2 + Tr(Σ + Σ -2(Σ Σ) 1/2 ),where T r is the trace of the matrix.Note, that w(t, x1, x) is the conditional velocity at the given point x.https://github.com/atong01/conditional-flow-matchingThe idea and common structure of the Table is taken from[20]Further reasoning is not absolutely rigorous, and in order not to introduce the axiomatics of generalized functions, we can assume that the delta function is the limit of the density of a normal distribution with mean 0 and variance tending to zero.https://github.com/gpapamak/mafhttps://github.com/atong01/conditional-flow-matching 


Decision: This paper advocates for a novel loss (Explicit Flow Matching, ExFM) with the goal of reducing variance during training. 

Below I summarize the main strengths and weaknesses of this paper. Overall, the reviewers find the goal of the paper useful, however the presented method is incremental, producing biased gradients, and the paper not sufficiently well written. For these reasons the paper cannot be accepted.  


Strengths:
+ Reducing variance in FM training is a useful goal. 
+ Importance sampling for FM loss. 

Weaknesses:
- Method proposed is incremental: the loss is not new - it is the original FM objective and the target velocity is replaced with a mini-batch estimate. 
- When velocity is estimated with mini-batch the method produce biased gradients, which can be a serious issue in high dimensions. 
- Paper is hard to read and follow.

Reviews: review: paper_summary: paper_summary: This works aims at producing a lower variance loss for flow matching. This is done by using the formula for the ground truth marginal velocity field, estimating it using self-normalized importance sampling, then regressing onto this estimated marginal velocity field.
strength_weakness: strength_weakness: - This paper lacks polish. I felt this paper is cumbersome to read, with a lot of heavy notation that can be drastically simplified, whereas the proposed algorithm is quite simple.
- The proposed approach is not exactly novel. The "Explicit Flow Matching" objective is just the original "Flow Matching" objective where we regress onto the optimal velocity field. The practical implementation being proposed is a simple application of self-normalized importance sampling.
- Importantly, the proposed approach leads to a biased objective (when estimated through a minibatch) where the optimum is not guaranteed to be the correct velocity field. This is not a problem for the low-dimensional experiments but importance sampling becomes more problematic in high dimensions.
questions: questions: The ExFM objective (Eq 8 in this paper) is the same as the FM objective [Ref. 1, Eq 5] where u_t is given by the marginal velocity field [Ref. 1, Eq 8], equivalent to Eq 10 in this paper. That this FM objective has the same gradient as CFM (Theorem 2.1 in this paper) is also stated in [Ref. 1, Theorem 2]. 

Given this, I don't think ExFM should be treated as a new loss.

The use of self-normalized importance sampling (SIS) for estimating the marginal velocity field could be interesting in its own right. However, it is misleading to state that this objective can reach zero because of the need for estimating the marginal velocity field. SIS introduces bias in favor of reducing variance. There should be some analysis on the bias introduced by the objective, comparing the learned model to the marginal velocity field.

[1] "Flow matching for generative modeling"
limitations: limitations: The proposed ExFM objective is equivalent to the original FM objective, and the paper should not over-claim this contribution. This objective is intractable, so the use of smart estimation methods is interesting in its own right.
paper_summary: paper_summary: The paper proposes a novel approach to training flow-based generative model by deriving the conditional flow matching objective function with respect to the flow function. The author argues that this new method of training will reduce variance, add stability, and ultimately lead to faster convergence. Additionally, the reformulation allow derivation of the exact vector field expression, and in some simple cases, enables the computation of the oracle trajectory solution.
strength_weakness: strength_weakness: - The paper is difficult to follow and lacks clear writing and organization. Specifically, the authors' use of notations is very confusing.
- The mathematical computations do not appear to be very rigorous and some assumptions seem very incorrect. I may have misunderstood some derivations, so please correct me if I'm wrong (see Questions section).
- The experimental results are not very robust or convincing. For instance, while the paper proposes that one of their contributions is the reduction of variance during training, many of the results demonstrate larger variance across numerous, different metrics.
- Overall, the paper feels like it requires substantial revisions and is far from being polished.
Typos:
- Line 65, rho_1$\rightarrow$ $\rho_1$.
- Line 130, practical $\rightarrow$ practical form.
questions: questions: - I find Eqn. 2 confusing. The authors should either replace the second term with the conditional vector field notation or explicitly state that the second term represents a derivative with respect to $t$.
- Eqn. 4 does not seem correct. Why can you rewrite the density function $\rho_j(x_1, x_t, t) = \rho_{x_1}(x_t, t)\rho_1(x_1)$? The distribution of the random variable $X_t$ is dependent on the distribution of $X_1$, so how can we write these terms as two independent variables.
- Also, in Eqn. 4 I believe $\rho_{x_1}(x_t, t)\rightarrow\rho_{t}(x_t, t)$. What does subscript $x_1$ mean? Does this mean it is some conditional distribution? This notation seems ambiguous.
- Why do you use the notation $p_m(\cdot)$ and not $p_t(\cdot)$. The distribution is different with respect to $t$, using $m$ makes it seems like a joint distribution of all $t$ values.
- Have the authors tried running experiment on non-Gaussian priors?
- Rather than explaining the training scheme in detail, I would suggest adding pseudo code.
limitations: limitations: The author has adequately addressed the limitations.
paper_summary: paper_summary: The paper proposes an analytic formula for the vector field satisfying the continuity equation for the given change of density which interpolates between two distributions in the sample space. This is a common setting in the Flow Matching model [6] (Rectified Flows [1], Stochastic Interpolants). 

The authors apply the formula for several special cases like linear interpolation between samples and propose to use this formula for the training. They study the proposed training procedure empirically for synthetic examples and CIFAR-10 images.

[1] Liu, Xingchao, Chengyue Gong, and Qiang Liu. "Flow straight and fast: Learning to generate and transfer data with rectified flow." *arXiv preprint arXiv:2209.03003* (2022).

[6] Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. "Flow matching for generative modeling." *arXiv preprint arXiv:2210.02747* (2022).
strength_weakness: strength_weakness: The proposed formula is already well-known in the community. It was proposed in [1] (see eq. 4, Def 3.1, Sec 5.1 for demonstration) and [2] (see Appendix D). In [3] (see eq. 4), the authors propose to use an analogous formula for training a generative model and conduct a much more exhaustive empirical study. Moreover, [4,5] develop different applied methods building upon this formula.

The downside of this formula is also very well-known in the community, i.e. one has to use large batch sizes to accurately estimate the vector field, which comes with a significant computational cost. Moreover, for any finite size, the estimation of the vector field is biased and the estimation of the loss is biased which would deteriorate generation quality when used to train large-scale models. This limitation is not adequately studied in the paper, e.g. the models are not compared in terms of the training time and memory used.

Given that the main contribution of the paper has already been proposed and the empirical study of this formula is unsatisfactory, I cannot recommend this paper for acceptance.

Minor comments:

- The authors refer to the original FM framework [6] as Conditional Flow Matching, which was proposed in [7].
- The objective proposed in Flow Matching is already tractable, so it is not clear what “a tractable form of the Flow Matching objective” means.
- There is a typo in line 65.
- The derivative notation used is inconsistent with its description in the text.
- From eq. 19, I assume that “dispersion” means “variance”.

[1] Liu, Xingchao, Chengyue Gong, and Qiang Liu. "Flow straight and fast: Learning to generate and transfer data with rectified flow." *arXiv preprint arXiv:2209.03003* (2022).

[2] Neklyudov, Kirill, Rob Brekelmans, Daniel Severo, and Alireza Makhzani. "Action matching: Learning stochastic dynamics from samples." In *International conference on machine learning*, pp. 25858-25889. PMLR, 2023.

[3] Xu, Yilun, Ziming Liu, Max Tegmark, and Tommi Jaakkola. "Poisson flow generative models." *Advances in Neural Information Processing Systems* 35 (2022): 16782-16795.

[4] Scarvelis, Christopher, Haitz Sáez de Ocáriz Borde, and Justin Solomon. "Closed-form diffusion models." *arXiv preprint arXiv:2310.12395* (2023).

[5] Xie, Tianyu, Yu Zhu, Longlin Yu, Tong Yang, Ziheng Cheng, Shiyue Zhang, Xiangyu Zhang, and Cheng Zhang. "Reflected Flow Matching." *arXiv preprint arXiv:2405.16577* (2024).

[6] Lipman, Yaron, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. "Flow matching for generative modeling." *arXiv preprint arXiv:2210.02747* (2022).

[7] Tong, Alexander, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector-Brooks, Kilian Fatras, Guy Wolf, and Yoshua Bengio. "Improving and generalizing flow-based generative models with minibatch optimal transport." *arXiv preprint arXiv:2302.00482* (2023).
questions: questions: I have no questions for the authors.
limitations: limitations: The authors do not provide a necessary literature review nor study the limitations of the proposed approach (see Weaknesses section above).
paper_summary: paper_summary: The paper proposes a loss for training flow matching (rectified flows, stochastic interoplants) that is based on integrating the target velocity over the available data as the regression target. This reduces the variance of the gradient estimator and can also be applied in the stochastic variant of flow matching.
strength_weakness: strength_weakness: I genuinely like the idea, but cannot recommend acceptance in the current presentation. I am happy to be corrected on any point and adjust my score.

Below I grouped the weaknesses by category.

## Theoretical contributions can be simplified and potentially reveal existing result (update: largely addressed, but presentation to be improved)

I think the theoretical derivation can be greatly simplified:

1. The notation is overly complicated. Why not use the standard notation for joint $p(a, b)$, conditional $p(a|b)$ and marginal $p(a)$ probabilities instead of index notation, where the index sometimes means "joint", "conditional" or "marginal", and sometimes indexes time $t$ or a condition such as $x_1$? (or $\\rho$ instead of $p$).
2. I think this reveals that the new loss is simply obtained by writing the target velocity in the ExFM loss as the expectation over the training data $\\rho_1$: Eqs (8, 10) say that the target velocity is given by $\\int w(t, x_1, x) \\rho(x|x_1, t) dx $, that is just average the velocities over the entire training data, weighted by the probability that the path actually  (where $x$ is sampled from the linear conditional paths). Given this observation, it seems to me that the actual contribution of the paper does not lie in this new loss, but how to efficiently estimate this integral, which is currently hidden in Appendix B.

In fact, I think that sections 2.1 and 2.2 can be merged to a simple importance sampling argument in the original flow matching argument.

I also think that the authors are missing that their third contribution has already been derived in the same form by their reference [10] in Eq. 4 and I think the second contribution is a simple extension to different conditional flow fields resulting from different inversions $\\varphi^{-1}$.

## Evaluation on tabular data is wrong (update: fixed)

The NLL defined in Appendix H.5.3 does not contain the volume change, which is an integral part of the negative log-likelihood. Did you use this equation for evaluation? My current score reflects the belief that volume change was accounted for.

Also, in Table 3, sometimes the highest values and sometimes the lowest values in each row are marked bold. Which model is better and does this use the incorrect formula? Please update without e-notation, adapting -1.29E+02 to -129, this is hard to read.

## Toy data evaluation (update: fixed)

It is easy to construct a very good approximation for the moons distribution by taking a Gaussian mixture of values for Table 4.

## Typos (has no influence on my recommendation)

Here is a list of what I found:

- l. 30: introduced -> introduce
- l. 33: base -> based
- l. 65: $rho$ -> $\\rho$
- l. 70: need -> needed
- l. 92/93: using map -> using the map
- l. 100: we return to end of the standard CFM loss representation -> ?
- l. 105: just (unknown) -> just the (unknown)
- eq. 7: the integral shares variables with the outside expression, e.g. add tilde on the integration variables
- l. 124: inevitable -?> invertible
- l. 126: have -> has
- eq. 16: consider moving numerical tricks like using softmax from the theory section of the paper, page 6 already introduces a lot of notation.
- throughout: dispersion -> variance
questions: questions: - Is the variance a problem in high dimensions? I would expect that overall there are few collisions of conditional flows paths.
- It is not fair to compare the MSE value of CFM and ExFM directly in Table 2, since they compute different quantities, so why list it?
- Why would the authors expect that their method performs beneficially in terms of the optimal transport distance (NPE)? OT-FM is explicitly built to reduce it, why would ExFM provide a benefit since it has the same minimum as CFM?
- What is the difference between Table 3 and 6?
limitations: limitations: I do not think that the limitations of the work are properly addressed in the theoretic part, in contrary to the statement of the authors in the paper checklist. In particular, I did not fully understand how accurately Eq. (10) (which is part of the loss) can be estimated. One question that can be a way towards addressing this is by expanding on when the assumption in line 172/173 is valid (and why this Jacobian is even a problem).

Comments: title: rebuttal acknowledgment
comment: We thank the reviewer. We acknowledge the areas identified for potential improvement and will carefully consider these points for future revisions.
We thank the reviewer for the time and effort. We are grateful for the increased score!

We will take the Reviewer’s comments into account as we work to enhance the overall clarity of the presentation. We aimed to maintain a high level of rigour in the derivations to align with the theoretical nature of the work. We understand that this might have hindered readability in some parts. In the revised version, we are going to incorporate intuitive, non-rigorous explanations (simplified derivations of the formulas) to provide a clearer overview before diving into the mathematical details.
Thank you for addressing the notation issue. While the improvement is noted, I believe the paper could benefit from further refinement in terms of overall organization, writing clarity, and consistent notation. As such, I will be maintaining my current score.
### Simplification

I read the answer as follows: The authors do not agree that the derivation can be simplified, and stick with the existing derivation although. This is in contrast to their earlier answers suggesting that they indeed found a simpler derivation:

> This algorithm itself could have been obtained in a simpler way.

> We agree that the loss itself could be obtained more easily and with less computation. We have adressed this comment in the global answer.

I would advise the authors to adopt a more direct communication style in the future.

Since I did not come up with an alternative derivation, I consider this point as addressed but strongly advise the authors to try and simplify the derivation for an improved version of the article. I think this will significantly improve the readability and chances of adoption. In my view, a promising starting point would be via identities of expectation values.

### Contributions

Thanks for this list of contributions. I agree that these results are in the paper and highlight that the two-batch-size algorithm is a useful training strategy, as evidenced by the reduced variance. I also think that the derived closed-form expressions are useful. 

### Conclusion

Overall, the paper makes a good case for two-batch-size training and has some interesting analytical results. However, the presentation also in the updated version is hard to follow. I therefore increase my score only slightly, to a borderline accept.
We thank the reviewer for pointing out weaknesses in our paper. At the same time, we respectfully disagree that our contribution consists only of an alogorithm with two different batch sizes. We listed the contribution of the paper in the first bullet in our general reply to all reviewers, and noted that it is much more than a single (already known) formula or a specific algorithm given.
We tried to address all of your concerns in our rebuttal attempt. So, if you specify which of our answers were unconvincing, it would be very helpful for us. 


We also apologize for misunderstanding with notation. The point is that we made a mistake with the word ``dash`` (not with the essence of the notations, as we thought at first). To avoid ambiguity, we replaced the corresponding text with the following "Hereinafter the symbol "${}^\prime$" indicates the time derivative:...".
Thank you for your response. My main concerns remain unaddressed. The main proposed formula is very well known and such minor contributions as using two different batch sizes for estimating two different expectations must be followed up by an extensive empirical study. As a minor comment, I suggest the authors check on the definition of dash https://en.wikipedia.org/wiki/Dash.
* ``Minor point: Jacobian independence
Can the authors say why the assumption in l. 172 is necessary and what it means?``

This condition is non-critical, it just simplifies the formula for the vector field. Consider the expression for the probability density function of the intermediate point L97 (in the old notations): 

$\rho_{x_1}(x,t)=[\phi_{t,x_1}]_*\rho_0(x):=\rho_0\big(\phi_{t,x_1}^{-1}(x)\big)\det\big[ \partial\phi_{t,x_1}^{-1}(x)\big/\partial x\big].$

When we substitute this expression into the conditional distribution density function (7), the Jacobian under consideration will be in both the numerator and the denominator. But in case it does not depend on $x_1$, in the denominator it can be taken out of the integral and it will reduce with the same expression from the numerator.
Since all important maps under consideration have this property, we considered mainly such simplified formulas (without explicitly writing out the Jacobian).
In the case of the general form of the Jacobian, we came out to the expressions like Eq. (33)--(34), where the Jacobian enters explicitly.
* ``Novelty of results
I remain confused about the similarity to existing work, so let us establish a common understanding.
The authors comment "that this formula is somewhat analogous to ours, as other similar formulas have also can be found in literature. However, we respectfully disagree that these are consistent with our contribution." So the existing results are analogous but inconsistent?
Starting from the list of contributions in the Introduction, can the authors for each theoretical contribution (i) give the central lines where this contribution is achieved and (ii) what the closest known result is, together with a list of novel contributions?``

Let us go through our contributions.

1. ``A tractable form of the FM loss is presented, which reaches...., but has a smaller variance``
To the best of our knowledge, the result in the form of loss, which could reach zero and which would contain integral expressions from distribution functions, has not been found in the literature (unlike the expression for the vector field, see the next point below). So, apparently, the closest result that corresponds to this contribution is the already mentioned paper by Lipman et al, Eq. (6), about which the authors also write "Flow Matching is a simple and attractive objective, but na¨ıvely on its own, it is intractable to use in practice since we have no prior knowledge for what an appropriate $p_t$ and $u_t$ are."
In addition, the rigorous analysis of variance that is mentioned in this contribution has not been encountered either.

2. ``The explicit expression in integral form for the vector field delivering the minimum to this
loss (therefore for Flow Matching loss) is presented.``
Partial cases, or analogs our Eq. (16), have been in various sources. The closest to ours is probably Eq. (4) and the unnumbered formula from Sec. 5.1 Liu et al, "Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow.". 
Also a formula similar to our Eq. (17) with softmax function is found in Scarvelis et al. "Closed-form Diffusion Models", Eq. (2) and others.
But we note that in all the works in which similar formulas are found a) this formulas are not strictly derived b) was given one specific formula; on the contrast, we have proposed a whole family of formulas for the vector field (and score), starting from the case where we evaluate the integral not with SIS but with rejection sampling (Appendix B) and ending with formulas for DDPM-like models, i.e., for ODEs containing a stochastic term, namely the expression for the vector field in Eq (44) or the expression for the score (46). We have not encountered such expressions in the literature.

3. ``As a consequence, we derive expressions for the flow matching vector field and score in several particular cases (when linear conditional mapping is used, normal distribution, etc.);``
We refer here, for example, to the results presented in Fig. 2, Figs. 4--6 (together with the analytical expressions by which they were obtained). We have not encountered similar formulas, similar studies on, for example, Gaussian Mixure separation we cited in the introduction ([15, 8]), but these studies are far from ours in terms of ideology.

4. ``Analytical analysis of SGD convergence showed that our formula have better training variance on several cases``.
In this case, we are talking about Theorems ans resultd from Sec. 2.4 "Irreducible dispersion of gradient for CFM optimization". To the best of our knowledge, no similar studies have been conducted.

5. ``Numerical experiments show that we can achieve better learning results in fewer steps``.
The topic of variance reduction in Flow-Matching-like models is raised most closely in the series of papers by Tong et al. on OT-CFM, where this heuristic technique was created specifically for practical implementation. We are compared with this technique in experiments in different dimensions. It turns out that in many cases we are not inferior to the results of this heuristic. However, our algorithm was rather a side effect of our theoretical calculations. It may well be accelerated if we use more efficient techniques (besides SIS or rejection sampling) for evaluating the integral.
We thank Reviewer for additional important questions and for the discussion. We will answer them one by one.

* ``Simplification
I cannot find an answer to my claim that the ExFM loss can be derived significantly simpler in the general answer. Can the authors please address my question explicitly? Since the authors agree that the derivation can be greatly simplified, I would in addition appreciate a high-level overview over the simplified variant and a proposed structure for the revised paper (in case the authors have already updated the manuscript available, it is often possible to post an updated PDF under anonymous link; ask the area chair for guidance).``

If one interested only in heuristics and want to derive some formula or method
only for the practical verification, then he could go down the following path.
Let us take the Eq. (8) from Lipman et al. "Flow Matching for Generative Modeling".
Formally, we may substitute $p_t(x|x_1)$ from Eq. (10) and $u_t(x|x_1)$ from Eq. (15)
from the same paper to this expression
to get a special case of our formula.
This is a straightforward way to derive our algorithm in the case of a Gaussian initial distribution.
Note that this is highly intuitive reasoning. 

In contrast, 
in our work we used the following strict reasoning (somewhere implicitly).
Our conclusions are based on

a) Theorems from the cited Lipman et al on the equivalence of CFM and FM gradients (Theroems 1--2, Appendix A) 

b) theorems from Tong et al, "Improving and Generalizing Flow-Based Generative Models with Minibatch Optimal Transport" 

which extend the previous result to more general cases (Theorems 3.1--3.4, Table 1).
Thus, these results allow us to start with the losses that are in CFM (in various modifications).
Then we consider an _invertible map_ of type $\phi_{t,x_1}(x_0)=(1-t)x_0+tx_1+\sigma_stx_0$ and in order to proceed to the case of a simple _non-invertible_ map $\phi_{t,x_1}(x_0)=(1-t)x_0+tx_1$ we strictly consider the limit $\sigma_s\to0$. Thus, we maintain the rigor of the reasoning of a theoretical paper.

Since there are a sufficient number of possible modifications of CFM (listed in Table 1 of Tong et al. or in Table 1 of our paper), and there are also many special cases, such as specific initial and target distributions as well as a stochastic modification of the equation in FM that is consistent with DDPM, we could not insert all these combinations in our paper.
Rather, we have described in some detail a method that can be used to further obtain similar results to those we have already given.
But to show the importance of the method, we have given some corollaries of the loss representation in a new form. For example a) a rigorous proof of variance reduction b) the comparability of numerical results of our algorithm, which directly follows from such a loss representation and the results of OT-CFM, which is sharpened for variance reduction c) an explicit record of trajectories for cases of Gaussian distributions (Eq. (36)), explicit recording of velocity for the case of Gaussian Mixture (Eq. (37)), a record of the score in integral form (Eq. (46)), the possibility of applying this technique to cases where both distributions are unknown and there are only samples from both distributions (Eq. (51)--(52)), etc.
All of these results were obtained in one way or another from various loss modifications, similar to the one in native CFM, using our method described in detail in the main body of the paper.

Because we hope that these results do not exhaust all possible applications of our technique, and because of page limitations, we have moved the specific results to Appendix, focusing in the main body of the paper on the method of obtaining them.

Thus, in the main text we have detailed calculations so that it is clear how we get the result step by step.

In the modified version of the article we added additional explanations that we discussed with other reviewers, tables, some figures we discussed in global answer and changed notations. The structure and the simplification of the article has not changed dramatically. We maintain all the nuances of the formulas' derivations as all these steps are the part of the strict analysis, which is also the aim of the paper. Revised version of the paper can be find in this anonymous link https://drive.google.com/file/d/17HDXtwWB505revPOXisi1F2lprECwUUu/view?usp=sharing.
I thank the authors for the comments. However, I still have concerns regarding central points: I think that their method can be derived significantly simpler, and I am confused about the novelty of some results.

- `Simplification`

I cannot find an answer to my claim that the ExFM loss can be derived significantly simpler in the general answer. Can the authors please address my question explicitly? Since the authors agree that the derivation can be greatly simplified, I would in addition appreciate a high-level overview over the simplified variant and a proposed structure for the revised paper (in case the authors have already updated the manuscript available, it is often possible to post an updated PDF under anonymous link; ask the area chair for guidance).

- `Novelty of results`

I remain confused about the similarity to existing work, so let us establish a common understanding.

The authors comment "that this formula is somewhat analogous to ours, as other similar formulas have also can be found in literature. However, we respectfully disagree that these are consistent with our contribution." So the existing results are analogous but inconsistent?

Starting from the list of contributions in the Introduction, can the authors for each theoretical contribution (i) give the central lines where this contribution is achieved and (ii) what the closest known result is, together with a list of novel contributions?


- Minor point: `Jacobian independence`

Can the authors say why the assumption in l. 172 is necessary and what it means?