Tackling the Curse of Dimensionality with Physics-Informed Neural Networks

Zheyuan Hu; Khemraj Shukla; George Em Karniadakis; Kenji Kawaguchi

Source-linked AI summary

Tackling the Curse of Dimensionality with Physics-Informed Neural Networks

Zheyuan Hu, Khemraj Shukla, George Em Karniadakis, Kenji Kawaguchi

arXiv:2307.12306v6Jul 23, 2023cs.LGcs.AImath.DSmath.NAstat.ML

TL;DR

The paper addresses the computational and memory costs that grow rapidly with PDE dimensionality and limit high-dimensional PINN training. It introduces SDGD, which samples dimension-specific gradient components, and proves unbiasedness and convergence. Experiments report training nonlinear PDEs in 100,000 dimensions in 12 hours on one GPU, while traditional PINN training exceeds memory limits.

Problem
High-dimensional PDEs face exponentially increasing computational costs, and traditional PINN training can incur prohibitive memory demands as dimensionality and second-order terms grow.
Method
SDGD decomposes PDE and residual-loss gradients into dimension-specific pieces and randomly samples subsets of those pieces during PINN training.
Results
100,000 dimensions were trained in 12 hours on one GPU, while traditional PINN training exceeded the GPU memory limit.
Takeaways & Limitations
SDGD provides a general PINN training methodology with mesh-free whole-domain prediction and extensions to physics-informed neural operators.
Takeaways & Limitations
Sampling more randomly can increase gradient variance and make convergence sub-optimal, while shared forward-backward samples reduce cost but introduce biased gradients.

Abstract

from arXiv · show

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.

1 Introduction

High-dimensional PDEs impose rapidly increasing computational and memory costs, while PINNs offer mesh-free, general-form solutions. SDGD addresses PINN scaling by sampling dimension-specific gradient pieces, with theoretical guarantees and experiments reaching 100,000 dimensions.

Motivation: High-dimensional PDEs incur exponentially increasing computational costs, creating a major challenge for accurate numerical solutions.The curse of dimensionality also creates memory challenges absent from low-dimensional problems.
Motivation: PINNs provide mesh-free, general-form PDE solutions and predictions throughout the domain, including complex geometries.Their neural-network interpolation capability distinguishes them from methods restricted to certain PDE types or single-point evaluation.
Method: SDGD decomposes residual gradients into dimension-specific pieces and randomly samples subsets during each training iteration.The sampled gradients form an unbiased estimator of the full gradient, supporting convergence across dimensions.
Method: SDGD supports parallel computation and gradient accumulation, reducing memory demands while enabling larger effective batch sizes.These mechanisms allow the method to use multiple devices or resource-limited machines more efficiently.
Theory: Unbiased stochastic gradients and convergence guarantees establish SDGD’s theoretical foundation, while suitable batch sizes can reduce variance and accelerate convergence.The analysis extends stochastic-gradient ideas beyond sampling collocation points alone.
Experiments: 100,000 dimensions were trained in 12 hours on one GPU, whereas traditional PINN training exceeded the GPU memory limit.Experiments also report stability comparable to SGD over collocation points and accelerated convergence on nonlinear high-dimensional PDEs.

2 Related Work

Prior work addresses high-dimensional PDEs through specialized deep-learning and numerical methods, while this paper develops a stochastic-gradient framework for scaling PINNs. The related literature includes theoretical PINN analyses and optimization guarantees motivating SDGD’s convergence treatment.

Physics-informed learning: PINNs use neural networks as surrogate PDE solutions and optimize boundary and residual losses, whereas DeepONet directly fits PDE operators from data.DeepONet can also incorporate physical information through residual loss.
PINN theory: Existing PINN theory studies generalization bounds, neural-tangent-kernel convergence, residual-based error estimates, and minimizer convergence with increasing samples.These analyses examine different theoretical aspects of PINN behavior.
High-dimensional PDE methods: High-dimensional PDE methods include specialized approaches for HJB equations, derivative-free PINN optimization, proper generalized decomposition, information fusion, ANOVA, and multilevel Picard methods.Several numerical approaches target particular structures or classes of high-dimensional equations.
Stochastic optimization: SDGD extends the SGD mini-batch idea from training examples to PDE terms and dimensions for accelerating and scaling PINNs.Its convergence analysis builds on conditions concerning gradient-estimator regularity and stochastic-gradient variance.

3 Method

SDGD scales PINN training by decomposing PDE-residual gradients into dimension-specific terms and sampling only subsets during each iteration. The approach reduces memory use, supports further acceleration, and preserves unbiasedness in its independently sampled forward/backward variant.

