Energetic Variational Neural Network Discretizations of Gradient Flows

In this subsection, we provide a detailed derivation of underlying PDEs for $L^{2}$-gradient flows and generalized diffusions by the general framework of EnVarA. We refer interested readers to [24, 74] for a more comprehensive review of the EnVarA. In both systems, the kinetic energy $\mathcal{K}=0$.

In this subsection, we briefly review some existing neural network-based algorithms for solving PDEs. We refer interested readers to [18, 52] for detailed reviews.

Considering a PDE subject to a certain boundary condition

and assuming the solution can be approximated by a neural network $\varphi_{\rm NN}(\bm{x};{\bm{\Theta}})$, where ${\bm{\Theta}}$ is the set of parameters in the neural network, the Physics-Informed Neural Network (PINN) [62] and similar methods [13, 68, 77] find the optimal parameter ${\bm{\Theta}}^{*}$ by minimizing a loss function defined as:

Here $\{\bm{x}_{i}\}_{i=1}^{N_{in}}$ and $\{\bm{s}_{j}\}_{j=1}^{N_{b}}$ are sets of samples in $\Omega$ and $\partial\Omega$ respectively, which can be drawn uniformly or by following some prescribed distributions. The parameter $\lambda$ is used to weight the sum in the loss function. Minimizers of the loss function (15) can be obtained by using some optimization methods, such as AdaGrad and Adam. Of course, the objective function of this minimization problem in general is not convex even when the initial problem is. Obtaining the global minimum of (15) is highly non-trivial.

In contrast to the PINN, which is based on the strong form of a PDE, the deep Ritz method (DRM) [21] is designed for solving PDEs using their variational formulations. If Eq. (14) is an Euler-Lagrangian equation of some energy functional

then the loss function can be defined directly by

The last term in (16) is a penalty term for the boundary condition. The original Dirichlet boundary condition can be recovered with $\lambda\rightarrow\infty$. Again, the samples $\{\bm{x}_{i}\}_{i=1}^{N_{in}}$ and $\{\bm{s}_{j}\}_{j=1}^{N_{b}}$ can be uniformly sampled from $\Omega$ and $\partial\Omega$ or sampled following some other prescribed distributions, respectively. Additionally, both the variational PINN [37] and the WAN [84] utilized the Galerkin formulation to solve PDEs. The variational PINN [37], stems from the Petrov–Galerkin method, represents the solution via a DNN, and keeps test functions belonging to linear function spaces. The WAN employs the primal and adversarial networks to parameterize the weak solution and test functions respectively [84].

The neural network-based algorithms mentioned above focus on elliptic equations. For an evolution equation of the form:

with a suitable boundary condition, the predominant approach is to treat the time variable as an additional dimension, and the PINN type techniques can be applied. However, these approaches are often expensive and may fail to capture the dynamics of these systems. Unlike spatial variables, there exists an inherent order in time, where the solution at time $T$ is determined by the solutions in $t<T$. It is difficult to generate samples in time to maintain this causality. Very recently, new efforts have been made in using NNs to solve evolution PDEs [8, 16]. The idea of these approaches is to use NNs with time-dependent parameters to represent the solutions. For instance, Neural Galerkin method, proposed in [8], parameterizes the solution as $\varphi_{h}(\bm{x};{\bm{\Theta}}(t))$ with ${\bm{\Theta}}(t)$ being the NN parameters, and defines a loss function in terms of ${\bm{\Theta}}$ and $\dot{\bm{\Theta}}$ through a residual function, i.e.,

where $\nu_{\Theta}(\bm{x})$ is a suitable measure which might depend on ${\bm{\Theta}}$. By taking variation of $J({\bm{\Theta}},{\bm{\eta}})$ with respect to ${\bm{\eta}}$, the neural Galerkin method arrives at an ODE of ${\bm{\Theta}}(t)$:

where ${\sf M}({\bm{\Theta}})=\int\nabla_{\bm{\Theta}}\varphi_{h}\otimes\nabla_{\bm{\Theta}}\varphi_{h}\mathrm{d}\nu_{\bm{\Theta}}(\bm{x}),$ and $F(t,{\bm{\Theta}})=\int\nabla_{\bm{\Theta}}\varphi_{h}F(\bm{x},t,\varphi_{h})\mathrm{d}\nu_{\bm{\Theta}}(\bm{x}).$ The ODE (20) can be solved by standard explicit or implicit time-marching schemes.

It’s important to note that, with the exception of DRM, all existing neural-network-based methods are developed based on either the strong or weak forms of PDEs. While DRM utilizes the free energy functional, it is only suitable for solving static problems, i.e., finding the equilibrium of the system. Additionally, most of these existing approaches are Eulerian methods. For certain types of PDEs, like generalized diffusions, these methods may fail to preserve physical constraints, such as positivity and the conservation of mass of a probability function.

