Mixture-of-Experts Operator Transformer for Large-Scale PDE Pre-Training
Authors: Hong Wang
Haiyang Xin
Jie Wang
Xuanze Yang
Fei Zha
huanshuo dong
Yan Jiang
Venue: NeurIPS 2025 poster ()
Keywords: AI4PDE
Neural Operator
Pre-Training
Primary area: machine_learning_for_sciences

Abstract: Pre-training has proven effective in addressing data scarcity and performance limitations in solving PDE problems with neural operators.
However, challenges remain due to the heterogeneity of PDE datasets in equation types, which leads to high errors in mixed training. 
Additionally, dense pre-training models that scale parameters by increasing network width or depth incur significant inference costs.
To tackle these challenges, we propose a novel **M**ixture-**o**f-**E**xperts **P**re-training **O**perator **T**ransformer (**MoE-POT**), a sparse-activated architecture that scales parameters efficiently while controlling inference costs.
Specifically, our model adopts a layer-wise router-gating network to dynamically select 4 routed experts from 16 expert networks during inference, enabling the model to focus on equation-specific features. 
Meanwhile, we also integrate 2 shared experts, aiming to capture common properties of PDE and reduce redundancy among routed experts. 
The final output is computed as the weighted average of the results from all activated experts.
We pre-train models with parameters from 30M to 0.5B on 6 public PDE datasets.
Our model with 90M activated parameters achieves up to a 40\% reduction in zero-shot error compared with existing models with 120M activated parameters.
Additionally, we conduct interpretability analysis, showing that dataset types can be inferred from router-gating network decisions, which validates the rationality and effectiveness of the MoE architecture.

Mixture-of-Experts Operator Transformer for Large-Scale PDE Pre-Training
Hong Wang 1233, Haiyang Xin'*, Jie Wangl.23t Xuanze Yang', Fei Zha', Huanshuo Dong Yan Jiang't University of Science and Technology of China
2 CAS Key Laboratory of Technology in GIPAS, University of Science and Technology of China
MoE Key Laboratory of Brain-inspired Intelligent Perception and Cognition, University of Sci- ence and Technology of China
vanghong17000mmail www.uuttc.... en, xhy28788mail ustc. edu
edu.cn, en, dlevangxxustc..ddu.c
Abstract
Pre-training has proven effective in addressing data scarcity and performance limitations in solving PDE problems with neural operators. However, challenges remain due to the heterogeneity of PDE datasets in equation types, which leads to high errors mixed training. Additionally, dense pre-training models that scale parameters by increasing network width depth incur significant inference costs. To tackle these challenges, we propose novel Mixture-of of-ererttss Pre-training Operator Transformer (MoE-POT), apttse-at architecture that scales parameters efficiently while controlling inference costs. Specifically, our model adopts a layer-we wise from 16 expert networks during inference, enabling the model to focus on equation- specific features. Meanwhile, also integrate 2 shared experts, aiming capture common properties of PDE and reduce redundancy among routed experts. The final output computed as the weighted average of the results from all activated experts. We pre-train models with parameters from 30M to 0.5B on 6 public PDE datasets. Our model with 90M activated parameters achieves up a 40% reduction in zero-shot error compared with existing models with 120M activated parameters. Additionally, we conduct interpretability analysis, showing that dataset types can be inferred from router- -at-ggttingg network decisions, which validates the rationality and effectiveness of the MoE architecture/
network dynamically select 4 routed experts
Introduction
Learning solution operators for partial differential equations (PDEs) has emerged as a fundamen- tal paradigm scientific machine learning, enabling data-driven modeling of complex physical systems through neural operators 662 22
28 [2]. These operators learn mappings between infinite-
dimensional function spaces, offering surrogate models that can outperform traditional numerical solvers by orders of magnitude in speed 2] To address the scarcity of PDE data and further enhance the performance of neural operators, recent studies have introduced pre-training techniques into neural operator frameworks [15]. However, their application to PDE learning remains nascent due to unique challenges in operator learning.
First, (challenge 1) PDE datasets demonstrate substantial variations across equation types, boundary conditions, and spatiotemporal resolutions. This diversity causes conflicting knowledge patterns when
"Equal contribution.
t Corresponding author.
'Our code is available at https//jithhbcccmmaaaaannnnnooop
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

NS (1e-5)
Multiple PDE Datasets
Zero-Sof
L2 Relestve Lavorr
F7:100
F130
Pre-Training
F Foundation Model
Espert
Expert2
Number Af MeE-POT Parameters (Millions)
Network
Bouter
Expert
NS (1e-3)
Expert
PNO
Rolatoee Droer
Downstream Tasks
Fine-Tuning
Activ (Millions)
Figure I: Left. An illustration of pre-training PDE foundation model using extensive data from diverse datasets. The pre-traue pre-trained model is subsequently fine-tuned for various downstream operator learning tasks, enabling the handling of complex scenarios. Right. (a) Comparison of errors across different numbers of fine-tuning epochs; (b) Comparison of zero-shot errors across different models.
merging different PDE types during training- direct data mixing frequently results in detrimental interference that restricts knowledge acquisition, rather than enhancing model generalization.
Second, (challenge 2) existing approaches to scaling model capacity through dense architectural expansion (increasing width/depth) incur prohibitive inference costs, making pre-training impractical for real-world applications.
To address these challenges, we propose the Mfxture-of-Expee Pre-training Operator Transformer (MoE-POT), novel sparse architecture. Our key insight is decouple capacity expansion from computational cost through dynamic expert activation.
Specifically, MoE-POT employs learnable router gatttrggating network that automatically selects 4 routed experts based on each layer's input data, working with fixed shared experts. The outputs from these 6 experts then weighted and aggregated produce the final result. The router-gating network dynamically selects routed experts that specialize in learning distinctive features of the current PDE category, effectively isolating interference from significantly different PDE data types. Meanwhile, the shared experts act as fixed computational modules for all data. They ensure consistent learning of fundamental dynamic evolution laws.
The key contributions and advantages of MoE-POT are summarized as follows:
Novel Architecture: We introduce MoE-POT, novel sparse architecture for neural operator pre-training. MoE-POT employs two types of expert networks: routed experts and shared experts, enabling balance between generalization and specialization.
Empirical Validation: As shown in Figure (right), we pre-train models on 6 public PDE datasets and design multiple versions of MoE-POT with total parameter scales ranging from 30M to 0.5B. Our model with 90M activated parameters achieves up to a 40% reduction in zero-shot error compared to existing models with 120M activated parameters.
Interpretable: We observe that the trained router-gating network can infer the PDE type of input data with 98% accuracy, showing MoE-POTTs ability to effectively handle diverse PDE datascts. This result validates both the rationale and effectiveness of the MoE architecture.

---

Similar Dunses
(b) Slanificcrryly Diftrant Datares
Natie51
KS(c-3
1212
101
120
bsciskal Traing
Anghl Stow
Reiatve 12 Love Ratio
13353
63-
Arg 73
181
1000
1000
100
Diraet
Eapar Tndes
Jales
Figure 2: Left. The impact of mixed training on multiple datasets using FNO on model performance. Right. Usage ratio of routed experts in different datasets in block 4.
Preliminaries
2.1 Time-dependent PDE Problem
We consider general form of parameterized time-dependent PDEs characterized by variables
u (x. 1) E IRm. These equations satisfy the following conditions:
du Ot Flu; 0 (x, = =0, (x,t) E E
T Rd+
(1)
(1x 0) u (x), x EEΩ,
Biu)(r, - 0, E 80.
Here, Flu: (z), F(t. 2r Oz u. U.... 0) represents differential operator involving spa- tial derivatives, while 8 E θ denotes unknown parameters that define the type and coefficients the PDE. The initial condition is given by u (z), and (1,](x,, t) specifies the boundary conditions. This general formulation encompasses a variety of fundamental PDEs.
In practical scenarios,
are often collected from multiple PDEs, represented as , E D is discretized on spatiotemporal meshes,
where D = (u:hsiss, Each solution function expressed as u, = (u......o u,) ), with u² = [(z), u/) 2j X,) for < t S T. The spatial meshes X, may consist of regular grids or irregular point clouds, depending on the geometry of the domain. The parameters 0 govern the type and specific characteristics of the PDF. However, in many real-world applications, such as climate modeling, only observational trajectories of data are available, while the detailed parameters 0 remain inaccessible. To predict future timesteps, it is essential to infer the most likely 9 implicitly from the observed sequence of T frames (u....., u.
2.2 Auto-regressive Denoising Pre-training
To effectively learn from temporal PDE datasets, we propose neural operator G. (u(tT) parameter- ized by weights which auto-regressively takes T frames as input and predicts the next frame based on the previous frames:
u = (w
u T1).
By predicting the next frame, the model learns internal representation of the underlying PDE dynamics. However, directly supervising the one-step loss has been shown to be suboptimal [3] Following DPOT [5]. we inject small-scale noise into the input frames. For Vt T, let s
denote (u ......... 2 ), and the noise E is sampled as E ~ N/O0, cl (TT). Noise injection improves robustness and reduces the discrepancy between training and inference.
We adopt the experimental setup of DPOT 15] and FNO 28 which focuses on the challenging scenario where models must infer system dynamics solely from solution trajectories, without access to the governing PDE parameters. In contrast to parameter-informet approaches, our core contribution is the integration of a MoE architecture into this auto-regressive paradigm.

---

3 Motivation
3.1 Challenges in Dense Neural Operator Pre-training
Current dense neural operator pre-training approaches face two challenges, which severely limit the model's ability to generalize across PDE tasks and hinder performance improvement.
(C1) Performance degradation due to dataset heterogeneity: As shown in Figure 2 (left), our preliminary experiments reveal that the heterogeneity of PDE datasets significantly impacts pre- training effectiveness. For instance, when training on different parameter configurations within the same equation family (e.g., fluid simulation data with varying Reynolds numbers), the average test error increases by only 80% compared to training on individual datasets. However, when mixing three entirely different equation types for training, the error increases by up 5000%.
The properties of different PDE types exhibit substantial variations, yet dense models enforce parameter sharing across all inputs, making it difficult to efficiently absorb diverse PDE knowledge within unified architecture. For example, when simultaneously learning PDEBench-SWE and PDEBench-DR, the model must encode two fundamentally different differential operators within the same parameter space, leading to negative transfer or inter-task interference
(C2) Scaling bottlenecks model capacity and performance: To better capture the diverse properties of PDEs, increasing model parameters is often necessary enhance the expressive power of neural operators. However, this approach introduces significant computational cost. As shown in Figure (right), experiments demonstrate that models follow diminishing returns scaling law: improving model performance requires substantial increase in parameters. Since all parameters are activated during inference, dense models generate high inference costs, further exacerbating computational challenges.
3.2 Sparse Pre-training with Mixture-offExpertt
To address these challenges, propose PDE pre-training model based on the MoE architecture.
Efficient scaling with sparse activation: As illustrated in Figure 3 the MoE architecture decomposes the fully connected computation of traditional dense layers into collaborative mechanism involving parallel expert networks and gated routing. Each Transformer layer consists of 16 routed experts and 2 shared experts. For a given input sample, the router-gating network activates only 4 routed experts, resulting in actual computational cost equivalent to only 33% of the total parameters. This mechanism effectively scales model capacity while controlling inference cost (C2).
Physics-driven gated routing: The dynamic selection capability of the router-gating network provides a natural solution for integrating heterogeneous PDE knowledge (C1). Shared experts: Through cross-task learning, the 2 shared experts constrained to capture universal physical principles (e.g., conservation laws, symmetry). 2. Routed experts: The remaining 16 experts autonomously develop distinct functional roles to learn the unique characteristics of different PDEs.
As shown in Figure (right), after training, the router-gating network decisions vary significantly across different PDE datasets. For example, NS(Ie-5) and NS(le-3), which are closely related datasets, exhibit similar gating patterns. In contrast, PDEBench-SY and PDEBench-DR, which differ substantially, show distinct gating behaviors. More importantly, (see Experiment (5.4) reveals that gating weights can serve as identifiers for PDE types. The expert selection for given input can be used to determine its dataset type with accuracy of 98%. This demonstrates that the MoE architecture not only improves performance but also enables implicit equation type recognition for neural operators, paving the way for building interpretable foundational models for PDEs.
interpretability analysis
Method
Overview. Our proposed model architecture is illustrated in Figure 3 It begins by processing raw data through a patchification layer and a temporal aggregation layer [S]. which reduces spatial-temppral resolution and extracts dynamic structures inherent to PDEs. The processed features are then passed through N blocks, each of which contains a Fourier layer efficient representation and specialization for diverse PDE tasks.
[3] and MoE layer, thereby achieving

---

Mixed
Weighlod Temporal MLP
A
Patche Positional Embaaking
Sampling Pudding
Noise
Temporal Aggregation Layer
Routed Export
Shared Expert
Inpal Layer
Layer
Rosler
Pre-tmined PDE Model
Mixture of Experts Layer
Blac x
Ouiput Projection Layer
Zut
Output Layer
Figure 3: An illustration of our model architecture. The process begins by sampling trajectories from mixed datasets of multiple PDEs. The model is optimized by predicting the next frame based on previous frames. The mixture-of-expprts layer consists of shared experts, routed experts, and a router-gating network. The nuter-gating network responsible for selecting the routed experts, enabling efficient and specialized processing while ensuring scalability and modularity.
Input Encoding and Temporal Aggregation. The input u E 1xxxxTXXG represents spa- tiotemporal signal with channels. To encode spatial features, we apply patchification layer with positional embeddings inspired by vision transformers [10]:
4 P(u' +'), t=
where a convolutional layer, and pl., = W, Myy t) denotes learnable positional encodings.
The output Z E #/ppx//pxC captures spatial features, with W, E Ik=*3, where is the feature dimension the positional encoding (e.g. n = C). To capture temporal dynamics, we employ a temporal aggregation layer that extracts information across adjacent time steps. For each local node feature pp R in Z we apply a learnable MLP transformation W, combined with Fourier feature constant 7 E IR:
Zagg Σw
W1 Wi-2 2/e
(4)
This aggregation enables the model to implicitly infer the underlying PDE governing parameters.
Fourier Layer. Using multi-head architecture, the Fourier layer is designed to learn complex kernel-based integral transformations that approximate PDE solutions [5). Let 2²(x) denote the feature at spatial location in the (-th block, and Z' its discretized representation. We apply kernll integral operator Ke parameterized by neural network:
(Kaa (x) = S. k(x, y: (y)dy
(5)
where s(I, y; Φ) is learnable kernel function. To reduce computational complexity, we constrain the kernel be translation-invaria s(x, y; Φ) x(I y: Φ). This reformulation allows efficient implementation in the Fourier domain:
(x)= RR F[2
(6)
Here, (x) E Rdd, (k)k) E Xd,xd, xd2 is frequency-dependent learnable transformation. and FIF-1
denote the Fourier transform and its inverse. To ensure memory efficiency and jointly attend to information from different representation subspaces, we first divide spatial features (x) into groups. The grouping is performed the channel dimension, where h is the number of heads, i.c..
Concat(z/, 2/,,.. .......)). where 2((k) E + we approximate (K²zz) using smaller MLPs:
2x F [W o(Wi, Flall by,,) (z),
5

---

where Wi, W2, E RR.//xd_/hh and b/ b2 Rd4//h are learnable parameters, and a is an activation function. We set z}} Concat 2 =bb)) that passed to the MoE layer.
Mixture Experts Layer. To facilitate sparse activation and enable expert specialization, we integrate a MoE layer specifically designed for PDE inputs. Both expert networks and router- gating networks are implemented using convolutional neural networks (CNNs) to preserve spatial information. For the reason behind the structural design, please Appendix B.2 These features are then passed to a guuee--gting network G¹ (x), which computes a vector of routing logits s' (x)( E R Nr, where N, the number of routed experts (e.g.. N, 16). The gating weights are computed via softmax function:
w/ (=b(xx)) Softmax (=)(()))) R
(8)
To maintain sparsity, only the Top-K entries in w/ (=b(())) are retained (e.g., K = 4), and the masked to zero:
TopK(w (opK(w(<<b<z))) {(k, w (26()))))
where i the index of the k-th selected routed expert and w (x) is the normalized routing weight.
Let E = (E((*) " Edii)) input. Let E = [E EN.)) denote the set of routed experts, from which the top-K are selected dynamically per input. Each expert E. or El(")) is a convolutional subnetwork that takes z (2r) as input and maps it to an output feature map of the same shape. Specifically, the final output of the MoE layer is computed as 51]
denote the set of shared experts, which are always activated for every
N.
K
(z)
1 N E (z)())) + w (z)())) 10) (=)()))
Σ
(10)
Load Balancing Objective. To encourage uniform utilization of all routed experts and avoid routing collapse [5I], we introduce load balancing loss during the training phase. Following prior work [1 8) we define the importance of each expert over batch B (IB) = B) as the sum of its routing weights:
B
Importance
Σ
wh, (x)).
(11)
We compute the coefficient of variation (CV) across all N, routed experts, and define the loss as:
Chalance Wbal CV C(Imporrtance)))
(12)
where What is a tunable scaling factor (e.g., What 0.1). This auxiliary loss regularizes the routing distribution to maintain a balanced expert load and improves overall training stability.
Loss Function. The primary objective of the model is to predict the one-step transition between samples from different datasets. The loss function defined as:
Σ (u <t +e) - r 112 +
Gaalanco:
(13)
IEET
where Gn represents the model's prediction function, <t denotes the input from previous timesteps, and e is a perturbation term. By predicting the next timestep data from previous frames, the model lears to implicitly infer the PDE's governing dynamics and propagate the solution forward in time.
5 Experiments
We conducted comprehensive experiments to evaluate the performance of MoE-POT. This section is organized as follows: Comparison with various small and pre-trained models on PDE datasets. 2. Testing knowledge transfer capabilities on downstream tasks. 3. Investigating scaling laws to understand performance trends. 4. Interpretable analysis of the router-gating network selection.
5. Analyzing model inference time 6. Ablation studies to assess the impact of hyperparameters.

