Sub-Sequential Physics-Informed Learning with State Space Model

We sincerely appreciate the constructive comments from reviewer PWyp and the time spent on reviewing this paper. We address the questions and clarify the issues accordingly as described below.

[C&E:] I have concerns on the novelty .... with a toy example in Figure 1 does not sound a challenging problem.

[Re:] To understand our work's contribution, one must first understand the current status of research on PINN failure modes, which is the core problem we are addressing. PINN failure mode is not as simple as a regular learnability problem. It is a problem unique to PINN and was first identified by Krishnapriyan et al in 2021 NeurIPS. PINN failure modes refer to a phenomenon where the residual loss of a PINN is extremely low, but its average error (even on training collection points) is extremely high. This phenomenon is distinguished from concepts like local minimal (PINN failure modes loss can be as small as ground truth), overfitting, and gradient vanishing, in which the core problem is optimization or generalization. The essential difficulty of this problem is that ground truth is unknowable, and the only known fully optimized PDE-based governing loss is unreliable. In real-world applications, this can result in serious misestimation of the system.

The central contribution of our paper is a unique understanding of this unique phenomenon. The fact that the continuous-discrete mismatch and simplicity bias are well known does not mean that their explanatory role for PINN failure modes is well established. The ideas in the traditional spatiotemporal modeling do not necessarily apply to PINN, since training a PINN is a data-free and ground-truth-unknown problem. A great deal of discussion has arisen in the research community regarding failure modes for PINNs, but to the best of our knowledge, including PINNsFormer, which introduced the Transformer to PINNs, has failed to intrinsically understand and address the PINN failure modes. The distortions in the propagation of the initial conditions we propose are empirically observable, as shown in Fig.1, 2, 6, 7, and 8. We propose that continuous-discrete mismatch and simplicity bias are the core causes of these propagation failures, which is a completely novel understanding of the subject.

We also need to point out that the convection problem shown in Fig.1 is a well-known challenging failure mode case, proposed by [1] and used by several related works.

[M&EC:] technical novelty is limited.

[Re:]: Our contribution lies not only in this simple and effective SSM-based approach but also in the deep understanding of PINN failure modes. Our proposed SSM-based approach is an embodiment of this understanding. Without this understanding, such an approach that can address failure modes in a wide range of equations would never have been produced. In general, we think the machine learning community always appreciates and prefers simple and effective methods.

Sub-sequence contrastive alignment differs from temporal contrastive learning in that, instead of a model learning contrastively for multiple similar data/frame features, it aims to allow different stages of time-varying SSM to form a consensus on the prediction of spatiotemporal collection points. Sub-sequence contrastive alignment emphasizes the ability of the model to inherit information at subsequent times and the ability to form an SMM-step-width agreement, while temporal contrastive learning focuses on representation learning of similar data features, which are not available in PINN.

[TC:] Theorem 3.1 ... does not directly indicate that continuous-discrete mismatch is the key challenge in propagating initial condition in training PINN.

[Re:] What we need to show through Theorem 3.1 is that continuous-discrete mismatch may lead to disconnections in the pattern propagation from the initial condition if only the losses at discrete collection points are optimized. This is because a pattern defined by a point may only act in its small neighborhood, and this neighborhood may not contain any other collection points, which can therefore cause the failure of propagation. Also, a concurrent work, ProPINN [2] observed a lower gradient correlation phenomenon, which is empirical evidence for our proposal that continuous-discrete mismatch is the cause of propagation failure.

[ED & A:] Detailed Ablation and Discussion.

[Re:] We add some ablation studies to discuss the combinatorial effect. See response to reviewer yZBS. Due to the rebuttal limit, we can't include a more detailed analytical discussion of these ablations but will do so in the next release.

Given these explanations, we sincerely invite you to reconsider the rating of our work.

[1] Krishnapriyan, et al. "Characterizing possible failure modes in physics-informed neural networks." NeurIPS 2021

[2] Wu, et al. "ProPINN: Demystifying Propagation Failures in Physics-Informed Neural Networks." arXiv:2502.00803

We sincerely appreciate the constructive comments from reviewer Ce7f and the time spent on reviewing this paper. We address the questions and clarify the issues accordingly as described below.

[W1]: Claims to resolve the continuous-discrete mismatch but still relies on a discretized version of Mamba..

[Response to W1]: We need to clarify that we do not claim to solve the continuous-discrete mismatch. In fact, we think that the computer's finite precision inherently makes it impossible to describe continuous systems (including continuous-time SSM).

Instead, we can only try to mitigate the effect of this mismatch. A discretized SSM is a description of the behavior of a continuous SSM. The point here is that even discrete-time SSM describes the dynamic in a continuous/differential manner since there is a set of rules for a direct mapping between a discretized SSM and a continuous SSM. Although it is actually an approximation of continuous dynamics, this ability to directly approximate/describe continuous dynamics is not available with traditional neural network models such as MLP and Transformer.

Indeed, SSM may not be the ultimate solution to this problem, but we have still taken a huge step towards solving continuous-discrete mismatch. We hope that our work will inspire more researchers to achieve a better solution to this problem.

[W2]: Comparsion with other methods.

[Response to W2]: Since our approach is an on model architecture, we have mainly listed multiple model architectures as baselines. But we do value your opinion and we will include the following results in the next release. Since the experimental settings of some of the papers differ from ours and the source code is not published, we report those works that we successfully reproduced.

[C&S1]: Theorem is not novel. There is also no theory justifying why PINNMamba outperforms PINN.

[Response to C&S1]: Theorem 3.1 is very intuitive, so we also tried to find a direct proof of Theorem 3.1 when preparing the manuscript, but failed. We would welcome some relevant references from the reviewers.

As for a theory justifying why PINNMamba outperforms PINN. We agree that an in-depth theoretical analysis of the proposed method is helpful to better understand our model. But to be honest, currently, we could not find a good theorem to prove rigorously such mechanisms in deep networks, since theoretically analyzing a complex system like PINNMamba is very difficult. However, we still have some intuitive and empirical analysis to show the methodology and philosophy behind our continuous-discrete mismatch and simplicity bias perspective. We hope such analysis can help readers to better understand our motivation and provide some guidance to design better models in future research.

[C&S2]: The proposal cannot fundamentally resolve simplicity bias.

[Response to C&S2]: As previous work[1,2,3] has pointed out, simplicity bias is an intrinsic problem for neural networks. Currently, there is no single way to completely solve this problem, only to evade or mitigate it.

[1] H Shah. The pitfalls.. NeurIPS 2020

[2] D Teney. Evading the simplicity bias.. CVPR 2022

[3]R Tiwari. Overcoming simplicity bias.. ICML 2023

[Q1]: Novelty: How is this work fundamentally different from PINNsFormer, given that it seemly replaces Transformers with Mamba...

[Response to Q1]: Our work is not a simple block-wise replacement. Although both works are based on model architecture, they are quite different in the following points.

1. Macro Architecture: PINNsFormer employs an Encoder-Decoder architecture, which we found unnecessary for PINN. So in PINNMamba, we use an Encoder-Only archi for better performance and efficiency.

2. Sequence Modeling: As shown in Fig.5, PINNsFormer uses a pseudo-sequence, which is not a real collection points sequence, instead, a regional extrapolation, which cannot propagate information. We build collection point sub-sequences and an alignment to realize info propagation across time.

3. Understanding: Our work provides a deeper understanding of PINN failure modes that it is caused by C-D mismatch and simplicity bias. We believe it can inspire more valuable work in the future.

[Q2]: typo

[Response to Q2]: It should be $(x,t+k\Delta t)$. Will fix.

[Q3]: Memory Usage

[Response to Q3]: We greatly optimized memory usage in a follow-up. The OOM is no longer an issue when using a reduced model size. See response to reviewer yZBS.

We sincerely appreciate the constructive comments from reviewer yZBS and the time spent on reviewing this paper. We address the questions and clarify the issues accordingly as described below.

[W1]: The proposed approach is computationally expensive, as shown in Table 5 of the supplementary material. The memory overhead is the second largest, and the training time is seven times slower than that of PINN.

[Q1]: From the sensitivity analysis, the proposed approach is very sensitive to the sub-sequence length, especially comparing the length of 3 and 5 (about 70 times smaller). Are there any explanations for such a big gap?

[Re W&Q1] We answer these two questions together. We start by fixing a bug in Table 6. In our follow-up experiments, we found an error in the experimental data for sequence lengths 3 and 5 in Table 6. The reason for this error is that we accidentally reduced the numerical precision of the calculations when performing these two sets of experiments, which led to an unexpected performance degradation of PINNMamba when using these two sets of parameters. So the conclusion of the original submission about the sensitivity of sequence length needs to be corrected, and we will fix this in the next release. Our latest experimental data are shown in the table below.

Convection:

Reaction:

This group of experiments shows that our model can reduce computational and memory overhead by reducing the sequence length. There is a slight accuracy degradation on the reaction problem, but the performance is even better on the convection problem. Nevertheless, they successfully address the failure modes in all cases.

On the other hand, we found that setting the MLP layer's width to 512 is unnecessary. Reducing the MLP width from 512 to 32 yields a model that still solves the failure modes.

Convection:

Reaction:

If we set sequence length to 3 and MLP width to 32, the memory usage is 1900MB, only slightly larger than 1605MB of vanilla PINN. Given the addressing of failure modes and an rMAE that is only 1/60th of vanilla PINN's, we consider this slight increase in memory and computational overhead to be insignificant. Our model is robust w.r.t sequence length and MLP width. Model with smaller sequence length and MLP width can also eliminate failure modes. This further enhances the generalizability of our approach.

The gap between 3 and 5 sequence length is in fact not large, but when it is 1, it degrades to PINN and failure modes present again. We sincerely apologize for this confusing.

[Q2]: The current ablation study only removes one component at a time, making it difficult to see the effect of other combinations of proposed components.

[Re Q2]: We added the following experiments.

These experimental data show that elimination of continuous-discrete mismatch with SSM is the most important factor, illustrating the significance of the PINNMamba model. Also, Sub Sequence Alignment and Time Varying SSM play an important role in eliminating simplicity bias, and the use of wavelet activation function leads to better numerical precision.

[Q3]: ... it would be interesting to see the performance comparison using a similarly sized model to PINN (i.e., with a smaller embedding size).

[Re Q3]: For the convection problem, we further reduce the width of MLP and Mamba to 8. This reduces the memory overhead of the model to 1300MB, and the average optimization time per iteration to 0.23s (less than PINN). The rMAE is 0.0262, the rRMSE is 0.0346. It addresses the failure modes and outperforms all baselines.

We sincerely appreciate the constructive comments from reviewer PwV2 and the time spent on reviewing this paper. We address the questions and clarify the issues accordingly as described below.

[W1]: The authors placed related works in the appendix, which I find unusual and not ideal. Integrating it into the main text would provide better context for their contributions.

[Response to W1]: We highly agree with your opinion. We will place the related works section in the main text in the next release since another page will be allowed in camera-ready according to ICML author guidance.

[Un RW]: ... Including a discussion of such methods (NN/Graph-based PDE solvers) would provide a more comprehensive view of the broader landscape of neural network-based PDE solvers.

[Response]: We will add the following paragraph in the related works section.

Learning-Based PDE Solvers. In addition to Physics-Informed Neural Networks (PINNs), other neural-network-based approaches have emerged for solving PDEs, each offering unique advantages. Graph Neural Networks (GNNs) like MeshGraphNet [1] excel at handling irregular domains by treating computational meshes as graphs, making them particularly effective for complex geometries. Neural operators [2], including Fourier Neural Operators (FNOs) [3] and Graph Neural Operators (GNOs) [4], learn mappings between function spaces, enabling generalization across different PDE parameters without retraining. Hybrid approaches combine neural networks with traditional numerical methods, such as neural finite element techniques [5], to enhance solver efficiency. While PINNs uniquely enable data-free solutions through direct physics constraint enforcement, GNNs, and neural operators provide complementary capabilities - GNNs for mesh-based problems and neural operators for parametric systems. These diverse approaches collectively demonstrate the expanding toolkit of neural-network-based PDE solvers across scientific computing applications.

[1] Pfaff, Tobias, et al. "Learning mesh-based simulation with graph networks." in ICLR. 2021.

[2] Kovachki, Nikola, et al. "Neural operator: Learning maps between function spaces with applications to pdes." in JMLR. 2023.

[3] Li, Zongyi, et al. "Fourier Neural Operator for Parametric Partial Differential Equations." in ICLR. 2021.

[4] Li, et al. "Multipole graph neural operator for parametric partial differential equations." in NuerIPS. 2020.

[5] Hennigh, Oliver, et al. "NVIDIA SimNet™: An AI-accelerated multi-physics simulation framework." International conference on computational science. 2021.

[W2]: reduce the length of the methods, introduction and remove some results to the appendix.

[Response to W2]: We will streamline these sections in the next release depending on space.

[Q]: About the memory usage.

[Response to Q]: The main reason for the large Memory Usage is that, initially, we set the sequence length to 7 and MLP width to 512 in our main experiment, which was made with the consideration of making a fair comparison with the then SOTA model PINNsFormer. High memory usage is one of the main drawbacks of sequence modeling on PINNs, as the gradient information needs to be preserved for every point in the sequence.

However, we find in a follow-up that our model is robust w.r.t sequence length and MLP width. Our model with a smaller sequence length and MLP width can also eliminate failure modes. This further enhances the generalizability of our approach. To scale to more complex problems, we suggest reducing the sequence length or reducing the MLP width.

Our follow-up experiments show that the model is not very sensitive to sequence length and that a sequence length of 3 makes PINNMamba very effective in combating failure modes. (This is contrary to our results in Table 6 because we found that we had incorrectly set the computational precision in the length 3 and 5 setting in our original sensitivity analysis, resulting in severe performance degradation. We will fix this in the next release.) Adjusting the sequence length to 3 leads to a reduction in the memory overhead of the model by about 57%.

In addition, we found that a large MLP width (512) is not necessary. Setting the width to 32 is sufficient to make PINNMamba effective. This saves up to 53% of memory usage. The combination of the two adjustment can lead to a reduction in memory consumption from 7899MB to 1900MB on the convection problem.

We added the following experiments. It is worth noting that the model successfully eliminates failure modes under all these settings.

When length is set to 1, PINNMamba degrades to PINN and failure modes present again.