On the Benefits of Memory for Modeling Time-Dependent PDEs

(Weakness) The motivation for using S4 as the model of choice is not convincing.

As observed by the reviewer, there may be other memory architectures that outperform Markovian ones. In this sense, we agree that S4 is not the only model that can incorporate memory, which we have emphasized in the updated version in the paper (lines 431-434). Nevertheless, S4FFNO shows strong performance against Markovian baselines and serves to prove our point that memory helps when there is partial observability. We also note that S4 models are have favourable (linear) scaling with the length of the input (as opposed to quadratic scaling with Transformers), which could be beneficial for long range modeling high dimensional PDEs, for example 3D turbulent flows. Empirically, we also found that LSTMs were unstable to train in our 2D Navier Stokes experiments, while S4 layers were stable to train and are consistently performant.

(Weakness) Lack of numerical comparison against benchmarks in 2D cases.

We demonstrate the effectiveness of our approach on both 1D and 2D partial differential equations (PDEs), including the 2D Navier-Stokes equations. We perform experiments in the 2D case where we show the effects of adding observation noise to the input PDEs and how S4FFNO consistently outperforms FFNO. For furhter ablation studies we focus on the 1D PDE cases due to the significantly longer training times for 2D problems (approximately 5 times slower). We note that our ablations suggest that even in the 1D cases there are clear benefits of explicitly modeling the memory!

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?

We do not think that increasing the complexity of FFNO can bridge the gap in performance with S4FFNO, because fundamentally FFNO is Markovian and does not use memory to compensate for the partial observation of high frequency modes. To illustrate this point, we increased the complexity of FFNO (hidden dimension and number of layers) in the KS experiment of Section 6.1 with viscosity $\nu=0.1$. The results are the following:

It can be seen that S4FFNO achieves a much better performance than FFNO (more than 3x smaller nRMSE), despite having up to 7x less parameters. Thus, the memory term is essential to achieve good performance, and the accuracy-efficiency of S4FFNO is much more favorable.

We also found this suggested experiment very valuable and have included in Appendix H.2 of the paper.

How well can the model perform for zero-shot super-resolution once the model is trained?

When the training resolution is low, as in our experiments, it is very hard to perform zero-shot super resolution. We have tried zero-shot super resolution by keeping only the Fourier modes seen during training (as in the original FNO paper [1]), however both FFNO and S4FFNO have very poor performance in this case. For example, when the training resolution is 64, S4FFNO and FFNO obtain a nRMSE of 0.011 and 0.033, respectively. However, when evaluated at a slightly higher resolution of 80, their nRMSE go to over 0.3, which is an order of magnitude higher, which indicates that zero-shot resolution is very difficult when there are high frequency components, regardless of memory. We believe that the achieved super-resoltuion in the original FNO paper is possible in the regime where the solutions do not contain high frequency components. For example, they perform super resolution in the Burgers' experiment with viscosity $\nu=0.1$ by training at resolution 256 and evaluating at higher resolutions. We note that most of the frequency spectrum of this Burgers' PDE is contained Fourier modes well below 128, as we show in our Figure 6. Thus, the first modes of the discrete fourier transform do not vary much when the observation resolution is changed from 256 to higher ones (intuitively, the fourier spectrum can be estimated accurately with resolution 256). This allows zero-shot super resolution, since FNO works with only the top K modes. Nevertheless, in our experiments the solutions of the PDEs contain a large fraction of high frequency components. Therefore, the discrete fourier transform changes a lot when the observation resolution is changed.

[1] Li, Zongyi, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2020. “Fourier Neural Operator for Parametric Partial Differential Equations.”

We find it encouraging that the reviewer found our work well motivated and the method to include memory easy to implement. Furthemore, we thank the reivewer for their suggestion on additional experiments and ablations on adding different Markovian operators and has led to two additional results: 1) we show that MemNO framework also improves the performance with the U-Net model as the Markovian model (Appendix G.3), and 2) S4FFNO still outperforms an FFNO model that is 7x bigger (Appendix H.2)! Please find our responses to your questions below:

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.

Yes! The key motivation of our work is to show when memory is helpful for modeling PDEs, and not a comment on which architecture should be used to model memory. We believe that having many different architectures with memory outperforming Markovian ones reinforces our claim. We explicitly included this point in Section 6.1 of the updated PDF (lines 431-434) and the Conclusion (lines 536-538).

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.

Following the suggestion of the reviewer, we have applied the MemNO framework to the U-Net architecture, and have found similar findings when perfoming the KS and Burgers experiments of Section 6.1 and Appendix C. We have included this experiment in Appendix G.3. Here we also present and analyze the results. The following table contains the nRMSE values in the same setup as Table 1:

It can be seen that enhacing the U-Net operator with memory brings significant benefits when the resolution is low, with S4U-Net outperforming U-Net in Burgers' and KS with resolution 32, achieving sometimes more than 2x lower error. However, in higher resolutions (128 in KS), the memory is not needed and S4U-Net suffers a degradation in performance. We hypothesize that with a more extensive tuning of the training hyperparamters this degradation can be minimized.

This experiment reinforces our main message, which is that including memory can be very helpful, regardless of the architecture. Moreover, it shows the flexibility of our proposed MemNO framework, as it can be applied to U-Net without modifications. We thank the reviewer for pointing out this important ablation and have added the experiments to our paper (Section G.3 of the Appendix).

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.

Regarding the concern about whether the superior performance of S4FFNO comes from modeling memory or simply from the expressivity of the S4 layer, we performed the suggested experiment. In the KS experiment with viscosity $\nu=0.1$ and resolution 32 (same setting as in Section 6.1), we used different temporal window lengths for the memory model S4: T=0, T=5, T=10, T=15, T=20 and T=25 (the model always predicts 25 timesteps in total, but uses the memory of the past T previous timesteps to predict the next one). It can be seen that the performance improves from having no memory (similar to FFNO) to having 25 timesteps of memory (4x lower error). Thus, the improved performance is related to the amount of past memory information provided to the model.

Additionally, we would like to point out to another experiment we included in Appendix H.2. In that experiment, instead of decreasing the window length of S4FFNO, we increased the number of parameters of FFNO (up to 8 layers and 256 hidden dimension). That makes FFNO make ~7x more parameters than S4FFNO. Nevertheless, FFNO has 3x larger nRMSE (Table 4). This indicates that the differences in performance cannot be explained from model expressivity alone, but rather about the inputs that are available to the model.

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).

While our work does not directly offer a solution to the issues found by McGreivy and Hakim (2024), we believe it is a valuable contribution towards desigining neural operators that improve their accuracy-efficiency tradeoff. S4FFNO was specifically designed to have an efficiency close to FFNO (see Table 2), yet we found that it can achieve up to 6x lower test error than FFNO with only an ~14% increase in computation time. Moreover, as we mentioned in another experiment we found that S4FFNO has more than 3x lower test error compared to an FFNO model that is 7x bigger (Table 4). Thus, in this sense we have provided a model with a better accuracy-efficiency tradeoff. More importantly, we provide an explanation of when using memory is necessary. This explanation can help the design of more cost-efficient architectures based on whether memory is required or not, especially given the community's push towards training large scale foundation models for complex high dimensional PDEs. We definitely agree with the sentiment that we need more challenging and meaningful benchmarks for machine-learned PDE solvers more broadly.

We thank the reviewer for their detailed feedback, and are glad that the reviewer finds our theoretical results valuable and that our experiments are well executed. Furthermore, we thank the reivewer for their suggestion on conducting ablation experiments that compare the outputs of varying window lenghts to understand the effects of memory, which we have added in Section H.1 of the Appendix, as well as for pointing out to an important observation regarding the modeling of spatial information, which has led to Appendix I.

Please find our response to the questions below:

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.

We agree that the memory term of the Mori-Zwanzig formalism includes spatial derivatives. We do not explicitly model such spatial information in the memory layer of the MemNO framework, but the hidden dimension can in principle encode this local information implicitly. We agree with the reviewer that this is an important observation which was not present in our original paper, so we have included it in the updated version (line 377 and Appendix I).

Moreover, following the suggestion of the reviewer, we have tried several modifications to the S4FFNO architecture to provide an inductive bias for modeling spatial information more explicitly. Concretely, we tested the following architecture modifications:

• S4FFNO + Input Gradients: The numerical gradients of the input $u_t$ are fed to the model encoder. Specifically, $u_t$ is represented as a vector of spatial shape $S$, and the spatial gradients are approximated using second-order accurate central differences method (using the method torch.gradient of the torch library). This would directly facilitate storing information about the first order spatial gradients in the hidden dimension.
• S4FFNO + Convolutional Encoder: Instead of having a linear layer as encoder, we substitute it with a 1D Convolution of kernel size 3 across the spatial dimension.
• S4FFNO + Convolution before memory layer: Before feeding the sequence of hidden presentations $[v_0, v_1, ..., v_t]$ to the memory layer S4, we apply a 1D Convolution across the spatial dimension with kernel size of 3. This is similar to modeling local interactions through Graph Neural Network, as the reviewer suggested, yet in a simple manner where each hidden dimension node only has two neighbours (its adjacent spatial points).

