Probability flow solution of the Fokker–Planck equation

Our contributions are both theoretical and computational:

• We provide a bound on the Kullback–Leibler (KL) divergence from the estimate ρ_t produced via an approximate velocity field v_t to the target $\rho_t^*$. 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 ρ_t .
• Based on this bound, we introduce two optimization problems that can be used to learn the velocity field (5) in the transport equation (4) so that its solution coincides with that of the FPE (1). 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 are not directly available from samples alone, like point-wise evaluation of ρ_t itself, the differential entropy, and the probability current.
• 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 [13, 56]. We show that our method can be used to accurately compute the entropy production rate, a quantity of interest in the active matter community [39], as it quantifies the out-of-equilibrium nature of the system's dynamics.

Throughout, we assume that the stochastic process (2) evolves over $\Omega = \mathbb{R}^d$, though our results can easily be extended to domains with either reflecting [30] or periodic boundary conditions. We let $|\cdot|:\mathbb{R}^d\rightarrow\mathbb{R}_{\geqslant 0}$ denote the Euclidean norm on vectors and $|\cdot|_F:\mathbb{R}^{d\times d}\rightarrow\mathbb{R}_{\geqslant 0}$ denote the Fröbenius norm on matrices. For our theory, we assume that the drift vector $b_t: \mathbb{R}^d \to \mathbb{R}^d$ and the diffusion tensor $D_t: \mathbb{R}^d \to \mathbb{R}^{d\times d}$ with $D_t(x) = \sigma_t(x)\sigma_t(x)^\mathsf{T}$ are both twice-differentiable in x for each t and satisfy, for some fixed C > 0, L > 0, and T > 0

$$\begin{equation} \begin{aligned} |b_t(x)| + |\sigma_t(x)|_F &\leqslant C(1 + |x|)\qquad \forall \: (x, t) \in \mathbb{R}^n\times [0, T],\\ |b(t, x) - b(t, y)| + |\sigma(t, x) - \sigma(t, y)|_F &\leqslant L|x - y| \qquad \forall \: (x, y, t)\in\mathbb{R}^n\times\mathbb{R}^n\times [0, T], \end{aligned} \end{equation} \tag{ 8 }$$

so that the solution to the SDE (2) is well-defined for $t \in [0, T]$ [40]. We further assume that the initial PDF ρ₀ is three-times differentiable, positive everywhere on Ω, and such that $H_0 = -\int_\Omega \log \rho_0(x) \rho_0(x) \mathrm{d}x \lt \infty$; $\rho^*_t$ then enjoys the same properties at all times $t \in [0, T]$. Finally, we assume that $\log \rho^*_t$ is K-smooth globally for $(t,x)\in [0,\infty)\times \Omega$, i.e.

$$\begin{equation} \exists K>0 \ : \quad \forall (t,x)\in [0,\infty)\times \Omega \quad |\nabla \log \rho^*_t(x) - \nabla \log \rho^*_t(y)| \leqslant K |x-y|. \end{equation} \tag{ 9 }$$

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 $\dot{y}_t = \frac{\mathrm{d}}{\mathrm{d}t}y_t$ interchangeably for a time-dependent quantity y_t .

Let $s_t:\Omega \to\mathbb{R}^d$ denote an approximation to the score of the target $\nabla\log\rho_t^*$, and consider the solution $\rho_t: \Omega \rightarrow \mathbb{R}_{\geqslant 0}$ to the transport equation

$$\begin{equation} \partial_t\rho_t(x) = - \nabla \cdot (v_t(x)\rho_t(x)) \qquad \textrm{with} \quad v_t(x) = b_t(x) - D_t(x) s_t(x), \end{equation} \tag{ 10 }$$

subject to the initial condition $\rho_{t = 0} = \rho_0$. Our goal is to develop a variational principle that may be used to adjust s_t so that ρ_t tracks $\rho_t^*$. 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 section 1.2 hold. Let ρ_t denote the solution to the transport equation (10), and let $\rho^*_t$ denote the solution to the FPE (1). Assume that $\rho_{t = 0}(x) = \rho^*_{t = 0}(x) = \rho_0(x)$ for all $x \in \Omega$. Then

$$\begin{equation} \frac{\mathrm{d}}{\mathrm{d}t} \mathsf{KL}(\rho_t\:\Vert\:\rho^*_t) \leqslant \frac{1}2 \int_\Omega\left|s_t(x) - \nabla \log \rho_t(x)\right|_{D_t(x)}^2 \rho_t(x)\mathrm{d}x, \end{equation} \tag{ 11 }$$

where $|\cdot|^2_{D_t(x)} = \langle \cdot, D_t(x) \cdot \rangle$.

In particular, (11) implies that for any $T\in[0,\infty)$ we have explicit control on the KL divergence

$$\begin{equation} \mathsf{KL}(\rho_T\:\Vert\:\rho^*_T) \leqslant \frac{1}2 \int_0^T \int_\Omega\left|s_t(x) - \nabla \log \rho_t(x)\right|_{D_t(x)}^2 \rho_t(x)\mathrm{d}x \mathrm{d}t. \end{equation} \tag{ 12 }$$

Remarkably, (12) only depends on the approximate ρ_t and does not include $\rho_t^*$: it states that the accuracy of ρ_t as an approximation of $\rho_t^*$ can be improved by enforcing agreement between s_t and $\nabla\log\rho_t$. This means that we can optimize (12) without making use of external data from $\rho_t^*$, which offers a self-consistent objective function to learn the score s_t using (10) alone.

The primary difficulty with this approach is that ρ_t must be considered as a functional of s_t , since the velocity v_t used in (10) depends on s_t . To render the resulting minimization of the right-hand side of (12) practical, we can exploit that (10) can be solved via the method of characteristics, as summarized in appendix A. Specifically, if $\dot{X}_t(x) = v_t(X_t(x))$ is the probability flow equation associated with the velocity v_t , then $\rho_t = X_t \sharp \rho_0$. This means that the expectation of any function $\phi(x)$ over $\rho_t(x)$ can be expressed as the expectation of $\phi_t(X_t(x))$ over $\rho_0(x)$. Observing that the score of the solution to (10) along trajectories of the probability flow $\nabla\log\rho_t(X_t(x))$ solves a closed equation leads to the following proposition.

Proposition 2 (Score-based transport modeling). Assume that the conditions listed in section 1.2 hold. Define $v_t(x) = b_t(x) - D_t(x) s_t(x)$ and consider

$$\begin{equation} \begin{aligned} &\dot X_{t}(x) = v_t(X_{t}(x)), && X_{0}(x) = x,\\ &\dot G_t(x) = - [\nabla v_t(X_{t}(x))]^\mathsf{T} G_t(x) - \nabla \nabla \cdot v_t(X_{t}(x)), \quad &&G_0(x) = \nabla \log \rho_0(x). \end{aligned} \end{equation} \tag{ 13 }$$

Then $\rho_t = X_t\sharp \rho_0$ solves (10), the equality $G_t(x) = \nabla \log \rho_t(X_t(x))$ holds, and for any $T\in[0,\infty)$

$$\begin{equation} \mathsf{KL}(X_T\sharp\rho_0\:\Vert\:\rho^*_T) \leqslant \frac{1}2 \int_0^T \int_\Omega\left|s_t(X_t(x)) - G_t(x)\right|_{D_t(X_t(x))}^2 \rho_0(x)\mathrm{d}x \mathrm{d}t. \end{equation} \tag{ 14 }$$

Moreover, if $s^*_t$ is a minimizer of the constrained optimization problem

$$\begin{equation} \min_{s} \int_0^T \int_\Omega \left|s_t(X_t(x)) - G_t(x)\right|_{D_t(X_t(x))}^2 \rho_0(x) \mathrm{d}x \mathrm{d}t \quad \textrm{subject to }(13) \end{equation} \tag{ 15 }$$

then $D_t(x) s^*_t(x) = D_t(x)\nabla \log \rho_t^*(x)$ where $\rho^*_t$ solves the FPE (1). The map $X^*_t$ associated to any minimizer is a transport map from ρ₀ to $\rho^*_t$, i.e.

$$\begin{equation} x \sim \rho_0 \qquad \textrm{implies that} \qquad X^*_{t}(x) \sim \rho^*_t, \qquad \forall t\in[0,T]. \end{equation} \tag{ 16 }$$

Proposition 2 is proven in appendix B.3. The result also holds with a standard Euclidean norm replacing the diffusion-weighted norm, in which case the minimizer is unique and is given by $s^*_t(x) = \nabla \log \rho_t^*(x)$. In the special case when the SDE is an Ornstein–Uhlenbeck (OU) process, the score and the equations for both X_t and G_t can be written explicitly; they are studied in appendix C.

In practice, the objective in (15) can be estimated empirically by generating samples from ρ₀ and solving the equations for $X_t(x)$ and $G_t(x)$ with $x\sim\rho_0$. The constrained minimization problem (15) 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-order 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 (1). Instead, we now develop a sequential timestepping procedure that avoids these difficulties entirely, and as a byproduct can scale to arbitrarily long time windows.

