remarkRemark hypothesisHypothesis claimClaim assumAssumption A PR-NF for
stochastic dynamical systemsM. Yang, P. Wang, D. del-Castillo-Negrete, Y. Cao,
G. Zhang

# A pseudo-reversible normalizing flow for stochastic dynamical systems with
various initial distributions

M. Yang Fusion Energy Division, Oak Ridge National Laboratory, Oak Ridge, TN
<EMAIL_ADDRESS><EMAIL_ADDRESS>P. Wang Department of Mathematics,
Auburn University, Auburn, AL<EMAIL_ADDRESS><EMAIL_ADDRESS>D.
del-Castillo-Negrete11footnotemark: 1 Y. Cao22footnotemark: 2 G. Zhang
Computer Science and Mathematics Division, Oak Ridge National Laboratory, Oak
Ridge, TN<EMAIL_ADDRESS>
###### Abstract

We present a pseudo-reversible normalizing flow method for efficiently
generating samples of the state of a stochastic differential equation (SDE)
with different 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 SDE, 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. Existing normalizing flows for SDEs depend on the initial
distribution, meaning the model needs to be re-trained 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.
We provide a rigorous convergence analysis of the pseudo-reversible
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.

###### keywords:

Stochastic differential equations, Normalizing Flows, Fokker-Planck equation

68Q25, 68R10, 68U05

##  1 Introduction

The numerical solution of stochastic differential equations (SDE) for large
ensembles of initial conditions is at the heart of modeling and simulation in
science and engineering. For example, computations of interest to plasma
physics rely on the solution of SDE tracking the evolution of electrons and
ions in which the deterministic (drift) part describes the interaction with
prescribed or self-consistent electromagnetic fields, and the stochastic part
describes collisional effects. In fluid mechanics, SDE describe the evolution
of passive tracers, e.g. pollutants in the oceans and the atmosphere, advected
by prescribed velocity fields under the effect of molecular and/or turbulent
diffusion. Further applications of SDE include the simulation of biological
and chemical systems. From a mathematical perspective, SDE naturally appear as
the discrete particle description formally equivalent to the Fokker-Planck
description of the particle distribution function in the continuum limit. This
equivalence is the foundation of particle-based Monte-Carlo methods for the
numerical solution of Fokker-Planck type partial differential equations (PDE).
Although particle based numerical methods offer advantages over continuum
methods (e.g., finite difference, fine elements, or spectral), they also face
limitations. In particular, to avoid numerical noise, Monte Carlo methods need
to integrate SDE for very large ensembles of particles to overcome the
unfavorable slow convergence $\sim 1/\sqrt{N}$ where $N$ is the number of
initial conditions.

To overcome numerical challenges faced by particle methods, in this paper we
propose an efficient and accurate algorithm for the solution of SDE. Our
approach is based on the use of pseudo-reversible normalizing flows (PR-NF) to
efficiently generate samples of the state of an SDE for various initial
distributions. Normalizing flows [12, 19, 3, 16, 5, 8, 9] are a class of
models that map an unknown probability distribution to a standard probability
distribution, e.g., the normal Gaussian distribution, such that one can
generate samples from the unknown distribution by sampling the standard
distribution and then propagating through the flow model. Normalizing flows
have been extended to learn stochastic dynamical systems [1, 13, 7]. However,
these models are not transferable, meaning they require model re-training
every time the initial distribution is changed.

We improve the normalizing flow models from three perspectives. First, we
incorporate the initial state of the SDE as an input variable of the
normalizing flow model, so that that our method aims at learning the
conditional distribution of the state at a given time instant. In this way, a
single normalizing flow model can be used to handle different initial
distributions by a simple convolution. In this sense, the proposed method can
be viewed as an approach to learn the Green’s function of the Fokker-Planck
equation associated with the SDE. Second, we utilize the pseudo-reversible
neural network architecture [10, 24, 20] to relax the strict reversibility
constraint. This architecture provides greater flexibility in designing flow
transformations and offers potential benefits in terms of performance and
applicability. The pseudo-reversible neural network architecture allows to use
simple feed-forward neural networks to model both the forward and inverse
flows, where an extra loss term is used to ensure pseudo-reversibility. Even
though the proposed PR-NF has $\mathcal{O}(d^{3})$ complexity in Jacobian
determinant computation, we observe that it is not a bottleneck, especially
with GPU-accelerated linear algebra libraries, for a wide range of physical
processes defined in the six-dimensional phase space. Third, we provide a
convergence analysis of the PR-NF model to the target probability density
function in the Kullback–Leibler divergence metric, where the simple structure
of the PR-NF model makes it feasible to exploit the universal approximation
theorem of fully-connected neural networks in our convergence analysis.

The outline of the rest of the paper is as follows. In Section 2 we formulate
the problem and define the SDE of interest. The details of the surrogate model
based on Normalizing Flows architecture are described in Section 3. Section 4
provides the convergence analysis of the PR-NF model. Section 5 is devoted to
examples including a Brownian motion validation model, a problem of interest
to magnetically confined fusion plasmas and an advection-diffusion transport
problem in fluid mechanics.

##  2 Problem setting

We consider the $d$-dimensional stochastic process $X_{t}$, corresponding to
the solution of the stochastic differential equation

(1)
$X_{t}=X_{0}+\int_{0}^{t}{\bm{b}}(s,X_{s})ds+\int_{0}^{t}{\bm{\sigma}}(s,X_{s})dW_{s}\;\;\text{
with }X_{0}\in\mathcal{D},\,\,t\in[0,T],$

where $T>0$, $\mathcal{D}$ is a bounded domain in $\mathbb{R}^{d}$,
${\bm{b}}:[0,T]\times\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$,
${\bm{\sigma}}:[0,T]\times\ \mathbb{R}^{d}\rightarrow\mathbb{R}^{d\times m}$
are drift and diffusion coefficients, respectively, and
$W_{s}:=(W_{s}^{1},\ldots,W_{s}^{m})^{\top}$ is a $m$-dimensional standard
Brownian motion. We assume that ${\bm{b}}$ and ${\bm{\sigma}}$ satisfy, for
some constant $C$,

(2) $|{\bm{b}}(t,{\bm{x}})|+|{\bm{\sigma}}(t,{\bm{x}})|\leq
C(1+|{\bm{x}}|),\text{ for }{\bm{x}}\in\mathbb{R}^{d},t\in[0,T],$

and

(3)
$|{\bm{b}}(t,{\bm{x}})-{\bm{b}}(t,{\bm{y}})|+|{\bm{\sigma}}(t,{\bm{x}})-{\bm{\sigma}}(t,{\bm{y}})|\leq
C|{\bm{x}}-{\bm{y}}|,\text{ for
}{\bm{x}},{\bm{y}}\in\mathbb{R}^{d},t\in[0,T].$

For an initial state $X_{0}\in\mathcal{D}$, the SDE in Eq. (1) has a unique
time continuous solution $X_{t}$ [14]. As the state $X_{t}$ in Eq. (1) depends
on the initial condition $X_{0}$, we will write $X_{t}$ in the conditional
form $X_{t}|X_{0}$ when we need to emphasize this dependence.

In applications one is typically interested in evaluating quantities of
interest, ${\rm QoI}$, of the form

(4) ${\rm
QoI}\,=\int_{\mathbb{R}^{d}}F({\bm{x}_{t}})\,p_{X_{t}}({\bm{x}_{t}})d{\bm{x}}_{t}\,,$

where $F$ denotes a physical property of the system, and $p_{X_{t}}$ is the
probability density of the final state. Because of the dependence of the final
state on the initial condition, normalizing flow strategies focusing on the
distribution $p_{X_{t}}({\bm{x}_{t}})$, e.g., Refs. [1, 13, 7], might not be
efficient since every time the initial condition changes, the evaluation of
${\rm QoI}$ requires re-training of the neural network for
$p_{X_{t}}({\bm{x}_{t}})$ corresponding to the new initial condition.

To circumvent this problem we propose here a new method that focuses on the
use of normalizing flows to create a surrogate model for the probability
density $p_{X_{t}|X_{0}}({\bm{x}_{t}|{\bm{x}_{0}}})$, which gives the
transition probability of being at ${\bm{x}_{t}}$ stating at ${\bm{x}_{0}}$.
Once a normalizing flow model is trained for
$p_{X_{t}|X_{0}}({\bm{x}_{t}|{\bm{x}_{0}}})$, the computation of the ${\rm
QoI}$ reduces to the direct evaluation of the integral

(5) ${\rm
QoI}=\int_{\mathbb{R}^{d}}\int_{\mathbb{R}^{d}}F({\bm{x}_{t}})\,p_{X_{t}|X_{0}}({\bm{x}_{t}}|{\bm{x}}_{0})p_{X_{0}}({\bm{x}}_{0})d{\bm{x}_{t}}d{\bm{x}}_{0}\,,$

where $p_{X_{0}}({\bm{x}}_{0})$ denotes the probability density of the initial
condition. In this approach the computational cost of evaluating ${\rm QoI}$
for different initial conditions reduces to the cost of performing a
quadrature without a retraining overhead. If $\\{{\bm{x}}_{0}^{(m)}\\}$ for
$i=1\,\ldots M$ is a set of $M$-samples drawn from $p_{X_{0}}({\bm{x}}_{0})$,
and $\\{{\bm{x}}_{t}^{(n)}\\}$ is the corresponding set of $N$-samples drawn
from $p_{X_{t}|X_{0}}({\bm{x}_{t}}|{\bm{x}}_{0})$, then the quadrature can be
approximated as

(6) ${\rm
QoI}\approx\frac{1}{MN}\sum_{m=1}^{M}\sum_{n=1}^{N}F\left({\bm{x}}_{t}^{(n)}|\bm{x}_{0}^{(m)}\right),$

##  3 The pseudo-reversible normalizing flow (PR-NF)

A normalizing flow model can be viewed as a nonlinear transformation that maps
an unknown target random variable, denoted by $X$, to a standard random
variable, e.g., the standard Gaussian, denoted by $Z$. Because we intend to
take into account the initial condition $X_{0}$ in our normalizing flow model,
we define the target random variable $X$ as

(7) $X:=(X_{0},X_{t}|X_{0})\in\mathbb{R}^{2d}.$

A normalizing flow model includes a forward transformation mapping $X$ to $Z$
and an inverse transformation mapping $Z$ to $X$, i.e.,

(8) $Z=\bm{h}(X;\theta_{h})\;\;\text{ and
}\;\;\widehat{X}=\bm{g}(Z;\theta_{g}),$

where $\theta_{h}$, $\theta_{g}$ are parameters of the transformations and
$\widehat{X}=X$ if $\bm{g}=\bm{h}^{-1}$. The function ${\bm{h}}$ moves (or
flows) in the normalizing direction, i.e., from an unknown distribution to a
standard distribution $p_{Z}$. The tradition of defining ${Z}$ as the standard
normal random variable $\mathcal{N}(0,\mathbb{I}_{2d})$ gives rise to the name
“normalizing flow”. Let $p_{X}(\bm{x})$ and $p_{Z}(\bm{z})$ denote the PDF of
$X$ and $Z$, respectively. The relation between the two PDFs can be obtained
by the change of variables formula, i.e.,

(9) $\displaystyle
p_{X}({\bm{x}})=p_{Z}({\bm{z}})\left|\frac{\partial{\bm{z}}}{\partial{\bm{x}}}\right|$
$\displaystyle=p_{Z}({\bm{h}}({\bm{x}}))|{\rm
det}{\mathbf{J}_{h}}({\bm{x}})|,$

where ${\rm det}{\mathbf{J}_{\bm{h}}}({\bm{x}})$ is the determinant of the
Jacobian matrix $\mathbf{J}_{\bm{h}}$ of the forward map $\bm{h}(\cdot)$.

###  3.1 The PR-NF model’s architecture

