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 study 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 show theoretically that the proposed approach controls the KL divergence from the learned solution to the target, while learning on external samples from the stochastic differential equation does not control either direction of the KL divergence. Empirically, we consider several high-dimensional Fokker-Planck equations from the physics of interacting particle systems. We find that the method accurately matches analytical solutions when they are available as well as moments computed via Monte-Carlo when they are not. Moreover, the method offers compelling predictions for the global entropy production rate that out-perform those obtained from learning on stochastic trajectories, and can effectively capture non-equilibrium steady-state probability currents over long time intervals.
翻译:高维含时福克-普朗克方程求解的首选方法是通过求解相关的随机微分方程生成样本解。本文研究了一种基于描述概率流动的常微分方程求解的替代方案。该方程作为传输映射,可确定性将初始密度的样本推至任意后续时刻的解样本。与随机动力学积分不同,该方法具有直接获取仅靠轨迹难以估计量的优势,例如概率流、密度本身及其熵。概率流方程依赖于解的对数梯度(即“分数”),因此先验未知。为解决此依赖关系,我们使用深度神经网络对分数进行建模,该网络根据瞬时概率流传播样本集在线学习。理论证明,所提方法能控制学习解与目标之间的KL散度,而基于随机微分方程外部样本的学习无法控制KL散度的任一方向。在实验方面,我们考虑来自相互作用粒子系统物理学的多个高维福克-普朗克方程。结果表明,该方法能精确匹配解析解(当存在时)以及通过蒙特卡洛计算的高阶矩(当解析解不存在时)。此外,该方法在全局熵产生率预测方面优于基于随机轨迹学习的传统方法,并能有效捕捉长时间间隔内的非平衡稳态概率流。