An alternative to the constrained minimization in proposition 2 is to consider an approach whereby the score s_t is obtained independently at each time to ensure that $\mathsf{KL}(\rho_t\:\Vert\:\rho_t^*)$ remains small. This suggests choosing s_t to minimize $\frac{\mathrm{d}}{\mathrm{d}t}\mathsf{KL}(\rho_t\:\Vert\:\rho_t^*)$, 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 score-matching that is equivalent to minimizing $\frac{\mathrm{d}}{\mathrm{d}t}\mathsf{KL}(\rho_t\:\Vert\:\rho_t^*)$ but obviates the calculation of G_t . Expanding the square in (11) and applying $\int_\Omega s_t(x)^\mathsf{T} \nabla \log \rho_t(x) \,\rho_t(x)\mathrm{d}x = - \int_\Omega \nabla\cdot s_t(x) \,\rho_t(x)\mathrm{d}x$, we may write

$$\begin{equation*} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t}\mathsf{KL}(\rho_t\:\Vert\:\rho^*_t) &\leqslant \frac{1}{2} \int_\Omega\big(|s_t(X_t(x))|_{D_t(X_t(x))}^2 + 2\nabla \cdot (D_t(X_t(x))s_t(X_t(x)))\big)\rho_0(x)\mathrm{d}x\\ & \quad + \frac{1}{2}\int |G_t(x)|^2 \rho_0(x)\mathrm{d}x. \end{aligned} \end{equation*}$$

Because $\nabla\log\rho_t(X_t(x)) = G_t(x)$ is independent of s_t , we may neglect the corresponding square term during optimization. This leads to a simple and comparatively less expensive way to build the pushforward $X^*_t$ such that $X^*_t\sharp \rho_0 = \rho_t^*$ sequentially in time, as stated in the following proposition.

Proposition 3 (Sequential SBTM). In the same setting as proposition 2, let $X_t(x)$ solve the first equation in (13) with $v_t(x) = b_t(x) - D_t(x)s_t(x)$. Let s_t be obtained via

$$\begin{equation} \min_{s_t} \int_\Omega \left(|s_t(X_t(x))|_{D_t(X_t(x))}^2 + 2 \nabla \cdot ( D_t(X_t(x))s_t(X_t(x))) \right) \rho_0(x) \mathrm{d}x. \end{equation} \tag{ 17 }$$

Then, each minimizer $s_t^*$ of (17) satisfies $D_t(x) s^*_t(x) = D_t(x) \nabla \log \rho^*_t(x)$ where $\rho^*_t$ is the solution to (1). Moreover, the map $X^*_t$ associated to $s_t^*$ is a transport map from ρ₀ to $\rho^*_t$.

Proposition 3 is proven in appendix B.4. Critically, (17) is no longer a constrained optimization problem. Given the current value of X_t at any time t, we can obtain s_t via direct minimization of the objective in (17). Given s_t , we may compute the right-hand side of (13) and propagate X_t (and possibly G_t ) forward in time. The resulting procedure, which alternates between self-consistent score estimation and sample propagation, is presented in algorithm 1 for the choice of a forward-Euler integration routine in time. The output of the method produces a feasible solution $\rho_t = X_t\sharp \rho_0$ for (15) because $\dot{X}_t$ satisfies the first constraint in (13) by construction. Moreover, because the method controls $\frac{\mathrm{d}}{\mathrm{d}t}\mathsf{KL}(\rho_t\:\Vert\:\rho_t^*)$ at each t, it also controls $\mathsf{KL}(\rho_t\:\Vert\:\rho_t^*)$ by integration; an a-posteriori bound can be obtained by calculating $G_t(x)$ according to the second equation in (13) and computing the loss in (15). A few remarks on algorithm 1 are now in order.

Higher-order integrators. Algorithm 1 is stated for choice of forward-Euler integration for simplicity. In practice, any off-the-shelf integrator can be used, such as an adaptive Runge–Kutta method, by temporal discretization of the dynamics

$$\begin{equation*} \begin{aligned} \dot{X}_t(x) & = v_t(X_t(x)) \\ s_t & = \underset{s}{\mathrm{argmin}} \int_\Omega \left(|s(X_t(x))|_{D_t(X_t(x))}^2 + 2 \nabla \cdot ( D_t(X_t(x))s(X_t(x))) \right) \rho_0(x) \mathrm{d}x \end{aligned} \end{equation*}$$

and spatial discretization of the expectation over a set of samples propagated according to the equation for $X_t(x)$. In practice, the minimization can be performed over a parametric class of functions such as neural networks via a few steps of gradient descent.

Divergence computation. To avoid computation of the divergence—which can be costly for neural networks with high input dimension—we can use the denoising score matching loss function introduced by [61], which we discuss in appendix B.6. Empirically, we find that use of either the denoising objective or explicit derivative regularization is necessary for stable training to avoid overfitting to the training data; the level of regularization (or the noise scale in the denoising objective) can be decreased as the size of the dataset increases.

Time-dependence. When optimizing over a parametric class of functions, the score can be taken to be explicitly time-dependent, or the time-dependence can originate only through the parameters. In either case, all required outputs can be computed on-the-fly to avoid saving the entire history of parameters, which could be memory-intensive for large neural networks. If a time-dependent architecture is used, the method is amenable to online learning by randomly re-drawing initial conditions and optimizing over the resulting trajectory. In the numerical experiments below, we consider time-independent models with time-dependent parameters, because we found them to be sufficient.

SBTM vs. Sequential SBTM. Given the simplicity of the optimization problem (17), one may wonder if (15) is useful in practice, or if it is simply a stepping stone to arrive at (17). The primary difference is that (15) offers global control on the discrepancy between s_t and $\nabla\log\rho_t$ over $t\in[0,T]$, in the sense that it directly minimizes the time-integrated error, while (17) controls a local truncation error that could lead to the accumulation of learning and time-discretization errors. In the numerical examples below, we took the timestep $\Delta t$ sufficiently small, and the number of samples n sufficiently large, that we did not observe any accumulation of error. Nevertheless, (15) may allow for more accurate approximation, because the loss is exactly minimized at zero. Moreover, the higher-order derivatives contained in $\nabla \nabla\cdot (D_ts_t)$ must remain well-behaved when using (15) because this term appears in the definition of $\dot{G}_t$, while (17) only contains $\nabla\cdot (D_ts_t)$.

An alternative to the sequential procedure outlined here would be to generate samples from the target $\rho_t^*$ via simulation of the associated SDE (2), and then to approximate the score $\nabla\log\rho_t^*$ via minimization of the loss

