M$^2$M: LEARNING CONTROLLABLE MULTI OF EXPERTS AND MULTI-SCALE OPERATORS ARE THE PARTIAL DIFFERENTIAL EQUATIONS NEED

## Abstract

Abstract Learning the evolutionary dynamics of Partial Differential Equations (PDEs) is critical in understanding dynamic systems, yet current methods insufficiently learn their representations. This is largely due to the multi-scale nature of the solution, where certain regions exhibit rapid oscillations while others evolve more slowly. This paper introduces a framework of multi-scale and multi-expert (M 2 M) neural operators designed to simulate and learn PDEs efficiently. We employ a divide-and-conquer strategy to train a multi-expert gated network for the dynamic router policy. Our method incorporates a controllable prior gating mechanism that determines the selection rights of experts, enhancing the model’s efficiency. To optimize the learning process, we have implemented a PI (Proportional, Integral) control strategy to adjust the allocation rules precisely. This universal controllable approach allows the model to achieve greater accuracy. We test our approach on benchmark 2D Navier-Stokes equations and provide a custom multi-scale dataset. M 2 M can achieve higher simulation accuracy and offer improved interpretability compared to baseline methods.

## 1 Introduction

Many challenges require modeling the physical world, which operates under established physical laws (Karniadakis et al., ; Brunton and Kutz, ). For example, the Navier-Stokes equations form the theoretical foundation of fluid mechanics and have widespread applications in aviation, shipbuilding, and oceanography (Vinuesa and Brunton, ). Various numerical approaches exist to tackle these equations. These include discretization methods such as finite difference (Godunov and Bohachevsky, ), finite volume (Eymard et al., ), finite element (Rao, ), and spectral methods (Shen et al., ). Although classical physical solvers based on first principles can achieve high accuracy, they must recalculate when faced with new problems, failing to generalize and resulting in inefficient solutions. Artificial intelligence-based surrogate models effectively address these issues by providing more adaptable and efficient solutions. Understanding and learning the data that embodies these physical laws is crucial for controlling and optimizing real-world applications (Lv et al., ; Kim and Boukouvala, ). Mastery of such data-driven insights enables more precise predictions, enhanced system performance, and significant advancements in how we interact with and manipulate the application in the fields of engineering and science (Noé et al., ).
The growing interest in efficient PDE solvers and the success of deep learning models in various fields has sparked significant attention, such as neural operator methods (Li et al., ; Kovachki et al., ; Bonev et al., ; Liu et al., ). Um et al. () proposed a spatial resolution solver to reduce the computational and accelerate physical simulations.
Wu et al. (); Sanchez-Gonzalez et al. () through reducing the dimensions of latent space to map the solution in the surrogate models.

Exploring how to integrate and fully leverage performance across different scales while controlling complex learning dynamics is a promising area of research.

Figure: Figure 1: Framework of the proposed Multi-scale and Multi-experts (M^2M). The Experts net has different models, $f_{gate}^{policy}$ decides which spatial domain is needed to allocate the different models in the roll-out predications. For more details, please refer to sec. 3.
Refer to caption: x1.png

Based on the frequency principle (Xu et al., ), our primary motivation is to enable smaller or more general models to learn low-frequency dynamic data, while delegating high-frequency data to more capable models. The router (distribution policy) regulates this allocation, which sets our approach apart from other methods. In this work, we introduce the multi-scale and multi-expert (M^2M) neural operators as the effective surrogate model to learn the dynamics of PDEs and optimize the appropriate allocation law for the different scales with different expert models. Our critical insight lies in leveraging the divide-and-conquer approach among models to learn the capabilities across different scales quickly. Divide and conquer is a fundamental algorithmic technique for solving complex problems by breaking them down into simpler, more manageable sub-problems (Smith, ; Huang et al., ; Ganaie et al., ; Emirov et al., ). This approach works on the principle that a large problem can often be divided into two or more smaller problems of the same or similar type. Each of these smaller problems is then solved independently. Once solutions are obtained for all the sub-problems, they are combined to form a solution to the original, more extensive problem. In addition, the model is designed to master the distribution of the most effective data while minimizing computational resources. To fairly evaluate the effectiveness of our framework, we standardized the internal models to Fourier Neural Operator (FNO) models with varying numbers of modalities.
This strategy enables the model to adaptively determine the best local spatial resolution to evolve the system. The M^2M is trained alternatingly, iterating between training the evolution model with supervised loss and allocation policy net. Together, the controllable routing mechanism effectively integrates prior knowledge with model capabilities. The implementation of PID control significantly aids in optimizing the training of Multi-of-Experts (MoE).

Our main contributions are as follows:

1. We propose a controllable multi-expert and multi-scale operator model to embed multiple models based on specific priors. The multi-expert system embodies the divide-and-conquer philosophy, while the multi-scale approach enables efficient learning.

2. By bridging the control theory-PID, this unified theory demonstrates its strong generalizability. It is a versatile and scalable method for the machine learning and science simulation community.

3. We validate the aforementioned method using the standard Navier-Stokes equations and a custom multi-scale dataset, ensuring a balance between speed and accuracy.

## 2 Problem Setting and Related Work

We consider temporal Partial Differential Equations (PDEs) w.r.t. time $t\in[0,T]$ and multiple spatial dimensions $\mathbf{x}=\left[x_{1},x_{2},\ldots x_{D}\right]\in\mathbb{X}\subseteq\mathbb{
R}^{D}$. We follow a similar notation as in (Brandstetter et al., ).

$$ $\displaystyle\partial_{t}\mathbf{u}=F\left(a(t),\mathbf{x},\mathbf{u},\mathbf{ u}_{x},\mathbf{u}_{xx},\ldots\right),$ $\displaystyle(t,\mathbf{x})\in[0,T]\times\mathbb{X}$ (1) $\displaystyle\mathbf{u}(0,\mathbf{x})=\mathbf{u}^{0}(\mathbf{x}),$ $\displaystyle\mathbf{x}\in\mathbb{X}$ $\displaystyle B[\mathbf{u}](t,\mathbf{x})=0,$ $\displaystyle(t,\mathbf{x})\in[0,T]\times\partial\mathbb{X}$ $$

where $\mathbf{u}:[0,T]\times\mathbb{X}\rightarrow\mathbb{R}^{n}$ is the solution, which is an infinite-dimensional function. $a(t)$ is a time-independent parameter of the system, which can be defined on each location $\mathbf{x}$, e.g. diffusion coefficient that varies in space but is static in time, or a global parameter. $F$ is a linear or nonlinear function. Note that in this work we consider time-independent PDEs where $F$ does not explicitly depend on $t$. $\mathbf{u}^{0}(\mathbf{x})$ is the initial condition, and $B[\mathbf{u}](t,\mathbf{x})=0$ is the boundary condition when $\mathbf{x}$ is on the boundary of the domain $\partial\mathbb{X}$ across all time $t\in[0,T]$. Here $\mathbf{u}_{x},\mathbf{u}_{xx}$ are first- and second-order partial derivatives, which are a matrix and a 3-order tensor, respectively (since $x$ is a vector). Solving such temporal PDEs means computing the state $\mathbf{u}(t,\mathbf{x})$ for any time $t\in[0,T]$ and location $\mathbf{x}\in\mathbb{X}$ given the above initial and boundary conditions.

The fundamental problem can be succinctly represented for tasks involving partial differential equations by the following formula.

$$ $(\partial\Omega,\mathbf{u}_{0})\xmapsto{f_{\theta}}(\partial\Omega,\mathbf{u}_ {1})\xmapsto{f_{\theta}}\cdots\xmapsto{f_{\theta}}(\partial\Omega,\mathbf{u}_{ T}),$ (2) $$

