On the Benefits of Memory for Modeling Time-Dependent PDEs

Authors: Ricardo Buitrago, Tanya Marwah, Albert Gu, Andrej Risteski
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: State Space Models, Partial Differential Equations

Abstract:
Data-driven techniques  have emerged as a promising alternative to traditional numerical methods for solving PDEs. For time-dependent PDEs, many approaches are Markovian---the evolution of the trained system only depends on the current state, and not the past states. In this work, we investigate the benefits of using memory for modeling time-dependent PDEs: that is, when past states are explicitly used to predict the future. Motivated by the Mori-Zwanzig theory of model reduction, we theoretically exhibit examples of simple (even linear) PDEs, in which a solution that uses memory is arbitrarily better than a Markovian solution. Additionally, we introduce Memory Neural Operator (MemNO), a neural operator architecture that combines recent state space models (specifically, S4) and Fourier Neural Operators (FNOs) to effectively model memory.  We empirically demonstrate that when the PDEs are supplied in low resolution or contain observation noise at train and test time, MemNO significantly outperforms the baselines without memory---with up to $6 \times$ reduction in test error. Furthermore, we show that this benefit is particularly pronounced when the PDE solutions have significant high-frequency Fourier modes (e.g., low-viscosity fluid dynamics) and we construct a challenging benchmark dataset consisting of such PDEs.

Decision: accept (oral)
Reviewer scores: 8, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1729111682810

thread_type:
official_review

<thread_object>
creator: reviewer_akxt

<content>
summary:
This work applies a state space model to model the memory term of a PDE that (under assumption) decomposed by the Mori-Zwanzig theory. The method is mostly well-motivated and demonstrates improvement in solving PDEs in low resolution. Overall the paper addresses an important issue in the field of neural operators.

strengths:
1. The paper is theoretically well motivated. 
2. The proposed method for memory operators can be implemented easily.
3. The method demonstrates improvement in performance for solving PDEs.

weaknesses:
1. The motivation for using S4 as the model of choice is not convincing.
2. Lack of numerical comparison against benchmarks in 2D cases.

questions:
1. The main question I have is regarding the flexibility of the proposed method. The author showcased that the memory term modeled by S4 achieves improved performance, but in Appendix G, the memory term modeled by LSTM can be almost just as good. I hope the authors can also highlight this in the main text as I think this shows that adding the memory term is important, no matter the architecture. 
2. Related to the above, can similar memory terms be added to the other benchmarks like U-Net neural operator? If so this can also be highlighted as an ablation study such that the importance of the memory term can be further recognized. 
3. When the high-frequency mode matters, would it simply be more effective to increase the complexity of FFNO rather than applying the memory term? Can the author briefly discuss how this trade-off affects the choice of using a memory term (or not) in this case?
4. How well can the model perform for zero-shot super-resolution once the model is trained?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652844247

thread_type:
official_review

<thread_object>
creator: reviewer_ocpa

<content>
summary:
This paper analyzes the advantages of incorporating memory into AI models for solving time-dependent partial differential equations (PDEs). Traditional AI-driven approaches often rely on a Markovian assumption, where future states are predicted solely from the current state. However, the authors argue that including memory (i.e., information from past states) improves model accuracy, particularly when solutions are low-resolution or contain noise.

The authors further present the Memory Neural Operator (MemNO), a novel architecture that combines state space models (namely S4) and Fourier Neural Operators (FNOs) to capture both spatial and temporal dynamics effectively. Theoretical foundations based on the Mori-Zwanzig theory suggest that memory-inclusive models can outperform Markovian models in specific scenarios, particularly with high-frequency Fourier components. Empirical results show MemNO achieves up to 6x lower test error compared to memoryless models.

strengths:
* This paper seeks to formalize and theoretically analyze the role of memory in learning time-dependent PDEs, distinguishing itself from prior studies that have approached memory in this context primarily from an empirical standpoint.
* The experiments and ablations are thorough and well-executed. Additionally, the empirical results are presented using relative errors, which is exemplary and should be the standard in this field. Unfortunately, this point deserves emphasis, as many other studies overlook this.
* the paper is well written and clearly structured
* the paper introduces recent state-of-the-art state-space models to the field of learning PDE solutions

weaknesses:
* The paper claims to improve learning the solution of time-dependent PDEs theoretically grounded in the 
    the Mori-Zwanzig formalism. However, equation 6 of the paper clearly states that the 'memory term' (i.e., second summand on the right hand side of the equation) involves the differential operator of the PDE, i.e., a local spatial operator. But the proposed method does not take into account spatial interactions and applies the memory cell (S4 layer) to each element of the spatial dimension independently. This is clearly not covered by the formalism and should be stated very clearly. That being said, it would be interesting to incorporate spatial interactions via Graph Neural Networks acting on the nearest neighbors for instance.