$$\begin{equation} \int_0^T \int_{\Omega} (|s_t(x) - \nabla\log\rho_t^*(x)|_{D_t(x)}^2 \rho_t^*(x)\mathrm{d}x\mathrm{d}t, \end{equation} \tag{ 18 }$$

similar to SBDM. ρ_t can be computed as in SBTM or sequential SBTM by simulation of the probability flow with the learned s_t . As we now show, neither $\mathsf{KL}(\rho_t\:\Vert\:\rho_t^*)$ nor $\mathsf{KL}(\rho_t^*\:\Vert\:\rho_t)$ are controlled when using this procedure.

Proposition 4 (Learning on external data). Let $\rho_t:\Omega\rightarrow\mathbb{R}_{\gt 0}$ solve (10), and let $\rho_t^*:\Omega\rightarrow\mathbb{R}_{\gt 0}$ solve (1). Then, the following equality holds

$$\begin{align} \mathsf{KL}(\rho_T^*\:\Vert\:\rho_T) & = \int_0^T\int_{\Omega}|s_t(x) - \nabla\log\rho_t^*(x)|_{D_t(x)}^2\rho_t^*(x)\mathrm{d}x\mathrm{d}t\nonumber\\ & \quad + \int_0^T \int_{\Omega} \left(\nabla\log\rho_t(x) - s_t(x)\right)^\mathsf{T} D_t(x)\left(s_t(x) - \nabla\log\rho_t^*(x)\right)\rho_t^*(x)\mathrm{d}x\mathrm{d}t. \end{align} \tag{ 19 }$$

Proposition 4 shows that minimizing the error between s_t and $\nabla\log\rho_t^*$ on samples of $\rho_t^*$ leaves a remainder term, because in general $\nabla\log\rho_t \neq s_t$. Young's inequality gives the simple upper bound

$$\begin{align} \mathsf{KL}(\rho_T^*\:\Vert\:\rho_T) & \leqslant \frac{3}{2}\int_0^T\int_{\Omega}|s_t(x) - \nabla\log\rho_t^*(x)|_{D_t(x)}^2\rho_t^*(x)\mathrm{d}x\mathrm{d}t\nonumber\\ &\quad + \frac{1}{2}\int_0^T \int_{\Omega} |s_t(x) - \nabla\log\rho_t(x)|_{D_t(x)}^2\rho_t^*(x)\mathrm{d}x\mathrm{d}t. \end{align} \tag{ 20 }$$

However, controlling the above quantity requires enforcing agreement between s_t and $\nabla\log\rho_t$ in addition to s_t and $\nabla\log\rho_t^*$, which is precisely the idea of SBTM. Empirically, we find in our numerical experiments that training on external data alone is significantly less stable than sequential SBTM. In particular, and importantly for the applications we consider, we could not stably estimate the trajectory of the entropy production rate using a score model learned from the SDE with the same number of samples as used for sequential SBTM.

Setup. Here we study a problem that admits a tractable analytical solution for direct comparison. We consider N two-dimensional particles ($\bar{d} = 2$) that repel according to a harmonic interaction but experience harmonic attraction towards a moving trap $\beta_t \in \mathbb{R}^2$. The motion of the physical particles is governed by the stochastic dynamics

$$\begin{equation} \mathrm{d}x^{(i)}_t = (\beta_t - x^{(i)}_t)\mathrm{d}t + \alpha \Big(x^{(i)}_t - \frac{1}{N}\sum_{j = 1}^N x^{(\,j)}_t\Big)\mathrm{d}t + \sqrt{2D}\,\mathrm{d}W_t^{(i)},\quad i = 1, \ldots, N \end{equation} \tag{ 21 }$$

where $\alpha \in (0,1)$ is a fixed coefficient that sets the magnitude of the repulsion and each $x^{(i)}_0 \sim \rho_0$. The dynamics (21) is an OU process in the extended variable $x \in \mathbb{R}^{\bar{d}N}$ with block components $x^{(i)}$. Assuming a Gaussian initial condition, the solution to the FPE associated with (21) is a Gaussian for all time and hence can be characterized entirely by its mean m_t and covariance C_t . These can be obtained analytically (appendices C and D), which facilitates a quantitative comparison to the learned model. The differential entropy S_t is given by

$$\begin{equation} H_t = \tfrac12 \bar d N\left(\log\left(2\pi\right) + 1\right) + \tfrac{1}{2}\log\det C_t. \end{equation} \tag{ 22 }$$

In the experiments, we take $\beta_t = a(\cos\pi\omega t, \sin\pi\omega t)^\mathsf{T}$ with a = 2, ω = 1, D = 0.25, α = 0.5, and N = 50, giving rise to a 100-dimensional FPE. The particles are initialized from an isotropic Gaussian with mean β₀ (the initial trap position) and variance $\sigma_0^2 = 0.25$.

Network architecture. We take $s_t(x) = -\nabla U_{\theta_t}(x)$, where the potential $U_{\theta_t}(\cdot)$ is given as a sum of one- and two-particle terms

$$\begin{equation} U_{\theta_t}\big(x^{(1)}, \ldots, x^{(N)}\big) = \sum_{i = 1}^N U_{\theta_t, 1}\big(x^{(i)}\big) + \frac{1}{N}\sum_{\substack{i,j = 1\\ i \neq j}}^N U_{\theta_t, 2}\big(x^{(i)}, x^{(\,j)}\big), \end{equation} \tag{ 23 }$$

which ensures permutation symmetry amongst the physical particles by direct summation over all pairs. 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, we show in appendix D.1 that the class of scores representable by (23) contains the analytical score for the harmonic problem considered in this section. To obtain the parameters $\theta_{t_k+\Delta t_k}$, we perform a warm start and initialize from $\theta_{t_k}$, which 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 $\texttt{swish}$ activation function [43]; 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 Σ over the samples and the entropy production rate $\frac{\mathrm{d}H_t}{\mathrm{d}t}$. Because an analytical solution is available for this system, we may also compute the target $\nabla \log \rho_t(x) =$ $-C_t^{-1}(x - m_t)$ and measure the goodness of fit via the relative Fisher divergence

$$\begin{equation} \frac{\int_\Omega |s_t(x) - \nabla\log\rho_t(x)|^2 \bar{\rho}(x) \mathrm{d}x}{\int_\Omega|\nabla\log\rho_t(x)|^2 \bar{\rho}(x)\mathrm{d}x}. \end{equation} \tag{ 24 }$$

In equation (24), $\bar{\rho}$ can be taken to be equal to the current empirical estimate of ρ_t (the training data), or estimated using samples from the SDE(the SDE data).

Results. The representation of the dynamics (21) 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 (6), the SDE (21), and the noise-free equation obtained by setting D = 0 in (21) are shown in figure 1(A).

Zoom In Zoom Out Reset image size

Results for this quantitative comparison are shown in figure 1(B). The learned model accurately predicts the entropy production rate of the system and minimizes the relative metric (24) to the order of 10⁻². 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 without an additional learning component; we study this possibility in the next example. In addition, the learned model accurately predicts the high-dimensional covariance Σ of the system (curves lie directly on 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 decaying to zero.

Setup. Here, we consider a system of N = 5 physical particles in an anharmonic trap in dimension $\bar{d} = 2$ that exhibit soft-sphere repulsion. This system gives rise to a 10-dimensional (1), which is significantly too high for standard PDE solvers. The stochastic dynamics is given by

$$\begin{equation} \begin{aligned} \mathrm{d}x^{(i)}_t & = 4B \big(\beta_t - x_t^{(i)} \big)|x_t^{(i)} - \beta_t|^2\mathrm{d}t \\ & \quad + \frac{A}{N r^2}\sum_{j = 1}^N \big(x^{(i)}_t - x^{(\,j)}_t\big)\exp\left(-\frac{|x^{(i)}_t - x^{(\,j)}_t|^2}{2r^2}\right)\mathrm{d}t + \sqrt{2D}\,\mathrm{d}W_t,\:\:\: i = 1, \ldots, N, \end{aligned} \end{equation} \tag{ 25 }$$

where β_t again represents a moving trap, A > 0 sets the strength of the repulsion between the spheres, r sets their size, B > 0 sets the strength of the trap, and each $x^{(i)}_0 \sim \rho_0$. We set $\beta(t) = a(\cos\pi \omega t, \sin \pi \omega t)^\mathsf{T}$ or $\beta(t) = a(\cos\pi \omega t, 0)^\mathsf{T}$ with $a = 2, \omega = 1, D = 0.25, A = 10$, and r = 0.5. We fix $B = D/R^2$ with $R = \sqrt{\gamma N}r$ and γ = 5.0. This ensures that the trap scales with the number of particles and that they have sufficient room in the trap to generate a complex dynamics. The circular case converges to a distribution $\rho_t^* = \rho^* \circ Q_t$ that can be described as a fixed distribution $\rho^*$ composed with a time-dependent rotation Q_t , and hence the entropy production rate converges to zero by change of variables. The linear case does not exhibit this kind of 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 section 4.1.

Results. Similar to section 4.1, an example trajectory from the learned system, the SDE (25), and the noise-free system obtained by setting D = 0 are shown in figure 2(A) in the circular case. The learned particle trajectories exhibit an intuitive circular motion when compared to the SDE trajectory. When compared to the noise-free system, the learned trajectories exhibit a greater amount of spread, which enables the deterministic dynamics to accurately capture the statistics of the stochastic dynamics. Numerical estimates of a single component of the covariance and of the entropy production rate are shown in figures 2(B)/(C), with all moments shown in appendix D.2. The learned and SDE systems accurately capture the covariance, while the noise-free system underestimates the covariance in both the linear and the circular case. The prediction of the entropy production rate via algorithm 1 is reasonable in both cases, exhibiting the expected convergence to and oscillation around zero in the circular and linear cases, respectively. In the inset, we show the prediction of the entropy production rate when learning on samples from the SDE; the prediction is initially offset, and later becomes divergent. We found that this behavior was generic when training on the SDE, but never observed it when training on self-consistent samples.

Zoom In Zoom Out Reset image size

Setup. We now consider a model from the physics of active matter, which describes the motion of a single motile swimmer in an anharmonic trap. The swimmer can be thought of as a run-and-tumble bacterium [59]; it travels in a fixed direction for a fluctuating duration before picking a new direction at random in which to swim. The system is two-dimensional, and is given by the SDE for the position x and velocity v

$$\begin{equation} \begin{aligned} \mathrm{d}x & = \left(-x^3 + v\right)\mathrm{d}t,\\ \mathrm{d}v & = -\gamma v \mathrm{d}t + \sqrt{2\gamma D}\mathrm{d}W_t. \end{aligned} \end{equation} \tag{ 26 }$$

While low-dimensional, (26) exhibits convergence to a non-equilibrium statistical steady state in which the probability current $j_t(x) = v_t(x)\rho_t(x)$ is non-zero. Here, we show that sequential SBTM is capable of accurately capturing such currents, which is necessary to resolve the dynamics of the FPE: if our goal were solely to sample at equilibrium, it would be sufficient to freeze the samples after an initial transient. Moreover, we show that the method preserves the stationary distribution over long times relative to the persistence time $1/\gamma$ of the swimmer, and does not display appreciable accumulation of error.

We set γ = 0.1 and D = 1.0. Because noise only enters the system through the velocity variable v in (26), the score can be taken to be one-dimensional, which is equivalent to learning the score only in the range of the rank-deficient diffusion matrix. Further details on the architecture can be found in appendix D.3.

Results. A phase portrait for the learned probability flow dynamics is shown in figure 3, computed by rolling out an additional set of 50 trajectories for time $5/\gamma$ with a fixed set of parameters (after learning for time $10/\gamma$). The phase portrait depicts closed limit cycles between and centered within the modes reminiscent of the classical phase portrait for the pendulum. Here, the closed limit cycles correspond to non-equilibrium currents that preserve the steady-state density.

Zoom In Zoom Out Reset image size

A kernel density estimate for the distribution of samples produced by the learned system, the stochastic system, and the noise-free systems are shown in figure 4, which demonstrate that the distribution of the learned samples qualitatively matches the distribution of the SDE samples. Comparatively, the noise-free system grows overly concentrated with time, ultimately converging to a singular dirac measure at the origin. A movie of the motion of the samples $(x_i(t), v_i(t))_{t\geqslant 0}$ over a duration $10/\gamma$ in phase space can be seen at this link. The movie highlights convergence of the learned solution to one with a non-zero steady-state probability current that qualitatively matches that of the SDE, but which enjoys more interpretable sample trajectories.

Zoom In Zoom Out Reset image size

We begin with a lemma.

Lemma 5. Let $\rho_t: \Omega \rightarrow \mathbb{R}_{\geqslant 0}$ satisfy the transport equation

$$\begin{equation} \partial_t \rho_t(x) = - \nabla \cdot \left(v_t(x) \rho_t(x)\right). \end{equation} \tag{ A.1 }$$

Assume that $v_t(x)$ is C² in both t and x for $t\geqslant0$ and globally Lipschitz in x. Then, given any $t,t^{^{\prime}} \geqslant 0$, the solution of (A.1) satisfies

$$\begin{equation} \rho_t(x) = \rho_{t^{^{\prime}}}(X_{t,t^{^{\prime}}}(x)) \exp\left( -\int_{t^{^{\prime}}}^t \nabla \cdot v_\tau(X_{t,\tau}(x)) \mathrm{d}\tau\right) \end{equation} \tag{ A.2 }$$

where $X_{\tau,t}$ is the probability flow solution to (6). In addition, given any test function $\phi: \Omega \to \mathbb{R}$, we have

$$\begin{equation} \int_\Omega \phi(x) \rho_t(x) \mathrm{d}x = \int_\Omega \phi(X_{t^{^{\prime}},t}(x)) \rho_{t^{^{\prime}}}(x) \mathrm{d}x. \end{equation} \tag{ A.3 }$$

In words, lemma 5 states that an evaluation of the PDF ρ_t at a given point x may be obtained by evolving the probability flow equation (6) backwards to some earlier time t^ʹ to find the point x^ʹ that evolves to x at time t, assuming that $\rho_{t^{^{\prime}}}(x^{^{\prime}})$ is available. In particular, for $t^{^{\prime}} = 0$, we obtain

$$\begin{equation} \rho_t(x) = \rho_0(X_{t,0}(x)) \exp\left( -\int_{0}^t \nabla \cdot v_\tau(X_{t,\tau}(x)) \mathrm{d}\tau\right), \end{equation} \tag{ A.4 }$$

and

$$\begin{equation} \int_\Omega \phi(x) \rho_t(x) \mathrm{d}x = \int_\Omega \phi(X_{0,t}(x)) \rho_{0}(x) \mathrm{d}x. \end{equation} \tag{ A.5 }$$

Since the probability current is by definition $v_t(x)\rho_t(x)$, using (A.4) to express $\rho_t(x)$ also gives the following equation for the current:

$$\begin{equation} v_t(x) \rho_t(x) = v_t(x) \rho_0(X_{t,0}(x)) \exp\left( -\int_0^t \nabla \cdot v_\tau(X_{\tau,t}(x)) \mathrm{d}\tau\right). \end{equation} \tag{ A.6 }$$

Proof. The assumed C² and globally Lipschitz conditions on v_t guarantee global existence (on $t\geqslant0$) and uniqueness of the solution to (6). Differentiating $\rho_t(X_{t^{^{\prime}},t}(x))$ with respect to t and using (6) and (A.1) we deduce

$$\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} \rho_t(X_{t^{^{\prime}},t}(x)) & = \partial_t \rho_t(X_{t^{^{\prime}},t}(x)) + \frac{\mathrm{d}}{\mathrm{d}t}X_{t^{^{\prime}},t}(x) \cdot \nabla \rho_t(X_{t^{^{\prime}},t}(x)) \\ & = \partial_t \rho_t(X_{t^{^{\prime}},t}(x)) + v_t(X_{t^{^{\prime}},t}(x)) \cdot \nabla \rho_t(X_{t^{^{\prime}},t}(x))\\ & = - \nabla \cdot v_t(X_{t^{^{\prime}},t}(x)) \, \rho_t(X_{t^{^{\prime}},t}(x)). \end{aligned} \end{equation} \tag{ A.7 }$$

Integrating this equation in t from $t = t^{^{\prime}}$ to t = t gives

$$\begin{equation} \begin{aligned} \rho_t(X_{t^{^{\prime}},t}(x)) & = \rho_{t^{^{\prime}}}(x) \exp\left( - \int_{t^{^{\prime}}}^t \nabla \cdot v_\tau(X_{t^{^{\prime}},\tau}(x)) \mathrm{d}\tau\right). \end{aligned} \end{equation} \tag{ A.8 }$$

Evaluating this expression at $x = X_{t,t^{^{\prime}}}(x)$ and using the group properties (i) $X_{t^{^{\prime}},t}(X_{t,t^{^{\prime}}}(x)) = x$ and (ii) $X_{t^{^{\prime}},\tau}(X_{t,t^{^{\prime}}}(x)) = X_{t,\tau}(x)$ gives (A.2). Equation (A.3) can be derived by using (A.2) to express $\rho_t(x)$ in the integral at the left hand-side, changing integration variable $x\to X_{t^{^{\prime}},t}(x)$ and noting that the factor $\exp( -\int_{t^{^{\prime}}}^t \nabla \cdot v_\tau(X_{t^{^{\prime}},\tau}(x)))$ is precisely the Jacobian of this change of variable. To see this, note that by definition the flow map satisfies

$$\begin{equation*} \frac{\mathrm{d}}{\mathrm{d}t}X_{t^{^{\prime}}, t}(x) = v_t(X_{t^{^{\prime}}, t}(x)), \qquad X_{\tau, \tau}(x) = x \:\: \forall \tau \geqslant 0. \end{equation*}$$

Hence,

$$\begin{equation*} \frac{\mathrm{d}}{\mathrm{d}t}\nabla X_{t^{^{\prime}}, t}(x) = \nabla v_t(X_{t^{^{\prime}}, t}(x)) \nabla X_{t^{^{\prime}}, t}(x), \qquad \nabla X_{\tau, \tau}(x) = I \:\: \forall \tau \geqslant 0. \end{equation*}$$

By Jacobi's formula for determinants,

$$\begin{equation*} \frac{\mathrm{d}}{\mathrm{d}t}\det\left(\nabla X_{t^{^{\prime}}, t}(x)\right) = \mathrm{Tr}\left(\nabla v_t(X_{t^{^{\prime}}, t}(x))\right) \det\left(\nabla X_{t^{^{\prime}}, t}(x)\right), \qquad \det\left(\nabla X_{\tau, \tau}(x)\right) = 1 \:\: \forall \tau \geqslant 0. \end{equation*}$$

Integrating this equation with respect to t and using that $\mathrm{Tr}(\nabla v_t(\cdot)) = \nabla\cdot v_t(\cdot)$, we obtain

$$\begin{equation*} \det\left(\nabla X_{t^{^{\prime}}, t}(x)\right) = \exp\left(\int_{t^{^{\prime}}}^t \nabla \cdot v_\tau(X_{t^{^{\prime}}, \tau}(x))\mathrm{d}\tau \right). \end{equation*}$$

The result is the integral at the right hand-side of (A.3).

Lemma 5 also holds locally in time for any $v_t(x)$ that is C² in both t and x. In particular, it holds locally if we set $s_t(x) = \nabla \log \rho_t(x)$ and if we assume that $\rho_0(x)$ is (i) positive everywhere on Ω and (ii) C³ in x. In this case, (A.1) is the FPEs (1) and (A.2) holds for the solution to that equation.

We now consider computation of the differential entropy, and state a similar result.

Lemma 6. Assume that $\rho_0: \Omega \rightarrow \mathbb{R}_{\geqslant 0}$ is positive everywhere on Ω and C³ in its argument. Let $\rho_t: \Omega \rightarrow \mathbb{R}_{\geqslant 0}$ denote the solution to the FPE (1) (or equivalently, to the transport equation (A.1) with $s_t(x) = \nabla \log \rho_t(x)$ in the definition of $v_t(x)$). Then the differential entropy $H_t = -\int_\Omega \log \rho_t(x) \, \rho_t(x) \mathrm{d}x$ can expressed as

$$\begin{equation} H_t = -\int_\Omega \log \rho_t(X_{0,t}(x))\, \rho_0(x) \mathrm{d}x = H_0 + \int_0^t \int_\Omega \nabla \cdot v_\tau(X_{0,\tau}(x)) \rho_0(x) \mathrm{d}x \mathrm{d}\tau \end{equation} \tag{ A.9 }$$

or

$$\begin{equation} H_t = H_0 - \int_0^t \int_\Omega s_\tau(X_{0,\tau}(x)) \cdot v_\tau(X_{0,\tau}(x)) \rho_0(x) \mathrm{d}x \mathrm{d}\tau. \end{equation} \tag{ A.10 }$$

Proof. We first derive (A.9). Observe that applying (A.5) with $\phi = \log \rho_t$ leads to the first equality. The second can then be deduced from (A.4). To derive (A.10), notice that from (A.1),

$$\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} H_t & = \int_\Omega \log \rho_t(x) \nabla \cdot \left(v_t(x) \rho_t(x)\right) \mathrm{d}x,\\ & = -\int_\Omega \nabla \log \rho_t(x) \cdot v_t(x) \rho_t(x) \mathrm{d}x,\\ & = -\int_\Omega s_t(x) \cdot v_t(x) \rho_t(x) \mathrm{d}x. \end{aligned} \end{equation} \tag{ A.11 }$$