where $f_{\theta}$ represents the model and $\partial\Omega$ denotes the boundary conditions.

Deep Learning-based Surrogate Methods.
There are two fundamental approaches:

- •
Autoregressive Model Approach: The model learns the mapping function $f_{\theta}$ from a given $\mathbf{u}_{t}$ to the next $\mathbf{u}_{t+1}$, acquiring discrete representations. This method involves learning the model to predict subsequent time steps based on previous inputs. Such frameworks include CNN-based models (Wang et al., ; Kemeth et al., ), GNN-based models (Pfaff et al., ; Li et al., ), and transformer-based models (Cao, ; Geneva and Zabaras, ; Takamoto et al., ).
- •
Neural Operator Approach: Unlike autoregressive models, the neural operator method (Lu et al., ) allows the model to map through multiple time steps, learning infinite-dimensional representations. This approach enables the model to handle more complex temporal dynamics by learning continuous representations. Apart from vanilla FNO, there are other operator learning methods such as U-FNO (U-Net Fourier Neural Operator, (Wen et al., )), UNO (U-shaped neural operators, (Azizzadenesheli et al., )), WNO (Wavelet Neural Operator, (Navaneeth et al., )), and KNO (Koopman Neural Operator, (Xiong et al., )).

In addition to these two conventional methods, researchers have developed several hybrid approaches that combine elements of both (Watters et al., ; Zhou et al., ; Keith et al., ; Hao et al., ; Kovachki et al., ; Wang et al., ). For multi-scale PDEs problems, Liu et al. () developed multi-scale deep neural networks, using the idea of radial scaling in the frequency domain and activation functions with compact support. Hu et al. () propose the augmented physics-informed neural network (APINN), which adopts soft and trainable domain decomposition and flexible parameter sharing to further improve the extended PINN further. Xu et al. () firstly find the deep neural network that fits the target functions from low to high frequencies. Liu et al. () demonstrate that for multi-scale PDEs form, the spectral bias towards low-frequency components presents a significant challenge for existing neural operators. Rahman et al. () study the cross-domain attention learning method for multi-physic PDEs by the attention mechanism. However, the aforementioned methods do not efficiently leverage frequency characteristics, and they lack a controllable mechanism for adjusting the learning process of partial differential equations. Compared with (Du et al., ; Chalapathi et al., ), M^2M directly optimizes for the PDEs objective, first using a universal controlled method and multi-scale to learn the policy of allocating experts to achieve a better accuracy vs. computation trade-off. In addition to simplifying the computation of attention, the MoE mechanism (Jacobs et al., ) has been incorporated into transformer architectures (Fedus et al., ; Chowdhery et al., ) to lower computational expenses while maintaining a large model capacity. The key distinction of our objective lies in its emphasis on controllability and multi-scale considerations, both of which are crucial factors for all fundamental partial differential equation data.

## 3 The Proposed Method

In this section, we detail our M^2M method. We first introduce its architecture in sec. 3.1. Then we introduce its learning method (sec. 3.2), including learning objective training, and a technique to let it learn to adapt to the varying importance of error and computation. The high-level schematic is shown in figure 1.

### 3.1 Model architecture

The model architecture of M^2M consists of three components: multi-scale segmentation and interpolation, Experts Net, and Gate router. We will detail them one by one.

Multi-scale Segmentation and Interpolation.
Multi-scale segmentation involves strategically decomposing the input into multiple scales or resolutions to facilitate detailed analysis and processing. This technique benefits applications that require fine-grained analysis on various scales, such as traditional image processing (Emerson, ; Sunkavalli et al., ) and deep learning methods (Zhong et al., ; Yuvaraj et al., ). Consider a discrete form input represented as $\mathbf{u}_{t}^{h\times w}$, where $h$ and $w$ denote the spatial domain resolution at time step $t$. In multi-scale segmentation, the $\mathbf{u}_{t}^{h\times w}$ first needs to be segmented into smaller, non-overlapping scale patches. For example, segmenting a tensor $\mathbf{u}_{t}^{h\times w}$ into $2\times 2$ patches results in four distinct segments. Each segment corresponds to a quarter of the original tensor, assuming that $h$ and $w$ are evenly divisible by 2. Secondly, suppose that we wish to perform an interpolation on these segmented patches to restore them to the original $h\times w$ dimensions.
Mathematically, this operation can be expressed as:

$$ $\mathbf{u}_{t}^{h\times w}\xrightarrow{\text{Segmentation}}\left\{\mathbf{u}_{ i,j}^{\frac{h}{2}\times\frac{w}{2}}\Big{|}i,j\in\{1,2\}\right\}\xrightarrow{ \text{Interpolation}}\left\{\mathbf{P}_{i,j}^{h\times w}\Big{|}i,j\in\{1,2\} \right\},$ (3) $$

where $\mathbf{P}^{h\times w}$ represents the tensor after interpolation, which combines the four patches back into the original size of $h\times w$. This segmentation approach effectively reduces the dimensionality of each patch and allows for localized processing, which is essential for tasks involving hierarchical feature extraction.

Experts Net.
In theory, an expert net is composed of multiple distinct models. However, our sub-expert networks are structured in a parallel configuration for rigorous comparison in this study. Importantly, we have opted for a non-hierarchical architecture. All constituent models are based on the Fourier Neural Operator (FNO), with potential variations in the number of modalities. Formally, let $E=\{E_{1},E_{2},\ldots,E_{n}\}$ represent the set of expert models, where each $E_{i}$ is an FNO. The input to each expert is a different patch $\mathbf{P}_{i}\in\mathbb{R}^{h\times w}$. The output of each expert maintains the same dimensionality as the input. The primary function of the expert system is to model the temporal evolution of the system state as shown in Eq. 2. We employ a divide-and-conquer strategy, where each expert $E_{i}$ operates on a subset of the input space:

$$ $E_{i}:\mathbf{P}_{i}^{h\times w}\rightarrow\mathbf{P^{\prime}}_{i}^{h\times w},$ (4) $$

where $\mathbf{P}_{i}^{h\times w}$ is a patch of the input and $\mathbf{P^{\prime}}_{i}^{h\times w}$ is the corresponding output patch. The predication solution $\mathbf{\hat{u}}_{t+1}^{h\times w}$ involves the aggregation of these individual patch predictions to reconstruct the full system state:

$$ $\mathbf{\hat{u}}_{t+1}^{h\times w}=A(\mathbf{P^{\prime}_{1}},\mathbf{P^{\prime }_{2}},\ldots,\mathbf{P^{\prime}_{n}}),$ (5) $$

where $A$ is an aggregation function that combines the individual patch predictions into a coherent global state, this approach allows for parallelization and potentially more efficient processing of complex spatio-temporal dynamics while maintaining consistency across all or sparse expert models.

Gate Router Mechanism in MoE.
The Gate Router Mechanism is a crucial component in the MoE architecture and is responsible for distributing input patches across expert models. The primary objectives of this mechanism are:

- 1.
To efficiently allocate different patches to different models and to route complex problems to more sophisticated networks. (Divide and Conquer)
- 2.
To avoid overloading a single model, which could lead to high computational complexity. (Simplicity is the ultimate sophistication)
- 3.
Optionally, the router could be set as the top-k and strong prior, which we encourage the sparse experts to apply the different regions.

Let $\mathcal{X}=\left\{\mathbf{u}_{1,1}^{h\times w},\mathbf{u}_{1,2}^{h\times w},
\mathbf{u}_{2,1}^{h\times w},\mathbf{u}_{2,2}^{h\times w}\right\}$ be a set of $4$ input scale domain. The router function $R$ is defined as:

$$ $R:\mathcal{X}\rightarrow[0,1]^{N\times M},$ (6) $$

where $R(x_{i})_{j}$ represents the probability of routing input $x_{i}$ to expert $E_{j}$. The ideal routing strategy aims to optimize the following objectives:

$$ $\min_{R}\mathbb{E}_{x\sim\mathcal{D}}\left[\sum_{j=1}^{N}R(x)_{j}\cdot\text{ Error}(E_{j},x)\right],$ (7) $$

where $\mathcal{D}$ is the data distribution and $\text{Error}(E_{j},x)$ is the error measure of each expert $E_{j}$ performs on input patch compared with the ground truth. This paper introduces an optional prior distribution $P(E_{j})$ over the experts to initialize the routing mechanism and gradually train it. This prior can be incorporated into the routing decision:

$$ $R(x)_{j}=\frac{\exp(r_{j}(x)+\log P(E_{j}))}{\sum_{k=1}^{N}\exp(r_{k}(x)+\log P (E_{k}))},$ (8) $$

where $r_{j}(x)$ is a learned function that scores the suitability of expert $E_{j}$ for input $x$. By combining these components, the router mechanism can efficiently distribute inputs across experts, adapt to the complexity of different inputs, and maintain a balanced computational load across the system.

### 3.2 Learning objective and control strategy

The training objective is defined as follows:

$$ $\mathcal{L}(t)=\lambda(t)\mathcal{L}_{\text{router}}+\mathcal{L}_{\text{ experts}},$ (9) $$

where the $\mathcal{L}_{\text{router}}$ and $\mathcal{L}_{\text{experts}}$ represent the training loss for the router and experts net, respectively. The $\lambda(t)$ is a hyperparameter related to the training epoch $t$. Our assumption is as follows: in the initial stage of the model, the router should allocate data evenly to the experts, allowing each expert to receive sufficient training. Once the expert networks have been adequately trained, if the router is not performing well, feedback should be used to train the router. This will enable the router to select the well-performing experts for further training, thereby fully leveraging the potential of the experts.

Router Loss.
The training objective for the router can be formulated as:

$$ $\mathcal{L}_{\text{router}}=\text{KL}(R(x)||P(E))+\mathcal{L}_{\text{load}},$ (10) $$

where KL is the Kullback-Leibler divergence, and this formulation allows the router to start from the prior distribution and gradually adapt to the optimal routing strategy as training progresses. The KL divergence term encourages the router to maintain some similarity to the prior, which can help prevent all inputs from being routed to a single expert. To promote the sparsity of the router and the computational tradeoff, we introduce a load-balancing entropy loss as $\mathcal{L}_{\text{load}}$:

$$ $\mathcal{L}_{\text{load}}=-\sum_{i=1}^{M}p_{ij}\log p_{ij},$ (11) $$

where $p_{ij}$ represents the probability $R(x_{i})_{j}$ of assigning the $i$-th data point to the $j$-th expert.

Expert Learning Loss.
Each expert model should be trained using supervised learning to approximate the solution of the PDE at a given time step. To achieve this, we define the loss function for each expert model using the Mean Squared Error (MSE) between the predicted solution and the true solution of the PDE at each time step.

For a given expert model $E_{j}$, the goal is to minimize the MSE between its prediction $\hat{u}_{j}(x,t)$ and the true solution $u(x,t)$ of the PDE over a set of input patches. The MSE loss for the $j$-th expert can be defined as:

$$ $\text{MSE}_{j}=\frac{1}{N}\sum_{i=1}^{N}\left(u(x_{i},t_{i})-\hat{u}_{j}(x_{i} ,t_{i})\right)^{2},$ (12) $$

where $N$ is the number of patches, $u(x_{i},t_{i})$ is the true solution at the $i$-th patch, and $\hat{u}_{j}(x_{i},t_{i})$ is the predicted solution by the expert $E_{j}$.

The total loss for all experts can be written as the sum of the MSE losses for each expert:

$$ $\mathcal{L}_{\text{experts}}=\sum_{j=1}^{M}\text{MSE}_{j}=\sum_{j=1}^{M}\frac{ 1}{N}\sum_{i=1}^{N}\left(u(x_{i},t_{i})-\hat{u}_{j}(x_{i},t_{i})\right)^{2}.$ (13) $$

By minimizing this total loss, each expert model learns to approximate the solution of the PDE over time, ensuring that their predictions become more accurate as the training stage.

Figure: Figure 2: Figure (a) shows the router policy in the training. Figure (b) shows the framework of the PI controller in the M^2M. By designing the target and feedback in the loop, $\lambda$ can be adjusted.
Refer to caption: x2.png

PID-Gate Control Connects the Expert and Router. The dispatch mechanism of the router presents a challenging issue. On the one hand, if the router initially has a strong prior, different experts may not receive sufficient training, and their specialized capabilities cannot be fully leveraged. On the other hand, if the router erroneously assigns tasks to less capable experts, the overall loss of the model may not decrease as expected. Inspired by automatic control theory (Åström et al., ; Wang et al., ), we design a non-linear PI controller in the loop as shown in figure 2, a variant of the PID control, to automatically tune the hyperparameter $\lambda(t)$ and use the desired loss or desired prior KL distribution as feedback during model training. we also demonstrate that PID-gate control improves the performance in the ablation study. To address this challenge, we propose two control strategies in Algorithm 1, and the proof is in the Appendix 6.5.

Figure: Algorithm 1 Two Dispatch Strategies with Desired Loss $\hat{L}$

## 4 Experiments

In the following experiments, we set out to answer the following questions on our proposed M^2M:

- •
Multi-scale effect and allocate mechanism: Can the M^2M model dynamically allocate the spatial domain to concentrate computational resources on regions with higher dynamics, thus enhancing prediction accuracy?
- •
Pareto frontier improvement: Does M^2M enhance the Pareto frontier of Error versus Computation compared to deep learning surrogate models (SOTA)?
- •
Controllable training: Is M^2M capable of adapting its learning results based on the dynamics of the problem, as indicated by the parameter $\lambda$?

We evaluate our M^2M on two challenging datasets: (1) a custom 2D benchmark nonlinear PDEs, which tests the generalization of PDEs with the different spatial frequencies; (2) a benchmark-based Naiver-Stokes simulation generated in (Li et al., ). Both datasets possess multi-scale characteristics where some domains of the system are highly dynamic while others are changing more slowly. We use the relative L2 norm (normalized by ground-truth’s L2 norm) as a metric, the same as in (Li et al., ). Since our research primarily focuses on control methods combined with multi-expert models, we aim to utilize the foundation modes of Fourier operators. In the following sections, we will consistently employ FNO32, FNO128, FNO32, and FNO16, with the goal of achieving a higher-order operator FNO256.

### 4.1 The comparison of custom multi-scale dataset v.s. PID-control effect

Data and Experiments. In this section, we test M^2M’s ability to balance error vs. computation tested on unseen equations with different parameters in a given family. For a fair comparison, we made the model size of the different methods as similar as possible. We use the custom dataset for testing the Multi-scale effect and Controllable training. The multi-scale dataset is given by

$$ $\nabla^{2}u(x,y)=f(x,y),$ (14) $$

where $u(x,y)$ is the unknown solution to be solved, and $f(x,y)$ represents the source term, which varies for different regions. More details about the multi-scale dataset are available in the Appendix 6.1.

##### Main Results.

The compared baseline methods are FNO (Li et al., ), UNO (Azizzadenesheli et al., ), CNO (Raonic et al., ), and KNO (Xiong et al., ). Please refer to the appendix for baseline visualization results in the appendix 6.4.3. The M^2M approach achieves Pareto optimality, as demonstrated in the Pareto frontier detailed in figure 10. As a heuristic choice, we set the target to 0 and defined the loss $L(t)$ as the RMSE in the training stage.

