[Paper Review] Understanding and mitigating gradient pathologies in physics-informed neural networks

Sifan Wang, Yujun Teng+1 more|arXiv (Cornell University)|Jan 13, 2020

Model Reduction and Neural Networks49 references197 citations

TL;DR

The paper identifies gradient pathologies arising from stiff gradient flow in physics-informed neural networks (PINNs) and proposes an adaptive learning-rate annealing scheme plus a new architecture to balance loss terms, improving predictive accuracy substantially.

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}.

Motivation & Objective

• Explain why PINNs suffer from gradient pathologies during training, especially due to stiffness in gradient flow.
• Diagnose the imbalance between data-fit and PDE-residual terms in PINN losses.
• Propose an adaptive learning-rate annealing algorithm to balance composite losses.
• Introduce a neural network architecture that is more resilient to gradient pathologies.
• Demonstrate improved predictive accuracy of PINNs across computational physics problems.

Proposed method

• Review PINNs formulation with a composite loss L(θ)=Lr(θ)+Σi λi Li(θ).
• Analyze gradient imbalance between boundary data fit and PDE residual terms using simple benchmarks like the Helmholtz and Poisson equations.
• Show that stiffness in gradient flow can cause instability with standard gradient descent, motivating adaptive methods.
• Propose a learning rate annealing algorithm that adapts the weights λi online based on gradient statistics (akin to adaptive optimizers).
• Introduce a novel fully-connected neural network architecture with reduced stiffness to improve training stability.
• Provide empirical demonstrations across problems in computational physics showing large predictive gains.

Experimental results

Research questions

• RQ1 What causes gradient pathologies and unbalanced back-propagated gradients in PINNs during training?
• RQ2 How does gradient stiffness affect the stability and convergence of gradient-based optimization in PINNs?
• RQ3 Can adaptive loss weighting based on gradient statistics stabilize training of PINNs?
• RQ4 Does a redesigned network architecture reduce gradient stiffness and improve accuracy for PINNs?
• RQ5 What level of predictive improvement is achievable by the proposed methods across computational physics problems?

Key findings

• Gradient pathologies in PINNs are linked to stiffness in the gradient flow, causing unbalanced gradients between loss terms.
• The PDE residual term can dominate training when boundary/initial-condition gradients vanish, leading to erroneous predictions.
• An adaptive learning-rate annealing algorithm that tunes loss-term weights online stabilizes training and balances contributions.
• A novel network architecture reduces gradient-flow stiffness compared to standard fully-connected nets.
• The proposed methods yield consistent predictive accuracy improvements in PINNs by a factor of 50–100x across multiple computational-physics problems.