PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks

Abstract

Physics-Informed Neural Networks (PINNs) have emerged as a promising deep learning framework for approximating numerical solutions to partial differential equations (PDEs). However, conventional PINNs, relying on multilayer perceptrons (MLP), neglect the crucial temporal dependencies inherent in practical physics systems and thus fail to propagate the initial condition constraints globally and accurately capture the true solutions under various scenarios. In this paper, we introduce a novel Transformer-based framework, termed PINNsFormer, designed to address this limitation. PINNsFormer can accurately approximate PDE solutions by utilizing multi-head attention mechanisms to capture temporal dependencies. PINNsFormer transforms point-wise inputs into pseudo sequences and replaces point-wise PINNs loss with a sequential loss. Additionally, it incorporates a novel activation function, Wavelet, which anticipates Fourier decomposition through deep neural networks. Empirical results demonstrate that PINNsFormer achieves superior generalization ability and accuracy across various scenarios, including PINNs failure modes and high-dimensional PDEs. Moreover, PINNsFormer offers flexibility in integrating existing learning schemes for PINNs, further enhancing its performance.

1 Introduction

Numerically solving partial differential equations (PDEs) has been widely studied in science and engineering. The conventional approaches, such as finite element method (Bathe, 2007) or pseudo-spectral method (Fornberg, 1998), suffer from high computational costs in constructing meshes for high-dimensional PDEs. With the development of scientific machine learning, Physics-informed neural networks (PINNs) (Lagaris et al., 1998; Raissi et al., 2019) have emerged as a promising novel approach. Conventional PINNs and most variants employ multilayer perceptrons (MLP) as end-to-end frameworks for point-wise predictions, achieving remarkable success in various scenarios.

Nevertheless, recent works have shown that PINNs fail in scenarios when solutions exhibit high-frequency or multiscale features (Raissi, 2018; Fuks & Tchelepi, 2020; Krishnapriyan et al., 2021; Wang et al., 2022a), though the corresponding analytical solutions are simple. In such cases, PINNs tend to provide overly smooth or naive approximations, deviating from the true solution.

Existing approaches to mitigate these failures typically involve two general strategies. The first strategy, known as data interpolation (Raissi et al., 2017; Zhu et al., 2019; Chen et al., 2021), employs data regularization observed from simulations, or real-world scenarios. These approaches face challenges in acquiring ground truth data. The second strategy employs different training schemes (Mao et al., 2020; Krishnapriyan et al., 2021; Wang et al., 2021; 2022a), which potentially impose a high computational cost in practice. For instance, Seq2Seq by Krishnapriyan et al. (2021) requires training multiple neural networks sequentially, while other networks suffer from convergence issues due to error accumulation. Another method, Neural Tangent Kernel (NTK) (Wang et al., 2022a), involves constructing kernels $K\in {ℝ}^{D\times P}$, where $D$ is the sample size and $P$ is the model parameter, which suffers from scalability issues as the sample size or model parameter increases.

While most efforts to improve the generalization ability and address failure modes in PINNs have focused on the aforementioned aspects, the crucial temporal dependencies in real-world physical systems have been largely neglected. Finite Element Methods, for instance, implicitly incorporate temporal dependencies by sequentially propagating the global solution. This propagation relies on the principle that the state at time $t+\mathrm{\Delta }t$ depends on the state at time $t$. In contrast, PINNs, being a point-to-point framework, do not account for temporal dependencies within PDEs. Neglecting temporal dependencies poses challenges in globally propagating initial condition constraints in PINNs. Consequently, PINNs often exhibit failure modes where the approximations remain accurate near the initial condition but subsequently fail into overly smooth or naive approximations.

To address this issue of neglecting temporal dependencies in PINNs, a natural idea is employing Transformer-based models (Vaswani et al., 2017), which are known for capturing long-term dependencies in sequential data through multi-head self-attentions and encoder-decoder attentions. Variants of transformer-based models have shown substantial success across various domains. However, adapting the Transformer, which is inherently designed for sequential data, to the point-to-point framework of PINNs presents non-trivial challenges. These challenges span both the data representation and the regularization loss within the framework.

Main Contributions. In this work, we introduce PINNsFormer, a novel sequence-to-sequence PDE solver built on the Transformer architecture. To the best of our knowledge, PINNsFormer is the first framework in the realm of PINNs that explicitly focuses on and learns temporal dependencies within PDEs. Our key contributions can be summarized as follows:

New Framework: We propose a novel yet straightforward Transformer-based framework named PINNsFormer. This framework equips PINNs with the capability to capture temporal dependencies through the generated pseudo sequences, thereby enhancing the generalization ability and approximation accuracy in effectively solving PDEs.

Novel Activation: We introduce a novel non-linear activation function Wavelet. Wavelet is designed to anticipate the Fourier Transform for arbitrary target signals, making it a universal approximator for infinite-width neural networks. Wavelet can also be potentially beneficial to various deep learning tasks across different model architectures.

Extensive Experiments: We conduct comprehensive evaluations of PINNsFormer for various scenarios. We demonstrate its advantages in optimization and approximation accuracy when addressing the failure modes of PINNs or solving high-dimensional PDEs. Additionally, we showcase the flexibility of PINNsFormer in incorporating variant learning schemes of PINNs. We show the outperformance of PINNsFormer than PINNs with the schemes.

2 Related Work