Figure: Figure 3: Results of one-step prediction on the multi-scale custom dataset at different epochs: 1, 10, and 100. Our scale is set to $4$, and the ablation on the multi-scale study is shown in appendix 6.1.
Refer to caption: x3.png

**Table 1: The comparison study in the custom dataset. This shows that M^2M can improve prediction error by selecting where to focus computation trade-offs, especially with more stringent computational constraints. All tests are conducted on an NVIDIA A800 GPU. The optimal scale of M^2M is set to $4$.**
| Models | Parameters (M) | Computation (ms) | L2 error |  |
| --- | --- | --- | --- | --- |
| FNO<br>(Li et al., ) | FNO16 | 0.047 | 1.9 | 0.54 |
| FNO32 | 0.16 | 2.2 | 0.13 |  |
| FNO64 | 0.61 | 2.3 | 0.050 |  |
| FNO128 | 2.4 | 2.5 | 0.038 |  |
| FNO256 | 9.5 | 2.6 | 0.036 |  |
| UNO<br>(Azizzadenesheli et al., ) | UNO16 | 5.2 | 6.2 | 0.080 |
| UNO32 | 19.2 | 8.9 | 0.075 |  |
| UNO64 | 74.2 | 18.2 | 0.042 |  |
| UNO128 | 292.6 | 33.7 | 0.026 |  |
| KNO<br>(Xiong et al., ) | KNO16 | 4.2 | 10.6 | 0.99 |
| KNO64 | 67.1 | 130.5 | 0.92 |  |
| CNO<br>(Raonic et al., ) | CNO4 | 2.0 | 18.3 | 0.12 |
| CNO64 | 14.2 | 148.6 | 0.010 |  |
| PID-M^2M (Ours) | Strategy 1,Topk=1 | 4.8 | 4.5 | 0.024 |
| Strategy 1,Topk=2 | 4.8 | 8.0 | 0.008^∗ |  |
| Strategy 1,Topk=3 | 4.8 | 11.2 | 0.012 |  |
| Strategy 2 | 4.8 | 14 | 0.008^∗ |  |

Figure: Figure 4: Dynamic weight distribution of router, the figure (a) and (b) are the distribution of the output on the 1th and 100th epoch. The prior is set to [0000] means none of any prior on the router. The TOPk is set 2.
Refer to caption: x4.png

From figure 4, M^2M can allocate models sparsely and only sends the region with a slower change (Patch number is 1-4) to lowest mode FNO16 for the computation efficiency and highest modes of FNO128 for accuracy both from zero-prior. The above results show that M^2M can focus computation on dynamic learning in the table 1. M^2M achieved notable improvements compared to baseline models, with Strategy 1 requiring less computation time than Strategy 2 while delivering the input to all experts. For PID-M^2M, we test our model on different initial $\lambda$ values and different TOPk by empirical selection of $K_{p}=0.001$, $\lambda_{min}=0$, $\lambda_{max}=1$ and $K_{i}=0.001$, where $\lambda$ focuses on router training in the initial training stage. The comparison study of the PI effect is shown in figure 3. The PI controller can speed up the error convergence in training and the value of $\lambda$ has been controlled. We see that with $\lambda$ (e.g., $\lambda(0)=0$ in the initial value) dynamic change, proving that the PI controller can control the router in the allocating process with different experts. To investigate whether M^2M can allocate the best expert on the most dynamic region according to different priors, we visualize which allocating experts on outputs of the router as shown in appendix 6.4.1.

### 4.2 The Naiver-Stokes (NS) dataset and comparison of SOTA

Here we evaluate our M^2M performance in a more challenging dataset in the Naiver-Stokes dataset, the description is shown in Appendix 6.2. In this experiment, to ensure a fair comparison and leverage the M^2M method’s ability to enhance FNO’s inherent capabilities, we selected the baseline model of FNO-3D instead of the auto-regressive style in FNO-2D. Since FNO-3D operates at least three times faster than FNO-2D, this choice significantly shortened the experimental cycle. As shown in table 2, our M^2M can allocate high FNO modes to the high-frequency region and achieve better accuracy than the baselines. Specifically, M^2M outperforms the strong baseline of FNO128 in the performance a little. This shows that the router could learn a proper allocation policy, allowing the evolution model to evolve the system more faithfully. However, the multi-scale effect did not perform well on this complex dataset, especially in the boundary. The reason is that there are strong temporal scale dependencies between patches, and as the spatial partitioning increases, the divide-and-conquer approach becomes less effective in the appendix 6.2. We applied our method to a cylinder wake flow in the appendix 6.4.4, close to real-world data which has the prior on the fluid mechanic. By incorporating prior distributions in regions where vortex shedding forms around the cylinder, the prediction is quite accurate.

Figure: Figure 5: Results of NS datasets in PID-M^2M. The number of scale is 1.
Refer to caption: x5.png

**Table 2: The comparison study in the NS dataset. This shows that M^2M can improve prediction error by selecting where to focus computation trade-offs, especially with more stringent computational constraints. All tests are conducted on an NVIDIA A800 GPU. The scale of M^2M is set to $1$.**
| Models | Parameters (M) | Computation (ms) | L2 error |  |
| --- | --- | --- | --- | --- |
| FNO | FNO16 | 0.050 | 2.1 | 0.29 |
| FNO32 | 0.16 | 4.7 | 0.26 |  |
| FNO64 | 0.61 | 4.9 | 0.25 |  |
| FNO128 | 2.4 | 6.2 | 0.24 |  |
| FNO256 | 9.5 | 6.9 | 0.22 |  |
| PID-M^2M (Ours) | Strategy 1, Topk=1 | 4.0 | 5.0 | 0.26 |
| Strategy 1, Topk=2 | 4.0 | 7.8 | 0.23^∗ |  |
| Strategy 1, Topk=3 | 4.0 | 13.0 | 0.25 |  |
| Strategy 2 | 4.0 | 14.9 | 0.23^∗ |  |

## 5 Conclusion and Limitation

In this work, we have introduced M^2M. This first intense learning-based surrogate model jointly learns the evolution of the physical system and optimizes assigning the computation to the most dynamic regions. In multi-scale and Naiver-Stokes datasets, we show that our PID method can controllably train the expert’s net and router, which improves long-term prediction error than strong baselines of deep learning-based surrogate models. The fitting error has been demonstrated to converge based on control theory in the Appendix 6.5. Furthermore, the PID-based M^2M could improve the convergence speed, showing that this intuitive baseline is suboptimal. Finally, M^2M outperforms its ablation without multi-scale segmentation, showing the divide-and-conquer strategy which can significantly reduce the prediction error. We hope M^2M can provide valuable insight and methods for machine learning and physical simulation fields, especially for applications requiring scalability and multi-physics models. The limitation of M^2M is shown in appendix 6.6.

## 6 Appendix

### 6.1 Custom Multi-scale Poisson equation dataset

The custom multi-scale dataset is designed to simulate a complex scenario, where data in some regions change slowly while changing more rapidly in others. We define the task as follows: the input is the solution $\mathbf{u}^{low}$ to a relatively low-frequency equation, while the output is the solution $\mathbf{u}^{high}$ to a corresponding high-frequency equation. To idealize this dataset, we adopted the form of the classical Poisson equation and used the finite difference method to solve the problem. The concise discrete form is $[1,128,128]\mapsto[1,128,128]$. The time step is set to $1$ and the spatial domain is set to $[128,128]$.

#### 6.1.1 Frequency Distribution on different regions

The source term for each region is a sinusoidal function with a systematically varying frequency. Specifically, the source term $f_{ij}(x,y)$ is defined as follows for $i,j=1,2,3,4$:

$$ $\displaystyle f_{11}(x,y)$ $\displaystyle=\sin(\pi\cdot(1\cdot\mu x))\sin(\pi\cdot(1\cdot\mu y)),$ (15) $\displaystyle f_{12}(x,y)$ $\displaystyle=\sin(\pi\cdot(2\cdot\mu x))\sin(\pi\cdot(2\cdot\mu y)),$ $\displaystyle f_{21}(x,y)$ $\displaystyle=\sin(\pi\cdot(3\cdot\mu x))\sin(\pi\cdot(3\cdot\mu y)),$ $\displaystyle f_{22}(x,y)$ $\displaystyle=\sin(\pi\cdot(4\cdot\mu x))\sin(\pi\cdot(4\cdot\mu y)).$ $$

The initial solution of PDEs will be decided by the dimensionless frequency $\mu$ and the other solution for the high frequency is $7\cdot\mu$. In this dataset, we sampled 1000 cases with different values of $\mu$, which were drawn from a normal distribution $\mathcal{N}(1,0.1)$ using Monte Carlo sampling. Out of the 1000 samples, 700 are allocated for the training dataset, while the remaining 300 are reserved for the test dataset.
To increase the complexity in the varying time PDEs, we assume that the solutions include a two-step solution and that the ground truth (the second time-solution) is the high spatial frequency to be predicted, $7\cdot\mu$ of each low-frequency domain corresponding to the input domain.

#### 6.1.2 Solver Implementation and setting of grids

The Poisson equation is solved numerically using a finite-difference method on each block. The boundary conditions and the source term $f(x,y)$ determine the solution $u(x,y)$ within each block. The computational grid is set into a $128\times 128$ grid and divided into $2\times 2$ blocks, each of size $64\times 64$. After calculation, the boundary condition $g(x,y)=0$ is assigned in each block boundary.

### 6.2 2D Naiver Stokes

#### 6.2.1 2D Naiver Stokes Datasets

The Navier-Stokes equation has broad applications in science and engineering, such as weather forecasting and jet engine design. However, simulating it becomes increasingly challenging in the turbulent phase, where multiscale dynamics and chaotic behavior emerge. In our work, we specifically test the model on a viscous, incompressible fluid in vorticity form within a unit torus. The concise discrete form is $[10,64,64]\mapsto[10,64,64]$. The input time step is set to $10$ and the spatial domain is set to $[64,64]$.

$$ $\displaystyle\partial_{t}w(t,x)+u(t,x)\cdot\nabla w(t,x)$ $\displaystyle=\nu\Delta w(t,x)+f(x),$ $\displaystyle x\in(0,1)^{2},\,t\in(0,T]$ (16) $\displaystyle\nabla\cdot u(t,x)$ $\displaystyle=0,$ $\displaystyle x\in(0,1)^{2},\,t\in[0,T]$ $\displaystyle w(0,x)$ $\displaystyle=w_{0}(x),$ $\displaystyle x\in(0,1)^{2}$ $$

where $w(t,x)=\nabla\times u(t,x)$ is the vorticity, $\nu\in\mathbb{R}_{+}$is the viscosity coefficient. The spatial domain is discretized into $64\times 64$ grid. The fluid is turbulent for $\nu=10^{-5}\left(Re=10^{5}\right)$.

#### 6.2.2 Results of M 2 M at different scales

Figure 6 and 7 below show the prediction results at two different scales. As can be seen, there are some sharp edges at the boundaries.

Figure: Figure 6: Model Performance at a scale factor of 2
Refer to caption: x6.png

Figure: Figure 7: Model Performance at a scale factor of 4
Refer to caption: x7.png

### 6.3 Detailed configuration of the M 2 M and baseline models

This section provides a detailed configuration of M^2M, baseline methods, and the hyperparameters used for training in Table 3 and Table 4.

#### 6.3.1 Hyper-parameters for training of M 2 M

The policy network is implemented as a classifier, with the output corresponding to the weights distribution of the expert network.

**Table 3: Hyperparameters for M^2M architecture and training.**
| Hyperparameter name | Custom dataset | 2D NS dataset |
| --- | --- | --- |
| Model architecture: Experts and Router |  |  |
| Experts architecture: FNO32, FNO128, FNO64, FNO16 |  |  |
| Router architecture: Transformer based classifier |  |  |
| Autoregressive roll-out steps | 1 | 1 |
| Hidden channels of FNO | 6 | 6 |
| $f_{gate}^{\text{policy}}$: Transformer embedding dim | 128 | 64 |
| $f_{gate}^{\text{policy}}$: Numbers of head | 4 | 4 |
| $f_{gate}^{\text{policy}}$: Numbers of layers | 2 | 2 |
| $f_{gate}^{\text{policy}}$: Encoder layers | 2 | 2 |
| $k_{p}$ of PID | 0.001 | 0.001 |
| $k_{i}$ of PID | 0.02 | 0.02 |
| Hyperparameters for training: |  |  |
| Learning rate | $1e^{-3}$ | $1e^{-3}$ |
| Optimizer | Adam | Adam |
| Batch size | 8 | 4 |
| Number of Epochs | 100 | 200 |

#### 6.3.2 Hyper-parameters for the training of Baseline models

The setting of our baseline methods is shown as follows. The hyperparameters used for training are the same as those used in the M^2M model above.

**Table 4: Setting for baseline models**
| Hyperparameter name | Custom dataset | 2D NS dataset |
| --- | --- | --- |
| Baseline Operators: FNO, UNO, KNO, and CNO |  |  |
| In channels | 1 | 10 |
| Modes of FNO: ${16,32,64,128}$ |  |  |
| Hidden channels of FNO | 6 | 6 |
| Modes of UNO: ${16,32,64,128}$ |  |  |
| Hidden channels of UNO | 6 | 6 |
| Scaling of UNO | [1,0.5,1,2,1] | [1,0.5,1,2,1] |
| Layers of UNO | 5 | 5 |
| Modes of KNO: ${16,128}$ |  |  |
| Operator size | 6 | 6 |
| Decompose Number | 15 | 15 |
| Modes of CNO: ${4,8}$ |  |  |
| Number of block | 4 | 4 |
| Channels | 16 | 16 |

### 6.4 Extra Visualization

#### 6.4.1 MoE and PID trajectory details

In this section, we present results concerning the two types of priors in the router during the initial phase, along with different PID parameters and scaling factors. One type of strong prior, such as $[0100]$ to add the output of the router, indicates that the router assigns each patch to four experts by incorporating the prior directly into the router’s output through hard constraints, followed by a softmax function. The other type of weak prior represented as $[0000]$, relies entirely on the router’s output without any prior constraints. As for the second prior, the results has shown in the figure 3.

Figure: Figure 8: Variation results of input, ground truth, prediction, and absolute difference along the training epochs. The columns from left to right represent training times of 1, 50, and 100th epoch units respectively. The prior distribution for FNO is set to $[0100]$
Refer to caption: x8.png

Figure: Figure 9: Variation of MoE Expert Weights with Prior and PID Parameters with scaling factor $4$. Figure (a) Output of the router during the first epoch of the training stage. Figure (b) Output of the router at the 50th epoch of training. Figure (c) Output of the router at the 100th epoch. Figure (d) Adjust PID model parameters with the target set to 0. The error is defined as the difference between the model’s loss function and the target. The green line represents the controllable value for $\lambda(t)$. The prior distribution for FNO is set to $[0100]$
Refer to caption: x9.png

#### 6.4.2 Ablation study on the multi-scale effect