---

Activated
FNO-V
PDEBench
CFDBench
Dataset
Params
NS (1e-5)
NS (Ie-3)
CNS(0.1.0001
SWE
DR
Small Model
FNO
0.5M
0.156
0.0128
0,170
0.00440
0.120
0.00761
UNet
25M
0.198
0.0245
0,357
0.0521
0.0971
0.0209
FENO
1,2M
0.161
0,0256
0,183
0.00458
0.161
0.0990
GK-T
1.IM
0.260
0.0148
0,919
0.0453
0.0120
0.419
Oformer
1.8M
0.289
0.00319
0.161
0.00474
0.991
0.00444
GNOT
2.2M
0.590
0.316
0.533
0.00199
0.930
0.0216
Pre-trained
FNO-T
10M
0.191
0.0245
0.0859
11.0
0.530
0.00601
FNO-S
30M
0.157
0.0225
0,357
0.184
0.385
0.00460
FNO-M
150M
0.141
0.00730
0.0104
0.0112
DPOT-T
7.5M
0.107
0.0155
0.0168
0.00631
0.0577
0.00673
DPOT-S
30.8M
0.0688
0.00781
0.0244
0.00392
0.0367
0.00870
DPOT-M
122M
0,0569
0,00708
0.0224
0.00247
0.0288
0.0113
Ours-T T
0.0682
0.00768
0.0105
0.00640)
0.0411
0.00529
Ours-S
90M
0.0552
0,00583
0.00959
0.00289
0.0342
0.00448
Ours-M
288M
0.0528
0.00570
0.00914
0.00299
0.0300
0.00513
Fine-tuned
DPOT-T
7.5M
0.0700
0.00725
0.0168
0.00313
0.0289
0.00391
DPOT-S
30.8M
0.0502
0.00635
0.0238
0.00315
0.0215
0.00586
DPOT-M
122M
0,0424
0.00593
0.0221
0.00260
0.0175
0,00653
Ours-T T
17M
0.0456
0.00493
0.00746
0.00305
0.0188
0.00330
Ours-S
90M
0.0361
0.00376
0.00742
0.00251
0.0182
0.00313
Ours-M
288M
0.0351
0.00388
0.00744
0.00193
0.0140
0,00398
Table 1: Results of main experiments are divided into three parts. We use L,RE as the evaluation metric, where lower L2RE indicates better performance. We bold the best results in each part. We highlight the globally best results using blue '-' indicates the error is greater than 20, signifying that failed completely. The first part is trained and evaluated individually on each dataset, the second part shows zero-shot results, and the last part shows results for fine-tuning each dataset.
Datasets. For pre- tre-training, we utilize 6 datasets sourced from 3 benchmark collections: FNO (8 PDEBench [53], and CFDBench (38]. These datasets encompass a wide range PDE types and parameters. The mathematical formulations of these PDEs are provided in Appendix B.5 To ensure consistency and compatibility across datasets, we applied preprocessing techniques such as padding and masking. Detailed descriptions of preprocessing steps can be found in Appendix B.1
Training and Evaluation. The model configurations for different scales are detailed in Appendix
B.3 Across all model sizes, we employed the Adam optimizer with a learning rate of x 10 and trained the models for 1000 epochs. Training was conducted on servers equipped with 8 RTX 4090 GPUs, each with 24 GB of memory. We use the l2 relative eтror (L2RE) as the primary metric to evaluate prediction quality, following the standard practice outlined in 28
Baseline. We selected the following influential methods as baselines for comparison, categorized
into two groups: Small models: This group includes FNO (along with Geo-FNO for irregular
datasets) 28 26
UNet (50), FFNO (55). Tr-Trranffrrmer (5), OFormer 27 and GNOT [16]
These models trained and evaluated individually each dataset.
Pre-trained models: FNO-(T/S/M): Larger-scale variants the original FNO model, with expanded parameter sizes for comparison; DPOT-(T/S/M): The tate--of-the-at pre-trained neural operator model [1]]. Both pre-trained baselines and (TESPOT-(T/S/M) are first pre-trained on six datasets and then evaluated for performance.

---

5.1 Main Experiments
the results of our main experiments. The parameter in the second row corresponds to the PDE dataset configuration, such as le - 5 for viscosity in the FNO NS dataset [28]. The activation parameter counts for MoF-POT-(T/SSM) are 17M, 90M, and 188M, respectively, with total parameter counts of 30M, 166M, and 489M.
The second part of the Table evaluates pre-trained models, including FNO variants. DPOT, and MoE-POT. Our model achieves the best zero-shot performance on 5 out of 6 datasets, with significant improvements in L.REE compared to DPOT and FNO-M. For example, on the PDEBench- C 0.01), Ours-S (with 90M activation parameters) reduces L2RE by 57% compared to DPOT-M (with 122M activation parameters). The FNO architecture, not specifically designed for pre-trane trainingg struggles to optimize datasets with highly diverse properties due to its dense network structure. This leads to instability during pre-training. resulting in large errors or training collapse on certain datasets. For instance, FNO-M fails to produce valid results on PDEBench-CNSS to excessively high L.2RE. The performance improvements of our model are primarily attributed to the expert network design within the MoE architecture, which effectively captures intrinsic features across datasets with differing properties without mutual interference. This demonstrates the effectiveness of the MoE architecture in pre-training scenarios for PDEs.
0.01) and CFDBench due
The last part of the Table [ presents the results of fine-tume fine-tuning pre-trained models on each subset for 200 epochs. Fine-tuning ttninn consistently improves performance across all datasets, with larger models yielding better results. For example, MoE-POT-M achieves the best fine-tuning results on out of datasets, reducing 1.2EE by over 50% compared to the zero-shot model on PDEBench-DR. Compared to DPOT, our MoE-based model achieves significant performance gains. These improvements stem from the ability of the MoF architecture to substantially expand the total model parameters while keeping activation parameters relatively constant, thereby enhancing model performance without increasing inference costs. Additionally, after fine-tuning, on most datasets, achieving state-of-the-ee results on 4 out of 6 datasets. This demonstrates that our model successfully learns from multiple PDE datasets simultaneously through pre me-ttrining.
model outperforms all small models
In summary, fine-tuning significantly enhances performance, suggesting that pre-training on large- scale PDE datasets is a promising and scalable approach for improving operator learning tasks. The results highlight the advantages our MoF architecture in handling complex and heterogeneous PDE data. Furthermore, for additional comparative experiments with large-scale models, including DPOT-L and Poseidon, please refer to Appendix C.2
5.2 Downstream Tasks Experiments
To evaluate the effectiveness of our pre-trained model in enhancing performance across diverse PDE downstream tasks, we conducted experiments to test its broader applicability. We selected three downstream tasks. NS (1e-4), closely related the pre-training datasets NS (Le-3), and PDEArena. which represents equations with entirely different mathematical structures. All models were trained or fine-tuned for 500 epochs, and the results are summarized in Table 2
Firstly, in all experiments, our model trained from scratch ('w/o Pre-train') consistently outperforms smaller models, demonstrating the effectiveness of the MoE-POT architecture for operator learn- ing. Secondly, across all tasks, both DPOT and MoE-POT models show significant performance improvements after pre-training and fine-tuning, far surpassing the performance of small models
DPOT
Ours
Dataset
Geo-FNO
U-Net
FFNO
w/o Pre-train
w/ Pre-train
w/o Pre-train
w/ Pre-train
NS (Ie-4)
0.107
0.413
0.220
0.0599
0.0264 0.0398
0,0480 0.0381
0.0160
CNS (1, 0.01)
0.0813
0.0827
0.390
0.0521
0.0307
PDEArema
0.154
0.167
0.161
0.111
0.0621
0.137
0.0618
Table 2: Experimental results of fine-tuning on downstream tasks. L,RE is used as the evaluation metric, where lower values indicate better performance. The first column shows the challenges associated with each downstream task. We bold the best results. "w/ Pre-train" refers to fine-tuning after pre-training, while "w/o Pre-train" refers to training from scratch.
8

---

7DEBeabh-1 OR
Beei Predettoe Accennyy for Eack Ditsurt
FNO
AcCuRRYY (1)
98
Nerther Art Domammossstsnimms....... toonmmess (Millon))
Acttraee Parmetee (Milloos)
M80-51 Mill-) #0thhmm FrEhbck rottidh c73boc
Figure 4: (a) Comparison of errors across different numbers of fine-tuning epochs; (b) Comparison of zero-shot errors across different models; (c) Dataset classification accuracy based on router-gatinn network selection.
trained from scratch. This indicates that pre-training enables the model to learn more effective and transferable representations. These results highlight the remarkable versatility of our model, which can be seamlessly extended to a wide range of downstream tasks.
5.3 Scaling Experiments
The relationship between performance and increasing model size a critical property of pre-trained models. In this study, we conduct scaling experiments to evaluate the scalability of our model. The results are presented in Figure 4 (a). We observe that as the model size increases, the zero-shot test error consistently decreases, approximately following scaling law. Furthermore, fine-tuning the model on specific datasets leads to improved performance. As shown in Figure while all models exhibit scaling properties, activated parameters. This advantage primarily stems from the MoE architecture, which significantly increases the total number of parameters without proportionally expanding the number of activated parameters. For example, MoE-POT-T with a total of 30M parameters, requires only 57% of the activated parameters (17M) compared to existing models with similar performance. Furthermore, further analytical experiments are presented in Appendices C.3 C.4. and CC.5) which include an analysis of error accumulation rollout steps, study the relationship between fine-tuning data size and performance, and an investigation into the impact of pre-training data heterogeneity.
model demonstrates better performance for given number of
5.4 Interpretable Analysis
We aim to determine which dataset given data point belongs to by analyzing the expert selection in the MoE router-gating network. For specific block, first compute the average expert selection values Y; for each dataset (forming a vector), where i 1,..., 6 represents the i-th dataset. Then, for the expert selection vector To of input, calculate its distance each Y. The Yi, with the smallest distance to it is obtained, indicating that the input belongs to the ig-th dataset. Detailed procedures can be found in Appendix B.4
(c), As shown Figure (c), the router-gating network in Block achieves an accuracy of 97.7% in classifying the imput dataset. Similar results observed in other blocks. This strongly demonstrates that the MoE architecture effectively learns the differences between PDE datasets and uses this information for classification. Furthermore, further interpretability analysis, including the emergence of classification ability and its generalization to out-of-distributiin (OOD) data, is available in Appendix C.6
DPOT
Ours Small
Model
Tiny
Small
Medium
Large
Tiny
Medium
Activated Parameters (M)
7.5
30
158
493
17
90
288
Total Parameters (M)
7.5
30
158
493
30
166
489
Inference Time (ms)
5.5
6.5
16.7
24.3
8.8
12.7
16.6
Table 3: Average single-step inference time of different models on the NS (1e 5) dataset.

---

%,-
S)
S)
CNSIOLLOO
SWE
CFDBeoch
Top
NS
NS(le
CNS601.000
SWE
CFDBench
030857
0/10111
0,00530
0J06229
0J88662
0.1046
0L00683
16
0106220
0,09922
041046
0.00539
BJ04994
0.00553
2
0.05R23
0.0777
001157
B.00429
D04142
0,00570
0.10105
0.22344
0L00KC
Table Results of ablation experiments on the influences of the number of routed experts N, (left part) and the number of expert selections Top- (right part). L2RE is used as the evaluation metric.
5.5 Inference Time Analysis
We evaluated the inference time of various models. As shown in Table 3 under the same total number of parameters, our model demonstrates significantly lower inference time compared to DPOT. For example, MoE-POT-M has total parameter count of 489M, comparable to DPOT-L, yet its inference time is only 68% of the latter, making it equivalent to DPOT-M with just 158M parameters. This efficiency is primarily attributed to the MoE structure, which activates far fewer parameters than the total parameter count. For certain PDE tasks, single computation may require 10 to 10 inference steps. The MoE structure effectively reduces inference time while preserving model performance.
5.6 Ablation Experiments
conduct ablation studies by training MoE-POT-T on pre-trained datasets and compare averaged zero-shot performance on the corresponding test datasets. As shown in Table the error remains stable when the number routed experts N, is sufficiently large. However, increasing N, leads to higher computational costs during training. Thus, select N, 16 balance between performance and efficiency. Increasing the number of expert selections (Top-K) reduces error, as more experts contribute to inference, thereby increasing the activated parameters. However, this improvement exhibits diminishing returns, so we set Top-K K = as an optimal trade-off. Additionally, analyzed the impact of the number of heads h and patch sizes the model's performance. Please refer to Appendix C.I for details.
Limitations and Conclusions
This paper introduces MoE-POT, sparse architecture designed for PDE pre-trale pre-training. By dynamically activating routed experts and leveraging fixed shared experts, MoE-POT achieves state-of of-the-art performance under the same activated parameters. Additionally, observe that the router- titter-gtinn network can effectively distinguish features of different PDEs and classify data, further illustrating the rationality of the MoE. structure.
However, we have yet to analyze the mathematical essence this classification mechanism. Exploring how PDE classification can guide the construction of more effective pre-training datasets remains an important direction for future work.
Acknowledgements
The authors would like to thank all the anonymous reviewers for their insightful comments and valuable suggestions. This work was supported by Smart-Grid National Science and Technology Major Project under contract 20252D0805500, the National Key R&D Program of China under contract 2022ZD01 19801, and the National Nature Science Foundations of China grants 124B1019, U23A20388, 62021001.
10

---

References
[1] Yannick Augenstein, Taavi Repan, and Carsten Rockstuhl. Neural operator baste surrogate solver for free-form electromagnetic inverse design. ACS Photonics, 2023.
[2] Kamyar Arzizadenesheli, Nikola Kovachki, Zongyi Li, Miguel Liu-Schiaffini, Jean Kossaifi, and Anima Anandkumar. Neural operators for accelerating scientific simulations and design. arXiv preprini arXiv:2309. 15325, 2023.
[3] Johannes Brandstetter, Daniel Worrall, and Max Welling, Message passing neural pde solvers. arXiv preprint urXiv:2202.03337, 2022.
[4] Tom Brown, Benjamin Mann, Nick Ryder, Mclanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1900 1901, 2020.
[5] Shuhao Cao. Choose a transformer: Fourier or galerkin. Advances in neural information processing systems, 34:24924-2499 24940, 2021.
[6] Rich Caruana. Multitask learning. Machine learning, 28(1):41-75, 1997.
[7] Salvatore Cuomo, Vincenzo Schiano Di Cola, Fabio Giampaolo, Gianluigi Rozza, Maziar Raissi, and Francesco Piccialli. Scientific machine learning through physics-informed neural networks: Where we are and what's next. Journal of Scientific Computing, 92(3):88, 2022.
[8] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Yu Wu, al. Deepseekmoe: Towards ultimate expert specializa- tion in mixture-of-experts language models. arXiv preprint arXiv:240I 002000060066,,... 2024.
[9] Huanshuo Dong, Hong Wang, Haoyang Liu, Jian Luo, and Jie Wang. Accelerating pde data generation via differential operator action in solution space. arXiv preprint arXiv:2402.05957, 2024.
[10] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1929, 2020.
[11] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
[12] Wenhan Gao, Ruichen Xu, Yuefan Deng, and Yi Liu. pnaretization-invariannee discretization mismatch errors in neural operators. The Thirteenth International Conference on Learning Representations, 2025.
the
[13] John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.13581 2021.
[14] Jayesh K Gupta and Johannes Brandstetter. Towards multi-sparioremporalsssaall generalized pde modeling. arXiv preprint 2tt.:22209.1161666. 2022.
[15] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying. Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training. In International Conference on Machine Learning, pages 17616-17635. PMLR. 2024.
[16) Zhongkai Hao, Zhengyi Wang, Hang Su, Chengyang Ying, Yinpeng Dong, Songming Liu, Ze Cheng, Jian Song, and Jun Zhu. Gnot: A general neural operator transformer for operator learning, In International Conference Machine Learning, pages 12556-12569, PMLR, 2023.
[17] Liu Haoyang, Wang Jie, Cai Yuyang. Han Xiongwei, Kuang Yufei, and HAO Jianye. Optitree: Hierarchical thoughts generation with tree search for L.M optimization modeling. In The Thitt Annual Conference on Neural Information Processing Systems, 2025.

---

[18] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pages 16000-16009, 2022.
[19] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 9729-9738, 2020.
[20]
Maximilian Herde, Bogdan Raonié, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Em- manuel de Bézenac, and Siddhartha Mishra. Poseidon: Efficient foundation models for pdes. In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, Tomczak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 72525-72624, Curran Associates, Inc., 2024.
[21] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7893):583--555)7 2021.
[22] George Em Kariiadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6))422-4444 2021.
[23] Ehsan Kharazmi, Zhongqiang Zhang, and George Em Karniadakis. Variational physics-informel neural networks for solving partial differential equations. arXiv preprint 11XXv:192200087. 2019.
[24] Yufei Kuang, Jie Wang, Haoyang Liu, Fangzhou Zhu, Xijun Li, Jia Zeng, HAO Jianye, Bin Li, and Feng Wu. Rethinking branching on exact combinatorial optimization solver: The first deep symbolic discovery framework. In The Twelfth International Conference Learning Representations, 2024.
[25] Xijun Li, Fangzhou Zhu, Hui-Ling Zhen, Weilin Luo, Meng Lu, Yimin Huang, Zhenan Fan, Zirui Zhou, Yufei Kuang, Zhihai Wang, et al. Machine learning insides optverse ai solver: Design principles and applications. arXiv preprint arXiv:2401.00 0006600 2024.
[26] Zhijie Li, Wenhui Peng, Zelong Yuan, and Jianchun Wang. Fourier neural operator approach large eddy simulation of dinee-dimensioall turbulence. Theoretical and Applied Mechanics Letters, 12(6),00389, 2022.
[27] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning, arXiv preprint arXiv:2205. 2022.
(28) Zongyi Li, Nikola Kovachki, Kamyar Arizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differen- tial equations. arXiv preprint arXiv:2010.0882) 2020.
[29] Zongyi Li, Nikola Borislavov Kovachki, Chris Choy, Boyi Li, Jean Kossaifi, Shourya Prakash Otta, Mohammad Amin Nabian, Maximilian Stadler, Christian Hundt, Kamyar Azizzade- nesheli, et al. Geometry-informed neural operator for large-scale 3d pdes. arXiv preprint arXiv:2309.00s: 00080333,0000...... 2023.
[30] Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar. Physics-informe neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03794, 2021.
[31] Haoyang Liu, Yufei Kuang, Jie Wang, Xijun Li, Yongdong Zhang, and Feng Wu. Promoting generalization for exact solvers via adversarial instance augmentation, 2023.
[32] Haoyang Liu, Jie Wang, Zijie Geng, Xijun Li, Yuxuan Zong, Fangzhou Zhu, Jianye HAO, and Feng Wu. Apollo-MIL... An alternating prediction-corredne neural solving framework for mixed-integer linear programming. In The Thirteenth International Conference on Learning Representations, 2025.
12

