Energetic Variational Neural Network Discretizations of Gradient Flows

TL;DR

The paper tackles structure-preserving computation of gradient flows by embedding neural-network spatial discretization within the energetic variational framework. EVNN discretizes the energy-dissipation law directly and employs a temporal-then-spatial strategy to achieve memory efficiency, enabling mesh-free simulations in high dimensions for both $L^2$-gradient flows and generalized diffusions. Key contributions include a semi-discrete energy law with an optimization-based time step (minimizing movement) for Eulerian problems and a diffeomorphism-based, Lagrangian EVNN for diffeomorphic transport, along with practical NN architectures (shallow/ResNet for Eulerian, planar/ICNN-based for Lagrangian) to preserve positivity, mass conservation, and energy decay. Numerical experiments across Poisson, heat, Allen–Cahn, Fokker–Planck, and porous medium equations demonstrate energy stability, accuracy, and the mesh-free advantage, including high-dimensional tests up to four dimensions. The work offers a scalable, physics-informed NN framework with potential impact on high-dimensional gradient flows, density estimation, and related multiscale PDEs in science and machine learning.

Abstract

We present a structure-preserving Eulerian algorithm for solving $L^2$-gradient flows and a structure-preserving Lagrangian algorithm for solving generalized diffusions. Both algorithms employ neural networks as tools for spatial discretization. Unlike most existing methods that construct numerical discretizations based on the strong or weak form of the underlying PDE, the proposed schemes are constructed based on the energy-dissipation law directly. This guarantees the monotonic decay of the system's free energy, which avoids unphysical states of solutions and is crucial for the long-term stability of numerical computations. To address challenges arising from nonlinear neural network discretization, we perform temporal discretizations on these variational systems before spatial discretizations. This approach is computationally memory-efficient when implementing neural network-based algorithms. The proposed neural network-based schemes are mesh-free, allowing us to solve gradient flows in high dimensions. Various numerical experiments are presented to demonstrate the accuracy and energy stability of the proposed numerical schemes.

Figures (8)

• Figure 1: Schematic illustration of a flow map $x(\mathbf{X}, t)$.
• Figure 1: Relative $l^2$ error with respect to CPU time for different methods for 2D Poisson equations.
• Figure 2: Numerical results for the heat equation. (a) - (d) are the neural network solutions at $t = 0.01, 0.2, 0.4$ and $0.6$. (e) - (h) are the absolute differences between the EVNN solutions and the FDM solutions. (i) The evolution of discrete free energy with respect to time for both the EVNN and FDM solutions.
• Figure 3: Numerical results for the Allen--Cahn equation with a volume constraint. (a)-(d) NN solutions at $t = 0, 0.05, 0.1$ and $0.3$, respectively. (e)-(h) Absolute differences between the EVNN solutions and the FEM solutions at $t = 0, 0.05, 0.1$, and $0.3$, respectively. (i) Evolution of numerical free energies with respect to time for both the FEM and EVNN solutions.
• Figure 4: Numerical results for the 2D Fokker--Planck equation. (a)-(d) Predicted solution on a $301 \times 301$ uniform mesh on $(-3, 3)^2$ at $t = 0, 0.1, 0.5$ and $1$, respectively. (e)-(f) The absolute and relative $l_2$-errors of the solution over time, evaluated on the uniform mesh.

Theorems & Definitions (6)

• Remark 2.1
• Remark 3.1
• Remark 3.2
• Remark 3.3
• Remark 3.4
• Remark 3.5