The PR-NF model uses the pseudo-reversible neural network architecture [10,
24, 20], i.e., $\bm{g}\approx\bm{h}^{-1}$ in Eq. (8), to replace the exact
reversible architectures used in most existing normalizing flow models, e.g.,
MAF [16], Real NVPs [5]. The typical exact reversible architectures rely on
functions with either triangular Jacobians or Lipschitz continuity, where the
Jacobian determinant can be efficiently approximated. However, the relaxation
of the reversibility condition increases the flexibility of the model design.
Specifically, we define $\bm{h}(\cdot;\theta_{h})$ and
$\bm{g}(\cdot;\theta_{g})$ as two independent fully-connected neural networks,
and the pseudo-reversibility is imposed by adding a soft constraint
$\|\widehat{X}-X\|_{2}^{2}$ to the loss function. This architecture is
different from autoencoders because the size of the bottleneck, i.e., the
dimension of $Z$, is the same as the dimension of $X$ and $\widehat{X}$.
Figure 1 illustrates the performance of the PR-NF on a standard two-
dimensional dataset. We observe that the left panel (i.e., the true
distribution) and the right panel (i.e., the PR-NF generated distribution)
have a good agreement, meaning the relaxation of the reversibility does not
cause a significant error.

Figure 1: Illustration of the PR-NF’s performance. The function ${\bm{h}}$
maps the true distribution (left) to the standard normal distribution
(middle), and the function $\bm{g}$ generates new samples of the true
distribution. We observe that the left panel and the right panel have a good
agreement, meaning the relaxation of the reversibility does not cause a
significant error.

The novelty of the PR-NF model is to learn the conditional distribution in Eq.
(5), such that we can use the trained PR-NF model to compute the QoI in Eq.
(5) for arbitrary initial conditions, $p_{X_{0}}(\bm{x}_{0})$, without re-
training. The key ingredient is to include the initial state $X_{0}$ into the
input of the PR-NF model, i.e., Eq. (7). The network architecture is shown in
Figure 2. Specifically, the mapping $\bm{h}(\cdot)$ consists of two
components, i.e.,
$\bm{h}(\bm{x}_{0},\bm{x}_{t}|\bm{x}_{0})=(\bm{h}_{0}(\bm{x}_{0}),\bm{h}_{1}(\bm{x}_{0},\bm{x}_{t}|\bm{x}_{0};\theta_{h}))$,
where
$\bm{h}_{1}(\bm{x}_{0},\bm{x}_{t}|\bm{x}_{0};\theta_{h}):\mathbb{R}^{2d}\mapsto\mathbb{R}^{d}$
is a fully connected neural network with tuning parameter $\theta_{h}$, and
$\bm{h}_{0}(\bm{x}_{0}):\mathbb{R}^{d}\mapsto\mathbb{R}^{d}$ is an identity
mapping depending only on $\bm{x}_{0}$. We emphasize that it is critical to
include $\bm{x}_{0}$ as one input of $\bm{h}_{1}$ to characterize the
conditional distribution $p_{X_{t}|X_{0}}(\bm{x}_{t}|\bm{x}_{0})$. We denote
by $\bm{z}_{0}$ the output of $\bm{h}_{0}(\bm{x}_{0})$ and $\bm{z}_{t}$ the
output of $\bm{h}_{1}(\bm{x}_{0},\bm{x}_{t}|\bm{x}_{0};\theta_{h})$. The
inverse map $\bm{g}$ has a similar structure as the forward map, i.e.,
$\bm{g}(\bm{z}_{0},\bm{z}_{t})=(\bm{g}_{0}(\bm{z}_{0}),\bm{g}_{1}(\bm{z}_{0},\bm{z}_{t};\theta_{g}))$,
where $\bm{g}_{0}$ is an identity map and $\bm{g}_{1}$ is a fully connected
neural network with tuning parameter $\theta_{g}$. We denote by
$\widehat{\bm{x}}_{0}$ and $\widehat{\bm{x}}_{t}$ the outputs of $\bm{g}_{0}$
and $\bm{g}_{1}$, respectively. Note that
$\widehat{\bm{x}}_{0}=\bm{z}_{0}=\bm{x}_{0}$ because $\bm{h}_{0}$ and
$\bm{g}_{0}$ are identity maps.

Figure 2: The network architecture of the proposed PR-NF model.

###  3.2 The loss function

The loss function is defined based on a training data set, i.e.,

(10) $\mathcal{V}_{\rm
train}=\left\\{{\bm{x}}^{(n)}\right\\}_{n=1}^{N}=\left\\{\bm{x}_{0}^{(n)},\bm{x}_{t}^{(n)}|\bm{x}_{0}^{(n)}\right\\}_{n=1}^{N},$

where $\\{\bm{x}_{0}^{(n)}\\}_{n=1}^{N}$ is generated from the uniform
distribution over the bounded domain $\mathcal{D}$ in Eq. (1). For each
initial state $\bm{x}_{0}^{(n)}$, we numerically solve the SDE in Eq. (1) to
obtain the sample $\bm{x}_{t}^{(n)}|\bm{x}_{0}^{(n)}$. The loss function
consists of two components, i.e.,

(11) $\mathcal{L}=\mathcal{L}_{1}+\lambda\mathcal{L}_{2},$

where $\mathcal{L}_{1}$ is the negative log-likelihood loss defined by

(12) $\mathcal{L}_{1}=-\frac{1}{N}\sum_{n=1}^{N}\left({\rm
log}p_{Z_{t}}({\bm{h}}_{1}({\bm{x}}^{(n)};\theta_{h}))+{\rm log}\left|{\rm
det}{\bf J}_{\bm{h}}({\bm{x}}^{(n)};\theta_{h})\right|\right),$

with $p_{Z_{t}}(\cdot)$ the probability density function of the standard
normal distribution, and $\mathcal{L}_{2}$ is the pseudo-reversibility loss
that measures the difference between ${\bm{x}}$ and $\widehat{\bm{x}}$, i.e.,

(13)
$\mathcal{L}_{2}=\frac{1}{N}\sum_{n=1}^{N}\left(\left\|{\bm{x}}^{(n)}-{\bm{g}}({\bm{h}}({\bm{x}}^{(n)};\theta_{h});\theta_{g})\right\|_{2}^{2}+\left|\det{\bf
J}_{\bm{g}}({\bm{h}}({\bm{x}}^{(n)}))\det{\bf
J}_{\bm{h}}({\bm{x}}^{(n)})-1\right|\right),$

and $\lambda$ is a hyperparameter that will be discussed in Section 3.3.

The Jacobian determinant in $\mathcal{L}_{1}$ can be written as

(14) $\left|{\rm
det}{\mathbf{J}_{\bm{h}}}({\bm{x}})\right|=\left|\text{det}\begin{pmatrix}\mathbb{I}_{d},\quad
0\\\
\frac{\partial\bm{z}_{t}}{\partial\bm{x}_{0}},\frac{\partial\bm{z}_{t}}{\partial\bm{x}_{t}}\end{pmatrix}\right|=\left|\text{det}\left(\frac{\partial\bm{z}_{t}}{\partial\bm{x}_{t}}\right)\right|,$

where $\mathbb{I}_{d}$ is a $d\times d$ identity matrix. The simplification in
Eq. (14) is based on the fact ${\bm{h}_{0}}$ is an identity map with respect
to $\bm{x}_{0}$. Thus, although we doubled the size of the input of the PR-NF
model, the size of the Jacobian matrix is still $d\times d$. The proposed PR-
NF has $\mathcal{O}(d^{3})$ complexity in Jacobian determinant computation.
However, we observe that it is not a bottleneck with GPU-accelerated linear
algebra libraries in Pytorch and TensorFlow, especially for a wide range of
physical processes, i.e., the examples in Section 5, in the phase space with
dimension up to six.

###  3.3 Hyperparameter tuning

The hyperparameter $\lambda$ in Eq. (11) determines the importance of the
reversibility loss $\mathcal{L}_{2}$. Thus, a hyperparameter tuning is needed
to achieve a good performance. In this work, we will perform grid search and
use the cross-entropy as a metric to find the best value for $\lambda$.
Specifically, we will generate a set of candidates, denoted by
$\\{\lambda_{j}\\}_{j=1}^{J}$. For each $\lambda_{j}$, we train one PR-NF
model, denoted by $(\bm{h}^{(j)}(\bm{x}),\bm{g}^{(j)}(\bm{z}))$. Then, we
generate $M$ samples of $\widehat{\bm{x}}$ by running the inverse mapping,
i.e.,

(15)
$\left\\{\widehat{\bm{x}}^{(j,m)}=\bm{g}^{(j)}(\bm{z}^{(m)}),m=1,\ldots,M\right\\},$

where
$\\{\bm{z}^{(m)}\\}_{m=1}^{M}=\\{\bm{z}_{0}^{(m)},\bm{z}_{t}^{(m)}\\}_{m=1}^{M}$
with $\bm{z}_{0}^{(m)}$ generated uniformly from $\mathcal{D}$ and
$\bm{z}_{t}^{(m)}$ generated from the standard normal distribution. Because it
does not require solving the SDE in Eq. (1), we can generate a large number of
samples. Next, we build a kernel density estimator using the samples in Eq.
(15), and compute the cross-entropy by

(16) $H(\lambda_{j})=-\frac{1}{N}\sum_{n=1}^{N}\log(p_{\rm
KDE}(\bm{x}_{t}^{(n)})),$

where $p_{\rm KDE}$ is the kernel density estimator based on the samples in
Eq. (15) and $\\{\bm{x}_{t}^{(n)}\\}_{n=1}^{N}$ is the training data set in
Eq. (10). The optimal hyperparameter $\lambda$ is obtained by minimizing the
cross-entropy, i.e., $\lambda=\text{argmin}\,H(\lambda_{j})$.

###  3.4 Discussion on the computational cost

The algorithmic cost of sampling methods consists of two parts, the “offline
cost” and the “online cost”. The offline cost of an algorithm represents the
running time for the training process. It is determined by the architecture of
the neural network, the optimizer (numerical scheme), the size of the training
dataset and the number of training epochs (“time” steps). An epoch in machine
learning means one complete pass of the training dataset through the
algorithm. We choose the number of epochs such that the loss function stably
attains the minimum. The Monte Carlo method for sampling has no training
process. However, it needs to be run repeatedly for each new initial
condition. The PR-NF model is a generative model where the offline cost
accounts for most of the total algorithmic cost. However, the offline cost is
a one-time cost (no additional training is needed for future simulations). We
denote the offline cost of the PR-NF method as $C_{\rm offline}$. It is easy
to utilize the GPU architecture to accelerate the computation for the training
process under the PyTorch machine learning framework.

The online cost is the computational cost for sampling. Although the Monte
Carlo method is conceptually simple and algorithmically straightforward, the
associated computational cost can be staggeringly high. In general, the method
requires many samples to get a good approximation, which may incur an
arbitrarily large total runtime if the processing time of a single sample is
high [22]. In contrast, the PR-NF model is a transformation that maps samples
from the standard normal distribution to the distribution of interest. For an
initial distribution, the computational cost of the MC method is significantly
greater than the online cost of the PR-NF method because it involves the
sampling and numerical integration of the SDE in Eq. (1). The computational
cost depends on how many samples need to generate. We denote $C_{\rm MC}$ and
$C_{\rm PR-NF}$ as the computational cost of an initial distribution for MC
and PR-NF methods, respectively. We have $C_{\rm MC}\gg C_{\rm PR-NF}$. To
test $N_{\rm init}$ different initial distributions, the cost of the PR-NF
method $\mathcal{O}(C_{\rm offline}+C_{\rm PR-NF}\times N_{\rm init})$ is
cheaper than the MC method $\mathcal{O}(C_{\rm MC}\times N_{\rm init})$ when
$N_{\rm init}$ is large.

##  4 Convergence analysis of the PR-NF model

This section provides the convergence analysis of the pseudo-reversible
normalizing flow (PR-NF) model in Section 3. For notational simplicity, the
target random variable and its probability density function are denoted by $X$
and $p_{X}({\bm{x}})$, respectively. Here, we limit our attention to single-
hidden-layer neural networks. The work in [15] shows that for any pair of
well-behaved111A well-behaved distribution $p_{X}({\bm{x}})$ means that
$p_{X}({\bm{x}})>0$ for all ${\bm{x}}\in\mathbb{R}^{d}$, and all conditional
probabilities ${\rm Pr}(X_{i}\leq x_{i}|{\bm{x}}_{<i})$ are differentiable,
for $i=1,2,\ldots,d$. distributions $p_{X}({\bm{x}})$ and $p_{Z}({\bm{z}})$,
there exists a diffeomorphism
$\bm{z}=\bm{f}(\bm{x})=(f_{1}(\bm{x}),\ldots,f_{d}(\bm{x}))$ that can
transform $p_{X}({\bm{x}})$ to $p_{Z}({\bm{z}})$, i.e.,