---

[33] Haoyang Liu, Jie Wang, Wanbo Zhang. Zijie Geng, Yufei Kuang, Xijun Li, Bin Li, Yongdong Zhang, and Feng Wu. MILP-studio: MILP instance generation via block structure decompo- sition. In The Thiht Annual Conference on Neural Information Processing Systems, 2024.
(34] Hongyu Liu, Haoyang Liu, Yufei Kuang, Jie Wang, and Bin Li. Deep symbolic optimization for combinatorial optimization: Accelerating node selection by discovering potential heuristics. In Proceedings of the Genetic and Evolutionary Computation Conference Companion, pages 2067-2075, 2024.
[35] Songming Liu, Zhongkai Hao, Chengyang Ying, Hang Su, Ze Cheng, and Jun Zhu. Nuno: general framework for learning parametric pdes with non-uniform data. arXiv preprint urXiv:2305.18694 2023.
[36] Lu Lu, Pengzhan Jin, Guofei Pang. Zhongqiang Zhang, and George Em Karniadakis. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218-22) 2021.
[37] Jian Luo, Jie Wang, Hong Wang, Zijie Geng, Hanzhu Chen, Yufei Kuang, et al. Neural krylov iteration for accelerating linear system solving. Advances in Neural Information Processing Systems, 37:128636-128661 2024.
[38] Yining Luo, Yingfa Chen, and Zhen Zhang. Cfdbench: A comprehensive benchmark for machine learning methods in fluid dynamics. arXiv preprint arXiv:2310.05963, 2023.
[39] Qitan Lv, Tianyu Liu, and Hong Wang. Exploiting edited large language models as general scientific optimizers. arXiv preprint arXiv:2503.09622 2025.
[40] Michael McCabe, Bruno Régaldo- Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, et al. Multiple physics pretraining for physical surrogate models. arXiv preprint arXiv:2310.02994, 2023.
[41] Katarzyna Michalowska, Somdatta Goswami, George Em Karniadakis, and Signe Riemer- Sørensen. Neural operator learning for long-time integration in dynamical systems with recurrent neural networks. arXiv preprint arXiv:2303.02243, 2023.
[42] S Karthik Mukkavilli, Daniel Salles Civitarese, Johannes Schmude, Johannes Jakubik, Anne Jones, Nam Nguyen, Christopher Phillips, Sujit Roy, Shraddha Singh, Campbell Watson, et al. Ai foundation models for weather and climate: Applications, design, and implementation. arXiv preprint 2xXiv:2399.10800 2023.
[43] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. Climax: A foundation model for weather and climate. arXiv preprint arXiv:2301. 10343, 2023.
[44] Oded Ovadia, Adar Kahana, Panos Stinis, Eli Turkel, and George Em Karniadakis. Vito: Vision transformerr-oeeratoo. arXiv preprint arXiv:2303.0880) 2023.
[45] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, al. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202.11214 2022.
[46] Tianle Pu, Zijie Geng, Haoyang Liu, Shixuan Liu, Jie Wang. Li Zeng, Chao Chen, and Changjun Fan. RoME: Domain-robust mixture-of-expprts for MILP solution prediction across domains. In The Thitty- Annual Conference on Neural Information Processing Systems, 2025.
[47] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training, 2018.
[48] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
13

---

