MgNO: Efficient Parameterization of Linear Operators via Multigrid

Response to Questions:

Question 1: Thank you for raising this interesting question. Indeed, Theorem 3.1 pertains solely to shallow neural operators. The theoretical impact of deep neural operator architecture on performance, as well as the role of depth in neural operators, remain open questions.

In fact, the role of depth in traditional deep Neural Networks is not yet fully understood. A vast amount of literature has explored this topic, primarily focusing on the following aspects:

1. Expressivity: 2. Approximation power; 3. Implicit bias and regularization effects; 4. The landscape of loss and its impact on training. I conjecture that depth in neural operators might play similar roles as in traditional neural networks. However, no rigorous theoretical framework has yet been established to confirm this.

In our experiments, we observed a rapid increase in performance when the depth of MgNO was increased from one hidden layer to 4, 5, or 6 layers. Beyond this point, however, the benefits of increased depth diminish, weighed against the costs of computational complexity and training challenges.

From both theoretical and practical standpoints, whether a significantly deeper neural operator architecture, such as one with 100 layers, can substantially enhance neural operator performance remains an open question.

Question 2: Thank you for drawing attention to this aspect. We acknowledge the need for a discussion on limitations or drawbacks in our concluding remarks section. Accordingly, we incorporate a discussion on the limitations or drawbacks there. One limitation of the current implementation of MgNO is its restriction to inputs and outputs on uniform meshes; otherwise, mesh data has to be integrated into the input such that the input is akin to pixel grids in image data. To address this, we plan to integrate ideas from algebraic multigrid methods to enhance our current MgNO framework. This enhancement will enable it to handle inputs and outputs on more general meshes, and even extend to arbitrary sampling data points or point clouds.

Thank you for your feedback on our paper. We appreciate your recognition of the strengths of our proposed Multigrid Neural Operator (MgNO) architecture and your comments on areas needing clarification. Below, we address each of your concerns and questions.

Addressing Weaknesses:

1. Discretization of Input:

• Clarification:

You are correct that the input requires discretization, as shown in Eq.(6). MgNO applies multigrid structures to the spatial domain, which inherently involves a discretization process. For irregular domains, the mesh structure data has to be further integrated into the input data.

• Impact of Mesh Size:

The performance and efficiency of MgNO are influenced by the mesh size. While discretization invariance, as observed in FNO, suggests that a model's performance is not sensitive to input resolutions for certain tasks, it's important to recognize the limitations of this concept. Discretization invariance can result in uniformly good or uniformly bad performance across different input resolutions. The uniformly good scenario implies that the underlying operator is low-dimensional, where coarse information suffices to capture its behavior. However, this is not typically the case in complex problems, such as fluid flow at high Reynolds numbers or wave propagation at high wave numbers. For instance, in turbulent flows, fine resolution is crucial to accurately represent small-scale phenomena, and failure to capture these details can significantly alter long-range predictions. Thus, expecting uniform performance across various resolutions is unrealistic, particularly for complex real-world problems.

To illustrate this, experiments in Table 6 of Appendix D, compare the performance of MgNO and FNO across different mesh sizes. These experiments reveal that for tasks where fine-scale information is crucial, both models exhibit deteriorated performance compared to the results shown in Table 1. This demonstrates the importance of resolution in such complex problems.

Furthermore, in the study "Representation Equivalent Neural Operators: a Framework for Alias-free Operator Learning," it is proven that discretization invariance is achievable only within a band-limited function space (low dimensional space), which is a rather restrictive condition. Our experiments, therefore, provide valuable insights into the practical implications of mesh size on model performance.

2. Parameterization with Multi-Channel Inputs:

We appreciate your query regarding the handling of multi-channel inputs in our model. This aspect, focusing on both expressivity and complexity, is thoroughly discussed in Section 4.1 of our manuscript.

In our experiments, the model indeed utilizes a multi-channel parametrization, and notably, the input for the Navier-Stokes task is multi-channel, demonstrating the model's capability to handle such data. This approach aligns with classical multigrid methods applied to solve elliptic partial differential equations (PDEs), including linear elasticity problems in 2D or 3D scenarios. In such cases, the V-cycle multigrid methods can be naturally represented as linear MgNO with multi-channel inputs.

Addressing Weakness:

Thank you for highlighting this aspect. We acknowledge the need for a discussion on limitations or drawbacks in our concluding remarks section. Accordingly, we incorporate a discussion on the limitations or drawbacks there. One limitation of the current implementation of MgNO is its restriction to inputs and outputs on uniform meshes; otherwise, mesh data has to be integrated into the input such that the input is akin to pixel grids in image data. To address this, we plan to integrate ideas from algebraic multigrid methods to enhance our current MgNO framework. This enhancement will enable it to handle inputs and outputs on more general meshes, and even extend to arbitrary sampling data points or point clouds.

Response to questions

Question 1: We thank the reviewer for raising concerns regarding the scalability of our proposed method. The MgNO architecture, central to our approach, is specifically tailored for efficient parameterization of linear operators within function spaces. This is achieved through a multigrid structure and multi-channel convolutional form, enabling scalability to high-resolution inputs typically encountered in complex problems.

Importantly, the multigrid structure of MgNO is recognized for its effectiveness in handling large-scale problems, particularly in numerical solutions to PDEs. This aspect of our architecture underpins our confidence in its scalability.

Empirically, MgNO has demonstrated superior performance in both accuracy and efficiency, as evidenced in our results. Notably, it features the smallest number of parameters as shown in Table 1 of our paper. To directly address the scalability concern, we conducted additional experiments to assess MgNO's performance at various input resolutions -- $64 \times 64$, $128 \times 128$, $256 \times 256$, and $512 \times 512$ -- while monitoring memory usage during both forward and backward passes. For MgNO, the recorded memory usages were 145.04 MB, 574.16 MB, 2290.66 MB, and 9156.64 MB, respectively. In comparison, for the Fourier Neural Operator (FNO), the memory usage were 184.25 MB, 630.16 MB, 2342.80 MB, and 9051.43 MB at the same resolutions. These results indicate a near-linear relationship between memory usage and the square of the resolution $(n^2)$, where $n$ represents the resolution. This linear trend in resource usage, coupled with the comparative efficiency over FNO, further substantiates our method's scalability.

In conclusion, both the architectural design of MgNO and the empirical evidence from our experiments reinforce our confidence in its scalability to more demanding scenarios, thereby addressing the concerns raised.

Question 2: Thank you for the insightful point. Indeed, we can extend the universal approximation theorem to encompass more general Banach or Hilbert spaces, which may not necessarily be spaces of functions. As evident from the neural operator architecture in Equation (1), the definitions of $\mathcal W_i$, $\mathcal B_i$, and $\mathcal A_i$ are independent of the structures of $\mathcal{X}$ and $\mathcal{Y}$. This allows for a direct extension to abstract Banach spaces.

However, there are two critical aspects to consider. Firstly, we must assume the existence of at least one Schauder basis, enabling the uniform approximation of $\mathcal{X}$ by finite-dimensional spaces. This assumption is crucial for achieving the approximation rate as outlined in the second step of the original proof in Appendix A.

Secondly, the definition of the activation function needs attention. In the initial version of our manuscript, we proposed a property for the nonlinear activation operator ($\widetilde{\sigma}: \mathcal{Y} \rightarrow \mathcal{Y}$): there exist two elements $\xi, \zeta \in \mathcal{Y}$ such that $\widetilde{\sigma}(a \xi) = \sigma(a) \zeta$, where $\sigma: \mathbb{R} \rightarrow \mathbb{R}$ is a real, non-polynomial function. This approach facilitates a proof analogous to the third step in Appendix A.

With these assumptions for $\mathcal{X}$, $\mathcal{Y}$, and the activation operator $\widetilde{\sigma}$, we can replicate the proof in Appendix A to achieve similar universal approximation results for highly general or abstract Banach spaces $\mathcal{X}$ and $\mathcal{Y}$. In the initial version of our manuscript, we included this result. However, we later decided to omit it, concerned that its generality might lead to misunderstandings among readers.

Response to Questions