(17) $p_{X}(\bm{x})=p_{Z}(\bm{f}(\bm{x}))|\det{\bf J}_{\bm{f}}(\bm{x})|.$

The proposed PR-NF model can be viewed as an approximation of the
diffeomorphism, where $\bm{h}(\bm{x})$ and $\bm{g}(\bm{z})$ in Eq. (8)
approximate $\bm{f}$ and $\bm{f}^{-1}$, respectively.

We study the convergence of the PR-NF model in three stages. Section 4.1 shows
the existence of a convergent neural network approximations to the target
diffeomorphism $\bm{f}$ by exploiting the universal approximation theorm of
fully connected neural networks. Section 4.2 shows that the convergence of the
loss functions $\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ in Eq. (11) is a
necessary condition to obtain the convergent approximation discussed in
Section 4.1, which verifies that appropriateness of the loss function used to
train the PR-NF model. Section 4.3 shows that the convergence of
$\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ plus some mild assumptions on the
$\bm{h}$ and $\bm{g}$ is sufficient to ensure the convergence of the KL
divergence between the true and the approximate probability density function.

###  4.1 Existence of a convergent approximation to the diffeomorphism

######  Assumption 1.

We assume that the determinant of the Jacobian matrix and any partial
derivative of $\bm{f}\in C^{1}(\mathbb{R}^{d})$ are bounded, i.e.,

(18) $\left|\frac{\partial f_{i}(\bm{x})}{\partial
x_{j}}\right|<A_{1}<\infty,\quad|\det\mathbf{J}_{\bm{f}}(\bm{x})|>A_{2}>0,$

where $A_{1}$ and $A_{2}$ are positive constants and $i,\
j\in\\{1,\ldots,d\\}$. Additionally, the probability density function
$p_{X}(\bm{x})$ is bounded and all the conditional probabilities are
differentiable. There exists a positive constant $K$ such that

(19) $p_{X}(\bm{x})<A_{3}\exp(-\alpha\|\bm{x}\|^{2})\;\text{ for
}\;\|\bm{x}\|>K,$

where $\alpha>0$ and $A_{3}>0$ are positive constants and $\|\bm{x}\|$ denotes
the $l^{2}$ norm.

Under Assumption 1, Lemma 4.1 and Lemma 4.2 demonstrate the existence and
reversibility of convergent flows $\bm{h}$ and $\bm{g}$, respectively. To
proceed, we introduce some multivariate notations: $\mathbb{Z}_{+}^{d}$ denote
the lattice of nonnegative multi-integers in $\mathbb{R}^{d}$. For
$\mathbf{m}=(m_{1},\ldots,m_{d})\in\mathbb{Z}_{+}^{d}$, we set
$|\mathbf{m}|=m_{1}+\ldots+m_{d}$, $\bm{x}^{\mathbf{m}}=x_{1}^{m_{1}}\cdots
x_{d}^{m_{d}}$, and $D^{\mathbf{m}}=\frac{\partial^{|\mathbf{m}|}}{\partial
x_{1}^{m_{1}}\cdots\partial x_{d}^{m_{d}}}$. We also define

$C^{k}(\Omega):=\big{\\{}{\bm{f}}:D^{\mathbf{m}}f_{i}\in C(\Omega)\ \text{for
all}\ i\in\\{1,\ldots,d\\},\,\text{and}\,|\mathbf{m}|\leq k\big{\\}}.$

######  Lemma 4.1 (Theorem 4.1 in [18]).

For any given $\varepsilon>0$, there exists a single-hidden-layer neural
network $\bm{h}=(h_{1},\ldots,h_{d})$ with sufficient number of neurons to
approximate a function $\bm{f}\in C^{k}(\mathbb{R}^{d})$ in a compact set
$\Omega\subset\mathbb{R}^{d}$ such that

(20)
$\max_{{\bm{x}}\in\Omega}|D^{\mathbf{m}}f_{i}(\bm{x})-D^{\mathbf{m}}h_{i}(\bm{x})|<\varepsilon,$

for all $i\in\\{1,\ldots,d\\}$ and $\mathbf{m}\in\mathbb{Z}_{+}^{d}$ with
$|\mathbf{m}|\leq k$.

######  Lemma 4.2.

Under Assumption 1 and Lemma 4.1, the flow $\bm{h}$ approximating $\bm{f}\in
C^{k}(\mathbb{R}^{d})$ is invertible for sufficiently small $\varepsilon>0$.

######  Proof 4.3.

Let $\mathcal{S}_{d}$ be the set of all permutation of $\\{1,\ldots,d\\}$, we
have

(21) $|\det{\bf J}_{\bm{f}}(\bm{x})-\det{\bf
J}_{\bm{h}}(\bm{x})|=\left|\sum_{\sigma\in\mathcal{S}_{d}}\left({\rm
sgn}(\sigma)\left(\prod_{i=1}^{d}\frac{\partial f_{i}}{\partial
x_{\sigma_{i}}}-\prod_{i=1}^{d}\frac{\partial h_{i}}{\partial
x_{\sigma_{i}}}\right)\right)\right|<C\varepsilon$

for any $\bm{x}\in\Omega$, where the last inequality holds by Lemma 4.1 with
$k=1$. According to Eq. (18), we have

(22) $A_{2}<|\det{\bf
J}_{\bm{f}}(\bm{x})|=\left|\sum_{\sigma\in\mathcal{S}_{d}}\left({\rm
sgn}(\sigma)\prod_{i=1}^{d}\frac{\partial f_{i}}{\partial
x_{\sigma_{i}}}\right)\right|<CA_{1}^{d}.$

We choose $\varepsilon<\frac{A_{2}}{2C}$ such that

(23) $\displaystyle|\det{\bf J}_{\bm{h}}(\bm{x})|\leq|\det{\bf
J}_{\bm{f}}(\bm{x})|+|\det{\bf J}_{\bm{f}}(\bm{x})-\det{\bf
J}_{\bm{h}}(\bm{x})|<\frac{A_{2}}{2}+CA_{1}^{d},$ (24) $\displaystyle|\det{\bf
J}_{\bm{h}}(\bm{x})|\geq|\det{\bf J}_{\bm{f}}(\bm{x})|-|\det{\bf
J}_{\bm{f}}(\bm{x})-\det{\bf J}_{\bm{h}}(\bm{x})|>\frac{A_{2}}{2}>0.$

By the inverse function theorem, $\bm{h}$ is invertible in $\Omega$.

Similarly, we can prove the invertibility of ${\bm{g}}$ for the function
${\bm{f}}^{-1}$ with Eq. (20). So far, we proved the existence of convergent
neural network $\bm{h}$ and $\bm{g}$ to approximate $\bm{f}$ and
$\bm{f}^{-1}$, respectively.

###  4.2 Convergence of the loss function

We now prove that the convergence of the loss functions $\mathcal{L}_{1}$ and
$\mathcal{L}_{2}$ in Eq. (11) is a necessary condition to obtain the
convergent approximation discussed in Section 4.1. To proceed, we define two
auxiliary random variables

(25) $\widetilde{X}={\bm{h}^{-1}}(Z)\;\text{ and }\;\widehat{X}={\bm{g}}(Z),$

where $Z$ follows the standard normal distribution. Using the change of
variables formula, the probability density functions of $\widetilde{X}$ and
$\widehat{X}$ are defined by

(26) $p_{\widetilde{X}}(\bm{x})=p_{Z}(\bm{h}(\bm{x}))|\det{\bf
J}_{\bm{h}}(\bm{x})|,\quad
p_{\widehat{X}}(\bm{x})=p_{Z}(\bm{g}^{-1}(\bm{x}))|\det{\bf
J}_{\bm{g}^{-1}}(\bm{x})|.$

For simplicity, we consider the loss functions in the continuous form, i.e.,

(27) $\mathcal{L}_{1}=-\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log
p_{\widetilde{X}}(\bm{x})d\bm{x},$

and

(28)
$\mathcal{L}_{2}=\int_{\Omega}p_{X}(\bm{x})\left(\|{\bm{g}}({\bm{h}}(\bm{x}))-\bm{x}\|^{2}+|\det{\bf
J}_{\bm{g}}({\bm{h}}(\bm{x}))\det{\bf J}_{\bm{h}}(\bm{x})-1|\right)d\bm{x},$

where $\Omega$ is a compact set. For the variables ${\widetilde{X}}$ and
${\widehat{X}}$, we have the following assumptions.

######  Assumption 2.

We assume the density functions $p_{X}(\bm{x})$ and
$p_{\widetilde{X}}(\bm{x})$ satisfy

(29)
$A_{4}\exp(-\alpha\|\bm{x}\|^{2})<\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}<A_{4}\exp(\alpha\|\bm{x}\|^{2}),$

where $\alpha$ and $A_{4}$ are positive constants.

Lemma 4.4 and Theorems 4.6, 4.8 below will demonstrate that both
$\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ will approach the optimal constants in
this case, with a difference between the loss and the optimal constant of the
order $O(\varepsilon)$.

######  Lemma 4.4.

For arbitrarily small $\varepsilon>0$, there exists a random variable
$\widetilde{X}$ with the density function $p_{\widetilde{X}}(\bm{x})$ defined
in Eq. (26) such that

(30) $|p_{X}(\bm{x})-p_{\widetilde{X}}(\bm{x})|<C\varepsilon,\,\,\text{ for
}\,\,\bm{x}\in\Omega,$

where $\Omega\subseteq\mathbb{R}^{d}$ is a compact set and $C>0$ is a
constant.

######  Proof 4.5.

By Lemma 4.1 with $k=0$, for any $\bm{x}\in\Omega$ we have

(31)
$\|\bm{f}(\bm{x})-\bm{h}(\bm{x})\|=\sqrt{\sum_{i=1}^{d}|f_{i}(\bm{x})-h_{i}(\bm{x})|^{2}}<C\varepsilon.$

The standard normal random variable $Z\in\mathbb{R}^{d}$ has bounded gradient,
i.e., $\|\nabla p_{Z}(\bm{z})\|<C$ for any $\bm{z}\in\mathbb{R}^{d}$. Applying
the mean value theorem, we have

(32) $|p_{Z}(\bm{f}(\bm{x}))-p_{Z}(\bm{h}(\bm{x}))|=|\nabla
p_{Z}(\xi)\cdot(\bm{f}(\bm{x})-\bm{h}(\bm{x}))|\leq\|\nabla
p_{Z}(\xi)\|\|\bm{f}(\bm{x})-\bm{h}(\bm{x})\|<C\varepsilon,$

where $\xi$ is between $\bm{f}(\bm{x})$ and $\bm{h}(\bm{x})$, $\varepsilon$ is
given in Lemma 4.1 and $C>0$ is a constant. Because $Z$ has bounded density
function, we can exploit Eqs. (21), (23), (32) to derive

(33) $\displaystyle\left|p_{X}(\bm{x})-p_{\widetilde{X}}(\bm{x})\right|$
$\displaystyle=$ $\displaystyle\,\left|p_{Z}(\bm{f}(\bm{x}))|\det{\bf
J}_{\bm{f}}(\bm{x})|-p_{Z}(\bm{h}(\bm{x}))|\det{\bf
J}_{\bm{h}}(\bm{x})|\right|$ $\displaystyle=$
$\displaystyle\,\left|p_{Z}(\bm{f}(\bm{x}))(|\det{\bf
J}_{\bm{f}}(\bm{x})|-|\det{\bf J}_{\bm{h}}(\bm{x})|)+|\det{\bf
J}_{\bm{h}}(\bm{x})|(p_{Z}(\bm{f}(\bm{x}))-p_{Z}(\bm{h}(\bm{x})))\right|$
$\displaystyle\leq$
$\displaystyle\,\left|p_{Z}(\bm{f}(\bm{x}))\right|\left|\det{\bf
J}_{\bm{f}}(\bm{x})-\det{\bf J}_{\bm{h}}(\bm{x})\right|+\left|\det{\bf
J}_{\bm{h}}(\bm{x})\right|\left|p_{Z}(\bm{f}(\bm{x}))-p_{Z}(\bm{h}(\bm{x}))\right|<C\varepsilon$