Physics-Informed Neural Networks (PINNs). Physics-Informed Neural Networks (PINNs) have emerged as a promising approach to tackling scientific and engineering problems. Raissi et al. (2019) introduced the novel framework that incorporates physical laws into the neural network training to solve PDEs. This pioneering work has inspired subsequent investigations, leading to applications across diverse domains, including fluid dynamics, solid mechanics, and quantum mechanics (Ling et al., 2016; Carleo et al., 2019; Yang et al., 2020). Researchers also investigate different learning schemes of PINNs (Mao et al., 2020; Wang et al., 2021; 2022a). These strategies have yielded substantial improvements in convergence, generalization, and interpretability than PINNs.

Failure Modes of PINNs. Despite the promise exhibited by PINNs, recent works have indicated certain failure modes inherent to PINNs, particularly when confronted with PDEs featuring high-frequency or multiscale features (Fuks & Tchelepi, 2020; Raissi, 2018; McClenny & Braga-Neto, 2020; Krishnapriyan et al., 2021; Zhao et al., 2022; Wang et al., 2022a). This challenge has prompted investigations from various perspectives, including designing variant model architectures, learning schemes, or using data interpolations (Han et al., 2018; Lou et al., 2021; Wang et al., 2021; 2022a; 2022b). A comprehensive understanding of PINNs’ limitations and the underlying failure modes is fundamental for applications in addressing complicated physical problems.

Transformer-Based Models. The Transformer model (Vaswani et al., 2017) has achieved significant attention due to its ability to capture long-term dependencies, leading to major achievements in natural language processing tasks (Devlin et al., 2018; Radford et al., 2018). Transformers have also been extended to other domains, including computer vision, speech recognition, and time-series analysis (Liu et al., 2021; Dosovitskiy et al., 2020; Gulati et al., 2020; Zhou et al., 2021). Researchers have also developed techniques aimed at enhancing the efficiency of Transformers, such as sparse attention and model compression (Child et al., 2019; Sanh et al., 2019).

3 Methodology

Preliminaries: Let $\mathrm{\Omega }$ be an open set in ${ℝ}^d$, bounded by $\partial \mathrm{\Omega }\in {ℝ}^{d-1}$. The PDEs with spatial input $𝒙$ and temporal input $t$ generally fit the following abstraction:

$$\begin{array}{c}𝒟[u(𝒙,t)]=f(𝒙,t),\forall 𝒙,t\in \mathrm{\Omega }\\ ℬ[u(𝒙,t)]=g(𝒙,t),\forall 𝒙,t\in \partial \mathrm{\Omega }\end{array}$$ (1)

where $u$ is the PDE’s solution, $𝒟$ is the differential operator that regularizes the behavior of the system, and $ℬ$ describes the boundary or initial conditions in general. Specifically, $\{𝒙,t\}\in \mathrm{\Omega }$ are residual points, and $\{𝒙,t\}\in \partial \mathrm{\Omega }$ are boundary/initial points. Let $\widehat{u}$ be neural network approximations, PINNs describe the framework where $\widehat{u}$ is empirically regularized by the following constraints:

$$\begin{array}{c}{ℒ}_{\text{𝙿𝙸𝙽𝙽𝚜}}={\lambda }_{\text{𝑟}}\sum _{i=1}^{N_{\text{𝑟}}}{\Vert 𝒟[\widehat{u}(𝒙,t)]-f(𝒙,t)\Vert }^2+{\lambda }_b\sum _{i=1}^{N_b}{\Vert ℬ[\widehat{u}(𝒙,t)]-g(𝒙,t)\Vert }^2\end{array}$$ (2)

where $N_b,N_r$ refer to the residual and boundary/initial points separately, ${\lambda }_r,{\lambda }_b$ are the regularization parameters that balance the emphasis of the loss terms. The neural network $\widehat{u}$ takes vectorized $\{𝒙,t\}$ as input and outputs the approximated solution. The goal is then to use machine learning methodologies to train the neural network $\widehat{u}$ that minimizes the loss in Equation 2.

Methodology Overview: While PINNs focus on point-to-point predictions, the exploration of temporal dependencies in real-world physics systems has been merely neglected. Conventional PINNs methods employ a single pair of spatial information $𝒙$ and temporal information $t$ to approximate the numerical solution $u(𝒙,t)$, without accounting for temporal dependencies across previous or subsequent time steps. However, this simplification is only applicable to elliptic PDEs, where the relationships between unknown functions and their derivatives do not explicitly involve time. In contrast, hyperbolic and parabolic PDEs incorporate time derivatives, implying that the state at one time step can influence states at preceding or subsequent time steps. Consequently, considering temporal dependencies is crucial to effectively address these PDEs using PINNs.

In this section, we introduce a novel framework featuring a Transformer-based model of PINNs, namely PINNsFormer. Unlike point-to-point predictions, PINNsFormer extends PINNs’ capabilities to sequential predictions. PINNsFormer allows accurately approximating solutions at specific time steps while also learning and regularizing temporal dependencies among incoming states. The framework consists of four components: Pseudo Sequence Generator, Spatio-Temporal Mixer, Encoder-Decoder with multi-head attention, and an Output Layer. Additionally, we introduce a novel activation function, named Wavelet, which employs Real Fourier Transform techniques to anticipate solutions to PDEs. The framework diagram is exhibited in Figure 1. We provide detailed explanations of each framework component and learning schemes in the following subsections.

Figure 1: Architecture of proposed PINNsFormer. PINNsFormer generates a pseudo sequence based on pointwise input features. It outputs the corresponding sequential approximated solution. The first approximation of the sequence is the desired solution $\widehat{u}(𝒙,t)$.

3.1 Pseudo Sequence Generator

While Transformers and Transformer-based models are designed to capture long-term dependencies in sequential data, conventional PINNs utilize non-sequential data as inputs for neural networks. Consequently, to incorporate PINNs with Transformer-based models, it is essential to transform the pointwise spatiotemporal inputs into temporal sequences. Thus, for a given spatial input $𝒙\in {ℝ}^{d-1}$ and temporal input $t\in ℝ$, the Pseudo Sequence Generator performs the following operations:

$$[𝒙,t]\stackrel{\text{𝙶𝚎𝚗𝚎𝚛𝚊𝚝𝚘𝚛}}{\Rightarrow }\{[𝒙,t],[𝒙,t+\mathrm{\Delta }t],\mathrm{\dots },[𝒙,t+(k-1)\mathrm{\Delta }t]\}$$ (3)

where $[\cdot ]$ is the concatenation operation, such that $[𝒙,t]\in {ℝ}^d$ is vectorized, and the generator outputs the pseudo sequence in the shape of ${ℝ}^{k\times d}$. In short, the Pseudo Sequence Generator extrapolates sequential time series by extending a single spatiotemporal input to multiple isometric discrete time steps. $k$ and $\mathrm{\Delta }t$ are hyperparameters, which intuitively determine how many steps the pseudo sequence needs to ’look ahead’ and how ’far’ each step should be. In practice, both $k$ and $\mathrm{\Delta }t$ should not be set to very large scales, as larger $k$ can cause heavy computational and memory overheads, while larger $\mathrm{\Delta }t$ may undermine the time dependency relationships of neighboring discrete time steps.

3.2 Model Architecture

In addition to the Pseudo Sequence Generator, PINNsFormer consists of three components of its architecture: Sptio-Temporal Mixer, Encoder-Decoder with multi-head attentions, and Output Layer. The Output Layer is straightforward to interpret as a fully-connected MLP appended to the end. We provide detailed insights into the first two components below. Notably, PINNsFormer relies only on linear layers and non-linear activations, avoiding complex operations such as convolutional or recurrent layers. This design preserves PINNsFormer’s computational efficiency in practice.

Spatio-Temporal Mixer. Most PDEs contain low-dimensional spatial or temporal information. Directly feeding low-dimensional data to encoders may fail to capture the complex relationships between each feature dimension. Hence, it is necessary to embed original sequential data in higher-dimensional spaces such that more information is encoded into each vector.

Instead of embedding raw data in a high-dimensional space where the distance between vectors reflects the semantic similarity (Vaswani et al., 2017; Devlin et al., 2018), PINNsFormer constructs a linear projection that maps spatiotemporal inputs onto a higher-dimensional space using a fully-connected MLP. The embedded data enriches the capability of information by mixing all raw spatiotemporal features together, so-called the linear projection Spatio-Temporal Mixer.

Figure 2: The architecture of PINNsFormer’s Encoder-Decoder Layers. The decoder is not equipped with self-attentions.

Encoder-Decoder Architecture. PINNsFormer employs an encoder-decoder architecture similar to Transformer. The encoder consists of a stack of identical layers, each of which contains an encoder self-attention layer and a feedforward layer. The decoder is slightly different from the vanilla Transformer, where each of the identical layers contains only an encoder-decoder self-attention layer and a feedforward layer. At the decoder level, PINNsFormer uses the same spatiotemporal embeddings as the encoder. Therefore, the decoder does not need to relearn dependencies for the same input embeddings. The diagram for the encoder-decoder architecture is shown in Figure 2

Intuitively, the encoder self-attentions allow learning the dependency relationships of all spatiotemporal information. The decoder encoder-decoder attentions allow selectively focusing on specific dependencies within the input sequence during the decoding process, enabling it to capture more information than conventional PINNs. We use the same embeddings for the encoder and decoder since PINNs focus on approximating the solution of the current state, in contrast to next state prediction in language tasks or time series forecastings.

3.3 Wavelet Activation

While Transformers typically employ LayerNorm and ReLU non-linear activation functions (Vaswani et al., 2017; Gehring et al., 2017; Devlin et al., 2018), these activation functions might not always be suitable in solving PINNs. In particular, employing ReLU activation in PINNs can result in poor performance, whose effectiveness relies heavily on the accurate evaluation of derivatives while ReLU has a discontinuous derivative (Haghighat et al., 2021; de Wolff et al., 2021). Recent studies utilize Sin activation for specific scenarios to mimic the periodic properties of PDEs’ solutions (Li et al., 2020; Jagtap et al., 2020; Song et al., 2022). However, it requires strong prior knowledge of the solution’s behavior and is limited in its applicability. Tackling this issue, we proposed a novel and simple activation function, namely Wavelet, defined as follows:

$$\text{𝚆𝚊𝚟𝚎𝚕𝚎𝚝}(𝒙)={\omega }_1\sin (𝒙)+{\omega }_2\cos (𝒙)$$ (4)

Where ${\omega }_1$ and ${\omega }_2$ are registered learnable parameters. The intuition behind Wavelet activation simply follows Real Fourier Transform: While periodic signals can be decomposed into an integral of sines of multiple frequencies, all signals, whether periodic or aperiodic, can be decomposed into an integral of sines and cosines of varying frequencies. It is evident that Wavelet can approximate arbitrary functions giving sufficient approximation power, which leads to the following proposition:

Proposition 1

Let $𝒩$ be a two-hidden-layer neural network with infinite width, equipped with Wavelet activation function, then $𝒩$ is a universal approximator for any real-valued target f.

Proof sketch: The proof follows the Real Fourier Transform (Fourier Integral Transform). For any given input $x$ and its corresponding real-valued target $f(x)$, it has the Fourier Integral:

$$f(x)={\int }_{-\mathrm{\infty }}^{\mathrm{\infty }}{F}_c(\omega )\cos (\omega x)𝑑\omega +{\int }_{-\mathrm{\infty }}^{\mathrm{\infty }}{F}_s(\omega )\sin (\omega x)𝑑\omega$$

