AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training
Abstract
Pre-training neural operators on diverse PDE datasets has emerged as a promising paradigm 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 difficult. Existing approaches address this mainly by enlarging model capacity, while the target solution operators themselves remain unchanged. Inspired by classical numerical analysis---where a problem-tailored operator transformation reformulates a complex PDE into a numerically easier equivalent---we propose to similarly reformulate the target solution operators, turning a heterogeneous set of complex, structurally divergent operators into a simpler and better-aligned equivalents. Since different PDEs admit different simplifications, the transformation must be adaptive and input-dependent, so that a single neural operator can approximate the entire family jointly. We instantiate this idea as AOT-POT (Adaptive Operator-Transformation for Pre-training Operator Transformer), which realizes such a transformation by expanding the hidden representation into multiple parallel streams, aggregating and redistributing them with input-dependent weights before and after each sub-layer, and mixing streams through Sinkhorn-projected doubly stochastic matrices for stable training. Together, these mechanisms reformulate the diverse, complex solution operators into a simpler equivalents that a single architecture can approximate jointly. Empirically, AOT-POT achieves state-of-the-art results on 12 PDE benchmarks with only 3% additional parameters, reducing the relative L2 error by up to 77.6% (40.9% on average). Fine-tuning further reduces L2 error by up to 92% on in-domain PDEs and 89% on out-of-domain PDEs, confirming that adaptive operator transformation is an orthogonal and effective axis for advancing PDE foundation models, beyond merely scaling model capacity.