[49] Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
[50] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted IneevenntionMIICAA 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015.
[51] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprini arXiv: (eee:7001....55 2017.
[52] Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. arXiv preprint arXiv:2306. 200.00028,0 2023.
[53] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning, Advances in Neural Information Processing Systems, 35:1596-1611 2022.
[54] Iya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp- mixer: An all-mlp architecture for vision. Advances in neural information processing systems, 34:24261-24272. 2021.
[55] Alasdair Tran, Alexander Mathews, Lexing Xie, and Cheng Soon Ong. Factorized fourier neural operators. arXiv preprint arXiv:2111.13802. 2021.
[56) Hong Wang, Zhongkai Hao, Jie Wang, Zijie Geng, Zhen Wang, Bin Li, and Feng Wu. Ac- celerating data generation for neural operators via krylov subspace recycling, arXiv preprint arXiv:2401.0955 2024.
[57] Sifan Wang, Yujun Teng, and Paris Perdikaris. Understanding and mitigating gradient flow pathologies in phforee-informed neural networks. SIAM Journal Scientific Computing, 43(5):43055-33881 2021.
[58] Sifan Wang, Hanwen Wang, and Paris Perdikaris. Learning the solution operator of para- metric partial differential equations with physics-informed deeponets. Science advances, 7(40):cabi8605. 2021.
[59] Haixu Wu, Tengge Hu, Huakun Luo, Jianmin Wang, and Mingsheng Long, Solving high- dimensional pdes with latent spectral models. arXiv preprint 2311::2300.166666 2023.
[60] Liu Yang, Siting Liu, Tingwei Meng, and Stanley Osher. In-context operator learning with data prompts for differential equation problems. Proceedings of the National Academy of Sciences, )20(39))=2100142120 2023.
[61] Zherui Yang, Zhehao Li, Kangbo Lyu, Yixuan Li, Tao Du, and Ligang Liu. Learning sparse approximate inverse preconditioners for conjugate gradient solvers on gpus. In Advances in Neural Information Processing Systems, 2025.
[62] Eleftherios C ZZchhmanglou and Dale W Thoe. Introduction to partial differential equations with applications. Courier Corporation, 1986.
[63] Xiaoyu Zhao, Xiaoqian Chen, Zhiqiang Gong, Weien Zhou, Wen Yao, and Yunyang Zhang. Recfno: a intluutioonnivrriin flow and heat field reconstruction method from sparse obser- vations via fourier neural operator. International Journal of Thermal Sciences, 195:108619 2024.
[64] Gengmo Zhou, Zhifeng Gao. Qiankun Ding. Hang Zheng, Hongteng Xu, Zhewei Wei, Linfeng Zhang, and Guolin Ke. Uni-mol: a universal 3d molecular representation learning framework. 2023.
14

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] pprovided proper justification is given (e.g., "error bars are not reported because would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The claims are put in the abstract and Section |
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
It is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? [Yes]
Answer:
Justification: For further details, please refer to the "Limitations and Conclusions" section, found in Section
15

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: For detailed information, please refer to the "Method" section Section Specific experimental parameters discussed in Section 5 Appendix B.3
Guidelines:
The answer NA means that the paper does not include experiments.
16

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. w
For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide the instruction and code in supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (nttps://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ 7/nips. /ttppp.cc/pblic/guideeee http://oeeuuiiiiiooooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
17

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
Specific experimental parameters are discussed in Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For a detailed comparison, please refer the experimental sections in Sec- tion
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-simma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Specific experimental environments can be found in Section
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
18

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We have checked the NeurlPS code ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
The work presented in this paper aims to develop a sparse architecture designed
for PDE pre-true pre-training, focused solely on the field of scientific computing. It does not have any direct positive or negative social impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We have checked this and confirmed the paper poses such risks.
19

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The original papers or URLs of the codes, models, and data sets used in this article have been cited in the paper.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each
For scraped data from a particular source (c.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pww.::s...aaseiiiia...... thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. this information is available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[Yes]
Justification:
We provide the instruction supplemental material about our code and data.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
20

---

Answer:
[NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing. editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: This paper only utilizes LLMs for writing and editing.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non ntanstandard components.
Please refer to LLM policy (https://neur)) http:///..ccooeeeeo/////////22 for what should or should not be described.
21

---

A Related Work
A.1 Neural Operators
Neural operators have emerged as powerful tools for learning solution operators of partial differential equations (PDEs) directly from data, demonstrating immense potential across diverse fields such as fluid dynamics [28], thermodynamies [63]. clectromagnetism []] and climate forecasting 45]. The wide-ranging mathematical properties and data formats associated with PDEs have driven substantial research into designing effective neural operator architectures.
For instance, DeepONet (36) employs dual-network structure consisting of branch and trunk network, while the Fourier neural operator (FNO) [28) introduces a frequency-domain approach to efficiently learn solution mappings. Building on FNO, extensions such Geo-FNO 226 NUNO
[35. and GINO [29] adapt the method to handle complex geometries. Other works [3] 41 focus on time-dependent PDEs, addressing next-time prediction challenges with specialized training strategies. Hybrid approaches like PINO [30] and PI-DeepONet [58] integrate operator learning with
physics-informed neural networks (PINNs) 49 23 57 7. leveraging physical constraints to enhance
generalization and reduce reliance on large datasets.
Further advancements have been achieved through tranffrmer-bbase architectures which incorporate techniques like patchification and linear attention mechanisms. For example, the GK-Transformer [5), OFormer (27), and GNOT [16] demonstrate strong performance on problems involving irregular geometries. AFNO [ combines the efficiency of Fourier transforms with attention mechanisms, inspired by FNO, to achieve low memory and computational costs akin to MLP-Mixer [54]. This approach has been further adapted for large-scale applications such as climate
16
forecasting 45 61].
Despite these advancements, existing neural operator methods often require task-specific training and large amounts of domain- domain-specific data, underscoring the need for more data-efficient approaches to broaden their applicability and scalability.
A.2 Pre-training Scientific Machine Learning
Pre-training has emerged highly effective paradigm for enhancing downstream tasks by training models in (sellf)superrised manner on large-scale datasets. This approach has achieved remarkable success in traditional domains such natural language processing 47 48 4 and computer vision
19 18), and is increasingly showing promise in scientific machine learning applications, including protein modeling 21). molecular representation learning [64], and climate and weather modeling
43
45 31
33 46
24
34
33].
In the context of learning PDE data, initial efforts have been made to explore pre-training across various physical systems [56 9 37 For instance, [52) designs relatively universal PDE model to collectively train data from multiple steady-state PDFs. (60] utilizes the MathGPT architecture to investigate in-context learning capabilities for PDE data. Additionally, MPP (40 introduces an auto-regressive approach for pre-training time-dependent PDE datasets.
[15) proposes auto-
regressive denoising pre-training strategy combined with a scalable Fourier-based model architecture, enabling efficient large-scale pre-training on PDE data.
Additionally, works such as Poseidon [20], which are based on a multiscale operator transformer, have achieved excellent pre-training effectiveness and generalization by employing a novel training strategy that leverages the semi-group property of time-dependent PDEs. These approaches primarily focus on mapping parameters to PDE solutions, which contrasts with the auto-regressive solution method discussed in this paper, giving each methodology distinct scope of application.
However, these works primarily rely on dense neural network architectures and do not explicitly consider the relationships between different PDE datasets, which could significantly impact model performance. There remains substantial room for exploration in pre-training models for more complex scenarios and larger parameter spaces.
22

---

B Details of Experiment Settings
B.1 Data Preprocessing and Sampling
We adopt the data preprocessing strategy proposed in DPOT 13 with modifications to ensure compatibility across diverse PDE datasets.
Data Padding and Masking. To standardize spatial resolution, we fix the resolution at / = 128, which aligns with a significant portion of the datasets. For datasets with lower resolutions, we upscale them to using interpolation. For datasets with higher resolutions, we downscale them to H using random sampling or interpolation.
To unify the number of variables (i.e., channels) across different PDEs, we pad all datasets along the channel dimension to match the dataset with the maximum number of channels, filling unused entries with a constant value (e.g.. 1). For datasets with irregular geometric shapes, we use an additional mask channel that encodes the geometric configuration of each PDR instance. This ensures consistent representation across datasets while preserving unique structural information.
Noise inserting, During the training process, we add noise to improve the stability of the training. We only insert noise during the pre training process, and do not insert noise in fine-tuning or downstream tasks. The insertion method of noise is as follows. For Vt < TT, denote u as (u", u and the noise as E [O, (1IT). Then the input is <t
Balanced Data Sampling. To balance the contribution of datasets with varying sizes, we assign an importance weight wi to each dataset. Let [Dx]] denote the number of data points in the k-th dataset, where < l: < K. The probability of sampling a data point from the k-th dataset is computed as:
Pk K|Da|
En Wi
This sampling strategy ensures that datasets with fewer samples or higher importance scores are appropriately represented during training.
Patchification Layer. We follow the patch-based tokenization ers [10]. Given spatiotemporal input tensor u E I Rx wx used in Vision Transform- we apply convolutional embedding layer with kernel size P P and stride P. This partitions the spatial domain into non-overlapping patches of size P x P. Each patch is mapped to a d-dimensional embedding vector via a shared linear projection implemented
convolution:
Co + d, kernel = P. stride : P).
This process generates sequence of patch tokens for each timestep, which subsequently fed into the downstream attention layers. By reducing spatial resolution while preserving local structure, this patchification approach enables efficient global modeling in the transformer blocks.
B.2 Design of the MoE Structure
As detailed in the Section 4 our MoE architecture built upon a CNN framework and utilizes both shared and routed experts in parallel. The design of the current MoE-POT architecture is the culmination of extensive experimentation. We explored numerous configurations, including combinations of MoE with other neural operators (such as MPP, FNO, and FFNO) and alternative MoE structural designs (e.g., omitting shared experts or employing MLP-based experts instead of CNNs). Many of these alternative designs proved unstable, yielding models with substantial errors on certain datasets that could not be corrected through fine-tuning. The final MoE-POT architecture presented in this work is the robust and effective design that emerged from this rigorous development process.
B.3 Model sizes and training details
Pre-training, We selected 3 models of varying sizes, i.e., MoE-POT- MoE-POT-SSall and Moo-POOT-Mediu. The specific parameters these models are shown Table 5 For the pre- training stage, we set the learning rate to 1 10-3 and used One-cycle learning rate schedule over 1000 epochs, with the first 200 epochs as the warm-up phase. The Adam optimizer was employed with weight decay of x 10 and momentum parameters (B1, B2) (0.9, 0.9). Training was
23

---

conducted using 8 RTX 4090 GPUs, with total batch size of 20. The patch size was set to 8. We set the weight w for all datasets. For our training process, we selected T 10 timesteps predict the next frame, aligning with the original settings of most datasets.
Attention dim
MLP dim
Routed experts
Shared experts
Model size
Activated
Size
Layers
Heads
Top-K
size
Tiny
512
512
4
4
16
2
30M
17M
Small Medium
1024
1024
6
16
166M
90M
1024
2048
8
8
16
2
489M
288M
Table 5: Configurations of MoE-POT with different sizes.
Fine-tuning, Our model supports fine-tuning across various downstream datasets while retaining the generalization capability learned during pretraining. Specifically, we freeze the parameters of the guter-ggtiigg network during fine-tuning to preserve the expert assignment strategy obtained from the joint training stage. This strategy allows the model to reuse the learned routing behavior, enabling different experts to specialize in different data distributions. Only the expert networks are updated to adapt to the target dataset, while the router-gating network continues to provide consistent and stable expert selection. This separation of routing and expert adaptation ensures more stable and efficient fine-tuning, particularly when transferring to tasks with limited dota.For the fine-tuning stage, we set the learning rate to 1 10-3 and used a one-cycle learning rate schedule over 200 epochs, with the first 40 epochs as the warm-up phase. And for the downstream tasks, we set the learning rate to 1 X 10 and used one-cycle learning rate schedule over 500 epochs, with the first 100 epochs as the warm-up phase.
Dataset size. The train and test dataset sizes used in the pre-training and fine-tuning stages are shown T And the train and test dataset sizes for downstream tasks are shown in Table It should
Size
FNO(1e-3) FNO(1e-5)
(O.,0..,0.01)
SWE
DR
CFDBench
Train
1000
1000
9000
900
900
9000
Test
200
200
200
60
60
1000
Fine-tuning
1000
1000
9000
900
900
9000
Table 6: Dataset size in pre-training and fine-tuning
be noted that NS (le-4) and CNS (1,0.01) have similar datasets in the pre-training dataset, while pdearea differs significantly from the pre-training dataset.
Size
NS(1e-4)
CNS(1, 0.01)
PDEArena
train
2000
2000
2000
test
200
200
200
Table 7: Dataset size in downstream
Details of inference. To learn from temporal PDE datasets, our network (u'TT) parameterized by weights that auto-regressively takes frames input and decodes the next frame from previous frames,
u iT =G. Gw
Vi.
By predicting the next frame, we can infer the numerical solution of the final time step based on regreseggess..n. For example, if we take the first 10 steps our input, we can predict the solution Ipred for the next 10 steps. And the ground truth is Zyt, then the loss is
| Iared zgalle
Rel-Gz
|
24

---

B.4 Interpretable Analysis Algorithms
To further investigate the router router-gating network within the MoE structure, we designed the following experiment. Our goal is to leverage the section of the router router-gating network to determine which dataset the input data belongs to.
Specifically, given an input sample X, the gating network outputs a probability vector E R16 representing the likelihood of selecting each of the 16 experts. Although only the top-4 experts used during inference, the full softmax output encodes meaningful distributional information about expert preferences.
For specific block, compute the average expert selection distribution Y; N. -j=1 Yu-- and
N; Y, Xú E I7 where N; is the number samples from i-th dataset, = 1,..., 6. Xi is the router- gating network output for the j-th sample in i-th dataset. Then, for any input X, compare its expert distribution Iu- = (Jo,1,.., 10,16) to all Y, (Y,...... Y(,16) using the cross-entropy loss function:
f(1,, Yi)
Σ lo, log(Yr,),
Suppose io represents the nearest dataset.
arg min fTo, Yi).
In this case, we classify the input data X as belonging to the ig-th dataset.
B.5 Mathematical Forms of Datasets
Here, we list the PDEs of the datasets we used for pre-training.
FNO-v 28 : The quantity of interest (Qol) the vorticity w(x, t), (x, t) [o, 1]a [0. T
and it satisfies, represents the viscosity coefficient. In paper are FNO (Ie-3), FNO (le-4) and FNO (1e-5).
Daw u Vw = vAw + f(z),
V v.u
0.
PDEBench-CNSSC O
53 We need to predict the velocity, pressure, and density fields
u(r, (), p(x, f), p(a,, t) where (ar, E [0. 1]2 [0, 1]. The PDEs are follows .y is dynamic shear viscosity and ( is bulk viscosity.In In paper are CNS (0.1,0.01), CNS (1,0.01).
0.p+ V (pw)
0,
pl0c + u Vu)
-Vp p + nAu + (s 1/33//3(VV u.,
=
pu²
= -v. (( (E+pp pu?
a 2
+
u u 0'
2
PDEBench-SSE 53:: We need to predict water depth M1,, 1) where the domain is [_1, [0, 5). The PDEs follows, In paper is SWE.
O,h + V (haa) 0.
(nmu² 2 g.hh²
a.hh)) V.
PDEBench-DR [53): We need to predict the density fields uxx The domain is [-2.5 2.5 x [0, 1] and the PDEs are as follows, paper is DR
dzu
PDEArena-NS12 14] We need to predict the velocity, pressure, and density fields
u(x, p(z, t), p(z, t) where (I, E [0, 32,2 x [0, 24]. The PDEs are as follows,in paper is PDEArena.
d. = - Vo+aVo Vp fp+f.
V. 0.
25

---

CFDBench : We need to predict the velocity and pressure fields u(x, p(x, t). The domains are different as this dataset with irregular geometries. The PDEs are follows, in paper is CFDBench.
dp( + V (pri) V (pr)
Vp V Vu Vu'),
O.
Experimental Data and Supplementary Experiments
C.1 Partial Hyperparameter Ablation Experiment
NScle-3)
NS(1o-5)
CFDBeach
NS(1e-3)
NS(1e-5)
CFDBaach
0.06748
0.04227
0,00559
0.06226
0,00819
0.00579
0.06920
0.00822
0.0/046
0.00879
0,04994
0.00663
0,06920
0,00762
0.01046
0,00639
0.04094
0,00663
0.06963
0J00709
0.01036
RJM322
0.04199
OJNISW
0.08964
0,00792
0.01301
0.00673
0.11671
0,00847
Table Results of ablation experiments on the influences of the number of heads h (left part) and patch sizes P (right part). L2RE is used as the evaluation metric.
BableBdemossttraeee that the number of heads has minimal impact on error but affects computational cost. Accordingly, choose h 4 for efficiency. Finally, medium patch sizes (P 4 or help reduce error, leading us to select P = for optimal performance.
C.2 More Comparative Experiments
We selected DPOT 15
as primary multi-physics baseline for the following reasons:
Clear Attribution Gains: Our MoE-POT architecture direct modification of the DPOT model, where replace the dense feed-forward network with proposed sparse MoE layer. This controlled comparison allows us to cleanly attribute any performance improvements directly to the MoE architecture, providing clear and rigorous validation of our core contribution.
Divergent Experimental Paradigms: Our work, following DPOT, operates under an auto-regressive paradigm, predicting future states based solely sequence of previous solution frames. In contrast, models like Poseidon [20] and MPP 40 are designed for a parameter-informel setting, where they take explicit problem parameters (e.g.. coefficients, boundary conditions) as input to predict future state. The public benchmark datasets used in our primary experiments (from FNO, PDEBench, and CFDBench) do not provide these explicit PDE parameters, making direct comparison with parameter-informed models infeasible under our main experimental protocol.
This section provides additional experimental results and analysis to supplement the main paper. We compare our MoE-POT architecture with the larger DPOT-L model and the Poseidon.
C.2.1 Experimental Results with DPOT-L
To provide a more comprehensive comparison against large-scale dense models, we evaluated the performance DPOT-L (493M parameters). The results, alongside our MoE-POT models and smaller DPOT variants, are presented in Table
Model & Activated Params
NS(le=5)
NS(Ie-3)
CN.(0.1,0.01)
SWE
DR
CFDBench
DPOT-S (31M)
0.0688
0.0078
0.0244
0.0039
0.0367
0.0087
DPOT-M (122M)
0.0569
0.0071
0.0224
0.0025
0.0288
0.01
DPOT-L (493M)
0.0576
0.0061
0.0113
0.0023
0.0219
0.0065
MoF-POT-T (17M)
0.0682
0.0077
0.0105
0.0064
0.041
0.0053
MoE-POT-S (90M)
0.0552
0.0058
0.0096
0.0029
0.0342
0.0045
MoE-POT-M (288M)
0.0528
0.0057
0.0091
0.0030
0.0300
0.0051
Table 9: Zero-shot L.2 Relative Error (L2RE) comparison, including DPOT-L.
26

---

The results in Table Table|9leat lead to two key observations: 1. Diminishing returns for dense models: The performance the dense DPOT architecture shows diminishing returns with scale. The improvement from DPOT-M (122M) to DPOT-L. (493M)-2 4x increase in parameters- is marginal on several datasets (e.g., (1e3)))) and modest on others. Competitive performance with higher efficiency: When comparing MoE-POT-M (288M activated) with DPOT-L (493M activated), our model achieves competitive, and in some cases superior, performance. MoE-POT-M outperforms DPOT-L on three of the six datasets (NS(Ie=3), CNS, CFDBench), while DPOT-L holds slight advantage on the other three.
C.2.2 Experimental Results with Poseidon
To address the interest in comparing with the latest models, we conducted supplementary fine-tuning experiments on two challenging downstream tasks from the Poseidon paper (20): Wave-Layer and Wave-Gauss. We evaluated both MoE-POT and Poseidon under two distinct settings to fairly assess their capabilities.
Setting 1: Auto-regressive (Our Native Setting) In this setting, models predict future states using only previous solution trajectories, without access to explicit PDE parameters. The results are shown in Table TO
Model (Activated Params)
Wave-Layer
Wave-Gauss
Poseidon-T (21M)
0.29
0.29
Poseidon-B (158M)
0.21
0.24
MoE-POT-T (17M)
0.07
0.07
MoE-POT-S (90M)
0.05
0.06
Table 10: 1.2 Relative Error downstream tasks in the auto-regressive setting. Lower is better.
Setting 2: Parameter-Informe (Poseidon's Native Setting) In this setting, models are provided with explicit PDE parameters additional input. We adapted our MoE-POT model accept these parameters to ensure a fair comparison. The results are shown in Table
Model (Activated Params)
Wave-Layer
Wave-Gauss
Poseidon-T (21M)
0.08
0.06
Poseidon-B (158M)
0.06
0.09
MoE-POT-T (17M)
0.11
0.14
MoE-POT-S (90M)
0.06
0.10
Table 11: L2 Relative Error on downstream tasks in the parameter-informed setting. Lower better.
The results from these supplementary experiments indicate that each model excels in its native operational setting. In the auto-regressive setting (Table 10). where PDE parameters are unknown, MoE-POT significantly outperforms Poscidon. This highlights our model's strength in implicitly learning system dynamics from solution trajectories alone.
Conversely, the parameter-informe setting (Table 11 Poseidon generally demonstrates superior performance, showcasing its effectiveness when explicit physical knowledge is available. These findings suggest that MoE-POT and Poseidon have different primary application scopes rather than one being definitively superior across all scenarios.
C.3 Rollout Error Different Timesteps
A key challenge in auto-regressive prediction is the accumulation of errors. Even small improvement in single-step prediction accuracy can lead to substantial reduction in the cumulative error over long rollout, prediction inaccuracies propagate through the sequence.
27

---

To illustrate this effect, we analyze the rollout error at different timesteps for the Shallow Water Equations (SWE) dataset. Table compares the L2RE of DPOT-S and MoE-POT-S frames 50, 70, and 100 of the rollout.
Model
Frame 50 L2RE
Frame 70 L2RE
Frame 100 L2RE
Average L2RE
DPOT-S MoE-POT-S
0.0031
0.0034
0.0051 0.0035
0.0039 0.0029
0.0024
0.0026
Table 12: Illustration of error accumulation on the SWE dataset. The table shows the L2RE at specific frames during a 100-step rollout.
As demonstrated in Table 12 the error for both models increases over the rollout period, confirming the effect of error accumulation. More importantly, the performance advantage of MoE-POT-- grows significantly over time. The relative error reduction compared to DPOT-S is approximately 23% at frame 50, but this gap widens to over 31% by frame 100. This super-linear divergence underscores the critical impact of achieving lower single-step prediction error, as its benefits amplified during long-term, multi-step rollouts.
C.4 Analysis of Fine- Sine-Tuning Sample Efficiency
This section investigates the relationship between the number of fine-tuning samples and model performance, thereby analyzing the data efficiency of the MoE-POT architecture. We conducted few- shot fine-tuning experiments on both an in-distribution and an out-of-distribution task to demonstrate how performance scales with data availability for both MoE-POT and dense counterpart, DPOT.
We compare the performance of MoE-POT-S and DPOT-S on two fine-tuning tasks: 1. In-Distribution Task: The NS (le-4) dataset, which is closely related to the data used during pre-training. 2. Out-of- Distribution Task: The Wave-Layer dataset from the Poseidon (20) which represents a novel physical system.
For each task, we fine-tuned both models for 500 epochs while varying the number of available training samples, and report the final L2 Relative Error. The results for the in-distribution and out-of-distribution tasks are presented in Table 13 and Table respectively.
Number of Samples
16
32
64
128
512
2000
DPOT-S MoE-POT-S
0.25
0.20
0.13
0.09
0.044
0.026
0.20
0.15
0.11
0.07
0.040
0.016
Table 13: L2 Relative Error on the in-distribution NS (le-4) task versus the number of fine-tuning samples. Lower values are better.
Number of Samples
16
32
64
128
DPOT-S
0.41
0.33
0.26
0.19
MoE-POT-S
0.34
0.26
0.20
0.14
Table 14: 1.2 Relative Error on the aut-of-distribution Wave-Layer task versus the number of fine- tuning samples. Lower values are better.
The results demonstrate a clear trend: while both models improve with more data, MoE-POT-S consistently outperforms DPOT-S across all sample sizes on both tasks. This highlights the superior data efficiency of the MoE-POT architecture. The larger capacity and specialized experts of the pre-trained MoE-POT model enable it to generalize more effectively from limited data. This means it can either achieve a target performance level with significantly fewer fine-tuning examples or deliver superior accuracy given the same amount of data.
28

---

C,5 Performance with Increasing Dataset
core motivation for our work is the challenge of negative transfer [6 dense neural operators when pre-trained on mixture of heterogeneous PDE datasets. A single. dense network struggles to learn conflicting physical laws, which can degrade performance more diverse data is added. This section presents an experiment designed to test this hypothesis and demonstrate the robustness of the MoF-POTT architecture in mitigating this issue.
To investigate the impact of increasing data heterogeneity, we pre-trained both the dense DPOT-S model and our sparse MoE-POT-S model on progressively larger and more diverse mixtures of datasets. For each experiment, the models were trained from scratch on the specified data mixture. We evaluated their zero-shot performance the original six pre-training datasets to measure how well they retained knowledge.
The dataset mixtures constructed as follows:
6 Datasets: The standard pre-training set used in our main experiments: NS(1e-5), NS(1e-3), CNS(0.1, 0.01), SWE, DR, and CFDBench.
10 Datasets: The base set plus four additional datasets from the DPOT paper 15 NS(1e-4),
CNS(L, 0.1), and two Navier-Stokee tasks from PDEArena.
12 Datasets: The 10-dataset mix plus two additional CNS variants: CNS(1, 0.01) and CNS(0.1, 0.1).
Model (Pre-trained on)
NS(1c-5)
NS(1c-3)
CNS(0.1,0.01)
SWE
DR
CFDBench
Dense Model
DPOT-S (6 Datasets)
0.0688
0.0078
0.0244
0.0039
0.0367
0.0087
DPOT-S (10 Datasets)
0.0663
0.0069
0.0224
0.0037
0.0575
0.0146
DPOT-S (12 Datasets)
0.0739
0.0079
0.0129
0.0105
0.0724
0.0075
Sparse Model (Ours)
MoE-POT-S (6 Datasets)
0.0552
0.0058
0.0096
0.0029
0.0342
0.0045
MoE-POT-S (10 Datasets)
0,0521
0,0053
0,0085
0.0029
0.0371
0.0047
MoE-POT-S (12 Datasets)
0.0533
0.0056
0.0062
0.0032
0.0383
0.0043
Table 15: Zero-shot L2RE on the six base evaluation datasets after pre-training on increasingly heterogeneous data mixtures (6, 10, and 12 datasets). Lower values are better.
The zero-shot L2 Relative Error (L2RE) for both models across the different pre-training configura- tions presented in Table
DPOT-S (Dense Model) The performance of DPOT-S is unstable and often degrades as more het- erogeneous data is introduced. While adding four datasets (from 6 10) yields minor improvements on some tasks, causes significant performance degradation on others (e.g., DR and CFDBench). Expanding 12 datasets results in notable performance collapse on several tasks (e.g.. NS(1e-3), SWE, DR) compared to the original 6-dataset training. This confirms that the dense architecture suffers from negative transfer, where the model's capacity overwhelmed by conflicting information from diverse physical systems.
MoE-POT-S (Sparse Model) In stark contrast, MoE-POT-S demonstrates remarkable robustness. As the number of pre-training datasets increases from 6 to 12, its performance remains stable or even improves most tasks (e.g., CNS and CFDBench). The slight variations in error minor compared to the drastic fluctuations observed with DPOT-S. This stability indicates that the MoE architecture effectively mitigates negative transfer by allowing different experts specialize in distinct physical dynamics, thereby preventing knowledge conflict.
C.6 Extended Interpretability Analysis: Emergence and Generalization of Router Specialization
This appendix expands on interpretabillty analysis in Section 5.4 which demonstrated the router- gating network's 's ability to classify input data by its source PDE. This capability is not explicitly
29

---

programmed; the router only processes tokenized inputs, and the model's loss function is based on prediction error, not classification objective. Here, we investigate two key questions: (I) How does this specialization emerge during pre-training? (2) Does this learned capability generalize to entirely new, OOD datasets?
Emergence of Specialization During Training We first examine how the router's 's classification ability develops. We tracked the dataset classification accuracy at different stages of pre-training, using the method described in Section 5.4 The results, shown in Table [T6] confirm that this is an emergent property. Initially (Epoch 50), the accuracy is low, but it rapidly improves and reaches 100% by Epoch 250. This demonstrates that the router learns to distinguish between data distributions as part of the end-to-end optimization process.
Epoch NS(1e-5) Accuracy CFDBench Accuracy
50
2%
70%
150
80%
78%
250
100%
100%
Table 16: Evolution of the router's dataset classification accuracy over training epochs. The accuracy steadily improves, eventually reaching 100% as the experts and router co-specialize.
Generalization to Out-of-Distriuution Datasets To test if this specialization generalizes beyond the pre-training distributions, we evaluated the router's classification performance on two OOD tasks from the Poseidon benchmark (00: Wave-Layer and Wave-Gauss. These datasets represent novel physical systems unseen during pre-training.
Unseen Dataset
Block-1 Accuracy
Block-2 Accuracy
Wave-Layer Wave-Gauss
100%
100%
100%
100%
Table 17: Router classification accuracy unseen (OOD) downstream tasks. The perfect accuracy demonstrates strong generalization.
As shown in Table 17. the router achieves 100% classification accuracy both unseen tasks. For instance, when processing the Wave-Layer dataset, the router in Block-2 consistently activated a sparse subset of experts (e.g.. Expert 11 at 100% usage, Expert 79%), with other experts receiving minimal or zero activation.
Taken together, these results provide strong evidence that the router's ability identify PDE types an emergent property of joint optimization. To minimize the global prediction loss across hetero- gencous datasets, the model learns to partition its knowledge, routing inputs with similar dynamics to specialized experts. This process effectively mitigates the negative transfer that hinders dense architectures. Crucially, the perfect classification accuracy OOD data demonstrates that the router not merely memorizing training distributions. Instead, it learns to recognize fundamental properties of the underlying physics from the tokenized solution data, capability that generalizes to novel systems.
30


Reviews:
[{"review_id": "aykXZacubJ", "reviewer": "Reviewer_J1rL", "summary": "This paper introduces MoE-POT, a Mixture-of-Experts Pre-training Operator Transformer, which modifies the MLP Layer of a DPOT[1] model using a router‐gating network that dynamically selects 4 routed experts from 16 total, alongside 2 shared experts, per layer. Pre-trained on six public PDE datasets with total parameters ranging from 30 M to 0.5 B (activated parameters 17 M–288 M), MoE-POT achieves up to 40% reduction in zero-shot L₂ relative error compared to existing dense DPOT models and infers input PDE types with 98% accuracy via gating patterns. \n\n[1] \"DPOT: Auto-regressive denoising operator transformer for large-scale pde pre-training\" Hao et al, ICML 2024", "questions": "* The paper claims that MoE-POT-M(same number of parameters as DPOT-L) which has a larger number of activated parameters than DPOT-M are equivalent. While inference time reported is similar, the accuracy improvement over DPOT-M seems to come at a significant memory and training cost. Moreover a comparison with DPOT-L is not provided. Can the authors provide a detailed analysis of the training and memory overhead and also add the DPOT-L results?\n* Does MoE effectively handle diverse PDE datasets as claimed by the paper? An ablation study showing performance gains with increasing number of datasets as compared to a normal DPOT would justify this claim.\n* It seems like the MLP layer of other transformer-based PDE foundation models like MPP[2] can also be replaced by an MoE layer resulting in a similar architecture. What makes MoE-POT unique?\n\n[2] \"Multiple Physics Pretraining for Physical Surrogate Models\" McCabe et al. NeurIPS 2024", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "### Strengths:\n* Interpretability: The router-gating network reliably encodes dataset identity showing that expert selections reflect underlying PDE types\n### Weakness:\n* Novelty: The architecture is just DPOT with the MLP layer replaced with an MoE, which is not enough novelty in the architecture.\n* I don't see a clear win for using an MoE layer in the architecture. The accuracy gain seems to come at a huge cost in terms of memory and training cost.\n* Lack of Theoretical Analysis: The paper claims that router-gating decisions yield dataset classification but does not analyze the mathematical underpinnings or guarantees of this mechanism which is noted in limitations.", "was_revised": true, "final_justification": "The rebuttal answers a lot of my questions. I have been following the other discussions as well and MoE-POT does seem to provide a small edge over DPOT. Thus I raise my score to acceptance", "rebuttal": "Dear Reviewer,\n\nThank you for your thoughtful and constructive feedback. We appreciate your time and have carefully considered your suggestions. We hope the following clarifications address your concerns and further highlight the significance of our work.\n\n---\n\n### **On the Novelty and Uniqueness of MoE-POT (Weakness 1 & Questions 3)**\n\nOur core innovation is not merely substituting a layer but systematically addressing the critical bottleneck of data heterogeneity in large-scale PDE pre-training. MoE-POT is presented as a well-validated and highly interpretable solution to this specific, previously underexplored challenge.\n\n- **Problem-Driven Innovation:** Prior works have overlooked the \"knowledge conflict\" that arises from mixing diverse PDEs. Our experiments show this is a critical issue (Figure 2, left), as dense models suffer a drastic performance drop when trained on heterogeneous datasets.\n- **Unique and Validated Architectural Solution:** The uniqueness of MoE-POT stems from its \"divide-and-conquer\" strategy, enabled by the combination of routed and shared experts. The routed experts specialize in the unique characteristics of specific PDEs, while the shared experts capture common physical principles. Crucially, we validate this mechanism's effectiveness through interpretability analysis (Figure 2, right; Figure 4c), demonstrating that the model learns to dynamically select experts based on the PDE type with an accuracy of up to 98%.\n- **Rigorous Experimental Exploration:** The current MoE-POT architecture is the result of extensive experimentation. We explored numerous combinations of MoE with other neural operators (including MPP, FNO, and FFNO) and various MoE structures (e.g., without shared experts, using MLP-based experts instead of CNNs). These alternative designs frequently failed, resulting in models with substantial errors on certain datasets that could not be rectified through fine-tuning. The final version of MoE-POT is the stable and effective architecture that emerged from this rigorous process.\n\nTherefore, the novelty of MoE-POT is not in an isolated architectural tweak but in its systematic demonstration that a sparse expert architecture is an effective and interpretable solution to the data heterogeneity challenge in PDE pre-training, supported by comprehensive evidence that prior models have not explored.\n\n### **On the Cost-Benefit Analysis and Comparison with DPOT-L (Weakness 2  & Question 1)**\n\nThank you for your detailed questions regarding the model's efficiency and training costs. We would like to clarify the cost-benefit trade-off of the MoE architecture and address the comparison with DPOT models.\n\n- **Clarifying the Comparison:** You rightly point out the comparison between our MoE-POT-M (489M total, 288M activated) and DPOT-L (493M total & activated). Our key argument is that the MoE architecture decouples total parameters (storage cost) from activated parameters (computational cost), achieving superior efficiency.\n- **Greater Parameter Efficiency & Faster Inference**: Our MoE-POT-S (90M activated) outperforms the larger DPOT-M (122M activated) on several tasks. Furthermore, MoE-POT-M (489M total) has an inference time comparable to the much smaller DPOT-M and is 32% faster than the similarly sized DPOT-L, a critical advantage for iterative PDE solving.\n- **Training Overhead and DPOT-L Results:** We acknowledge that the large size of the L-series models made it challenging to complete the full pre-training for MoE-POT-L within our available computational resources, which is why a direct comparison was not included in the main table. We also recognize that MoE requires more memory to store the total parameters. However, we argue this is a highly valuable trade-off for achieving superior model performance, stronger generalization, and critically, higher inference efficiency. This rationale is widely accepted in other fields, such as large language models. In practical, large-scale deployments, memory usage can be optimized by loading shared expert networks that are used across multiple parallel models.\n\nTo further address your concerns, we provide supplementary experimental results (L2RE) for DPOT-L below:\n\n| Model and Activation Parameter | NS(1e-5) | NS(1e-3) | CNS(0.1,0.01) | SWE    | DR     | CFDBench |\n| ------------------------------ | -------- | -------- | ------------- | ------ | ------ | -------- |\n| DPOT-S 31M                     | 0.0078   | 0.0688   | 0.0244        | 0.0039 | 0.0367 | 0.0087   |\n| DPOT-M 122M                    | 0.0071   | 0.0569   | 0.0224        | 0.0025 | 0.0288 | 0.0113   |\n| DPOT-L 493M                    | 0.0061   | 0.0576   | 0.0113        | 0.0023 | 0.0219 | 0.0065   |\n| MoE-POT-T 17M                  | 0.0077   | 0.0682   | 0.0105        | 0.0064 | 0.0411 | 0.0053   |\n| MoE-POT-S 90M                  | 0.0058   | 0.0552   | 0.0096        | 0.0029 | 0.0342 | 0.0045   |\n| MoE-POT-M 288M                 | 0.0057   | 0.0528   | 0.0091        | 0.003  | 0.03   | 0.0051   |\n\nAs the table shows: 1) DPOT's performance improves marginally when scaling from 122M to 493M parameters. 2) When comparing DPOT-L (493M activated) with MoE-POT-M (288M activated), DPOT-L performs best on two datasets, while MoE-POT is superior on the other four. This indicates that MoE-POT achieves comparable, and in many cases better, performance with nearly half the activated parameters, which aligns with our claims. We will clarify this cost-benefit analysis and incorporate more detailed results for L-sized models in the final version of the paper.\n\n### **On the Lack of Theoretical Analysis (Weakness 3)**\n\nWe fully agree with your assessment and have candidly noted this in Section 6, \"Limitations and Conclusions\". The current work focuses on empirically demonstrating the feasibility and effectiveness of the MoE architecture for PDE pre-training. Our results, particularly the 98% dataset classification accuracy, provide strong empirical evidence for the mechanism that the router network can identify PDE types. Investigating the underlying mathematical principles and theoretical guarantees is a valuable and important direction for future research, which we are excited to pursue.\n\n### **On Effectively Handling Diverse PDE Datasets (Question 2)**\n\nThank you for this insightful question. We believe the existing evidence in our paper robustly demonstrates that MoE-POT handles heterogeneous datasets more effectively than dense models like DPOT, due to its inherent architectural advantages.\n\n- **Fundamental Limitation of Dense Models:**  Our work highlights a core weakness of dense models. As shown in Figure 2 (left), when training on PDEs with disparate properties, the performance of dense model FNO degrades sharply, with errors increasing by up to 5000%.\n- **MoE-POT's Architectural Advantage**: MoE-POT is explicitly designed to resolve this \"knowledge conflict,\" and its router network intelligently identifies different PDE tasks and assigns them to specialized expert networks. Our interpretability analysis provides direct evidence for this mechanism: the router's selection weights alone can infer the input data's PDE type with 97.7% accuracy. This confirms that MoE-POT can effectively distinguish and process heterogeneous data in a way that dense architectures like DPOT cannot.\n- **Superior Performance as a Result:** This architectural superiority translates directly into leading performance. In the zero-shot comparison in Table 1, our MoE-POT-S (90M activated params) reduces the error on the PDEBench-CNS dataset by 57% compared to the larger DPOT-M (122M activated params) after pre-training on six heterogeneous datasets. This clearly shows MoE-POT learns more effectively from diverse data with lower computational cost.\n\nTo further address your concern and illustrate the relationship between the number of datasets and performance for a dense model, we conducted the following supplementary experiment with DPOT-S, measuring its performance on the six core datasets as we increased the total number of pre-training datasets.\n\n|                        | NS(1e-5) | NS(1e-3) | CNS(0.1,0.01) | SWE    | DR     | CFDBench |\n| ---------------------- | -------- | -------- | ------------- | ------ | ------ | -------- |\n| DPOT-S (6 datasets)    | 0.0078   | 0.0688   | 0.0244        | 0.0039 | 0.0367 | 0.0087   |\n| DPOT-S (10 datasets)   | 0.0069   | 0.0663   | 0.0224        | 0.0037 | 0.0575 | 0.0146   |\n| DPOT-S (12 datasets)   | 0.0079   | 0.0739   | 0.0129        | 0.0105 | 0.0724 | 0.0075   |\n| MoE-POT-S (6 datasets) | 0.0058   | 0.0552   | 0.0096        | 0.0029 | 0.0342 | 0.0045   |\n\nThese results show that: 1) Regardless of the number of training datasets, DPOT-S does not match the performance of MoE-POT-S. 2) For a fixed-capacity dense model, simply increasing the number of training datasets does not guarantee better performance. In fact, performance degraded on four of the six datasets when scaling from 6 to 12. This reinforces our central claim that mixing disparate PDE datasets in a dense model leads to performance degradation and highlights the rationale for our MoE architecture. We will add a more detailed analysis and these experiments to the final version.\n\n---\n\n#### **Thanks again** \n\nWe sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly."}, {"review_id": "4mdKk0ganf", "reviewer": "Reviewer_wk8j", "summary": "This paper trains a mixture of experts model for surrogate PDE models. Since different PDEs have different types of dynamics, initial conditions as well as boundary conditions, the authors show that training a mixture of experts model where a router routes different pdes to different experts (along with having two shared experts). The authors show that their model with 90M active parameters improves upon the zero-shot performance of baselines, and show that different experts do end up taking in PDEs of similar form, and therefore the type of the PDE data can potentially be inferred from the expert that the model ends up going to. \n\nThe architecture in this paper is similar to the DPOT architecture with Fourier attention layers with appropriate temporal aggretation layers. For the mixture of experts and to maintain sparsity, the Top-K activations are used. Appropriate regularization is applied to ensured that not all the examples go through the same tokens (Load balancing objective). \n\nThe authors train their method on 5 datasets from PDEBench and CFDBench., and show that their model consistently outperforms multiphysics baselines like DPOT, as well as single model per dataset type of baselines like FNO.", "questions": "the authors mention something to the form of “model leads to feature competition” and “knowledge illusion” in lines 95-96. What do they mean by that, and more importantly, can they cite relevant works that coin these terms?", "limitations": "Yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper is easy to read and follow. \n\nI think that this is perhaps the first MOE papers for PDE solvers, something that makes a lot of sense, since as the authors mention, each PDE has different dynamics, and variability through initial and boundary conditions. \n\nThe fact that similar PDEs are routed through same experts, and that this emerges is quite promising, and interesting. \n\nHowever, there are few points that make me question a lot of the contributions of the paper. \n\n- First of all, the multi-physics models that they try are DPOT, and not models like Poseidon and MPP (especially Poseidon since they are currently the best performing multiphysics models).\n- The authors compare their numbers with FNO, but its a very old model and since then a lot of improvements have been suggested on top of FNO (such as FFNO, and others) that could have better performance over FNO.\n- The datasets that the authors train on, like PDE bench are not very difficult, and it is quite easy to get a good models on these datasets, therefore it is unclear how good the architecture is here.\n- The authors mention that they finetune the model for 200 epochs, that is quite a lot, esp since the model has been pre-trained. The authors do show that training from scratch does not work as much (in Table 2), but for single model per dataset baselines like FNO (and its better performing derivatives like FFNO), it would be interesting to see how well it performs if it is given the overall compute budget (of pretraining + finetuning).", "was_revised": true, "final_justification": "The authors have address all my concerns and also added extra experiments, and have comparisons with more recent multi-physics models. I thank the authors for the effort and have also updated my scores accordingly.", "rebuttal": "Dear Reviewer,\n\nThank you for your thorough review and positive assessment of our work. We are grateful for your insightful feedback and for recognizing the novelty and promise of applying a Mixture-of-Experts (MoE) architecture to PDE solvers. Your questions have been instrumental in helping us clarify our contributions. We have carefully considered your points and provide the following detailed responses.\n\n--- \n\n### **On the Choice of Baselines (DPOT vs. Poseidon/****MPP****) (Weakness 1)**\n\nThank you for this valuable suggestion. We chose DPOT as our primary multi-physics baseline for three main reasons:\n\n- **Clear Attribution of Gains:** Our MoE-POT architecture is built upon the DPOT model, where we specifically replace the dense feed-forward network with our sparse expert layer. This design allows us to cleanly attribute performance differences to the advantages conferred by the MoE architecture itself, providing a direct and compelling validation of our core method. Comparing with structurally disparate models like Poseidon would make it difficult to isolate the source of any performance gains.\n- **Differences in Experimental Setup:** Our work, following DPOT, uses an auto-regressive paradigm: predicting the next frame from a sequence of previous frames. In contrast, models like MPP and Poseidon have different setups. For instance, Poseidon takes problem parameters, boundary conditions, and initial conditions as input to predict a future state. As noted in the original DPOT paper, DPOT outperformed MPP in most scenarios. Furthermore, Poseidon requires PDE parameters as input, which are not available in the public datasets we use (FNO, PDEBench, CFDBench), making a direct comparison infeasible under our experimental protocol.\n- **Contemporaneous Work:** Poseidon is an excellent and cutting-edge model. At the time we began our experimental cycle, this work had not yet been publicly released and was developed concurrently with our own. Therefore, we were unable to include it in our initial baseline comparisons.\n\nWe agree that comparison with the latest state-of-the-art models is crucial. To address your concern, we conducted supplementary fine-tuning experiments on two challenging downstream tasks from the Poseidon paper, Wave-Layer and Wave-Gauss, using the $L_2$ relative error metric. We considered two settings:\n\n1. **Our Setting (Auto-regressive, no parameters):**\n\n| Model (Activated Params) | Wave-Layer | Wave-Gauss |\n| ------------------------ | ---------- | ---------- |\n| Poseidon-T (21M)         | 0.29       | 0.29       |\n| Poseidon-B (158M)        | 0.21       | 0.24       |\n| MoE-POT-T (17M)          | 0.07       | 0.07       |\n| MoE-POT-S (90M)          | 0.05       | 0.06       |\n\n2. **Poseidon's Setting (Parameter-informed):**\n\n| Model (Activated Params) | Wave-Layer | Wave-Gauss |\n| ------------------------ | ---------- | ---------- |\n| Poseidon-T (21M)         | 0.08       | 0.06       |\n| Poseidon-B (158M)        | 0.06       | 0.09       |\n| MoE-POT-T (17M)          | 0.11       | 0.14       |\n| MoE-POT-S (90M)          | 0.06       | 0.10        |\n\nThese results indicate that each model excels in its native experimental setting, suggesting they have different scopes of application rather than one being definitively superior. Due to the time constraints of the rebuttal period, we could not complete an exhaustive comparison. However, if the paper is accepted, **we commit to including a more thorough comparison with Poseidon and other recent** **SOTA** **models in the final version**, along with citations to these important works.\n\n### **On Comparison with Advanced FNO Variants (Weakness 2)**\n\nThank you for your close reading. We completely agree that a comparison limited to the original FNO would be insufficient. In fact, in **Table 1 of our paper, we have already conducted a comprehensive comparison against a range of more advanced single-task baselines.** This includes not only the **FFNO** you mentioned but also **GK-Transformer, OFormer, and GNOT**. The results demonstrate that our fine-tuned MoE-POT model outperforms these state-of-the-art single-task baselines on the vast majority of datasets, validating the effectiveness of our pre-training framework.\n\n### **On the Difficulty of Pre-training Datasets (Weakness 3)**\n\nWe appreciate your concern regarding the choice of datasets. We selected them based on the following considerations:\n\n- **Standard and Fair Benchmarking:** The six datasets we use are sourced from three widely-accepted public benchmarks: FNO, PDEBench, and CFDBench. They are also the datasets used for pre-training by DPOT. Using these standard benchmarks ensures that our results are fair, reproducible, and comparable with existing and future research. \n- **The Core Challenge is \"****Heterogeneity****,\" Not \"Single-Task Difficulty\":** The central thesis of our paper is to solve the performance bottleneck that arises when a single model must handle a *heterogeneous mixture* of PDEs. As we explicitly demonstrate in **Figure 2 (left)**, even on these individually \"less difficult\" datasets, mixing them during training causes the performance of a traditional dense model to collapse. Therefore, the advantage of our architecture is its ability to effectively manage data diversity and heterogeneity, a property that is best demonstrated through the diverse dataset combination we selected.\n\n### **On Fine-Tuning Epochs and the Overall Compute Budget (Weakness 4)**\n\nThis is a very insightful and fair point. We offer two clarifications:\n\n- **Relative Cost of Fine-tuning:** Our pre-training phase runs for 1000 epochs, making the 200-epoch fine-tuning a reasonable fraction of the total training. More importantly, the value of fine-tuning lies in its ability to rapidly adapt general knowledge to specific tasks. The results in **Table 2** clearly show that pre-trained models (\"w/ Pre-train\") significantly outperform models trained from scratch (\"w/o Pre-train\") on all downstream tasks, proving the immense value and knowledge transfer capability of pre-training.\n- **Model** **Capacity** **vs. Training Duration:** The idea of aligning the total compute budget is an excellent one. However, we argue that smaller models like FFNO, even if granted equivalent training time, would quickly hit a performance ceiling due to their limited model capacity, potentially even degrading from overfitting. Indeed, all baseline results in our paper represent the best performance achieved during their training until full convergence. In contrast, our MoE architecture's success stems from having a very large total parameter count (ensuring high capacity) while keeping computational costs low via sparse activation. Our scaling experiments (Sec 5.3) further support the conclusion that larger model capacity leads to better performance.\n\n### **Clarification of \"Feature Competition\" and \"Knowledge Illusion\" (Question)**\n\nThank you for pointing out this imprecise phrasing. These terms were intended as intuitive descriptors for the well-established phenomenon known as **\"****Negative Transfer****\"** in multitask learning.\n\n- **\"Feature Competition\"** refers to the scenario where a dense model, forced to learn fundamentally different tasks (e.g., fluid dynamics of NS equations vs. reaction-diffusion of DR equations) with a single set of shared parameters, experiences conflicting gradient updates. This \"tug-of-war\" between tasks prevents the model from fully learning any single task, thereby harming overall performance.\n- **\"Knowledge Illusion\"** describes the outcome of this phenomenon: the model appears to have learned something about all tasks, but this knowledge is shallow and compromised, failing to capture the core physics of each task and leading to poor generalization.\n\nWe acknowledge that we did not provide direct citations for these descriptive terms. In the final version of the paper, we will replace them with the more standard academic terms **\"****negative transfer****\"** or **\"inter-task interference\"** and cite foundational multitask learning literature, such as [1], to make our argument more rigorous.\n\n**References:** \n\n[1] Multitask learning, Caruana, 1997\n\n--- \n\n#### **Thanks again** \n\nWe sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly."}, {"review_id": "riEvhNv4vV", "reviewer": "Reviewer_fgJS", "summary": "The paper introduces a mixture of experts module at the end of a SOTA neural operator (DPOT) to deal with two challenges of building large SciML foundation models: (a) performance degradation due to heterogeneity of PDE datasets and (b) too many parameters necessary for dense models to alleviate the previous challenge. The authors pre-train the DPOT with MOE on 6 PDE datasets and demonstrate zero-shot and fine-tuned performance on several downstream tasks. They show that the MOE addition uniformly helps increase the accuracy of DPOT. Further, they show that their MOE activations correspond to learning the underlying PDE task distribution, motivating its creation. Some ablations on number of experts demonstrate that a balance is being struck w.r.t accuracy vs speed. Finally, the authors show the practical utility of their model by comparing inference times - the setting where these models will be used in production.", "questions": "Clarity:\n* The results are unclear in some places. It would be useful to demonstrate clearly what is the train/val/test split for each dataset used, which datasets are used for pre-training vs zero-shot vs finetuning, size of the datasets for each case (pre-training/finetuning), degree of difficulty for the downstream tasks w.r.t out-of-distribution generalization (this can be as simple as quantifying the effects of the physics coefficients on the difficulty of the tasks -  some tasks such as diffusion would be easier to generalize rather than more hyperbolic systems and hence the error metrics are not sufficient to gauge the improvements). \n* Tab1. would be more beneficial as some sort of figure (error vs fine-tuning dataset size) - it's then clear what value the MOE adds in addition to the dataset volume. \n* Fig 2 (right) would be beneficial to see it as one figure (with different color bars) so it's easier to compare the MOE activations\n\nExperiments: \n* The error over DPOT seems very comparable. How many finetuning examples are taken? How much care is given to optimization when finetuning to a new dataset (learning rates, batch size, frozen weights)? Will DPOT eventually catch-up to MOE with more examples and at what point (how many examples)? Finally, I understand that dropping the error from 0.02 to 0.01 is a 100% improvement, but how valuable is that if your error is already low? It seems like DPOT does get low errors already in most cases. I think it would be beneficial to stress-test on more out-of-distribution examples.\n* I appreciate the interpretability experiment. I feel like the authors could do more with this line of analysis for this paper. For example, when you look at classification accuracy, is it just for a training sample X or validation/test? How does the accuracy change as you pick X from a downstream task and systematically make that downstream task out-of-distribution? Do the MOE logits still point to some pre-training dataset or does it become more uniform?\n* Related to above, the authors mention that part of the MOE is to learn common representations (conservation laws) and the other part is to learn the PDE task distribution? Is there any evidence to suggest the former? Such as high activation for the shared experts for all datasets? Similarly, if you picked X (input) from an out-of-distribution example, would you still see the shared experts being activated? That would be a nice result.\n* Another related thread: if part of the MOE is just to distinguish the PDE, why not just one-hot encode that information at the beginning (or have the MOE as well at the beginning). It's understandable in fields like language modeling that the source of the tokens is pretty unclear but the science, data is often from structured simulation+experiments and hence it should be possible to encode that information directly?\n* The inference time experiment is not convincing. DPOT-M and MOE-M have roughly the same inference time. Tab 1 only compares the accuracy of these two models (as the best models), so I'm wondering what the advantage of MOE is? Also what is the accuracy of DPOT-L in Tab. 1?\n* How does the MOE block impact rollout accuracy? This probably is more important than one-step prediction accuracy for spatiotemporal models. It would be good to discuss the multistep performance as well.\n* [Minor] What is the performance of the FNO with the MOE? The FNO numbers seem disentangled from the whole experiment and seem to only emphasize that DPOT is a better architecture choice?\n\nMinor:\n* Fig labels and tick sizes could be made larger so it's clear\n* Tab 1. The caption says error greater than 20. I did not follow what that means\n* What is n in W_p (after Eqn (3))? embedding dimension?", "limitations": "There is a minimal discussion on limitations.\n* It seems like a limitation that the MOE is evaluated on simple and small pre-training volumes (6 datasets from PDEBench). I would expect that evaluations on larger volumes such as The Well (https://polymathic-ai.org/the_well/datasets_overview/) would show up more interesting ways the MOE could benefit the training or show some instabilities with the large data distributions. Some discussion on this would be useful.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: \n* The paper is well-motivated, the topic is of high interest to the NeurIPS audience\n* Performance of the MOE model is better than DPOT (a SOTA pre-training architecture for SciML)\n* Experiments demonstrate that the experts learn the PDE data in-distribution structure by looking at the MoE activation distributions\n\nWeaknesses:\n* The paper's clarity could be improved (see comments). While the motivation is clear, the results section could benefit greatly with more clarifications\n* The improvement from DPOT seems marginal and limited to one-step improvement. It's unclear how much this translates to rollout improvements\n* There is very little discussion on out-of-distribution generalization properties of the MOE model (or others)\n* The interpretability analysis (while appreciated) could be made more extensive to stress-test the motivation behind the MOE (see comments)", "was_revised": true, "final_justification": "The paper presents a novel contributed to PDE foundation models using MoE. The performance increase is marginal but is still an interesting analysis to the community. There are some clarity revisions that the authors have committed too along with new results from the discussion phase from all reviews. Overall, I recommend a borderline accept", "rebuttal": "Dear Reviewer,\n\nThank you for your detailed and insightful review. We are very grateful for your positive feedback on our paper's motivation and the novelty of our approach. Your comments are exceptionally helpful, and we agree that they point to clear ways we can improve the paper. We have carefully considered all your questions and provide the following responses, including new experimental results, which we hope will address your concerns.\n\n---\n\n### **On Clarity and Experimental Details**\n\nWe completely agree that presenting experimental details with utmost clarity is crucial. Thank you for these excellent suggestions. We have uploaded our code for full transparency.\n\n- **Experimental Setup & Data Splits:** As stated on Page 6 (Line 186), our datasets are from three widely-used public benchmarks: FNO, PDEBench, and CFDBench. We used the standard, publicly available train/validation/test splits and sample sizes for all benchmarks to ensure fair and reproducible comparisons, following the standard procedures of prior works like DPOT and FNO.\n- **Commitment to Improvement:** We recognize that these details could be presented more clearly. In the final version, we will:\n  - **Add a comprehensive summary table** in the appendix, detailing the train/val/test splits, sample counts, fine-tuning procedures, and data sources for all pre-training and downstream datasets. We will also include quantitative metrics of task difficulty (e.g., Reynolds number for NS equations) and discuss their implications for generalization.\n  - **Revise Figure 2 (right)** into a single, unified plot with grouped bar charts. This will allow for a direct and easy comparison of expert activation patterns across different datasets, as you wisely suggested.\n  - **Improve Table 1's presentation** by adding a summary figure in the main text that plots performance against the number of fine-tuning samples for MoE-POT vs. DPOT, visually highlighting the value added by our MoE architecture. The full table will be retained in the appendix.\n  - **Enlarger all figure labels and ticks** to improve readability.\n\n### **On Performance Improvement and Multi-Step Rollout Accuracy**\n\nYou have raised a crucial point about the significance of our performance gains and the fairness of our comparisons. We would like to offer a detailed clarification.\n\n- **Reported Errors are for Multi-Step Rollouts:** We would like to clarify a critical point regarding our evaluation: the reported errors in our tables are **not for single-step prediction but are the average errors over long-term, auto-regressive rollouts.** Our experimental setup follows the standard protocol used in DPOT and FNO. For instance, on the PDEBench-SWE dataset, we input the first 10 frames and auto-regressively predict the subsequent 90 frames, with the reported error being the average over these 90 predicted steps.\n- **The Value of \"Marginal\" Error Reduction:** For spatio-temporal dynamics, even a small reduction in single-step prediction error is vital.\n  - **Significant Relative Improvement:** As shown in Tables 1 and 2, our model achieves a substantial *relative* performance improvement on multiple tasks. In the field of neural operators, where current SOTA precision (1e-2 to 1e-4) still lags behind traditional solvers (<1e-7), every bit of error reduction is a meaningful step forward.\n  - **Importance of Compounded Error:** A significant reduction in single-step error (e.g., 50% relative improvement) drastically suppresses the accumulation of errors during long-term rollout predictions. This leads to far more stable and accurate simulations in practical applications, directly addressing your concern about multi-step performance.\n- **Fair Comparison with DPOT:**\n  - **Fine-tuning Details:** As described on Page 7 (Lines 221-222) and Page 8 (Line 240), our fine-tuning protocol was applied consistently across all pre-trained models (DPOT and ours) to ensure a fair comparison. For Table 1, models were fine-tuned for 200 epochs; for Table 2, 500 epochs. The training set of the corresponding dataset was used for fine-tuning.\n  - **Convergence** **and Architectural Advantage:** In all fine-tuning experiments, both MoE-POT and DPOT models reached convergence (typically before 180 epochs for Table 1 tasks and before 400 epochs for Table 2 tasks). Therefore, there was no possibility of DPOT \"catching up\" with more training. We argue the performance gap stems from a fundamental architectural advantage. As a dense model, DPOT's capacity is limited by its activated parameters. MoE-POT, through sparse activation, leverages a much larger total parameter count, giving it a higher capacity to absorb heterogeneous knowledge more effectively.\n\n### **On Inference Efficiency and DPOT-L Comparison**\n\n- **Advantage in Inference Efficiency:** You correctly observe that MoE-POT-M and DPOT-M have similar inference times. This is precisely where our model's advantage lies. As shown in Table 3, MoE-POT-M (288M activated, 489M total params) has nearly the same inference time (16.6ms) as DPOT-M (158M activated, 16.7ms). However, as shown in Table 1, MoE-POT-M is far more accurate for this same inference cost. This demonstrates a significantly better efficiency-to-performance ratio and highlights the scalability advantage of the MoE architecture.\n- **DPOT-L Experimental Results:** Due to computational constraints, we were unable to complete the full pre-training for a corresponding MoE-POT-L model. To address your question, we provide the zero-shot results for the official DPOT-L model below.\n\n| Model and Activation Parameter | NS(1e-5) | NS(1e-3) | CNS(0.1,0.01) | SWE    | DR     | CFDBench |\n| --- | --- | --- | --- | --- | --- | --- |\n| DPOT-S 31M      | 0.0078   | 0.0688   | 0.0244        | 0.0039 | 0.0367 | 0.0087   |\n| DPOT-M 122M      | 0.0071   | 0.0569   | 0.0224        | 0.0025 | 0.0288 | 0.0113   |\n| DPOT-L 493M       | 0.0061   | 0.0576   | 0.0113        | 0.0023 | 0.0219 | 0.0065   |\n| MoE-POT-T 17M    | 0.0077   | 0.0682   | 0.0105        | 0.0064 | 0.0411 | 0.0053   |\n| MoE-POT-S 90M      | 0.0058   | 0.0552   | 0.0096        | 0.0029 | 0.0342 | 0.0045   |\n| MoE-POT-M 288M      | 0.0057   | 0.0528   | 0.0091        | 0.003  | 0.03   | 0.0051   |\n\nThese results show that **MoE-POT-M achieves comparable or superior performance to DPOT-L on 4 out of 6 datasets, using only about half the activated parameters (and thus, computational cost).** This strongly supports our claims. We will add a more detailed comparison with DPOT-L in the final version.\n\n### **On Interpretability and the MoE Mechanism**\n\nThank you for appreciating the interpretability analysis and for pushing us to explore it further. These are excellent suggestions.\n\n- **Role of Shared Experts:** The two shared experts are, by design, activated for every input. In our initial experiments, models without shared experts often ended up training unstable. \n- **Stress-Testing with Out-of-Distribution Data:** In our original analysis (Appendix B.3), we used the training set to compute the mean expert preference vector for each dataset and the test set to calculate classification accuracy. To address your excellent question about OOD data, we conducted a **new experiment** testing the classification accuracy on two challenging downstream tasks from the Poseidon paper [1] that were **not seen during pre-training**.\n\n| Dataset (Unseen) | Block-1 Accuracy | Block-2 Accuracy |\n| ---------------- | ---------------- | ---------------- |\n| Wave-Layer       | 100%             | 100%             |\n| Wave-Gauss       | 100%             | 100%             |\n\nThe router achieves **100% classification accuracy** on these unseen datasets. Furthermore, in Block-2, the expert network ratios for the Wave-Layer dataset were: Expert 1 (79%), 9 (31%), 11 (100%), 13 (100%), 15 (100%), others (0%). This strongly demonstrates that the gating network learns to distinguish fundamental properties of different PDEs and generalizes this capability to new, unseen tasks.\n\n### **Why Not One-Hot Encode?:**\n\n- Our work follows the experimental paradigm and dataset of DPOT and FNO, which aims to simulate real-world scenarios (e.g., weather forecasting) where only solution trajectories are available, and the exact PDE parameters are unknown. The goal is for the model to implicitly infer the system's dynamics from observed data, which is precisely the role our router network successfully plays. While parameter-informed modeling via encoding is a valid and important research direction, our core contribution is the effective introduction of the MoE structure into this auto-regressive pre-training paradigm.\n\n### **On Minor Questions**\n\n- **FNO with MoE:** During our research, we explored numerous combinations, including MoE with FNO and FFNO. These attempts were unsuccessful, leading to unstable training and large errors that could not be resolved by fine-tuning. The current MoE-POT architecture was the final stable and effective design that emerged from this rigorous process.\n- **Table 1 Caption:** Thank you for pointing this out. The \"-\" indicates that the model's L2RE on that task exceeded 20, signifying that training diverged or failed completely. We will clarify this in the caption.\n- **Wp in Eqn. (3):** You are correct. 'n' is the feature dimension of the positional encoding, which we typically set to be the same as the number of channels C. We will explicitly define this in the final version.\n\n**References:** \n\n[1] Poseidon: Efficient foundation models for pdes, NeurIPS 2024.\n\n---\n\n#### **Thanks again** \n\nWe sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly."}, {"review_id": "4ig0XvLNZd", "reviewer": "Reviewer_frLV", "summary": "Context: PDEs are very important as mathematical models in physics and engineering and given the slow speed of physics-based PDE solvers, the design of neural PDE surrogates (also termed Operator Learning) is very important. Most of the research in operator learning as been in the task specific context. However, it is well known that the number of training examples need per task can be too high to be practically useful. Hence, the design of Foundation models has been pursued recently as data across PDE operators might be leveraged to  reduce sample complexity per task through finetuning. Some recently developed PDE Foundation models are MPP, DPOT and Poseidon (which the authors strangely omit from their literature review). However, there is certainly enormous scope for further development of Foundation models for PDEs. \n\nGoal: The article proposes a PDE Foundation Model, called MoE-POT, aiming to reduce model size and improve scaling. \n\nMethod: The basic goal is operator learning for time-dependent PDEs. The task is \"given the PDE solution at a set of time steps, predict the solution at the next time step\". To do so, MoE-POT reuses most elements of DPOT -- patchification, temporal aggregation, noise addition, Fourier layers etc. The main new element is a mixture-of-experts (MoE) Module. At each processor layer, the outputs of the Fourier layer are fed into a set of shared and routing experts. The shared experts are always active and the routing experts are activated through a gating network. The expert and gating network are CNNs. During training, a standard load balance loss is added to promote balance between the experts and prevent routing collapse. \n\nResults: The pretraining dataset are 6 operators (Incompressible Navier-Stokes with 2 different viscosities from the original FNO paper, A compressible Navier-Stokes, Shallow-water and a Reaction-Diffusion dataset from PDEbench and another forced Incompressible Navier-Stokes example from CFD bench. All are 2-D fairly academic problems. Most of the evaluation is on the Pretraining dataset itself, either in the form of zero-shot test loss or with additional fine-tuning. 3 Downstream tasks are also considered: Incompressible Navier-Stokes with a different viscosity from the FNO paper, Compressible Navier-Stokes with different parameters from PDEbench and a Navier-Stokes dataset from PDEarena. MoEPOT is shown to be either competitive or better than baselines (small neural operators and DPOT). Additional interpretability is shown through the importance values of gating weights.", "questions": "1. Fig 2 (Left): Why would the performance of FNO deteriorate so sharply on a mixed dataset ? If so, why does DPOT work in the same setting as it only time aggregates and patches, when compared to FNO. Anyway, there is too little information about this figure in the caption or the text. \n\n2. L79-81: Where exactly is the noise injected ? Is it uniform over time steps ? \n\n3. L90-91: See relation to Q1.\n\n4. L96: Explain feature competition and knowledge illusion ? or else drop jargon !\n\n5. L102: Long Inference time is a misnomer: DPOT-L takes 25 milliseconds to infer compared to 17 milliseconds for MoEPOT-L. This has be compared to the min-to-hour time scale for PDE solvers. Inference is not an argument for MoEs in this setting. \n\n6. Sec 3.2 should be moved ? It does not fit here ?\n\n7. L108: Explain parallel expert networks and gated routing, else drop jargon till you explain it in the next section. \n\n8. L136: Should the spatial dimensions of Z^t_p be HxW ? Isn't patchification supposed to reduce spatial size ? \n\n9. L135: The notation W_p(x_i,y_j,t) is unclear, I am assume that W_p is just a fixed matrix or is it a MLP ? \n\n10. L138: What is W_t ? Is it a MLP ? \n\n11. L139: What are Fourier features \\gamma -- are they learnable or not ? \n\n12. L141: Why do you call Eq (7) as a Fourier Attention layer ? It is just  Fourier layer from FNO and no attention mechanism is used here.\n\n13. L151: Is the group and concatanation over channels or over spatial points, it is unclear from Eq (7) as the weight matrices have dimensions of the space with patches ? Is h the \"patch size\" ? \n\n14. L156-157: Why choose CNNs here -- shouldn't FNO be a natural choice ? \n\n15. L171: Is the Batching over samples only or over spatial points -- I ask as you use x (spatial coordinate) in Eqn (11) below. \n\n16. L172: What is importance in Eq (11) and what is e_i ? Aren't they the same ? \n\n17. L175: The phrase \"one step transition between samples from different datasets\" is misleading here -- you just want to predict the next state, given a sequence of states. \n\n18. L187: The claim on \"wide range of PDEs\" is disputable as you only consider fluid equations with 1 Reaction-Diffusion equation -- PDEs have a much wider variety. \n\n19. L194: L2RE but for what time ? Is it average one-step prediction error ? or final time evaluation error ? \n\n20. L201-202: The claim \"SOTA\" pre-trained neural operator (DPOT) can be contested, see arXiv:2405.19101 Sec D.5. \n\n21. Table 1: What does 0.1 and 0.01 in CNS stand for ? \n\n22. L221: Why do you need to fine-tune your pretrained model on the same tasks as in the pretraining dataset. Is it taking the model checkpoint and then task-specific training ? What is the fine-tuning procedure ? \n\n23. L238-239: \"represents entirely different equations\" is simply wrong - it is just forced incompressible Navier-Stokes for which you already have 2 datasets in the pretraining data. \n\n25. L254: Claim on superior performance is untrue as DPOT has a sharper gradient in Fig. 4 (b). \n\n26. Fig 4 a and b: please provide averages over all the tasks or provide all the tasks in the Supplement. \n\n27. L262 and L263: Y and I are introduced here out of thin air -- what are they ? At least provide link to Supplement \n\n28. L543 -- what is T ?\n\n29. L546: What are the constants -- see Q 21 ?\n\n30. L550: Define what u, D and R are ?\n\n31. L551-552: What are the initial conditions ? \n\n32. L554: What are the geometric shapes ? How do you carry geometric information ? \n\n33. Please state how you normalize different spatial and temporal domains in your datasets ?\n\n34. How do you set the number of input time steps $t$ in your operator learning task ? Is it the same for all the datasets. \n\n35. To reach a certain time from the initial data (and the initial time step sequence), do you use an autoregressive rollout ? If so, how does the error accumulate over time ? Please present a plot to show this. \n\n36. Fig 4 (c): Please provide average classification accuracy over all blocks ? There is presumably nothing special about block 2. \n\n37. It would be great if Fig 4 (c) can be shown for each of the downstream tasks, by comparing with the pretraining dataset -- it make sense that the cross-entropy loss for FNO is close to the FNO datasets, Cross-entropy loss for CNS is close to the CNS pretraining datasets and the one for PDEarena is also close to the FNO dataset. The exact values of the cross-entropy losses should be compared with the ones obtained for Fig 4 (c) and I am speculating that the results will be consistent with my intuition that the downstream tasks are not really that out-of-distribution with respect to the pretraining dataset -- this will be a very interesting test of the authors' method.\n\n38. Although Fig 4 (c) shows that the expert networks are able to separate the pretraining datasets -- the authors do not explain why this is the case ? After all, the experts have only access to processed tokens and not to the datasets directly and there is nothing in either the model or in the loss function so that the experts can distinguish between different datasets ? Is this an emergent property ? If so, can the authors indicate if this arises after a certain model size, certain number of training epochs or certain size of the dataset ? \n\n39. What is the size of the total pretraining dataset in terms of the number of training samples ? \n\n40. What is the size of the dataset for each downstream task ?", "limitations": "Limited Novelty: See Weaknesses.\n\nLimited scope of operator learning task: Next time-step prediction from a sequence, unclear how to use parameters explicitly, unclear how to directly predict from initial data, unclear how to do steady state predictions and large time predictions. Error accumulation is possible due to long-time autoregressive rollouts. Restricted to 2D, mostly cartesian. \n\nLimited Evaluation: particularly at the level of downstream tasks, unclear if there is any gain with respect to DPOT on genuine downstream tasks. No comparison to Poseidon and MPP. \n\nEvaluation: Although very natural in this context, given its widespread use in LLMs, the authors are the first to use a standard MoE module in a PDE foundation model. However, they have not demonstrated that it leads to any benefits, particularly on downstream tasks involving unseen physics. hence, this reviewer cannot recommend acceptance in current form but I am willing to modify my score if the authors can rebut the criticism, particularly if they provide evidence for their model directly predicting solution trajectories from initial data, at least at the level of downstream tasks, they test on unseen downstream tasks such as the Wave-Layer and Wave-Gauss experiments of arXiv:2405.19101 (datasets are publicly available) and benchmark performance over DPOT and Poseidon, provide data scaling plots to show that MoEPOT is competitive in the low task-specific sample regime  and answer the minor questions above satisfactorily.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths: \n\n--- The problem being addressed is of great importance in Scientific machine learning. \n--- First direct application of MoE in a PDE Foundation Model setting\n--- MoEPOT outperforms DPOT at least marginally.\n--- It is always an engineering challenge to train PDE foundation models and the authors have done so.\n--- The classification of datasets using the routing networks (Fig 4(c)) is a very interesting observation.  \n\nWeaknesses: \n\n--- Limited Novelty: The only novel model contribution is the use of MoE in this setting as all the other choices in architecture and dataset selection are based on DPOT. Nevertheless, MoEs have been widely used in Foundation models for text and vision. \nAlso, MoE have been widely used in scientific machine learning:\nRafael Bischof and Michael A Kraus. Mixture-of-experts-ensemble meta-learning for physics-informed neural networks. In Proceedings of 33. Forum Bauinformatik, 2022, \nZheyuan Hu, Ameya D Jagtap, George Em Karniadakis, and Kenji Kawaguchi. Augmented\n360 physics-informed neural networks (apinns): A gating network-based soft domain decomposition, Engineering Applications of Artificial Intelligence, 126:107183, 2023, \nand see in particular arXiv:2404.09101 for even a theoretical analysis of MoEs for operator learning. \nGiven this context and the absence of theory, the reviewer has to judge this paper based on the empirical results. \n\n--- Incomplete literature review and context: The authors do not cite or contextualize let alone compare with Poseidon (arXiv:2405.19101) published in Neurips last year, which in the opinion of the reviewer, is the current state-of-the-art PDE foundation model (it is compared and shown to be superior to DPOT in Sec D.5 of arXiv:2405.19101 at least on a wide range of downstream tasks. \n\n--- Lack of Clarity: This paper is badly written. The narrative does not flow smoothly. For instance, Sec 3.2 comes in abruptly with a lot of jargon that is totally unhelpful for the reader. Just look at the number of questions that this reviewer poses below to see why the paper lacks clarity. \n\n---- Incorrect Operator learning task: Given a time-dependent PDE such as Eq. [1], the solution operator of a PDE is very well defined as mapping from the initial conditions $u^0$, PDE parameters $\\theta$ and boundary conditions and to the solution $u(t)$ of (1) for any time t >0. This operator is well-posed and is exactly the mapping that is learned by a numerical Physics-based PDE solver. \n\nOn the other hand, the authors misinterpret this learning task to: Given the trajectory $\\{ u(\\tau) \\}$ for $0 \\leq \\tau \\leq t$, map this input into the PDE solution $u(t+dt)$ for some lead time $dt$. This is different from the PDE solution operator and it is over-determined and may not be mathematically correct for many PDEs. From a practical perspective, unless $t=0$, it will require the end user to first run a numerical solution of the underlying PDE (access to knowledge about all the PDE parameters) for some time $t$ and then use your neural operator to update the solution. However, such a solver may not be available for inference. Moreover, the upfront cost of setting up such a solver if available (for instance in mesh generation) can be very significant, vastly limiting utility of the PDE surrogate.\n\nAlthough, the authors are not the first to setup this operator learning task, this incorrect or limited utility task has crept into the SciML community from surrogates for weather forecasting, where it is totally justified as weather forecasting is an online streaming task whereas solving PDEs is rarely so. \n\nIn particular, it does not really matter what the authors do at Pretraining as they have all the trajectory data but for downstream tasks, you have to run your model such that it can learn from the initial data to the solution trajectory. See arXiv:2405.19101 OLT on Page 3 for more context. \n\n---- Very Limited Evaluation: As common in text/vision, evaluating on pretraining dataset is not of much importance for a Foundation model, except perhaps for model and dataset size scaling. In fact, (V)-LLMs do not even have the same pretraining objective as that in the downstream tasks (for instance masked modeling). The whole point of a foundation model is how it performs on downstream tasks and NOT on the pretraining dataset as an end-user will have very different inference tasks (PDE types or IC/BC) than the Foundation model designers. This issue was carefully analyzed in arXiv:2405.19101 and its authors clearly compared model performance on downstream tasks in terms of two reasonable metrics i) what is the gain in terms of how few samples are needed to finetune a Foundation model on a downstream task to obtain the same error as that of a task-specific baseline with a given number of samples (see E_G in Eqn (11) of arXiv:2405.19101). This tells the user how few samples that they can get away with to obtain reasonable errors and ii) for a given number of downstream task-specific examples, the gain in accuracy over the baseline model (A_G in Eqn (11) of arXiv:2405.19101). \n\nThus, all the evaluation of a Foundation Model should be at the level of downstream tasks. This paper does exactly the opposite -- bulk of its results are on the pretraining dataset (in-distribution) and there is very limited evaluation on the downstream tasks. Thus, the whole point of the model and its purported gains in performance are very incomplete. \n\n--- Poor Choice of Downstream tasks: Given the above point, any Foundation Model has to be evaluated on a challenging set of downstream tasks. The authors make a very poor selection of such tasks. The first downstream task is just the same NS example from the original FNO paper but with a different choice of viscosity, which conveniently lies between the viscosity coefficients considered in the pretraining dataset. The second downstream is just CNS with a different (10 times larger) coefficient from the CNS dataset considered in the pretraining data. Again, this is very related and even simpler than the one similar task in the pretraining data. The final downstream task is again Incompressible forced Navier-Stokes (already seen twice in the pretraining dataset) with perhaps different initial conditions. This reviewer can buy this example as a genuine downstream task but here, MoEPOT has exactly the same error as DPOT (6.21 % vs 6.18 %). None of the downstream tasks represent physics unseen during the pretraining, just contrast this with the Poseidon evaluation where 9 of the 15 downstream tasks were for unseen physics. After all, no Foundation model can be pretrained on all the physics in the universe and needs to be tested on unseen physics. \n\n--- Questionable choice of Pretraining dataset:  Out of the 6 operators in the pretraining dataset, 4 are very simplistic as even small models (less than 2M params) are able to attain less than 1% relative error. Only the CNS and FNO task with very low viscosity are relatively challenging. It is unlikely that a Foundation model trained on such simple operators can be expected to generalize well to complicated physics. \n\n--- Underwhelming results: Given the above critique on zero-shot test losses, even here, on the 4 easy datasets, very small models are as good as MoEPOT and on the difficult ones, the gains wrt DPOT are marginal (on FNO low-viscosity: MoEPOT-M 5.28% vs. DPOT-M 5.69% and on CNS: MoE-POT-M 1% approx vs. DPOT 2% approx) are hardly worth writing about. The authors claim that their model scales better than DPOT but this is not true on the evidence of Fig. 4(b) where DPOT has a sharper gradient than MoEPOT and is already catching up for moderate model sizes. \n\n--- Lack of comparison with Poseidon: It would certainly be worthwhile to compare your model with DPOT and Poseidon on a neutral downstream task, something not seen in either of the pretraining datasets -- a very good choice for that would be the wave-layer and wave-gauss datasets of arXiv:2405.19101, which are publicly available and where both Poseidon and DPOT have already been tested. \n\n--- Lack of data scaling: The whole point of a Foundation model is to require less task-specific examples downstream. Please provide such results for MoEPOT -- see arXiv:2405.19101 Figs. 7-21 for how to perform these scaling tasks. \n\n--- Lack of visualization: The authors do not provide any examples of how exactly their model performs in terms of sample visualization.", "was_revised": true, "final_justification": "I appreciate the effort to train a PDE foundation model and the integration of the MoE approach in this context. On the negative side, the MoE approach brings in marginal though consistent gains on downstream tasks over DPOT. Moreover, the camera-ready version need a significant revision, where the authors need to provide the correct context, admit their limitations, compare with Poseidon and convince the reader that their approach leads to a gain in downstream tasks.\n\nOn balance, this paper is a solid contribution to this emerging field and my score will reflect this assessment.", "rebuttal": "Dear Reviewer,\n\nThank you for your incredibly thorough and insightful review. We are grateful for your feedback and for acknowledging our work's motivation and novelty. Your critiques regarding clarity, evaluation, and literature context are invaluable and provide a clear roadmap for improving our paper. We sincerely apologize for these shortcomings and are committed to a full revision. Below are our responses and new experimental results.\n\n---\n\n### On Literature Review, Novelty, and Comparison with Poseidon \n\nOur initial choice of DPOT as a baseline was threefold: to ensure clear attribution of gains from our MoE module; the divergent experimental setups (our auto-regressive, \"trajectory-to-trajectory\" paradigm differs significantly from Poseidon's excellent parameter-informed approach); and the fact that Poseidon was contemporaneous work released after our experiments began.\n\n- We agree that testing on challenging, unseen tasks is essential. To address this, we conducted new experiments on the Wave-Layer and Wave-Gauss datasets you proposed.\n\n|                          | Poseidon's Setting |            |               |    Our Setting        |            |\n| ------------------------ | ------------------ | ---------- | ------------------------ | ---------- | ---------- |\n| Model (Activated Params) | Wave-Layer         | Wave-Gauss | Model (Activated Params) | Wave-Layer | Wave-Gauss |\n| Poseidon-T  21M          | 0.08               | **0.06**   | Poseidon-T  21M          | 0.29       | 0.29       |\n| Poseidon-B  158M         | **0.06**           | 0.09       | Poseidon-B  158M         | 0.21       | 0.24       |\n| MoE-POT-T  17M           | 0.11               | 0.14       | MoE-POT-T  17M           | 0.07       | 0.07       |\n| MoE-POT-S  90M           | 0.06               | 0.10       | MoE-POT-S  90M           | **0.05**   | **0.06**   |\n\n\n\n- These results compellingly show that each model excels in its intended paradigm. This highlights that they are both powerful tools with different, valuable scopes of application.\n\n- Commitment: We commit to including a thorough discussion of Poseidon and other related MoE works in our literature review. We will also incorporate these new experimental results to provide a more complete and honest picture of our model's performance and positioning within the field.\n\n### On Task Definition, Evaluation, and Clarity  (Q4, Q6, Q7, Q17, Q18, Q23)\n\nYour critiques on our operator learning task definition, limited evaluation, and writing clarity are spot on. We sincerely apologize for these shortcomings and are committed to a full revision, including Section 3, to ensure clarity and logical flow.\n\n- Task Definition & Motivation: You are correct that our \"trajectory-to-trajectory\" task differs from the classic formulation. Our setup, following DPOT, is motivated by real-world scenarios (e.g., weather forecasting) where explicit PDE parameters are unknown and must be inferred from observed data. While parameter-informed approaches like Poseidon are powerful, our core contribution is introducing the MoE architecture to this specific auto-regressive paradigm.\n- Dataset Choice & Heterogeneity: We chose our datasets for fair comparison with DPOT. Our central argument is that the key challenge is not single-task difficulty but handling the heterogeneity of mixed datasets. This \"knowledge conflict\" is where dense models fail, as shown in Figure 2, validating the need for our architecture. Even for tasks with similar equations like PDEArena, our router distinguishes them with 100% accuracy, proving it captures significant distributional differences beyond just the equation type.\n\n### On Dense Model Challenges & MoE's Rationale (Q1, Q3)\n\nThe performance collapse of dense models like FNO on mixed datasets is due to \"negative transfer,\" where a single shared kernel struggles to learn conflicting physical laws. To validate this, we tested a dense DPOT-S model as more heterogeneous datasets were added. The results below show that adding more data often degrades DPOT's performance, while our MoE-POT remains superior. This confirms that a sparse architecture is crucial for handling knowledge conflict in diverse PDE pre-training.\n\n|               | NS(1e-5) | NS(1e-3) | CNS(0.1,0.01) | SWE    | DR     | CFDBench |\n| ------------- | -------- | -------- | ------------- | ------ | ------ | -------- |\n| DPOT-S (6 datastes)    | 0.0078   | 0.0688   | 0.0244        | 0.0039 | 0.0367 | 0.0087   |\n| DPOT-S (12 datastes)   | 0.0079   | 0.0739   | 0.0129        | 0.0105 | 0.0724 | 0.0075   |\n| MoE-POT-S (6 datastes) | 0.0058   | 0.0552   | 0.0096        | 0.0029 | 0.0342 | 0.0045   |\n\n### On Interpretability with Unseen  Datasets (Q36)\n\nTo stress-test our router's generalization, we tested its ability to classify data from two challenging, unseen downstream tasks (Wave-Layer and Wave-Gauss).\n\n| Unseen Downstream Task | Classification Accuracy |\n| ---------------------- | ----------------------- |\n| Wave-Layer             | 100%                    |\n| Wave-Gauss             | 100%                    |\n\nFurthermore, the router exhibited distinct and highly confident selection patterns for these new tasks. For example, in Block-2, Wave-Layer data triggered experts 1 (79%), 9 (31%), 11 (100%), 13 (100%), and 15 (100%). The router not only classified them perfectly but also exhibited distinct and confident expert activation patterns. This proves the router learns fundamental physical dynamics rather than simply memorizing the training set, and that the experts have learned specialized knowledge.\n\n### On the Emergence of Classification Accuracy during Training (Q37)\n\nTo show how the router's classification ability arises, we tracked its accuracy during training. The results confirm this is an emergent property: accuracy starts low but steadily improves, eventually reaching 100%. This aligns with the intuition that the model learns to separate data distributions to optimize overall performance.\n\n| Epoch | NS(1e-5) Acc. | CFDBench Acc. |\n| ----- | ------------- | ------------- |\n| 50    | 2%            | 70%           |\n| 150   | 80%           | 78%           |\n| 250   | 100%          | 100%          |\n\n### Responses to Specific Questions\n\n#### On Experimental Setup\n\n- Q1: FNO's performance collapses due to \"negative transfer\" from its shared kernel learning conflicting PDE patterns. DPOT mitigates this but doesn't fully solve it, as supported by our new experiments showing its performance wavers when more datasets are added.\n- Q2: Noise is injected into all input frames, with variance proportional to the input's norm and applied uniformly across time.\n- Q32: We followed DPOT's preprocessing, normalizing spatial and temporal domains; details will be added to the appendix.\n- Q33: A fixed input length of T=10 was used for all experiments, consistent with the source datasets.\n- Q34: Yes, we use auto-regressive rollouts. The reported L2RE is the average error over the entire sequence. We will add a requested error accumulation plot comparing MoE-POT and DPOT to the final paper.\n- Q39, Q40: The pre-training dataset has ~30,000 samples; each downstream task has 2,000. We will provide a detailed table.\n\n#### On Model & Methods\n\n- Q8: You are correct; the output dimension after patchification should be $(H/P)\\times(W/P)\\times d$. We will correct this typo.\n- Q9-11: $W_t$ is a learnable MLP for positional encoding, $W_t$ is a learnable linear transformation, and $\\gamma$ are fixed, non-learnable Fourier features.\n- Q12: You are correct; \"Fourier Attention\" is a misnomer. We will correct this to \"Fourier Layer\" throughout the paper.\n- Q13: The grouping is performed on the channel dimension, where h is the number of heads.\n- Q14: We chose lightweight CNNs for the experts as they effectively preserve spatial structure and proved more stable than alternatives like FNO in our early experiments.\n- Q15: Batching is over samples, so `x` in Eq. (11) represents a single sample in the batch B. We will correct this confusing notation.\n- Q16: `e_i` and `Importance(e_i)` are not the same. `e_i` is the i-th expert network, while `Importance(e_i)` is a scalar value—the total routing weight for that expert in a batch—used for the load-balancing loss.\n\n#### On Results & Inference\n\n- Q5: The per-step millisecond difference aggregates to a significant advantage in long rollouts or optimization loops. Our point is MoE's higher capacity and accuracy for a comparable computational cost.\n- Q19, Q34: L2RE is the average error over the full auto-regressive rollout, not a single step.\n- Q21, Q28: The values (0.1, 0.01) are the physical coefficients $\\zeta$ and $\\eta$ in the CNS equations; we will define all such constants in the appendix.\n- Q22: During fine-tuning, we freeze the router and update only the expert networks to preserve the routing strategy while adapting expert knowledge.\n- Q24: We agree the slope in Fig. 4(b) can be misleading. Our claim is about efficiency: better performance for a given number of activated parameters. We will rephrase this for clarity.\n- Q25, Q26: We will add supplementary figures for all tasks and clearly define Y (mean expert preference) and I (input's expert preference) in the appendix.\n- Q27, Q29-31: All requested details on T, u, D, R, initial conditions, and geometry will be added to the appendix.\n\n---\n\n#### Thanks again \n\nWe sincerely thank you again for your constructive and detailed feedback, which has helped us identify clear pathways to improve our paper. Should you have any further questions or require additional discussion, please don't hesitate to reach out. If we have adequately addressed your concerns, we would be grateful for your consideration in adjusting your evaluation score accordingly."}]

Author rebuttal:

Meta-review: This manuscript proposes a new architecture that combines the concept of MoE with neural operators. It is to address the challenge of building foundational models capable of fitting multiple PDEs simultaneously. Whereas a naive approach would substantially increase computation, this work employs two types of expert networks to mitigate that overhead. Empirical evaluation demonstrates that the proposed architecture outperforms baselines and operates as expected.

The AC recommends: Accept as a poster.

**Strengths:** The main strength of the paper is being the first to employ the MoE concept within neural operators in the context of foundational models. All reviewers highlighted this point, and the AC concurs. There was extensive discussion during the review process regarding clarity and additional experiments, and the reviewers were ultimately satisfied with the rebuttals. Although a minor concern remains about the relatively modest performance improvement over prior work, the AC believes that the novelty lies in the architectural design and conceptual contribution rather than the raw performance gains.

**Final concern.** During the discussion, the authors provided substantial additional experimental results:

1. Comparison of computational overheads to DPOT
2. Performance analysis across different numbers of datasets
3. Model and activation parameter analysis
4. Stress-testing experiments

If incorporated directly into the main body of the camera-ready, these additions would go far beyond minor changes. The camera-ready version should not constitute a substantial revision. Therefore, the AC recommends that the authors restrict changes to minimal edits (e.g., typos, clarifications) in the main text, while placing additional experimental results in the Appendix.
Decision: Accept (poster)