where $F_c$ and $F_s$ are the coefficients of Sines and Cosines respectively. Second, by Riemann sum approximation, the integral can be approximated by the infinite sum such that:

$$f(x)\approx \sum _{n=1}^N\left[F_c({\omega }_n)\cos ({\omega }_{n}x)+F_s({\omega }_n)\sin ({\omega }_{n}x)\right]\equiv W_2(\text{𝚆𝚊𝚟𝚎𝚕𝚎𝚝}(W_{1}x))$$

where $W_1$ and $W_2$ are the weights of $𝒩$’s first and second hidden layer. As $W_1$ and $W_2$ are infinite-width, we can divide the piecewise summation into infinitely small intervals, making the approximation arbitrarily close to the true integral. Hence, $𝒩$ is a universal approximator for any given $f$. In practice, most PDE solutions contain only a finite number of major frequencies. Using a neural network with finite parameters would also lead to proper approximations of the true solutions.

Although Wavelet activation function is primarily employed by PINNsFormer to improve PINNs in our work, it may have potential applications in other deep-learning tasks. Similar to ReLU, $\sigma (\cdot )$, and Tanh activations, which all turn infinite-width two-hidden-layer neural networks into universal approximators (Cybenko, 1989; Hornik, 1991; Glorot et al., 2011), we anticipate that Wavelet can demonstrate its effectiveness in other applications beyond the scope of this work.

3.4 Learning Scheme

While conventional PINNs focus on point-to-point predictions, there is an unexplored realm in adapting PINNs to handle pseudo-sequential inputs. In PINNsFormer, each generated point in the sequence, i.e., $[{𝒙}_i,t_i+j\mathrm{\Delta }t]$, is forwarded to the corresponding approximation, i.e., $\widehat{u}({𝒙}_i,t_i+j\mathrm{\Delta }t)$ for any . This approach allows us to compute the $n$th-order gradients with respect to $𝒙$ or $t$ independently for any valid $n$. For instance, for any given input pseudo sequence $\{[{𝒙}_i,t_i],[{𝒙}_i,t_i+\mathrm{\Delta }t],\mathrm{\dots },[{𝒙}_i,t_i+(k-1)\mathrm{\Delta }t]\}$, and the corresponding approximations $\{\widehat{u}({𝒙}_i,t_i),\widehat{u}({𝒙}_i,t_i+\mathrm{\Delta }t),\mathrm{\dots },\widehat{u}({𝒙}_i,t_i+(k-1)\mathrm{\Delta }t)\}$, we can compute the first-order derivatives w.r.t. $𝒙$ and $t$ separately as follows:

$$\begin{array}{c}\frac{\partial {\{\widehat{u}({𝒙}_i,t_i+j\mathrm{\Delta }t)\}}_{j=0}^{k-1}}{\partial {\{t_i+j\mathrm{\Delta }t\}}_{j=0}^{k-1}}=\{\frac{\partial \widehat{u}({𝒙}_i,t_i)}{\partial t_i},\frac{\partial \widehat{u}({𝒙}_i,t_i+\mathrm{\Delta }t)}{\partial (t_i+\mathrm{\Delta }t)},\mathrm{\dots },\frac{\partial \widehat{u}({𝒙}_i,t_i+(k-1)\mathrm{\Delta }t)}{\partial (t_i+(k-1)\mathrm{\Delta }t)}\}\\ \frac{\partial {\{\widehat{u}({𝒙}_i,t_i+j\mathrm{\Delta }t)\}}_{j=0}^{k-1}}{\partial {𝒙}_i}=\{\frac{\partial \widehat{u}({𝒙}_i,t_i)}{\partial {𝒙}_i},\frac{\partial \widehat{u}({𝒙}_i,t_i+\mathrm{\Delta }t)}{\partial {𝒙}_i},\mathrm{\dots },\frac{\partial \widehat{u}({𝒙}_i,t_i+(k-1)\mathrm{\Delta }t)}{\partial {𝒙}_i}\}\end{array}$$ (5)

This scheme for calculating the gradients of sequential approximations with respect to sequential inputs can be easily extended to higher-order derivatives and is applicable to residual, boundary, and initial points. However, unlike the general PINNs optimization objective in Equation 2, which combines initial and boundary condition objectives, PINNsFormer distinguishes between the two and applies different regularization schemes to initial and boundary conditions through its learning scheme. For residual and boundary points, all sequential outputs can be regularized using the PINNs loss. This is because all generated pseudo-timesteps are within the same domain as their original inputs. For example, if $[{𝒙}_i,t_i]$ is sampled from the boundary, then $[{𝒙}_i,t_i+j\mathrm{\Delta }t]$ also lies on the boundary for any $j\in {\mathbb{N}}^{+}$. In contrast, for initial points, only the $t=0$ condition is regularized, corresponding to the first element of the sequential outputs. This is because only the first element of the pseudo-sequence exactly matches the initial condition at $t=0$. All other generated time steps have $t=j\mathrm{\Delta }t$ for any $j\in {\mathbb{N}}^{+}$, which fall outside the initial conditions.

By these considerations, we adapt the PINNs loss to the sequential version, as described below:

