HotPOT: Operator Transformer with Heterogeneous Physics Experts for Large-Scale PDE Pre-training

Jiayi Li; Guancheng Zhou; Yisi Luo; Ruihua Chen; Penghao Jiang; Deyu Meng; Flora D. Salim; Piotr Koniusz

cs.LGOct 28, 2026

HotPOT: Operator Transformer with Heterogeneous Physics Experts for Large-Scale PDE Pre-training

Authors: Jiayi Li, Guancheng Zhou, Yisi Luo, Ruihua Chen, Penghao Jiang, Deyu Meng, Flora D. Salim, Piotr Koniusz

Organizations: University of New South Wales · Xi’an Jiaotong University · Commonwealth Scientific and Industrial Research Organisation

Abstract

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 $\mathbf{q}^{n}$ at step $n$ , with history length $T=10$ . (b) Pre-trained 17M HotPOT and MoE-POT on the same NS test trajectory ( $\nu=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 $L_{2}$ error in this region decreases from 0.165 for MoE-POT to 0.057 for HotPOT (65.4% lower). Here $\nu$ 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
Model	Activated Params	NS ( $10^{-5}$ )	NS ( $10^{-3}$ )	CNS $(0.1,0.01)$	SWE	DR	CFDBench
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 (

\downarrow

). CNS lists Mach number

M

and shear viscosity

\eta

, with bulk viscosity

\zeta=\eta

. 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.

	CNS	Cavity			Cylinder
Model	$(M,\eta)=(1,10^{-2})$	BC	Geo †	Prop	BC	Geo †	Prop
HotPOT-T	0.189	0.00568	0.0169	0.00467	0.00119	0.00180	0.000577
MoE-POT-T	$0.423$	$0.0276$	$0.0439$	$0.0116$	$0.00802$	$0.0117$	$0.00223$

Table 4: Zero-shot H10 L2RE (

\downarrow

) using fixed 17M pre-training checkpoints without target adaptation. CFDB uses case-uniform masked aggregation.

	NS $\nu=10^{-4}$			CNS $(M,\eta)=(1,10^{-2})$
Model / shots	16	64	256	16	64	256
HotPOT-T	0.124	0.0964	0.0819	0.0983	0.0599	0.0365
MoE-POT-T	$0.304$	$0.206$	$0.108$	$0.365$	$0.153$	$0.0735$

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	$C_{\mathrm{pred},r}$	$L_{\mathrm{roll}}$	Fine-tune LR
NS $10^{-5}$	$64^{2}$	1,000	200	1	10	$10^{-5}$
NS $10^{-3}$	$64^{2}$	1,000	200	1	20	$10^{-3}$
CNS	$128^{2}$	9,000	200	4	11	$10^{-4}$
SWE	$128^{2}$	900	60	1	91	$10^{-3}$
DR	$128^{2}$	900	60	2	91	$10^{-3}$
CFDB	$64^{2}$	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.

C_{\mathrm{pred},r}=|\mathcal{O}_{r}|

counts scored physical channels, and

L_{\mathrm{roll}}

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

L_{2}

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 $\nu=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\text{\times}{10}^{-3}$
NS $10^{-3}$	Relative enstrophy error	2.16e-3	$2.51\text{\times}{10}^{-3}$
CNS	Normalized mass drift	7.63e-8	$2.04\text{\times}{10}^{-4}$
	Normalized energy drift	8.15e-8	$2.06\text{\times}{10}^{-4}$
SWE	Relative mass drift	1.19e-9	$9.31\text{\times}{10}^{-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.

Panel of Figure 9 — 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.

PDE Foundation Models

See all 12 papers

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
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
NESTOR: A Nested MOE-based Neural Operator for Large-Scale PDE Pre-Training
Feb 25, 2026Dengdi Sun, Xiaoya Zhou, Xiao Wang +4Neural Operators PDE Surrogate Modeling
School of Artificial Intelligence, Anhui University, Hefei, China · School of Computer Science and Technology, Anhui University, Hefei, China

May 15, 2026J/K move · Enter open · S save

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

Qitan Lv, Hong Wang, Zhongkai Hao, Wen Wu+4

University of Science and Technology of China · 2Shanghai AI Laboratory · Department of Electronic Engineering, Tsinghua University, China

Transformer PDE Operator Learning Neural Operators

Open paper AskPDF arXiv

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.

Abstract

PDE Foundation Models

Explore similar work

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

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

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