Title: Neural McKean-Vlasov Processes: Inferring Distributional Dependence
Venue: NeurIPS 2023
Accepted: no
Keywords: McKean-Vlasov

stochastic process

ito diffusion
Abstract: McKean-Vlasov stochastic differential equations (MV-SDEs) provide a mathematical description of the behavior of an infinite number of interacting particles by imposing a dependence on the particle density.
These processes differ from standard Ito-SDEs to the extent that MV-SDEs include distributional information in their individual particle parameterization.
As such, we study the influence of explicitly including distributional information in the parameterization of the SDE.
We first propose a series of semi-parametric methods for representing MV-SDEs, and then propose corresponding estimators for inferring parameters from data based on the underlying properties of the MV-SDE.
By analyzing the properties of the different architectures and estimators, we consider their relationship to standard Ito-SDEs and consider their applicability in relevant machine learning problems.
We empirically compare the performance of the different architectures on a series of real and synthetic datasets for time series and probabilistic modeling.
The results suggest that including the distributional dependence in MV-SDEs is an effective modeling framework for temporal data under an exchangeability assumption while maintaining strong performance for standard Ito-SDE problems due to the richer class of probability flows associated with MV-SDEs. 

Abstract

McKean-Vlasov stochastic differential equations (MV-SDEs) provide a mathematical description of the behavior of an infinite number of interacting particles by imposing a dependence on the particle density. These processes differ from standard Itô-SDEs to the extent that MV-SDEs include distributional information in their individual particle parameterization. As such, we study the influence of explicitly including distributional information in the parameterization of the SDE. We first propose a series of semi-parametric methods for representing MV-SDEs, and then propose corresponding estimators for inferring parameters from data based on the underlying properties of the MV-SDE. By analyzing the properties of the different architectures and estimators, we consider their relationship to standard Itô-SDEs and consider their applicability in relevant machine learning problems. We empirically compare the performance of the different architectures on a series of real and synthetic datasets for time series and probabilistic modeling. The results suggest that including the distributional dependence in MV-SDEs is an effective modeling framework for temporal data under an exchangeability assumption while maintaining strong performance for standard Itô-SDE problems due to the richer class of probability flows associated with MV-SDEs.

1 Introduction

Stochastic differential equations (SDEs) model the evolution of a stochastic process through two functions known as the drift and diffusion functions. Beginning with Itô-SDEs, where individual sample paths are assumed to be independent, neural representations of the drift and diffusion have achieved high performance in many applications, such as time series and generative modeling [Song et al., 2020, Tashiro et al., 2021].

On the other hand, interacting particle systems are also used to model stochastic processes using many of the same characteristics as an Itô-SDE, but they additionally dictate an interaction between the different sample paths [Liggett, 1997]. When 31 the number of particles approaches infinity, these processes generalize Itô-SDEs to nonlinear SDEs 32 known as McKean-Vlasov SDEs (MV-SDEs). The nonlinearity arises from the individual particle 33 dependence on the whole particle density, often in the form of a mean-field term represented by an 34 expectation with respect to the particle density. This distributional dependence allows for greater 35 flexibility in the time marginal distributions that the MV-SDE can represent versus the Itô-SDE. An 36

(a) Itô-SDE

(b) MV-SDE

example of the differences between the two frameworks is illustrated in Figure 1 where Figure 1a 37 depicts an Itô-SDE where the sample paths are independent and Figure 1b depicts a MV-SDE where 38 the sample paths interact through distributional dependence. While these models appear in a variety 39 of disciplines such as in finance [Feinstein and Søjmark, 2021], biology [Keller and Segel, 1971], and 40 social sciences [Carrillo et al., 2020], relatively few works have considered the problem of estimating 41 parameters from observations or their application in machine learning tasks. 42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

This brings us to a motivating question:

(Q1) Can we develop theoretically justified neural architectures to represent MV-SDEs?

To answer (Q1) , we use the relationship between the approximation capabilities of neural networks and properties of MV-SDEs. We consider two ideas: (i) expressing a layer in a neural network as an expectation with respect to a density and (ii) using generative models to capture distributions and generate samples.

Our second question relates the theoretical generality of MV-SDEs to Itô-SDEs:

(Q2) Does including explicit distributional dependence empirically affect modeling capabilities?

We discuss a few theoretical properties that motivate this question and answer the question empirically by comparing different architectures for applications in time series and in probabilistic modeling.

1.1 Related work

Methods that estimate MV-SDEs from observations often assume known interaction kernels and drift parameters. They then rely on a large number of samples at regularly spaced time intervals to empirically approximate the expectation in the mean-field term [Messenger and Bortz, 2022, Della Maestra and Hoffmann, 2022, Yao et al., 2022, Della Maestra and Hoffmann, 2023]. In Pavliotis and Zanoni [2022], the authors describe a method of moments estimator for the parameters of the MV-SDE. Other approaches concerned analyzing the partial differential equation (PDE) associated with MV-SDEs as in Gomes et al. [2019]. In our work, we are primarily concerned with inference in regions where we have limited time-marginal data and the number of samples is not large. Other applications of MV-SDEs in machine learning topics include estimating optimal trajectories in scRNASeq data [Chizat et al., 2022] and stochastic control problems relating to mean-field games [Han et al., 2022]. Ruthotto et al. [2020] considered a machine learning approach for solving certain kinds of mean field games and mean field control problems. Inverse problems can also be solved by deriving an appropriate MV-SDE as the authors describe in Crucinio et al. [2022]. Extensive analysis of the dynamics of the parameters of a neural network under stochastic gradient descent has been conducted using the theory from MV-SDEs, e.g. [Hu et al., 2021]. These methods use a pre-described form of the drift to conduct their analyses whereas we're interested in learning a representation of the drift.