We use the setup of the KS experiment with viscosity $\nu=0.1$ and resolutions $f=16, 32, 48, 64$ and $80$, and present the nRMSE values in the following table:

It can be seen that none of the methods provides better performance than S4FFNO, which we believe is because training becomes more difficult with these architectures. It would a great direction for future work to consider other ways to explicitly model spatial information and further improve the performance of memory architectures. We have included this experiment in Appendix I.

We are encouraged to see that the reviewer finds that our empirical and theoretical results are concise and clearly show the regimes where using memory is helpful for neural operators. Please find our responses to the questions below:

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.

Thank you for your suggestion, we agree that a discussion on training cost will help better understand the performance of different models! We provide a table with the empirical training times (forward + backward) of the different architectures for a batch size of 32, resolution 64 and 25 timesteps (the rest of the architecture details are the same as in Section 6.1), as well as the parameter counts. We note that S4FFNO is only ~14% slower than FFNO, yet it achieves a loss that is 3x lower in the KS experiment with viscosity $\nu=0.075$ (see Table 1). Moreover, we have included another experiment in Appendix H.2, where we show that S4FFNO achieves 3x lower error than an FFNO model that has 7x more parameters (see Table 4).

We have included these values in Table 2 of appendix B.1. We used a NVIDIA RTX L40S for these measurements.

As for the theoretical algorithmic complexity, S4FFNO has a smaller cost than FFNO. Let $S$ be the spatial resolution, $T$ the number of timesteps, $H$ the hidden dimension and $N$ the state dimension of S4. The core spectral convolution of FFNO (Eq. 16) has a Discrete Fourier Transform across the space dimension and a matrix multiplication in the frequency space, which have complexities $O(T H S \text{log}S)$ and $O(T S H^2)$ respectively. In contrast, the S4 layer has a Discrete Fourier Transform across the time dimension and it requires building the convolution kernel, which have complexities $O(S H T \text{log} T)$ and $O(S H (N+T) \text{log}(N+T))$ respectively (see [1] for details). In our cases, $S$ ranges from 32 to 128, and $T$ is either 20, 25 or 32, so $O(T H T \text{log}T) < O(T H S \text{log}S)$. As for the other term, we use $H=128$ and $N=64$, thus we also have $O(S H (N+T) \text{log}(N+T)) < O(T S H^2)$. Thus, in theory the added S4 model to S4FFNO contributes to a factor smaller than one layer of FFNO to the total architecture. We have included this analysis in Appendix B.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?

This experiment is concerned with the practical situation where a physical measurement device has observation noise. We believe this is a situation of practical importance since the measurement devices generally come with an observation error, especially when the observation grid is very large (as in the case of 2D systems), where obtaining accurate measurements is very costly. In this case, we showed that using memory helps the neural operator. We checked that the observation resolution (64x64) is high enough to observe all relevant Fourier modes, and thus it is a different setting than Section 6.1. But still, memory helps the neural operator. We believe this reinforces one of our main claims, which is that memory helps when the solution of the PDE is partially observed.

lines 1119-1121: this is a strange training setup, but maybe necessary due to the GPU memory limitations.

Indeed this setup was due to GPU memory limitations. When the training length was split into two chunks of length 16, the GPU memory requirement was 34GB for this experiment. Using the full length for training would require a memory beyond 48GB, which was beyond our GPU capacity (we used NVIDIA A6000/A6000-Ada GPUs). We have made this fact more explicit in our paper.

[1] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022

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)

