Probability flow solution of the Fokker-Planck equation
Authors: Nicholas Matthew Boffi
Eric Vanden-Eijnden
Venue: Submitted to ICLR 2023 (2023)
Keywords: score-based diffusion
high-dimensional scientific computing

Abstract: The method of choice for integrating the time-dependent Fokker-Planck equation in high-dimension is to generate samples from the solution via integration of the associated stochastic differential equation. Here, we introduce an alternative scheme based on integrating an ordinary differential equation that describes the flow of probability. Acting as a transport map, this equation deterministically pushes samples from the initial density onto samples from the solution at any later time. Unlike integration of the stochastic dynamics, the method has the advantage of giving direct access to quantities that are challenging to estimate from trajectories alone, such as the probability current, the density itself, and its entropy. The probability flow equation depends on the gradient of the logarithm of the solution (its "score"), and so is a-priori unknown. To resolve this dependence, we model the score with a deep neural network that is learned on-the-fly by propagating a set of samples according to the instantaneous probability current. We consider several high-dimensional examples from the physics of interacting particle systems to highlight the efficiency and precision of the approach; we find that the method accurately matches analytical solutions computed by hand and moments computed via Monte-Carlo.

PROBABILITY FLOW SOLUTION OF THE FOKKER- PLANCK EQUATION
Anonymous authors Paper under double-blind review
ABSTRACT
The method of choice for integrating the time-dependent Fokker-Planck equation in high-dimension is generate samples from the solution via integration of the associated stochastic differential equation. Here, we introduce an alternative scheme based on integrating an ordinary differential equation that describes the flow of probability. Acting as transport map, this equation deterministically pushes samples from the initial density onto samples from the solution at any later time. Unlike integration of the stochastic dynamics, the method has the advantage of giving direct access to quantities that are challenging to estimate from trajectories alone, such the probability current, the density itself, and its entropy. The probability flow equation depends on the gradient the logarithm of the solution (its "score"), and so is a-priori unknown. To resolve this dependence, we model the score with a deep neural network that is learned on-the-fly by propagating set of samples according to the instantaneous probability current. We consider several high-dimensional examples from the physics of interacting particle systems to highlight the efficiency and precision of the approach; we find that the method accurately matches analytical solutions computed by hand and moments computed via Morlo-Carlo.
INTRODUCTION
The time evolution of many dynamical processes occurring in the natural sciences, engineering, economics, and statistics naturally described in the language of stochastic differential equations (SDE) (Gardiner, 2009; Oksendal, 2003; Evans, 2012). Typically, one interested in the probability density function (PDF) of these processes, which describes the probability that the system will occupy a given state at a given time. The density can be obtained as the solution to a Fokker-Planck equation (FPE), which can generically be written as (Risken, 1996; Bass, 2011)
OepE (x) = V (be(x)p (x) - (x
x ZENÇR. Rd.
(FPH)
where pf (x)) E R20 denotes the value of the density at time 6, bx(x) E Rdd is a vector field known as the drift, and Dx(x) E dRdx is a dositive-semidefinie tensor known as the diffusion matrix. (FPE) must be solved for > 0 from some initial condition (=) po(x)) but in all but the simplest cases, the solution is not available analytically and can only be approximated via numerical integration.
digh--dimensionality. For many systems of interest- - such as interacting particle systems in statistical physics (Chandler, 1987; Spohn, 2012), stochastic control systems (Kushner al., 2001), ), and models in mathematical finance (Oksendal, 2003) the dimensionality d can be very large. This renders standard numerical methods for partial differential equations inapplicable, which become infeasible for d as small as five or six due to an exponential scaling of the computational complexity with d. The standard solution to this problem a Monte-Carlo approach, whereby the SDE associated with (FPE)
dx. = (r(x))d + V D.(xz)d )tt + v2,, (z, )dW.
(1)
is evolved via numerical integration to obtain large number of trajectories (Kloeden & Platen, 1992). In (1), (z) satisfies o, [ (z)) = (x) and W, is a standard Brownian motion on R Assuming that we draw samples (=t [=) from the initial PDF Pa, simulation of (1) enables the

---

estimation of expectations via empirical averages
J (()) ed Σ o(()),
where Φ - IR is an observable of interest. While widely used, this method only provides samples from Pū. and hence other quantities of interest like the value of pt itself or the time- dependent differential entropy of the system H,- = [g log Pi (x)pp (x)dæ require sophisticated interpolation methods that typically do not scale well to high-dimension.
transport map approach. Another possibility, building on recent theoretical advances that connect transportation of measures to the Fokker-Plande equation (Jordan et al., 1998), is to recast (FPE) as the transport equation (Villani, 2009; Santambrogio, 2015)
O.pi (2) = - (x))
(3)
where we have defined the velocity field
(x) bx(x) (z (z))
(4)
This formulation reveals that Pi can be viewed the pushforward of pa under the flow map X x-, of the ordinary differential equation
d = (X;;, dt
(2) = 2.
(5)
Equation (5) is known as the probability flow equation, and its solution has the remarkable property that if x is a sample from Pu. then x will be a sample from př. Viewing X;, : Ω + Ω as transport map, P. X60 can be evaluated any position in via the change of variables formula (Villani, 2009; Santambrogio, 2015)
(-/"
- (o exp
V oz (X),
(6)
where (o is obtained by solving (5) backward from some given Importantly, access to the PDF as provided by (6) immediately gives the ability compute quantities such as the probability current or the entropy; by contrast, this capability is absent when directly simulating the SDE.
Learning the flow. The simplicity of the probability flow equation (5) is somewhat deceptive, because the velocity or depends explicitly on the solution P to the Fokker-Planck equation (FPE). Nevertheless, recent work in generative modeling via score-based diffusion (Song & Ermon, 2020a;b; Song& & Kingma, 2021) has shown that it is possible use deep neural networks to estimate vf. or equivalently the so-called score V log P. of the solution density. Here, we introduce a variant of score-based diffusion modeling in which the score is learned on-the-fly over samples generated by the probability flow equation itself. The method is self-contained and enables us to bypass simulation of the SDE entirely; moreover, we provide both empirical and theoretical evidence that the resulting self-consistent training procedure offers improved performance when compared to training via samples produced from simulation of the SDE.
1.1 CONTRIBUTIONS
Our contributions are both theoretical and computational:
We provide a bound on the Kullback-Leible divergence from the estimate P: produced via an approximate velocity field v, to the target pī. This bound motivates our approach, and shows that minimizing the discrepancy between the learned score and the score of the push-forward distribution systematically improves the accuracy of Pt.
Based on this bound, we introduce two optimization problems that can be used to learn the velocity field (4) in the transport equation (3) so that its solution coincides with that of the Fokker Planck equation (FPE), Due to its similarities with score-based diffusion approaches in generative modeling (SBDM), we call the resulting method score-based transport modeling (SBTM).
We provide specific estimators for quantities that can be computed via SBTM but not directly available from samples alone, like wise evaluation of Pt itself, the differential entropy, and the probability current.
2

---

We test SBTM on several examples involving interacting particles that pairwise repel but are kept close by common attraction to a moving trap. In these systems, the FPE is high-dimensional due to the large number of particles, which vary from 5 to 50 in the examples below. Problems of this type frequently appear in the molecular dynamics of externally driven soft matter systems (Frenkel & Smit, 2001; Spohn, 2012). We show that our method can be used to accurately compute the entropy production rate of system, a quantity of interest in the active matter community (Nardini et al., 2017), quantifies the uutoof-qquilirrimm nature of the system's 's dynamics.
1.2 NOTATION AND ASSUMPTIONS.
Throughout, we assume that the stochastic process (1) evolves over domain S C Rd in which remains at all times t 2 O. We assume that the drift vector b, 52 - R and the diffusion tensor
D₂ : Q - Rddd are wiice-iiffrrentiabee and bounded in both x and t, so that the solution to the SDE (1) is well-defined at all times t 2 0. The symmetric tensor (x) = D. (a) is assumed to be positive semi-definite for each r), with Cholesky decomposition Dx(x)) = or (x)ar (x)). We further assume that the initial PDF po is three-times differentiable, positive everywhere on St and such that Hg = Jg log po po())po(( po(()) < 0o. This guarantees that p enjoys the same properties at all times t 0. Finally, we assume that log P. is K-smooth globally for (t, x) E 00) x S2, i.e.
[
K > v(t,2)) (0) x Q
log (z) Vog (b) S K| - yl.
(7)
This technical assumption is needed to guarantee global existence and uniqueness of the solution of the probability flow equation. Throughout, we use the shorthand notation ÿ/ de Vt interchangeably for a time-dependent quantity yt-
d
2 RELATED WORK
Score matching Our approach builds directly on the toolbox of score matching originally developed by Hyvärinen (Hyvärinen. 2005; Hyvarinen, 2007; Hyvärinen, 2007: 2008) and more recently extended in the context of diffusion-based generative modeling (Song & Ermon, 2020a;b; Song et al., 2021; De Bortoli al., 2021; Dockhorn et al., 2022; Mittal et al., 2021). These approaches assume access to training samples from the target distribution (e.g., in the form of examples of natural images). Here, we bypass this need and use the probability flow equation to obtain the samples needed to learn an approximation of the score. Lu et al. (2022) recently showed that using the transport equation (TE) with velocity field learned via SBDM can lead to inaccuracies the likelihood unless higher-order score terms are cell-aaproximate. Proposition 1 shows that the self-consistent approach used in SBTM solves these issues and ensures a systematic approximation of the target p
Density estimation and Bayesian inference Our method shares commonalities with transport map- based approaches (Marzouk et al., 2016) for density estimation and variational inference (Zhang et al., 2019; Blei et al., 2017) such normalizing flows (Tabak & Vanden Eijnden,, 2010; Tabak & Turner, 2013; Rezende & Mohamed, 2016; Huang et al., 2021; Papamakarios et al., 2021; Kobyzev al., 2021). Moreover, because expectations are approximated over a set of samples according to (2), the method also inherits elements of classical "particle-based" approaches for density estimation such as Markov chain Monte Carlo (Robert & Casclla, 2004) and sequential Monte Carlo (Dai et al., 2020; Del Moral et al., 2006).
Our approach is also reminiscent of recent line of work in Bayesian inference that aims combine the strengths of particle methods with those of variational approximations (Dai et al., 2016; Saeedi et al., 2017). In particular, the method we propose bears some similarity with Stein variational gradient descent (SVGD) (Liu, 2017; Liu & Wang, 2018: 2019) (see also (Lu et al., 2018; 8; Li et al., 2020)), in that both methods approximate the target distribution via deterministic propagation of set of samples. The key differences that (i) our method learns the map used to propagate the samples, while the map in SVGD corresponds to optimization of the kernelized Stein discrepancy, and (ii) the methods have distinct goals, than sampling at equilibrium.
interested in capturing the dynamical evolution of Pī rather
Approaches for solving the FPE Most closely connected to our paper are the works by Maoutsa et al. (2020) and Shen et al. (2022), who similarly propose to bypass the SDE through use of the probability flow equation, building on earlier work by Degond & Mustieles (1990) and Russo (1990). The critical
3

---

differences between Maoutsa et al. (2020) and our approach are that they perform estimation linear space or a reproducing kernel Hilbert space rather than over the significantly richer class of neural networks, and that they train using the original score matching loss of Hyvärinen (2005), while the use of neural networks requires the introduction of regularized variants. Because of this, Maoutsa et al. (2020) studies systems of dimension less than or equal to five; in contrast, study systems with dimensionality as high as 100.
Concurrently to our work, Shen al. (2022) proposed a variational problem similar SBTM. A key difference is that SBTM is not limited to Fokker-Planck equations that can be viewed as a gradient flow in the Wasserstein metric over some energy (i.e., the drift term in the SDE (1) need not be the gradient of potential), and that it allows for spatially-dependeer and rank-deficient diffusion matrices; moreover, our theoretical results are similar but avoid the use of costly Sobolev norms.
Neural-network solutions to PDEs Our approach can also be viewed as an alternative to recent neural network-bar based methods for the solution of partial differential equations (see e.g. Yu (2017); Raissi et al. (2019); Han et al. (2018); Sirignano & Spiliopoulos (2018); Bruna et al. (2022)). Unlike these existing approaches, our method is tailored to the solution of the Fokker-Planck equation and guarantees that the solution is valid probability density. Our approach is fundamentally Lagrangian in nature, which has the advantage that it only involves learning quantities locally at the positions of a set of evolving samples; this is naturally conducive to efficient scaling for high sigh-dimensional systems.
3 METHODOLOGY
3.1 SCORE- SCORE-BASED BASED TRANSPORT MODELING
Let 8t Ω Rd denote approximation to the score of the target V log pi, and consider the solution Pr : f2 + R20 to the transport equation
(x),(x) V (vx(x))p.x))) with ve(x) bxx) -(())))(((()s(((())
(TE)
Our goal is to develop variational principle that may be used to adjust 8 so that Pt tracks Pi. Our approach is based on the following inequality, whose proof may be found in Appendix B.1: Proposition 1 (Control of the KL divergence). Assume that the conditions listed in Sec. 1.2 hold.
Let Pr. denote the solution to the transport equation (TE), and let Pł denote the solution to the Planck equation (FPR). Assume that pt=o(x) (i=o(())= = po(a)) for all x E Q. Then
d Drc 1 (p)) = 2 2 (s((x)- - V (()) B.(o) (z())dx
(8)
dt
where - -.B.(e)) = (1 D.
In partieular, (8) implies that for any T [0, oo) we have explicit control on the KI. divergence T
(KL (pt) S (" [<((xx) V (l pe(x)drdt.
Remarkably, (9) only depends on the approximate At and does not include pf: it states that the accuracy of as an approximation of P can be improved by enforcing agreement between and V log Pt. This means that we can optimize (9) without making use of external data from of, which offers self-consistent objective function to learn the score s, using (TE) alone.
The primary difficulty with this approach is that Pt must be considered as a functional of 8;, since the velocity i used in (TE) depends 8t. To render the resulting minimization of the right-hand side of (9) practical, we can exploit that (TE) can be solved via the method of characteristics, as summarized in Appendix A. Specifically, if X, (x) = 0e (Xx(z)) is the probability flow equation associated with the velocity et, then p: XiiPo- This means that the expectation of any function
((x) over (x) can be expressed as the expectation of Φi (xx()))) over (o(r). Observing that the
score of the solution to (TE) along trajectories of the probability flow V log pa (x)()) solves a closed equation leads to the following proposition.
Proposition 2 "Score-based transport modeling). Assume that the conditions listed in Sec. 1.2 hold. Define ve (x)) = (z) - (z) and consider
(x)x) = (x)()))
Xo(x) = x, Go(z) = V log log))(x).
(10)
C,(x) (Vo. G.(x)- - VV (x)(())..,
4

---

Then Pt XeAPu solves (TE), the equality Ge(z) = V log (()((() holds, and for any T [0, 00)
//2 1.² [
1
DKL(X pT) - 2
Pol
(11)
Moreover, st minimizer of the constrained optimization problem
min [-" J ( (x.(=)) po(x))xdt subject to (10)
(SBTM)
then (x) = Vz())Voo log (1 where Ph solves the Lokker-Plancee equation (FPE). The map Xi associated to any minimizer transport map from pu to Pe: i.e.
implies that
(z) ~ pi.
VI E [0.7].
(12)
Proposition is proven in Appendix B.3. The result also holds with a standard Euclidean norm replacing the Enffusion-weighhed norm, in which case the minimizer is unique and is given by (2) = V log Pf (æ). In the special case when the SDE is an and the equations for both X, and G, can be written explicitly; they studied in Appendix C.
process, the score
In practice, the objective in (SBTM) can be estimated empirically by generating samples from Po and solving the equations for (x) and Gz(z) with ~ Po. The constrained minimization problem (SBTM) can then in principle be solved with gradient-based techniques via the adjoint method. The corresponding equations are written in Appendix B.3, but they involve fourth- spatial derivatives that are computationally expensive to compute via automatic differentiation. Moreover, each gradient step requires solving a system of ordinary differential equations whose dimensionality is equal to the number of samples used to compute expectations times the dimension of (FPE). Instead, we now develop a sequential procedure that avoids these difficulties entirely.
3.2 SEQUENTIAL -BASED TRANSPORT MODELING
An alternative the constrained minimization in Proposition 2 is to consider an approach whereby the score 8; is obtained independently at each time to ensure that DKL (Pt /t) remains small. This suggests choosing to minimize dt DKL (Pt p.)), which admits a simple closed-form bound, as shown in Proposition 1. While this explicit form can be used directly, an application of Stein's identity recovers an implicit objective analogous to Hyvärinen sctre-matching that is equivalent to minimizing at (KL(P: pt)) but obviates the calculation of G,. Expanding the square in (8) and applying Sa s((x) 7 log Pa (x) pe(r)dr = - -Sa V st(z) p (z))dz, we may write
d.
/ (lsx(X DKL (p: p.) 1 2 dt
1
+ + (D((X(())ss + (G
Because log Pe (X, = (x) is independent of st, we may neglect the corresponding square term during the optimization. This leads to simple and comparatively less expensive way to build the pushforward X; such that Xp = pf sequentially in time, as stated in the following proposition.
Proposition 3 (Sequential SBTM). In the same setling Proposition 2, let Xx(x)) solve the first equation in (10) with Uy (x)) be(x) - Dz(x)s((z). Let 8t be obtained via
min ( (1s(Xx((e))) +2V (D.(((()))s (o)))))))))ppodd. po(x)dd.
(SSBTM)
Then, each minimizer St of (SSBTM) satisfies Di(x)s;()) D.(x)V log p (x) where Pt is the solution to (FPE), Moreover, the map X; associated to st is a transport map from to pi.
Proposition is proven in Appendix B.4. Critically, (SSBTM) no longer a constrained optimization problem. Given the current value of X, at any time f, we can obtain s, via direct minimization of the objective in (SSBTM). Given $t, we may compute the right-hand side of (10) and propagate X, (and possibly G₂) forward in time. The resulting procedure, which alternates between self-consistent score estimation and sample propagation, is presented in Algorithm 1. The output of the method produces feasible solution for (SBTM) with an a-posteriori bound on the loss obtained via integration.
5

---

Algorithm 1 Sequential score-based transport modeling.
1: Input: An initial time to E R20. A set of n samples {2i)i=1 from Pto set Nr timesteps
(Alx) {A Nr-!
2: Initialize sample locations x = 24 for =1= 1....,,
3: for k = 0.....,NN N₂ 1 do
4:
Optimize: Stz- = arg min, & 2i=1 ) (xi,) (xi,, + 2V (Du,
5: Propagate samples:
Xix+1 = Xix + Ats (bry (Xix) - Dra (Xi (xx
6: Set 4-11 = lx + Ale-
Output: A set of n samples {== from Pte and the score {st, (xi))== for all {tx Nir
7:
Practical considerations To avoid computation of the divergence V often costly for neural networks, we can use the denoising score matching loss function introduced by Vincent (2011), which we discuss in Appendix B.6. Empirically, we find that the use of either the denoising objective or explicit derivative regularization necessary for stable training.
which
Why not use the SDE? An alternative to the sequential procedure outlined here would be to generate samples from the target of via simulation of the associated SDE, and to approximate the score V log P via minimization of the loss So' [2(((s(22)) + 2V st(x))p; (ar)drdt, similar to SBDM. As shown in Appendix B.5 neither (PL pt) nor DKL pt) are controlled when using this procedure, where Pt XeePo the density the probability flow equation. Empirically, we find in the numerical experiments that this approach is significantly less numerically stable than sequential SBTM. In particular, we could not estimate the entropy using the score learned from the SDE.
SBTM vs. Sequential SBTM Given the simplicity of the optimization problem (SSBTM), one may wonder if (SBTM) is useful in practice, or if it simply a stepping stone to arrive at (SSBTM). The primary difference is that (SBTM) offers global control on the discrepancy between 8; and V log Pt over E [0, 2] that unavoidably arises in practice due to learning and time-discretization errors. By contrast, because (SSBTM) proceeds sequentially, these errors could accumulate over time in way that is harder to control. In the numerical examples below, we took the timestep At sufficiently small, and the number of samples sufficiently large, that did not observe any accumulation of error. Nevertheless, (SBTM) may allow for more accurate approximation, because the loss is exactly minimized at zero and high-order derivatives 8 must be controlled through the calculation of G.-
4 NUMERICAL EXPERIMENTS
In the following, we study two high dimensional examples from the physics of interacting particle systems, where the spatial variable of the Fokker-Planee Planck equation (FPE) can be written as æ = (z(1), (2) ...., ....,.(N)) (N)) with each (0(i) Rdd Here, d describes lower-dimensional ambient space, e.g. d d= 2, so that the dimensionality of the Fokker- Planck equation d Nd will be high if the number of particles N is even moderate. The still figures shown below do not do full justice to the complexity of the particle dynamics, and we encourage the reader to view the movies available here. With a timestep At = 10 3 horizon T = 10, and fixed n. = 10a, we find that the sequential SBTM procedure takes around two hours for each simulation on a single NVIDIA RTX8000 GPU.
4.1 HARMONICALLY INTERACTING PARTICLES IN A HARMONIC TRAP
Setup. Here we study a problem that admits a tractable analytical solution for direct comparison. We consider N two-dimensional particles (d 2) that repel according to harmonic interaction but experience harmonic abut experience ttraction towards a moving trap 3, E RR.. The motion of the particles is governed by the stochastic dynamics
dx(()) (B. x. + +a 2a(xx
1 N
Σ
x(0) dt /2D i=1,..,,
(13)

---

A
learned
SDE
noise free
2
1
10-2
1000
learned
entropy production
analytics tical
50
noise froe
ideal Joes
Tr(5)
Trarel
SDE
20
froe
0-
analytical
10-3
training data
SDE data
-50
0
0
0
0
2
time
time
time
Figure 1: A system of N 50 particles in harmonic trap with harmonic interaction: (A) A single sample trajectory. The mean of the trap 3, shown with red star, while past positions of the particles are indicated by fading trajectory. The noise-free system (right) is concentrated, and fails to capture the variance of the stochastic dynamics (center). The learned system (left) accurately captures the variance, and in addition generates physically interpretable trajectories for the particles. (B) Quantitative comparison to the analytical solution. The learned solution matches the entropy production rate, score, and covariance well. Movie can be found here.
where a E (0, 1) is a fixed coefficient that sets the magnitude of the repulsion. The dynamics (13) an Drnstein-Uhlenbee
process in the extended variable x E RdN with block components ((()
Assuming Gaussian initial condition, the solution the Fokker-Planck equation associated with (13)
a Gaussian for all time and hence can be characterized entirely by its mean m, and covariance Cr.
These can be obtained analytically (Appendices and D), which facilitates quantitative comparison to the leamned model. The differential entropy S₃ is given by (see Appendix D)
J.N (log (2m) + 1) log det C.
(14)
In the experiments, we take 3, = a(cos mul, sin nw)) with a 2, w = =, = 0.25, 0.5, and N 50, giving rise to 100-dimensional Fokker- Planckk equation. The particles are initialized from an isotropic Gaussian with mean 3n (the initial trap position) and variance o = 0.25.
Network architecture. We take (zx))= - -U, (æ), where the potential Uo, (-) is given sum of one- and two-particle terms
N Σ Uo (2(0))
Us. (1)) x (N))
Uo (9) (0)))
(15)
N
To ensure permutation symmetry amongst the particles, we require that Ua,2 (x. y) U0, (y, x:) for each x, y E I Modeling at the level of the potential introduces an additional gradient into the loss function, but makes it simple to enforce permutation symmetry: moreover, by writing the potential as a sum of one- and two-particle terms, the dimensionality of the function estimation problem is reduced. As motivation for this choice of architecture, show in Appendix D.I that the class of scores representable by (15) contains the analytical score for the harmonic problem considered in this section. To obtain the parameters Or, Aty perform warm start and initialize from O.., which
7

---

reduces the number of optimization steps that need to be performed at each iteration. All networks are taken to be multi-layer perceptrons with the swish sh activation function (Ramachandran et al., 2017); further details on the architectures used can be found in Appendix D.
Quantitative comparison. For a quantitative comparison between the learned model and the exact solution, we study the empirical covariance E over the samples and the entropy production rate ds. Because an analytical solution is available for this system, we may also compute the target m²) and measure the goodness of fit via the relative Fisher divergence
di
V Poog - (x
So (s((z) V log (c(x)) Sa V log (Fx))))((((()x
(16)
In Equation (16), á can be taken to be equal to the current particle estimate of Pt (the training data), or estimated using samples from the stochastic differential equation (the SDE data).
The representation of the dynamics (13) in terms of the flow of probability leads to an intuitive deterministic motion that accurately captures the statistics of the underlying stochastic process. Snapshots of particle trajectories from the learned probability flow (5), the SDE (13), and the noise-free equation obtained by setting = 0 in (13) are shown in Figure 1A.
Results.
Results for this quantitative comparison are shown in Figure IB. The learned model accurately predicts the entropy production rate of the system and minimizes the relative metric (16) to the order of 10-2 The noise-free system incorrectly predicts a constant and negative entropy production rate, while the SDE cannot make a prediction for the entropy production rate. In addition, the learned model accurately predicts the high- sigh-dimensional covariance E of the system (curves lie directly top of the analytical result, trace shown for simplicity). The SDE also captures the covariance, but exhibits more fluctuations in the estimate; the noise-free system incorrectly estimates all covariance components as converging to zero.
4.2 SOFT SPHERES IN AN ANHARMONIC TRAP
Setup. Here, we consider a system of N = 5 particles in an anharmonic trap in dimension d d= that exhibit soft-sphere repulsion. This system gives rise to 10-dimensional (FPE), dimensionality that is significantly too high for standard PDE solvers. The stochastic dynamics is given by
dx (1)
4B(B,
) -Bel2dt
A
N
(x (
x(x) exp
(x(0) - x((0)2
dt + v2D i= L=
2
2r
j=1
where 3, again represents moving trap, A >> 0 sets the strength of the repulsion between the spheres, r' sets their size, and B > 0 sets the strength of the trap. We set 3() = a(cos mwl, sin Tw]) or 2, D. 0.25, A = 10, and r 0.5. We fix B D/R² with R = V7 and y = 5.0. This ensures that the trap scales with the number of particles and that they have sufficient room in the trap generate complex dynamics. The circular case converges to a distribution P3= = p² o Qe that can be described as a fixed distribution p" composed with a time-dependent rotation Qt, and hence the entropy production rate should converge to zero. The linear case does not exhibit such convergence, and the entropy production rate should oscillate around zero as the particles are repeatedly pushed and pulled by the trap. We make use of the same network architecture as in Sec. 4.1.
8(t) a(cos Twt, 0)T with
Results. Similar Section 4. an example trajectory from the learned system, the SDE (17) and the noise-free system obtained by setting D = O in (17) are shown in Figure 2A in the circular case. The learned particle trajectories exhibit an intuitive circular motion with increased disorder relative to the frree- system that accurately captures the statistics of the stochastic dynamics. Numerical estimates of a single component of the covariance and of the entropy production rate are shown in Figure 2B/C, with all moments shown in Appendix D.2. The learned and SDE systems accurately capture the covariance, while the noise-free -rre system underestimates the covariance in both the linear and the circular case. The prediction of the entropy production rate from Algorithm 1 is reasonable in both cases, exhibiting convergence to zero and oscillation around zero as expected. In the inset, show the prediction of the entropy production rate when learning on samples from the SDE; the
8

---

A
SDE
learned
noise free
2
-2
-2
-2
B
0.6
0.6
2i
0.4
0.4
SDE
learned
Dolse
Iree
SDE
learued
Doise Free
0.2
40
40
20
20
10.0
0.0
2.5
5.0
7.5
10.0
time
time
Figure 2: A system of N = 5 soft-spheres in an anharmonic trap: (A) Example particle trajectories in the case of a rotating trap. Trap position shown with a red star. (B/C) A single component of the covariance of the samples, in the case of rotating trap (B) and linearly oscillating trap (C). The system learned system agrees well with the SDE, while the noise-free systems under-predicts the moments. (D/E) Prediction of the entropy production rate for rotating trap (B) and linearly
oscillating trap (C). Main figure depicts prediction from SBTM, while the inset depicts the prediction
when learning SDE samples. SBTM captures the temporal evolution of the entropy production rate, while learning on the SDE is initially offset and later divergent. Movies of the circular and linear motion be viewed here and here, respectively.
prediction initially offset, and later becomes divergent. We found that this behavior was generic when training the SDE, but never observed it when training self-consistent samples
OUTLOOK AND CONCLUSIONS
Building on the toolbox of score-based diffusion recently developed for generative modeling, we introduced related approach score-based transport modeling (SBTM) - that gives an alternative to simulating the corresponding SDE solve the Lokker-Plandee equation. While SBTM is more costly than integration of the SDE because it involves learning component, it gives access to quantities that are not directly accessible from the samples given by integrating the SDE, such as pointwise evaluation of the PDF, the probability current, or the entropy. Our numerical examples indicate that SBTM scalable to systems in high dimension where standard numerical techniques for partial differential equations
inapplicable. The method can be viewed integration method for the Fokker- Planckr-- equation, and our results show that its trajectories are more easily interpretable than the corresponding trajectories of the SDE.
deterministic Lagrangian

---

REFERENCES
Richard F Bass. Stochastic processes, volume 33. Cambridge University Press, 2011.
David M. Blei, Alp Kucukelbir, and Jon D. McAuliffe. Variational Inference: A Review for Statisticians. Journal of the American Statistical Association, 11 112(518):859-877, April 2017.
ISSN 0162-1459, 1537-274X. doi: 10. (0080000622499221111125116
Joan Bruna, Benjamin Peherstorfer, and Eric Vanden-Eijnden Neural galerkin scheme with active learning for high-dimensional evolution equations. arXiv:2203.01 2022.
David Chandler. Introduction to modern statistical. Mechanics. Oxford University Press, Oxford, UK 5, 1987.
Bo Dai, Niao He, Hanjun Dai, and Le Song. Provable Bayesian Inference via Particle Mirror Descent. arXiv: 1506.655 00/00 May 2016.
Chenguang Dai, Jeremy Heng, Pierre E. Jacob, and Nick Whiteley. An invitation to sequential Monte Carlo samplers. 2rXiv:2007.119336 2020.
Valentin De Bortoli, James Thortoon, Jeremy Heng, and Armadd Doucet. Diffusion Schrödinger
Bridge with Applications to Score-Based Generative Modeling. arXiv:2106.01337, 2021.
Pierre Degond and Francisco-José Mustieles. A deterministic approximation of diffusion equations using particles. SIAM Journal Scientific and Statistical Computing, 11(2):993311, 1990. doi: 10.1137/0911013
Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society: Series (Statistical Methodology), 68(3):411-44 2006.
Tim Dockhorn, Arash Vahdat, and Karsten Kreis. Score-Based Generative Modeling with Critically- Damped Langevin Diffusion. 1rXiv:21 2022.
Weinan E and Bing Yu. The deep ritz method: A deep learning based numerical algorithm for solving variational problems, 2017. URL httpt::/tta/tttti/ wwrrxivv.org/aaae org/abs/////..oo.oo.o..........r /1710. 00211.
Lawrence C Evans. An introduction to stochastic differential equations, volume 82. American Mathematical Soc., 2012.
Daan Frenkel and Berend Smit. Understanding molecular simulation: from algorithms to applications, volume 1. Elsevier, 2001.
Crispin Gardiner. Stochastic Methods. Springer- Verllg Berlin Heidelberg, 4th edition, 2009. ISBN 978-3-540-70712-7.
Jiequn Han, Arnulf Jentzen, and Weinan E. Solving high-dimensional partial differential equa- tions using deep learning. Proceedings of the National Academy of Sciences, 115(34):8505- 8510, 2018. doi: 10.1073/pnas.17 1073/pnas 1718942115.
17tpppan.....99... 1718942115. URL httph:ttp/// /www /twp//ppp.... pnw..passooom ottp:////////////////////////// abb/10.
Chin- Wei Huang, Ricky T. Q. Chen, Christos Tsirigotis, and Aaron Courville. Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization. arXiv:2012.05942,, February 2021.
Aapo Hyvarinen. Connections Between Score Matching, Contrastive Divergence, and Pseudolikeli- hood for Continuous- Vaauuu Variables. IEEE Transactions on Neural Networks, 18(5), 2007. doi: 10.1109/TNN 200..N.200..998111
Aapo Hyvärinen. Estimation of Non-Normalized Statistical Models by Score Matching. Journal of Machine Learning Research, 6(24), 2005. ISSN 1533-7928.
Aapo Hyvärinen. Some extensions of score matching. Computational Statistics & Data Analysis, 51 (5):2499-2511 2512, 2007. doi: 10. 101 2016//.csda, www.............ooooo 2006.09.003.
10

---

Aapo Hyvärinen. Optimal Approximation of Signal Priors. Neural Computation, 20(12), 2008. ISSN 0899-7667, 1530-888X. doi: 10.11622meco.2008, 10-00088-----44
Sergey loffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. arXiv:/502.03167, 2015.
Richard Jordan, David Kinderlehrer, and Felix Otto. The Variational Formulation of the Fokker- Planck Equation. SIAM Journal on Mathematical Analysis, 29(1):1-17, 1998. ISSN 1410- 1095-7154. doi: 10. 00.1377/00036141093303333
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv: 2017. arviv: 1412.6980.
Peter E Kloeden and Eckhard Platen. Stochastic differential equations. In Numerical Solution of Stochastic Differential Equations, pp. 103-160. Springer, 1992.
Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker. Normalizing Flows: An Introduction and Review of Current Methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43 (11):3964 (9999994-3979, 2021. ISSN 1939-3539. doi: 10. 10.1 10.1109TTPMII.. 202009999020
Harold Joseph Kushner Kushner, Harold Kushner, Paul G Dupuis, and Paul Dupuis. Numerical methods for stochastic control problems in continuous time, volume 24. Springer Science & Business Media, 2001.
Lei Li, Yingzhou Li, Jian-Guo Liu, Zibu Liu, and Jianfeng Lu. A stochastic version of Stein variational gradient descent for efficient sampling. Communications in Applied Mathematics and Computational Science, (511:37--3, 2020. ISSN 1559-3940, 2157-5452.
Qiang Liu. Stein Variational Gradient Descent Gradient Flow, arXiv:1704.07522 2017.
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent as Moment Matching. 1rXiv:1810.1169.. 2018
Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: General Purpose Bayesian
Inference Algorithm. (tti:1600804477// 2019.
Cheng Lu, Kaiwen Zheng, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Maximum likelihood training for score-based diffusion odes by high-order denoising score matching, 2022. URL pettss:/aaree org/ als 2206 08265.
Jianfeng Lu, Yulong Lu, and James Nolen. Scaling limit of the Stein variational gradient descent: the mean field regime. arXiv:1805.044
November 2018
Dimitra Maoutsa, Sebastian Reich, and Manfred Opper. Interacting particle solutions of Fokker- Planck equations through gradient-log-densin estimation. Entropy, 22, 2020. ISSN 1099-4300.
Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini. An introduction to sampling via measure transport. In Handbook Uncertainty Quantification: R. Ghanem, D. Higdon, and H. Owhadi, editors, pp. 1-41. Springer, 2016. doi: 1111007997833319-11255
Gautam Mittal, Jesse Engel, Curtis Hawthorne, and lan Simon. Symbolic Music Generation with Diffusion Models. arXiv:2 103. 16091, 2021. arXiv:
Cesare Nardini, Étienne Fodor, Elsen Tjhung, Frédéric van Wijland, Julien Tailleur, and Michael E. Cates. Entropy production in field theories without time-reversal symmetry: Quantifying the non-equilibrium character of active matter. Phys. Rev. X, 7:021007, Apr 2017.
Bernt Oksendal. Stochastic Differential Equations. Springer- Verlag Berlin Heidelberg, 6 edition, 2003. ISBN 198-3-642-11394--
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing Flows for Probabilistic Modeling and Inference. Journal of Machine Learning Research, 22(57):1-66) 2021. ISSN 1533-7928.
11

---

M. Raissi, P. Perdikaris, and G.E. Karniadakis. Physics-informet neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential
equations. Journal of Computational Physics, 378:686-707, 2019. ISSN 0021-9991. doi: https://
dbg..orgg10.10011ijj.pp..11111. 10.044 10.045. URL ht.tst:::: /www. sci enceednedi rect. www.t.cmmette..ee..........eetoe coom/scence/
articcleepii/500211999 9991
Prajit Ramachandran, Barret Zoph, and Quoe V. Le.
Searching for Activation Functions.
arXiv:1710.05944, 2017.
Danilo Jimenez Rezende and Shakir Mohamed. Variational Inference with Normalizing Flows. arXiv:1505.05770, 2016.
Hannes Risken. Fokker-planin planck equation. In The Fokker- Fokker-Planck Equation, pp. 63-95. Springer, 1996.
Christian P. Robert and George Casella. Monte Carlo Statistical Methods. Springer, 2004.
Giovanni Russo. Deterministic diffusion of particles. Communications on Pure and Applied Mathe-
matics, 43(6):697-733, 1990. ISSN 1097-0312.
Ardavan Saeedi, Tejas D. Kulkarni, Vikash K. Mansinghka, and Samuel Gershman. Variational
Partiele Approximations. Journal of Machine Learning Research, 18(69):1-29, 2017. ISSN 1533-7928.
Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58-63):94, 2015.
Zebang Shen, Zhenfu Wang, Satyen Kale, Alejandro Ribeiro, Aim Karbasi, and Hamed Hassani. Self-consistency of the fokker-planck equation. orXiv:2206.00800 2022.
Justin Sirignano and Konstantinos Spiliopoulos. Dgm: A deep learning algorithm for solving partial
differential equations. Journal of Computational Physics, 375:1339-1364, 2018. ISSN 0021-9991. doi: https://doi.org/10.. http:////00..1001......220111010 sssiance/arriicee///////29999999
URL https: / www.://iieeeireee...o...o com/
Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution.
arXiv: /00::/9070000 907.05000, 2020a.
Yang Song and Stefano Ermon. Improved Techniques for Training Score-Based Generative Models. arXiv:2006.0900 1, 2020b.
Yang Song and Diederik P. Kingma. How to Train Your Energy- Models. arX.v:2101.03288 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:2011.13456 2021.
Herbert Spohn. Large scale dynamics of interacting particles. Springer Science & Business Media, 2012.
E. G. Tabak and Cristina V. Turner. A Family of Nonparametric Density Estimation Algorithms. Communications on Pure and Applied Mathematics, (6(2)):15-164, 2013. ISSN 1097-0312. doi: 10.1002//pa.21423.
Esteban G. Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent of the log-likelihood.
Communications in Mathematical Sciences, 8(1):217-233, 2010. ISSN 15396746, 19450796. doi:
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
Pascal Vincent. Connection Between Score Matching and Denoising Autoencoders. Neural
Computation, 23(7):1661-1674 2011. ISSN 0899-7667, 1530-888X. doi: 10.1162/NECO 00142.
Cheng Zhang. Judith Bütepage, Hedvig Kjellström, and Stephan Mandt. Advances in Variational Inference. IEEE Transactions on Pattern Analysis and Machine Intelligence, 41(8):2008-200 August 2019. ISSN 1939-3539. doi: 10.1 209/TPAM..201.... 2.2889777.
12

---

A SOME BASIC FORMULAS
Here, we derive some results linking the solution of the transport equation (TE) with that of the probability flow equation (5).
A.1
PROBABILITY DENSITY AND PROBABILITY CURRENT
We begin with lemma.
Lemma A.1. Let Pt: : Ω Rzo satisfy the transport equation
O.pe(a)) - (v)(z))x<z)).
(A.1)
Assume that ve(x) is C² in both t and x for 2 0 and globally Lipschitz in x. Then, given any t, 2 the solution of (A.1) satisfies
(-f"v U. (x (-
px(x) = pe (x, (x)) exp
(A.2)
where Nx is the probability flow solution to (5). In addition, given any test function 6: : + R, we have
[ ((z)) = J. (x)) (x)da..
(A.3)
In words, Lemma A. states that an evaluation of the PDF at given point may be obtained by evolving the probability flow equation (5) backwards to some earlier time t' to find the point x' that evolves to x at time assuming that (x()
( is available. In particular, for t' = O, we obtain
(x(x) - (o(x,(())) exp
v-
(A.4)
and
J (()) /
(A.5)
Since the probability current is by definition (x()))((x). using (A.4) to express pr(x) also gives the follwing equation for the current:
(x)))((()) = (())))((X exp (- vr
(A.6)
Proof. The assumed C2 and globally Lipschitz conditions on de guarantee global existence (on 2 0) and uniqueness of the solution to (5). Differentiating (((e,,(()) with respect to t and using (5) and (A.1) we deduce
d
te (((,,(())) (
di (x) vp((Xe,(())
(A.7)
-e, (x), vpe(Xe:(z))
- - -v (x)) (X),
Integrating this equation in from t = t' tot = t gives
pr p+ (x)r,(x)) = Pr (2r) exp (- (-f.'v ur (xx,(((()
(A.8)
Evaluating this expression x e (I) and using the group properties (i) X ((x,, (r (z)) r and (ii) (X,, (x)) = Xz,(xx) gives (A.2). Equation (A.3) can be derived by using (A.2) to express (e(x:) in the integral at the left hand-side, changing integration variable x - (x) and noting
that the factor exp (-S, f V (x)) is precisely the Jacobian of this change of variable. The result is the integral at the right hand-side of (A.3).
I1
Lemma A.I also holds locally in time for any (x) that is C2 in both & and x. In particular, it holds locally if we set (x) V log (z) and if we assume that Po(xz) is (i) positive everywhere Ω and (ii) 8 in x. In this case, (A.1) is the Fokker-Planne equation (FPE) and (A.2) holds for the solution that equation.
13

---

A.2 CALCULATION OF THE DIFFERENTIAL ENTROPY
We now consider computation of the differential entropy, and state similar result.
Lemma A.2. Assume that po S2 + Rzo is positive everywhere on Q and C in its argument. Let the Fokker Planck equation (FPE) (or equivalently, to the transport equation (A.1) with St (x)) = V log pt(z) in the definition of Up (x))). Then the differential
P: : Ω + Rzo denote the solution
entropy 8t Sa log
pe(z)) dæ can expressed as (x) log
( H1 1 log (x), po(x)) Ho (-' V (x)( 12
(A.9)
or
H. = Ho - S o (x(oo.((1)) ur
(A.10)
Proof. We first derive (A.9). Observe that applying (A.5) with a = log Pe leads to the first equality. The second can then be deduced from (A.4). To derive (A.10), notice that from (A.1),
d. H, log -ogp(l)VV((u(())000 (v.(=)p((z) de,
dt
Vlogm(()) (x(r))
(A.11)
se(z) ue (z))
Above, we used integration by parts to obtain the second equality and s; = V log Pt to get the third. Now, using (A.5) with @ St integrating the result gives (A.10).
A.3 RESAMPLING OF Pt AT ANY TIME
If the score RJ V log pt. is known sufficient accuracy, Pr can be resampled at any time / using the dynamics
dX, =s(Xx)dt + dW,.
(A.12)
In (A.12), r is an artificial time used for sampling that is distinct from the physical time in (1). For 8, V log Pt, the equilibrium distribution of (A.12) exactly Pt. practice, 8; will imperfect and will have error that increases away from the samples used to learn it; as result, (A.12) 12.12) should be used near samples for a fixed amount of time to avoid the introduction of additional errors.
B FURTHER DETAILS ON SCORE-BASE TRANSPORT MODELING
B. BOUNDING THE KL DIVERGENCE
Let us restate Proposition I for convenience:
Proposition (Control of the KL divergence). Assume that the conditions listed in Sec. 1.2 hold. Let Pt denote the solution the transport equation (TE), and let pł denote the solution to the Fokker-Plande equation (FPE). Assume that P==u(z) = (i-o(z) = po(x) for all x E Ω. Then
d Dra (pr l < di
1 |s((x) (z) P, (x)dx,
(8)
2 n
where 112,00) (.
14

---

Proof. By assumption, Ps solves (TE) and pi solves (FPE). Denote by vx(x) b,(x) (i(x)s((a))
and (z(x)) be(x)) (z)))) (x) with s7 (z) = V log (r)) Then, we have d dt
d
/ log
pe(x))
(pL pt)) dt n
Pr pr(2))de,
p1(x))
pr(a))
/ log (c(T)) depi(a)dz, pl(x)
Dept (1)dd
(z)
/ (x)) v (Pi(e)
(p((z) Pi (i())do / v1(x)) Vlog [c((z)
pe(x)dd,
pt(x))
/ (x) - (x)) (Vlog)t(1) - V Vlog pi (z)) (c())dr
(x) - sx)) D2 (z) (Vogg((>)) of (x)) pr (z)dz..
Above, we used integration by parts to obtain the third equality. Now, dropping function arguments for simplicity of notation, have that
V log - st = V log - si + s - 8t selD.
= V log si + |s - stI.. + 22Vlog - 8t) De(s -st), (B.1)
2 2(V log - si) De(s - ss).
Hence, deduce that
d dt
DKL (pe p.) $ o (s((x) v log po(z) (m)
(B.2)
B.2 SBTM EULERIAN FRAME
The Eulerian equivalent of Proposition 2 can be stated as:
Proposition B.1 (SBTM in the Eulerian frame). Assume that the conditions listed in Sec. 1.2 hold. Fix T (0, oo] and consider the optimization problem
min L'" [sc(x)) v log perpe()) Polz)dodt (s
(SBTM2)
subject to: (ep((z)= = -
x E
with vx(x) (x)) - (x(x)s((x). Then every minimizer of (SBTM2) satisfies D,(x)s²(x) Dx(x)V log P (x:) where Pt Ω - Ryo solves (FPE).
In words, this proposition states that solving the constrained optimization problem (SBTM2) equivalent to solving the Fokker-Plance equation (FPE).
Proof. The constrained minimization problem (SBTM2) can be handled by considering the extended objective
J." J/ (1ss(0)) V log (lgg((()))a (x) (c(x) (x) (0) + (v:()))(=))) drdt (B.3)
where u, (x) = (x)))-111 (x)ss, and μ Rd Rzo is a Lagrange multiplier. The Euler-Lagrange equations associated with (B.3) read
(ep((x) V (x)
0(m((z) =x) (p((() (sx(z) (.(x) - V Pi (.))))
+2V [D(()) (s)(x) Vlogpt(z))
(B.4)
=Pr(x),
(x(x) (x) v logg((())) (c(x) D
15

---

Clearly, these equations will be satisfied if (z) = V Vog (x) for all E Ω2 p4 (z) for all a, and P solves (FPE). This solution is also global minimizer, because it zeroes the value of the objective. Moreover, all global minimizers must satisfy (x) Da log p (æ) (Pe-almost everywhere), as this is the only way to zero the objective.
It is also easy that there other local minimizers. To check this, can the fourth equation to write
De(x)(s((a)) - Vlog -
Then,
(s,(1)), - log (x)) (s,(2) Vl) (.(x)V ((())))))
This reduces the first three equations to
@xp((z) = - (b.(x))((z) - Di(1) - (x(D)(()VV((tt)
0:11: = -bb(xx) - -
(M((T)
(B.5)
+ V (D.(=)Vm(z)) + (s) V
V log (x))) (r())V
(T(x)= 0.
Since the equation for μ homogeneous in μr and pr 0, we must have μr for all & [0, T) and the equation for Pr reduces to (FPE).
B.3 SBTM IN THE LAGRANGIAN FRAME
As stated, Proposition B. is not practical, because is phrased in terms of the density pr. The follow- ing result demonstrates that the transport map identity (6) can be used to re-express Proposition B.1 entirely in terms of known quantities.
Proposition 2 (Score-based transport modeling). Assume that the conditions listed in Sec. 1.2 hold. Define Ue (x)) = (z) - (x(r)s((z) and consider
(x(x) = (x))),
Xo(x) = x. Go(æ) = V log po(x))
(10)
G.(x) -Vxxx(XX - VV -V-w (x((x)))
Then Pt XiáPo solves (TE), the equality Gr(z) = Vlog holds, and for any T E [0, 0o)
(-' -Gl Polr))rr
T
DxL(X p S
(11)
Moreover, st a minimizer of the constrained optimization problem
min f' S (s((X((z)) - G.( (x.te)) po(z)dzdt subject to (10)
(SBTM)
then (z( (x(z) log p (x) where Pt solves the Fokker-Planne Planck- equation (FPE). The map X; associated to any minimizer a transport map from Po to Pł, i.e.
x ~ Pn implies that
(i(x) ~ Pi,
E
(12)
Proof. Let us first show that G.(x) = log pt (x)()) satisfies (10) if P. XetPo. i.e. if pt satisfies the transport equation (TE). Since (TE) implies that
0 log (z) + (x) Vog -V -v-
(B.6)
taking the gradient gives
0, log + [V, (z) + VV logp:(x) (x(x) -V (x(x).
(B.7)
Therefore G,(x) (og solves
-. (X (Cce))=AAAT0agaXXX00)))T)Ta7aaa4 (x,(x).
di
0. (x) o((x),
(B.8)
=vV ----
16

---

which recovers the equation for Gz(z) in (10). Hence, the objective in (SBTM) can also be written as
J.' / (s((X((z)) Vlogg((X(()))) po(x)dxdt
(B.9)
= =Io J. (s,(()) V Vlog (x)) pr(d)drdt
where the second equality follows from (A.5) if (z) satisfies (A.1), Hence, (SBTM) is equivalent to (SBTM2). The bound on DKL (Xrrtpo PT) follows from (9).
Adjoint equations. In terms of a practical implementation, the objective in (SBTM2) can be evaluated by generating samples {2r] from Pu and solving the equations for X₁ and G2 using the initial conditions (z;);) = 2, and Go = V log (a(z)). Note that evaluating this second initial condition only requires to know Po up normalization factor. To evaluate the gradient of the objective, we can introduce equations adjoint to those for X, and G.- They read, respectively
d di 0x()) - (Vv,
(())-VVVV(X(()))))))
+n((2) VVVV((X 0r(x)=
(B.10)
n((x) (x) ((G((z) (x(())))
mT( O.
In terms of these functions, the gradient of the objective is the gradient with respect to sx(x)) (or the parameters this function when is modeled by neural network) of the extended objective:
Lise] = L/" / (x(()(())) Go(z)) -G(z)) ( po(x)dzdt
{/" [ Oc(z)) (x:(:)- po(x))xdtt
T
(B.11)
[-¹
T
[ m(()) (G.(s)) + (Vr((X
- oV, p
where uz(x) (z) Dz(z)s((z).
B.4 SEQUENTIAL SBTM
Let us restate Proposition 3 for convenience:
Proposition 3 (Sequential SBTM). In the same setting as Proposition 2, let (x()) solve the first equation in (10) with Ue (xx) be(x)) (x(x))((x)). Let 8t be obtained via
min /
Sq (Ix((x(()))
+ 2V (D.(X(()))s (o()))
(SSBTM)
Then, each minimizer st of (SSBTM) satisfies (x(z)s((x) = Do(z) where pf is the solution to (FPF). Moreover, the map X; associated to st is a transport map from Po to pt.
Proof. If X,EPo Pe. then by definition we have the identity
/ (Ixx + +2 (xx(X)()))s pola)do
= =/ (1s(6))) + +2 (D.(z)s((z)) (c(x)dr.
(B.12)
17

---

This means that the optimization problem in (SSBTM) is equivalent to
min / (Ia(E))BB.c) + (D((z)s((x))) pclz)dz. #t
All minimizers of of this optimization problem satisfy (,(r)s logr))) Hence,
by (TE),
depr(x) - V. ( which recovers (FPE), that pr(x) (x) solves (FPE).
(B.13)
B.5 LEARNING FROM THE SDE
In this section, we show that learning from the SDE alone iie.,, avoiding the use of self-consistent samples does not provide a guarantee on the accuracy of We have already seen in (9) that it is sufficient to control Jo Jo (x)) V log pl (x)) (2) (x)d to control DKL (pr p7-). The proof of Proposition shows that control on
[' (s((x)) V log pi (x) (x)d
(B.14)
as would be provided by training on samples from the SDE, does not ensure control on DKL (pr PT))- The following proposition shows that control on 14) does not guarantee control on DKL (pr PP) either. An analogous result appeared in Lu et al. (2022) in the context of SBDM for generative modeling; here, we provide a self-contained treatment to motivate the use the sequential SBTM procedure discussed in the main text.
Proposition B.2. Let p: Ω + Rzo solve (TE), and let P - R solve (FPE). Then, the following equality holds
1-" / (s((x) Veaii--logppaae (o) (x)dxdt
DKL(PT (PT)
(B.15)
1/" / (Vlogp(())+s(()) si(x)) Di(x) (s((x) V V pi (z)) pt (x)drdt.
Proposition B.2 shows that minimizing the error between 8 and V log Pr on samples of p leaves remainder term, because in general V log + st. The proof shows that we may obtain the simple upper bound
1 /'/ / [s((x) V (x)
DKL (PT) S -
(a)) (x) (x)dxdt
(B.16)
1 + 1/ / [x(()) (o) (z)drdt.
However, controlling the above quantity requires enforcing agreement between and V log in addition to s, and V log pf; this is precisely the idea of SBTM.
Proof. By symmetry, we may replace pt by P in the proof of Proposition I to find
d
DiL I / (Vlog V log pf (x)) (s(=)) V log pf (+)) pt
dt
Adding and subtracting (x) the first term in the inner product and expanding gives
dt DKL(p? pb) = / [s(( v Poo
(B.17)
+ #/6 (Vlogg(())) sx(x)) D+(x) (s<(x) V (())) (x)dd,
Integrating from 0 to I completes the proof.
18

---

B.6 DENOISING LOss
The following standard trick can be used to avoid computing the divergence of (x
Lemma B.3, Given Ç NO, I), we have
lim a (sx a ac) .E) V st(),
a40
(B.18)
lim (()) ox = (Di(())sss(tt) a40
Proof. We have
8e + - aa=a-= + (VV((())))
5 5+
(B.19)
The expectation of the first term the right-hand side of this equation is zero; the expectation the second gives the result in (B.18). Hence, taking the expectation of (B.19) and evaluating the result in the limit as a gives the first equation in (B.18). The second equation in (B.18) can be proven similarly using o (x) = Dz(r).
Replacing V st(z) in (SSBTM) with the first expression in (B.19) for a fixed a >> 0 gives the loss
[ = Eg V/ (ss((x(()))) 2 s (X)(x) + a 5) Po (x))dx
(B.20)
Evaluating the square term at perturbed data point recovers the denoising loss of Vincent (2011)
[ Ee J sz (xx((x)))++++))+)))))))))))))( + a + (o(x)d
(B.21)
We can improve the accuracy of the approximation with a "doubling trick" that applies two draws of the noise of opposite sign reduce the variance. This amounts replacing the expectations in (B.18) with
2 B (x ag) E st(x - ag)) E).
(B.22)
a - [s((1 (1)()) (x(x) se(a (0, (x)
whose limits as + are (ar) and (>)))) respectively. In practice, we observe that this approach always helps stabilize training. Moreover, we observe that use of the denoising loss also stabilizes training, so that it is preferable to full computation V (x() even when the latter is feasible.
C GAUSSIAN CASE
Here, consider the case of an analytically, thereby providing a benchmark for our approach. The example treated in Section 4. with details in Appendix D.1 is special case of such an OU process with additional symmetry arising from permutations of the particles.
(OU) process where the score be written
The SDE reads
dx, = -T, -b) + d2,,W,
(C.1)
where X₂ E Rd, 1', E pdxd is a time-dependent positive-definite tensor (not necessarily symmetric), by E Rd time-dependent vector, and o; E IRdxd is dimendepndeent tensor. The Fokker-Plame equation associated with (C.1) is
(z) - b.)
(C.2)
where D, = T.oo. Assuming that the initial condition is Gaussian, Po = N(mp, Co) with Co = CJ dxd dositive-definite, the solution is Gaussian at all times > 0. P N(m,, C.) with me and C, C solutions to
m (m, b.)
(C.3)
C = -1,C, - -CT + +2D
19

---

This implies in particular that
V Plog -r- - m²).
(C.4)
so that the probability flow equation for X, and the equation for G5 written in (SBTM2) read
xx(x) (D,C, -T) Iibe - DDC
(C.5)
Gx(x)= - (P. - -C,
with initial condition Xo(z) = and Go(æ) = Vlogpo(z) = -c -ma). It is easy to see that with 2 ~ Po = N(mo: Cb) we have X((x) ~ Pi = N(m,, Ci) since, from the first equation in (C.5), the mean and variance of X2 satisfy (C.3). Similarly, when r Po =N(mo, Co), Go(z) ~ (0, 1) that Gz(r) N(O, CT') because the second equation in (C.5) is linear and hence preserves Gaussianity. Moreover, (x) O and B, B, = Eo (r) (æ)] satisfies
d B, = (T. -T D,),, + B. D,C,")
(C.6)
dt
The solution to this equation B, = Cτ¹ for C, that we can deduce from (C.3)
since substituting this ansatz into (C.6) gives the equation
d c. C. 0 - 2 20,'D,,C
(C.7)
Note that I', = I, b, = b, and D. = are all time-independent, then lim Pt = N(moo Coo) with m.x b and Coc the solution to the Lyapunov matrix equation
TC + C 2D.
(C.8)
This means that at long times the coefficients the right-hand sides of (C.5) also settle on constant values. However, X₁ and G2 do not necessarily stop evolving; one situation where they too tend to fix values is when the OU process in detailed balance, i.c. when I' = DA for some A AT Rxxd positive-definite In that case, the solution to (C.8) is Co = A and it is easy to see that at long times the right hand sides of (C.5) tend to zero.
Remark C.I. This last conclusion is actually more generic than for a simple OU process. For any SDE in detailed balance, i.c. that can be written as
dx,- DUX,) + V D(Xi)dt + (2o((X))dW
(C.9)
where U Rd + R5o is a -o-pttnttall such that Z = Jg (U(()) dx < 0o, we have that , (x( = --() and the corresponding flows X, and G, eventually stop L - co. In this case, Pr follows gradient descent in W2 over the energy
{ (U(x) log p(z)))(()dd J
Elpl
(C.10)
The unique PDF minimizing this energy 2 -v( and as I + 00 X, converges towards a transport map between the initial Po and - UU())
D EXPERIMENTAL DETAILS AND ADDITIONAL EXAMPLES
All numerical experiments were performed in jax using the Amnhaiku package to implement the networks and the optax package for optimization.
D.1 HARMONICALLY INTERACTING PARTICLES A HARMONIC TRAP
Network architecture Both the single-particle energy Uo Rd R and two-particle interaction energy U ddd x dd + R are parameterized as single hidden-layer neural networks with the swish wissh activation function (Ramachandran et al., 2017) and n hidden = 100 hidden neurons. The hidden layer biases initialized to zero while the hidden layer weights are initialized from
20

---

truncated normal distribution with variance 1/fan in, following the guidelines recommended in (loffe & Szegedy, 2015).
Optimization The Adam (Kingma & Ba, 2017) optimizer is used with an initial learning rate of 7 10 and otherwise default settings. At time = 0, the analytical relative loss
|s((2)) (z) V log po ( (2)
[iso]
(D.1)
V log
minimized value less than 10-4 using knowledge of the initial condition Pu N (So. T with do - 0.25. In (D.1), the expectation with respect to Po is approximated by an initial set of samples x = 2 (2) ..) T
with j = .,...," n drawn from po- the experiments, we set n 100. We set the physical timestep Al = 10-3 and take n_opt_steps 25 steps of Adam until the norm of the gradient is below gtol 0.1.
Analytical moments First define the mean, second moment, and covariance according to
m (i) = E [x
M (1) = E [ (x J
c M (i)
m (1) (m())
It is straightforward to show that the mean and covariance obey the dynamics
N
( m = (m, - 31) +
Σ
m
- m
(D.2)
N
N
(1() = =2(1-a)) 2D16:
Σ
(, c((h))
(D.3)
N
k=1
Because the particles are indistinguishable so long as they are initialized from a distribution that symmetric with respect permutations of their labeling, the moments will satisfy the ansatz
m (4) m(t); I,...,N
(D.4)
((1)) Ca(());; + C.(()( -
6zz), i,j I...,.
(D.5)
The dynamics for the vector m R20 - Rd. well as the matrices C. R20 - Rdxd and C, R20 + gdxd can then be obtained from (D.2) and (D.3) as
m = 3 m,
č. = 2(a 1)Ca -
a
(Ca (n 1)C%) + 2DI,
N
C = 2( 1 - (Ca (n - 1)Co).
a
N
For given R- Rd, these equations can be solved analytically in Mathematica as function of time, giving the mean m, mtt) @ IN E Nd and covariance C, (Ca(t) - C.(t)) INxN C,()) & (Iv1k)= E Ndxdd Because the solution is Gaussian for all t, we then obtain the analytical solution the Fokker-Planne equation pi N (m,, C.) and the corresponding analytical score -V pf (x) C, (x-m)..
Potential structure Here, we show that the potential for this example lies in the class of potentials described by (15). From Equation D.5, we have characterization the structure of the covariance matrix Ct for the analytical potential (xr) circulant, and hence block diagonalized by the roots of unity (the block discrete Fourier transform).
(x m²) C (Ir m²). ). In particular, C2 is block
That is, may take "block eigenvector" of the form Wh (Iaxao 1 ..... I (N with exp(-2mi/ for l = O....N 1. By direct calculation, this block diagonalization leads two distinct block eigenmatrices,
Ca(t) + (N -))C.(!) Call)
0
0
C.(1)
0
V-1
0
0
0
0
0
21

---

where V R N Nå denotes the matrix with block columns wy. The inverse matrix CT then must similarly have only two distinct block eigenmatrices given by (Ca(t) + (N (t)) and (Ca(l) - - By inversion of the block Fourier transform, we then find that
(CT1) (9) Cadi + Co(1 -ig)
for some matrices Cd, Co. Hence, by direct calculation
(x - m.)² Cr (x m²) = Σ
(0 - m())² (CT1)(6) (c)1)(00(s5)-
(s0))mm(²) -mmm mlm))
Σ
8ij)) (200) (0()) -m()) (Caōej Co(1-6ij)
m(())
N
(2(+) - m(t)) Ea (z(0) - -m(tt)
(1 1 (m1)-m())) Co (z.00) ml())
igj
Above, we may identify the first term in the last line as Eiz (1(1)) and the second term in the last line as N Enyy 1 z()).. Moreover, U2(-,-) is symmetric with respect to its arguments.
Analytical Entropy For this example, the entropy can be computed analytically and compared directly the learned numerical estimate. By definition,
81
JzNv
lo (())),
(_Nd
.Nd
1 7(( (m-m)) Tm (x mi) p (x()) 1 log det C. -
2 2
log(22 (2T)
RNd
Nd
(log (2Tr) + +) + log det Cr.
Additional figures Images of the learned velocity field and potential in comparison to the correspond- ing analytical solutions can be found in Figures D. and D.2, respectively. Further detail can be found in the corresponding captions. We stress that the two-dimensional images represent single-particle slices of the high-dimensional functions.
D.2 SOFT SPHERES AN ANHARMONIC TRAP
Network architecture Both potential terms Ue..1 and Ve. are modeled as four hidden-layer deep fully connected networks with hidden = 32 neurons in each layer. The initialization is identical to Appendix D.2.
Optimization and initialization The Adam optimizer is used with an initial learning rate of r) 5 6x00 3 and otherwise default settings, At time 0, the loss (D.1) is minimized to a value less than 10-1 over n samples æo.j N(3a, of with ao 0.5 and = 1000, similar to Appendix D.2. After this initial optimization, 100 steps of the SDE (17) are taken in artificial time T with fixed physical t = 0 to ensure that no spheres are overlapping at initialization. Past this initial stage, the denoising loss is used with a noise scale o = 0.025. The loss is minimized by taking opt_steps = 25 steps of Adam until the norm of the gradient is below gtol 0.5. The physical timestep set to At 10
Additional figures A depiction of the one- mne-particle potential, estimated as the negative logarithm of the one- one-partion particle PDF obtained via kernel density estimation, can be found in Figure D.3 (for further details, the caption).
22

---

learned
snalytical
noise free
2
2
2
2
0
2
10
15
2
Figure D.I: A system of N = 50 harmonically interacting particles in a harmonic trap: slices of the high-dimensional velocity field. Cross sections of the velocity field for N = 50 harmonically interacting particles in a moving harmonic trap. Columns depict the learned, analytical, noise-free free, and error between the learned and analytical velocity fields, respectively. Rows indicate different time points, corresponding to / = 1.25, 2.5, 3.75, and 5.0, respectively. Fach velocity field is plotted function of single particle's 's coordinate (denoted fixed to be at the location of sample. Color depicts the magnitude of the velocity field while arrows indicate the direction. Learned, analytical, and noise-free share colorbar for direct comparison; the error occurs on different scale and is plotted with its own colorbar. White circles in the error plot indicate samples projected onto the zy plane; locations of low error correlate well with the presence of samples.
and y); all other particle coordinates
23

---

learned
analytical
CiTor
2
-2
2
2
2
01
-2
2
S 0
-2
2
0
2
2
2
-2
x
5
10
15
0.0
2.5
5.0
Figure D.2: system of 50 harmonically interacting particles harmonic trap: slices of the dimenniinai potential. Cross sections of the potential field Ug, (r) computed via (15). Columns depict the learned, analytical, and error between the learned and analytical, respectively. Rows indicate distinct time points, corresponding to 1.25, 2.5, 3.75, and 5.0, respectively. As Figure D.I, each potential field is plotted as a function of single particle's s coordinate (denoted and with other particle coordinates fixed sample. All potentials are normalized via overall shift so that the minimum value is zero. White circles the error plot indicate samples from the learned system projected onto the ry plane.
24

---

learned
SDE
noisc free
R
0
0
1
1
1
T
T
T
Figure D.3: A system of N = 5 soft-sphere particles in an anharmonic trap: one-particle potential. Cross sections of the one-particle potential field U(x) - log (KDE(#)) where PKDE denotes kernel density estimate of the one-particle density obtained by pooling all particles and treating them as equivalent woo-dimensional samples, shown relative the moving mean. Columns depict the learned, SDE, and noise free systems, respectively. Purple dots indicate samples from the corresponding system. Rows indicate distinct time points, corresponding to 1.25, 2.5, 3.75, and 4.95, respectively. All potentials are normalized via an overall shift that the minimum value is zero, and are clipped to a maximum value 15. The learned and SDE potentials match well, while the noise free KDE becomes too peaked and develops spurious maximum that causes the particles to align in a ring.
25

---

600
WWWMWM
0.5
0.0
WwMWW
0.5
89999898
Avmmmw
0.5
010
0.5
47
10
10
18ce
Figure D.4: A system of = 5 soft-sphere particles an anharmonic trap: moments. All components of the covariance matrix over time for the circular trap motion.
26

---

MMMMMM
53
27
.7
10
18mm
13km
Figure D.5: A system of = 5 soft-sphere particles an anharmonic trap: moments. All components of the covariance matrix over time for the linear trap motion.
27

---

learned
SDE
noise free
2
33
2
22
3
3
Figure D.6: An active swimmer: sample trajectories. Samples in the plane. Columns denote solution type and rows indicate snapshots in time (t = 0, 0.25, 0.5, 3.0, respectively). The learned and SDE systems develop bimodality while the noise free system collapses with time and does correctly capture the variance.
28

---

D.3 AN ACTIVE SWIMMER
Here, we study an "active swimmer" model that describes the motion of a particle in an anharmonic trap with preference travel in noisy direction. The system is tvo-dimensional, and is given by the stochastic differential equation for the position x and velocity u
dx = (-x² + +0) dt,
(D.6)
dc -uddt + DdW,
Despite its low-dimensionality, (D.6) exhibits convergence to non-equilibrium statistical steady state in which the probability current je(x) = (r)) is non-zero.
Setup We set 7 0.1 and D 1.0. Because noise only enters the system through the velocity variable v in (D.6), the score can be taken to be one-dimensional. This is equivalent to learning the score only in the range of the rank-deficient diffusion matrix. We parameterize the score directly s, I22 + R using a threc-hidden layer neural network with hidden = 32 neurons per hidden layer.
Optimization and initialization The network initialization is identical to the previous two exper- iments. The physical timestep is set to At 10-3 The Adam optimizer is used with an initial learning rate of 10-4, At time / = 0 the loss (D.1) is minimized to a tolerance of 10-4 over n 5000 samples drawn from an initial distribution N(O, oğ!) with dg 1. The denoising loss used with a noise scale = 0.05, using __opt_stepp__pe steps 25 steps of Adam until the norm of the gradient is below gtol 0.5.
Results Depictions of the sample trajectories (z,(t), v, (t) in phase space are shown in Figure D.6. The trajectories demonstrate that the distribution the learned samples qualitatively matches the distribution of the SDE samples. The noise- free system grows increasingly and overly compressed with time. The learned velocity field effectively captures a non-zero rotational steady-state current that qualitatively matches the current of the SDE but enjoys more interpretable sample trajectories.
movie the motion of the samples (It,, U;) in phase space can be seen here. The movie highlights convergence of the learned solution to non-zero steady-state probability current that qualitatively matches that of the SDE. By contrast, the noise-free system becomes increasingly concentrated with time, failing accurately capture the current. Figure D.7 depicts the learned velocity field (x(2)) (z) - (x). The figure highlights the structure of the steady-state current, which contains an elliptical region with closed orbits. The elliptical region remains roughly fixed in size as time proceeds, while the orbits of the noise-free free system in Figure D.8 become increasing compressed.
Kernel density estimation demonstrates that an estimated PDF for the samples of learned solution qualitatively matches that of the SDE (Figure D.9).
29

---

1
0-
-2
2
1
0
2
2
-2
-2
2
-2
1
0
2
2
0
1
2.5
5.0
7.5
Figure D.7: An active swimmer: learned velocity. The learned velocity field (right-hand) hightthandd side of (5)) for the active swimmer example. Color indicates the magnitude of the velocity field computed on a grid, while arrows indicate the direction of the velocity field on samples. Time corresponds to progressing in the grid along columns from the top-left to the bottom-right image (f = k x .75 with the image number, zero-indexed). The learned velocity field converges to closed streamlines that enforce nonzero steady-state current.
30

---

-1
2
2
1
2
2
0
1
0.5
1.0
Figure D.8: An active swimmer: noise free velocity. Noise free velocity field. As in Figure D.7, color indicates the magnitude the velocity field while arrows indicate the direction, and time corresponds to progressing in the grid along columns from the top-left the bottom-right image (t x .75 with & the image number, zero-indexed). The velocity field in the noise-free case incorrectly pushes the swimmers lie along thin band.
31

---

learned
SDE
noise free
2-
0-
-2
2
0
-2
-2
0.0
0.2
0.4
0.6
Figure D.9: An active swimmer: density. PDFs computed via kernel density estimation in the plane. Columns denote solution type and rows denote snapshots time (t O, 0.5, 1.5, 6.0, respectively). Similar to the samples presented in Figure D.6, the KDE reveals bimodality in the probability density due to the presence of the particle velocity field. The noise free system becomes too concentrated and does accurately capture the shape of the SDE and learned solutions, while the SDE and learned solutions are nearly identical.
32


Reviews:
[{"review_id": "vzAn3w2kt3L", "reviewer": "Reviewer_1NXr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In high dimensional setting, it is typically impossible to approximate the solution $\\rho_t$ to a Fokker Planck Equation (FPE) with conventional grid-based methods. A  standard method consists in considering the associated SDE a simulate many trajectories from this SDE to collect statistics of the solution at any future time $t$. This methods allows to compute any moment $\\int \\mathcal{O}(x) \\rho_t(dx)$ of the solution (eg. Monte-Carlo method), but the evaluation of $\\rho_t(x)$ is typically not available with this basic approach.\n\nThe article proposes to express the solution of the FPE as a transport equation $\\partial_t \\rho_t = -\\nabla \\cdot (v_t \\rho_t)$ with a velocity field $v_t$ that is expressed as a function of the score $s_t$ of $\\rho_t$ itself, $s_t = \\nabla \\rho_t$. This motivates the following particle approach:\n\n1. start from $N$ samples $x_1, \\ldots, x_N$ samples from $\\rho_0$\n2. use (some variant of) score matching to evaluate the score of $\\rho_t$ and construct the velocity field $v_t$\n3. push forward the particles through the transport equation to get a particle approximation of $\\rho_{t + \\Delta t}$\n4. iterate: go back to 2\n\nThe advantage of this approach is that, because $\\rho_t$ is expressed as transport equation, is is straightforward to evaluate $\\rho_t$ along trajectories of the transported particles (contrarily to the SDE approach).", "summary_of_the_review": "The proposed method is natural and very interesting. I would like to encourage the authors to work on their simulations in order to more convincingly demonstrate the appeal of the approach when compared to more standard methodologies.", "strength_and_weaknesses": "**Strength:**\nThe method is very intuitive and appears to work well on simulated examples.\n\n**weakness:**\n1. the proposed scheme does not seem to very stable in the sense that errors in estimating the score may amplify as the particles are propagated. In some sense, the vanilla approach consisting in simulating forward the associated SDE appears to be much more stable in that respect. Is this true? Can the authors discuss and possibly illustrate this empirically.\n\n2. I think that the simulation should take compute-time into account. The vanilla Monte-Carlo approach (ie. simulate forward the SDE) is extremely scalable. The proposed approach requires implementing score-matching at each step.\n\n3. Evaluating $\\rho_t$ from samples is basically a density estimation problem. I think the authors should consequently consider comparing their method to Vanilla-Monte-Carlo + density estimation (eg. normalising flow, or something of that sort). Indeed, there are also quite a few ways to estimate differential entropies, and mutual information, from samples!\n\n4. The network architectures are exploiting quite a lot the structure of the problem. One may argue that this structure could also be used to estimate the density out of samples much more efficiently than a completely problem-agnostic approach ... \n\n\n**minor question**:\n1. is it $-\\alpha$ in Equation (11) or $+\\alpha$", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "uIZBcOWxwh", "reviewer": "Reviewer_LTVr", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes score-based transport modeling (SBTM), a particle-based method to solve the Fokker-Planck equation. This is based on the transport map approach (TE) where the key challenge is to estimate the score of the distribution at the current time. Such a challenge is resolved by learning the score with neural networks using the current set of particles, using an objective derived from Stein's identity. Two synthetic experiments are done to demonstrate the effectiveness of the proposed method.\n", "summary_of_the_review": "Given the lack of motivation for the proposed method and the somewhat misleading theoretical development, I am leaning toward rejection of this paper.\n", "strength_and_weaknesses": "## Strengths:\n* The idea of learning the score from particles while simulating particles along the probability flow of the Fokker-Planck equation using the score is a very interesting one. \n* Theoretically there are some interesting perspectives on viewing SBTM in Lagrangian or Eulerian frames, and the KL divergence with respect to the true flow can be controlled by the error of the score.\n\n\n## Weaknesses:\n* The motivation for learning the score during the flow is not super clear to me when you can simply simulate the SDE. The authors suggest you can obtain more quantities like the density of the flow, yet to me, the only new quantity you obtain is the score. And to obtain the score, you could also just simulate the SDE using Euler, and learn the score as a post-processing step using (7), i.e., I don't see any practical benefit of learning the flow while simulating the particles which can be very costly. The only way I can think of to obtain the density is to keep the *entire history* of the particles (which can be very memory intensive) and then compute the density using (5), but this is not demonstrated in any experiment. \n* I found Proposition 1 quite misleading and not relevant to the actual algorithm proposed. What is the role of $G_t$? Isn't it just the true score? The constrained optimization in (SBTM) formulation is never actually done (it is mentioned only in passing that it can be solved using the adjoint method). The function $\\lambda$ appears to not matter as it is simply taken to be a dirac delta in the sequential SBTM (wouldn't this also break the assumptions in Proposition 1? Dirac deltas are not positive functions). \n* Equation (seqSBTM) seems identical to me as (7). It is possible that I entirely misunderstood, but to me Algorithm 1 has very little to do with the theories developed in Section 3 and can be simply obtained by applying (7) to get the score and then integrating the ODE (3)+(4).\n* I found the experiments section quite weak. Despite the main selling point of the proposed algorithm being its ability to access quantities like the density, it is never demonstrated in what occasions these quantities are useful. It is only verified that these quantities are accurate. What downstream applications can these quantities facilitate? The authors should also compare with the naive solution which is to simulate the corresponding SDE using Euler scheme and then estimate these quantities as postprocessing steps, such as using kernel density estimation. I would also hope to see more interesting instances of Fokker-Planck equations rather than the two toyish examples shown here.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TLGBABJKJv", "reviewer": "Reviewer_qgFR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper considers the transport-map approach to learn the score function along the solution trajectory of the Fokker-Planck equation. In contrast to the score-matching approach, the proposed proposed sequential SBTM method does not require to simulate the underlying SDE. Instead, the particles are propagated via a sequentially constructed estimation of the score function. Experiments are conducted to show the empirical advantage of the proposed method.", "summary_of_the_review": "Overall I think the self-consistency based approach is a promising direction for score estimation and this paper makes a good contribution. ", "strength_and_weaknesses": "Strength: While the score-based transport model is similar to a recent work (Shen et al. 2022), as acknowledged by the authors, the authors generalizes the previous art by considering a non-constant diffusion coefficient and show that the objective of SBTM controls the KL divergence between the hypothesis density and the ground-truth density, which improves over the previous Wasserstein-type bound.\n\nWeakness: \n1. Algorithm 1 can be memory consuming since it requires a individual neural network model for the score function at every time-stamp. Please correct me if I am wrong. \n2. I do not see how Proposition 3 differs from the standard score matching technique. This seems the exact score-matching technique for approximating $\\nabla \\log \\rho_t$.\n3. It is not clear how the optimization problem (SBTM) is solved and the analysis on how Algorithm 1 solves SBTM, e.g. the convergence guarantee, is missing.\n4. Since SBTM is proposed as an improvement over SBDM, it seems reasonable to include SBDM in the empirical study, but I do not seem the comparison between SBTM and SBDM.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject