Paper: PINNsFormer: A Transformer-Based Framework For Physics-Informed Neural Networks (ICLR-2024)

Reviewer comment (Experiments):
As a follow-up comment, we sample many random collocation points at each iteration, and MLP can see many time coordinates with spatial coordinates during training. Hence, I believe the original PINN is capable of modeling time dependency.

Author response:
We respectfully disagree with the reviewer’s statements regarding capturing temporal dependencies with MLP-based PINNs. While it's true that MLP-based PINNs incorporate both spatial and temporal inputs, MLP-based PINNs process different inputs (i.e., batch data) independently. For standard MLP architectures, each spatio-temporal input (even with random samplings of collocation points) is treated as a distinct entity, and the model processes them separately without explicitly capturing the temporal dependencies between consecutive time steps. In contrast, PINNsFormer explicitly suggests capturing temporal dependencies within local consecutive time steps when processing each spatiotemporal input. In our experiment section, we highlight the function of PINNsFormer in explicitly capturing temporal dependencies rather than MLP-based PINNs, with clear performance benefits by empirical evaluations.

Reviewer comment (Writing):
If I understand correctly, it seems that Spatio-Temporal Mixer is just one layer MLP. I might have missed something, but the technical details are not properly described in explaining each component. The formal definition of each module would be appreciated.

Author response:
Yes, Spatio-Temporal Mixer in our implementation is just a one-layer MLP. The Spatio-Temporal Mixer module serves as an embedding process as in the vanilla Transformer. We show that a one-layer MLP is enough and effective to embed low-dimensional spatiotemporal information in high-dimensional spaces based on empirical evaluations.

Reviewer comment (Novelty):
The proposed Wavelet activation function seems to be a simplified version of the positional encoding. In addition, IMHO, the sine activation function, followed by MLP, which attaches weights to each neuron, could do the same thing. cos function can also be easily expressed by a bias term from the previous layer e.g., sin(x) = cos(90-x).

Author response:
We respectfully disagree with the reviewer’s statement. 

First, Positional Encoding (PE) is a typical technique used in combination with embedding methods, it serves different purposes as activation functions, including Wavelet. Also, PE and Wavelet have different mathematical expressions. 

Second, Wavelet activation is not the same as Sin activation. We assign two learnable parameters $\omega_1$ and $\omega_2$ to weigh the sin and cos components in Wavelet activation. Indeed, Wavelet can be simplified to $k \sin (x+ \phi)$ for some $k$ and $\phi$ (note $\phi$ does not necessarily need equal to $\pi/2$). However, on the one hand, intuitively, compared to pure Sin activation, Wavelet introduces additional phase and magnitude information. The arbitrary learnable combination of phase and magnitude gives a larger capability than the pure Sin activation function with fixed phase and magnitude information. On the other hand, precisely, let assume $\sin (x_1) = k_1 \sin(x_2)+ k_2 \cos(x_2)$, where $x_1$ and $x_2$ are the arbitrary outputs from two independent linear layers. To reach the equality, we need have $x_2 = \arcsin(\frac{k_1}{\sqrt{k_1^2+k_2^2}}) + \arccos(\frac{\sin x}{\sqrt{k_1^2+k_2^2}})$, the translation from $x_1$ to $x_2$ involves non-linear operations, which cannot be achieved by simply add additional weights/bais to neurons. Our ablation study 1 in the global common response also empirically demonstrates the difference between Sin and Wavelet activation function.

Reviewer comment (Theory):
The system is using a time discretised set of function values to predict the next time step, which reminds the reviewer of the finite difference method. In this case, the stencil is 100 elements long, so the optimum stencil would have a very high order in accuracy.

Author response:
**(1) Similarities and Differences with Finite Difference Method:** 

Despite the sequential input in PINNsFormer, which can be similar to the finite difference method, they are actually different. The key idea of the finite difference method is that it discretizes the whole domain and propagates the solution globally and sequentially. PINNsFormer, instead, focuses on the temporal dependency by looking at additional collocation points near the original collocation point. In PINNsFormer, the predictions within a sequence do not have a specific sequential propagation order, but all predictions utilize temporal dependencies from each other for better predictions in parallel.

Reviewer comment (Experiments):
Very large stencils are generally not used because they bear a computational cost, and the same issue occurs when using attention, although most of the computing is parallel.

Author response:
**(2) Computational Cost:** 

Using attention mechanisms or complicated architectures can indeed raise additional computational costs. However, there are existing solutions, such as Informer [1], which uses prob-sparse attention to reduce the computation cost from $\mathcal{O} (L^2)$ to $\mathcal{O}(L\log L)$. However, the main purpose of this paper is to bring the temporal dependency idea to PINNs with a natural adaptation of transformer architectures. Investigating such mechanisms in PINNsFormer is worth future explorations but yet out of our scope.

Reviewer comment (Experiments):
The manuscript does not provide a baseline using a normal, discrete PDE solver that has a computational load comparable to the attention mechanism.

Author response:
**(3) Adding Normal, Discrete PDE Solver as Baseline:** 

Here, we recognize the normal, discrete PDE solver as the finite element method (FEM) as the reviewer mentioned. However, FEM is typically not a baseline in the context of PINNs. Instead, research works on PINNs [2] use simulation results from FEM as the ground truth to evaluate the neural networks’ performance. Hence as such, we did not include FEM as a baseline when comparing to PINNs and their variations.

