[Paper Review] Tackling the Curse of Dimensionality with Physics-Informed Neural Networks

Zheyuan Hu, Khemraj Shukla+2 more|arXiv (Cornell University)|Jul 23, 2023

Model Reduction and Neural NetworksPhysics and Astronomy62 references17 citations

TL;DR

The paper introduces Stochastic Dimension Gradient Descent (SDGD) to scale physics-informed neural networks (PINNs) to arbitrarily high-dimensional PDEs, with theoretical convergence guarantees and fast, memory-efficient experiments on very high-dimensional problems.

ABSTRACT

The curse-of-dimensionality taxes computational resources heavily with exponentially increasing computational cost as the dimension increases. This poses great challenges in solving high-dimensional PDEs, as Richard E. Bellman first pointed out over 60 years ago. While there has been some recent success in solving numerically partial differential equations (PDEs) in high dimensions, such computations are prohibitively expensive, and true scaling of general nonlinear PDEs to high dimensions has never been achieved. We develop a new method of scaling up physics-informed neural networks (PINNs) to solve arbitrary high-dimensional PDEs. The new method, called Stochastic Dimension Gradient Descent (SDGD), decomposes a gradient of PDEs into pieces corresponding to different dimensions and randomly samples a subset of these dimensional pieces in each iteration of training PINNs. We prove theoretically the convergence and other desired properties of the proposed method. We demonstrate in various diverse tests that the proposed method can solve many notoriously hard high-dimensional PDEs, including the Hamilton-Jacobi-Bellman (HJB) and the Schrödinger equations in tens of thousands of dimensions very fast on a single GPU using the PINNs mesh-free approach. Notably, we solve nonlinear PDEs with nontrivial, anisotropic, and inseparable solutions in 100,000 effective dimensions in 12 hours on a single GPU using SDGD with PINNs. Since SDGD is a general training methodology of PINNs, it can be applied to any current and future variants of PINNs to scale them up for arbitrary high-dimensional PDEs.

Motivation & Objective

• Motivate and address the curse of dimensionality in solving high-dimensional PDEs with PINNs.
• Develop a gradient-splitting training method (SDGD) that reduces memory use and accelerates convergence.
• Provide theoretical guarantees for unbiased gradients and convergence of SDGD.
• Demonstrate the method on diverse high-dimensional PDEs including HJB and Schrödinger equations.

Proposed method

• Decompose the PDE residual gradient into per-dimension terms and sample a subset to form an unbiased stochastic gradient.
• Define SDGD as the process of selecting a random set of PDE terms I and backpropagating only through those terms to reduce memory.
• Prove SDGD gradients are unbiased and establish convergence guarantees.
• Extend SDGD with gradient accumulation and parallel computing to further speed up training.
• Propose forward-backward sampling strategies (Algorithms 2 and 3) to trade variance for speed.
• Show that SDGD enables mesh-free PINN training and scales to extremely high dimensions (up to 100k effective dimensions) on a single GPU.

Figure 5: Convergence plot obtained from data-parallel algorithm for HJB-Lin equation in 100,000 dimensions. To perform this experiment we use DIM=101 and $N_{f}=100$ . Here we denote the batch size for dimension in Algorithms 1 and/or 2 of SDGD as DIM and the number of collocation points used in ea

Experimental results

Research questions

• RQ1 Can SDGD provide unbiased stochastic gradients for PINN training when gradients are decomposed by PDE terms across dimensions?
• RQ2 Does sampling over PDE terms reduce memory and accelerate convergence without sacrificing accuracy?
• RQ3 How does SDGD compare to standard SGD over collocation points in high-dimensional PDEs in terms of stability, convergence, and memory usage?
• RQ4 Can SDGD solve high-dimensional PDEs such as Hamilton-Jacobi-Bellman and Schrödinger equations in very large dimensions on practical hardware?
• RQ5 What theoretical guarantees (unbiasedness, convergence) accompany SDGD, and how should batch sizes be chosen for optimal variance under fixed memory budgets?

Key findings

• SDGD yields unbiased stochastic gradients for PINNs by sampling PDE-term gradients, enabling memory-efficient training.
• With SDGD, training can scale to extremely high-dimensional PDEs and, in some cases, solve nonlinear PDEs with nontrivial, anisotropic, inseparable solutions in up to 100,000 effective dimensions within hours on a single GPU.
• SDGD demonstrates comparable stability to SGD over collocation points and can even outperform it under the same memory constraints in several nonlinear high-dimensional PDEs.
• The method enables mesh-free training and full-domain prediction, and when combined with adversarial training, significantly accelerates high-dimensional learning.
• The approach leverages parallel computing and gradient accumulation, achieving substantial speed-ups and expanding the practical applicability of PINNs to very high dimensions.

Figure 6: Schematic of tensor parallel approach. Here, $w$ is a weight matrix that is split into column axis and deployed to different GPUs.