$$\begin{array}{c}{ℒ}_{\text{𝑟𝑒𝑠}}=\frac{1}{k{N}_{\text{𝑟𝑒𝑠}}}\sum _{i=1}^{N_{\text{𝑟𝑒𝑠}}}\sum _{j=0}^{k-1}{\Vert 𝒟[\widehat{u}({𝒙}_i,t_i+j\mathrm{\Delta }t)]-f({𝒙}_i,t_i+j\mathrm{\Delta }t)\Vert }^2\\ {ℒ}_{\text{𝑏𝑐}}=\frac{1}{k{N}_{\text{𝑏𝑐}}}\sum _{i=1}^{N_{\text{𝑏𝑐}}}\sum _{j=0}^{k-1}{\Vert ℬ[\widehat{u}({𝒙}_i,t_i+j\mathrm{\Delta }t)]-g({𝒙}_i,t_i+j\mathrm{\Delta }t)\Vert }^2\\ {ℒ}_{\text{𝑖𝑐}}=\frac{1}{N_{\text{𝑖𝑐}}}\sum _{i=1}^{N_{\text{𝑏𝑐}}}{\Vert ℐ[\widehat{u}({𝒙}_i,0)]-h({𝒙}_i,0)\Vert }^2\\ {ℒ}_{\text{𝙿𝙸𝙽𝙽𝚜𝙵𝚘𝚛𝚖𝚎𝚛}}={\lambda }_{\text{𝑟𝑒𝑠}}{ℒ}_{\text{𝑟𝑒𝑠}}+{\lambda }_{\text{𝑖𝑐}}{ℒ}_{\text{𝑖𝑐}}+{\lambda }_{\text{𝑏𝑐}}{ℒ}_{\text{𝑏𝑐}}\end{array}$$ (6)

where $N_{\text{𝑟𝑒𝑠}}=N_r$ refers to the residual points as in Equation 2, $N_{\text{𝑏𝑐}},N_{\text{𝑖𝑐}}$ represent the number of boundary and initial points, respectively, with $N_{\text{𝑏𝑐}}+N_{\text{𝑖𝑐}}=N_b$. ${\lambda }_{\text{𝑟𝑒𝑠}}$, ${\lambda }_{\text{𝑏𝑐}}$, and ${\lambda }_{\text{𝑖𝑐}}$ are regularization weights that balance the importance of the loss terms in PINNsFormer, similar to the PINNs loss.

During training, PINNsFormer forwards all residual, boundary, and initial points to obtain their corresponding sequential approximations. It then optimizes the modified PINNs loss ${ℒ}_{\text{𝙿𝙸𝙽𝙽𝚜𝙵𝚘𝚛𝚖𝚎𝚛}}$ in Equation 6 using gradient-based optimization algorithms such as L-BFGS or Adam, updating the model parameters until convergence. In the testing phase, PINNsFormer forwards any arbitrary pair $[𝒙,t]$ to observe the sequential approximations, where the first element of the sequential approximation corresponds exactly to the desired value of $\widehat{u}(𝒙,t)$.

3.5 Loss Landscape Analysis

Figure 3: Visualization of the loss landscape for PINNs (left) and PINNsFormer (right) on a logarithmic scale. The loss landscape of PINNsFormer is significantly smoother than conventional PINNs.

While achieving theoretical convergence or establishing generalization bounds for Transformer-based models can be challenging, an alternative approach to assess optimization trajectory is through visualization of the loss landscape. This approach has been employed in the analysis of both Transformers and PINNs (Krishnapriyan et al., 2021; Yao et al., 2020; Park & Kim, 2022). The loss landscape is constructed by perturbing the trained model along the directions of the first two dominant Hessian eigenvectors. This technique is often more informative than random parameter perturbations. In general, a smoother loss landscape with fewer local minima indicates an easier case for the model to converge to the global minimum. We visualize the loss landscape for both conventional PINNs and PINNsFormer. The visualizations are presented in Figure 5.

The visualizations clearly reveal that PINNs exhibit a more compliated loss landscape than PINNsFormer. To be specific, we estimate the Lipschitz constant for both loss landscapes. We find that $L_{\text{𝙿𝙸𝙽𝙽𝚜}}=776.16$, which is significantly larger than $L_{\text{𝙿𝙸𝙽𝙽𝚜𝙵𝚘𝚛𝚖𝚎𝚛}}=32.79$. Furthermore, the loss landscape of PINNs exhibits several sharp cones near its optimal point, indicating the presence of multiple local minima in close proximity to the convergence point (zero perturbation). The rugged loss landscape and multiple local minima of conventional PINNs suggest that optimizing the objective described in Equation 6 for PINNsFormer offers an easier path to reach the global minimum. This implies that PINNsFormer has advantages in avoiding the failure modes associated with PINNs. The analysis is further validated by empirical experiments, as shown in the following section.

4 Experiments

4.1 Setup

Goal. Our empirical evaluations aim to demonstrate three key advantages of PINNsFormer. Firstly, we show that PINNsFormer improves generalization abilities and mitigates failure modes compared to PINNs and variant architectures. Secondly, we illustrate the flexibility of PINNsFormer in incorporating various learning schemes, resulting in superior performance. Thirdly, we provide evidence of PINNsFormer’s faster convergence and improved generalization capabilities in solving high-dimensional PDEs, which can be challenging for PINNs and their variants.

PDE Setup. We use a selection of PDEs, including convection, 1D-reaction, 1D-wave, and Navier–Stokes PDEs. The setups for the convection, 1D-reaction, and 1D-wave equations follow past work (Chen et al., 2018). For baseline model training, including PINNs, QRes (Bu & Karpatne, 2021), and First-Layer Sine (FLS) (Wong et al., 2022), we uniformly sampled $N_{\text{𝑖𝑐}}=N_{\text{𝑏𝑐}}=101$ initial and boundary points, as well as a uniform grid of $101\times 101$ mesh points for the residual domain, resulting in total $N_{\text{𝑟𝑒𝑠}}=10201$ points. In the case of training PINNsFormer, we reduce the collocation points, with $N_{\text{𝑖𝑐}}=N_{\text{𝑏𝑐}}=51$ initial and boundary points and a $51\times 51$ mesh for residual points. The reduction in fewer training samples serves two purposes: it enhances training efficiency and allows us to demonstrate the generalization capabilities of PINNsFormer with limited training data. For testing, we employed a $101\times 101$ mesh within the residual domain.