Motivation: PINNs solve general PDEs with mesh-free predictions across complex geometries, but high-dimensional models incur memory costs that grow with network size and derivative count.Traditional point sampling cannot reduce the cost below one collocation point plus all PDE dimensions.
Core method: SDGD decomposes the full residual gradient into PDE-term or dimension-specific pieces and samples a subset of them for stochastic optimization.The sampled index set can be much smaller than the total number of dimensions, reducing backpropagation cost.
Batching and memory: The main backpropagation cost matches that of the corresponding sampled-dimensional PDE, and the resulting stochastic gradient is unbiased.Terms not selected for backpropagation can be detached to save GPU memory; full gradients can also be accumulated sequentially, although this is time-consuming.
Batching and memory: SDGD has the same memory cost for point-term batches with equal products, such as 50 points and 100 terms versus 500 points and 10 terms.This combines sampling over PDE terms with the conventional sampling of residual points.
Further speed-up: Independent sampling in forward and backward passes keeps the gradient unbiased, whereas sampling once is faster but introduces bias controlled by the dimension batch size.The sampling-once variant can converge well in practice when its dimension batch is not too small.
Further speed-up: SDGD trades convergence quality against speed when increased randomness raises gradient variance, particularly in extremely large-scale PINN training.The authors state that imprecise gradients can still reduce loss substantially during initialization.

4 Theory

The theory establishes unbiased stochastic gradients, variance-aware batch selection, and convergence guarantees for SDGD under boundedness assumptions. It also characterizes local convergence and compares SDGD favorably with RS-PINN and HTE in generality and high-dimensional variance behavior.

Unbiasedness: SDGD’s stochastic gradients are unbiased estimators of the full-batch gradient when sampling PDE terms, or both PDE terms and residual points.The expected sampled gradients equal the gradient obtained using all PDE terms.
Assumptions and stability: Normalizing by residual-point and PDE-term counts prevents residual-loss and gradient scales from becoming exceptionally large as dimensionality and term count increase.The stated purpose is to mitigate gradient explosions.
Variance and batching: Under a fixed memory budget |B| × |I|, choosing point and PDE-term batch sizes appropriately can minimize gradient variance and accelerate convergence.SDGD therefore complements SGD over residual points by adding stochasticity over PDE terms.
Variance and batching: With |B||I| = 100 units and zero point-induced gradient variance, the minimum-variance choice is |B| = 100 and |I| = 1.Increasing the number of PDE terms does not reduce variance in this stated case.
Assumptions and stability: Bounded neural-network derivatives, bounded PDE operators, and bounded optimization trajectories imply bounded PINN residuals and stochastic-gradient variance.These conditions support the convergence analysis.
Convergence: The convergence rate is O(1/n^p) under bounded gradient variance, with larger variance producing slower convergence and smaller variance producing faster convergence.The result concerns the SGD trajectory under the theorem’s assumptions.
Convergence: The convergence guarantee is local and requires initialization inside a specified neighborhood; Xavier initialization is given as a way to mitigate unstable starting parameters.The initialization discussion links excessively large or small weights to exploding or vanishing loss and gradients.
Comparisons: Compared with RS-PINN, SDGD applies to arbitrary high-order PDEs, retains a regular universal-approximator network, avoids the cited nonlinear-PDE sampling bias, and is reported faster.RS-PINN is described as limited to lower-than-second-order variance-reduction settings and smooth solutions.

5 Experiments

Experiments show that SDGD remains stable and efficient across nonlinear PDEs and scales PINNs to 100,000 dimensions despite the memory limitations of conventional training.

Core results: SDGD achieves approximately 1e-3 relative L2 errors across dimensions and nonlinear PDEs, with remarkably similar convergence curves under matched settings.The tested equations include Poisson, Allen-Cahn, and Sine-Gordon PDEs.
Core results: SDGD’s memory and optimization time costs grow linearly with dimensionality; from 10K to 100K dimensions, both increase by less than 10×.
Core results: At 100K dimensions, SDGD uses 100 sampled dimensions, whereas full-batch gradients would require over 10× more memory and exceed an A100 GPU’s 80GB limit.
Core results: SDGD convergence curves become smoother in higher dimensions, which the authors attribute to the blessing of dimensionality.
Comparison with vanilla PINNs: Vanilla PINNs incur quadratic computational growth from network inputs and Laplacian derivatives, exceeding A100 memory at 10,000 dimensions.Below 5,000 dimensions, SDGD has comparable performance with acceleration and lower memory use.
Conclusion: Overall, SDGD with PINNs handles nonlinear high-dimensional PDEs with inseparable and complicated solutions.
Batch-size ablation: Across batch-size choices, small sampled point and dimension batches preserve convergence while improving speed, but extremely unbalanced or tiny batches can slow computation.Balanced batches perform faster than extreme allocations with the same total batch size.
Algorithm ablation: Among Algorithms 1–3, Algorithm 3 is fastest and remains effective, while additional randomness produces marginally worse performance than more precise alternatives.The algorithms use equal memory settings for fair comparison.

