Title: Explicit Flow Matching: On The Theory of Flow Matching Algorithms with Applications
Venue: NeurIPS 2024
Accepted: no
Keywords: Flow Matching

Deep Learning Theory

Generative modeling

Variance Reduction

Stochastic Differential Equations
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.

1

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.

1 Introduction

In recent years, there has been a remarkable surge in Deep Learning, wherein the advancements 16 have transitioned from purely neural networks to tackling differential equations. Notably, Diffusion 17 Models [16] have emerged as key players in this field. This models transform a simple initial 18 distribution, usually a standard Gaussian distribution, into a target distribution via a solution of 19 Stochastic Differentiable Equation (SDE) [1] or Ordinary Differentiable Equation (ODE)[2] with 20 right-hand side representing a trained neural network. The Conditional Flow Matching (CFM) [9] 21 technique, which we focus on in our research, is a promising approach for constructing probability 22 distributions using conditional probability paths, which is notably a robust and stable alternative for 23 training Diffusion Models. The development of the CFM-based approach includes various techniques 24 and heuristics [4, 7, 13] aimed at improving convergence or quality of learning or inference. For 25 example, in the works [19, 20, 10] it was proposed to straighten the trajectories between points by 26 different methods, which led to serious modifications of the learning process. We refer the reader 27 for, example, to the paper [20] where different FM-based approaches are summarised, and to the 28 paper [9] for the connection between Diffusion Models and CFM. 29

In our work, we introduced an approach which we called Explicit Flow Matching (ExFM), to consider 30

the Flow Matching framework theoretically by modifying the loss and writing the explicit value of 31

the vector field. Strictly speaking, the presented loss is a tractable form of the FM loss, see Eq. (5) 32

of [9]. Base on this methods we can improve the convergence of the method in practical examples 33

reducing the variance of the loss, but the main focus of our paper is on theoretical derivations. 34

Our method allows us to write an expression for the vector field in closed form for quite simple 35 cases (Gaussian distributions), however, we note that Diffusion Models framework in the case of 36 a Gaussian Mixture of two Gaussian as a target distribution is still under investigation, see recent 37 publications [15, 8]. 38

39

40

41

42

43

44

45

46

47

48

49

Our main contributions are:

Atractable 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;

The explicit expression in integral form for the vector field delivering the minimum to this loss (therefore for Flow Matching loss) is presented.

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.);

Analytical analysis of SGD convergence showed that our formula have better training variance on several cases;

Numerical experiments show that we can achieve better learning results in fewer steps.

1.1 Preliminaries

Flow matching is well known method for finding a flow to connect samples from two distribution 50 with densities ρ 0 and ρ 1 . It is done by solving continuity equation with respect to the time dependent 51 vector field v ( x, t ) and time-dependent density ρ ( x, t ) with boundary conditions: 52

Function ρ ( x, t ) is called probability density path . Typically, the distribution ρ 0 is known and it is 53 chosen for convenience reasons, for example, as standard normal distribution ρ ( x ) = N ( x | 0 , I ) . 54 The distribution ρ 1 is unknown and we only know the set of samples from it, so the problem is to 55 approximate the vector field v ( x, t ) ≈ v ( x, t ) using these samples. To make problem (1) well defined, 56 one usually imposes additional regularity conditions on the densities, such as smoothness. The 57 rigorous justification of the obtained results we put in the Appendix, leaving the general formulations 58 of theorems and ideas in the main text. 59

60

61

62

63

64

65

66

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.

67

68

69

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 θ

where ϕ t,x 1 ( x 0 ) is some flow, conditioned on x 1 (one can take ϕ t,x 1 ( x 0 ) = (1 -t ) x 0 + tx 1 + σ s tx 0 70 in the simplest case, where σ s > 0 is a small parameter need for this map to be invertable at 71 any 0 ≤ t ≤ 1 ). Hereinafter the dash indicates the time derivative. Time variable t is uniformly 72 distributed: t ∼ U [0 , 1] and random variables x 0 and x 1 are distributed according to the initial and 73 final distributions, respectively: x 0 ∼ ρ 0 , x 1 ∼ ρ 1 . Below we omit specifying of the symbol E the 74 distribution by which the expectation is taken where it does not lead to ambiguity. 75

1.2 Why new method? 76

Model training using loss (2) have the following disadvantage: during training, due to the randomness 77 of x 0 and x 1 , significantly different values can be presented for model as output value at close model 78

CFM

ExFM

ExFM

P1

P1

y(xt,t|x1)

u(x

CFM

x1~p1(x1)

1Dx,x△v(x,t）)1

[0x]

[0xo]

Po

P

10x1

0

t

argument values ( x t , t ) . Indeed, a fixed point x t = ϕ t,x 1 ( x 0 ) can be obtained by an infinite set of x 0 79 and x 1 pairs, some of which are directly opposite, and at least for small times t the probability of these 80 different directions may not be significantly different. At the same time, data ϕ ′ t,x 1 ( x 0 ) on which the 81 model learns significantly different for such different positions of pairs x 0 and x 1 . Thus, the model is 82 forced to do two functions during training: generalize and take the mathematical expectation (clean 83 the data from noise). 84

90

91

2 Main idea

2.1 Modified objective

Lets expand the last two mathematical expectations in the loss (2) and substitute variables using 92 map ϕ t,x 1 , passing from the point x 0 to its position x t = ϕ t,x 1 ( x 0 ) at time t : 93

We assume, that the map ϕ t,x 1 is invertible at each 0 < t < 1 , i. e. that ϕ -1 t,x 1 ( x t ) exits on this 94 time interval and for all x t = { ϕ t ( x 0 ) | ∀ x 0 : ρ ( x 0 ) > 0 } . Eq. (3) can be seen as a transition 95 from expectation on the variable x 0 ∼ ρ 0 to expectation on the variable x t ∼ ρ x 1 ( · , t ) , where 96 ρ x 1 ( x, t ) = [ ϕ t,x 1 ] ∗ ρ 0 ( x ) := ρ 0 ( ϕ -1 t,x 1 ( x ) ) det [ ∂ϕ -1 t,x 1 ( x ) / ∂x ] . See paper [5] for details about the 97

To obtain the modified loss, we return to end of the standard CFM loss representation in (3). It is 100 written as the expectation over two random variables x 1 and x t having a common distribution density 101

which, 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:

102

103

104

while the marginal distribution of x 1 is just (unknown) function ρ 1 . Let for convenience w ( t, x 1 , x ) = 105 ϕ ′ t,x 1 ( ϕ -1 t,x 1 ( x ) ) 1 . We have 106

107

where we introduce a conditional distribution

The key feature of the representation (6) is that the integration variables x 1 and x are independent. 108 Thus, we can evaluate them using Monte Carlo-like schemes in different ways. However, we go 109 further and make a modification to this loss to reduce the variance of Monte Carlo methods. 110

111

112

113

114

115

116

117

118

119

120

121

2.2 New loss and exact expression for vector field

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 loss

Theorem 2.1. Losses L CFM in Eq. (2) and L ExFM in Eq. (8) have the same derivative with respect to model parameters:

Proof 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 have

We can obtain the exact form of this vector field given the particular map ϕ t,x 1 . For example, the 122 following statement holds: 123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

Corollary 2.2. Consider the linear conditioned flow ϕ t,x 1 ( x 0 ) = (1 -t ) x 0 + tx 1 which is inevitable as 0 ≤ t < 1 . Then w ( t, x 1 , x ) = x 1 -x 1 -t , ρ x 1 ( x, t ) = ρ 0 ( x -x 1 t 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 form

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 simpler

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,x 1 = (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)

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

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.

Simple examples Consider the case of Standard Normal Distribution as ρ 0 and Gaussian Mixture of two Gaussians as ρ 1 . Vector field have a closed form (37) in this case, and we can fast numerically solve ODE for trajectories. Random generated trajectories and plot of the vector field are shown on Fig. 2 (a)-(b). Detailed explanation of this case is in the Sec. D.2. Another example is related to the case of a stochastic map in the form of Brownian Bridge, which briefly described in the last paragraph and considered in Sec. E.3.2 in details, see Fig. 2 (c)-(f). Note that at some σ e values the trajectories are a little bit straightened in this case compared to the usual linear map, if we compare cases on the Fig. 6.

2.3 Training scheme based on the modified loss

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

161

162

163

164

165

166

167

168

169

170

171

172

173

0.2

0.4

0.6

0.8

1.0

t

-

3

-

2

-

1

1

2

3

x

(

t

)

(a)

GM trajecto-

ries

(b) GM VF

0.2

0.4

0.6

0.8

1.0

t

-

5

5

10

x

(

t

)

(c)

BB Trajecto-

ries,

σ

e

= 3

(d) BB VF,

σ

e

= 3

0.2

0.4

0.6

0.8

1.0

t

-

5

5

10

x

(

t

)

(e)

BB Trajecto-

ries,

σ

e

= 10

(f) BB

VF,

σ

e

= 10

the unknown distribution ρ 1 . Then we pair them (according to some scheme), and get n samples as x 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 samples

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)

For example, if we use self-normalized importance sampling and assume that the Jacobian det [ ∂ϕ -1 t,x 1 ( x ) / ∂x ] do not depend on x 1 , we can write

Theorem 2.4. Under mild conditions, the error variance of the integral gradient (9) using the Monte 174 Carlo method (14) is lower than using formula (13) with the same number n · m of samples for { x } . 175

176

177

Sketch of the proof is in the Appendix A.2. The steps of our scheme are formally summarized in Algorithm 1.

178

179

180

181

182

183

184

185

186

187

Particular case of linear map and Gaussian noise Let ϕ t,x 1 be the linear flow: ϕ t,x 1 ( 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 have

Here, 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.

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

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.

2.4 Irreducible dispersion of gradient for CFM optimization

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):

For 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:

where ˆ x i 0 = x j,i -t j x i 1 1 -t j . We define the dispersion D x,x 1 f ( x, x 1 ) for x ∼ ρ x 1 ( · , t ) and x 1 ∼ ρ 1 as:

Proposition 2.5. At the time t = 0 , the dispersion of update in the form (18) have the following element-wise lower bound:

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:

where x 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:

Equality is reached when the model v ( x j,i , 0) has exact values equal to (12) . 212

In comparison to CFM, the dispersion of the update is N times smaller than the dispersion of the 213 target distribution and could be controlled without impeding convergence by adjusting the number of 214 samples N . In Figure 1(b), we visually compare the dispersions of CFM and ExFM. The illustration 215 aligns a standard normal distribution N (0 , I ) with a shifted and scaled variant N ( µ, Iσ 2 ) . ExFM 216 yields lower dispersion throughout the range t ∈ [0 , 1] . Detailed analytical calculations of the optimal 217 velocity v ( x, t ) and dispersion are provided in the Appendix G. 218

(a) swissroll

(b) moons

(c) circles

(d) 2spirals

(e)

checkerboard

(f) pinwheel

(g) rings

S

S

O

S

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

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

Tabular data We conducted unconditional density estimation on five tabular datasets, namely power , gas , hepmass , minibone , and BSDS300 . Additional details of the experiments see in the Appendix H. The empirical findings obtained from the numerical experiments from Table 3 indicate a statistically significant improvement in the performance of our proposed method. Notably, ExFM demonstrates a notable acceleration in convergence rate.

High-dimensional data and additional experiments We conducted experiments on highdimensional data, among them experiments on CIFAR10 and MNIST dataset. FID results on CIFAR10 shows slightly better score among sampled images.

Additional details of the experiments and sampled images see in the Appendix H.

Stochastic ExFM (ExFM-S) on toy 2D data We evaluated the performance of the stochastic version of ExFM (ExFM-S) with use of expressions given in Sec. E.3.2 on four standard toy datasets. 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.

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

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

References

Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Eijnden. 'Stochastic Interpolants: A Unifying Framework for Flows and Diffusions'. In: arXiv preprint 2303.08797 (2023).

Michael S. Albergo and Eric Vanden-Eijnden. 'Building Normalizing Flows with Stochastic Interpolants'. In: International Conference on Learning Representations (ICLR) (2023).

Gabriel Cardoso et al. 'BR-SNIS: Bias Reduced Self-Normalized Importance Sampling'. In: Advances in Neural Information Processing Systems . Ed. by S. Koyejo et al. Vol. 35. Curran Associates, Inc., 2022, pp. 716-729. URL: https://proceedings.neurips.cc/ paper_files/paper/2022/file/04bd683d5428d91c5fbb5a7d2c27064d- PaperConference.pdf .

Ricky T. Q. Chen and Yaron Lipman. 'Riemannian Flow Matching on General Geometries'. In: arXiv:2302.03660 (2023).

Ricky T. Q. Chen et al. 'Neural Ordinary Differential Equations'. In: Advances in Neural Information Processing Systems . Ed. by S. Bengio et al. Vol. 31. Curran Associates, Inc., 2018. URL: https://proceedings.neurips.cc/paper_files/paper/2018/file/ 69386f6bb1dfed68692a24c8686939b9-Paper.pdf .

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models . 2020. arXiv: 2006.11239 [cs.LG] .

Alexia Jolicoeur-Martineau, Kilian Fatras, and Tal Kachman. 'Generating and Imputing Tabular Data via Diffusion and Flow-based Gradient-Boosted Trees'. In: arXiv:2309.09968 (2023). arXiv: 2309.09968 [cs.LG] .

Puheng Li et al. 'On the Generalization Properties of Diffusion Models'. In: Advances in Neural Information Processing Systems . Ed. by A. Oh et al. Vol. 36. Curran Associates, Inc., 2023, pp. 2097-2127. URL: https://proceedings.neurips.cc/paper_files/paper/ 2023/file/06abed94583030dd50abe6767bd643b1-Paper-Conference.pdf .

Yaron Lipman et al. 'Flow Matching for Generative Modeling'. In: The Eleventh International Conference on Learning Representations . 2023. URL: https://openreview.net/forum? id=PqvMRDCJT9t .

Xingchao Liu, Chengyue Gong, and qiang liu. 'Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow'. In: The Eleventh International Conference on Learning Representations . 2023. URL: https://openreview.net/forum?id=XVjTT1nw5z .

D. Martin et al. 'A database of human segmented natural images and its application to evaluating segmentation algorithms and measuring ecological statistics'. In: Proceedings Eighth IEEE International Conference on Computer Vision. ICCV 2001 . Vol. 2. 2001, 416-423 vol.2. DOI: 10.1109/ICCV.2001.937655 .

Gaurav Parmar, Richard Zhang, and Jun-Yan Zhu. 'On Aliased Resizing and Surprising Subtleties in GAN Evaluation'. In: CVPR . 2022.

Aram-Alexandre Pooladian et al. 'Multisample Flow Matching: Straightening Flows with Minibatch Couplings'. In: Proceedings of the 40th International Conference on Machine Learning . Ed. by Andreas Krause et al. Vol. 202. Proceedings of Machine Learning Research. PMLR, July 2023, pp. 28100-28127. URL: https://proceedings.mlr.press/v202/ pooladian23a.html .

Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. 'On wasserstein two-sample testing and related families of nonparametric tests'. In: Entropy 19.2 (2017), p. 47.

Kulin Shah, Sitan Chen, and Adam Klivans. 'Learning Mixtures of Gaussians Using the DDPM Objective'. In: Thirty-seventh Conference on Neural Information Processing Systems . 2023. URL: https://openreview.net/forum?id=aig7sgdRfI .

Jascha Sohl-Dickstein et al. 'Deep Unsupervised Learning using Nonequilibrium Thermodynamics'. In: Proceedings of the 32nd International Conference on Machine Learning . Ed. by Francis Bach and David Blei. Vol. 37. Proceedings of Machine Learning Research. Lille, France: PMLR, July 2015, pp. 2256-2265. URL: https://proceedings.mlr.press/v37/ sohl-dickstein15.html .

Yang Song and Stefano Ermon. 'Generative Modeling by Estimating Gradients of the Data Distribution'. In: Neural Information Processing Systems (NeurIPS) (2019).

Gábor J Székely. 'E-statistics: The energy of statistical samples'. In: Bowling Green State University, Department of Mathematics and Statistics Technical Report 3.05 (2003), pp. 1-18.

Alexander Tong et al. 'Improving and generalizing flow-based generative models with mini323 batch optimal transport'. In: Transactions on Machine Learning Research (2024). Expert 324 Certification. ISSN: 2835-8856. URL: https://openreview.net/forum?id=CD9Snc73AW . 325

Alexander Tong et al. 'Simulation-Free Schrödinger Bridges via Score and Flow Matching'. 326 In: The 27th International Conference on Artificial Intelligence and Statistics . 2024. URL: 327 https://virtual.aistats.org/virtual/2024/poster/6691 . 328

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Theoretical things are proved in theorems in the main text and in the Appendix, and numerical experiments have been performed

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.

2. Limitations

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.

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

4. Experimental Result Reproducibility

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

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

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

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

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

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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

8. 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).

