Understanding and mitigating gradient pathologies in physics-informed neural networks

TL;DR

The paper investigates why physics-informed neural networks struggle to converge when constrained by PDE residuals and data terms. It identifies gradient pathologies stemming from stiffness in the gradient flow, and introduces an adaptive learning-rate annealing procedure plus a novel neural-architecture (ADGM) to mitigate these issues. Across Helmholtz, Klein-Gordon, and lid-driven cavity benchmarks, the proposed methods yield substantial accuracy gains (up to 50–100x) and more robust training, with improved boundary-condition handling and incompressibility enforcement. The work offers practical strategies to balance multi-term losses in constrained neural networks and highlights directions for future stable optimization and architecture design in scientific machine learning.

Abstract

The widespread use of neural networks across different scientific domains often involves constraining them to satisfy certain symmetries, conservation laws, or other domain knowledge. Such constraints are often imposed as soft penalties during model training and effectively act as domain-specific regularizers of the empirical risk loss. Physics-informed neural networks is an example of this philosophy in which the outputs of deep neural networks are constrained to approximately satisfy a given set of partial differential equations. In this work we review recent advances in scientific machine learning with a specific focus on the effectiveness of physics-informed neural networks in predicting outcomes of physical systems and discovering hidden physics from noisy data. We will also identify and analyze a fundamental mode of failure of such approaches that is related to numerical stiffness leading to unbalanced back-propagated gradients during model training. To address this limitation we present a learning rate annealing algorithm that utilizes gradient statistics during model training to balance the interplay between different terms in composite loss functions. We also propose a novel neural network architecture that is more resilient to such gradient pathologies. Taken together, our developments provide new insights into the training of constrained neural networks and consistently improve the predictive accuracy of physics-informed neural networks by a factor of 50-100x across a range of problems in computational physics. All code and data accompanying this manuscript are publicly available at \url{https://github.com/PredictiveIntelligenceLab/GradientPathologiesPINNs}.

Figures (18)

• Figure 1: Helmholtz equation: Exact solution versus the prediction of a conventional physics-informed neural network model with 4 hidden layers and 50 neurons each layer after 40,000 iterations of training with gradient descent (relative $L^2$-error: 1.81e-01).
• Figure 2: Helmholtz equation: Histograms of back-propagated gradients $\nabla_\theta \mathcal{L}_r(\theta)$ and $\nabla_\theta \mathcal{L}_{u_b}(\theta)$ at each layer during the $40,000$th iteration of training a standard PINN model for solving the Helmholtz equation.
• Figure 3: Possion equation: Histograms of back-propagated gradients $\nabla_\theta \mathcal{L}_r(\theta)$ and $\nabla_\theta \mathcal{L}_{u_b}(\theta)$ at each layer during the $40,000$th iteration of training a standard PINN model for solving the one-dimensional Poisson equation for different values of the constant $C$, see equation \ref{['eq:Poisson']}.
• Figure 4: Helmholtz equation: (a) All eigenvalues of $\nabla^2_\theta \mathcal{L}_r (\theta)$ and $\nabla^2_\theta \mathcal{L}_{u_b} (\theta)$, respectively, arranged in increasing order. (b) Loss curves of $\mathcal{L}_r(\theta)$ and $\mathcal{L}_{u_b}$, respectively, after 40,000 iterations of gradient decent using the Adam optimizer in full batch mode.
• Figure 5: Stiffness in the gradient flow dynamics: Largest eigenvalue of the Hessian $\nabla_{\theta}^2\mathcal{L}(\theta)$ during the training of a physics-informed neural network model for approximating the solution to the two-dimensional Helmholtz problem (see equation \ref{['eq:Helmholtz']}) and for different values of the control parameters $a_1$ and $a_2$.

Theorems & Definitions (1)

• proof