[Paper Review] An Expert's Guide to Training Physics-informed Neural Networks

Sifan Wang, Shyam Sankaran+2 more|arXiv (Cornell University)|Aug 16, 2023

Model Reduction and Neural Networks64 citations

TL;DR

This paper presents a comprehensive training pipeline for PINNs, including non-dimensionalization, advanced architectures (Fourier features, random weight factorization), and training algorithms (causal, curriculum, adaptive loss weighting), plus fully reproducible JAX implementations and extensive ablations.

ABSTRACT

Physics-informed neural networks (PINNs) have been popularized as a deep learning framework that can seamlessly synthesize observational data and partial differential equation (PDE) constraints. Their practical effectiveness however can be hampered by training pathologies, but also oftentimes by poor choices made by users who lack deep learning expertise. In this paper we present a series of best practices that can significantly improve the training efficiency and overall accuracy of PINNs. We also put forth a series of challenging benchmark problems that highlight some of the most prominent difficulties in training PINNs, and present comprehensive and fully reproducible ablation studies that demonstrate how different architecture choices and training strategies affect the test accuracy of the resulting models. We show that the methods and guiding principles put forth in this study lead to state-of-the-art results and provide strong baselines that future studies should use for comparison purposes. To this end, we also release a highly optimized library in JAX that can be used to reproduce all results reported in this paper, enable future research studies, as well as facilitate easy adaptation to new use-case scenarios.

Motivation & Objective

• Provide a structured training pipeline that mitigates PINN training pathologies (spectral bias, causality, gradient imbalance).
• Develop robust, transferable best practices for PINN architecture and optimization.
• Offer challenging benchmarks and extensive ablation studies to establish strong baselines.
• Release a high-performance JAX library to reproduce results and facilitate future research.

Proposed method

• Apply PDE non-dimensionalization to place variables on a similar scale.
• Use MLPs with Fourier feature embeddings and random weight factorization to reduce spectral bias and speed convergence.
• Formulate a weighted PINN loss with adaptive temporal and global weights to balance terms.
• Incorporate causal training by sequentially weighting temporal residuals to respect causality.
• Incorporate curriculum training and loss-balancing strategies (gradient-based and NTK-inspired) to improve stability and accuracy.
• Provide a reproducible JAX library to reproduce results and enable future work.

Experimental results

Research questions

• RQ1 How do non-dimensionalization and network architecture choices affect PINN training efficiency and accuracy?
• RQ2 What training strategies (causal training, curriculum learning, adaptive loss weighting) most effectively address PINN pathologies?
• RQ3 How do ablation studies across benchmarks inform robust baselines and best practices for PINNs?
• RQ4 Can a standardized, reproducible library reproduce state-of-the-art results and facilitate future research in PINNs?

Key findings

• Non-dimensionalization and architectural choices (Fourier features, random weight factorization) significantly improve training convergence and accuracy.
• Self-adaptive loss balancing and causal training reduce gradient imbalance and enforce temporal causality, leading to more robust optimization.
• Curriculum training and NTK-based weighting schemes further enhance stability and performance across benchmarks.
• Extensive ablation studies establish practical guidelines and strong baselines that future PINN work should compare against.
• A highly optimized JAX library is released to reproduce results and support future applications.