Above, we used integration by parts to obtain the second equality and $s_t = \nabla \log \rho_t$ to get the third. Now, using (A.5) with $\phi = s_t \cdot v_t$ integrating the result gives (A.10).

If the score $s_t \approx \nabla \log \rho_t$ is known to sufficient accuracy, ρ_t can be resampled at any time t using the dynamics

$$\begin{equation} \mathrm{d}X_\tau = s_t(X_\tau) \mathrm{d}\tau + d W_\tau. \end{equation} \tag{ A.12 }$$

In (A.12), τ is an artificial time used for sampling that is distinct from the physical time in (2). For $s_t = \nabla \log \rho_t$, the equilibrium distribution of (A.12) is exactly ρ_t . In practice, s_t will be imperfect and will have an error that increases away from the samples used to learn it; as a result, (A.12) should be used near samples for a fixed amount of time to avoid the introduction of additional errors.

Let us restate proposition 1 for convenience:

Proposition 1 (Control of the KL divergence). Assume that the conditions listed in section 1.2 hold. Let ρ_t denote the solution to the transport equation (10), and let $\rho^*_t$ denote the solution to the FPE (1). Assume that $\rho_{t = 0}(x) = \rho^*_{t = 0}(x) = \rho_0(x)$ for all $x \in \Omega$. Then