for any $\bm{x}\in\Omega$ which concludes the proof.

Based on the results from Lemma 4.4, we have the convergence of the continuous
loss function $\mathcal{L}_{1}$ defined in Eq. (27).

######  Theorem 4.6.

For arbitrarily small $\varepsilon>0$, there exists a random variable
$\widetilde{X}$ with the density function $p_{\widetilde{X}}(\bm{x})$ in Eq.
(26) such that $\mathcal{L}_{1}$ has the bound

(34) $-\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log
p_{X}(\bm{x})dx\leq\mathcal{L}_{1}<-\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log
p_{X}(\bm{x})d\bm{x}+\varepsilon.$

######  Proof 4.7.

Using the inequality $\log t\leq t-1$ for all $t>0$, we have

(35) $\displaystyle-\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log
p_{X}(\bm{x})d\bm{x}-\mathcal{L}_{1}=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{\widetilde{X}}(\bm{x})}{P_{X}(\bm{x})}d\bm{x}$
$\displaystyle\leq$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left(\frac{p_{\widetilde{X}}(\bm{x})}{P_{X}(\bm{x})}-1\right)d\bm{x}$
$\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{\widetilde{X}}(\bm{x})-p_{X}(\bm{x})d\bm{x}=0.$

To prove the right side of Eq. (34), we need to show that

(36)
$\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}<\varepsilon.$

We define $B_{d}(r):=\\{\bm{x}\in\mathbb{R}^{d}:\|\bm{x}\|\leq r\\}$ and
$V_{d-1}(r)$ be the volume of $\\{\bm{x}\in\mathbb{R}^{d}:\|\bm{x}\|=r\\}$.
For every $\varepsilon_{1}<\min\\{A_{3}\exp(-\alpha),A_{3}\exp(-\alpha
K^{2})\\}$, $K$ is the constant in Assumption 1, we let
$K_{1}=\sqrt{-\frac{1}{\alpha}\log\frac{\varepsilon_{1}}{A_{3}}}$, then
$K_{1}>\max\\{1,K\\}$. According to Eq. (19) in Assumption 1,
$p_{X}(\bm{x})<A_{3}\exp(-\alpha\|\bm{x}\|^{2})\leq\varepsilon_{1}$ over the
domain $B_{d}(K_{1})^{c}$. Define
$\Omega_{\varepsilon_{1}}:=\\{\bm{x}:p_{X}(\bm{x})\geq\varepsilon_{1}\\}$,
then $\Omega_{\varepsilon_{1}}$ is contained in $B_{d}(K_{1})$. Since
$p_{X}(\bm{x})$ is continuous, $\Omega_{\varepsilon_{1}}$ is closed. Moreover,
as a bounded closed set in $\mathbb{R}^{d}$, $\Omega_{\varepsilon_{1}}$ is
compact. The KL divergence between distributions $p_{X}$ and
$p_{\widetilde{X}}$ has the following bound

(37) $\displaystyle D_{\rm KL}(p_{X}\,\|\,p_{\widetilde{X}})=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle\leq$
$\displaystyle\int_{p_{X}(\bm{x})\geq\varepsilon_{1}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle+\int_{\\{p_{X}(\bm{x})<\varepsilon_{1}\\}\cap
B_{d}(K_{1})}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle+\int_{\\{p_{X}(\bm{x})<\varepsilon_{1}\\}\cap
B_{d}(K_{1})^{c}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle:=$ $\displaystyle I_{1}+I_{2}+I_{3}.$

Next we investigate all terms on the right side of Eq. (37). By Lemma 4.4, for
every $\varepsilon_{2}>0$, there exists variable $\widetilde{X}$ such that

(38) $|p_{X}(\bm{x})-p_{\widetilde{X}}(\bm{x})|<C\varepsilon_{2},$

for any $\bm{x}\in\Omega_{\varepsilon_{1}}$. Choosing
$\varepsilon_{2}<\frac{\varepsilon_{1}}{2C}$, we have
$p_{\widetilde{X}}(\bm{x})>\varepsilon_{1}/2$ for the case
$p_{X}(\bm{x})\geq\varepsilon_{1}$. Applying the mean value theorem, $I_{1}$
has the bound

(39) $\displaystyle I_{1}=$
$\displaystyle\int_{p_{X}(\bm{x})\geq\varepsilon_{1}}p_{X}(\bm{x})\left(\log(p_{X}(\bm{x}))-\log(p_{\widetilde{X}}(\bm{x}))\right)d\bm{x}$
$\displaystyle\leq$
$\displaystyle\int_{p_{X}(\bm{x})\geq\varepsilon_{1}}p_{X}(\bm{x})\left|\frac{1}{{\xi}}\right|\left|p_{X}(\bm{x})-p_{\widetilde{X}}(\bm{x})\right|d\bm{x}$
$\displaystyle\leq$
$\displaystyle\frac{2C\varepsilon_{2}}{\varepsilon_{1}}\int_{p_{X}(\bm{x})\geq\varepsilon_{1}}p_{X}(\bm{x})d\bm{x}\leq\,\frac{2C\varepsilon_{2}}{\varepsilon_{1}},$

where $\left|\frac{1}{\xi}\right|\leq\frac{2}{\varepsilon_{1}}$ because $\xi$
is between $p_{X}(\bm{x})$ and $p_{\widetilde{X}}(\bm{x})$. Since we define
$\varepsilon_{2}<\frac{\varepsilon_{1}}{2C_{1}}$, we are able to let
$\varepsilon_{2}/\varepsilon_{1}\rightarrow 0$ to guarantee $I_{1}\rightarrow
0$.

Positive density functions $p_{X}(\bm{x})$, $p_{\widetilde{X}}(\bm{x})$ are
continuous, thus $\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}$ is
continuous and bounded over $B_{d}(K)$, i.e.,
$\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}<C_{1}$ for ${\bm{x}}\in
B_{d}(K)$. By Eq. (29) in Assumption 1, it follows that
$\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}<C_{2}K_{1}^{2}$ for
$\bm{x}\in B_{d}(K_{1})\setminus B_{d}(K)$. Hence,
$\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}<\max\\{C_{1},\
C_{2}K_{1}^{2}\\}$ for $\bm{x}\in B_{d}(K_{1})$. The term $I_{2}$ has the
bound

(40) $\displaystyle I_{2}$
$\displaystyle=\int_{\\{p_{X}(\bm{x})<\varepsilon_{1}\\}\cap
B_{d}(K_{1})}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle\leq\int_{\\{p_{X}(\bm{x})<\varepsilon_{1}\\}\cap
B_{d}(K_{1})}\left|p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}\right|d\bm{x}$
$\displaystyle\leq C\varepsilon_{1}\int_{B_{d}(K_{1})}\max\\{1,\
K_{1}^{2}\\}d\bm{x}$
$\displaystyle=C\varepsilon_{1}\frac{\pi^{d/2}}{\Gamma(d/2+1)}\max\\{K_{1}^{d},\
K_{1}^{d+2}\\}$ $\displaystyle\leq
C\max\\{\varepsilon_{1}(\log\frac{\varepsilon_{1}}{A_{1}})^{d/2},\
\varepsilon_{1}(\log\frac{\varepsilon_{1}}{A_{1}})^{d/2+1}\\},$

where $C$ is a constant and $I_{2}\rightarrow 0$ as
$\varepsilon_{1}\rightarrow 0$. The last inequality holds due to
$K_{1}=\sqrt{-\frac{1}{\alpha}\log\frac{\varepsilon_{1}}{A_{1}}}$ . Lastly we
have the bound for $I_{3}$

(41) $\displaystyle I_{3}$
$\displaystyle=\int_{\\{p_{X}(\bm{x})<\varepsilon_{1}\\}\cap
B_{d}(K_{1})^{c}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle=\int_{B_{d}(K_{1})^{c}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}d\bm{x}$
$\displaystyle\leq
C\int_{B_{d}(K_{1})^{c}}\exp(-\alpha\|\bm{x}\|^{2})\|\bm{x}\|^{2}d\bm{x}.$

According to the identity $\int_{B_{d}(R_{2})\setminus
B_{d}(R_{1})}y(\|\bm{x}\|)d\bm{x}=V_{d-1}(1)\int_{R_{1}}^{R_{2}}y(r)r^{d-1}dr$,
we have

(42) $\displaystyle I_{3}$ $\displaystyle\leq
CV_{d-1}(1)\int_{K_{1}}^{+\infty}\exp(-\alpha r^{2})r^{2}r^{d-1}dr$
$\displaystyle=C\int_{K_{1}}^{+\infty}\exp(-\alpha r^{2})r^{d+1}dr$
$\displaystyle=C\left(\left[\frac{\exp(-\alpha
r^{2})r^{d}}{-2\alpha}\right]_{K_{1}}^{+\infty}+\frac{d+1}{2\alpha}\int_{K_{1}}^{+\infty}\exp(-\alpha
r^{2})r^{d}dr\right)$ $\displaystyle=C\left(\left[\frac{\exp(-\alpha
r^{2})r^{d}}{-2\alpha}\right]_{K_{1}}^{+\infty}+\ldots+\frac{(d+1)!}{(2\alpha)^{d+1}}\int_{K_{1}}^{+\infty}\exp(-\alpha
r^{2})dr\right),$

where $I_{3}\rightarrow 0$ as $K_{1}\rightarrow+\infty$. Since
$K_{1}=\sqrt{-\frac{1}{\alpha}\log\frac{\varepsilon_{1}}{A_{1}}}$, we have
$I_{3}\rightarrow 0$ as $\varepsilon_{1}\rightarrow 0$.

In summary, Theorem 4.6 shows that the loss function $\mathcal{L}_{1}$
converges to the entropy of $p_{X}(\bm{x})$ as $\bm{h}$ converges to $\bm{f}$.
The following Theorem 4.8 proves the convergence of the continuous loss
function $\mathcal{L}_{2}$ as defined in Eq. (28).

######  Theorem 4.8.

For arbitrarily small $\varepsilon>0$, there exists a single-hidden-layer
neural network $\bm{g}$ such that

(43) $\mathcal{L}_{2}<\varepsilon,$

in any compact set $\Omega\subseteq\mathbb{R}^{d}$.

######  Proof 4.9.

By Assumption 1, $p_{X}$ is bounded. $\mathcal{L}_{2}$ has the following
bound:

(44) $\displaystyle\mathcal{L}_{2}=$
$\displaystyle\int_{\Omega}p_{X}(\bm{x})\left(\|\bm{g}(\bm{h}(\bm{x}))-\bm{x}\|^{2}d\bm{x}+\int_{\Omega}|\det{\bf
J}_{\bm{g}}(\bm{h}(\bm{x}))\det{\bf J}_{\bm{h}}(\bm{x})-1|\right)d\bm{x}$
$\displaystyle\leq$ $\displaystyle
C\int_{\Omega}\|\bm{g}(\bm{h}(\bm{x}))-\bm{x}\|^{2}d\bm{x}+\int_{\Omega}|\det{\bf
J}_{\bm{g}}(\bm{h}(\bm{x}))\det{\bf J}_{\bm{h}}(\bm{x})-1|d\bm{x}$
$\displaystyle=$ $\displaystyle
C\Big{(}\int_{\Omega}\|\bm{g}(\bm{h}(\bm{x}))-\bm{h}^{-1}(\bm{h}(\bm{x}))\|^{2}d\bm{x}$
$\displaystyle+\int_{\Omega}|\det{\bf J}_{\bm{g}}(\bm{h}(\bm{x}))(\det{\bf
J}_{\bm{h}^{-1}}(\bm{h}(\bm{x})))^{-1}-1|d\bm{x}\Big{)}$ $\displaystyle=$
$\displaystyle
C\Big{(}\int_{\bm{h}(\Omega)}\|\bm{g}(\bm{y})-\bm{h}^{-1}(\bm{y})\|^{2}|\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}$
$\displaystyle+\int_{{\bm{h}}(\Omega)}|\det{\bf J}_{\bm{g}}(\bm{y})-\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}\Big{)}:=C(I_{1}+I_{2}).$