For the Navier–Stokes PDE, our experiment follows the established setup in Raissi et al. (2017). We sampled 2500 points from the 3D mesh within the residual domain for training purposes. The evaluation was performed by testing the predicted pressure at the final time step $t=20.0$.

Training and Testing. We build PINNs, QRes, and FLS as the baselines, along with the proposed PINNsFormer. We maintain approximately close numbers of parameters across all models to highlight the advantages of PINNsFormer from its ability to capture temporal dependencies rather than relying solely on model overparameterization. We train all models using the L-BFGS optimizer with Strong Wolfe linear search for 1000 iterations. For simplicity, we set ${\lambda }_{\text{𝑟𝑒𝑠}}={\lambda }_{\text{𝑖𝑐}}={\lambda }_{\text{𝑏𝑐}}=1$ for the optimization objective in Equation 6. Detailed hyperparameters are provided in Appendix A.

In terms of evaluation, we adopted commonly used metrics in related works (Krishnapriyan et al., 2021; Raissi et al., 2019; McClenny & Braga-Neto, 2020), including the relative Mean Absolute Error (rMAE or relative ${\mathrm{\ell }}_1$ error) and the relative Root Mean Square Error (rRMSE or relative ${\mathrm{\ell }}_2$ error). The detailed formulations of the metrics are provided in Appendix A.

Reproducibility. All models are implemented in PyTorch (Paszke et al., 2019), and are trained separately on a single NVIDIA Tesla V100 GPU. All code and demos are included and reproducible at https://github.com/AdityaLab/pinnsformer.

4.2 Mitigating Failure Modes of PINNs

Our primary evaluation focuses on demonstrating the superior generalization ability of PINNsFormer in comparison to PINNs, particularly on PDEs that are known to challenge PINNs’ generalization capabilities. We focus on solving two distinct types of PDEs: the convection equation and the 1D-reaction equation. These equations pose significant challenges for conventional MLP-based PINNs, often resulting in what is referred to as ”PINNs failure modes” (Mojgani et al., 2022; Daw et al., 2022; Krishnapriyan et al., 2021). In these failure modes, optimization gets stuck in local minima, leading to overly smooth approximations that deviate from the true solutions.

The objective of our evaluation is to showcase the enhanced generalization capabilities of PINNsFormer when compared to standard PINNs and their variations, specifically in addressing PINNs’ failure modes. The evaluation results are summarized in Table 1, with detailed PDE formulations provided in Appendix B. We showcase the prediction and absolute error plots of PINNs and PINNsFormer on convection equation in Figure 4, all prediction plots available in Appendix C.

Model	Convection			1D-Reaction
	Loss	rMAE	rRMSE	Loss	rMAE	rRMSE
PINNs	0.016	0.778	0.840	0.199	0.982	0.981
QRes	0.015	0.746	0.816	0.199	0.979	0.977
FLS	0.012	0.674	0.771	0.199	0.984	0.985
PINNsFormer	3.7e-5	0.023	0.027	3.0e-6	0.015	0.030

Table 1: Results for solving connection and 1D-reaction equations. PINNsFormer consistently outperforms all baseline methods in terms of training loss, rMAE, and rRMSE.

The evaluation results demonstrate significant outperformance of PINNsFormer over all baselines for both scenarios. PINNsFormer achieves the lowest training loss and test errors, distinguishing PINNsFormer as the only approach capable of mitigating the failure modes. In contrast, all other baseline methods remain stuck at global minima and fail to optimize the objective loss effectively. These results show the clear advantages of PINNsFormer in terms of generalization ability and approximation accuracy when compared to conventional PINNs and existing variants.

Figure 4: Prediction (left) and absolute error (right) of PINNs (up) and PINNsFormer (bottom) on convection equation. PINNsFormer shows success in mitigating the failure mode than PINNs.

An additional consideration when assessing PINNsFormer is its computational and memory overheads relative to PINNs. While MLP-based PINNs are known for their efficiency, PINNsFormer, with Transformer-based architecture in handling sequential data, naturally incurs higher computational and memory costs. Nonetheless, our empirical evaluation indicates that the overhead is tolerable, benefitting from the reliance on only linear hidden layers with non-linear activation functions, avoiding complicated operators such as convolutional or recurrent layers. For instance, when setting the pseudo-sequence length $k=5$, we observe an approximate 2.92x rise in computational cost and a 2.15x rise in memory usage (as detailed in Appendix A). These overheads are deemed acceptable in exchange for the substantial performance improvements achieved by PINNsFormer.

While PINNs and their various architectural adaptations may encounter challenges for certain scenarios, prior research has explored sophisticated optimization schemes to mitigate these issues, including learning rate annealing (Wang et al., 2021), augmented Lagrangian methods (Lu et al., 2021), and neural tangent kernel approaches (Wang et al., 2022a). These modified PINNs have shown significant improvement of PINNs under certain scenarios. Notably, when these optimization strategies are applied to PINNsFormer, they can be easily incorporated to achieve further performance improvements. For instance, the Neural Tangent Kernel (NTK) method to PINNs has shown success in solving the 1D-wave equation. In such context, we demonstrate that when combining NTK with PINNsFormer, we can achieve further outperformance in approximation accuracy. Detailed results are presented in Table 2, and comprehensive PDE formulations are available in Appendix B, with prediction plots in Appendix C.