$$\begin{equation} \frac{\mathrm{d}}{\mathrm{d}t} \mathsf{KL}(\rho_t\:\Vert\:\rho^*_t) \leqslant \frac{1}2 \int_\Omega\left|s_t(x) - \nabla \log \rho_t(x)\right|_{D_t(x)}^2 \rho_t(x)\mathrm{d}x, \end{equation} \tag{ 11 }$$

where $|\cdot|^2_{D_t(x)} = \langle \cdot, D_t(x) \cdot \rangle$.

Proof. By assumption, ρ_t solves (10) and $\rho_t^*$ solves (1). Denote by $v_t(x) = b_t(x) - D_t(x) s_t(x)$ and $v^*_t(x) = b_t(x) - D_t(x) s^*_t(x)$ with $s^*_t(x) = \nabla \log \rho_t^*(x)$. Then, we have

$$\begin{equation*} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} \mathsf{KL}(\rho_t\:\Vert\:\rho^*_t) & = \frac{\mathrm{d}}{\mathrm{d}t} \int_\Omega \log\left(\frac{\rho_t(x)}{\rho_t^*(x)}\right) \rho_t(x) \mathrm{d}x,\\ & = -\int_\Omega \frac{\rho_t(x)}{\rho_t^*(x)} \partial_t \rho^*_t(x) \mathrm{d}x + \int_\Omega \log \left(\frac{\rho_t(x)}{\rho_t^*(x)}\right) \partial_t \rho_t(x)\mathrm{d}x,\\ & = -\int_\Omega v_t^*(x) \cdot \nabla \left(\frac{\rho_t(x)}{\rho_t^*(x)}\right) \rho^*_t(x) \mathrm{d}x + \int_\Omega v_t(x) \cdot \nabla \log \left(\frac{\rho_t(x)}{\rho_t^*(x)}\right) \rho_t(x) \mathrm{d}x,\\ & = -\int_\Omega \left(v_t^*(x) - v_t(x)\right) \cdot \left(\nabla \log \rho_t(x)- \nabla \log \rho^*_t(x)\right) \rho_t(x) \mathrm{d}x,\\ & = \int_\Omega \left(s_t^*(x) - s_t(x)\right) \cdot D_t(x) \left(\nabla \log \rho_t(x)- s^*_t(x)\right) \rho_t(x) \mathrm{d}x. \end{aligned} \end{equation*}$$

Above, we used integration by parts to obtain the third equality. Now, dropping function arguments for simplicity of notation, we have that

$$\begin{equation*} \begin{aligned} |\nabla\log\rho_t - s_t|^2_{D_t} & = |\nabla\log\rho_t -s^*_t + s^*_t - s_t|^2_{D_t},\\ & = |\nabla\log\rho_t - s^*_t|^2_{D_t} + |s^*_t - s_t|^2_{D_t} + 2 (\nabla\log\rho_t -s^*_t)\cdot D_t(s^*_t - s_t),\\ & \geqslant 2 (\nabla\log\rho_t - s^*_t)\cdot D_t (s^*_t - s_t). \end{aligned} \end{equation*}$$

Hence, we deduce that

$$\begin{equation} \begin{aligned} &\frac{\mathrm{d}}{\mathrm{d}t} \mathsf{KL}(\rho_t\:\Vert\:\rho^*_t) \leqslant \frac12 \int_\Omega |s_t(x) - \nabla\log\rho_t(x)|^2_{D_t(x)} \rho_t(x) \mathrm{d}x. \end{aligned} \end{equation} \tag{ B.1 }$$

The Eulerian equivalent of proposition 2 can be stated as:

Proposition 7 (SBTM in the Eulerian frame). Assume that the conditions listed in section 1.2 hold. Fix $T\in(0,\infty]$ and consider the optimization problem

$$\begin{equation} \begin{aligned} &\min_{\{s_t:t\in [0,T)\}} \int_0^T \int_\Omega \left|s_t(x) - \nabla \log \rho_t(x)\right|_{D_t(x)}^2 \rho_t(x) \mathrm{d}x \mathrm{d}t\\[6pt] & \textrm{subject to:} \quad \partial_t \rho_t(x) = - \nabla \cdot \left( v_t(x) \rho_t (x) \right), \:\:\: x \in \Omega \end{aligned} \end{equation} \tag{ B.2 }$$

with $v_t(x) = b_t(x) - D_t(x) s_t(x)$. Then every minimizer of (B.2) satisfies $D_t(x)s_t^*(x) = D_t(x)\nabla \log \rho_t^*(x)$ where $\rho^*_t: \Omega \to\mathbb{R}_{\gt 0}$ solves (1).

In words, this proposition states that solving the constrained optimization problem (B.2) is equivalent to solving the FPE (1).

Proof. The constrained minimization problem (B.2) can be handled by considering the extended objective

$$\begin{equation} \begin{aligned} \int_0^T \int_\Omega \left(\left|s_t(x) - \nabla \log \rho_t(x)\right|_{D_t(x)}^2 \rho_t(x) + \mu_t(x) \left(\partial_t \rho_t(x) + \nabla \cdot \left(v_t(x) \rho_t(x)\right)\right) \right) \mathrm{d}x \mathrm{d}t \end{aligned} \end{equation} \tag{ B.3 }$$

where $v_t(x) = b_t(x) - D_t(x) s_t(x)$ and $\mu_t:\mathbb{R}^d\rightarrow\mathbb{R}_{\geqslant 0}$ is a Lagrange multiplier. The Euler–Lagrange equations associated with (B.3) read

$$\begin{align} \partial_t \rho_t(x) & = - \nabla \cdot \left(v_t(x)\rho_t(x)\right)\nonumber\\ \partial_t \mu_t(x) & = -v_t(x)\cdot \nabla \mu_t(x) + |s_t(x)|_{D_t(x)}^2 - |\nabla \log \rho_t|_{D_t(x)}^2 \nonumber\\ &\qquad+ 2 \nabla\cdot\left[D_t(x)\left(s_t(x) - \nabla \log \rho_t(x)\right)\right] , \\ 0 & = \mu_T(x),\nonumber\\ 0 & = 2D_t(x)\left(s_t(x) - \nabla\log\rho_t(x)\right)\rho_t(x) - D_t(x) \nabla\mu_t(x)\rho_t(x).\nonumber \end{align} \tag{ B.4 }$$

Clearly, these equations will be satisfied if $s^*_t(x) = \nabla \log\rho^*_t(x)$ for all $x \in \Omega$, $\mu^*_t(x) = 0$ for all x, and $\rho_t^*$ solves (1). This solution is also a global minimizer, because it zeroes the value of the objective. Moreover, all global minimizers must satisfy $D_t(x)s^*_t(x) = D_t(x)\nabla \log\rho^*_t(x)$ ($\rho_t-$almost everywhere), as this is the only way to zero the objective.

It is also easy to see that there are no other local minimizers. To check this, we can use the fourth equation to write

$$\begin{equation*} D_t(x)(s_t(x) - \nabla\log\rho_t(x)) = \tfrac{1}{2}D_t(x)\nabla\mu_t(x). \end{equation*}$$

Then,

$$\begin{align*} |s_t(x)|_{D_t(x)}^2 - |\nabla\log\rho_t(x)|_{D_t(x)}^2 = \tfrac{1}{2}\left(s_t(x) + \nabla\log\rho_t(x)\right)^\mathsf{T} D_t(x) \nabla\mu_t(x). \end{align*}$$

This reduces the first three equations to

$$\begin{equation} \begin{aligned} \partial_t \rho_t(x) & = - \nabla \cdot \left(b_t(x)\rho_t(x) - D_t(x) \nabla \rho_t(x) - \tfrac{1}{2}\rho_t D_t(x) \nabla \mu_t(x)\right) \\ \partial_t \mu_t & = \left(b_t(x) - D_t(x)\nabla\log\rho_t(x) - \tfrac{1}{2}D_t(x)\nabla\mu_t(x)\right)^\mathsf{T}\nabla\mu_t(x)\\ &\qquad + \nabla\cdot\left(D_t(x)\nabla\mu_t(x)\right) + \tfrac{1}{2}\left(s_t(x) + \nabla\log\rho_t(x)\right)^\mathsf{T} D_t(x)\nabla\mu_t(x).\\ \mu_T(x) & = 0. \end{aligned} \end{equation} \tag{ B.5 }$$