Since $\bm{h}$ is continuous in a compact set $\Omega$, by the extreme value
theorem, $\bm{h}$ is bounded in $\Omega$, i.e., there exists a constant $M$
such that $|\bm{h}(\bm{x})|\leq M$ for all $\bm{x}\in\Omega$. Therefore,
$\bm{h}(\Omega)\subseteq B_{d}(M)$. Let $\bm{g}$ be the approximation of
$\bm{h}^{-1}$, and the term $I_{1}$ has the bound

(45) $\displaystyle I_{1}=$
$\displaystyle\int_{\bm{h}(\Omega)}\|\bm{g}(\bm{y})-\bm{h}^{-1}(\bm{y})\|^{2}|\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}$ $\displaystyle\leq$
$\displaystyle\int_{B_{d}(M)}\|\bm{g}(\bm{y})-\bm{h}^{-1}(\bm{y})\|^{2}|\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}$ $\displaystyle=$
$\displaystyle\int_{B_{d}(M)}\left(\sum_{i=1}^{d}|g_{i}(\bm{y})-h^{-1}_{i}(\bm{y})|^{2}\right)|\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}\leq\int_{B_{d}(M)}d\varepsilon_{1}^{2}Cd\bm{y}\leq
C\varepsilon_{1}^{2},$

where the last inequality holds by Lemma 4.1 with $k=0$ and the determinate of
Jacobian of ${\bm{h}}^{-1}$ has the upper bound in Eq. (24). Let
$\mathcal{S}_{d}$ be the set of all permutation of $\\{1,\ldots,d\\}$, we have

(46) $|\det{\bf J}_{\bm{g}}(\bm{y})-\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|=\left|\sum_{\sigma\in\mathcal{S}_{d}}\left({\rm
sgn}(\sigma)\left(\prod_{i=1}^{d}\frac{\partial g_{i}(\bm{y})}{\partial
x_{\sigma_{i}}}-\prod_{i=1}^{d}\frac{\partial h^{-1}_{i}(\bm{y})}{\partial
x_{\sigma_{i}}}\right)\right)\right|<C\varepsilon_{1}$

for any $\bm{y}\in B_{d}(M)$, where the last inequality holds by Lemma 4.1
with $k=1$. We have the bound for $I_{2}$

(47) $I_{2}=\int_{\bm{h}(\Omega)}|\det{\bf J}_{\bm{g}}(\bm{y})-\det{\bf
J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}\leq\int_{B_{d}(M)}|\det{\bf
J}_{\bm{g}}(\bm{y})-\det{\bf J}_{\bm{h}^{-1}}(\bm{y})|d\bm{y}\leq
C\varepsilon_{1}.$

The convergence of $\mathcal{L}_{2}$ to 0 as $\varepsilon_{1}$ approaches 0 is
established by Eqs. (44), (45) and (47). By choosing $\varepsilon_{1}$
sufficiently small, we complete the proof.

###  4.3 Convergence of the KL divergence

The analysis in Section 4.2 shows that the convergence of $\mathcal{L}_{1}$
and $\mathcal{L}_{2}$ is a necessary condition necessary condition to obtain
the convergent approximation discussed in Section 4.1, which verifies that
appropriateness of the loss function used to train the PR-NF model. Here we
show that the convergence of $\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ plus some
mild assumptions on the $\bm{h}$ and $\bm{g}$ is sufficient to ensure the
convergence of the KL divergence between $p_{X}$ and $p_{\widehat{X}}$.

######  Assumption 3.

We assume that the target random variable $X$ has finite second-order moment,
and there exists positive constant $A_{5}$ such that

(48) $\nabla p_{\widehat{X}}(\bm{x})\leq
A_{5}(\|\bm{x}\|+1)p_{\widehat{X}}(\bm{x}),$

where $p_{\widehat{X}}(\bm{x})$ is defined in Eq. (26).

######  Theorem 4.10.

For any given $\varepsilon>0$, under the assumptions of Theorems 4.6, 4.8 and
Assumption 3, the KL divergence between target variable $X$ and the
approximation $\widehat{X}$ defined in Eq. (25) satisfies $D_{\rm
KL}(p_{X}\|p_{\widehat{X}})<\varepsilon$.

######  Proof 4.11.

For arbitrarily small $\varepsilon_{1}>0$, under Theorems 4.6, we have

(49) $\displaystyle\mathcal{L}_{1}\leq-\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log
p_{X}(\bm{x})d\bm{x}+\varepsilon_{1},$

and for arbitrarily small $\varepsilon_{2}>0$, under Theorems 4.8, we have

(50)
$\mathcal{L}_{2}(\mathbb{R}^{d})=\int_{\mathbb{R}^{d}}p_{X}({\bm{x}})\left(\|{\bm{g}(\bm{h}(\bm{x}))-\bm{x}}\|^{2}+|\det{\bf
J}_{\bm{g}}(\bm{h}({\bm{x}}))\det{\bf
J}_{\bm{h}}({\bm{x}})-1|\right)d{\bm{x}}<\varepsilon_{2}.$

To simplify the notation, we denote $\bm{M}(\bm{x}):=\bm{g}(\bm{h}(\bm{x}))$,
Then Eq. (50) is rewritten as

(51)
${\mathcal{L}_{2}(\mathbb{R}^{d})=\int_{\mathbb{R}^{d}}p_{X}(\bm{x})(\|\bm{M}(\bm{x})-\bm{x}\|^{2}+|\det{\bf
J}_{\bm{M}}(\bm{x})-1|)d\bm{x}}<\varepsilon_{2}.$

Since
$\widehat{X}=\bm{g}(Z)=\bm{g}(\bm{h}(\widetilde{X}))=\bm{M}(\widetilde{X})$,
by the change of variables formula we have

(52) $p_{\widetilde{X}}(\bm{x})=p_{\widehat{X}}(\bm{M}(\bm{x}))|\det{\bf
J}_{\bm{M}}(\bm{x})|.$

The KL divergence between distributions $P_{X}$ and $P_{\widehat{X}}$ has the
following bound

(53) $\displaystyle D_{\rm KL}(p_{X}\|p_{\widehat{X}})\leq$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widehat{X}}(\bm{x})}d\bm{x}$
$\displaystyle\leq$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}dx+\int_{\mathbb{R}^{d}}p_{X}(x)\log\frac{p_{\widetilde{X}}(\bm{x})}{p_{\widehat{X}}(\bm{x})}d\bm{x}$
$\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}dx+\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{\widehat{X}}(\bm{M}(\bm{x}))|\det{\bf
J}_{\bm{M}}(\bm{x})|}{p_{\widehat{X}}(\bm{x})}d\bm{x}$ $\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{X}(\bm{x})}{p_{\widetilde{X}}(\bm{x})}dx+\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{\widehat{X}}(\bm{M}(\bm{x}))}{p_{\widehat{X}}(\bm{x})}d\bm{x}$
$\displaystyle+\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log|\det{\bf
J}_{\bm{M}}(\bm{x})|d\bm{x}:=I_{1}+I_{2}+I_{3}.$

According to Eq. (49), we have $I_{1}<\varepsilon_{1}$. Apply the mean value
theorem, $I_{2}$ has the bound

(54) $\displaystyle I_{2}=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log\frac{p_{\widehat{X}}(\bm{M}(\bm{x}))}{p_{\widehat{X}}(\bm{x})}d\bm{x}$
$\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\Big{(}\log(p_{\widehat{X}}(\bm{M}(\bm{x})))-\log(p_{\widehat{X}}(\bm{x}))\Big{)}d\bm{x}$
$\displaystyle=$
$\displaystyle\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left(\frac{\nabla
p_{\widehat{X}}(\bm{\xi})}{p_{\widehat{X}}(\bm{\xi})}\left\|\bm{M}(\bm{x})-\bm{x}\right\|\right)d\bm{x},$

where $\bm{\xi}$ is between $\bm{M}(\bm{x})$ and $\bm{x}$. By Eq. (48) in
Assumption 1, we have

(55) $\displaystyle I_{2}\leq$ $\displaystyle
C\int_{\mathbb{R}^{d}}p_{X}(\bm{x})(\|\bm{\xi}\|+1)\left\|\bm{M}(\bm{x})-\bm{x}\right\|d\bm{x}$
$\displaystyle\leq$ $\displaystyle
C\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left(\|\bm{M}(\bm{x})-\bm{x}\|+\|\bm{x}\|+1\right)\left\|\bm{M}(\bm{x})-\bm{x}\right\|d\bm{x}$
$\displaystyle=$ $\displaystyle
C\bigg{(}\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left\|\bm{M}(\bm{x})-\bm{x}\right\|^{2}d\bm{x}+\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\|\bm{x}\|\left\|\bm{M}(\bm{x})-\bm{x}\right\|d\bm{x}$
$\displaystyle+\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left\|\bm{M}(\bm{x})-\bm{x}\right\|d\bm{x}\bigg{)}$
$\displaystyle\leq$ $\displaystyle
C\bigg{[}\varepsilon_{2}+\left(\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\|\bm{x}\|^{2}d\bm{x}\right)^{\frac{1}{2}}\left(\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left\|\bm{M}(\bm{x})-\bm{x}\right\|^{2}d\bm{x}\right)^{\frac{1}{2}}$
$\displaystyle+\left(\int_{\mathbb{R}^{d}}p_{X}(\bm{x})d\bm{x}\right)^{\frac{1}{2}}\left(\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\left\|\bm{M}(\bm{x})-\bm{x}\right\|^{2}d\bm{x}\right)^{\frac{1}{2}}\bigg{]}\leq
C(\varepsilon_{2}+\varepsilon_{2}^{\frac{1}{2}}).$

Since $\log|t|\leq|t-1|$ for every $t\in\mathbb{R}$, $I_{3}$ has the bound

(56) $I_{3}=\int_{\mathbb{R}^{d}}p_{X}(\bm{x})\log|\det{\bf
J}_{\bm{M}}(\bm{x})|d\bm{x}\leq\int_{\mathbb{R}^{d}}p_{X}(\bm{x})|\det{\bf
J}_{\bm{M}}(\bm{x})-1|d\bm{x}\leq\varepsilon_{2}.$

Therefore,

(57) $D_{\rm
KL}(p_{X}\|p_{\widehat{X}})\leq\varepsilon_{1}+C(\varepsilon_{2}+\varepsilon_{2}^{\frac{1}{2}}).$

Let $\varepsilon_{1}<\varepsilon/2$ and
$\varepsilon_{2}<\min\\{\varepsilon^{2}/(16C^{2}),1\\}$, then we have

(58) $\begin{split}D_{\rm
KL}(p_{X}\|p_{\widehat{X}})\leq\varepsilon_{1}+C(2\varepsilon_{2}^{\frac{1}{2}})\leq\frac{\varepsilon}{2}+C\frac{\varepsilon}{2C}\leq\varepsilon.\end{split}$

Thus, $D_{\rm KL}(p_{X}\|p_{\widehat{X}})\leq\varepsilon$. We complete the
proof.

##  5 Numerical examples and applications

In this section we present numerical experiments demonstrating the superior
performance of the PR-NF model in comparison with the standard MC method. The
example in Sec.5.1 benchmarks the accuracy of our method for a problem with
known analytical ground-truth solution. In Sec. 5.2 we present an application
to plasma physics and in Sec.5.3 an application to fluid mechanics. Since in
these cases there are not known analytical solutions, the ground-truth
corresponds to MC simulations with sufficiently large samples. In all
numerical simulations, the PyTorch machine learning framework has been
implement with CUDA GPU. The PR-NF model uses the Adam optimizer [11] and the
Tanh activation function.

###  5.1 Verification of algorithm accuracy

In this subsection, we use a one-dimensional nonlinear SDE and a ten-
dimensional linear SDE to verify the accuracy of the proposed algorithm. We
first consider the following one-dimensional SDE

(59)
$X_{t}=X_{0}+\int_{0}^{t}{\bm{b}}(s,X_{s})ds+\int_{0}^{t}{\bm{\sigma}}(s,X_{s})dW_{s}\;\;\text{
with }X_{0}\in[0,L],$

where ${\bm{b}}(s,X_{s})=2\sqrt{X_{s}}+1$,
${\bm{\sigma}}(s,X_{s})=2\sqrt{X_{s}}$, $L=5$, $t=0.1$. The analytical
solution of this SDE is

