Alias-Free Mamba Neural Operator

1. Most of the hyperparameters, including encoder layer, upsampling factor, scanning direction, and integration depth, have been ablated with results and practical settings given in subsection E of supplemental material. Note that for most NO applications, the dimension of state space is simply set as 16 or 32, hence we omitted its ablation in our previous submission. As suggested, the concerned experiments are provided in the uploaded PDF. The practical selections are dependent on the balance of efficacy and efficiency.

2. Yes, handling different scales is a key issue in solving PDEs. MambaNO addresses the problem of different behavioral scales not only through the up- and down- samplings of its U-shaped architecture. Additionally, MambaNO combines global (Mamba integration) and local (convolution integration) treatments to learn holistic features, thereby better handling multi-scale phenomena. Besides, the ablations on U-net layers and the replacement of Mamba or convolution integration with naive modules are also given in the subsection E of supplemental material. However, the texts (deep analysis) on how effectively MambaNO can handle problems with significantly different scales are indeed insufficient, which will be fulfilled in later versions. In addition, we conjecture that introducing adaptive mesh refinement or incorporating physics-guided priors would tap the potential of enhancing the capabilities in multi-scale applications.

3. Boundary conditions play a crucial role in solving PDEs because they define the behavior of the solution at the boundaries of the computational domain. We believe MambaNO's unique scanning mechanism can better handle the initial and boundary conditions of PDEs, allowing the model to utilize this information during the forward inference process. Reducing the scanning directions would also decrease the model's effectiveness, as demonstrated in our ablation study in supplemental material. In fact, techniques such as data augmentation, inputting more boundary information, adaptive boundary conditions, and multi-task learning can significantly enhance the model ability to handle different boundary conditions. We will continue to explore the limits of MambaNO's capability to handle various boundary conditions in the near future.

4. To extend the cross-scan operation to 3D or higher-dimensional PDEs, we need to generalize the scanning process to handle additional dimensions. This involves iterating along multiple spatial axes and potentially managing increased computational complexity. In 3D case, this means scanning along the x, y, and z axes while ensuring the consistency and accuracy of the numerical solution. Anticipated challenges Computational Complexity: The increase in dimensions significantly raises computation and memory demands. Data Handling: Managing large datasets in higher dimensions can be complex. This includes efficient storage, retrieval, and manipulation of data. Accuracy and Stability: Increased dimensions would heighten the potential for numerical errors and instabilities. Unfortunately, we have not yet conducted experiments related to 3D PDEs. Inspired by GINO [1], which combines GNO and FNO for 3D PDEs, we may next consider integrating GNO and MambaNO to address the 3D challenges.

5. The state-space model is a mathematical framework that represents a physical system in terms of input, output, and state variables, which are related through first-order differential equations. PDEs typically need to be solved within specific spatial and temporal domains, requiring precise handling of boundary and initial conditions. The state-space model is particularly well-suited for managing these conditions. Specifically, the state-space model uses a scanning-like mechanism to handle information from initial and boundary conditions, integrating this information into the state equations to determine the system's state at the next time point. This approach is logical and intuitive because the state-space model systematically tracks the time evolution of the system's states, ensuring that initial and boundary conditions are correctly processed at each time step.

6. As suggested, we have conducted experiments on larger dataset such as 2D CFD in PDEBench [2] and compared MambaNO with CNO[3] and other two O(N) competitors, including GNOT[4] and OFormer[5]. The numerical and visual results are given in Table 3 and Fig. 1 of the uploaded PDF, respectively. Evidently, the performance of MambaNO is ahead of other three competitors. In terms of memory usage, stopping epoch, and training time, our MambaNO ranks third, second, second, respectively. Note that all these four are with O(N) complexity.

7. We believe that generalizing to very different PDEs would be a tremendous advance in this field. In the original texts, we provide both in-distribution and out-of-distribution test results for the same PDE, verifying that MambaNO has optimal generalization across different distributions. However, we have not yet conducted experiments on significantly different PDEs. We consider that the generalization capability for irregular or higher-dimensional PDEs still needs improvement. In the future, we will further explore transfer learning to enhance the ability of the algorithm in handling different PDEs.

[1]. Geometry-informed neural operator for large-scale 3d pdes. in NeurIPS 2024.

[2]. Pdebench: An extensive benchmark for scientific machine learning. in NeurIPS 2022.

[3]. Convolutional neural operators for robust and accurate learning of PDEs. in NeurIPS 2024.

[4]. Gnot: A general neural operator transformer for operator learning. in ICML 2023.

[5]. Transformer for Partial Differential Equations’ Operator Learning. Transactions on Machine Learning Research, 2022.