We compared the performance of multi-scale models on the custom dataset, where the model is directly routed to different experts by a router, with the prior set to $[0000]$. It is worth noting that these comparisons were made without the inclusion of the PID algorithm, to ensure fairness in the table 5. Both interpolation and extrapolation methods in the multi-scale stage were chosen to be linear for the sake of computational efficiency.

**Table 5: Ablation study on the multi-scale effect in M^2M. The prefix number represents the scale factor $S$, and Topk is set to 4. All tests were conducted without a controller.**
| Models | RMSE | MAE |
| --- | --- | --- |
| 1-Scale M^2M | 0.015 | 0.004 |
| 2-Scale M^2M | 0.010 | 0.004 |
| 4-Scale M^2M | 0.008 | 0.005 |
| 8-Scale M^2M | 0.008 | 0.007 |

#### 6.4.3 Baseline results

Here, we show the Pareto Frontier with different models in the figure 10. It can be observed that our M^2M model, represented by the blue stars, lies on the Pareto frontier, demonstrating that our computational speed and accuracy are quite competitive. Performances of baseline models CNO, FNO, UNO, and KNO on the custom dataset are presented in figure 11, figure 12, figure 13, and figure 14.

Figure: Figure 10: Pareto Frontier in the multi-scale dataset. The larger shape in the legend means larger mode numbers and larger parameters.
Refer to caption: extracted/5893354/Fig/Pareto.png

Figure: Figure 11: CNO performance on multi-scale datasets. left: CNO4, right:CNO16
Refer to caption: x10.png

Figure: Figure 12: FNO performance on multi-scale datasets. Five Columns from left to right: FNO16, FNO32, FNO64, FNO128, and FNO256
Refer to caption: x11.png

Figure: Figure 13: UNO performance on multi-scale datasets. Four Columns from left to right: UNO16, UNO32, UNO64, UNO128
Refer to caption: x12.png

Figure: Figure 14: KNO performance on multi-scale datasets. From left to right: KNO16, KNO128
Refer to caption: x13.png

#### 6.4.4 Performance on the Dataset of flowing around a cylinder

We also consider the general application of our method on natural datasets, the dataset is shown in the following figure 15, such as the flowing around a cylinder [Tencer and Potter, ].
The input shape $[1,192,112]$ needs to be mapped to the next time step with the shape $[1,192,112]$. The model simulation is conducted at a Reynolds number of 160, which leads to the formation of a vortex.

The allocation strong prior between different experts’ models is relatively reasonable, as the vortex is generated behind the cylinder in the patch $5-12$. During training, it is necessary to assign the high-frequency components to the higher modes of the FNO. It is shown that M^2M can learn based on the strong prior like the figure 16, which will promote artificial intelligence in scientific discoveries and simulations.

Figure: Figure 15: Dataset of flow past a cylinder, left tensor shape:$[1,192,112]$, right tensor shape:$[1,192,112]$
Refer to caption: x14.png

Figure: Figure 16: MOE weight on the strong prior distribution and predication result on the flow past a cylinder
Refer to caption: x15.png

### 6.5 Control theory on PID

A PID controller regulates the control input $u(t)$ by combining three terms: proportional ($P$), integral ($I$), and derivative ($D$). The controller aims to minimize the tracking error $e(t)$, defined as the difference between the desired reference signal $r(t)$ and the system output $y(t)$:

$$ $e(t)=r(t)-y(t)$ $$

The PID control law is given by:

$$ $u(t)=K_{p}e(t)+K_{i}\int_{0}^{t}e(\tau)d\tau+K_{d}\frac{de(t)}{dt}$ $$

where $K_{p}$, $K_{i}$, and $K_{d}$ are the proportional, integral, and derivative gains, respectively.

#### Proof of Convergence

The behavior of the error $e(t)$ as $t\to\infty$. It will diminish over time due to the combined effect of the PID control effect.

#### 1. Proportional Term

The proportional term $K_{p}e(t)$ provides an immediate response to the current error. The larger the error, the stronger the control input. This term ensures that the error decreases proportionally to its magnitude, reducing the error in time:

$$ $\frac{de(t)}{dt}=-K_{p}e(t)$ $$

This equation indicates that the error decreases exponentially for values of $K_{p}$.

#### 2. Integral Term

The integral term $K_{i}\int_{0}^{t}e(\tau)d\tau$ corrects accumulated error over time. It helps eliminate steady-state error by ensuring that small but persistent errors are corrected:

$$ $\frac{d}{dt}\left(\int_{0}^{t}e(\tau)d\tau\right)=e(t)$ $$

As long as $e(t)$ remains nonzero, the integral term grows, increasing the control input $u(t)$ until the error converges to zero.

#### 3. Derivative Term

The derivative term $K_{d}\frac{de(t)}{dt}$ anticipates future error based on the rate of change of the error. It provides a damping effect that helps reduce overshoot and oscillations in the system’s response. The term is proportional to the velocity of the error, thus slowing down the system’s response as the error decreases.

#### 4. Combined Dynamics and Stability

The overall system dynamics, taking all three terms into account, can be modeled as:

$$ $\frac{de(t)}{dt}=-K_{p}e(t)-K_{i}\int_{0}^{t}e(\tau)d\tau-K_{d}\frac{de(t)}{dt}$ $$

For a properly tuned system, the combination of the proportional, integral, and derivative terms ensures that the error will decrease over time. Specifically, the integral term guarantees that any steady-state error will be driven to zero, while the proportional and derivative terms ensure fast response and stability. Thus, as $t\to\infty$, $e(t)\to 0$.

### 6.6 Limitation

Our approach may have certain limitations in the following areas:

- 1.
It may require some prior knowledge of physics, such as frequency decomposition and domain-specific knowledge embedding;
- 2.
There may be competitive interactions between expert models, where the quality of initialization plays a decisive role in model performance;
- 3.
The PID approach may not be suitable for more complex models and datasets, and methods like Model Predication Control or reinforcement learning might need to be explored in the future;
- 4.
To reduce patch boundary effects, especially in scenarios where performance degrades with a high time and spatial dynamic (e.g. Turbulence) dataset;

### 6.7 Mathematical Notations

This section lists the mathematical notations used in the paper for reference.