Before we discuss the neural network discretization, we first briefly review the discrete energetic variational approach proposed in [50]. Given a continuous energy-dissipation law (1), the discrete energetic variational approach first constructs a finite-dimensional approximation to a continuous energy-dissipation law by introducing a spatial discretization of the state variable $\varphi$, denoted by $\varphi_{h}(\bm{x};{\bm{\Theta}}(t))$, where ${\bm{\Theta}}(t)\in\mathbb{R}^{K}$ is the parameter to be determined. By replacing $\varphi$ by $\varphi_{h}$, one obtain a semi-discrete energy-dissipation law (here we assume $\mathcal{K}$ in the continuous model for simplicity) in terms of ${\bm{\Theta}}$:

where $\mathcal{F}_{h}\left[{\bm{\Theta}}(t)\right]=\mathcal{F}[\varphi_{h}(\bm{x};{\bm{\Theta}})]$, and $\triangle_{h}[{\bm{\Theta}}(t),{\bm{\Theta}}^{\prime}(t)]=\int\eta(\varphi_{h})|\nabla_{\bm{\Theta}}\varphi_{h}\cdot{\bm{\Theta}}^{\prime}(t)|^{2}\mathrm{d}\bm{x}.$ For example, in Galerkin methods, one can let $\varphi_{h}(\bm{x},t)$ be $\varphi_{h}(\bm{x},t)=\sum_{i=1}^{K}\gamma_{i}(t)\psi_{i}(\bm{x})$ with $\{\psi_{i}\}_{i=1}^{K}$ being the set of basis functions. Then ${\bm{\Theta}}(t)$ is a vector given by ${\bm{\Theta}}(t)=(\gamma_{1}(t),\gamma_{2}(t),\ldots,\gamma_{K}(t))^{\rm T}\in\mathbb{R}^{K}$.

By employing the energetic variational approach in the semi-discrete level (21), one can obtain an ODE system of ${\bm{\Theta}}(t)$. Particularly, in the linear response regime, $\mathcal{D}_{h}({\bm{\Theta}}(t),{\bm{\Theta}}^{\prime}(t))=\frac{1}{2}\int\eta(\varphi_{h})|\nabla_{\bm{\Theta}}\varphi_{h}\cdot{\bm{\Theta}}^{\prime}(t)|^{2}\mathrm{d}\bm{x}$ is a quadratic function of ${\bm{\Theta}}^{\prime}$. The ODE system of ${\bm{\Theta}}(t)$ can then be written as

where

The ODE (22) is the same as the ODE (20) in the neural Galerkin method [8] for $L^{2}$ gradient flows, although the derivation is different.

Since (22) is a finite-dimensional gradient flow, one can then construct a minimizing movement scheme for ${\bm{\Theta}}$ [25]: finding ${\bm{\Theta}}^{n+1}$ such that

Here $\mathcal{S}_{ad}^{h}$ is the admissible set of ${\bm{\Theta}}$ inherited from the admissible set of $\varphi$, denoted by $\mathcal{S}$, and ${\sf D}^{n}_{*}$ is a constant matrix. A typical choice of ${\sf D}^{n}_{*}={\sf D}({\bm{\Theta}}^{n})$. An advantage of this scheme is that

if ${\sf D}^{n}_{*}$ is positive definite, which guarantees the energy stability for the discrete free energy $\mathcal{F}_{h}({\bm{\Theta}})$. Moreover, by choosing a proper optimization method, we can assure that $\varphi^{n+1}_{h}$ stays in the admissible set $\mathcal{S}$.

Although neural networks can be used to construct $\varphi_{h}(\bm{x};{\bm{\Theta}}(t))$, it might be expensive to compute $\nabla_{\bm{\Theta}}\varphi_{h}\otimes\nabla_{\bm{\Theta}}\varphi_{h}$ in ${\sf D}({\bm{\Theta}})$. Moreover, ${\sf D}({\bm{\Theta}})$ is not a sparse matrix and requires a lot of computer memory to store when a deep neural network is used.

To overcome these difficulties, we propose an alternative approach by introducing temporal discretization before spatial discretization. Let $\tau$ be the time step size. For the $L^{2}$-gradient flow (5), given $\varphi^{n}$, which represents the numerical solution at $t^{n}=n\tau$, one can obtain $\varphi^{n+1}$ by solving the following optimization problem: finding $\varphi^{n+1}$ in some admissible set $\mathcal{S}$ such that

Let $\varphi_{h}(\bm{x};{\bm{\Theta}})$ be a finite-dimensional approximation to $\varphi$ with ${\bm{\Theta}}\in\mathbb{R}^{K}$ being the parameter of the spatial discretization (e.g., weights of linear combination in a Galerkin approximation) yet to be determined, then the minimizing movement scheme (26) can be written in terms of ${\bf\Theta}$: finding ${\bm{\Theta}}^{n+1}$ such that