* The experiments are interesting and extensive. However, there are some issues that need to be addressed. It's not clear up to which extend the performance increase is due to the memory or simply due to the very expressive S4 model. One good way to analyze this empirically would be to plot the performance of the proposed method on the y-axis, while on the x-axis the size of the look-back window is varied, i.e., how many past instances are fed into the memory cell. Currently, all available past instances are fed to the memory cell. If the performance does not change significantly by varying that, this would indicate that memory is not as important and the performance difference can be explained by the different architectures. Figure 7 in the appendix already suggests that the architecture choice influences the performance.
* A general weakness: Why do we need an advanced AI model with 5M parameters only to solve a simple 1D Burgers equation that can be solved in milliseconds with traditional methods? I acknowledge that the aim of the paper is to improve AI models applied to PDEs. However, it would still be important to draw some connections to traditional methods, which in particular for the experiments considered in this paper are vastly superior. A good starting point to look into this is the referenced paper by McGreivy and Hakim (2024).

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730689551370

thread_type:
official_review

<thread_object>
creator: reviewer_vkvf

<content>
summary:
The authors propose a variant of a factorized Fourier neural operator with memory and show how it can improve over the standard Markovian neural operator for modeling PDEs with small grid sizes or noise added to the observations. Additionally, some theory motivated by the Mori-Zwanzig theory of model reduction is given to motivate and explain the success of the addition of memory (Eq. 6). This success in certain regimes is observed in experiments.

strengths:
1) The problem of whether to use memory to model time-dependent PDEs is well-defined and well-motivated.
2) The theoretical motivation (Eq. 6) and example (Section 4) are clear, concise, and show the possible situations where adding memory in Neural Operators could improve performance.
3) The two experiments varying resolution and additive noise nicely demonstrate regimes where memory can be helpful. Figure 6 also provides good context in this regard.

weaknesses:
1) There should be a brief discussion about the algorithmic cost of adding memory to the FNO. The authors use the diagonal matrix parameterization of S4, so this should be a reasonable increase in cost. Empirical timing would also be helpful to make this point.
2) In the added noise experiments (Section 6.2) it is expected that memory improves performance because the memory unit can approximately filter the input data. The discussion surrounding this experiment is less fleshed out than that of the resolution experiments in Section 6.1 and does not have an obvious connection to the theory and example presented in sections 3 and 4. What practical situation is this intended to model?
3) lines 1119-1121: this is a strange training setup, but maybe necessary due to the GPU memory limitations.

questions:
See Weakness 2.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730694236781

thread_type:
official_review

<thread_object>
creator: reviewer_8n1k

<content>
summary:
The paper proposes a new form of neural operators (NOs) that are inspired by the Mori–Zwanzig (MZ) formalism. This operator is to model spatio-temporal dynamics that arise in time-dependent (nonlinear) partial differential equations (PDEs). The paper points out that the previous roll-out-of-NOs which is caused by memorylessness and, thus, proposes to implement the memory term similarly in the MZ formalism. The specific example architecture has been investigated: Fourier NO + S4, which are empirically shown to be more effective in capturing the dynamics. Moreover, the paper points that the previous benchmark PDEs are easy (including relatively only frequency) and, thus, proposes a new set of benchmark datasets that gives challenge to previous NOs.

strengths:
The paper is well-motivated and proposes solution methods as well as a new set of datasets.

The paper is well-written, providing necessary details without going into too much detail and explaining why including the memory term can be beneficial in modeling complex dynamics. 

The paper introduces a new challenging dataset (although not extensive), the dataset could be used as a next important benchmark for future works.

The empirical results seem to show that the proposed method is more effective in modeling the dynamics. Interestingly, the results are somewhat aligned with the effect of the MZ formalism in heterogeneous systems; that is, on a coarse mesh (on a coarse-grained or homogenized system in MZ), the Markovian surrogate model can be aided by the memory term to capture the influence the state defined on a finer mesh (the eliminated variables in MZ).

weaknesses:
It is not easy to see exactly how the architecture is designed. Including diagrams showing the forward computational path will be very helpful. 

Although taking a cue from the MZ formalism is good, some discussions on the related topic such as closure modeling for surrogate models seems to be missing.

FNOs are known to select low-frequency modes and are expected to struggle with high-frequency solutions. Are transformer-based neural operators expected to behave similarly? Could the authors provide why GKT performs poorly? 

Also, could the authors provide experimental results with some advanced transformers-based neural operators? such as OFormer (Li, et al, TMLR, 2023) and their variants (for example, FactFormer, Li et al, NeurIPS 2023, 2023)

questions:
Please see the questions in the weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732152900408

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response to reviews and updated PDF