**Table 6: List of Mathematical Notations**
| Symbol | Description |
| --- | --- |
| $\mathbf{u}$ | State variable (solution of the PDE), $\mathbf{u}:[0,T]\times\mathbb{X}\rightarrow\mathbb{R}^{n}$ |
| $t$ | Time variable in $[0,T]$ |
| $\mathbf{x}$ | Spatial coordinate in domain $\mathbb{X}\subseteq\mathbb{R}^{D}$ |
| $a$ | Time-independent parameter of the system, possibly varying with $\mathbf{x}$ |
| $F$ | Linear or nonlinear function representing PDE dynamics |
| $w(t,\mathbf{x})$ | Vorticity in the Navier-Stokes equations |
| $u(t,\mathbf{x})$ | Velocity field in the Navier-Stokes equations |
| $f(\mathbf{x})$ | Source term in PDEs |
| $\mathbb{X}$ | Spatial domain |
| $\partial\mathbb{X}$ | Boundary of the spatial domain $\mathbb{X}$ |
| $\Omega$ | Domain of the PDE problem |
| $\partial\Omega$ | Boundary of the domain $\Omega$ |
| $\mathbf{u}^{0}(\mathbf{x})$ | Initial condition of the PDE at $t=0$ |
| $B[\mathbf{u}]$ | Boundary operator for boundary conditions |
| $\partial_{t}\mathbf{u}$ | Partial derivative of $\mathbf{u}$ with respect to time $t$ |
| $\mathbf{u}_{\mathbf{x}}$ | First-order partial derivative(s) of $\mathbf{u}$ with respect to $\mathbf{x}$ |
| $\mathbf{u}_{\mathbf{xx}}$ | Second-order partial derivative(s) of $\mathbf{u}$ with respect to $\mathbf{x}$ |
| $\nabla^{2}\mathbf{u}$ | Laplacian operator (second-order differential operator of $\mathbf{u}$) |
| $f_{\theta}$ | Mapping function of the model, parameterized by $\theta$ |
| $E$ | Set of expert models, $E=\{E_{1},E_{2},\ldots,E_{n}\}$ |
| $E_{i}$ | Expert model $i$ in the multi-expert network |
| $\mathbf{P}_{i}$ | Segmented and interpolated patch $i$ |
| $\lambda(t)$ | Hyperparameter controlling training focus between router and experts |
| $R(\mathbf{x})_{j}$ | Probability of routing input $\mathbf{x}$ to expert $E_{j}$ |
| $p_{ij}$ | Probability of assigning data point $i$ to expert $E_{j}$ |
| $Top_{k}$ | Number of top experts selected by the router |
| $\mathcal{D}$ | Data distribution |
| $\text{Error}(E_{j},\mathbf{x})$ | Error measure of expert $E_{j}$ on input $\mathbf{x}$ |
| $\mathcal{L}(t)$ | Total loss function at time $t$ |
| $\mathcal{L}_{\text{experts}}$ | Loss function for the experts net |
| $\mathcal{L}_{\text{router}}$ | Loss function for the router mechanism |
| $\mathcal{L}_{\text{load}}$ | Load balancing loss for the router |
| $\text{KL}(R(\mathbf{x})||P(E))$ | Kullback-Leibler divergence between routing distribution $R(\mathbf{x})$ and prior $P(E)$ |
| $K_{p}$, $K_{i}$, $K_{d}$ | Proportional, Integral, and Derivative gains in the PID controller |
| $\hat{L}$ | Desired loss or target value in PID control |
| $e(t)$ | Error signal at time $t$ in PID control, $e(t)=\hat{L}-\mathcal{L}(t)$ |
| $P(t)$ | Proportional term in PID controller at time $t$, $P(t)=K_{p}e(t)$ |
| $I(t)$ | Integral term in PID controller at time $t$, $I(t)=I(t-1)+K_{i}e(t)$ |
| $\lambda_{\min}$, $\lambda_{\max}$ | Minimum and maximum values for $\lambda(t)$ |
| $\nu$ | Viscosity coefficient in Navier-Stokes equations |
| $Re$ | Reynolds number |
| $N$ | Number of training epochs |
| $S$ | Scale factor in multi-scale segmentation |
| $B$ | Batch size |
| $T_{\text{in}}$, $T_{\text{out}}$ | Number of input and output time steps |
| $H$, $W$ | Height and width of spatial domain grid |

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: Can the PID controller be replaced by the learning-based controller?
soundness: 2
strengths: 1. The motivation of this paper is clear.
2. The three modules make sense to me.
3. The authors have created numerous figures and visualizations to clearly illustrate the algorithm framework and experimental results.
summary: This paper introduces a framework of multi-scale and multi-expert (M2M) neural operators designed to simulate and learn PDEs efficiently.
There are three main modules in the M2M framework.  1) a divideand- conquer strategy to train a multi-expert gated network 2) a controllable prior gating mechanism and 3) a PI (Proportional, Integral) control strategy to adjust the allocation rules precisely. Experimental results on multi-scale and Naiver-Stokes datasets show the superior performance over deep learning-based surrogate models.
weaknesses: 1. It would be better if experiments could be conducted in more real-world scenarios to validate the results of the article.
2. There are lots of hyperparameters in the M2M method. Please give some sensitivity analysis of them.
3. The length limitation is 10 rather than 9 this year. Some important information in the Appendix can be moved to the main paper. 
4. sec. -> Sec, fig. -> Fig, etc.
5. The contribution parts are not in the enumerate environment.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: The implementation of PID control significantly aids in optimizing the training of Multi-of-Experts (MoE) -- optimizing how ?
a(t) is a time-independent parameter of the system, which can be defined at each location x, --> a(t) is time dependent not independent, explain
By bridging the control theory-PID, this unified theory demonstrates its strong generalizability -- did not demonstrate generalizability adequately
Please provide provide a specific section explaining your approach for handling boundary conditions between segments, including any techniques used to ensure continuity and consistency across segment boundaries. Also discuss potential issues that may arise from inter-segment dependencies and how their method addresses these.
Please consider including a complexity analysis section, detailing the computational and memory requirements of your approach compared to existing methods. Scaling experiments on larger problem sizes or more complex PDEs to demonstrate broader applicability will also be useful.