Your additional explanations and experiments have addressed many of the concerns raised. The ablation studies on hyperparameters and the new results on larger datasets like 2D CFD in PDEBench are particularly valuable, demonstrating MambaNO's competitive performance against other O(N) complexity models. Your insights on handling multi-scale phenomena and boundary conditions are appreciated, though further exploration would be beneficial. The explanation of extending the cross-scan operation to higher dimensions is clear, and the potential integration with GNO for 3D PDEs is a natural next step, and I would be excited to see your follow-up work: )! Overall, thanks for the detailed explanation! I will however keep my score: )! Thank you

Thank you very much for your acknowledgment!

In the original texts, we presented the results of eight representative two-dimensional partial differential equations, demonstrating that MambaNO enjoys better accuracy yet with O(N) complexity. Thanks for your reminding of the new datasets and the two outstanding models. As suggested, we have added several experiments on 2DCFD dataset from PDEBench[1], comparing with the mentioned OFormer[2] and GNOT[3], as well as one original competing method CNO[4]. The numerical and visualized results are given respectively in Table 3 and Fig. 1 of the uploaded PDF file. It is shown that MambaNO still leads the performance among the four O(N) competitors, while ranking second in terms of used epochs and training time, as well as third in terms of GPU memory usage (close to CNO).

Note that the main contribution of MambaNO is not simply complexity reduction. In fact, we propose a novel form of kernel integral that conducts holistic feature learning both locally and globally. On that basis, the theoretical analysis on the obeyance of discretization invariance and continuous-discrete equivalence is also provided. Moreover, the performance again is achieved with only O(N) complexity.

Finally, we will cite this mentioned dataset and models in later version. The experimental results on the complete PDEBench dataset will be also added, given chance of revision. Thanks for your valuable suggestions!

[1]. Pdebench: An extensive benchmark for scientific machine learning. in NeurIPS 2022.

[2]. Transformer for Partial Differential Equations’ Operator Learning. Transactions on Machine Learning Research, 2022.

[3]. Gnot: A general neural operator transformer for operator learning. in ICML 2023.

[4]. Convolutional neural operators for robust and accurate learning of PDEs. in NeurIPS 2024.

Due to the imminent closure of the discussion period, we kindly request the reviewer to provide us with their valuable feedback on our rebuttal. We are at their disposal to answer any further questions in this regard.

1. As suggested, we have provided the experimental data and their p-values in the uploaded PDF, showing that the differences between the experimental results are statistically significant. In other words, these differences are not due to any random fluctuations. Note that such tables are too redundant and often contain many repetitions since most p-values are <0.001, so we have omitted the presentation of p-values in the original submission. For all competing methods, we used the best parameters suggested in their original papers. All the selected parameters and the trained models will be released at GitHub for comparisons and evaluations. For the proposed MambaNO, most of the hyperparameters, including encoder depth, upsampling factor, scanning direction, and integration layers, have been ablated with results and practical settings given in subsection E of supplemental material. However, we omitted the training parameters in our previous submission. As suggested, they are now provided in Table 2 of the uploaded PDF. In terms of different PDEs, most of them are fixed in all the experiments, yet with partial of them fine-tuned for better performance.

2. The phenomenon of occasionally performance declines with increasing model capacity can be attributed to underfitting (insufficient data samples). Thanks for your reminding, we have re-conducted comparative experiments on a larger dataset (10,000 samples, with 7,000 used for the training set) among four algorithms of the same O(N) complexity, i.e., Oformer[1], GNOT[2], CNO[3], and our MambaNO, the first two of which are suggested by another reviewer. The experimental results are provided in Table 3 of the uploaded PDF. As seen, the accuracy of our model has been further boosted. Note the efficiency metrics are also provided for a deeper comparison.

3. As suggested, we have provided the training loss curves for FNO[4], CNO, and MambaNO in Fig. 2 of the uploaded PDF. As seen, FNO suffers from poorer performance that results from early convergence. The convergence trends of MambaNO and CNO, both with a time complexity of O(N), are basically consistent. With a deeper inspection, we can see that the performance fluctuation of MambaNO is more stable than CNO, with finally lower errors given same epochs. Therefore, neither more epochs nor more training time is required for MambaNO. As suggested, more training details are given in Table 2 of uploaded PDF file. As for better reproducibility and transparency, we indeed only provided the trained model for evaluation. During current stage, the attachment can not be updated, nor any link to website is allowed to provide. However, we promise to release the complete codes at GitHub in the coming days.

