NESTOR: A Nested MOE-based Neural Operator for Large-Scale PDE Pre-Training
Authors: Dengdi Sun, Xiaoya Zhou, Xiao Wang, Hao Si, Wanli Lyu, Jin Tang, Bin Luo
Organizations: School of Artificial Intelligence, Anhui University, Hefei, China · School of Computer Science and Technology, Anhui University, Hefei, China
Neural operators have emerged as an efficient paradigm for solving PDEs, overcoming the limitations of traditional numerical methods and significantly improving computational efficiency. However, due to the diversity and complexity of PDE systems, existing neural operators typically rely on a single network architecture, which limits their capacity to fully capture heterogeneous features and complex system dependencies. This constraint poses a bottleneck for large-scale PDE pre-training based on neural operators. To address these challenges, we propose a large-scale PDE pre-trained neural operator based on a nested Mixture-of-Experts (MoE) framework. In particular, the image-level MoE is designed to capture global dependencies, while the token-level Sub-MoE focuses on local dependencies. Our model can selectively activate the most suitable expert networks for a given input, thereby enhancing generalization and transferability. We conduct large-scale pre-training on twelve PDE datasets from diverse sources and successfully transfer the model to downstream tasks. Extensive experiments demonstrate the effectiveness of our approach.
Figures & tables
Figure 1: Comparison of two different network architectures. (a) Traditional single-network architecture; (b) our proposed nested MoE architecture, where image-level MoE experts learn global diversity across different PDE types, while token-level Sub-MoE experts capture complex local features within equations.
Figure 2: Overview architecture. We train on twelve mixed PDE datasets, predicting the next frame based on the preceding frames. We design a nested MoE architecture: (1) the top shows the overall model architecture; (2) the bottom right illustrates the nested Sub-MoE architecture; and (3) the bottom left depicts the improved Flash Attention architecture.
L2RE
Activated
FNO- ν
PDEBench CNS-( η,ζ ), DR, SWE
PDEArena
CFDBench
Model
Params
1e-5
1e-4
1e-3
1,0.1
1,0.01
Avg.(1)
0.1,0.1
0.1,0.01
Avg.(0.1)
DR
SWE
NS
NS-cond
-
Pre-trained
FNO
0.5M
0.116
0.0922
0.0156
0.151
0.108
0.130
0.230
0.076
0.153
0.0321
0.0091
0.210
0.384
0.0274
UNet
25M
0.198
0.119
0.0245
0.334
0.291
0.313
0.569
0.357
0.463
0.0971
0.0521
0.102
0.337
0.209
FFNO
1.3M
0.121
0.0503
0.0099
0.0212
0.052
0.0366
0.162
0.0452
0.104
0.0571
0.0116
0.0839
0.602
0.0071
GK-T
1.6M
0.134
0.0792
0.0098
0.0341
0.0377
0.0359
0.0274
0.0366
0.0320
0.0359
0.0069
0.0952
0.423
0.0105
GNOT
1.8M
0.157
0.0443
0.0125
0.0325
0.0420
0.0373
0.0228
0.0341
0.0285
0.0311
0.0068
0.172
0.325
0.0088
Table 1: The experiments are divided into two parts: one reports the pre-training performance of the model, and the other shows the fine-tuning results on each task. Here, “-200” denotes fine-tuning for 200 epochs, and “-500” for 500 epochs. The evaluation metric is L2RE. The best result within each part is highlighted in bold , while the overall best result is emphasized in blue bold.
Figure 3: Visualization of 2D high-resolution turbulence prediction results. (1) The first column shows the true values, the second shows the model predictions, and the third shows the corresponding errors. (2) The predicted physical quantities are horizontal velocity, vertical velocity, density field, and pressure field.
Figure 4: Performance comparison of different models on the 2D high-resolution turbulence task. We use L2RE as the evaluation metric, where Vanilla denotes training from scratch, and -FT indicates results after 500 fine-tuning epochs on the downstream task.
Setting
Num. of experts
FNO
PDEBench
SWE
Avg. L2
FT-200
2
0.0575
0.0182
0.0024
0.0262
4
0.0577
0.0240
0.0579
0.0466
6
0.0563
0.0150
0.0025
0.0246
12
0.0575
0.1896
0.0025
0.0832
FT-500
2
0.0519
0.0126
0.0022
0.0222
4
0.0504
0.0114
0.0025
0.0214
Table 2: The impact of the number of non-shared experts. We use L2RE as the evaluation metric.
Num. of datasets
FNO
PDEBench
SWE
Avg. L2
3
0.0512
0.0165
0.0026
0.0234
12
0.0505
0.0094
0.0026
0.0208
Table 3: The impact of the size of the pre-training data on performance. We use L2RE as the evaluation metric.
Method
1,0.1
1,0.01
0.1,0.1
0.1,0.01
DR
SWE
Avg. L2
Promotion
Ours
0.0144
0.0355
0.0135
0.0178
0.0282
0.0045
0.0173
-
w/o Sub-MoE
0.0157
0.0393
0.0130
0.0209
0.0245
0.0049
0.0197
0.0024
w/o Load Balance Loss
0.0135
0.0335
0.0109
0.0159
0.0265
0.0062
0.0178
0.0005
FlashAttn + AFNO Sum
0.0149
0.0363
0.0136
0.0178
0.0304
0.0046
0.0196
0.0023
Table 4: Ablation experiments of our proposed model on the PDEBench datasets. “w/o” denotes the removal of the corresponding component. We use L2RE as the evaluation metric.
Figure 5: Visualization of spatial activation patterns produced by token-level experts in the Sub-MoE layer. For each input sample, the expert activation probabilities are projected onto heatmaps.
Dataset
Expert 0
Expert 1
Expert 2
Expert 3
Expert 4
Expert 5
M1(-1,-1)
22.39
50.00
6.98
5.37
10.69
4.58
M-1(-1,-1)
22.88
50.00
3.51
11.58
10.22
1.81
SWE
0.00
0.00
50.00
49.74
0.00
0.25
DR
28.77
0.00
2.31
12.64
6.28
50.00
Table 5: Expert selection distribution of the MoE router across different datasets. The top two experts for each dataset are highlighted in bold (%).
DPOT-T
MoE-POT-T
Ours
Activated Parameters
7.5M
17M
13M
Total Parameters
7.5M
30M
83M
Activation Ratio
100%
56.67%
16.67%
Table 6: Comparison of activation parameters, total parameters, and activation rates of different models.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Dim
Ratio
Layers
Heads
Routed 1
Shared 1
Top- k1
Routed 2
Shared 2
Top- k2
Model Size
Activated Size
512
1
2
4
1
6
2
1
6
2
83M
13M
Appendix
Table 7: Setting of the Attention Module.
FNO- ν
PDEBench CNS-( η,ζ ), DR, SWE
PDEArena
CFDBench
1e-5
1e-4
1e-3
1,0.1
1,0.01
0.1,0.1
0.1,0.01
DR
SWE
NS
NS-cond
-
Train set size
100
9800
1000
9000
9000
9000
9000
900
900
6500
3100
9000
Test set size
200
200
200
1000
1000
1000
1000
100
100
1300
600
1000
Appendix
Table 8: Train and test set sizes of the PDE datasets used for pre-training.
L2RE
Activated
FNO- ν
PDEBench
CFDBench
Model
Params
1e-5
1e-3
0.1,0.01
SWE
DR
-
MoE-POT
17M
0.0682
0.00768
0.0105
0.00640
0.0411
0.00529
Ours
13M
0.0674
0.00763
0.0159
0.00449
0.0184
0.00911
Appendix
Table 9: Comparison with MoE-POT in pre-training across six datasets. The evaluation metric is L2RE. The best result within each part is highlighted in bold .
Figure 6: FNO series of result visualizations. (1) The first column shows the true value, the second column shows the model prediction value, and the third column shows the corresponding error. (2) Each row is the predicted physical quantity.
Figure 7: PDEBench series of result visualizations. (1) The first column shows the true value, the second column shows the model prediction value, and the third column shows the corresponding error. (2) Each row is the predicted physical quantity.
Figure 8: DR series of result visualizations. (1) The first column shows the true value, the second column shows the model prediction value, and the third column shows the corresponding error. (2) Each row is the predicted physical quantity.
Figure 9: SWE series of result visualizations. (1) The first column shows the true value, the second column shows the model prediction value, and the third column shows the corresponding error. (2) Each row is the predicted physical quantity.
Figure 10: PDEArena series of result visualizations. (1) The first column shows the true value, the second column shows the model prediction value, and the third column shows the corresponding error. (2) Each row is the predicted physical quantity.
Figure 11: CFDBench series of result visualizations. (1) The first column shows the true value, the second column shows the model prediction value, and the third column shows the corresponding error. (2) Each row is the predicted physical quantity.
Large-scale pre-training has improved the generalization of neural operators across diverse PDEs. However, existing PDE foundation models still struggle with heterogeneous dynamics, where shared representations may cause knowledge interference, while mixture-of-experts (MoE) architectures suffer from increasing expert redundancy. We propose SPEAR, a spectral-disentangled MoE neural operator with knowledge-guided expert aggregation for large-scale PDE pre-training. SPEAR decouples latent features into low- and high-frequency components, enabling shared modeling of transferable dynamics and specialized learning of PDE-specific patterns. To address expert redundancy, we design a knowledge-guided expert aggregation strategy that measures expert similarity from dataset-specific learned knowledge and routing preferences, enabling the identification and consolidation of similar experts. Experiments on twelve PDE datasets and multiple downstream benchmarks demonstrate superior performance in pre-training, fine-tuning, and transfer learning. Furthermore, our aggregation strategy reduces the number of experts by 50% while maintaining or improving prediction accuracy, achieving a balance between model efficiency and generalization for PDE foundation models.
Dengdi Sun, Xiaoya Zhou, Xiao Wang +3
School of Artificial Intelligence, Anhui University, Hefei 230601, China · School of Computer Science and Technology, Anhui University, Hefei 230601, China
Pre-training neural operators on diverse partial differential equation (PDE) datasets has emerged as a promising direction for building general-purpose surrogate models in scientific machine learning. However, the inherent complexity and structural diversity of PDE solution operators make multi-PDE pre-training fundamentally challenging. Existing methods mainly address this by increasing model capacity, while leaving the target solution operators unchanged. Inspired by classical numerical analysis, we instead propose to transform complex and diverse solution operators into simpler, better-aligned forms that are easier to model jointly. Since the optimal transformation varies across PDE types, it must be adaptive and input-dependent, allowing a single neural operator to approximate an entire family of operators. We instantiate this idea as AOT-POT (adaptive operator-transformation for pre-training operator transformer), which expands hidden representations into multiple parallel streams, adaptively aggregates and redistributes them before and after each sub-layer, and mixes streams through Sinkhorn-projected doubly stochastic matrices for stable training. These mechanisms together reshape diverse solution operators into a unified form that can be effectively modeled by a single architecture. Empirically, AOT-POT achieves state-of-the-art performance on 12 PDE benchmarks with only 3% additional parameters, reducing relative L2 error by up to 77.6% (40.9% on average). Fine-tuning AOT-POT further reduces L2 error by up to 92% on in-domain PDEs and 89% on out-of-domain PDEs (unseen types during pre-training), demonstrating that adaptive operator transformation is an effective and complementary direction for advancing PDE foundation models beyond simply scaling model capacity.
Qitan Lv, Hong Wang, Zhongkai Hao +5
University of Science and Technology of China · 2Shanghai AI Laboratory · Department of Electronic Engineering, Tsinghua University, China
Transformer-based neural operators have achieved substantial progress in solving Partial Differential Equations (PDEs) by projecting spatial observations into compact latent tokens and learning physical interactions in latent spaces. However, we reveal that existing learnable projection mechanisms cannot ensure stable and balanced assignments from observation points to latent tokens, causing some latent tokens to be over-assigned while others remain underutilized. This limitation further restricts the design of hierarchical architectures, as assignment imbalance is continuously inherited and amplified across latent spaces, eventually causing severe token collapse in deeper spaces. To address these issues, we propose MoNo (Multiscale Optimal Transport Neural Operator), a progressive multiscale neural operator that efficiently solves PDEs on general geometries through stable latent-space construction. At its core is CoTAP (Cross-scale Optimal Transport Assignment and Projection), a novel latent-space construction method that formulates cross-space assignment between adjacent spaces as an entropy-regularized optimal transport problem, thereby constructing balanced bidirectional projections and stable latent spaces. CoTAP also ensures stable information transfer across multiple latent spaces, further enabling multiscale architectures on general geometries, which in turn support more efficient learning of long-range physical interactions. Extensive experiments demonstrate that MoNo outperforms existing state-of-the-art neural operators in both prediction performance and computational efficiency. Code is available at https://github.com/ZijiangY1116/MoNo.
Zijiang Yang, Xiaomeng Wu, Dongmei Fu
School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging