Stiff Transfer Learning for Physics-Informed Neural Networks

TL;DR

Stiff differential equations present a major challenge for data-driven solvers, particularly PINNs, due to disparate time scales that cause training instability and inaccuracies. The authors introduce STL-PINN, which trains a Multi-Head PINN in a low-stiff regime and then transfers to a high-stiff regime to obtain a final, one-shot solution, addressing stiffness-induced failure modes and retraining costs. The method leverages a fixed latent representation $H$ to enable fast transfer via a closed-form update for the head weights and uses a perturbation framework for nonlinear terms, achieving superior accuracy and competitive speed against implicit methods across stiff ODEs/PDEs and scalable reparametrizations. The results demonstrate practical impact for rapid parametric simulations and optimization in stiff systems, with STL-PINN outperforming state-of-the-art PINN baselines and offering a viable alternative to traditional solvers in many settings.

Abstract

Stiff differential equations are prevalent in various scientific domains, posing significant challenges due to the disparate time scales of their components. As computational power grows, physics-informed neural networks (PINNs) have led to significant improvements in modeling physical processes described by differential equations. Despite their promising outcomes, vanilla PINNs face limitations when dealing with stiff systems, known as failure modes. In response, we propose a novel approach, stiff transfer learning for physics-informed neural networks (STL-PINNs), to effectively tackle stiff ordinary differential equations (ODEs) and partial differential equations (PDEs). Our methodology involves training a Multi-Head-PINN in a low-stiff regime, and obtaining the final solution in a high stiff regime by transfer learning. This addresses the failure modes related to stiffness in PINNs while maintaining computational efficiency by computing "one-shot" solutions. The proposed approach demonstrates superior accuracy and speed compared to PINNs-based methods, as well as comparable computational efficiency with implicit numerical methods in solving stiff-parameterized linear and polynomial nonlinear ODEs and PDEs under stiff conditions. Furthermore, we demonstrate the scalability of such an approach and the superior speed it offers for simulations involving initial conditions and forcing function reparametrization.

Figures (34)

• Figure 1: MH-PINNs architecture with $h$ heads. The base network constructs $H \in \mathbb{R}^{n \times m}$ from input $\boldsymbol{x}$ that passes through multiple hidden layers. $m$ is the dimension of the last hidden layer and $n$ is the dimension of the system of equations. The $h$ head layers consist of separate linear layers $W_i$ without any activation function. Each of them produces individual outputs $u_i=HW_i$ and their loss. These losses are then summed up to yield to the final loss $L_{tot}$.
• Figure 2: STL-PINN diagram process: 1. Linearize non-linear PDE-ODE. 4. Substitute $u=HW$ in the loss function. 2. Train MH-PINN in a low-stiff regime. 5. Compute $\hat{W}$ in a high-stiff regime with Eq.\ref{['eq:3']}. 3. Extract and fix $H$ after training. 6&7. Compute final solution $u$.
• Figure 3: Results for (a) OHO, NCFF, Duffing ODEs (b) AR PDE. For each equation, an example of both training low-stiff and transfer high-stiff STL-PINN solutions is shown. Low-stiff regimes take $\alpha = \{10, 10, 19, 6\}$ while high-stiff regimes take $\alpha = \{150, 110, 200, 50\}$ for OHO, NCFF, Duffing, AR respectively. Additionally for ODEs, the Radau solution is plotted as an accurate approximation and the vanilla PINN solution is provided for the high-stiff regimes.
• Figure 4: $L_2$ relative error of vanilla PINN, PINNsFormer and STL-PINN solving stiff-parameterized ODEs-PDEs as stiffness regime increases, with $\alpha \in [10, 200]$. The reported errors are the average of three separate runs, along with 90% confidence intervals.
• Figure 5: $L_2$ relative error of STL-PINN over increasing ranges of $\alpha$ during training, where $\alpha_{max}$ denotes the maximum value. The x-axis is the $\alpha$ transfer range. The reported errors are the average of three separate runs, along with 90% confidence intervals.