Since the equation for µ_t is homogeneous in µ_t and $\mu_T = 0$, we must have $\mu_t = 0$ for all $t\in[0,T)$, and the equation for ρ_t reduces to (1).

As stated, proposition 7 is not practical, because it is phrased in terms of the density ρ_t . The following result demonstrates that the transport map identity (7) can be used to re-express proposition 7 entirely in terms of known quantities.

Proposition 2 (Score-based transport modeling). Assume that the conditions listed in section 1.2 hold. Define $v_t(x) = b_t(x) - D_t(x) s_t(x)$ and consider

$$\begin{equation} \begin{aligned} &\dot X_{t}(x) = v_t(X_{t}(x)), && X_{0}(x) = x,\\ &\dot G_t(x) = - [\nabla v_t(X_{t}(x))]^\mathsf{T} G_t(x) - \nabla \nabla \cdot v_t(X_{t}(x)), \quad &&G_0(x) = \nabla \log \rho_0(x). \end{aligned} \end{equation} \tag{ 13 }$$

Then $\rho_t = X_t\sharp \rho_0$ solves (10), the equality $G_t(x) = \nabla \log \rho_t(X_t(x))$ holds, and for any $T\in[0,\infty)$

$$\begin{equation} \mathsf{KL}(X_T\sharp\rho_0\:\Vert\:\rho^*_T) \leqslant \frac{1}2 \int_0^T \int_\Omega\left|s_t(X_t(x)) - G_t(x)\right|_{D_t(X_t(x))}^2 \rho_0(x)\mathrm{d}x \mathrm{d}t. \end{equation} \tag{ 14 }$$

Moreover, if $s^*_t$ is a minimizer of the constrained optimization problem

$$\begin{equation} \min_{s} \int_0^T \int_\Omega \left|s_t(X_t(x)) - G_t(x)\right|_{D_t(X_t(x))}^2 \rho_0(x) \mathrm{d}x \mathrm{d}t \quad \textrm{subject to }(13) \end{equation} \tag{ 15 }$$

then $D_t(x) s^*_t(x) = D_t(x)\nabla \log \rho_t^*(x)$ where $\rho^*_t$ solves the FPE (1). The map $X^*_t$ associated to any minimizer is a transport map from ρ₀ to $\rho^*_t$, i.e.

$$\begin{equation} x \sim \rho_0 \qquad \textrm{implies that} \qquad X^*_{t}(x) \sim \rho^*_t, \qquad \forall t\in[0,T]. \end{equation} \tag{ 16 }$$

Proof. Let us first show that $G_t(x) = \nabla \log \rho_t(X_t(x))$ satisfies (13) if $\rho_t = X_t\sharp \rho_0,$ i.e. if ρ_t satisfies the transport equation (10). Since (10) implies that

$$\begin{equation} \partial_t \log \rho_t(x) + v_t(x) \cdot \nabla \log \rho_t(x) = - \nabla \cdot v_t(x), \end{equation} \tag{ B.6 }$$

taking the gradient gives

$$\begin{equation} \partial_t \nabla \log \rho_t(x) + [\nabla v_t(x)]^\mathsf{T} \nabla \log \rho_t(x) + \nabla \nabla \log \rho_t(x) \cdot v_t(x) = - \nabla \nabla \cdot v_t(x) . \end{equation} \tag{ B.7 }$$

Therefore $G_t(x) = \nabla \log \rho_t(X_t(x))$ solves

$$\begin{equation} \begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} G_t(x) & = \partial_t \nabla \log \rho_t(X_t(x)) + \nabla \nabla \log \rho_t(X_t(x))\cdot \frac{\mathrm{d}}{\mathrm{d}t} X_t(x), \\ & = \partial_t \nabla \log \rho_t(X_t(x)) + \nabla \nabla \log \rho_t(X_t(x))\cdot v_t(x),\\ & = - \nabla \nabla \cdot v_t(X_{t}(x)) - [\nabla v_t(X_{t}(x))]^\mathsf{T} \nabla \log \rho_t(X_t(x)) , \end{aligned} \end{equation} \tag{ B.8 }$$

which recovers the equation for $G_t(x)$ in (13). Hence, the objective in (15) can also be written as

$$\begin{equation} \begin{aligned} &\int_0^T \int_\Omega \left|s_t(X_t(x)) - \nabla \log \rho_t(X_t(x))\right|_{D_t(X_t(x))}^2 \rho_0(x) \mathrm{d}x \mathrm{d}t\\ &\quad= \int_0^T \int_\Omega \left|s_t(x) - \nabla \log \rho_t(x)\right|_{D_t(x)}^2 \rho_t(x) \mathrm{d}x \mathrm{d}t \end{aligned} \end{equation} \tag{ B.9 }$$

where the second equality follows from (A.5) if $\rho_t(x)$ satisfies (A.1). Hence, (15) is equivalent to (B.2). The bound on $\mathsf{KL}(X_T\sharp\rho_0\:\Vert\:\rho_T^*)$ follows from (12).

Adjoint equations. In terms of a practical implementation, the objective in (B.2) can be evaluated by generating samples $\{x_i\}_{i = 1}^n$ from ρ₀ and solving the equations for X_t and G_t using the initial conditions $X_0(x_i) = x_i$ and $G_0(x_i) = \nabla \log \rho_0(x_i)$. Note that evaluating this second initial condition only requires one to know ρ₀ up to a normalization factor. To evaluate the gradient of the objective, we can introduce equations adjoint to those for X_t and G_t . To do so, we can consider the extended objective

$$\begin{equation} \begin{aligned} L[s_t] = &\int_0^T \int_{\Omega}|s_t(X_t(x)) - G_t(x)|_{D_t(X_t(x))}^2\rho_0(x)\mathrm{d}x\mathrm{d}t \\ & + \int_0^T \int_{\Omega}\theta_t(x)\left(\dot{X}_t(x) - v_t(X_t(x))\right)\rho_0(x)\mathrm{d}x\mathrm{d}t\\ & + \int_0^T \int_{\Omega}\eta_t(x)\left(\dot{G}_t(x) + \nabla v_t(X_t(x))^\mathsf{T} G_t(X_t(x)) + \nabla\nabla\cdot v_t(X_t(x))\right)\rho_0(x)\mathrm{d}x\mathrm{d}t. \end{aligned} \end{equation} \tag{ B.10 }$$

Taking the first variation with respect to $G_t(x)$ and $X_t(x)$, respectively, gives the equations

$$\begin{equation} \begin{aligned} \partial_t \eta_t(x) & = \nabla v_t(X_t(x))\eta_t(x) + 2 D_t(X_t(x))\left(G_t(x) - s_t(X_t(x))\right)\\ \eta_T(x) & = 0\\ \partial_t \theta_t(x) + \nabla v_t(X_t(x))^\mathsf{T} \theta_t(x) & = 2 \nabla s_t(X_t(x))^\mathsf{T} D_t(X_t(x))\left(s_t(X_t(x)) - G_t(x)\right)\\ &\quad + \left(s_t(X_t(x)) - G_t(x)\right)\cdot \nabla D_t(X_t(x)) \left(s_t(X_t(x)) - G_t(x)\right)\\ &\quad + \eta_t(x)\cdot\left(\nabla\nabla v_t(X_t(x))^\mathsf{T} G_t(x)\right) + \eta_t\cdot \nabla\nabla\nabla\cdot v_t(X_t(x))\\ \theta_T(x) & = 0. \end{aligned} \end{equation} \tag{ B.11 }$$

Let us restate proposition 3 for convenience:

Proposition 3 (Sequential SBTM). In the same setting as proposition 2, let $X_t(x)$ solve the first equation in (13) with $v_t(x) = b_t(x) - D_t(x)s_t(x)$. Let s_t be obtained via

$$\begin{equation} \min_{s_t} \int_\Omega \left(|s_t(X_t(x))|_{D_t(X_t(x))}^2 + 2 \nabla \cdot ( D_t(X_t(x))s_t(X_t(x))) \right) \rho_0(x) \mathrm{d}x. \end{equation} \tag{ 17 }$$

Then, each minimizer $s_t^*$ of (17) satisfies $D_t(x) s^*_t(x) = D_t(x) \nabla \log \rho^*_t(x)$ where $\rho^*_t$ is the solution to (1). Moreover, the map $X^*_t$ associated to $s_t^*$ is a transport map from ρ₀ to $\rho^*_t$.

Proof. If $X_t \sharp \rho_0 = \rho_t$, then by definition we have the identity

$$\begin{align} &\int_\Omega \left(|s_t(X_t(x))|_{D_t(X_t(x))}^2 +2 \nabla \cdot \left(D_t(X_t(x))s_t(X_t(x))\right) \right) \rho_0(x) \mathrm{d}x \nonumber\\ & \quad= \int_\Omega \left(|s_t(x)|_{D_t(x)}^2 +2 \nabla \cdot \left(D_t(x) s_t(x) \right)\right) \rho_t(x) \mathrm{d}x. \end{align} \tag{ B.12 }$$

This means that the optimization problem in (17) is equivalent to

$$\begin{equation*} \min_{s_t} \int_\Omega \left(|s_t(x)|_{D_t(x)}^2 +2 \nabla \cdot \left(D_t(x) s_t(x) \right)\right) \rho_t(x) \mathrm{d}x. \end{equation*}$$

All minimizers $s_t^*$ of this optimization problem satisfy $D_t(x)s^*_t(x) = D_t(x)\nabla\log \rho_t(x)$. Hence, by (10),