Here ${\bm{\Theta}}^{n}$ is the value of ${\bm{\Theta}}$ at time $t^{n}$. $\mathcal{F}_{h}[{\bm{\Theta}}]=\mathcal{F}[\varphi_{h}(\bm{x};{\bm{\Theta}})]$, and $\eta^{n}=\eta(\varphi_{h}(\bm{x};{\bm{\Theta}}^{n}))$.

By choosing a certain neural network to approximate $\varphi$, denoted by $\varphi_{\rm NN}(\bm{x};{\bm{\Theta}})$ (${\bm{\Theta}}$ is used to denote all parameters in the neural network), we can summarize the EVNN scheme for solving a $L^{2}$-gradient flow in Alg. 1.

It can be noticed that both Eq. (29) and Eq. (30) involve integration in the computational domain $\Omega$. This integration is often computed by using a grid-based numerical quadrature or Monte–Carlo/Quasi–Monte–Carlo algorithms [65]. It is worth mentioning that due to the non-convex nature of the optimization problem and the error in estimating the integration, it might be difficult to find an optimal ${\bm{\Theta}}^{n+1}$ at each step. But since we start the optimization procedure with ${\bm{\Theta}}^{n}$, we’ll always be able to get a ${\bm{\Theta}}^{n+1}$ that lowers the discrete free energy at least on the training set. In practice, the optimization problem can be solved by either deterministic optimization algorithms, such as L-BFGS and gradient descent with Barzilai–Borwein step-size, or stochastic gradient descent algorithms, such as AdaGrad and Adam.

In this subsection, we show how to formulate an EVNN scheme in the Lagrangian frame of reference for generalized diffusions.

As discussed previously, a generalized diffusion can be viewed as an $L^{2}$-gradient flow of the flow map $\bm{x}(\mathbf{X},t)$ in the space of diffeomorphisms. Hence, the EVNN scheme for the generalized diffusion can be formulated in terms of a minimizing movement scheme of the flow map given by

where $\Phi^{n}(\mathbf{X})$ is a numerical approximation of the flow map $\bm{x}(\mathbf{X},t)$ at $t^{n}=n\tau$, $\rho_{0}(\mathbf{X})$ is the initial density. $\mathcal{F}[\rho]$ is the free energy for the generalized diffusion defined in Eq. (9), and

One can parameterize $\Phi:\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$ by a suitable neural network. The remaining procedure is almost the same as that of the last subsection. However, it is often difficult to solve this optimization problem directly, and one might need to build a large neural network to approximate $\Phi^{n+1}$ when $n$ is large.

To overcome this difficulty, we proposed an alternative approach, Instead of seeking an optimal map $\Phi^{n+1}$ between $\rho^{0}$ and $\rho^{n+1}$, we seek an optimal $\Psi^{n+1}$ between $\rho^{n}$ and $\rho^{n+1}$. More precisely. we assume that

Given $\rho^{n}$, one can compute $\Psi^{n+1}$ by solving the following optimization problem

The corresponding $\rho^{n+1}$ can then be computed through $\rho^{n+1}(\bm{x})=(\Psi^{n+1}_{\verb|#|}\rho^{n})(\bm{x})$. An advantage of this approach is that we only need a small size of neural network to approximate $\Psi^{n+1}$ at each time step when $\tau$ is small.

The numerical algorithm for solving the generalized diffusion is summarized in Alg. 2.

The next question is how to accurately evaluate the numerical integration in (34). Let $\bm{x}_{i}^{n}=\Phi^{n}(\bm{x}_{i}^{0})$, for the general free energy (9), one way to evaluate the integrations in Eq. (34) is using

where $|\Omega_{i}^{n}|$ is the volume of the Voronoi cells associated with the set of points $\{\bm{x}_{i}^{n}\}$, $\rho^{n}_{i}$ stands for $\rho^{n}(\bm{x}_{i}^{n})$. Here the numerical integration is computed through a piece-wisely constant reconstruction of $\rho^{n}$ based on its values at $\{\Phi^{n}(\bm{x}_{i}^{0})\}_{i=1}^{N}$. However, it is not straightforward to compute $|\Omega_{i}^{n}|$, particularly for high dimensional cases. In the current study, we assume that the initial samples are drawn from $\rho_{0}$, one can roughly assume that $\{\bm{x}_{i}^{n}\}$ follows the distribution $\rho^{n}$, then according to the Monte–Carlo approach, the numerical integration can be evaluated as

where $f_{\omega}(\rho)=\omega(\rho)/\rho.$ The proposed numerical method can be further improved if one can evaluate the integration more accurately, i.e., have an efficient way to estimate $|\Omega_{i}^{n}|$. Alternatively, as an advantage of the neural network-based algorithm, we can get $\rho^{n}(x)=\rho_{0}\tilde{\circ}(\Phi^{n})^{-1}(\bm{x}),~{}\forall\bm{x}$, which enables us to estimate $J(\Psi)$ by resampling. More precisely, assume $\mu$ is a distribution that is easy to sample