(60) $X_{t}=(\sqrt{X_{0}}+t+W_{t})^{2},$

where $X_{0}$ is the initial condition, and the corresponding conditional
distribution is

(61) $p_{X_{t}|X_{0}}(x|x_{0})=\frac{1}{{2\sqrt{2\pi
tx}}}\left[\exp{\left(\frac{-(\sqrt{x}-\sqrt{x_{0}}-t)^{2}}{2t}\right)}+\exp{\left(\frac{-(\sqrt{x}+\sqrt{x_{0}}+t)^{2}}{2t}\right)}\right].$

The distribution of $X_{t}$ subject to an initial distribution
$p_{X_{0}}(x_{0})$ can be computed by the convolution with
$p_{X_{t}|X_{0}}(x|x_{0})$. Since the analytical solution is known, we use
this example to test the accuracy of the propose method.

The training dataset consists of $N_{\rm train}=20000$ samples with initial
positions $\\{X_{0}^{(n)}\\}_{n=1}^{N_{\rm train}}$ sampled from a uniform
distribution over the domain $\mathcal{D}=[0,5]$, and terminal positions
$\\{X_{t}^{(n)}\\}_{n=1}^{N_{\rm train}}$ generated by MC simulations of Eq.
(59). The neural network has $N_{\rm layer}=1$ hidden layer with $N_{\rm
neuron}=256$ neurons. The hyperparameter $\lambda$ in the loss function is
selected according to the method described in Sec. 3.3. The left panel of Fig.
3 shows the cross entropy $H(\lambda)$ defined in Eq. (16) with different
$\lambda$. As it shown, the PR-NF model attains the minimum of $H(\lambda$)
when $\lambda=50$. The middle panel and right panel of Fig. 3 correspond to
the accuracy performance of the proposed PR-NF model with two cases
$\lambda=1$ and $\lambda=50$, respectively. The red curve is the exact density
function and the histogram is synthesized by the PR-NF model. A good agreement
is shown in the right panel ($\lambda=50$). In contrast, $\lambda=1$ is too
small to guarantee the reversibility. Without the tuning process of $\lambda$,
the PR-NF model may fail to balance the $\mathcal{L}_{1}$ and
$\mathcal{L}_{2}$, e.g., the middle panel of Fig. 3. The minimization of cross
entropy $H(\lambda)$ in Section 3.3 is an effective and appropriate approach
to choose $\lambda$. An optimal parameter $\lambda$ is necessary to obtain an
accurate surrogate model.

|  |
---|---|---
Figure 3: The left panel shows the cross entropy $H(\lambda)$, which attains
the minimun when $\lambda=50$. The middle and right panels show the fitting
performance for $\lambda=1$ and $\lambda=50$, respectively. It is observed
that the accuracy of PR-NF model with $\lambda=50$ is better than the case
$\lambda=1$. Without the tuning process of $\lambda$, the PR-NF model may fail
to balance the $\mathcal{L}_{1}$ and $\mathcal{L}_{2}$ (e.g., $\lambda=1$
case). The minimization of the cross entropy $H(\lambda)$ is an effective and
correct approach to choose $\lambda$. The PR-NF model with an optimal
parameter $\lambda$ ($\lambda=50$ case) achieves a good fitting performance.

We now use the well-trained PR-NF model ($\lambda=50$) to sample $X_{t}$ with
variable initial distributions. The top row of Fig. 4 shows four different
initial distributions, denoted as “$\delta$ func”, “bar”, “sin2”, and
“ricker”. The bottom row of Fig. 4 shows corresponding histogram of $X_{t}$
generated by the PR-NF model, where a good agreement with the exact density is
observed. The results demonstrate that the well-trained PR-NF model can handle
different initial conditions without additional training process. Figure 5
shows the decay of the loss function of training dataset (left panel) and the
decay of the Kullback–Leibler (KL) divergence of the four initial
distributions (right panel). The KL divergence formula is defined as the
relative entropy from the approximate density (generated from PR-NF) $p_{\rm
approx}$ to the exact density $p_{\rm exact}$, i.e.,

(62) $D_{\rm KL}(p_{\rm exact}\,\|\,p_{\rm
approx})=\int_{-\infty}^{\infty}p_{\rm exact}(x)\log{\left(\frac{p_{\rm
exact}(x)}{p_{\rm approx}(x)}\right)}dx,$

where $D_{\rm KL}$ can be approximated by the Riemann sum over a uniform mesh
or the mean value of $\log{\left(\frac{p_{\rm exact}(x)}{p_{\rm
approx}(x)}\right)}$ over the training dataset $\mathcal{X}$ based on the
exact solution in Eq. (60). We monitor the KL divergence for different
distributions during the training process. The consistency of decays between
loss function and KL divergence illustrates that the loss function
$\mathcal{L}$ defined in Eq. (11) involving the reversibility error is
effective to be regarded as a loss function for normalzizing flow.

Figure 4: The accuracy performance of the well-trained PR-NF model for test
dataset. The top row shows the initial distributions of four dataset: $\delta$
func, bar, sin2, and ricker. The corresponding plots on the bottom row present
the fitting between the exact density of $X_{t}$ (red curve) and the histogram
(generated by PR-NF model) at $t=0.1$. The results demonstrate that the well-
trained PR-NF model can handle various initial conditions without re-training.
Figure 5: The decay of loss function of training dataset (left panel) and the
decay of KL divergence of four test dataset (right panel). We monitor the KL
divergence during the training process. The consistency of decays between loss
function and KL divergence demonstrates that the loss function $\mathcal{L}$
defined in Eq. (11) involving the reversibility error is effective to be
regarded as a loss function for normalzizing flow.

To illustrate the application of our method to high-dimensional problems,
consider the ten-dimensional SDE

(63) $X_{t}=X_{0}+\int_{0}^{t}X_{s}ds+\int_{0}^{t}KX_{s}dW_{s}\;\;\text{ with
}X_{0}\in\mathcal{D},$

where $t=1$, $\mathcal{D}=[0,1]^{10}$, $K=\frac{1}{2}\begin{pmatrix}1&1&&\\\
&\ddots&\ddots&\\\ &&\ddots&1\\\ &&&1\end{pmatrix}_{10\times 10}.$ In this
case the analytical solution $X_{t}$ is given by

(64) $X_{t}=\exp((\mathbb{I}_{10}-K^{2}/2)t+KW_{t}))X_{0}.$

The training dataset consists of $N_{\rm train}=20000$ samples where $X_{0}$
are sampled from an uniform initial distribution over the domain
$\mathcal{D}$. The neural network has $N_{\rm hidden}=1$ hidden layer with
$N_{\rm neuron}=256$ neurons. The left panel of Fig. 6 shows the cross entropy
as function of $\lambda$, with a minimum at $\lambda=100$.

We use the well-trained PR-NF model to test four initial distributions, one is
the normal distribution $x_{\rm
test}\sim\mathcal{N}(0.5,0.1\cdot\mathbb{I}_{10})$, other three are nonlinear
transformations of the normal distribution, square:= $x_{\rm test}^{2}$, log:=
$\ln(|x_{\rm test}|+1)$, and sin:= $\sin{(x_{\rm test}^{2})}$. The middle
panel of Fig. 6 shows the decay of loss function of the training dataset. And
the decay of KL divergence of four test dataset is shown in the right panel. A
good consistency is observed.

|
---|---
Figure 6: Numerical results for the ten-dimensional Brownian motion case. The
left panel shows the cross entropy $H(\lambda)$ with different $\lambda$,
where it attains the minimum when $\lambda=100$. Moreover, we evaluated
$H(\lambda)$ for $\lambda=1,2000$ (not shown). It is large when $\lambda=2000$
and is relative small when $\lambda=1$. However, the loss decay is noisy for
$\lambda=1$. The middle panel shows the decay of loss function and the right
panel shows the decay of KL divergence of four test dataset using the well-
trained PR-NF model. Since we are facing a ten-dimensional problem, the
density values for samples are very small. That leads to a relatively large KL
divergence that happens to the high-dimensional density estimation [19]. We
also tested the marginal KL divergence, which is on the order of $10^{-2}$.

###  5.2 Generation of runaway electrons in magnetically confined nuclear
fusion plasmas

Modeling and simulation of plasmas is an extremely complex computational
problem involving, among other aspects, the trajectories of the plasma
particles (electrons and ions), the evolution of external and self-generated
electromagnetic fields, as well as energy and particle sources and sinks.
Particle-based simulations (usually known as Monte Carlo simulations in the
plasma physics literature) are based on the solution of stochastic
differential equations of the form in Eq.(1) with $X_{t}$ denoting the
particles’ coordinates in the $d$-dimensional phase space. In the full-orbit
description $d=6$, but reduced models with $d<6$ are commonly used. The
deterministic term, ${\bm{b}}\,ds$, is given by the electromagnetic forces,
and the stochastic term, ${\bm{\sigma}\,dW_{s}}$, is used to model collisional
and diffusive processes. This is a strongly multi-scale problem since there is
a big disparity of spatiotemporal scales in the evolution of the orbits and
the electromagnetic fields. In addition, to avoid statistical sampling errors,
these equations need to be solved for very large number of particles and large
sets of different initial conditions. The goal of this section is to discuss
how the proposed PR-NF method can be used to overcome some of these
computational challenges in the case of magnetically confined high temperature
plasmas in toroidal geometry. This configuration (known as tokamak) is the
leading approach to achieve the elusive goal of controlled nuclear fusion for
electricity production. The specific problem we consider is the simulation of
relativistic runaway (RE) electrons produced during magnetic disruptions, see
for example Ref. [2] and references therein. Understanding this problem is
critical for the safe operation of nuclear fusion reactors.

The model of interest consists of the following set of SDEs