comment:
We thank reviewers for their feedback! We are encouraged to see that the reviewers find our study on the benefits of memory on neural operators important and well motivated [8n1K, vkVf, oCpa, akXT], and find our empirical analysis and ablations to be thorough [oCpa,vkVf]. Similarly, we are glad that the reviewers appreciated our theoretical treatment of memory based  on the Mori-Zwanzig equation, which distinguishes us from prior works and is strongly backed up by our experiments [8n1K, vkVf, oCpa]. Lastly, we appreciate that the reviewers found valuable our new challenging benchmark of PDEs in low resolutions [8n1K] and our proposed MemNO as an effective and easy to implement architecture to model memory [akXT].

Based on the suggestions of the reviewers, we have modified the original version of the paper with more baselines, experiments and ablations. We believe these new results reinforce our main claims of the paper: 
 - Our S4FFNO has much superior performance than another newly introduced baseline (Factformer) and a more extensively finetuned GKT baseline, showing that **modeling memory brings significant performance gains compared to an even wider range of Markovian baselines.** *(Table 1 and Figure 1)*
 - The MemNO framework with U-Net as the Markovian model and S4 as the memory model outperforms the purely Markovian U-Net in low resolutions. **This illustrates the flexibility of our framework and shows that the benefits of memory do not apply to only one specific architecture.** *(Table 4)*
 - S4FFNO outperforms the Markovian FFNO by a large margin even when FFNO has 7x more parameters, proving that **the superior performance of S4FFNO comes from the ability to model memory**. *(Table 5)*
 - S4FFNO improves performance as it has access to a larger memory window, again showing that **the performance gap between S4FFNO and FFNO is due to the memory of the full history of previous timesteps.** *(Figure 8)*

The changes in the updated version are in green, we briefly summarize them here:
1. **Factformer as an additional baseline:** We included Factformer [1] as an additional Markovian baseline in our KS and Burgers' experiments (Sections 6.1 and C, respectively). The results are shown in Table 1 and Figure 1. **Our proposed S4FFNO architecture outperforms Factformer in low resolutions**. Details on the architecture are given in Appendix B.
4. **S4U-Net as another effective memory model**: We showed the flexibility of our proposed MemNO framework by combining S4 as the memory model and U-Net as the neural operator. As in the case of FFNO, **S4U-Net outperforms U-Net in low resolutions**, see Appendix G.3. Thus, this experiments shows that the benefits of memory are not unique to a specific archiecture, and the MemNO framework is effective when applied to other neural operators apart from FFNO. We further emphasized this point in our paper by explaining that several memory architectures are able to outperform Markovian ones (lines 431-433), which we believe reinforces our main contributions regarding the effect of memory. 
5. **Comparison between S4FFNO and FFNO with up to 7x more parameters**: We conducted an experiment (Appendix H.2) **where we increased the model size (hidden dimension and number of layers) of FFNO up to 7x times, and still the original S4FFNO outperformed it and achieved 3x less error**. 
6. **Memory window length of S4FFNO:** We show that S4FFNO increases performance as it has access to the memory of more previous timesteps, going from the FFNO performance (no previous timesteps) to 4x lower error. This eliminates confounding factors and indicates that **memory is the reason for the improved performance of S4FFNO**. The ablation is showed in Appendix H.1.
7. **Computational cost and algorithmic complexity of neural operators:**  We provided empirical timings for all architectures, as well as an analysis on the algorithmic complexity of FFNO and S4FFNO. The results are shown in Appendix B.1 and B.2.
8. **Improved GKT baseline**: The Markovian GKT baseline was strenghtened by more extensive hyperparameter tuning (Figure 1). The details on the training setup are included in Appendix B.1. We note that these new hyperaparameters do not improve the performance on the results previously shown in Table 1, and still S4FFNO and other Markovian baselines outperform GKT.
9. **References on closure modeling in Mori-Zwanzig discussion**: We included references related to the topic of closure modeling of surrogate models in Section 3.2 (lines 189-194). 
10. **Ablation on the modeling of local spatial information**: On Appendix I, we have included an ablation on architecture modifications that model spatial information more explicitly to approximate the memory term.
11. **Diagram of MemNO**: We have included a diagram of the MemNO framework in Figure 9.


[1] Li, Z., Shu, D., & Barati Farimani, A. (2023). Scalable Transformer for PDE Surrogate Modeling
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a novel neural operator with memory for modeling time-dependent PDEs. By combining state-space models and Fourier neural operators, the authors establish a sound framework that extends beyond standard Markovian neural operators. The approach is not only novel, but also demonstrates its significance through strong experimental results. In addition, Mori-Zwanzig theory provides a theoretical motivation for the proposed model. Overall, the key strengths of this paper lie in its innovative architecture, theoretical underpinnings, and strong empirical findings. 

All reviewers agree that this is a very strong submission. In particular, I believe that this paper has the potential to inspire interesting future research. Thus, I strongly recommend accepting this paper.