I am not convinced that your error dynamics developments for the PID controller are correct.  The derivations are mathematically inconsistent, please refer to one of your citations Angstorm for clarity.   
Also from Fig2. your P control is Kp/(1 + exp(e(t))  ==> for large error this -->P- term =  0 ??

You may considering fixing the title,  "...MULTI OF EXPERTS '
soundness: 2
strengths: The proposed approach appears fresh compared to existing methods like FNO, UNO, and CNO
Literature review is adequate
Paper is written well and easy to read, 
Experiments are provided to support results
Supplementary material is adequate
summary: The paper, M2M...,introduces Multi-Scale and Multi-Expert neural operators, designed to simulate PDEs efficiently. The authors propose  using domain specific expert models and a divide-and-conquer strategy combined with a controllable gating mechanism that determines the allocation of computational resources across models. 
The use of multiple models and the PI control strategy is innovative over standard approaches. The approach appears fresh compared to existing methods like FNO, UNO, and CNO. However, while the experimental evidence is strong, the theoretical underpinnings and the scope of generalizability could be expanded. A deeper critical evaluation of assumptions and additional testing across diverse conditions would enhance its impact.  Specifically in combining domain local solutions, considerations on boundary conditions is very critical. In multi-scale and divide-and-conquer approaches, ensuring smooth and consistent transitions between patches is critical, particularly when different models are used.  These issues are not discussed in the paper.  
The dynamic allocation mechanism, while innovative, might oversimplify the complexities of real-world PDEs. The M2M framework introduces considerable architectural complexity, including multiple models, a dynamic router, and a control mechanism. While this complexity is justified in the paper as necessary for handling multi-scale problems efficiently, there is limited discussion on the practical implications, such as computational cost, scalability, and ease of implementation.
The experiments focus on a few benchmark datasets, such as custom multi-scale and Navier-Stokes simulations. While these examples are relevant, the paper does not test the model on a wider variety of PDEs or real-world scenarios that could demonstrate claims on generalizability.
Additionally the discussions in section 6. 5. on the error dynamics of the PID control related discussions is mathematically inconsistent.  Specifically 
Section 4. COMBINED DYNAMICS AND STABILITY has the combined error dynamics defined such that de/dt = -u(t), this and other developments are inconsistent and incorrect with respect to control and system dynamics.
Overall I think the paper needs more work and theoretical development to provide convincing arguments toward soundness of the proposed solution.
weaknesses: The most significant weakness  is the lack of detailed explanation of how boundary conditions between segments are handled and the potential issues with inter-segment dependencies. This omission raises concerns about the physical accuracy and robustness of the method when applied to real-world problems. Coupled with the incorrect or oversimplified error dynamics analysis of the PID control law undermine the overall impact of this paper, especially for audiences that require rigorous theoretical and practical validation as is typically the case with PDE based problems that  

1. Lack of Detailed Explanation on Boundary Handling. The paper does not provide any detailed explanations on how the solutions from different models or segments are combined, especially as they relate to  the handling of boundaries between segments.
2. Insufficient Theoretical Analysis, especially related to combining solutions and model selection criteria
3. Complexity considerations for practical applications and scaling, insufficient assessment to convince on broader applicability and scaling 
4. Limited Scope of Experiments presented, results are based on only a  few benchmark datasets, ex.custom multi-scale and Navier-Stokes simulations. While relevant, the paper does not test the solution on a wider variety of PDEs or real-world scenarios to demonstrate generalizability
5. incorrect error dynamics for PID control. Derivations are inconsistent and appear to be derived to force asymptotic convergence to 0 error.   the derivation relates to a  very specific and uncommon case where the control input is directly designed to be the negative of the rate of change of the error, which is not typical in a PID controller setup. Instead, the error dynamics and the control input are related, but not in such a direct or simple manner.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 1
questions: Can you quantify the contribution of the different components of your architecture?

Can you provide a strong argument why your approach is supperior to the state of the art?

One major advantage of Neural Operators is being mesh independent, while the method presented seems mesh dependent. Can it be adapted to be mesh-independent?

Does your method generalize to time horizons that are longer than the trained sequence length?

How can the developments of states in the different patches interact with each other if they are always segmented in the same way?
soundness: 1
strengths: The idea to split the problem to different expert models as well as a regularization for computational complexity is interesting.
summary: The paper proposes a method to learn surrogates for PDE models. The proposed method builds on Fourier Neural Operators and divides the physical field to be simulated in segments. Each segment is assigned by a router to different expert models. The authors claim that their method can deal with multiple scales and multiple frequencies, offering advantages like increased flexibilty, better generalizability, more accuracy than baseline models and gain of system knowledge.
weaknesses: Motivation: The motivation is not presented well. The authors fail to adequately situate their approach within the existing body of literature. E.g. in line 14 "current methods insufficiently learn their representations" is not supported by examples or according literature. In line 50 it is claimed that "Exploring how to integrate and fully leverage performance across different scales while controlling complex learning dynamics is a promising area of research", however a reference or argumentation supporting the claim is missing. 
In section 2, autoregressive and neural operators are presented as two different approaches, however the authors do not clearly position their method within both. In line 154, the authors claim “However, the aforementioned methods do not efficiently leverage frequency characteristics, and they lack a controllable mechanism for adjusting the learning process of partial differential equations.”, which is not supported by citations, experiments or further arguments. Please provide specific examples or citations from the literature that demonstrate the limitations of current methods in learning representations of multi-scale PDEs. Furthermore, references or a brief discussion of existing work on multi-scale modeling including connecntions to the proposed approach would better support their claim about the promise of this research direction and provide more structure in the literature review. 

Results: The method does not increase the prediction quality for the Navier Stokes data set compared to the Fourier Neural Operator (FNO) (line  448: "Specifically, M2M outperforms the strong baseline of FNO128 in the performance a little."), which is not significant enough to support that this is achieved by the architecture design of the method. A more detailed discussion (especially on hyperparameter tuning) is missing.


The custom dataset: The custom dataset consists of fields of 4 squares with different dynamics. The proposed methodology seems to segment this in 16 squares, inducing an alignment of these squares. This is most probable not a fair comparison as 1) in physics, such quadratically structures field do not exist and 2) the alignment of dataset segments and segments that are distributed within the method induce an unfair bias to this method being superior. Experiments on publicly available datasets would provide a fairer evaluation. Usage of e.g. the PDEBench from the University of Stuttgart should be considered. 

Unsupported claims: E.g. Contribution 2, Line 97 “By bridging the control theory-PID, this unified theory demonstrates its strong generalizability. It is a versatile and scalable method for the machine learning and science simulation community. “ The method is only validated for two different datasets, and a reason why PID controller integration unifies a theory is not given. 
E.g. in line 442 "As shown in table 4.2, our M2M can allocate high  FNO modes to the high-frequency region" - but table 4.2 does not show this. To provide evidence on the claim of generalizability more corresponding experiments (see comment w.r.t. to datasets above) should be added. Moreover, a more detailed discussion and analysis on how the PID controller integration contributes to the method's generalizability and versatility should be added, e.g. how do the distinct controller parts (P, I and D contribute). 

Clarity: The explanation of expert nets in Line 191-193 is too short and thus hard to understandable in the current status. This section should be detailed to describe the function and concept of the expert nets. Moreover, the aggregation function in eq. 6 is should be specified. Also, a more detailed discussion on the routing should be added since it is hard to follow currently. The same holds for the explanation of PI control (252, 253).  Algorithm 1 is described very imprecisely and uses not introduced variable names. According symbols and variables should be introduced and explained. 

Distinction between the method presentation and its parametrization for the application is not given; this raises confusion about how the model is later used, e.g. starting in line 179, 4 patches are used, and then again n in equation 5, and then again 4 in line 219.

Graphics and tables with error metrics and terms that are not introduced/explained: e.g. caption of Figure 3 (“scale is set to 4”, how is the error metric defined), Figure 4 (“Prior [0000]”, “TOP_k”, how is the expert specified), Figure 5 (scale again, values plotted are not described, graphics quality is poor), Table 1 (what does the subsets mean). 

Submitted source code is not usable: The required Python environment is not specified, there is no explanation how to use models and how to get datasets, all comments are Chinese such that they are not understandable to the worldwide community. Please provide a description or HowTo specifying the Python environment, including instructions for using the models and obtaining datasets, and translate comments to English to make the code accessible to a wider audience.

Mathematical notation unclear: use of introduced subsets (N, M, equation 6), in line 235 x misses the subscript, use of same variable for different things (e.g. x in line 107 and 222), not clear where equation 7 is used, meaning of variables in equation 10 is not clear.

Misleading introduction of PI control: The authors claim that they have a proof for the application of PI control in their method and link to the appendix. But what is claimed to be a proof is not a proof, but an introduction to PID control that can be found in many teaching books. Furthermore, intensified by the not defined error metric used in Figure 3, it remains unclear if the realized convergence increase is due to a different scaling of the components of the loss function, or a real benefit. Anyhow, to achieve the best result in Figure 3, both models still need to be trained for the same number of epochs. Additionally, the authors mismatch feedback and reference in line 297 and are not clear about what they control with PI (KL or loss?)

Paper quality: Grammar, sentence structure and spelling are frequently faulty. It seems that the quality of explanations, language and presentation gradually degrades to the end of the paper.

I think the paper introduces too many new elements of a model structure, such that the implications of the single changes cannot be separated. Here an ablation study should be added.

Important parts of results are in the appendix, as well as the limitations. This should be fixed.

Minor points:

The first two sentences in the abstract are missing arguments or citations which prove the statement “Learning the evolutionary dynamics of Partial Differential Equations (PDEs) is critical in understanding dynamic systems, yet current methods insufficiently learn their representations. This is largely due to the multi-scale nature of the solution, where certain regions exhibit rapid oscillations while others evolve more slowly.”
•	The meaning of M2M is not explained.
•	Contribution 3, line 99, is not a real contribution as validation is a must in machine learning papers “We validate the aforementioned method using the standard Navier-Stokes equations and a custom multi-scale dataset, …”
•	FNO-3D allows for zero shot super resolution, while your approach does not. This should be mentioned.
•	The title of 4.1 is not exact.
•	The numbering in the appendix is wrong; it should start with letters.

Meta-review: This paper proposes M$^2$M, a framework combining multi-expert and multi-scale operators for PDE dynamics, incorporating a gating mechanism and PI control strategy. While the approach shows promise with improved accuracy on limited datasets, the paper lacks clarity in its novelty compared to prior work, has a narrow experimental scope, and provides insufficient theoretical analysis of key components like the PI control strategy. Due to these limitations, I recommend rejection at this time.

Decision: Reject