Question 1: Thank you for your valuable suggestion. In the revised version of our manuscript, we have now included a comparison with FNO in Figure 7 (Helmholtz). Additionally, we have incorporated a comparison with sFNO-v2 as referenced in the paper you mentioned [https://arxiv.org/pdf/2301.11509]. For this comparison, we employed the default hyperparameter settings as provided in the paper, using their official implementation.

However, it's important to note that the results we obtained did not precisely mirror the performance reported in the paper. This discrepancy could potentially be attributed to the hyperparameters not being fine-tuned for the specific task in our study.

Performance on Helmholtz.

Question 2: Thank you for your careful review. We have revised the sentence accordingly to "In our training setup, we implemented a comprehensive approach."

Addressing Weakness

Weakness 1: We appreciate your valuable feedback. For initial remarks on MgNet, please refer to the 'Background and Related Work' section, particularly the subsection titled 'Multigrid'. We include a more comprehensive discussion here and also in the revised manuscript. MgNet is the pioneering work to introduce multigrid to neural network architecture. The authors rigorously demonstrated that the linear V-cycle multigrid structure for the Poisson equation, can be represented as a convolutional neural network, despite their experimental focus on nonlinear multigrid structures for vision-related tasks. Our main idea is that the key of operator learning is how to efficiently parametrize a rich family of linear operators integrated with specific boundary conditions. The impressive performance outcomes we've recorded offer robust evidence in favor of this approach. The insight of MgNet forms a cornerstone in the implementation level of our approach in utilizing multigrid to parameterize such operators. However, as we have discussed in the "Background and Related Work" section under "Multigrid," the original MgNet and its subsequent adaptations in the context of numerical PDEs and forecasting scenarios, have not provided a simple but effective solution in the realm of operator learning.

Weakness 2: Thank you for highlighting this aspect. We should have added a few more words for this in our paper. The constant $c$ in Equation (9) is determined by the dimensions and properties of $a(x)$ but remains independent of the mesh size $h$. Typically, a classical iterative method (denoted as $\mathcal{W}$) for solving discrete elliptic PDEs (such as $A_h u_h = f_h$) will exhibit a convergence rate of $1-\frac{1}{c_h}$, where $c_h = \mathcal{O}(h^{-2})$ with mesh size (resolution) $h$. This indicates that the approximation rate of $\mathcal{W}$ for $A_h^{-1}$ (i.e., the inverse of discrete elliptic operators, which is the discrete Green's function) is not uniform and deteriorates when $h$ is small. However, Equation (9) reveals that if we employ the operator implemented by a multigrid iterative method (denoted as $W_{Mg}$), the convergence rate becomes $1-\frac{1}{c}$. This implies that the approximation rate of $W_{Mg}$ to $A_h^{-1}$ (Green's function) is uniformly bounded by $1-\frac{1}{c}$, independent of the mesh size $h$. Therefore, Equation (9) demonstrates that $W_{Mg}$ has significantly enhanced expressivity for general linear operators between two finite element spaces. Specifically, it achieves a uniform approximation rate for a special type of kernel functions, namely, Green's functions, which are the inverses of elliptic operators. This insight is crucial in motivating the use of multigrid methods to parameterize the linear operator $W_i$.

In Section C of the Appendix, we provide a detailed numerical example to quantify the constant $c$ and approximation rate.

Weakness 3: We agree that other neural operators can always zeros pad the input to a larger enough (for example padding size 7 for the baseline FNO) domain such that the output functions space is redefined on larger domain. Without this strategy, FNO would be constrained to learning mappings to periodic function space on the target domain. However, it's important to note that padding (always zero-padding for any boundary conditions), is not reflective of an operator's intrinsic ability to learn specific boundary conditions and should be considered more as a workaround than a solution.

In contrast, our distinct parametrization method enables the explicit expression of the multigrid solver. Consequently, the inverse of elliptic operators $\mathcal L$ under a linear FEM framework can be approximated by MgNO operators rigorously. This includes, but is not limited to, the inverse Laplacian $\Delta^{-1}$ with any boundary conditions. Furthermore, the constructive approximation of the inverse Laplacian $\Delta^{-1}$ with any boundary conditions using MgNO explicitly requires the padding mode to be zero, reflect, or periodic (with a padding size of 1), corresponding to Dirichlet, Neumann, and periodic boundary conditions, respectively. This is achieved without the requirement to expand the domain via padding. In Section C of the Appendix, we have further provided a detailed numerical example illustrating the construction of a linear MgNO operator $\mathcal{W}_{Mg}$, showcasing the integration of boundary conditions into MgNO and its approximation capabilities. In summary, we assert our approach with MgNO inherently integrates boundary conditions into its structure without resorting to padding to larger domain, offering a more natural and mathematically sound method for handling various boundary conditions.