$$\begin{equation} \partial_t \rho_t(x) = -\nabla\cdot \left(b_t(x)\rho_t(x) - D_t(x)\nabla\rho_t(x)\right) \end{equation} \tag{ B.13 }$$

which recovers (1), so that $\rho_t(x) = \rho_t^*(x)$ solves (1).

In this section, we show that learning from the SDE alone—i.e. avoiding the use of self-consistent samples—does not provide a guarantee on the accuracy of ρ_t . We have already seen in (12) that it is sufficient to control $\int_0^{T}\int_{\Omega} |s_t(x) - \nabla\log\rho_t(x)|_{D_t(x)}^2\rho_t(x)\mathrm{d}x\mathrm{d}t$ to control $\mathsf{KL}(\rho_T\:\Vert\:\rho_T^*)$. The proof of proposition 1 shows that control on

$$\begin{equation} \int_0^T \int_{\Omega}|s_t(x) - \nabla\log\rho_t^*(x)|_{D_t(x)}^2\rho_t^*(x)\mathrm{d}x\mathrm{d}t, \end{equation} \tag{ B.14 }$$

as would be provided by training on samples from the SDE, does not ensure control on $\mathsf{KL}(\rho_T\:\Vert\:\rho_T^*)$. The following proposition shows that control on (B.14) does not guarantee control on $\mathsf{KL}(\rho_T^*\:\Vert\:\rho_T)$ either. An analogous result appeared in [34] in the context of SBDM for generative modeling; here, we provide a self-contained treatment to motivate the use of the sequential SBTM procedure discussed in the main text.

Proposition 4 (Learning on external data). Let $\rho_t:\Omega\rightarrow\mathbb{R}_{\gt 0}$ solve (10), and let $\rho_t^*:\Omega\rightarrow\mathbb{R}_{\gt 0}$ solve (1). Then, the following equality holds

$$\begin{align} \mathsf{KL}(\rho_T^*\:\Vert\:\rho_T) & = \int_0^T\int_{\Omega}|s_t(x) - \nabla\log\rho_t^*(x)|_{D_t(x)}^2\rho_t^*(x)\mathrm{d}x\mathrm{d}t\nonumber\\ &\quad + \int_0^T \int_{\Omega} \left(\nabla\log\rho_t(x) - s_t(x)\right)^\mathsf{T} D_t(x)\left(s_t(x) - \nabla\log\rho_t^*(x)\right)\rho_t^*(x)\mathrm{d}x\mathrm{d}t. \end{align} \tag{ 19 }$$

Proof. By an analogous argument as in the proof of proposition 1, we find

$$\begin{align*} \frac{\mathrm{d}}{\mathrm{d}t}\mathsf{KL}(\rho_t^*\:\Vert\:\rho_t) & = \int \left(\nabla\log\rho_t(x) - \nabla\log\rho_t^*(x)\right)^\mathsf{T} D_t(x)\left(s_t(x) - \nabla\log\rho_t^*(x)\right)\rho_t^*(x)\mathrm{d}x. \end{align*}$$

Adding and subtracting $s_t(x)$ to the first term in the inner product and expanding gives

$$\begin{align} \frac{\mathrm{d}}{\mathrm{d}t}\mathsf{KL}(\rho_t^*\:\Vert\:\rho_t) & = \int_{\Omega}|s_t(x) - \nabla\log\rho_t^*(x)|_{D_t(x)}^2\rho_t^*(x)\mathrm{d}x\nonumber\\ & \quad + \int_{\Omega} \left(\nabla\log\rho_t(x) - s_t(x)\right)^\mathsf{T} D_t(x)\left(s_t(x) - \nabla\log\rho_t^*(x)\right)\rho_t^*(x)\mathrm{d}x. \end{align} \tag{ B.15 }$$

Integrating from 0 to T completes the proof.

The following standard trick can be used to avoid computing the divergence of $s_t(x)$:

Lemma 8. Given $\xi = N(0,I)$, we have

$$\begin{equation} \begin{aligned} &\lim_{\alpha\downarrow0}\alpha^{-1} \mathbb{E} \big(s_{t}(x+\alpha \xi) \cdot \xi \big) = \nabla\cdot s_t(x),\\ &\lim_{\alpha\downarrow0}\alpha^{-1} \mathbb{E} \big(s_{t}(x+ \alpha \sigma_t(x) \xi) \cdot \sigma_{t} (x) \xi \big) = \mathrm{tr}\left(D_t(x) \nabla s_t(x)\right). \end{aligned} \end{equation} \tag{ B.16 }$$

Proof. We have

$$\begin{equation} \begin{aligned} \alpha^{-1} s_{t}(x+\alpha \xi) \cdot \xi & = \alpha^{-1} s_{t}(x) \cdot \xi + (\nabla s_t(x) \xi ) \cdot \xi + o(\alpha). \end{aligned} \end{equation} \tag{ B.17 }$$

The expectation of the first term on the right-hand side of this equation is zero; the expectation of the second gives the result in (B.16). Hence, taking the expectation of (B.17) and evaluating the result in the limit as $\alpha\downarrow 0$ gives the first equation in (B.16). The second equation in (B.16) can be proven similarly using $\sigma_t(x) \sigma_t(x)^\mathsf{T} = D_t(x)$.

Replacing $\nabla\cdot s_t(x)$ in (17) with the first expression in (B.17) for a fixed α > 0 gives the loss

$$\begin{equation} \mathcal{L}[s_t] = \mathbb{E}_\xi\left[\int_{\Omega}\left(|s_t(X_t(x))|^2 + \frac{2}{\alpha}s_t(X_t(x) + \alpha \xi)\cdot\xi\right)\rho_0(x)\mathrm{d}x\right]. \end{equation} \tag{ B.18 }$$

Evaluating the square term at a perturbed data point recovers the denoising loss of Vincent [61]

$$\begin{equation} \mathcal{L}[s_t] = \mathbb{E}_\xi\left[\int_{\Omega}\left|s_t(X_t(x) + \alpha\xi) + \frac{\xi}{\alpha}\right|^2\rho_0(x)\mathrm{d}x\right]. \end{equation} \tag{ B.19 }$$

We can improve the accuracy of the approximation with a 'doubling trick' that applies two draws of the noise of opposite sign to reduce the variance. This amounts to replacing the expectations in (B.16) with

$$\begin{equation} \begin{aligned} &\tfrac12\alpha^{-1} \mathbb{E} \big[ s_{t}(x+\alpha \xi) \cdot \xi - s_{t}(x- \alpha \xi) \cdot \xi \big],\\ &\tfrac12\alpha^{-1} \mathbb{E} \big[s_{t}(x+ \alpha \sigma_t(x) \xi) \cdot \sigma_{t} (x) \xi - s_{t}(x-\alpha \sigma_t(x) \xi) \cdot \sigma_{t} (x) \xi\big], \end{aligned} \end{equation} \tag{ B.20 }$$

whose limits as α → 0 are $\nabla\cdot s_t(x)$ and $\mathrm{tr}\left(D_t(x) \nabla s_t(x)\right)$, 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 of $\nabla \cdot s_t(x)$ even when the latter is feasible.

Network architecture. Both the single-particle energy $U_{\theta_t, 1}: \mathbb{R}^d\rightarrow \mathbb{R}$ and two-particle interaction energy $U_{\theta_t, 2}: \mathbb{R}^d\times\mathbb{R}^d \rightarrow \mathbb{R}$ are parameterized as single hidden-layer neural networks with the swish activation function [43] and $\mathtt{n\_hidden} = 100$ hidden neurons. The hidden layer biases are initialized to zero while the hidden layer weights are initialized from a truncated normal distribution with variance $1/\mathtt{fan\_in}$, following the guidelines recommended in [21].

Optimization. The Adam [23] optimizer is used with an initial learning rate of $\eta = 10^{-4}$ and otherwise default settings. At time t = 0, the analytical relative loss

$$\begin{equation} L[s_0] = \frac{\int |s_0(x) - \nabla \log \rho_0(x)|^2\rho_0(x) \mathrm{d}x}{\int |\nabla \log \rho_0(x)|^2\rho_0(x) \mathrm{d}x} \end{equation} \tag{ D.1 }$$

is minimized to a value less than 10⁻⁴ using knowledge of the initial condition $\rho_0 = \mathsf{N}\left(\beta_0, \sigma_0^2I\right)$ with $\sigma_0 = 0.25$. In (D.1), the expectation with respect to ρ₀ is approximated by an initial set of samples $x_{j} = (x^{(1)}_{j}, x^{(2)}_{j}, \ldots, x^{(N)}_{j})^\mathsf{T}$ with $j = 1, \ldots, n$ drawn from ρ₀. In the experiments, we set n = 100, which we found to be sufficient to obtain a few digits of relative accuracy on various quantities of interest. We set the physical timestep $\Delta t = 10^{-3}$ and take $\mathtt{n\_opt\_steps} = 25$ steps of Adam on the standard sequential SBTM loss function (17) until the norm of the gradient is below $\mathtt{gtol} = 0.1$.

Analytical moments. First define the mean, second moment, and covariance according to

$$\begin{align*} m_t^{(i)} & = \mathbb{E}\big[X_t^{(i)}\big],\\ M_t^{(ij)} & = \mathbb{E}\big[X_t^{(i)} \big(X_t^{(\,j)}\big)^\mathsf{T}\big],\\ C_t^{(ij)} & = M^{(ij)} - m^{(i)}\big(m^{(\,j)}\big)^\mathsf{T}. \end{align*}$$

