PROBABILITY FLOW SOLUTION OF THE FOKKER-PLANCK EQUATION

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.

1. INTRODUCTION

The time evolution of many dynamical processes occurring in the natural sciences, engineering, economics, and statistics are naturally described in the language of stochastic differential equations (SDE) (Gardiner, 2009; Oksendal, 2003; Evans, 2012) . Typically, one is 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)  ∂ t ρ * t (x) = -∇ • (b t (x)ρ * t (x) -D t (x)∇ρ * t (x)) , x ∈ Ω ⊆ R d , (FPE) where ρ * t (x) ∈ R ≥0 denotes the value of the density at time t, b t (x) ∈ R d is a vector field known as the drift, and D t (x) ∈ R d×d is a positive-semidefinite tensor known as the diffusion matrix. (FPE) must be solved for t ≥ 0 from some initial condition ρ * t=0 (x) = ρ 0 (x), but in all but the simplest cases, the solution is not available analytically and can only be approximated via numerical integration. High-dimensionality. For many systems of interest -such as interacting particle systems in statistical physics (Chandler, 1987; Spohn, 2012) , stochastic control systems (Kushner et 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 is a Monte-Carlo approach, whereby the SDE associated with (FPE) dx t = b t (x t )dt + ∇ • D t (x t )dt + √ 2σ t (x t )dW t , is evolved via numerical integration to obtain a large number n of trajectories (Kloeden & Platen, 1992) . In (1), σ t (x) satisfies σ t (x)σ T t (x) = D t (x) and W t is a standard Brownian motion on R d . Assuming that we can draw samples {x i 0 } n i=1 from the initial PDF ρ 0 , simulation of (1) enables the

