A Pseudoreversible Normalizing Flow for Stochastic Dynamical Systems with Various Initial Distributions

Abstract

SIAM Journal on Scientific Computing, Volume 46, Issue 4, Page C508-C533, August 2024.
Abstract. We present a pseudoreversible normalizing flow method for efficiently generating samples of the state of a stochastic differential equation (SDE) with various initial distributions. The primary objective is to construct an accurate and efficient sampler that can be used as a surrogate model for computationally expensive numerical integration of SDEs, such as those employed in particle simulation. After training, the normalizing flow model can directly generate samples of the SDE’s final state without simulating trajectories. The existing normalizing flow model for SDEs depends on the initial distribution, meaning the model needs to be retrained when the initial distribution changes. The main novelty of our normalizing flow model is that it can learn the conditional distribution of the state, i.e., the distribution of the final state conditional on any initial state, such that the model only needs to be trained once and the trained model can be used to handle various initial distributions. This feature can provide a significant computational saving in studies of how the final state varies with the initial distribution. Additionally, we propose to use a pseudoreversible network architecture to define the normalizing flow model, which has sufficient expressive power and training efficiency for a variety of SDEs in science and engineering, e.g., in particle physics. We provide a rigorous convergence analysis of the pseudoreversible normalizing flow model to the target probability density function in the Kullback–Leibler divergence metric. Numerical experiments are provided to demonstrate the effectiveness of the proposed normalizing flow model. Reproducibility of computational results. This paper has been awarded the “SIAM Reproducibility Badge: Code and data available” as a recognition that the authors have followed reproducibility principles valued by SISC and the scientific computing community. Code and data that allow readers to reproduce the results in this paper are available at https://github.com/mlmathphy/PRNF and in the supplementary materials (PRNF-main.zip [27.4MB]).