It is straightforward to show that the mean and covariance obey the dynamics

$$\begin{align} \kern36.4pt \dot{m}_t^{(i)} & = -(m_t^{(i)} - \beta_t) +\frac{\alpha}{N}\sum_{k = 1}^{N}\left(m_t^{(i)} - m_t^{(k)}\right), \end{align} \tag{ D.2 }$$

$$\begin{align} \dot{C}_t^{(ij)} & = -2 (1 - \alpha)C_t^{(ij)} + 2DI\delta_{ij} - \frac{\alpha}{N}\sum_{k = 1}^{N} \left(C_t^{(kj)} + C_t^{(ik)}\right). \end{align} \tag{ D.3 }$$

Because the particles are indistinguishable so long as they are initialized from a distribution that is symmetric with respect to permutations of their labeling, the moments will satisfy the ansatz

$$\begin{align} m_t^{(i)} & = \bar m(t), \:\:\: i = 1, \ldots, N \end{align} \tag{ D.4 }$$

$$\begin{align} C_t^{(ij)} & = C_d(t) \delta_{ij} + C_o(t) (1 - \delta_{ij}), \:\:\: i, j = 1, \ldots, N. \end{align} \tag{ D.5 }$$

The dynamics for the vector $\bar m: \mathbb{R}_{\geqslant 0} \rightarrow \mathbb{R}^{\bar{d}}$, as well as the matrices $C_d: \mathbb{R}_{\geqslant 0} \to \mathbb{R}^{\bar{d} \times \bar{d}}$ and $C_o: \mathbb{R}_{\geqslant 0} \to \mathbb{R}^{\bar{d} \times \bar{d}}$ can then be obtained from (D.2) and (D.3) as

$$\begin{align*} \dot{\bar m} & = \beta_t - \bar m,\\ \dot{C}_d & = 2(\alpha - 1) C_d - 2\frac{\alpha}{N}\left(C_d + (n-1)C_o\right) + 2DI,\\ \dot{C}_o & = 2(\alpha - 1)C_o - 2\frac{\alpha}{N}\left(C_d + (n-1)C_o\right). \end{align*}$$

For a given $\beta:\mathbb{R}\rightarrow\mathbb{R}^{\bar{d}}$, these equations can be solved analytically in Mathematica as a function of time, giving the mean ${m}_t = \bar m(t)\otimes 1_N \in \mathbb{R}^{N\bar{d}}$ and covariance $C_t = (C_d(t) - C_o(t)) \otimes I_{N\times N} + C_o(t) \otimes \left(1_N 1_N^\mathsf{T}\right)\in \mathbb{R}^{N\bar{d}\times N\bar{d}}$. Because the solution is Gaussian for all t, we then obtain the analytical solution to the FPE $\rho^*_t = \mathsf{N}\left(m _t, C_t\right)$ and the corresponding analytical score $-\nabla \log \rho^*_t(x) = C_t^{-1}(x - m_t)$.

Potential structure. Here, we show that the potential for this example lies in the class of potentials described by (23). From equation (D.5), we have a characterization of the structure of the covariance matrix C_t for the analytical potential $U_t(x) = \frac{1}{2}(x - m_t)^\mathsf{T} C_t^{-1} (x - m_t)$. In particular, C_t is block circulant, and hence is block diagonalized by the roots of unity (the block discrete Fourier transform). That is, we may take a 'block eigenvector' of the form $\omega_k = \left(I_{\bar{d}\times \bar{d}} \rho^k, I_{\bar{d}\times \bar{d}} \rho^{2k}, \ldots, I_{\bar{d}\times \bar{d}} \rho^{(N-1)k}\right)^\mathsf{T}$ with $\rho = \exp(-2\pi i /N)$ for $k = 0, \ldots N-1$. By direct calculation, this block diagonalization leads to two distinct block eigenmatrices,

$$\begin{equation*} C_t = V \begin{pmatrix} C_d(t) + (N-1)C_o(t) & 0 & 0 & \ldots & 0\\ 0 & C_d(t) - C_o(t) & 0 & \ldots & 0\\ 0 & 0 & \ddots & \ldots & 0 \\ 0 & 0 & 0 & \ldots & C_d(t) - C_o(t) \end{pmatrix} V^{-1} \end{equation*}$$

where $V \in \mathbb{R}^{N\bar{d} \times N\bar{d}}$ denotes the matrix with block columns ω_k . The inverse matrix $C_t^{-1}$ then must similarly have only two distinct block eigenmatrices given by $\left(C_d(t) + (N-1)C_o(t)\right)^{-1}$ and $\left(C_d(t) - C_o(t)\right)^{-1}$. By inversion of the block Fourier transform, we then find that

$$\begin{equation*} \left(C_t^{-1}\right)^{(ij)} = \bar{C}_d\delta_{ij} + \bar{C}_o(1 - \delta_{ij}) \end{equation*}$$

for some matrices $\bar{C}_d, \bar{C}_o$. Hence, by direct calculation

$$\begin{align*} \left(x - m_t\right)^\mathsf{T} C_t^{-1} \left(x - m_t\right) & = \sum_{i, j}^N \left(x^{(i)} - m_t^{(i)}\right)^\mathsf{T}\left(C_t^{-1}\right)^{(ij)}\left(x^{(\,j)} - m_t^{(\,j)}\right)\\ & = \sum_{i, j}^N \left(x^{(i)} - \bar m(t)\right)^\mathsf{T}\left(\bar{C}_d\delta_{ij} + \bar{C}_o (1 - \delta_{ij})\right)\left(x^{(\,j)} - \bar m(t)\right)\\ & = \sum_{i}^N \left(x^{(i)} - \bar m(t)\right)^\mathsf{T}\bar{C}_d\left(x^{(i)} - \bar m(t)\right)^\mathsf{T} \\ & \quad + \sum_{i\neq j}^N\left(x^{(i)} - \bar m(t)\right)^\mathsf{T}\bar{C}_o\left(x^{(\,j)} - \bar m(t)\right). \end{align*}$$

Above, we may identify the first term in the last line as $\sum_{i = 1}^N U_1(x^{(i)})$ and the second term in the last line as $\frac{1}{N}\sum_{i \neq j}^N U_2(x^{(i)}, x^{(\,j)})$. Moreover, $U_2(\cdot, \cdot)$ is symmetric with respect to its arguments.

Analytical entropy. For this example, the entropy can be computed analytically and compared directly to the learned numerical estimate. By definition,

$$\begin{align*} s_t & = -\int_{\mathbb{R}^{N\bar d}} \log \rho_t(x) \rho_t(x)\mathrm{d}x,\\ & = - \int_{\mathbb{R}^{N\bar d}} \left(-\frac{N\bar d}{2}\log(2\pi) - \frac{1}{2}\log\det C_t - \frac{1}{2} (x-m_t)^\mathsf{T} C_t^{-1}(x-m_t)\right)\rho_t(x)\mathrm{d}x,\\ & = \frac{N\bar d}{2}\left(\log\left(2\pi\right) + 1\right) + \frac{1}{2}\log\det C_t. \end{align*}$$

Additional figures. Images of the learned velocity field and potential in comparison to the corresponding analytical solutions can be found in figures D1 and D2, 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.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Network architecture. Both potential terms $U_{\theta_t, 1}$ and $U_{\theta_t, 2}$ are modeled as four hidden-layer deep fully connected networks with $\mathtt{n\_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 $\eta = 5\times 10^{-3}$ and otherwise default settings. At time t = 0, the loss (D.1) is minimized to a value less than 10⁻⁶ over n samples $X_0^{(i)} \sim \otimes_{j = 1}^N \mathsf{N}(\beta_0, \sigma_0^2I)$, $i = 1, \ldots, n$ with $\sigma_0 = 0.5$ and $n = 10^4$. Past this initial stage, the denoising loss is used with a noise scale σ = 0.1; we found that a higher noise scale regularized the problem and led to a smoother prediction for the entropy, at the expense of a slight bias in the moments. By increasing the number of samples n, the noise scale can be reduced while maintaining an accurate prediction for the entropy. The loss is minimized by taking $\mathtt{n\_opt\_steps} = 25$ steps of Adam at each timestep. The physical timestep is set to $\Delta t = 10^{-3}$.

Additional figures. Figures D3 and D4 show the full grid of covariance components for the SDE, learned, and noise free systems. The noise free underestimates the moments, while the learned and SDE agree well.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Setup. We parameterize the score directly $s_t: \mathbb{R}^{2}\rightarrow\mathbb{R}$ using a three hidden layer neural network with $\mathtt{n\_hidden} = 32$ neurons per hidden layer. Because the dynamics is anti-symmetric, we impose that $s(x, v) = -s(-x, -v)$.

Optimization and initialization. The network initialization is identical to the previous two experiments. The physical timestep is set to $\Delta t = 10^{-3}$. The Adam optimizer is used with an initial learning rate of $\eta = 10^{-4}$. At time t = 0 the loss (D.1) is minimized to a tolerance of 10⁻⁴ over $n = 10^4$ samples drawn from an initial distribution $\mathsf{N}(0, \sigma_0^2I)$ with $\sigma_0 = 1$. The denoising loss is used with a noise scale σ = 0.05, using $\mathtt{n\_opt\_steps} = 25$ steps of Adam until the norm of the gradient is below $\mathtt{gtol} = 0.5$.