5.2 Nonlinear Fokker-Planck PDEs and Comparison with Strong Baselines

The experiments compare SDGD-PINNs with strong non-PINN baselines on nonlinear PDEs without analytical solutions. Across the tested equations and dimensions, PINN surpasses the competing methods while evaluating solutions at a test point.

Baseline comparison: SDGD-PINNs outperform strong non-PINN competitors on several nonlinear PDEs without analytical solutions.The evaluation uses reference values computed by the multilevel Picard method.
Test problems: The tested problems include Allen-Cahn, semilinear heat, and sine-Gordon PDEs with highly nontrivial, nonseparable solutions.
Results: PINN surpasses other methods across all tested dimensions and PDE types using relative L1 error on one test point.The comparison attributes this capability to PINNs’ mesh-free formulation and interpolation properties.

5.3 SDGD Accelerates PINN’s Adversarial Training in HJB PDEs

SDGD enhances adversarial PINN training for high-dimensional HJB equations by sampling PDE-term gradients, addressing the computational cost of adversarial training and outperforming prior methods in accuracy and speed.

The HJB-Rosenbrock benchmark has anisotropic, asymmetric, highly coupled, inseparable solutions with effective dimension d + 1.
Adversarial training approximates the L∞ loss but is slow and inefficient for high-dimensional HJB equations with LQG control.
SDGD enhances adversarial PINN training by applying stochastic gradient descent over PDE terms to improve scalability and efficiency.
In the 250D HJB-LQG case, SDGD achieves a superior L2 error in 3.5 hours, whereas baseline requires 1 day to reach relative L2 error 1e-2.
In the 1000D case, SDGD reaches relative L2 error 5.852E-3 in 217 minutes while full-batch adversarial training requires over 5 days.
SDGD with PINNs predicts 20,000 test points after one training process, whereas traditional methods require 20,000 separate inference rounds.
Across the tested HJB equations and dimensions, SDGD outperforms the baselines in accuracy and speed while using relatively lower memory.

5.4 Schr¨odinger Equation

The Schrödinger experiments combine a separable Tensor Neural Network with SDGD to address high-dimensional eigenvalue problems and their derivative-related computational bottlenecks. Gradient accumulation reaches 2 ∗104 dimensions within limited GPU memory while maintaining low relative error, although the approach retains important memory and sampling limitations.

High-dimensional Schrödinger equations are treated as eigenvalue problems using a specialized Tensor Neural Network because the traditional PINN framework is unsuitable.
Traditional quadrature methods incur exponentially growing costs and cannot be used even in the lowest 100-dimensional experiment.
TNN’s separable structure computes numerical integrals along individual axes, reducing the integration cost from exponential to linear in dimension.
TNN uses separate subnetworks for dimensions, making first-order derivatives the primary computational bottleneck and increasing memory requirements.
SDGD samples PDE-term-level gradient pieces using a random index set, producing an unbiased stochastic gradient while reducing computational and memory costs.
Gradient accumulation trains the CQHO problem in 2 ∗104 dimensions within limited GPU memory, with SDGD maintaining relative error below 1e-6 across dimensions.
The approach remains limited for large-scale non-differential bottlenecks: separate subnetworks consume substantial memory, and sampling denominator integrals would bias the stochastic gradient.

6 Conclusions

The conclusions present SDGD-PINNs as a general mesh-free framework for high-dimensional PDEs, with theoretical guarantees, broad applicability, and an explicit batch-size limitation.

Contributions: SDGD generalizes stochastic gradient descent by sampling subsets of dimension-specific gradient pieces while targeting arbitrary high-dimensional PDEs.The sampling estimator is unbiased and is associated with convergence guarantees.
Generality to Arbitrary PDEs: The PINN-based approach handles wave, biharmonic, Schrödinger, and other PDEs that existing compared approaches cannot address.The claimed scope contrasts with methods restricted to specific parabolic PDE forms.
Prediction on the Entire Domain: PINNs provide predictions across the entire domain in one training instance rather than only at a single test point.The conclusions describe this as a “once for all” capability.
Dependency on the Mesh: Because PINNs are mesh-free, training points can be randomly sampled and long-time PDE problems can be handled without temporal discretization.The passage attributes generalization on long-time problems to neural-network interpolation.
Extensions: SDGD can extend beyond PINNs to physics-informed neural operators, general regression, and possibly classification in very high dimensions.The passage specifically names DeepONet and FNO as compatible neural operators.
Limitations: For extremely high-dimensional PDEs, SDGD may require relatively large dimension batches; small batches can increase gradient variance and slow convergence.For PDEs needing special loss functions, the minimum achievable memory may grow faster with dimension.