4.3 Flexibility in Incorporating Variant Learning Schemes

Model	1D-Wave
	Loss	rMAE	rRMSE
PINNs	1.93e-2	0.326	0.335
PINNsFormer	1.38e-2	0.270	0.283
PINNs + NTK	6.34e-3	0.140	0.149
PINNsFormer + NTK	4.21e-3	0.054	0.058

Table 2: Results for solving the 1D-wave equation, incorporating the NTK method. PINNsFormer combined with NTK outperforms all other methods on all metrics.

Our evaluation results show both the flexibility and effectiveness of incorporating PINNsFormer with the NTK method. In particular, we observe a sequence of performance improvements, from standard PINNs to PINNsFormer and from PINNs+NTK to PINNsFormer+NTK. Essentially, PINNsFormer explores a variant architecture of PINNs, while many learning schemes are designed from an optimization perspective and are agnostic to neural network architectures. This inherent flexibility allows for versatile combinations of PINNsFormer with various learning schemes, offering practical and customizable solutions for accurate solutions in real-world applications.

4.4 Generalization on High-Dimensional PDEs

Figure 5: Training loss vs. Iterations of PINNs and PINNsFormer on the Navier-Stokes equation.

In the previous sections, we demonstrated the clear benefits of PINNsFormer in generalizing the solutions for PINNs failure modes. However, those PDEs often have simple analytical solutions. In practical physics systems, higher-dimensional and more complex PDEs need to be solved. Therefore, it’s important to evaluate the generalization ability of PINNsFormer on such high-dimensional PDEs, especially when PINNsFormer is equipped with advanced mechanisms like self-attention.

We evaluate the performance of PINNsFormer compared to PINNs on the 2D Navier-Stokes PDE, a problem previously investigated by Raissi et al. (2019). The training loss is shown in Figure 5, and the evaluation results are shown in Table 3. The detailed formulations of the 2D Navier-Stokes equation can be found in Appendix B, and the prediction plots are provided in Appendix C.

Model	Navier-Stokes
	Loss	rMAE	rRMSE
PINNs	6.72e-5	13.08	9.08
QRes	2.24e-4	6.41	4.45
FLS	9.54e-6	3.98	2.77
PINNsFormer	6.66e-6	0.384	0.280

Table 3: Results for solving 2D Navier-Stokes equation, PINNsFormer outperforms all baselines on all metrics.

The evaluation results demonstrate clear advantages of PINNsFormer over PINNs on high-dimensional PDEs. Firstly, PINNsFormer outperforms PINNs and their MLP-variants in terms of both training loss and validation errors. Firstly, PINNsFormer exhibits significantly faster convergence during training, which compensates for the higher computational cost per iteration. Secondly, while PINNs and their MLP-variants predict the pressure with good shapes, they exhibit increasing magnitude discrepancies as time increases. In contrast, PINNsFormer consistently aligns both the shape and magnitude of predicted pressures across various time intervals. This consistency is attributed to PINNsFormer’s ability to learn temporal dependencies through Transformer-based model architecture and self-attention mechanism.

5 Conclusion

In this paper, we introduced PINNsFormer, a novel Transformer-based framework, designed as an extension of PINNs, aimed at capturing temporal dependencies when approximating solutions to PDEs. To adapt conventional PINNs to Transformer-based models, we introduced the Pseudo Sequence Generator, a mechanism that translates vectorized inputs into pseudo time sequences, and incorporated a modified Encoder-Decoder layer along with a novel Wavelet activation function. Our empirical evaluations demonstrate that PINNsFormer consistently outperforms conventional PINNs across various scenarios, including handling PINNs’ failure modes and addressing high-dimensional PDEs. Furthermore, PINNsFormer retains computational simplicity by using only linear layers with non-linear activation functions, making it a practical choice for real-world applications. It can flexibly integrate with existing learning schemes for PINNs, leading to superior performance.

Beyond its application in PINNsFormer, the novel Wavelet activation function holds promises for the broader machine learning community. We provided a sketch proof demonstrating Wavelet’s ability to approximate arbitrary target solutions using a two-hidden-layer infinite-width neural network, leveraging the Fourier decomposition of these solutions. We encourage further exploration, both theoretically and empirically, of the Wavelet activation function’s potential. Its applicability extends beyond PINNs and can be leveraged in various architectures and applications.

Appendix A Appendix A: Model Hyperparameters

Model Hyperparameters. We provide a detailed set of hyperparameters used to obtain the experiment results, shown in Table 4.

Model	Hyperparameters	Value	Model Parameters
PINNs & FLS	hidden layer	4	527k
	hidden size	512
QRes	hidden layer	4	397k
	hidden size	256
PINNsFormer	$k$	5	454k
	$\mathrm{\Delta }t$	1e-3, 1e-4
	# of encoder	1
	# of decoder	1
	embedding size	32
	head	2
	hidden size	512

Table 4: Hyperparameters for Main Results

Training Overhead. We compare the training overhead of PINNsFormer over PINNs, as PINNs are known as an efficient framework while Transformer-based models are known for being computationally costly. The comparison relies on solving the Convection PDEs, which are detailed in Table 5. Here, we vary the hyperparameter of pseudo-sequence length $k$ for validation purposes. In practice, we set $k=5$ for all the empirical experiments in this paper.

Model		Training Time (sec/epoch)	Computational Overhead	GPU Memory (MiB)	Memory Overhead
PINNs		0.80	/	1311	/
PINsFormer	$k$=3	2.10	2.62x	2207	1.68x
	$k$=5	2.34	2.92x	2827	2.15x
	$k$=10	3.10	3.87x	4803	3.66x