Our Contributions To address the lack of non-parametric MV-SDE estimators in the existing litera70 ture, this paper contributes the following: First, we present two neural architectures for representing 71 MV-SDEs based on learned measures and generative networks; then, we present three estimators, 72 based on maximum likelihood, used in conjunction with the architectures without prior knowledge 73 on the structure of the drift; next, we characterize the properties of implicit regularization and richer 74 probability flows of these architectures; finally, we empirically demonstrate the applicability of the 75 architectures on time series and generative modeling. 76

2 Properties of MV-SDEs 77

We begin by describing the background and properties of the transition densities of MV-SDEs. 78

Figure 2 illustrates some of these concepts qualitatively where we first consider non-local dynamics 79

and then consider jumps in the sample paths. 80

2.1 Background 81

Consider a domain D ⊂ R d and let P k ( D ) be the space of all probability distributions supported on 82 D with finite k th moment. Let W t ∈ R d be a d -dimensional Wiener process and let X t ∈ R d be a 83

84

solution to the following MV-SDE

where p t denotes the law of X t at time t and √ Σ denotes the Cholesky decomposition of Σ . We 85 assume that the drift vector b : R d × P 2 ( D ) × R + → R d and the diffusion matrix Σ : R d × P 2 ( D ) × 86 R + → SPD ( R d × d ) are globally Lipshitz for the existence and uniqueness of the solution, with SPD 87 denoting the space of symmetric, positive definite matrices. 88

We focus on the case where the diffusion coefficient is a known constant, σ , and focus on estimating 89 the drift, b , from data. In addition, for simplicity in analysis, we suppose that b factors linearly 90 into a non-interacting component, and an interacting component, where the mean-field term with 91 dependence on p t is often written in terms of an expectation, specifically 92

where f : R d × R + → R d can be seen as the Itô drift, the expectation as the mean-field drift, and ϕ : R d × R d → R k as the interaction function describing the interaction between particles, e.g. attraction with ϕ ( x, y ) = -( x -y ) in Figure 1b and the left side of Figure 2. We also assume that all coefficients exhibit sufficient regularity such that the empirical law converges to the true law of the system (i.e. 1 N ∑ N i =1 δ X ( i ) t → N →∞ p t ( X t ) ), i.e. propagation of chaos holds [Méléard, 1996]. As mentioned, unlike Itô-SDEs which only consider dependence on X t and t , MV-SDEs also depend on the marginal time distribution p t . By introducing a dependence on the marginal law, the transition density of the process satisfies a richer class of functions.

2.2 Non-locality of the transition density

Following the background, we describe a favorable property of the MV-SDE that induces non-local dependencies in the state space. The transition density of (2) can be written as the non-linear PDE

This non-local behavior has a variety of implications. For example, the distribution of particles 'far away' from a reference particle can affect the behavior of the reference particle. This property is illustrated in the left side of Figure 2 with an example from the mean-field FitzHugh-Nagumo model used to model spikes in neuron activation, leading to interactions between distinct spikes [Crevat et al., 2019]. Notably, this is not possible when considering only the Itô drift, since that operator acts locally on the density.

2.3 Discontinuous sample paths

The richer class of densities modeled by MVSDEs has direct influence on individual sample paths. In a modeling scenario, we may wish to approximate a process that exhibits jumps. For example, in finance, a number of related entities may have common exposure and experience failure simultaneously [Nadtochiy and Shkolnikov, 2019, Feinstein and Søjmark, 2021]. Similarly, in neuroscience, a number of neurons spiking simultaneously results in discontinuities in the sample paths [Carrillo et al., 2013]. The fact

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

that the interaction of many particles can cause blowups leads to a remarkable property of MV-SDEs that allows discontinuous paths. The major benefit of this property is that we do not need to consider an additional jump noise process - we only need to specify a particular interaction between the particles to induce the jump behavior. A simple proof for the case of positive feedback is given in Hambly et al. [2019, Theorem 1.1].

Having described the theoretical advantages of MV-SDEs as compared to Itô-SDEs, we will proceed 127 to discuss the neural architectures for representing these processes. 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

Figure 3: Schematic comparing neural architectures for modeling MV-SDEs. Implicit measure (IM) architecture uses a mean-field layer that represents particles as learned weights; the empirical measure (EM) architecture computes the expectation with the observed particles; the marginal law (ML) estimates the particle density, computing the expectation with samples from the estimated density.

3 Mean-field Architectures

We now describe methods for representing the mean-field drift of a MV-SDE in (2). We first consider a modification of the cylindrical architecture [Pham and Warin, 2022] that empirically computes the expectation using observations, and denote it as the empirical measure (EM) architecture. We then propose two architectures - an architecture based on representing a learned measure with neural weights, denoted as the implicit measure (IM) architecture, and a generative architecture based on representing the marginal law of the samples (ML). Figure 3 provides a schematic of the different architectures and mean-field representations. We denote a function f parameterized by parameters θ as f ( · ; θ ) .

3.1 Empirical measure architecture

Suppose we observe N particles at each time t given by { X ( i ) t } N i =1 and denote the discrete measure associated with these observations as p δ t = 1 N ∑ N i =1 δ X ( i ) t . Then, we can use p δ t to approximate the expectation in (2) as

for a neural network ϕ ( · , · ; θ ) describing the interaction function between the particles [Pham and Warin, 2022]. Suppose the non-mean field component f is also represented with a neural network f ( · , t ; θ ) . Assuming that ϕ and f are well learned, this architecture can represent the true MV-SDE drift in the limit as the number of observations N → ∞ . We refer to this architecture as the empirical measure (EM) architecture since at each time step the expectation is taken with respect to the empirical measure derived from the observations.

3.2 Implicit measure architecture

While the EM architecture in (4) explicitly defines the relationship between the law p t and the interaction ϕ , it relies on obtaining the empirical measure at each time point. This may be difficult in practice for a variety of reasons such as having few samples or the lack of data at some time points.

Instead, let us first recall that a single layer in a multilayer perceptron (MLP) can be written in terms of an expectation as

where the expectation is taken over ν ( · ) , a measure over the space of parameters y = ( W,b ) , and σ is an activation function.