9. 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 ?

Answer: [Yes]

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

10. 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).

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

12. 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 work

Guidelines:

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.

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

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

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

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

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

641

642

643

644

645

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 Proof of the theorems 646

647

A.1 Proof of the Theorem 2.1

648

649

650

Proof. We need to proof, that d L CFM ( θ ) d θ = d L ExFM ( θ ) 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):

Utilizing the bilinearity of the 2-norm, we can rewrite L CFM as: 651

Here, T denotes transposed vector, dot denotes scalar product, C represents a constant independent 652 of θ . 653

Noting that E x 1 ρ c ( x | x 1 , t ) /ρ 1 ( x 1 ) = 1 : 654

we can simplify the first term in the expansion (21): 655

For our loss L ExFM in the form (8) we also use the bilinearity of the norm: 656

657

658

659

660

661

662

663

664

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

Proof. We need to prove that D d L d ExFM ( θ ) d θ ≤ D d L 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:

Note that in this expression, values x j,i as well as t j , which are included in the argument of the 665 function v , are fixed (our goal to calculate the variance with fixed model arguments). Thus, we need 666 to consider the variance of the remaining expression arising from the randomness of x k 1 . 667

Recall (below we will omit the indices at variables x and t ),

Note, that if N = 1 , ( i. e. we do not sample any additional points other than the ones we have 668 already sampled) this expression is exactly the same as the derivative of the common discretized 669 CFM loss d L d CFM ( θ ) d θ . 670

Moreover, recall that one of the points (without loss of generality, we can assume that its index is 1) 671 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 672 to x 1 1 ) 673

Thus, we can rewrite expression for v d : 674

Thus, our task was reduced to evaluating how well the additional terms (for k starting from 2 ) improve 675 approximate of the original integrals that are in loss (8). 676

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:

The smaller coefficient k D is, the better the proposed loss ExFM works. 677

Formally, we can write our problem as an importance sampling problem for the following integral:

This integral we estimate by sample mean of the following expectation over some random variable with density function q ( x ) :

with

We replace the exact value of I with the value

It follows from the strong law of large numbers that in the limit N →∞ , I → I almost surely. From 678

the central limit theorem we can find the asymptotic variance: 679

In our case (loss L ExFM), we have q ( x 1 ) = ρ 1 ( x 1 ) , f ( x 1 ) = w ( t, x 1 , x ) and w ( x 1 ) = ρ 0 ( ϕ -1 t,x 1 ( x ) ) . 680

Despite the fact that the equation (25) for the variance contains N in the denominator, it is rather 681 682

683 times t noticeably less than 1 . 684

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

But in the case, when t is closed to 1 we have, for example, for the linear map, that

and this function has a sharp peak near the point x/t if it is considered as a function of x 1 . Thus, at 685 such values of t , only a small number of summands will give a sufficient contribution to the sum 686 compared to the first term. 687

Finally, inequality k D < 1 is formally fulfilled, but how much k D is less than one depends on many 688 factors. 689

690

691

692

693

694

695

696

697

