Stabilize physics-informed neural networks for stiff differential equations: Re-spacing layer

Abstract

Approximating the solution of stiff differential equations, which exhibit abrupt changes in certain regions, using physics-informed neural networks (PINNs) is challenging. Typically, training PINNs involves using a larger number of samples concentrated around regions of rapid changes to resolve the sharp gradients. However, this strategy leads to data imbalance, resulting in slower convergence and reduced solution quality. Here, we propose Re-spacing layer (RS-layer) to mitigate these challenges. RS-layer is a pre-trained encoding layer designed to map the skewed distribution of sampling points onto a uniform distribution, maintaining the desirable statistical properties of the input data for effective PINN training. We demonstrate that RS-layer improves PINN training by regularizing the solution gradient in the transformed space. The efficacy of our method is validated through numerical experiments on one-dimensional singularly perturbed equations, the ROBER problem, and the Akzo Nobel problem. Our results show that RS-layer not only accelerates convergence, but also enhances accuracy.