(65) $\left\\{\begin{aligned} dp&=\left[E\xi\,-\frac{\gamma
p}{\tau}(1-\xi^{2})-C_{F}(p)+\frac{1}{p^{2}}\frac{\partial}{\partial
p}\left(p^{2}C_{A}\right)\right]dt+\sqrt{2C_{A}}\,dW_{p},\\\
d\xi&=\left[\frac{E\left(1-\xi^{2}\right)}{p}+\frac{\xi(1-\xi^{2})}{\tau\gamma}-2\xi\frac{C_{B}}{p^{2}}\right]dt+\frac{\sqrt{2C_{B}}}{p}\,\sqrt{1-\xi^{2}}\,dW_{\xi}\,.\end{aligned}\right.$

This two-dimensional reduced model tracks the evolution of the magnitude of
the relativistic momentum, $p$, and $\xi=\cos\eta$, where $\eta$ is the pitch
angle between the particle velocity and the magnetic field which is assumed
fixed and given. $dW_{p}$ and $dW_{\xi}$ are independent standard Brownian
motions, $E$ is the electric field, and $\tau$ the radiation damping time
scale. The collision operator includes the momentum diffusion, $C_{A}$, pitch
angle scattering, $C_{B}$, and Coulomb drag, $C_{F}$, which are defined as

$C_{A}(p)=\bar{\nu}_{ee}\,\bar{v}_{T}^{2}\,\,\frac{\psi(y)}{y},\quad
C_{F}(p)=2\,\bar{\nu}_{ee}\,\bar{v}_{T}\,\psi(y),$
$C_{B}(p)=\frac{1}{2}\,\bar{\nu}_{ee}\,\bar{v}_{T}^{2}\,\,\frac{1}{y}\left[Z+\phi(y)-\psi(y)+\frac{y^{2}}{2}\delta^{4}\right],$
$\phi(y)=\frac{2}{\sqrt{\pi}}\int_{0}^{y}e^{-s^{2}}ds,\quad\psi(y)=\frac{1}{2y^{2}}\left[\phi(y)-y\frac{d\phi}{dy}\right],\quad
y=\frac{1}{\bar{v}_{T}}\frac{p}{\gamma},$
$\gamma=\sqrt{1+(\tilde{\delta}p)^{2}},\quad\tilde{\delta}=\frac{\tilde{v}_{T}}{c}=\sqrt{\frac{2\tilde{T}}{mc^{2}}},\quad{\delta}=\frac{\hat{v}_{T}}{c}=\sqrt{\frac{2\hat{T}_{f}}{mc^{2}}},\quad\bar{v}_{T}=\sqrt{\frac{\hat{T}_{f}}{\tilde{T}}},$
$\bar{\nu}_{ee}=\left(\frac{\tilde{T}}{\hat{T}_{f}}\right)^{3/2}\frac{\rm{ln\hat{\Lambda}}}{\rm{ln\tilde{\Lambda}}},\quad{\rm
ln}\tilde{\Lambda}=14.9-\frac{1}{2}{\rm ln}0.28+{\rm ln}\tilde{T},\quad{\rm
ln}\hat{\Lambda}=14.9-\frac{1}{2}{\rm ln}0.28+{\rm ln}\hat{T}_{f},$

with $Z$ and $c$ denoting the ion effective charge and the speed of light,
respectively. Further details on this model can be found in Ref. [4, 25] and
references therein.

We are interested in the hot-tail acceleration of electrons due to the rapid
cooling of the plasma (thermal quench) during a magnetic disruption [23, 17].
We use a fast cooling model for the temperature and an electric field from
Ohm’s law with a Spitzer resistivity temperature dependent model
$E=E_{0}\left[\frac{\tilde{T}}{\hat{T}_{f}}\right]^{3/2}$, with typical model
parameters $Z=1$, $\tau=6\times 10^{3}$, $E_{0}=1/2000$, $\tilde{T}=3$,
$\hat{T}_{f}=0.05$ and $mc^{2}=500$. The integration domain is
$p\in(p_{\min},p_{\max})$ and $\theta\in(0,\pi)$, where
$(p_{\min},p_{\max})=(0.5,5)$, and $t_{\max}=26$. The different cases of
initial conditions are obtained by sampling a family of Maxwellian
distributions

(66)
$f(p,\xi,t_{0})=\frac{2p^{2}}{\pi^{1/2}p_{0}^{3}}e^{-(p/p_{0})^{2}},\quad\text{where}\,\,p_{0}=\sqrt{\frac{\hat{T}_{0}}{\tilde{T}}},$

normalized as
$\int_{0}^{1}\int_{-1}^{1}\int_{p_{\min}}^{p_{\max}}f_{0}(p,\xi)dpd\xi=1$, and
parameterized by $\hat{T}_{0}$. Exploration of the hot-tail generation of RE
in this case requires the solution of this problem for different values of
$\hat{T}_{0}$. In the standard Monte Carlo approach this is a computational
intense problem because each time $\hat{T}_{0}$ changes the problem needs to
be recomputed for the new initial condition. However, as it will be shown
bellow, this is not the case with the proposed PR-NF method.

In this problem the training dataset $\mathcal{V}$ consists of $N_{\rm
train}=30000$ samples with initial positions, $\\{(p_{0},\xi_{0})^{(n)}\\}$,
$n=1,\ldots,N_{\rm train}$, uniformly distributed in the integration domain
and terminal positions, $(p_{t_{\max}},\xi_{t_{\max}})^{(n)}$, obtained by a
direct Monte Carlo simulation of Eqs. (65). The neural network architecture of
the PR-NF model uses $N_{\rm hidden}=1$ hidden layer and $N_{\rm neuron}=512$
neurons. To choose the hyperparameter $\lambda$, which controls the relative
importance of the pseudo-reversivility and the negative log-likelihood losses,
we follow the same procedure as in Section 5.1. Namely, we select the value of
$\lambda$ for which the cross entropy, $H(\lambda)$, has a minimum, which in
this case corresponds to $\lambda=100$.

To check the accuracy of the PR-NF model we consider a test set of $N_{\rm
test}=20000$ samples drawn from the Maxwellian distribution in Eq. (66) with
$\hat{T}_{0}=10$. Fig. 7 compares the final probability density function (pdf)
at $t_{\rm max}=26$ computed with the PR-NF method and the direct Monte Carlo
simulation. To construct the pdf from the $N_{\rm test}$ data points, we use a
Gaussian kernel density estimation. The plots on the top row show the one-
dimensional marginal pdfs for the pitch angle $\theta$ (left) and momentum $p$
(right). The plots on the bottom row show the two-dimensional contour plots of
the $\log_{10}$ of the pdfs for the PR-NF (left) and the MC (right)
simulations. The local pointwise differences between PR-NF and MC methods are
of the order $10^{-2}$, which implies a small error in the evaluation of
quantities of interest like the production of runaway electrons shown in Fig.
8.

Figure 7: Final state of runaway electrons probability density function,
$f(\theta,p,t_{max})$, according to the model in Eq. (65) with initial
condition, $f(\theta,p,t_{0})$ in Eq. (66) with $\hat{T}_{0}=10$. The top row
compares the $1d$ marginal distributions for the pitch angle (left) and the
momentum (right) computed using the PR-NF and the MC methods. The inset in the
the top plot on the right highlights, in logarithmic scale, the bump in the
distribution resulting from the acceleration of the hot tail. The contour
plots at the bottom compare the $2d$ distributions using a $\log_{10}$ scale.
The local pointwise differences between PR-NF and MC methods are of the order
$10^{-2}$. Such agreement between MC and the PRNF methods is sufficient to
ensure the accuracy of quantity of interest, for example, the runaway electron
production rate in Fig. 8.

To illustrate the performance of the PR-NF method in the computation of the
final state for different initial conditions, we consider the quantity of
interest

(67) $n_{\rm
RE}=\int_{-1}^{1}\int_{p^{*}}^{p_{\max}}f_{t_{\max}}(p,\xi)dpd\xi,$

representing the total fraction of runaway electrons produced by the hot-tail
mechanism during the thermal quench, where $f_{t_{\max}}$ is the density
function of electrons at $t_{\max}=26$, $p^{*}=1.75$ is the RE energy
threshold and $p_{\max}$ is a numerical upper bound. We consider a family of
Maxwellian initial conditions of the form in Eq. (66) with
$\hat{T}_{0}=1,2,3,4,5,6,7,8,9,10.$ Figure 8 compares $n_{\rm RE}$ computed
using a direct MC simulation that requires to solve the whole problem from the
start for each value of $\hat{T}_{0}$, with the PR-NF method that reduces the
computation to the direct evaluation of the final state using the trained
neural network. This numerical simulation is based on $N_{\rm sample}=20000$
samples for different $\hat{T}_{0}$. By using a single PR-NF model, we can
obtain the orbits of particles for ten initial conditions without any
additional simulations or training. However, the Monte Carlo method has to
repeatedly simulate the particles’ orbits for ten times. For the simulation,
the time step size of the MC method is $\Delta t=0.001$ and the number of
epochs of the PR-NF method is $N_{\rm epoch}=20000$. We record the running
wall-clock time (measured using the python function “time.time()”) between
these two methods. For the results in Fig. 8, the error between the PR-NF and
MC method (viewed as the ground-truth) is acceptable. The accuracy of PR-NF
model is enough to predict and analyze the profile of fraction runaway
electrons production rate $n_{\rm RE}$. Moreover, the total running time of
the MC method of these ten cases of Maxwellian distribution is around $C_{\rm
MC}=4000$ sec. The total running time of the PR-NF method consists of two
parts, the offline cost is around $C_{\rm offline}=1354$ sec and the online
cost is around $C_{\rm online}=12$ sec. The PR-NF model is faster that the MC
method, especially on the online cost. This efficiency advantage is more
valuable when we handle a large number of Maxwellian distributions.

Figure 8:  Quantity of interest $n_{\rm RE}$ in Eq. (67) represents the
production of runaway electrons as function of the parameter $\hat{T}_{0}$ in
Eq. (66). The red curve shows the result using the PR-NF surrogate model and
the blue curve the result using the direct MC method. A single PR-NF model can
compute $n_{\rm RE}$ for ten initial conditions without any additional
simulations or training, where the online cost is much cheaper than MC method.
The accuracy of PR-NF model is sufficient to predict the profile of the
production rate $n_{\rm RE}$ as the MC method (around 1% absolute error is
acceptable here).

###  5.3 Passive scalar advection-diffusion transport in a 3D chaotic flow

Understanding passive scalar transport is a problem of significant interest in
fluid dynamics in general and environmental engineering, oceanography, and
atmospheric sciences in particular. By passive we mean that the transported
scalar exactly follows the given velocity field without modifying it.

As a specific example to illustrate the application of the PR-NF method to
this problem, we consider the ABC (Arnold-Beltrami-Childress) flow [6] which
is a three-dimensional incompressible velocity field which is an exact
solution of Euler’s equation. In Cartesian coordinates this velocity field has
components ${\bf v}=(v_{x},v_{y},v_{z})$ with $v_{x}=A\sin{z}+C\cos{y}$,
$v_{y}=B\sin{x}+A\cos{z}$ and $v_{z}=C\sin{y}+B\cos{x}$. For a wide range of
the $A$, $B$ and $C$ parameters, this flow is not integrable. That is, it
exhibits chaotic advection and the sensitive dependence on initial condition
gives rise to rapid mixing and long-range transport [21].

In addition to advection, passive scalar transport is affected by diffusion.
In particle based simulations this important effect can be modeled by adding a
stochastic term. Following this approach we consider the following set of SDEs

(68) $\left\\{\begin{aligned}
dx&=P_{e}\,[A\sin{z}+C\cos{y}]\,dt+\,dW_{x}\,,\\\
dy&=P_{e}\,[B\sin{x}+A\cos{z}]\,dt+\,dW_{y}\,,\\\
dz&=P_{e}\,[C\sin{y}+B\cos{x}]\,dt+\,dW_{z}\,.\end{aligned}\right.$

where $dW_{x}$, $dW_{y}$ and $dW_{z}$ are independent Brownian motions, and
the Peclé number, $P_{e}$ is a dimensionless parameter controlling the
relative importance of advection and diffusion. When, $P_{e}\ll 1$, transport
is dominated by diffusion and when, $P_{e}\gg 1$, transport is dominated by
advection.

The initial conditions,
$\\{x_{0}^{(n)},y_{0}^{(n)},z_{0}^{(n)}\\}_{n=1}^{N_{\rm train}}$, of the
training dataset consists of an ensemble of $N_{\rm train}=30000$ random
points uniformly distributed in the cubic domain
$\mathcal{D}=[0,2\pi]\times[0,2\pi]\times[0,2\pi]$. The terminal positions,
$\\{x_{t_{\max}}^{(n)},y_{t_{\max}}^{(n)},z_{t_{\max}}^{(n)}\\}_{n=1}^{N_{\rm
train}}$, of the training dataset are obtained by a direct Monte Carlo
simulation of Eq. (68). In the calculation presented we use $t_{\rm max}=2$,
and model parameters $P_{e}=3$, $A=1$, $B=1$, and $C=0.25$. For the PR-NF
neural network model we use $N_{\rm hidden}=1$ hidden layer with $N_{\rm
neuron}=512$ neurons. Once the training process is done, the PR-NF model can
be used as a surrogate model to predict the final state for various initial
conditions, without the need to integrate Eqs. (68).

As a first test, we apply the PR-NF to predict the final state of an initial
condition of the form

(69)
$f(x,y,z,t_{0})=H(x,y,z)\exp{\left[-\left(\frac{x-x_{c}}{\sigma_{x}}\right)^{2}-\left(\frac{y-y_{c}}{\sigma_{y}}\right)^{2}-\left(\frac{z-z_{c}}{\sigma_{z}}\right)^{2}\right]},$

modeling, for example, a Gaussian-shaped cloud of a pollutant, with
$\sigma_{x}=\pi/3$, $\sigma_{x}=\pi/5$, $\sigma_{x}=\pi/4$, localized at
$(x_{c},y_{c},z_{c})\in\mathcal{D}$, where $H(x,y,z)=1$ if
$(x,y,z)\in\mathcal{D}$ and $H(x,y,z)=0$ otherwise. Figure 9 shows a good
agreement between the final pdf of the passive scalar, $f(x,y,z,t_{\rm max})$,
computed with the PR-NF surrogate model and the pdf computed using direct
Monte Carlo simulation for the case $(x_{c},y_{c},z_{c})=(\pi/2,\pi,\pi)$.
Since both methods are particle-based, once the discrete particle data have
been obtained, the corresponding pdfs are constructed using a standard
Gaussian-kernel estimation. To ease the comparison of these $3d$ pdfs, Fig. 9
shows the $1d$ marginal distributions, e.g., $g(x)=\int dy\int
dzf(x,y,z,t_{\rm max})$, and the $2d$ marginal distributions, e.g.,
$g(x,y)=\int dzf(x,y,z,t_{\rm max})$. The mixing due to the combination of
chaotic advection and diffusion rapidly displaces the scalar outside the
$\mathcal{D}$ domain.

Figure 9:  Final state of the probability density function,
$f(x,y,z,t_{max})$, of the scalar resulting from the advection-diffusion in
the ABC fluid velocity model with diffusion in Eq. (68) for the initial
condition, $f(x,y,z,t_{0})$, in Eq. (69). Since the pdf is $3d$, to ease the
comparison between the PR-NF and the MC methods, we show the $1d$ marginal
distributions in the top row and the $2d$ marginal distribution if the bottom
rows. The controur plots in one column share the same level, and the scale is
of the order $10^{-3}$. The agreement between the two methods is sufficient to
ensure the accuracy of the quantity of interest shown in Fig. 10.

As a second application we consider a “target” transport problem of interest,
for example, to environmental fluid dynamics. The problem consists of finding
the concentration of the scalar in a target domain, ${\cal
T}=[x_{min},x_{max}]\times[-\infty,\infty]\times[z_{min},z_{max}]$, resulting
from the transport of a concentrated pollutant cloud released inside the
$\mathcal{D}=[0,2\pi]\times[0,2\pi]\times[0,2\pi]$ domain. The initial cloud
will be described using the Gaussian model in Eq. (69) with
$\sigma_{x}=\pi/3$, $\sigma_{x}=\pi/5$, $\sigma_{x}=\pi/4$. The $y$-coordinate
of the cloud will be fixed at $y_{c}=\pi$ and the goal is to compute the total
scalar density in ${\cal T}$ as function of the $(x_{c},z_{c})$ coordinates of
the center of the initial cloud in ${\cal D}$. The quantity of interest in
this case is

(70) $n_{{\cal
T}}(x_{c},z_{c})=\int_{x_{\min}}^{x_{\max}}\int_{-\infty}^{\infty}\int_{z_{\min}}^{z_{\max}}f(x,y,z,t_{\rm
max})dxdydz\,,$

and for the calculation presented here we will use $x_{\min}=0$,
$x_{\max}=\pi$, $z_{\min}=2\pi$, and $z_{\max}=3\pi$.

To perform this calculation a $N_{x}\times N_{z}$ uniform grid is constructed
in the $(x_{c},z_{c})\in[0,2\pi]\times[0,2\pi]$ space and the
$N_{x}N_{z}=N_{\rm ic}$ grid nodes are used to define the ensemble
$\\{x^{i}_{c},z^{i}_{c}\\}_{i=1}^{N_{\rm ic}}$. For each element of the
ensemble, the standard direct Monte Carlo approach requires the solution of
the system of SDE in Eq. (68) for $N_{\rm test}$ particles with initial
conditions sampled from the Gaussian cloud in Eq. (69) centered at
$(x^{i}_{c},y_{c}=\pi,z^{i}_{c})$. In total, this approach requires $N_{\rm
ic}N_{\rm test}$ solutions of the SDEs. On the other hand, the PR-NF method
only requires to solve the SDE in Eq. (68) for $N_{\rm train}$ initial
conditions and, once the training is done, the computation of the $N_{\rm
ic}N_{\rm test}$ initial conditions can be done by simply evaluating the
surrogate neural network model.

In the case discussed here we use $N_{x}=N_{z}=21$, and $N_{\rm test}=20000$.
With the Monte Carlo method this requires solving $8.82\times 10^{6}$ SDEs.
For $t_{\rm final}=2$, with step size $\Delta t=0.001$, the total run-time to
produce the results reported in Fig. 10 was $C_{\rm MC}=5320$ sec. On the
other hand, the training of the PR-NF neural network, like in the previous
calculation, is done using $N_{\rm train}=30000$, $N_{\rm hidden}=1$ hidden
layer with $N_{\rm neuron}=512$ neurons, and $N_{\rm epoch}=20000$. For these
parameters, the offline cost of the training is around $C_{\rm offline}=2200$
sec. But, once the training is done, the evaluation of the different initial
conditions has minimal cost. In particular, the total PR-NF run-time to
produce the results reported in Fig. 10 is of the order $C_{\rm online}=50$
sec. The contour plots of $n_{\cal T}$ shown in Fig. 10 provided evidence that
the much more efficient PR-NF surrogate model reproduces well the results
obtained with the direct, time-consuming, MC method for the target transport
problem.

Figure 10: Contour plots of quantity of interest $n_{\tau}(x_{c},z_{c})$ in
Eq. (70) computed using the PR-NF surrogate model (left panel) and the direct
MC method (right panel). For a given $(x_{c},z_{c})$, $n_{\tau}(x_{c},z_{c})$
denotes the total density of the scalar in the target region ${\cal
T}(x,y,z)=[0,\pi]\times[-\infty,\infty]\times[2\pi,3\pi]$ resulting from the
transport of an initial Gaussian cloud, Eq. (69), centered at
$(x_{c},y_{c}=\pi,z_{c})\in\mathcal{D}(x,y,z)=[0,2\pi]\times[0,2\pi]\times[0,2\pi]$.
The agreement of the two calculation provides evidence of the accuracy of the
significantly faster PR-NF method (absolute error between two methods in
$n_{\tau}$ is of the order $10^{-2}$). The results of PR-NF (left panel) is
sufficient to simulate the profile of $n_{\tau}$. Moreover, the online cost of
PR-NF method is about 100 times faster than MC method.

##  6 Conclusion

We proposed an accurate and efficient normalizing-flow algorithm for the
solution of stochastic differential equations with arbitrary initial
conditions. Our approach leverages the pseudo-reversible normalizing flow
neural network framework. The novelty of our normalizing flow approach can be
viewed from two perspectives. Firstly, it enables learning the distribution of
the final state based on any given initial state. This allows the model to be
trained just once, and then applied to handle various initial distributions
effectively. Secondly, we develop a pseudo-reversible neural network
architecture that relaxes the strict reversibility constraint. This feature
simplifies the implementation of feed-forward neural networks, which model
both forward and inverse flows.

Our subsequent objective is to expand the applicability of the proposed method
to tackle more complex problems. For example, the two-dimensional runaway
electron model discussed in Section 5.2 is a reduced 2-D plasma transport
model. Going beyond this simplified description, we envision extending the
capabilities of the PR-NF model to encompass the six-dimensional full-orbit
transport model, which poses significant challenges when employing PDE-based
methods. Additionally, we aim to explore the complete dynamics of stochastic
differential equations, not limited to a single time instant.

## Acknowledgement

This material is based upon work supported by the U.S. Department of Energy,
Office of Science, Office of Advanced Scientific Computing Research, Applied
Mathematics program under the contract ERKJ387, Office of Fusion Energy
Science, and Scientific Discovery through Advanced Computing (SciDAC) program,
at the Oak Ridge National Laboratory, which is operated by UT-Battelle, LLC,
for the U.S. Department of Energy under Contract DE-AC05-00OR22725.

## References

  * [1] G.-J. Both and R. Kusters, Temporal normalizing flows, arXiv preprint arXiv:1912.09092, (2019).
  * [2] B. N. Breizman, P. Aleynikov, E. M. Hollmann, and M. Lehnen, Physics of runaway electrons in tokamaks, Nuclear Fusion, 59 (2019), p. 083001.
  * [3] A. Creswell, T. White, V. Dumoulin, K. Arulkumaran, B. Sengupta, and A. A. Bharath, Generative adversarial networks: An overviewa, IEEE signal processing magazine, 35 (2018), pp. 53–65.
  * [4] D. Del-Castillo-Negrete, M. Yang, M. Beidler, and G. Zhang, Generation and mitigation of runaway electrons: spatio-temporal effects in dynamic scenarios, tech. report, Oak Ridge National Lab.(ORNL), Oak Ridge, TN (United States), 2021.
  * [5] L. Dinh, J. Sohl-Dickstein, and S. Bengio, Density estimation using real nvp, arXiv preprint arXiv:1605.08803, (2016).
  * [6] T. Dombre, U. Frisch, J. M. Greene, M. Hénon, A. Mehr, and A. M. Soward, Chaotic streamlines in the abc flows, Journal of Fluid Mechanics, 167 (1986), pp. 353–391.
  * [7] L. Feng, Xiaodong Zeng and T. Zhou, Solving time dependent fokker-planck equations via temporal normalizing flow, Communications in Computational Physics, 32 (2022), pp. 401–423.
  * [8] W. Grathwohl, R. T. Chen, J. Bettencourt, I. Sutskever, and D. Duvenaud, Ffjord: Free-form continuous dynamics for scalable reversible generative models, arXiv preprint arXiv:1810.01367, (2018).
  * [9] L. Guo, H. Wu, and T. Zhou, Normalizing field flows: Solving forward and inverse stochastic differential equations using physics-informed flow models, Journal of Computational Physics, 461 (2022), p. 111202.
  * [10] T. A. Keller, J. W. Peters, P. Jaini, E. Hoogeboom, P. Forré, and M. Welling, Self normalizing flows, in International Conference on Machine Learning, PMLR, 2021, pp. 5378–5387.
  * [11] D. P. Kingma and J. Ba, Adam: A method for stochastic optimization, arXiv preprint arXiv:1412.6980, (2014).
  * [12] I. Kobyzev, S. J. Prince, and M. A. Brubaker, Normalizing flows: An introduction and review of current methods, IEEE transactions on pattern analysis and machine intelligence, 43 (2020), pp. 3964–3979.
  * [13] Y. Lu, R. Maulik, T. Gao, F. Dietrich, I. G. Kevrekidis, and J. Duan, Learning the temporal evolution of multivariate densities via normalizing flows, Chaos: An Interdisciplinary Journal of Nonlinear Science, 32 (2022), p. 033121.
  * [14] B. Øksendal and B. Øksendal, Stochastic differential equations, Springer, 2003.
  * [15] G. Papamakarios, E. Nalisnick, D. J. Rezende, S. Mohamed, and B. Lakshminarayanan, Normalizing flows for probabilistic modeling and inference, The Journal of Machine Learning Research, 22 (2021), pp. 2617–2680.
  * [16] G. Papamakarios, T. Pavlakou, and I. Murray, Masked autoregressive flow for density estimation, Advances in neural information processing systems, 30 (2017).
  * [17] Y. V. Petrov, P. Parks, and R. Harvey, Numerical simulation of the hot-tail runaway electron production mechanism using cql3d and comparison with smith–verwichte analytical model, Plasma Physics and Controlled Fusion, 63 (2021), p. 035026.
  * [18] A. Pinkus, Approximation theory of the mlp model in neural networks, Acta numerica, 8 (1999), pp. 143–195.
  * [19] D. Rezende and S. Mohamed, Variational inference with normalizing flows, in International conference on machine learning, PMLR, 2015, pp. 1530–1538.
  * [20] O. Rippel and R. P. Adams, High-dimensional probability estimation with deep density models, arXiv preprint arXiv:1302.5125, (2013).
  * [21] M. F. Shlesinger, G. M. Zaslavsky, and J. Klafter, Strange kinetics, Nature, 363 (1993), pp. 31–37.
  * [22] R. W. Shonkwiler and F. Mendivil, Explorations in monte carlo methods, Springer Science & Business Media, 2009.
  * [23] A. Stahl, O. Embréus, G. Papp, M. Landreman, and T. Fülöp, Kinetic modelling of runaway electrons in dynamic scenarios, Nuclear Fusion, 56 (2016), p. 112009.
  * [24] Y. Teng, Z. Wang, L. Ju, A. Gruber, and G. Zhang, Level set learning with pseudo-reversible neural networks for nonlinear dimension reduction in function approximation, arXiv preprint arXiv:2112.01438, (2021).
  * [25] M. Yang, G. Zhang, D. del Castillo-Negrete, and M. Stoyanov, A feynman-kac based numerical method for the exit time probability of a class of transport problems, Journal of Computational Physics, 444 (2021), p. 110564.