When ν = 1 N ∑ N i =1 δ y ( i ) , a discrete measure with N particles, the expectation is exactly a single layer of width N , suggesting a correspondence between an empirical measure with N samples and a single layer of width N . Building on this correspondence, we propose a mean-field layer:

Definition 3.1 (Mean-field Layer) . Define the weight of the mean-field layer with width n as the 159 matrix W 0 ∈ R n × d and denote its i th row as W ( i ) 0 . The mean-field layer then is defined by the 160

operation 161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

The mean-field layer (MF) can be thought of as another layer within the network architecture that approximates the law p t . Each row W ( i ) 0 is of size R d , corresponding to the dimensions of X ( i ) t ∈ R d . The activation function of the mean-field layer is the average over the augmented dimension over which MF operates. The change of measure d P t d P 0 can be learned as part of the estimator of the interaction function, ϕ ( · , · , t ; θ ) . Importantly, the above representation allows modeling mean-field interactions without the need for a full set of observations at each time point and without the need to explicitly represent the distribution p t at each time point. Assuming that ϕ and MF are well learned, this architecture can represent the true MV-SDE drift in the limit as the width n → ∞ . We note empirically that a finite n is sufficient and we provide examples of ablations in the appendix.

A similar analysis can be made for the standard MLP architecture. However, the explicit separation of f and ϕ is not enforced in this case. This leads us to the following remark:

Remark 3.2 (Itô-SDEs with drift represented using MLPs can model MV-SDEs) . From the above discussion, the expectation with respect to the law p t may be implicitly represented by a MLP.

Our motivation is then concerned with how a relatively more explicit distribution dependence with ϕ and MF affect modeling capabilities. This explicit structure lends to an implicit regularization that promotes a smaller norm of the mean-field component under a maximum likelihood estimation framework, which we detail later in Section 5.1.

3.3 Marginal law architecture

A solution to the MV-SDE is the pair ( X,p ) such that p t = Law ( X t ) . In addition, if p is a solution to the SDE in (2), it is also a weak solution to the PDE in (3), and the converse holds. For this reason, p is often itself the main object of study. In the marginal law (ML) architecture, in conjunction with the drift, we introduce a generative model for representing the time-varying density. In this case, we approximate the expectation in (2) as

where the expectation is taken with respect to the discrete measure derived from samples { ˆ X ( i ) t } n i =1 185 from the generative model ˆ P t . The parameter estimation problem then requires optimizing both the 186 generative model ˆ P t and the networks f and ϕ representing the drift, while ensuring consistency 187 between the two. Using knowledge of the PDE in (3), we regularize ˆ P t such that it matches the flow 188 relating to the drift. Additional details regarding the PDE and its relationship to the ML architecture 189 are in the appendix. 190

191

192

193

194

195

196

197

198

4 Parameter Estimation

Having presented the relevant architectures, we now describe the procedures for estimating the parameters of the different architectures. We first describe the likelihood function for use in cases with regularly sampled data. We then describe a bridge estimator for cases of irregularly sampled data. Finally, we describe an estimator for the generative architecture based on both the likelihood function and the transition density. For this section, we assume that we observe multiple paths, i.e., { { X t j } ( i ) j =1 ...K } i =1 ...N . Full details of all algorithms are in the appendix.

4.1 Maximum likelihood estimation

We use an estimator based on the path-wise likelihood derived from Girsanov's theorem and an 199 Euler-Maruyama discretization for the likelihood, considered in Sharrock et al. [2021]. The likelihood 200

function is given as 201

202

203

204

205

206

207

208

209

210

211

212

Following discretization, with the approximations ∆ X t j = X t j +1 -X t j and ∆ t j = t j +1 -t j , the log-likelihood is approximated by

If the time interval ∆ t is large, then this likelihood loses accuracy, as is a property of the EulerMaruyama discretization. Optimization is performed using standard gradient based optimizers with the drift b represented as one of the presented architectures.

4.2 Estimation with Brownian bridges

Often data are not collected at uniform intervals in time, but rather, the time marginals may be collected at irregular intervals. In that case, we consider an interpolation approach to maximizing the likelihood following the results of Lavenant et al. [2021] and Cameron et al. [2021] in the Itô-SDE case. We can write the likelihood conditioned on the set of observations (dropping the particle index for ease of notation) as

where { Z s : s ∈ [ t j , t j +1 ] } is a Brownian bridge from X t j to X t j +1 and Q is the Wiener measure. 213 Brownian bridges can easily be sampled and reused for computing the expectation, which reduces 214 the variance of the estimator. By applying Jensen's inequality, we can write an evidence lower bound 215 (ELBO) as 216

217

218

219

220

221

222

223

The ELBO in this case aims to fit the observed marginal distributions exactly while penalizing deviations in regions without data that deviate from the Brownian bridge paths.

4.3 Estimation with explicit marginal law ˆ P t

Returning to the ML architecture described in Section 3.3, where we explicitly model the density p t with a generative network ˆ P t , our estimator should enforce the regularity of p t through its PDE in (3). Let the parameters of the drift be θ and the parameters of the generative model be φ , then we solve the optimization problem

for time intervals indexed by j = 1 . . . K -1 , the state space x ∈ supp( X t ) , and where the 224 trajectories of ˆ X t follow the dynamics of the ML architecture, specifically 225

The likelihood at the observed margins is first maximized in (10). In (11), the marginals at previous 226 times are regularized using the correspondence between the PDE and its associated SDE via the 227 nonlinear Kolomogorov backwards equation [Buckdahn et al., 2017], which describes p t as an 228 expectation of trajectories at a terminal time, i.e. p t ( x ) = E [ p T ( X T ) | X t = x ] for t < T . 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

267

268

269

270

271

5 Modeling Properties

Having discussed the architectures and estimators, we now discuss specific properties of the modeling framework, which follow from the theoretical discussion presented in Section 2. We first discuss how the factorization into ϕ and MF lends to an implicit regularization of the IM architecture. We then compare the gradient flows of Itô-SDEs and MV-SDEs.

5.1 Implicit regularization of the implicit measure architecture