A.3 Expressions for the regularized map

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,x 1 ( 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,x 1 ( 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 true

By direct substitution we make sure that for this vector field 698

and where we perform change of the variables y ← x 1 -x σ s t .

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 :

and

It follows that both integrals in expression (26) converge absolutely and uniformly. So, we can swap 705 the operations of taking the limit and integration, and we can take the limit σ s → 0 in the integrand 706 for any time t ∈ [0 , t 0 ] for arbitrary t 0 < 1 . 707

Now, let us consider the case t = 1 . From Assumption A.1 the boundedness of the density ρ 1 follows: ρ 1 ( x ) < C 2 , ∀ x ∈ ( -∞ , ∞ ) . Thus, integrand functions in the numerator and denominator in the Eq. (28) can be bounded by the following integrable functions independent of σ s :

and

The existence of the limit

follows from Assumption A.1. 708

Finally, we conclude that formula (11), regarded as the limit σ s → 0 of the (26) at any t ∈ [0 , 1] , is 709 true. 710

699

700

701

702

703

704

711

712

713

714

715

716

Theorem A.3. The vector field in Eq. (11) delivers minimum to the Flow Matching objective (see the work [9]),

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 ) 717

where p t ( x | x 1 , σ s ) is conditional probability paths obtained by regularized linear conditional map. 718 Distribution p t in the time t = 0 is equal to standard normal distribution p 0 ( x | x 1 , σ s ) = N ( x | 0 , 1) 719 and at the time t = 1 it is a stretched Gaussian centered at x 1 : p 1 ( x | x 1 , σ s ) = N ( x | x 1 , σ s I ) . 720

Substituting p 1 into the Eq. (29) and considering that there exists a limit σ s → 0 due to Assumption A.1, we obtain

This finish the proof. 721

722

A.3.2 Learning procedure for σ s > 0

Using standard normal distribution as initial density ρ 0 , and the regularized map ϕ t,x 1 ( x 0 ) = 723 (1 -t ) x 0 + tx 1 + σ s tx 0 we obtain the following approximation formula 724

In practical applications, the exponent calculation is replaced by the SoftMax function calculation, 725 which is more stable. 726

727

B Estimation of integrals

In general, we need to estimate the following expression 728

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 .

729

730

731

732

733

This estimate is biased in theory, but there several methods to reduce this bias and improve this 735 estimate, see, for example, [3]. Our numerical experiments generally show that the estimation (30) in 736 the form is already sufficient for stable results; we don not observe any bias. 737

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

then we put the point x j k to the set Y ; otherwise we reject it. 738

Having formed the set Y , we evaluate the integral as 739

To justify the last estimation, we note, that the points from the set Y are distributed according 740 to (non-normalized) density ρ ( x ) f ( x, η ) ρ 1 ( x ) . One can show it using the proof of the rejection 741 sampling method. This is the same density as in Eq. (7) and thus we estimate the expression (10) 742 using Important Sampling without any additional denominator. 743

744

745

746

747

748

749

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 .

750

751

752

753

754

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 .

C The main Algorithm and extensions and generalization of the exact expression

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 do

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

6: 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 θ .

8: Update model parameters θ ← SGD ( θ, G )

9: end while

General form of the proposed Algorithm is given in Alg 1. 755

When using other maps, formula (11) is modified accordingly. For example, if we use the regularized 756 map ϕ t,x 1 ( x 0 ) = (1 -t ) x 0 + tx 1 + σ s tx 0 , we get the formula (26).Note, that in this case the 757 final density ρ ( x, 1) , obtained from the continuity equation is not equal to ρ 1 , but is its smoothed 758 modification. 759

When using a different initial density ρ 0 (not the normal distribution), an obvious modification will 760 be made to formula (16). 761

Diffusion-like models We can treat so-called Variance Preserving [6] model as CFM with the map

and ρ 0 as standard normal distribution: ρ 0 = N ( · ∣ ∣ 0 , 1 2 ) In this case, the common expression (10) 762 for vector filed transforms to 763

where α ′ s = d α s d s . 764

Similarity we can treat so-called Variance Exploding [17] model as CFM with the map

and ρ 0 also as standard normal distribution: ρ 0 = N ( · ∣ ∣ 0 , 1 2 ) In this case, the common expres765 sion (10) for vector filed transforms to 766

where σ ′ s = d σ s d s . 767

Joint Distribution Moreover, in addition to the independent densities x 0 ∼ ρ 0 and x 1 ∼ ρ 1 , we 768 can use the joint density { x 0 , x 1 } ∼ π ( x 0 , x 1 ) . In the papers [20, 19], optimal transport (OT) 769 and Schrödinger's bridge are taken as π . In this case the expression for the vector field changes 770 insignificantly: the conditional probability ρ c from Eq. (7) is subject to change: 771

Then, Eq. (10) remains the same in general case. In the case of linear ϕ , the extension of Eq. (11) 772 reads 773

In all of the above cases, the essence of Algorithm 1 does not change (except that in the case of 774 dependent x 0 and x 1 we should be able either to calculate the value of π ( ϕ -1 t,x 1 ( x ) , x 1 ) /ρ 1 ( x 1 ) or to 775 estimate it). 776

D Several analytical results, following from the explicit formula 777

778

779

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.

780

D.1 Exact path from one Gaussian to another Gaussian

Consider the flow from a one-dimensional Gaussian distribution ρ 0 ∼ N ( · ∣ ∣ µ 0 , σ 2 0 ) into another (with 781 other parameters) Gaussian distribution ρ 1 ∼ N ( · ∣ ∣ µ 1 , σ 2 1 ) . Note that in this case the generalization 782 to the multivariate case is done directly, so the spatial variables are separated. 783

From the general formula (11) we have: 784

Both integrals in the last expression are taken explicitly: 785

Note that the last relation can be obtained as a distribution of two Gaussian random variables with 786 corresponding parameters. 787

The second integral: 788

Thus, in the considered case we can explicitly write the expression for the vector field v : 789

For this vector field we can explicitly solve the equation for the path x ( t ) starting from the arbitrary 790 point x 0 791

Note that although this solution does not correspond to the Optimal Transport joint distribution, since 793 the obtained path is not a straight line in general, ( i. e. we do not have a solution to the Kantorovich's 794 formulation of the OT problem) the endpoint x (1) = µ 1 +( x 0 -µ 0 ) σ 1 σ 0 falls exactly in the one that 795 is optimal if we solve the OT problem in the Monge formultation. Thus, the map x (0) → x (1) is the 796 OT map for the case of 2 Gaussian. 797

See the Fig. 4 for the examples of the paths for the obtained solution. 798

799

D.2 From one Gaussian to Gaussian Mixture

Let initial distribution be standard Gaussian ρ 0 = N ( · ∣ ∣ 0 , 1 2 ) , and the target distribution be Gaussian 800 Mixture (GM) of two symmetric Gaussians: ρ 1 ( x ) = 1 / 2( N ( x ∣ ∣ µ, σ 2 ) ) + N ( x ∣ ∣ -µ, σ 2 ) ) , In this 801

The solution is: 792

0.2

0.4

0.6

0.8

1.0

t

-

2

2

4

6

8

x

(

t

)

(a) Trajectories

(b) Vector field

10

5

0

0

-5

0.0

0.2

0.4

0.6

0.8

1.0

case, we can obtain exact form for v 802

but the expression for the path x ( t ) is unknown. 803

0.2

0.4

0.6

0.8

1.0

t

-

3

-

2

-

1

1

2

3

x

(

t

)

(a) Trajectories

(b) Vector field

6

4

2

2

0

0

2

- 4

6

0.0

0.2

0.4

0.6

0.8

1.0

806

807

808

809

810

811

812

813

D.3 From Gaussian to Gaussian with stochastic

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.

D.3.1 Gaussian to Gaussian with noise

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.

We can solve ODE with this field and get the expression for the trajectories, starting from the given 814 point x 0 : 815

These trajectories, for different x 0 are depicted in Fig. 6.

0.2

0.4

0.6

0.8

1.0

t

-

5

5

10

x

(

t

)

(a) Trajectories,

σ

e

= 0

.

3

(b)

Vector

field,

σ

e

= 0

.

3

0.2

0.4

0.6

0.8

1.0

t

-

5

5

10

x

(

t

)

(c) Trajectories,

σ

e

= 1

(d)

Vector

field,

σ

e

= 1

0.2

0.4

0.6

0.8

1.0

t

-

5

5

10

x

(

t

)

(e) Trajectories,

σ

e

= 3

(f)

Vector

field,

σ

e

= 3

0.2

0.4

0.6

0.8

1.0

t

-

5

5

10

x

(

t

)

(g) Trajectories,

σ

e

= 10

(h)

Vector

field,

σ

e

= 10

816

At the limit σ e → 0 expressions (38) and (39) turn into expressions (35) and (36) as expected. 817

For the score s in the considered case we have 818

Thus, we can explicitly write expressions for the stochastic process for the evolution from the initial 819 distribution rho 0 (standard Gaussian) to the final distribution ρ 1 : 820

Here g ( t ) is arbitrary smooth function. In the case of Shrödinger Bridge we take g ( t ) = σ e √ t (1 -t ) . 821

E Detail on the SDE case 822

823

824

825

826

827

828

829

830

831

832

833

834

835

836

837

838

839

840

841

E.1 Optimal vector field and score for stochastic map

Following [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 ) as

Thus, the conditional distribution p ( t, x | x 0 , x 1 ) conditioned on the starting x 0 and end point x 1 is Gaussian:

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 form

where parameter η is small enough. Then we can use the above Theorems and immediately write

After 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:

where π ( x 1 , x 0 ) is the density of the joint distributions with the marginal equal to the two given probabilities:

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 842 minimum by performing the same calculations as in the derivation of formula (10): 843

As in the work [20] we can also train score network. Namely, as marginals for Brownian bridge are 844 Gaussian, we can write explicit conditional score for conditional probabilistic path 845

In the work [20] the following loss is introduced to train a model for this score 846

847

Similar to (44), for the optimal score s we have:

where p is given in (41). 848

849

850

851

852

853

854

855

856

860

861

862

863

864

865

866

867

868

869

870

871

872

E.2 Use stochastic

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 ODE

recovers the marginal ρ ( x, t ) (with the given initial conditions) of the stochastic process P ( t ) which 857 is obtained by marginalization conditional Brownian bridge (40) over initial and target distribution 858

As the second consequence of this Theorem, the SDE 859

generates so-called Markovization of the process P ( t ) . Indeed, we can rewrite PDE Eq. (47) in the form

where nabla operator is defined as ∆ = div ∇ . Thus, we get the Fokker-Planck equation for the density of the stochastic process (48).

E.3 Particular cases

In particular case of Brownian bridge when σ e ( t ) = σ ϵ √ t (1 -t ) , then σ ′ e ( t ) = σ ϵ (1 -2 t ) / ( 2 √ t (1 -t ) ) . In this section we consider simple case of separable variables π ( x 0 , x 1 ) = ρ 0 ( x 0 ) ρ 1 ( x 1 ) .

E.3.1 Gaussian initial distribution

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:

where w ( x, t | x 1 ) is the conditional velocity, generated by the conditional map ϕ t,x 1 ( x ) = 873 √ σ 2 e t (1 -t ) + (1 -t ) 2 + tx 1 : 874

Thus, note that in the case of Gaussian distributions, all the difference between this expression and 875 the expression without the stochastic part is the appearance of additional (time-dependent, in general) 876

variance. Marginal distributions are still Gaussian's. 877

Similar, using Eq. (46) we have for the score s : 878

E.3.2 Samples instead of distributions 879

Consider the case where we only have access to the samples { x i 0 } N 0 i =1 and { x i 1 } N 1 i =1 from both 880 distributions, ρ 0 and ρ 1 , but do not know their explicit expressions. In this case, we can estimate the 881 vector field using by a method similar to the one we used to estimate the vector field in (15): 882

In addition, we can also use the importance sampling method in this case. Namely we can use 884 both approaches: self-normalized importance sampling and rejection sampling, similar to what is 885 described in Sec. B 886

F Consistency of Eq. (24) in the case of optimal transport 887

Let us analyze what happens if in formula (24) the joint density π represents the following Dirac delta-function 4 :

i. e. we have a deterministic mapping F from x 1 to x 0 . Then, the Eq. (34) come to 888

Let y ( x, t ) be the unique solution of the equation 889

Now, let us use linear mapping ϕ t,x 1 ( x ) = x 1 t + x (1 -t ) , with inverse ϕ -1 t,x 1 ( x ) = x -tx 1 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, namely

Similar for the score 883

considered as an equation on y . Then

890

891

892

893

894

895

896

897

898

899

900

Here and further by index i we denote i th component of the corresponding vector. Then, the Eq. (53) reads as with the solution

Then the expression for the vector field is

Now, knowing the expression for velocity, we can write the equations for the trajectories x ( t ) :

This equation have closed-form solution:

Analyzing the obtained solution, we conclude that, first, the trajectories obey the given mapping F :

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.1 CFMdispersion

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 , σ 0 = 1 , µ 1 = µ , σ 1 = σ , to the exact expression (35) to get

where

and C is constant independent of x . We then redefine the dispersion based on Eq. (19) using 901 x = (1 -t ) x 0 + tx 1 with x 0 ∼ ρ 0 and x 1 ∼ ρ 1 : 902

This leads us to the final expression: 903

This provides a comprehensive representation of the updated dispersion for the CFM objective at any 904 given time t . 905

906

907

908

909

910

911

912

913

914

915

916

917

918

919

920

921

922

923

924

925

926

927

928

929

930

931

932

933

934

935

936

937

938

939

940

Algorithm 2 Computation ExFM dispersion algorithm

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 objective

1: Sample ( M · N ) samples x from ρ 1 , where i ∈ [1 , M ] and j ∈ [1 , N

2: Sample ( M ) samples x from ρ 0 , where i ∈ [1 , M

3: Compute points x i as (1 -t ) x i + tx

5: Compute and return dispersion D ( v ( x i , t ) -v d ( x i , t

G.2 ExFM dispersion

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 = 200 k , N = 128 , ρ 0 = N (0 , I ) , ρ 1 = N ( µ, σ 2 I ) , and the optimal model v ( x, t ) was derived from equation (54).

H Additional Experiments

H.1 2D toy examples

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 , 000 points. 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 1 e -5 .

H.2 Tabular

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 = 8, train size = 852174, test size = 105206) recorded readings from 16 chemical sensors exposed to gas mixtures. The hepmass dataset (dimension = 21, train size = 315123, test size = 174987) described Monte Carlo simulations for high energy physics experiments. The minibone (dimension = 43, train size = 29556, test size = 3648) dataset contained examples of electron neutrino and muon neutrino. Furthermore, we utilized the BSDS300 dataset (dimension = 63, train size = 1000000, test size = 250000), which involved extracting random 8 x 8 monochrome patches from the BSDS300 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. To train the model we use the same MultiLayer Perceptron (1024 x 3) model with ReLu activations, Adam as optimizer with learning rate of 10 -3 and no learning rate scheduler. As in the pretrained step, we use separately training and testing sets for training the model and calculating metrics. We train the models on the full dataset (of size train_set_size ) with batch size N = 5000 ( batch_size )

949

Training loss,power

Training loss, gas

Training loss, hepmass

2.00

ExFM

CFM

ExFM

CFM

2.0

ExFM

CFM

1.75 -

OT-CFM

OT-CFM

1.8

1.50-

1.6

1.25

1.4

SSOI

0.75

1.0

0.50

0.8

0.25

0.6

0.00

2000

4000

6000

8000

10000

2000

4000

6000

8000

10000

0.4

2000

4000

6000

8000

10000

step

step

step

Training loss, bsds300

Training loss,miniboone

1.0

ExFM

2.0

+ ExFM

OT-CFM

1.8-

CFM

OT-CFM

0.8-

1.6

1.4

0.6

loss

0.4

1.0

0.8

0.2

0.6

0.4

0.0

2000

4000

6000

8000

10000

2000

4000

6000

8000

10000

step step

(except miniboone dataset, here we used 2000 since the smaller size of the dataset) and mini 941 batches n = 256 elements ( mini_batch_size ), the number of epochs and steps for each dataset 942 is adaptive num_epochs = train_set_size // batch_size and num_steps = batch_size 943 // mini_batch_size . 944

For both 2D-toy an tabular data: we take m = n time variable, individual value of variable t 945 corresponds to its pair ( x 0 , x 1 ) . The notations N , n and m corresponds to those in Algorithm 1. To 946 perform sampling, we employed the function odeint with dopri5 method from the python package 947 torchdiffeq import odeint with atol and rtol equal 1 e -5 . 948

H.3 ExFM-S evaluation

The models were assessed using four toy datasets of two dimensions each. A three-layer MLP 950 network was utilized, featuring SeLU activations and a hidden dimension of 64. Optimization was 951 carried out using the AdamW optimizer with a learning rate of 10 -3 and a weight decay of 10 -5 . 952

NLL, N(0, 1) → power

NLL, N(O, ) → gas

NLL, N(O, 1) → hepmass

1.0

ExFM

X

ExFM

ExFM range

ExFM range

ExFM range

0.8

CFM range

CFM range

26

CFM

CFM range

0.6

OT-CFM

OT-CFM range

OT-CFM

OT-CFM range

25

OT-CFM range

OT-CFM

0.4

0.2-

23

0.0

22

0.2

0

2000

4000

step

6000

8000

10000

2000

4000

step

6000

8000

10000

21

2000

4000

step

6000

8000

10000

NLL, N(0, 1) → bsds300

NLL, N(0,1) → miniboone

110

22

ExFM

ExFM

ExFM range

21

ExFM range

115-

*+*.

CFM

CFM range

20

CFM range

OT-CFM

OT-CFM

OT-CFM range

19

OT-CFM range

-120

77N

18-

125

16

130

15

14

135

0

2000

4000

6000

8000

10000

13

0

2000

4000

6000

8000

10000

step step

Training loss, CIFAR10

0.25

ExFM

CFM

0.20

OT-CFM

S

S

0.15

0.10

50000

100000

150000

200000

250000

300000

350000

400000

The model was trained over 2,000 iterations with a batch size of 128. Inference was conducted using 953 the Euler solver for Ordinary Differential Equations (ODE) with 100 steps. To validate the models, 954 the POT library was employed to compute the Wasserstein distance based on 4,000 samples. The 955 experiments were performed on a single Nvidia H100 GPU with 80gb memory. 956

957

H.4 CIFAR 10 and MNIST

We conducted experiments related to high dimensional data, the parameters for training were taken 958 from the open-source code 6 from the works [20, 19]. We saved the leverage of additional heuris959 tics(EMA, lr scheduler). 960

961

962

963

964

965

966

967

968

969

Training loss, MNIST

0.25

ExFM

CFM

0.20

OT-CFM

S

S

0.15

0.10

20000

40000

60000

80000

100000

H.5 Metrics

For evaluating 2D toy data we use Energy Distance and W2 metricis, for Tabular datasets we use Negative Log Likelihood, for CIFAR10 we took Fréchet inception distance (FID) metrics. This choice is connected with an instability and poor evaluation quality of Energy Distance metrics and W2 among high-dimensional data .

H.5.1 Energy Distance

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: µ = ν .

For statistical samples from X and Y :

the following arithmetic averages of distances are computed between the X and the Y samples: 970

FID, N(0, 1) → CIFAR10

6.0

ExFM

ExFM range

CFM

5.5

CFM range

OT-CFM

OT-CFM range

5.0-

4.5

4.0

3.5

0

0000000000000000000000000000000000000000000009

step

44981/2241

73800700

5139978659

9

797305C1

93

2

7

792667315

The E-statistic of the underlying null hypothesis is defined as follows: 971

H.5.2 2-Wasserstein distance (W2) 972

The 2-Wasserstein distance [14], also called the Earth mover's distance or the optimal transport 973 distance W is a metric to describe the distance between two distributions, representing two different 974 subsets A and B . For continuous distributions, it is: 975

where F A and F B are the corresponding cumulative distribution functions and F -1 A and F -1 B the 976 respective quantile functions. 977

H.5.3 Negative Log Likelihood (NLL) 978

To compute the NLL, we first sampled N = 5000 samples { x s i } N i =1 from the target distribution. Then 979 we solved the following inverse flow ODE: 980

for t from 1 to 0 . For simplicity, changing time variable τ = 1 -t we solve the following ODE: 981

982

983

984

985

986

987

988

989

for τ from 0 to 1 . Thus we obtained N solutions { x 0 i } N i =1 which are expected to be distributed according to the standard normal distribution N ( x | 0 , I ) . So we calculate NLL as

H.5.4 Fréchet inception distance (FID)

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:

where Tr is the trace of the matrix.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nTo the best of my knowledge, the idea is novel to use a larger batch of data for the velocity target in the flow matching objective than as the input for the network. This reduces the variance in the loss estimate.\n\nThe additional compute over the FM loss seems negligible, since only the target velocity field is adapted, but the network still receives batches. This makes it an attractive improvement to flow matching training.\n\nWEAKNESSES:\nI 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.\n\nBelow I grouped the weaknesses by category.\n\n## Theoretical contributions can be simplified and potentially reveal existing result (update: largely addressed, but presentation to be improved)\n\nI think the theoretical derivation can be greatly simplified:\n\n1. 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$).\n2. 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.\n\nIn 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.\n\nI 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}$.\n\n## Evaluation on tabular data is wrong (update: fixed)\n\nThe 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.\n\nAlso, 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.\n\n## Toy data evaluation (update: fixed)\n\nIt is easy to construct a very good approximation for the moons distribution by taking a Gaussian mixture of values for Table 4.\n\n## Typos (has no influence on my recommendation)\n\nHere is a list of what I found:\n\n- l. 30: introduced -> introduce\n- l. 33: base -> based\n- l. 65: $rho$ -> $\\\\rho$\n- l. 70: need -> needed\n- l. 92/93: using map -> using the map\n- l. 100: we return to end of the standard CFM loss representation -> ?\n- l. 105: just (unknown) -> just the (unknown)\n- eq. 7: the integral shares variables with the outside expression, e.g. add tilde on the integration variables\n- l. 124: inevitable -?> invertible\n- l. 126: have -> has\n- eq. 16: consider moving numerical tricks like using softmax from the theory section of the paper, page 6 already introduces a lot of notation.\n- throughout: dispersion -> variance\n\nQUESTIONS:\n- Is the variance a problem in high dimensions? I would expect that overall there are few collisions of conditional flows paths.\n- 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?\n- 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?\n- What is the difference between Table 3 and 6?\n\nLIMITATIONS:\nI 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).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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). \n\nThe 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.\n\n[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).\n\n[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).\n\nSTRENGTHS:\nThe presentation of the proposed method is clear.\n\nWEAKNESSES:\nThe 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.\n\nThe 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.\n\nGiven 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.\n\nMinor comments:\n\n- The authors refer to the original FM framework [6] as Conditional Flow Matching, which was proposed in [7].\n- The objective proposed in Flow Matching is already tractable, so it is not clear what “a tractable form of the Flow Matching objective” means.\n- There is a typo in line 65.\n- The derivative notation used is inconsistent with its description in the text.\n- From eq. 19, I assume that “dispersion” means “variance”.\n\n[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).\n\n[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.\n\n[3] Xu, Yilun, Ziming Liu, Max Tegmark, and Tommi Jaakkola. \"Poisson flow generative models.\" *Advances in Neural Information Processing Systems* 35 (2022): 16782-16795.\n\n[4] Scarvelis, Christopher, Haitz Sáez de Ocáriz Borde, and Justin Solomon. \"Closed-form diffusion models.\" *arXiv preprint arXiv:2310.12395* (2023).\n\n[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).\n\n[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).\n\n[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).\n\nQUESTIONS:\nI have no questions for the authors.\n\nLIMITATIONS:\nThe authors do not provide a necessary literature review nor study the limitations of the proposed approach (see Weaknesses section above).", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- If the derivations are correct, this methodology could potentially add some innovations in the field of flow matching.\n\nWEAKNESSES:\n- The paper is difficult to follow and lacks clear writing and organization. Specifically, the authors' use of notations is very confusing.\n- 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).\n- 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.\n- Overall, the paper feels like it requires substantial revisions and is far from being polished.\nTypos:\n- Line 65, rho_1$\\rightarrow$ $\\rho_1$.\n- Line 130, practical $\\rightarrow$ practical form.\n\nQUESTIONS:\n- 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$.\n- 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.\n- 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.\n- 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.\n- Have the authors tried running experiment on non-Gaussian priors?\n- Rather than explaining the training scheme in detail, I would suggest adding pseudo code.\n\nLIMITATIONS:\nThe author has adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2024-07-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nVariance reduction for CFM is a useful avenue of research\n\nWEAKNESSES:\n- 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.\n- 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.\n- 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.\n\nQUESTIONS:\nThe 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]. \n\nGiven this, I don't think ExFM should be treated as a new loss.\n\nThe 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.\n\n[1] \"Flow matching for generative modeling\"\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-17"}