A.1 Proof of Theorem 4.1

The proof section states that the unbiasedness of the stochastic gradients in Algorithms 1 and 2 follows from the displayed derivations.

Algorithm 1: The appendix derives unbiasedness for the stochastic gradient generated by Algorithm 1.The supplied passage introduces the derivation but does not include its mathematical steps.
Algorithm 2: The appendix also derives unbiasedness for the stochastic gradient generated by Algorithm 2.The supplied passage similarly provides only the derivation heading.

A.2 Proof of Theorem 4.2

The proof analyzes stochastic-gradient variance by modeling random samples of PDE terms and residual points, then decomposing expectations across matching and nonmatching indices.

Sampling lemma: Lemma A.1 considers independent random selections from N fixed numbers and the variance of their unbiased estimator.The lemma samples k random numbers, with each selected uniformly from the N values.
Full-batch gradient: The full-batch gradient is normalized over both residual points and PDE terms and contains NrNL terms.The proof frames the full batch as a mean over residual-point and PDE-term combinations.
Stochastic gradient: Algorithm 1’s stochastic gradient is defined from sampled residual points and PDE terms.The supplied passage gives the stochastic-gradient definition without its complete displayed expression.
Variance control: With fixed product |B| · |J|, batch sets B and J can be selected to minimize stochastic-gradient variance and accelerate convergence.The analysis treats variance as depending on the sizes of the residual-point and PDE-term batches.
Expectation decomposition: The expectation is decomposed into four cases according to whether residual-point and PDE-term indices match.The cases separately distinguish equality and inequality for each index pair.

A.3 Proof of Lemma 4.1

The proof develops derivative bounds for the PINN neural network by counting vector terms and controlling their matrix-valued factors through induction.

Derivative bounds: Lemma 4.1 bounds derivatives of the defined neural network using its parameters, depth, width, and input.The norms are vector or matrix 2-norms, and h denotes the maximal network width.
Low-order derivatives: The first derivative contains one R^d vector term, while the second derivative contains d(L −1) such terms.The second-order terms arise from differentiating the nested network structure.
Second-order case: The second-order case is explicitly identified as satisfying the lemma.The supplied passage states this result without reproducing the complete bound.
Inductive counting: For the nth derivative, the proof counts (n −1)!d^(n−1)(L −1)^(n−1) R^d vector terms.Each term contains at most n(L −1) matrix-valued functions depending on x.
Matrix-factor control: The derivative bounds use activation-function derivatives and matrix products whose norms are controlled by submultiplicativity.The proof introduces Φ terms for activation derivatives and uses ∥AB∥≤∥A∥∥B∥.
Induction step: The induction step yields n!d^n(L −1)^n vector terms for the (n + 1)th derivative.The count combines existing terms, product-rule factors, and derivative-tensor size.
Parameter derivatives: Differentiating with respect to network parameters preserves the induction logic and adds dependence on parameter dimension, width, and input norm.The stated bound includes h^2, derivative order, network depth, weight norms, and max{∥x∥,1}.

A.4 Proof of Theorem 4.3

Theorem 4.3 establishes convergence of SDGD under bounded gradient variance, assuming bounded SGD trajectories and a decaying step size. The proof bounds network derivatives, PDE outputs, gradients, and gradient variance before applying an external SGD convergence theorem.

Assumptions: Under the theorem’s assumptions, SGD trajectories remain bounded throughout optimization, enabling a uniform bound on model parameters.The proof denotes the bound by M(l), defined using the maximum trajectory norm and 1.
Assumptions: The SGD step size follows γn = γ/(n + m)^p with p ∈ (1/2, 1] and sufficiently large γ and m.
Conclusion: Conditioned on the event E∞, the expected squared distance to the regular minimizer satisfies E[∥θn − θ∗∥2|E∞] ≤ O(1/np).
Proof strategy: Bounded trajectories and Lemma 4.1 place high-order neural-network derivatives in a compact domain throughout optimization.
Proof strategy: Assumption 4.2 bounds the PDE operator outputs and their parameter derivatives, which are the optimization gradients.
Conclusion: The resulting gradient variance for SDGD is universally bounded and finite, independently of epoch n.This follows from bounding the gradients produced by the physics-informed loss functions and the stochastic gradients of Algorithms 1 and 2.