Closely related to the IM architecture are neural Itô-SDEs, where we previously remarked can model MV-SDEs. On the other hand, the factorization of the IM architecture into ϕ and MF leads to a type of implicit regularization when the parameters are estimated using gradient descent.

Proposition 5.1 (Implicit Regularization) . Suppose f , ϕ known and fixed. Further, assume that ϕ is twice differentiable. Then, for each time step t , the minimizing finite width MF with weight matrix W 0 ∈ R n × d and i th row W ( i ) 0 under gradient descent satisfies the following optimization problem

Proof. We follow the blueprint in Belabbas [2020] and give full details in the appendix.

Proposition 5.1 effectively says that the mean-field system approximated is the one that has the least influence from the other particles under perfectly matched marginals. In the case where ϕ can be decomposed as a norm, this amounts to finding the drift parameterized by weight W 0 with smallest norm while still matching the marginals.

5.2 Gradient flows of the MV-SDE

To illustrate the difference between the MV-SDE and Itô-SDE particle flows, we invoke the analysis in Santambrogio [2017, Section 4.6] to describe the functionals that are minimized by each.

Remark 5.2 (Functional Minimizer) . Consider two drifts B = ∇ f ( X ) and B MF = B E [ ∇ ϕ ( X -y )] . Consider a functional F [ p ] = ∫ log p d p + ∫ f ( X )d p for some measure p absolutely continuous with respect to the Lebesgue measure. Then, the gradient flow satisfying the linear Fokker-Planck equation with drift B minimizes F . On the other hand, the nonlinear Fokker-Planck associated with drift B F [ p ] = F [ p ] + ϕ ( X Y )d p ( X )d p ( Y ) .

This has an important implication, for example, if we take ϕ ( · ) = 2 ‖ · ‖ d q d p -‖ · ‖ 2 -‖ · ‖ 2 ( d q d p ) 2 then the functional is minimizing the squared energy distance between a target measure q as well as

+ MF minimizes the functional MF ∫ - the entropy. We use this example to motivate some of the experiments on probabilistic modeling.

6 Numerical Experiments

We discussed Q1 on modeling and inferring distributional dependence. We now wish to answer Q2 and quantify the effect of distributional dependence in machine learning tasks. To do this, we test the methods on synthetic and real data for time series estimation and sample generation. The main goal is to determine the difference between standard Neural Itô-SDE and the proposed Neural MV-SDEs under different modeling scenarios. In that sense, the baseline we consider is the Itô-SDE parameterized using an MLP. However, we also consider other deep learning based methods for comparison in a broader context. We abbreviate the different architectures as the Empirical Measure (EM) in Section 3.1, Implicit Measure (IM) in Section 3.2, and Marginal Law (ML) in Section 3.3. Full descriptions of the models, baselines, and datasets are given in the appendix.

Synthetic data experiments Motivated by the application of MV-SDEs in physical, biological, social, and financial settings, we benchmark the proposed methods on 4 canonical MV-SDEs: the Kuramoto model which describes synchronizing oscillators [Sonnenschein and Schimansky-Geier, 2013], the mean-field FitzHugh-Nagumo model which characterizes spikes in neuron activations

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

[Mischler et al., 2016], the opinion dynamic model on the formation of opinion groups [Sharrock et al., 2021], and the mean-field atlas model for pricing equity markets [Jourdain and Reygner, 2015]. We additionally benchmark the proposed methods on two Itô-SDEs: an Ornstein-Uhlenbeck (OU) process and a circular motion equation to determine the performance on Itô-SDEs. Finally, to understand the performance on discontinuous paths, we benchmark the proposed methods on an OU process with jumps. We focus on recovering the drift from observations.

Since the true drifts of the synthetic data are known, we directly compare the estimated drifts to the true drifts. The performance on five different datasets with three different levels of added observational noise is presented in Figure 4. The proposed meanfield architectures outperform the standard MLP in modeling MV-SDEs; moreover, our experiments on OU and circular process suggest that incorporating explicit distributional depedence does not diminish the performance in estimating non-interacting Itô-SDEs. When modeling processes with

jump discontinuities, Figure 5 highlights the flexibility of the proposed methods, IM, ML, to match such models. The EM likely does not perform as well due to the high variance of the empirical measure, leading to difficulties in learning. Additionally, the MLP does not have an explicit decomposition between the MV and Itô components, resulting in issues when estimating the feedback between the particles inducing jumps.

Real data experiments Extending from the synthetic examples, we consider two real examples: brain activity recorded by electroencephalograms (EEG), which is closely related to the Kuramoto model [Nguyen et al., 2020]; and chemically stimulated movement of organisms (chemotaxis), which can be modeled by the Keller-Segel model [Toma š evi´ c, 2021, Keller and Segel, 1971].

We evaluate the proposed architectures in these modeling tasks by comparing the goodness-of-fit 300 of generated path samples to the observed path samples. We compute the Continuous Ranked 301 Probability Score (CRPS) defined in Gneiting and Raftery [2007] (see appendix for details) for 302 the 1-dimensional EEG data, and the normalized MSE (normalized with sample variance) for the 303 3-dimensional chemotaxis data with respect to the held out data. We also benchmark against the 304 DeepAR probabilistic time series forecaster [Salinas et al., 2020] with RNN, GRU, LSTM, and 305 Transformer (TR) backbones as another baseline model to compare the goodness-of-fit. 306

The performances of different architectures are presented in Table 1. For EEG, the proposed 307 architectures generally perform better than the baselines in generating paths within the training time 308 steps, and on par with the DeepAR architectures for forecasting (full results presented in appendix). 309 For chemotaxis data, the MV-SDE based architectures all outperform the DeepAR baselines. 310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

