NESTOR: A Nested MOE-based Neural Operator for Large-Scale PDE Pre-Training

Dengdi Sun; Xiaoya Zhou; Xiao Wang; Hao Si; Wanli Lyu; Jin Tang; Bin Luo

cs.CVFeb 25, 2026

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

Abstract

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- $\nu$			PDEBench CNS-( $\eta,\zeta$ ), 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. $\mathcal{L}_{2}$
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. $\mathcal{L}_{2}$
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. $\mathcal{L}_{2}$	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- $k_{1}$	Routed 2	Shared 2	Top- $k_{2}$	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- $\nu$			PDEBench CNS-( $\eta,\zeta$ ), 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- $\nu$		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.

Neural Operators

Neural operators learn mappings between functions, enabling fast surrogates for PDEs and other continuous physical systems.

See all 310 papers

SPEAR: A Spectral-Disentangled MoE Neural Operator with Knowledge-Guided Expert Aggregation for Large-Scale PDE Pretraining
Oct 2, 2026Dengdi Sun, Xiaoya Zhou, Xiao Wang +3Neural Operators Mixture-of-Experts Models
School of Artificial Intelligence, Anhui University, Hefei 230601, China · School of Computer Science and Technology, Anhui University, Hefei 230601, China
AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training
May 15, 2026Qitan Lv, Hong Wang, Zhongkai Hao +5Transformer PDE Operator Learning
University of Science and Technology of China · 2Shanghai AI Laboratory · Department of Electronic Engineering, Tsinghua University, China
MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries
Aug 10, 2026Zijiang Yang, Xiaomeng Wu, Dongmei FuNeural PDE Solvers PDE Operator Learning
School of Automation and Electrical Engineering, University of Science and Technology Beijing · 2Beijing Engineering Research Center of Industrial Spectrum Imaging

Oct 2, 2026J/K move · Enter open · S save

SPEAR: A Spectral-Disentangled MoE Neural Operator with Knowledge-Guided Expert Aggregation for Large-Scale PDE Pretraining

Dengdi Sun, Xiaoya Zhou, Xiao Wang, Wanli Lyu+2

School of Artificial Intelligence, Anhui University, Hefei 230601, China · School of Computer Science and Technology, Anhui University, Hefei 230601, China

Neural Operators Mixture-of-Experts Models Mixture-of-Experts Pruning

Open paper AskPDF arXiv

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.

Abstract

Neural Operators

Explore similar work

SPEAR: A Spectral-Disentangled MoE Neural Operator with Knowledge-Guided Expert Aggregation for Large-Scale PDE Pretraining

AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training

MoNo: Multiscale Optimal Transport Neural Operator for Solving PDEs on General Geometries