Large-scale PDE pre-training trains a single model on the data from multiple partial differential equations (PDEs). However, adding heterogeneous PDEs with different governing equations, physical variables, and variable interactions does not always improve predictive performance across tasks. Existing PDE mixture-of-experts models mitigate cross-task interference through learned routing, but typically use homogeneous expert architectures, making it difficult to directly reflect differences in state evolution across equations. We introduce HotPOT, an operator transformer with heterogeneous physics experts for large-scale PDE pre-training. HotPOT uses a shared encoder to capture information common across PDEs, while each PDE family uses a neural-numerical expert that pairs a learnable predictor with a fixed numerical update containing no trainable parameters. Each predictor outputs the quantities required by its paired numerical update, which then computes the next state. This design allows different PDEs to share a common encoder while retaining their own learnable predictors and numerical updates. Experiments across five PDE families and six settings show that HotPOT outperforms MoE-POT in all 24 matched-scale comparisons before and after fine-tuning, reducing rollout error by 50.7% on average. Notably, the 17M model also surpasses substantially larger MoE-POT models, and the performance gains extend to zero-shot prediction under unseen physical conditions. Code is available at https://github.com/cruiseresearchgroup/HotPOT-code.
Figures & tables
Figure 1: HotPOT: framework overview and an NS prediction example. (a) Shared history encoding connects PDE-specific predictor–update pairs for pre-training and transfer. The state is qn at step n , with history length T=10 . (b) Pre-trained 17M HotPOT and MoE-POT on the same NS test trajectory ( ν=10−5 ). The boxed region is compared at forecast step 10; gray dashed and red solid curves denote reference and predicted contours, respectively. The relative L2 error in this region decreases from 0.165 for MoE-POT to 0.057 for HotPOT (65.4% lower). Here ν is viscosity, M is Mach number, and L is domain length.
Figure 2: Architecture of HotPOT. (a) Prediction pipeline with sampling, padding, and training-time noise. (b) A shared encoder connects PDE-specific neural–numerical experts in one model. (c) Temporal aggregation and five alternative numerical updates. Blue, orange, green, and gray indicate shared computation, system-specific prediction, fixed numerical updates, and physical states, respectively. Stage-state queries, history statistics, and repeated readouts are omitted. FFN denotes feed-forward network. Field images are illustrative. Notation and module details are given in Appendices A.1 and A.2 .
Model
Activated Params
NS benchmarks
PDEBench
CFDBench
NS ( 10−5 )
NS ( 10−3 )
CNS (0.1,0.01)
SWE
DR
Small models
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
FFNO
1.2M
0.161
0.0256
0.183
0.00458
0.161
0.0990
GK-T
1.1M
0.260
0.0148
0.919
0.0453
0.0120
0.419
Table 1: Full-horizon mean L2RE on six PDE settings ( ↓ ). CNS lists Mach number M and shear viscosity η , with bulk viscosity ζ=η . Bold: best within each training block; blue: best overall. HotPOT uses final E1000/E200 checkpoints. Baselines follow MoE-POT Table 1 ( Wang et al., 2025 ) , with matched splits, histories, horizons, resolutions, and metrics. Dashes denote errors above 20.
Model
NS 10−5
NS 10−3
CNS
SWE
DR
CFDB
Mean
Full HotPOT
0.0315
0.00325
0.00671
0.00177
0.0150
0.00219
0.0101
Learned private history
0.0322
0.00420
0.00790
0.000962
0.0149
0.00222
0.0104
Late-shared evolution
0.0331
0.00413
0.00941
0.00421
0.119
0.00217
0.0286
Table 2: Cross-PDE sharing ablations. Full-horizon L2RE; bold marks minima before rounding.
Model
NS 10−5
NS 10−3
CNS
SWE
DR
CFDB
Mean
Full HotPOT
0.0315
0.00325
0.00671
0.00177
0.0150
0.00219
0.0101
Direct-state predictor
0.0525
0.00625
0.0104
0.00190
0.122
0.00211
0.0324
Generic iteration
0.0320
0.00404
0.00720
0.00177
0.0150
0.00234
0.0104
Table 3: Numerical-expert ablations. Full-horizon L2RE; bold marks minima before rounding.
Table 5: Within-family adaptation: model-feedback L2RE on 200 test trajectories per target (NS/CNS: 20/11 steps). Full-rollout validation selects checkpoints during 500-epoch adaptation.
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Component
HotPOT-T (17M active)
HotPOT-S (90M active)
Shared width / depth / FFN width
256 / 2 / 512
512 / 3 / 2048
System-specific width / depth / FFN width
512 / 5 / 2048
1024 / 7 / 3968
Decoded context channels
32
64
Typed-head hidden width
128
256
Task-feature channels / hidden width
16 / 64
32 / 128
Task-embedding hidden width
128
128
Appendix
Table 6: Backbone and system-specific predictor dimensions. The neural template and numerical programs are common to both scales.
Setting
Native grid
Train
Test
Cpred,r
Lroll
Fine-tune LR
NS 10−5
642
1,000
200
1
10
10−5
NS 10−3
642
1,000
200
1
20
10−3
CNS
1282
9,000
200
4
11
10−4
SWE
1282
900
60
1
91
10−3
DR
1282
900
60
2
91
10−3
CFDB
642
9,000
1,000
2
10
10−5
Appendix
Table 7: Main evaluation populations and fine-tuning schedule for the pre-training tasks at both model scales. Train is the trajectory count used per epoch; test is the fixed evaluated row count. Cpred,r=∣Or∣ counts scored physical channels, and Lroll counts model-feedback prediction steps after ten input frames. The learning rate is the maximum during E200 fine-tuning.
Quantity
HotPOT-T
HotPOT-S
All-PDE shared history encoder
2,011,648
11,901,952
Active NS/SWE path
16,770,577
89,584,673
Active DR path
16,903,188
90,112,036
Active CFDB geometry path
16,905,234
90,116,130
Active CNS path
17,178,004
91,185,956
Stored total
97,902,720
518,480,368
Appendix
Table 8: Exact parameter accounting at both scales. Storage includes all system-specific expert paths; an active path includes the shared encoder and the selected system or CFDB geometry path.
Model
Activated
Stored total
MoE-POT-T
17
30
HotPOT-T
16.77–17.18
97.90
MoE-POT-S
90
166
HotPOT-S
89.58–91.19
518.48
MoE-POT-M
288
489
Appendix
Table 9: Activated and stored parameters, in millions. MoE-POT counts are the published nominal values ( Wang et al., 2025 , Table 5) ; HotPOT ranges cover its active routes.
Model
Placement
Weights
Stored (M)
Mean L2RE
Full
History
Shared
97.90
0.0101
Learned private history
History
Private
105.95
0.0104
Late-shared evolution
Evolution
Shared
97.90
0.0286
Late-private evolution
Evolution
Private
105.95
0.0287
Appendix
Table 10: History/evolution placement and sharing controls. The mean equally weights six settings; all results use fixed E1000 checkpoints.
Figure 3: Space–time errors and two-dimensional cuts for the sharing-location comparison. Rows show SWE height and both DR fields for fixed example 0. The first two columns render absolute errors in two spatial dimensions and forecast time. The last two columns show HotPOT absolute error minus Late-shared evolution absolute error: an x – y slice at step 91 and an x –time slice at zero-based y=64 . Blue denotes lower error for HotPOT; orange denotes lower error for Late-shared evolution. Within each field, the error volumes share the camera, full-range color mapping, and opacity mapping; opacity increases with absolute error. The two cuts share a symmetric range determined by the full space–time error difference. Volume colorbars describe scalar values before compositing, whereas the cuts display local scalar differences directly.
Figure 4: Mean per-step errors across the complete test populations. Each curve averages per-trajectory, scored-channel relative L2 error at each forecast step. All models use fixed E1000 checkpoints and their own autoregressive feedback. CFDB uses the same row-weighted population as the main table. Vertical axes are logarithmic.
Figure 5: NS numerical execution and generic iteration. The first four test examples at ν=10−3 use matched reference trajectories. The first two columns show final-step absolute errors; the third shows their difference. The fourth is an x –time cut at the central y index through the difference volume. Blue favors Full, orange favors Generic iteration. Error scales are shared across all four examples; difference scales are symmetric and use the complete forecast volumes. The task has two spatial dimensions; the rightmost vertical axis is time.
Target
Model
16
32
64
128
Wave-Layer
MoE-POT full fine-tuning
0.576
0.517
0.435
0.365
HotPOT scratch
0.553
0.428
0.331
0.235
HotPOT encoder transfer
0.506
0.395
0.304
0.220
Wave-Gauss
MoE-POT full fine-tuning
0.514
0.452
0.394
0.325
HotPOT scratch
0.542
0.372
0.284
0.225
HotPOT encoder transfer
0.453
0.370
0.279
0.218
Appendix
Table 11: Transfer to one unseen wave-equation family. Mean per-step L2RE on 200 custom held-out trajectories, averaged over 11/6 forecast frames for Layer/Gauss. All methods update their complete target models and select checkpoints by full-rollout validation within 500 epochs.
System
Diagnostic
Full
Direct
NS 10−5
Relative enstrophy error
1.78e-3
3.63\times10−3
NS 10−3
Relative enstrophy error
2.16e-3
2.51\times10−3
CNS
Normalized mass drift
7.63e-8
2.04\times10−4
Normalized energy drift
8.15e-8
2.06\times10−4
SWE
Relative mass drift
1.19e-9
9.31\times10−5
Appendix
Table 12: Physical diagnostics over the complete test rollouts of fixed E1000 Full and Direct models. Values average all trajectories and forecast steps; lower is better.
Figure 6: Numerical execution along the fixed E1000 HotPOT rollouts. (a) Mean absolute spatial-mean drift at each program stage, averaged over all recorded trajectory–step pairs for SWE height, CNS density, and CNS energy. (b,c) Per-step drift for SWE height and CNS density. Drift is the absolute change in spatial mean from the current state to the candidate next state, in released field units. Lines show arithmetic means and bands the 25th–75th percentiles across trajectories. All stages use the same current state on the model’s own feedback trajectory.
Figure 7: Predictions and absolute errors: NS, CNS, and SWE. All scored fields are shown at the final forecast step (NS/SWE: row 0; CNS: row 1). Reference and prediction fields share color limits within each row; error maps share a zero lower bound and the full-rollout maximum. Signed fields use zero-centered scales.
Figure 8: Predictions and absolute errors: DR and CFDB. All scored fields appear at the last forecast step (DR/cavity: row 0; cylinder: row 110). Within each row, reference and prediction fields share color limits; error maps share a zero lower bound and the full-rollout maximum.
Figure 9: Temporal errors for the seven illustrated trajectories. Both E1000 models use autoregressive feedback from the same ten-frame histories. Each unsmoothed curve averages spatial L2RE over scored channels.
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
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
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.
Dengdi Sun, Xiaoya Zhou, Xiao Wang +4
School of Artificial Intelligence, Anhui University, Hefei, China · School of Computer Science and Technology, Anhui University, Hefei, China