Generative modeling experiments We focus on applying the bridge estimator discussed in Section 4.2 to map between a Gaussian and a target distribution. We are interested in two aspects: 1) the properties of the learned mapping, and 2) the generated trajectories. We first study the properties of the learned mapping using a synthetic eight Gaussian mixture with increasing dimensionality. We compare the performance of different architectures through the ELBO of the sample paths generated by the optimal transport (OT) mapping between the initial distribution and held out target samples. We next evaluate the generated trajectories through the energy distance (see appendix for details) between generated and held-out data for 5 real data density estimation experiments. In addition, we compare to common density estimators of variational autoencoder (VAE) [Kingma and Welling, 2013], Wasserstein generative adversarial network (W-GAN) [Gulrajani et al., 2017], masked autoregresive flow (MAF) [Papamakarios et al., 2017] and score-based generative modeling through SDEs, which corresponds to a constrained form of the MLP [Song et al., 2020]. The MV-SDE architectures not only outperform the Itô architecture for all dimensions in the eight Gaussian experiment, as shown in Figure 6, but also for the 5 real data density estimation experiments, as shown in Table 2, while outperforming common baselines. All sampling is performed using standard Euler-Maruyama, with full details of the sampling and inference algorithms in the appendix. This again suggests the MV-SDE provides a more amenable probability flow for modeling compared with the Itô case.

7 Discussion

In this paper we discuss an alternative viewpoint of the standard Itô-SDE parameterization. In particular, we focus on MV-SDEs and discuss how neural networks can represent a process that depends on the distribution, and we describe ways of making this dependence more explicit. We demonstrated the efficacy of the proposed architectures on a number of synthetic and real benchmarks. The results suggest that the proposed architectures provide an improvement over baselines in certain generative modeling and time series applications.

Limitations We only studied the implicit regularization of the IM architecture under gradient descent, but the extension of the analysis to the other proposed architectures is important to understand the corresponding regularization. Additionally, computing expectations incurs additional computational cost. Improving the computational accuracy using a multilevel scheme as proposed in Szpruch et al. [2019] could improve the performance of the methods.

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

References

Mohamed Ali Belabbas. On implicit regularization: Morse functions and applications to matrix factorization. arXiv preprint arXiv:2001.04264 , 2020.

Rainer Buckdahn, Juan Li, Shige Peng, and Catherine Rainer. Mean-field stochastic differential equations and associated pdes. The Annals of Probability , 45(2):824-878, 2017. ISSN 00911798, 2168894X. URL http://www.jstor.org/stable/44245559 .

Scott Cameron, Tyron Cameron, Arnu Pretorius, and Stephen Roberts. Robust and scalable sde learning: A functional perspective. arXiv preprint arXiv:2110.05167 , 2021.

JA Carrillo, RS Gvalani, GA Pavliotis, and A Schlichting. Long-time behaviour and phase transitions for the mckean-vlasov equation on the torus. Archive for Rational Mechanics and Analysis , 235 (1):635-690, 2020.

José A Carrillo, María d M González, Maria P Gualdani, and Maria E Schonbek. Classical solutions for a nonlinear fokker-planck equation arising in computational neuroscience. Communications in Partial Differential Equations , 38(3):385-409, 2013.

Lénaïc Chizat, Stephen Zhang, Matthieu Heitz, and Geoffrey Schiebinger. Trajectory inference via mean-field langevin in path space. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=Mftcm8i4sL .

Joachim Crevat, Grégory Faye, and Francis Filbet. Rigorous derivation of the nonlocal reactiondiffusion fitzhugh-nagumo system. SIAM Journal on Mathematical Analysis , 51(1):346-373, 2019.

Francesca R Crucinio, Valentin De Bortoli, Arnaud Doucet, and Adam M Johansen. Solving fredholm integral equations of the first kind via wasserstein gradient flows. arXiv preprint arXiv:2209.09936 , 2022.

Laetitia Della Maestra and Marc Hoffmann. Nonparametric estimation for interacting particle systems: Mckean-vlasov models. Probability Theory and Related Fields , 182(1):551-613, 2022.

Laetitia Della Maestra and Marc Hoffmann. The lan property for mckean-vlasov models in a mean-field regime. Stochastic Processes and their Applications , 155:109-146, 2023. ISSN 03044149. doi: https://doi.org/10.1016/j.spa.2022.10.002. URL https://www.sciencedirect.com/ science/article/pii/S0304414922002113 .

Kai Du, Yifan Jiang, and Jinfeng Li. Empirical approximation to invariant measures for mckeanvlasov processes: mean-field interaction vs self-interaction. arXiv preprint arXiv:2112.14112 , 2021.

Zachary Feinstein and Andreas Søjmark. Dynamic default contagion in heterogeneous interbank systems. SIAM Journal on Financial Mathematics , 12(4):SC83-SC97, 2021.

Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association , 102(477):359-378, 2007. doi: 10.1198/016214506000001437. URL https://doi.org/10.1198/016214506000001437 .

Susana N Gomes, Andrew M Stuart, and Marie-Therese Wolfram. Parameter estimation for macroscopic pedestrian dynamics models from microscopic data. SIAM Journal on Applied Mathematics , 79(4):1475-1500, 2019.

Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id=rJxgknCcK7 .

Marianne Grognot and Katja M Taute. A multiscale 3d chemotaxis assay reveals bacterial navigation mechanisms. Communications biology , 4(1):1-8, 2021.

387

388

389

390

391

392

393

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron C Courville. Improved training of wasserstein gans. Advances in neural information processing systems , 30, 2017.

Ben Hambly, Sean Ledger, and Andreas Søjmark. A mckean-vlasov equation with positive feedback and blow-ups. The Annals of Applied Probability , 29(4):2338-2373, 2019.

Jiequn Han, Ruimeng Hu, and Jihao Long. Learning high-dimensional mckean-vlasov forwardbackward stochastic differential equations with general distribution dependence. arXiv preprint arXiv:2204.11924 , 2022.

Kurt Hornik. Approximation capabilities of multilayer feedforward networks. Neural networks , 4(2): 251-257, 1991.

Kaitong Hu, Zhenjie Ren, David Š i š ka, and Ł ukasz Szpruch. Mean-field langevin dynamics and energy landscape of neural networks. In Annales de l'Institut Henri Poincaré, Probabilités et Statistiques , volume 57, pages 2043-2065. Institut Henri Poincaré, 2021.

Chin-Wei Huang, Jae Hyun Lim, and Aaron C Courville. A variational perspective on diffusion-based generative models and score matching. Advances in Neural Information Processing Systems , 34: 22863-22876, 2021.

Benjamin Jourdain and Julien Reygner. Capital distribution and portfolio performance in the meanfield atlas model. Annals of Finance , 11(2):151-198, 2015.

Evelyn F. Keller and Lee A. Segel. Model for chemotaxis. Journal of Theoretical Biology , 30(2): 225-234, 1971.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Hugo Lavenant, Stephen Zhang, Young-Heon Kim, and Geoffrey Schiebinger. Towards a mathematical theory of trajectory inference. arXiv preprint arXiv:2102.09204 , 2021.

Thomas M. Liggett. Stochastic models of interacting systems. The Annals of Probability , 25(1):1 29, 1997. doi: 10.1214/aop/1024404276. URL https://doi.org/10.1214/aop/1024404276 .

Sylvie Méléard. Asymptotic behaviour of some interacting particle systems; mckean-vlasov and boltzmann models. Probabilistic models for nonlinear partial differential equations , pages 42-95, 1996.

Daniel A. Messenger and David M. Bortz. Learning mean-field equations from particle data using wsindy. Physica D: Nonlinear Phenomena , 439:133406, 2022. ISSN 0167-2789. doi: https: //doi.org/10.1016/j.physd.2022.133406. URL https://www.sciencedirect.com/science/ article/pii/S0167278922001543 .

Stéphane Mischler, Cristóbal Quininao, and Jonathan Touboul. On a kinetic fitzhugh-nagumo model of neuronal network. Communications in mathematical physics , 342(3):1001-1042, 2016.

Sergey Nadtochiy and Mykhaylo Shkolnikov. Particle systems with singular interaction through hitting times: application in systemic risk modeling. The Annals of Applied Probability , 2019.

Phuong Thi Mai Nguyen, Yoshikatsu Hayashi, Murilo Da Silva Baptista, and Toshiyuki Kondo. Collective almost synchronization-based model to extract and predict features of EEG signals. Scientific Reports , 10(1):16342, 2020. URL https://doi.org/10.1038/s41598-020-73346-z .

Bernt Karsten Øksendal and Agnes Sulem. Applied stochastic control of jump diffusions , volume 498. Springer, 2007.

George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems , 30, 2017.

Grigorios A Pavliotis and Andrea Zanoni. A method of moments estimator for interacting particle systems and their mean field limit. arXiv preprint arXiv:2212.00403 , 2022.

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

Huyên Pham and Xavier Warin. Mean-field neural networks: learning mappings on wasserstein space. arXiv preprint arXiv:2210.15179 , 2022.

Lars Ruthotto, Stanley J Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung. A machine learning framework for solving high-dimensional mean field game and mean field control problems. Proceedings of the National Academy of Sciences , 117(17):9183-9193, 2020.

David Salinas, Valentin Flunkert, Jan Gasthaus, and Tim Januschowski. Deepar: Probabilistic forecasting with autoregressive recurrent networks. International Journal of Forecasting , 36(3): 1181-1191, 2020.

Filippo Santambrogio. { Euclidean, metric, and Wasserstein } gradient flows: an overview. Bulletin of Mathematical Sciences , 7(1):87-154, 2017.

Louis Sharrock, Nikolas Kantas, Panos Parpas, and Grigorios A Pavliotis. Parameter estimation for the mckean-vlasov stochastic differential equation. arXiv preprint arXiv:2106.13751 , 2021.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Bernard Sonnenschein and Lutz Schimansky-Geier. Approximate solution to the stochastic kuramoto model. Physical Review E , 88(5), nov 2013.

Lukasz Szpruch, Shuren Tan, and Alvin Tse. Iterative multilevel particle approximation for mckeanvlasov sdes. The Annals of Applied Probability , 29(4):2230-2265, 2019.

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems , 34:24804-24816, 2021.

Milica Toma š evi´ c. A new mckean-vlasov stochastic interpretation of the parabolic-parabolic kellersegel model: The two-dimensional case. The Annals of Applied Probability , 31(1):432-459, 2021.

Rentian Yao, Xiaohui Chen, and Yun Yang. Mean-field nonparametric estimation of interacting particle systems. In Po-Ling Loh and Maxim Raginsky, editors, Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 2242-2275. PMLR, 02-05 Jul 2022. URL https://proceedings.mlr.press/v178/yao22a. html .

Xiao Lei Zhang, Henri Begleiter, Bernice Porjesz, Wenyu Wang, and Ann Litke. Event related potentials during object recognition tasks. Brain Research Bulletin , 38(6):531-538, 1995.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe submission considers McKean-Vlasov Stochastic Differential Equation models, which are a generalization of the more-familiar Ito processes. The difference is that the former additionally feature the evolution of the law of the process $X_t$ (denoted $p_t$) as well as $X_t$ itself. Such processes are the limit as the number of particles in a system approach infinity, and their finite-particle approximation (where $p_t$ becomes some empirical distribution) give rise to a model that exhibits temporal as well as between-particle interactions. \n\nA key aspect of the work is to posit a model whereby $p_t$ enters only through a term $\\mathbb{E}_{y_t \\sim p_t}[\\varphi(X_t, y_t)]dt$ for some *interaction function* $\\varphi$. The advantage of such a formulation is that the transition density of the process then reduces to the solution of a PDE.  The work discusses the properties of MV-SDEs of the aforementioned form that make them a desirable class of models (e.g., non-local interactions between paths, and the ability to incorporate jumps). \n\nThree neural architectures are proposed, differing in particular for their approach as it relates to the modelling of $p_t$: \n* Implicit Measure: which recasts the empirical measure as a single layer of a neural network (though I did not completely understand how this happens). The implicit measure architecture has the advantages of being able to cope with the missing data setting.\n* Empirical Measure: where $\\varphi(\\cdot, \\cdot)$ takes the form of a trainable neural network. \n* Marginal Law: which involves learning a generative model.  \n\nMethods for parameter estimation are presented (Section 4), including methods for the missing data setting based on Brownian bridges. A result is proven (Proposition 5.1) regarding implicit regularization properties. A numerical study is conducted on a number of simulated data examples, real data (for which forecasting is also explored), as well as a generative modelling task.\n\nSTRENGTHS:\n* The ideas contained in the submission are highly non-trivial, yet explored in impressive depth from a number of different angles in a sophisticated manner.\n* The study of such techniques is impressively comprehensive in terms of both theory and different approaches (giving a number of methodological approaches, as well as some theory). The submission appears to represent the outcome of a significant body of work and investigation. \n* The submission is very well written and presented. The numerical experiments appear well-executed. \n\nDespite some familiarity with concepts in the submission, I am not an expert on SDE modelling, so can not particularly comment on the novelty in that regard. However, the ideas in the submission appear like ones that are very natural to explore, have been explored well, and are worthy of publication in my opinion.\n\nWEAKNESSES:\n* As much as I enjoyed the abstract presentation, it would be beneficial to have additional clarity as to settings where MV-SDE modelling may be of interest, so to contextualise the work, or where different approaches would be preferred. It is mentioned that few works have considered such models in machine learning tasks, and it would be beneficial to have a small discussion what the contribution could potentially be there (at a high level, that is). \n\n* Related to the above, a clearer motivation of the task at hand would be beneficial. The paper takes an SDE-first viewpoint, but isn't the overall goal to fit some sort of particle approximation to the SDE? Some additional discussion and background would be beneficial. \n \n* Some applications are mentioned at present, but it would be nice to have a small discussion something along the lines of \"MV-SDEs are most useful when the goal of interest is to model... \".\n\n* The derivation of the implicit measure architecture was something that I could not properly parse as currently written.\n\nQUESTIONS:\n* The simplified case of fixing the diffusion coefficient is considered, and only looking at learning the drift. How strong is this restriction in terms of modelling potential? Does it exclude certain behaviour? \n* Similar to the above, but for the factorization imposed on $b(X_t, p_t, t)$ (though I do understand the impetus for such is the PDE representation, I am curious about the other aspects). \n* p4, l130: Additional background should be given regarding what \"mean-field drift\" is, and mean-field approximations of MV-SDEs in general. Those without a strong background in MV-SDE modelling will likely not understand the paper otherwise. Also, should \"in (2)\" not be \"of the form in (2)\", as (2) is the limiting case? \n* p4, l153: It is not clear what the intended meaning of this equation or the sentence preceding it is. Is the intention to say that one is using an MLP with stochastic neurons?  This section would benefit for additional details and clarity. \n* p4, l151: Why does having few samples make it difficult to obtain an empirical measure? I would have thought this would simply made the computation easier. \n* Regarding the title: \"inferring distributional dependence\" is a little vague, and perhaps could be replaced by something else that does the paper more justice and gives those without knowledge of what a MV-SDE is more of a clue what the paper is about (this is a very minor point, and just a suggestion).  \n\nMinor things:\n * p5, l180: should $(X,p)$ be $(X_t, p_t)$?\n\nLIMITATIONS:\nA little additional discussion would be beneficial (see other parts of this review for specifics).", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers the problem of parameter estimation from data when the underlying dynamical system is modeled by the MV-SDE. To represent the target MV-SDE, the authors propose two strategies: (i) expressing a layer in a neural network as an expectation with respect to a density and (ii) using generative models to capture distributions that generate observations at different time stamp. With these strategies, the authors then propose to conduct parameter estimation via 1. maximum likelihood estimation, 2. using Brownian bridge for estimation, 3. explicit marginal law estimation.\n\nSTRENGTHS:\nPlease see the discussion below.\n\nWEAKNESSES:\n1. Poor presentation. Overall this paper is obscure and hard to follow. Here are some detailed examples.\n* When stating the underlying MV-SDE model in Eq.(2), it is not clear which terms are known and which terms are to be learned. Specifically, do we know $f$ and $\\phi$? Since we are interested in the problem of parameter estimation, both terms should be learned from the observations.\n* It is now clear how the proposed implicit measure architecture can be carried out: In Eq.(6), it is not clear what $\\mathbb{P}_t$ and $\\mathbb{P}_0$ are and why we can estimate the Radon–Nikodym derivative $d \\mathbb{P}_t/d\\mathbb{P}_0$. Moreover, the authors are motivating the implicit measure architecture as an alternative to the standard empirical measure approach since it can handle the situation when only samples from irregular time stamps are available. However, it is not clear why this is the case.\n* In section 4.2, where the authors propose to estimate parameters using the Brownian bridge, it is not clearly how the Brownian bridge comes into play. \n* The maximum likelihood estimation in section 4.1 is proposed in previous work [Sharrock et al. 2021], but is presented as a strategy proposed by this work, which is misleading.\n\n2. Lack concrete contributions. Partly due to the poor presentation of this work, this paper presents no clear contributions. For example, the proposed implicit measure architecture and the marginal law architecture are just two ways to rephrase the standard empirical measure.  The results presented in section 5 are well-known results in the literature, e.g. the Wasserstein gradient flow structure of the MV-SDE.\n\n3. A contradiction between the goal of this project and its fundamental assumption. While the authors motivate the research of MV-SDE to model jump (discontinuous behavior) in time sequence data, they also assume that the drifting term of the MV-SDE is sufficiently regular. This is a clear contradiction. For MV-SDE with regular drifting term and interaction term, it can be proved that the characteristic flow is also regular.\n\nQUESTIONS:\nPlease see the comments above.\n\nLIMITATIONS:\nPlease see the comments above.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 1, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors proposed two new methods of modelling McKean--Vlasov SDEs using a neural network, and studied its empirical performance. \n\nSince I'm not an expert this exact topic, I would like to ask the authors some questions first. I would be happy to raise my score further once I understand the paper better.\n\nSTRENGTHS:\nThe authors proposed methods that do not model a finite population of the particles, which is a very interesting alternative.\n\nWEAKNESSES:\nSeveral parts of the paper seem unclear to me at the moment, and I will ask specific questions next.\n\nQUESTIONS:\nAgain, I would be happy to raise my score once my questions are adequately addressed. \n\n1. Let me start with a basic question: under what type of conditions can we factor the drift $b$ into $f$ and $\\varphi$ components as in equation 2? This seems to be an important simplification and I would like the authors to motivate this further. Perhaps the authors can provide some important examples in applications where this simplification is available. \n\n2. The authors seem to also occasionally write $\\varphi(x,y) = \\varphi(x-y)$ fairly interchangeably, e.g. in equation 3. Can the authors clarify if this is intentional in any of the circumstances? For example, it seems like in the MLP representation case of Appendix A.1, it is necessary to use $\\sigma = \\varphi(x-y)$. \n\n3. On a related note, when is $\\varphi$ known a priori? It seems like if this is known, then we wouldn't need a neural network to estimate $\\varphi$ in the empirical measure approach, we would just directly simulate the particle system? \n\n4. The authors wrote that MLPs can model McKean--Vlasov dynamics, but in the derivations it seems like it would require the weights to be identity matrix, which fundamentally restricts the complexity of the bias $b$ since it can only be $d$-dimensional. So doesn't this mean that $\\nu$ can at most be an average of $d$ Dirac-delta's, and not necessarily capable of representing a general distribution $p_t$? \n\nI found the definition and discussion about the mean field layer quite confusing. I have several questions specifically dedicated to this. \n\n5. What is the distribution $\\mathbb{P}_t$, and how does the Radon--Nikodym derivative $\\frac{d \\mathbb{P}_t}{d \\mathbb{P}_0}$ show up in equation 6? \n\n6. When written along with $\\varphi( X_t, W_0^{ (i) } )$, is the Radon--Nikodym derivative a function of $X_t$ or $W_0^{(i)}$? As in which measure is being changed here? \n\n7. Can we interpret each of the $W_0^{(i)}$ as a hypothetical particle? \n\n8. Most importantly, the authors suggested that the Radon--Nikodym derivative can be learned, and that as $n\\to\\infty$ this can also represent the drift. Can the authors provide more details on this part? I don't think either claims are clear at all, and I would like to understand the arguments behind this critical step. \n\nWith respect to the marginal law, I also wanted to ask a few questions. \n\n9. Are the authors estimating the marginal law $P_t$ at each time of $t$? Does this imply that if the authors were to increase the number of time steps, this would require more estimates? \n\n10. Can the authors provide more details about how $P_t$ is being estimated? I can't seem to find anywhere the authors described the procedure to modelling this density at all. \n\nOn a high level, I also have questions regarding the empirical measurements of errors. \n\n11. While MSE going to zero of course implies the method is correctly modelling the underlying dynamics, but it doesn't provide a relative scale of how the methods are performing. Can the authors provide a measure in terms of the Kolmogorov--Smirnov distance, i.e. $L^\\infty$ norm of the empirical CDFs between true $p_t$ and estimations? \n\n12. Can the authors also demonstrate how the method improves as a function of computational power, e.g. in terms of the width of the network and maybe the number of particles for the empirical measures case?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a methodology for simulating McKean-Vlasov (mean-field) equations using standard function approximation techniques, e.g. neural networks. It provides mathematical intuition for these algorithms and evaluates them on a broad suite of benchmarks.\n\nSTRENGTHS:\nThe paper is very well written with clear exposition of its main points. The proofs of key claims seem broadly correct as well.\n\nThe methodology is broadly well justified and uses very intuitive ideas from stochastic analysis. \n\nIn particular, the adaptation of standard neural network techniques for simulating ODEs/SDEs is not entirely applicable here, and so the derived techniques need to account for the estimation of the particle density. The resulting algorithm is novel and an important independent contribution.\n\nThe experiment evidence, especially in the Gaussian case, seems to vindicate the intuition of this algorithm and clearly outperforms the chosen baselines.\n\nWEAKNESSES:\nI would say that the ultimate idea is rather simple, i.e. approximating both the drift function (both interactive and interaction-free), and possibly the particle density with some kind of learned approximations.\n\nI appreciate the inclusion of standard deviations in the Tables, however these values seem, particularly in Table 1, to be quite large relative to the proposed gains.\n\nI have some additional questions about the methodology. To summarize, I think this paper makes fairly solid contributions to the simulation of McKean-Vlasov equations, which is an important problem. If my questions are addressed, I would be amenable to raising my score.\n\nQUESTIONS:\nHow is the Radon-Nikodym change-of-measure learnt in the implicit model? This seems like a difficult task and I would be skeptical that it would be done without large amounts of reference data. Consequently, I am not sure what the clear advantage of IM over EM would be. Could the authors elaborate?\n\nThe usage of sample trajectories from Eq. (12) in the objective (11) seems quite non-trivial if the constraint in (11) is exact. In particular, I don’t see how (11) could be easily enforced. It seems that this is not being done exactly from Algorithms 3 and 4 in the appendix, so this should probably be clarified in the main text.\n\nI would argue that Methods 3.2 and 3.3 are more similar than given in Figure 3, since both are essentially proposing to learn the densities (but one using a Radom-Nikodym w.r.t. Time $0$ and the other a generative model). However the emphasis on neural networks in the implicit method is quite important, so perhaps the current organization is OK. \n\nThe first line in sentence 180 is redundant given the definitions.\n\nRemark 3.2 is probably better cited from Villani, particularly Section 8.3 of Topics or 9.6 of Topics.\n\n**Typos:**\nL. 591 -> I don’t think this should be pointing to section C.2.3, but rather to C.3 or something like that.\n\n**References**\nVillani, C., 2021. Topics in optimal transportation (Vol. 58). American Mathematical Soc..\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-23"}