We haved provided results for Factformer and included it in the paper as an additional baseline for the experiments in Section 6.1 (KS) and Appendix B (Burgers'). The results are shown in Table 1 and Figure 1. Here we also present the results for Factformer, GKT and S4FFNO to facilitate the comparison:

Factformer outperforms GKT in low resolutions of KS and in the Burgers' equation, yet has worse performance in high resolutions. As seen in Figure 1, Factformer is a strong Markovian baseline in low resolutions, yet it is still outperformed by our memory model S4FFNO. We believe that showing the superiority of S4FFNO against an even stronger Markovian baseline reinforces our claim that memory helps in low resultion.

As for the difference between GKT and Factformer, a possible explanation for the superior performance of Factformer in low resolutions might be related to how GKT and Factformer introduce positional information for the transformer layer. GKT encodes position by concatenating the grid into the queries, keys and values; while Factformer introduces position through cross attention and additionally has an additive positional encoding. This latter method dedicates more parameters to the positional information, which might be helpful in low resolution to give a different treatment to each spatial position and compute a better approximation of the Fourier spectrum. In contrast, in higher resolutions this might be too fine-grained information and Factformer fails to generalize.

Nevertheless, there are several other confounders in the architecture. For example, GKT uses a FNO-based regressor after the transformer encoding layers, while Factformer uses an MLP regressor. Additionally, Factformer uses an MLP to build the queries and keys, while GKT only uses a linear projection. We hope that our proposed benchmark can illustrate new tradeoffs between architectures and can help design better ones that achieve high performance in both low and high resolutions.

Lastly, we mention that we also implemented Oformer (taking it from the repository https://github.com/BaratiLab/OFormer). However, in our KS experiments with low viscosities and low resolutions, the test loss diverged because rolling out Oformer for 25 timesteps lead to very large values. Just like GKT and Factformer, Oformer was tested in a different experimental setup and has several architectural differences (for example, in [3] Oformer is not rolled out for several timesteps from a single initial condition in 1D experiments, while in [2] Factformer was only applied to 2D problems). We hypothesize that making architectural changes would solve the training instabilities. But given that the architecture of Factformer is similar to Oformer, we believe that having Factformer as a baseline is already representative of the performance of more advanced transformer architectures.

[2] Li, Z., Shu, D., & Barati Farimani, A. (2023). Scalable Transformer for PDE Surrogate Modeling

[3] Li, Zijie, Meidani, Kazem, and Barati Farimani, Amir. "Transformer for Partial Differential Equations’ Operator Learning." Transactions on Machine Learning Research, 2023.

Could the authors provide why GKT performs poorly?

In our experiments, it is the case that GKT has worse performance than FFNO, especially at resolutions 36-64 in Figure 1. In general, we found GKT harder to train, but after extensive finetuning we were able to improve its performance in this range of resolutions by making two modifications: 1) adding 0.05 dropout in the linear attention layer and 0.025 dropout in the FFN layer; and 2) reducing the training epochs from 200 to 50. However, we note that this change made the performance drop considerably in resolutions higher than 64. We present the results in the following table:

In order to strenghten the baseline, we have included this hyperparameter sweep over the dropout in our experiment (Section 6.1), and have chosen the performance of the best model (with or without dropout). Although the GKT curves in Figure 1 are now smoother, the performance of FFNO is still better.

However, we do not believe this is necessarily a limitation of Transformer-based neural operators. As the reviewer also suggested, we provide results for Factformer in the answer to the next question. We show that Factformer has superior performance than FFNO in lower resolutions, although worse performance in higher resolutions.

We are pleased that the reviewer finds our work well-motivated and well-written, and that they find it interesting that the empirical results are aligned with the Mori Zwanzig theory, which we believe is a key insight of our work. We provide responses to their questions:

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

We appreciate the suggestion of adding a diagram. We have included it in Figure 9, along with the pseudocode for MemNO in Figure 10, which we hope can help facilitate the understanding of the architecture.

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.

We agree with the reviewer that there are other ways to model the unresolved dynamics beyond the Mozi-Zwanzig formalism. In our work, we explain when modeling the memory term of the Mori-Zwanzig equation is expected to bring benefits, we propose a neural architecture to build surrogate models and we back up our study with empirical and theoretical claims. We hope that our methodology stimulates more interest in the topic of closure modeling more broadly, the design of neural surrogates and the study of when the closures are expected to outperform Markovian approximations.

Additionally, based on this suggestion we have included more citations of works that discuss the closure of surrogate models in the last paragraph of Section 3.2 (see lines 189-195).

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?

Fundamentally, we believe that all Markovian operators will struggle when the solution of the PDE has high frequency components beyond the observation resolution. In our experiments, Markovian operators take inputs on a low resolution grid, which cannot correctly capture high frequency modes. Thus, learning the solution operators exactly is not possible because the operator does not have information on the full spectrum. In the original GKT paper, it is also stated that GKT is not expected to perform well in the presence of high frequency components (see the Conclusion in [1]). In contrast, the Mori-Zwanzig formalism states that we can compensate for this lack of information by leveraging the memory of past states (Equation 6). Thus, we generally expect Markovian operators to underperform their memory-augmented counterparts, as we observe in Table 1.

[5] Shuhao Cao. Choose a transformer: Fourier or Galerkin. Advances in Neural Information Processing Systems (NeurIPS 2021), 34, 2021.