4. For irregular geometries, GNOT [2] encodes shape features using an encoder into K and V, which are then combined with cross-attention. RecFNO [5] uses MLP to learn irregular or sparse observations and reshapes the obtained features into a regular shape. These approaches provide insights into how our model can handle irregular problems. We need to use an encoder to integrate irregular shape features and transform them into a form that can be used for subsequent processing. GINO [6] combines GNO and FNO to solve 3D problems. I believe that our MambaNO can also compensate for shortcomings in irregular shapes and 3D problems in this way. Additionally, for 3D problems, designing an efficient integral scanning scheme is also an issue to be discussed. So far, we have not attempted to pre-train the model on multiple PDEs and then fine-tune it. However, these three fine-tuning related papers have broadened our perspective. Thanks. Perhaps pre-training on different PDEs would be more beneficial for our model's performance on downstream tasks. We believe this is an interesting and worthwhile direction to explore. Thanks again for your valuable insights!

5. Yes, fixed-grid convolutions may not be suitable for irregular discrete grids. However, as mentioned in response to point 4, there are already related works that address irregular shapes [2][5][6]. We believe that MambaNO also owns potential for handling irregular shapes, which will be a focus of our future research.

[1]. Transformer for Partial Differential Equations’ Operator Learning. Transactions on Machine Learning Research, 2022.

[2]. Gnot: A general neural operator transformer for operator learning. In ICML 2023.

[3]. Convolutional neural operators for robust and accurate learning of PDEs. in NeurIPS 2024.

[4]. Fourier Neural Operator for Parametric Partial Differential Equations. in ICLR 2021.

[5]. RecFNO: A resolution-invariant flow and heat field reconstruction method from sparse observations via Fourier neural operator. International Journal of Thermal Sciences, 2024.

[6]. Geometry-informed neural operator for large-scale 3d pdes. in NeurIPS 2024.

Due to the imminent closure of the discussion period, we kindly request the reviewer to provide us with their valuable feedback on our rebuttal. We are at their disposal to answer any further questions in this regard.

1. Alias-free is truly an important property of the framework. However, two reasons lead us to move it into the appendix. First, our innovation lies in proposing MambaNO that follows alias-free property. The alias-free framework is innovated by other work [1], not ours. Therefore, we put more effort in clarifying our own ideas, yet with most of the detailed discussions on “alias-free” left in the appendix. Note that a brief introduction of this framework was already given in the main texts. Second, an alias-free NO means it is a representation-equivalent neural operator (ReNO), enjoying discrete-continuous equivalence. Therefore, we have also provided an explicit proposition claiming that MambaNO is actually a ReNO. This implies that the discrete implementations within the neural network are equivalent to the underlying continuous operations, minimizing aliasing errors as much as possible. As you stated, it does mean zero alias error and can be viewed as the reconstruction error for function approximation.

2. Most previous works [2][3][4] have used the fluid-like patterns to visualize the variables of PDEs. In this work, we simply follow this typical scheme to demonstrate the performance of our methods. As suggested, more explanations will be given in latter version.

3. Note that a new Mamba variant is not our core effort. Instead, we attempt to propose a new neural operator for PDEs. Along this line, we found that the kernel integral in traditional NOs bears certain similarities to state-space models, leading to the proposal of MambaNO. Compared with the original NOs, such as FNO[2] and CNO[3], rather than MAMBA, the proposal enjoys both local and global dependencies. The advantages of memory benefit and linear computation are also inherited.

4. In fact, we have provided ablation experiments on the number of stacking MAMBA layers, as shown in Table 6 of supplementary material. Besides, the ablations of varying inherent parameters by using different scanning mechanisms are also given, as shown in Table 5 of supplementary material. As suggested, we would move these contents into the main texts for clarity.

5. As in most previous works [1][2][3][4], the primary application of NOs is solving PDEs. The constraints that need to be considered during application include, but are not limited to, the initial and boundary conditions of PDEs [5], Continuous-Discrete Equivalence in band-limited function spaces [1], and discretization invariance [2]. We believe these constraints are more theoretically dependent, different from LLMs those are more data driven. To the best of our knowledge, LLM has never been used in solving PDEs. The performance of MambaNO stems from multi-scale information and a scanning mechanism that handles PDE boundary and initial conditions. We believe that causal dependency is essential for the interaction between boundary and initial conditions.

6. The primary application is the generalized solution of PDEs. Considering that 1D PDEs are relatively simple, we provide experimental data and visualization results for eight representative 2D PDEs in the original text to demonstrate MambaNO's potential for solving PDEs. Due to time and computational resource constraints, irregular shapes and higher-dimensional problems remain to be explored.

[1]. Representation equivalent neural operators: a framework for alias-free operator learning. in NeurIPS 2024.

[2]. Fourier Neural Operator for Parametric Partial Differential Equations. in ICLR 2021.

[3]. Convolutional neural operators for robust and accurate learning of PDEs. in NeurIPS 2024.

[4]. Transformer for Partial Differential Equations’ Operator Learning. Transactions on Machine Learning Research, 2022.

[5]. Pdebench: An extensive benchmark for scientific machine learning. in NeurIPS 2022.

We sincerely thank the reviewer for acknowledging our response and for raising our score.