[1] Zhou, Haoyi, et al. "Informer: Beyond efficient transformer for long sequence time-series forecasting." Proceedings of the AAAI conference on artificial intelligence. Vol. 35. No. 12. 2021.

[2] Raissi, Maziar, Paris Perdikaris, and George E. Karniadakis. "Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations." Journal of Computational physics 378 (2019): 686-707.

Reviewer comment (Reproducibility):
The lack of published code or detailed hyperparameters makes reproducibility difficult. Providing an implementation would strengthen the paper's contributions.

Author response:
We have included the detailed implementation and hyperparameters in our paper. We include model hyperparameters in Table 4, Appendix A. We include all implementation, including demos for all experiments included in the paper (with fixed random seed for reproducibility purposes), through an anonymous GitHub repo: https://anonymous.4open.science/r/pinnsformer_iclr-4D09, as mentioned on page 7.

Reviewer comment (Experiments):
While the overall approach is promising, some ablation studies would help determine the impact of different components like the pseudo-sequence generation and Wavelet activation.

Author response:
Thanks. We are happy to provide additional ablation studies on (1) the effects of different hyperparameters in pseudo-sequence generators, and (2) the effects of using different activation functions. Please check the common response for details. Overall, our additional ablations show that (i) using Wavelet activation shows constantly better performance than ReLU, Sigmoid, and Sin activations, (ii) $k$ and $\Delta t$ show non-sensitivity to PINNsFormer's performance in general and can be chosen from a wide range.

Reviewer comment (Evaluation):
Since PINNs are notoriously slow in training, the computational overhead of PINNsFormer could be prohibitive for some large-scale applications. Analysis of model complexity and efficiency could help elucidate this issue.

Author response:
We include the computational and memory overheads in Table 5, Appendix A. Empirically, we demonstrate that the computation and memory overheads increase sublinearly as the pseudo-sequence length increases. In particular, using $k=5$ is already good enough for all experiments included, which yields tolerably approximately 3x computational cost and 2x memory cost. Theoretically, existing works such as informer [1] have been shown to reduce the complexity of attention from $\mathcal{O}(L^2)$ to $\mathcal{O}(L\log L)$ by replacing full-attention to prob-sparse attention. Investigating such mechanisms in PINNsFormer is worth future explorations but yet out of our scope.

Reviewer comment (Theory):
The reliance on introducing a discrete timestep risks undermining the automatic differentiation advantage of PINNs. Justification for this design choice could be expanded.

Author response:
We do use automatic differentiation rather than relying on finite difference approximations, so the advantage of PINNs is also inherited in PINNsFormer. In particular, despite a Transformer architecture to consider the temporal dependencies over multiple discrete time steps, the gradient calculation, as specified in Equation 5, is still executed point-wisely through automatic differentiation. We are happy to add more detailed explanations in the final version. The investigation of different $\Delta t$s is included in Response to W2.

Reviewer comment (Experiments):
One major advantage of PINNs is that it leverages automatic differentiation rather than relying on finite difference approximations. On the other hand, the use of discrete pseudo-sequences means temporal dependencies are modeled in the fashion of finite difference approximations rather than pure automatic differentiation. Would it create any difficulties in picking parameters for this differentiation (e.g., the Δt)?

Author response:
We do use automatic differentiation rather than relying on finite difference approximations, so the advantage of PINNs is also inherited in PINNsFormer. In particular, despite a Transformer architecture to consider the temporal dependencies over multiple discrete time steps, the gradient calculation, as specified in Equation 5, is still executed point-wisely through automatic differentiation. We are happy to add more detailed explanations in the final version. The investigation of different $\Delta t$s is included in Response to W2.

Reviewer comment (Experiments):
The impact of the timestep granularity Δt seems worth further analysis. Is there a study on model sensitivity to this parameter? Does the performance degrade if Δt is too small or large? Are there any guidelines for setting Δt?

Author response:
We do use automatic differentiation rather than relying on finite difference approximations, so the advantage of PINNs is also inherited in PINNsFormer. In particular, despite a Transformer architecture to consider the temporal dependencies over multiple discrete time steps, the gradient calculation, as specified in Equation 5, is still executed point-wisely through automatic differentiation. We are happy to add more detailed explanations in the final version. The investigation of different $\Delta t$s is included in Response to W2.

Reviewer comment (Experiments):
Other sampling schemes may also play an important role [1], especially for mitigating the temporal propagation failure in PINNs. Since PINNFormer relies on discretization on temporal dimensions, I am curious how PINNFormer can adapt to non-fixed sampling. I believe the transformer architectures have such flexibility.

[1] Daw, Arka & Bu, Jie & Wang, Sifan & Perdikaris, Paris & Karpatne, Anuj. (2022). Rethinking the Importance of Sampling in Physics-informed Neural Networks. 10.48550/arXiv.2207.02338.

Author response:
The adaptive sampling strategy can be easily incorporated into PINNsFormer, as we have demonstrated the flexibility of incorporating PINNsFormer with existing training schemes in Section 4.3. For instance, the mentioned paper uses adaptive samplings according to the computed risks. One can also incorporate this to PINNsFormer by (1) input collocation point (x,t), (2) output sequential output $[ \hat{u}(x,t+i\Delta t)]_{i=0}^{k-1}$, (3) compute the weight based on either the sum of the risk of all elements of the sequence or the risk of the first element of the sequence.