Table 5: Overhead of PINNsFormer than PINNs in varying pseudo-sequence length. Both computational and memory overhead are tolerable and grow approximately linearly as $k$ increases

Evaluation Metrics. We present the detailed formula of rMAE and rRMSE as the following:

$$\begin{array}{c}\text{𝚛𝙼𝙰𝙴}=\frac{{\sum }_{n=1}^N|\widehat{u}(x_n,t_n)-u(x_n,t_n)|}{{\sum }_{n=1}^{N_{\text{𝑟𝑒𝑠}}}|u(x_n,t_n)|}\\ \text{𝚛𝚁𝙼𝚂𝙴}=\sqrt{\frac{{\sum }_{n=1}^N{|\widehat{u}(x_n,t_n)-u(x_n,t_n)|}^2}{{\sum }_{n=1}^N{|u(x_n,t_n)|}^2}}\end{array}$$ (7)

where $N$ is the number of testing points, $\widehat{u}$ is the neural network approximation, and $u$ is the ground truth.

Appendix B Appendix B: PDEs setups

We provide detailed PDE setups for convection, reaction-diffusion, and 1D-reaction equations.

Convection PDE. The one-dimensional convection problem is a hyperbolic PDE that is commonly used to model transport phenomena. The system has the formulation with periodic boundary conditions as follows:

$$\begin{array}{c}\frac{\partial u}{\partial t}+\beta \frac{\partial u}{\partial x}=0,\forall x\in [0,2\pi ],t\in [0,1]\\ \mathtt{\text{IC:}}u(x,0)=\sin (x),\mathtt{\text{BC:}}u(0,t)=u(2\pi ,t)\end{array}$$ (8)

where $\beta$ is the convection coefficient. As $\beta$ increases, the frequency of its solution goes higher, and it becomes harder for PINNs to approximate. Here, we set $\beta =50$.

1D-Reaction PDE. The one-dimensional reaction problem is a hyperbolic PDE that is commonly used to model chemical reactions. The system has the formulation with periodic boundary conditions as follows:

$$\begin{array}{c}\frac{\partial u}{\partial t}-\rho u(1-u)=0,\forall x\in [0,2\pi ],t\in [0,1]\\ \mathtt{\text{IC:}}u(x,0)=\exp (-\frac{{(x-\pi )}^2}{2{(\pi /4)}^2}),\mathtt{\text{BC:}}u(0,t)=u(2\pi ,t)\end{array}$$ (9)

where $\rho$ is the reaction coefficient. Here, we set $\rho =5$. The equation has a simple analytical solution:

$$u_{\text{𝚊𝚗𝚊𝚕𝚢𝚝𝚒𝚌𝚊𝚕}}=\frac{h(x)\exp (\rho t)}{h(x)\exp (\rho t)+1-h(x)}$$ (10)

where $h(x)$ is the function of the initial condition.

1D-Wave PDE. The 1D-Wave equation is a hyperbolic PDE that is used to describe the propagation of waves in one spatial dimension. It is often used in physics and engineering to model various wave phenomena, such as sound waves, seismic waves, and electromagnetic waves. The system has the formulation with periodic boundary conditions as follows:

$$\begin{array}{c}\frac{{\partial }^{2}u}{\partial t^2}-\beta \frac{{\partial }^{2}u}{\partial x^2}=0\forall x\in [0,1],t\in [0,1]\\ \mathtt{\text{IC:}}u(x,0)=\sin (\pi x)+\frac{1}{2}\sin (\beta \pi x),\frac{\partial u(x,0)}{\partial t}=0\\ \mathtt{\text{BC:}}u(0,t)=u(1,t)=0\end{array}$$ (11)

where $\beta$ is the wave speed. Here, we are specifying $\beta =3$.The equation has a simple analytical solution:

$$\begin{array}{c}u(x,t)=\sin (\pi x)\cos (2\pi t)+\frac{1}{2}\sin (\beta \pi x)\cos (2\beta \pi t)\end{array}$$ (12)

2D Navier-Stokes PDE. The 2D Navier-Stokes equation is a parabolic PDE that consists of a pair of partial differential equations that describe the behavior of incompressible fluid flow in two-dimensional space. They are widely used in fluid dynamics to model the motion of fluids, such as air and water, in various engineering and scientific applications. The system has the formulation as follows:

$$\begin{array}{c}\frac{\partial u}{\partial t}+{\lambda }_1(u\frac{\partial u}{\partial x}+v\frac{\partial u}{\partial y})=-\frac{\partial p}{\partial x}+{\lambda }_2(\frac{{\partial }^{2}u}{\partial x^2}+\frac{{\partial }^{2}u}{\partial v^2})\\ \frac{\partial v}{\partial t}+{\lambda }_1(u\frac{\partial v}{\partial x}+v\frac{\partial v}{\partial y})=-\frac{\partial p}{\partial y}+{\lambda }_2(\frac{{\partial }^{2}u}{\partial x^2}+\frac{{\partial }^{2}u}{\partial v^2})\end{array}$$ (13)

where $u(t,x,y)$ and $v(t,x,y)$ are the $x$-component and $y$-component of the velocity field separately, and $p(t,x,y)$ is the pressure. Here, we set ${\lambda }_1=1$ and ${\lambda }_2=0.01$. The system does not have an explicit analytical solution, while the simulated solution is given by Raissi et al. (2019).

Appendix C Appendix C: Additional Results

We here present the plots of ground truth solutions, neural network predictions, and absolute errors for all evaluations included in the experimental section. The plots on convection, 1D-reaction, 1D-wave, and 2D Navier-Stokes equations are shown in Figure separately.