🤖 machine learning

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

AOT-POT introduces an adaptive operator transformation framework that reshapes diverse PDE solution operators into a unified form through parallel stream aggregation and Sinkhorn-projected mixing, achieving state-of-the-art performance across 12 benchmarks with minimal parameter overhead and demonstrating that transforming operators is a more effective strategy for PDE foundation models than simply scaling model capacity.

Original authors: Qitan Lv, Hong Wang, Zhongkai Hao, Wen Wu, Xuenan Xu, Bowen Zhou, Feng Wu, Chao Zhang

Published 2026-05-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Qitan Lv, Hong Wang, Zhongkai Hao, Wen Wu, Xuenan Xu, Bowen Zhou, Feng Wu, Chao Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: One Size Doesn't Fit All

Imagine you are trying to teach a single robot to solve every type of math problem in the universe: predicting how water flows, how heat spreads, how air moves around a plane, and how chemicals react.

In the world of science, these are called Partial Differential Equations (PDEs). Each one is a different "language" with its own rules.

The Old Way: Scientists tried to build bigger and bigger robots (neural networks) hoping that if the robot was just huge enough, it could memorize the rules for every single language. This is like trying to learn every language in the world by just memorizing a dictionary with a million pages. It works, but it's slow, expensive, and the robot still gets confused when switching between languages.
The Paper's Insight: The authors realized that instead of making the robot bigger, we should give it a universal translator. If we can translate every complex math problem into a simpler, common "dialect" that the robot already knows how to speak, the robot can solve them all much easier.

The Solution: AOT-POT (The Adaptive Translator)

The paper introduces a new architecture called AOT-POT. Think of it as a smart, adaptive translator built right into the robot's brain.

Here is how it works, step-by-step:

1. The "Parallel Streams" (The Multitasking Team)

Imagine the robot's brain isn't just one long hallway of neurons. Instead, AOT-POT splits the information into four parallel streams (like four different teams of workers).

Team A might be good at seeing the big picture.
Team B might be good at spotting fine details.
Team C and Team D have other special skills.
This gives the robot a "super-basis" to look at the problem from multiple angles at once.

2. The "Adaptive Transformation" (The Chameleon Filter)

This is the magic part. Before the robot tries to solve a problem, it looks at the input (e.g., "Is this a wave? Is this a fluid?").

Based on what it sees, it dynamically mixes the four teams.
If the problem is a Wave, it might say, "Team A and Team B, you do 80% of the work; Team C and D, you just listen."
If the problem is Heat, it might say, "Team C and D, take the lead now."
The Analogy: Imagine a chef who has four different knives. When cutting a tomato, they use the sharp serrated knife. When chopping onions, they switch to the heavy cleaver. They don't use the same knife for everything. AOT-POT does this automatically for math problems. It reshapes the problem into a form that is easiest for the robot to solve right now.

3. The "Sinkhorn" Safety Net (The Traffic Cop)

When you have four teams constantly swapping jobs and mixing information, things can get chaotic. The math could blow up, and the robot could get confused or crash.

The paper uses a mathematical trick called a Sinkhorn projection.
The Analogy: Think of this as a strict traffic cop at a busy intersection. No matter how much the teams want to rush or change lanes, the cop ensures that the total amount of "traffic" (information) stays balanced. It guarantees that the robot doesn't accidentally amplify a tiny error into a huge disaster. This keeps the training stable and safe.

Why This is a Big Deal (The Results)

The authors tested this new "translator" on 12 different types of physics problems (from weather forecasting to fluid dynamics).

Efficiency: They didn't make the robot bigger. In fact, they only added 3% more parameters (tiny bits of memory).
Performance: Despite being almost the same size as the old models, AOT-POT was significantly better.
- It reduced errors by up to 77.6% on average.
- When they fine-tuned it for a specific job, it reduced errors by up to 92%.
Stability: When the robot had to predict the future for a long time (like predicting the weather for 500 steps ahead), the old models would eventually spiral out of control and become nonsense. AOT-POT stayed accurate and stable for much longer.

The "Aha!" Moment

The paper proves that the best way to build a "Foundation Model" for science isn't just to throw more computing power at the problem (scaling up). Instead, it's to build a smarter architecture that can adapt its own internal structure to fit the specific problem it's facing.

In summary: AOT-POT is like giving a general-purpose robot a set of magic glasses. When it looks at a fluid problem, the glasses turn the fluid into a simple line drawing. When it looks at a wave, the glasses turn the wave into a simple rhythm. The robot doesn't need to be a genius; it just needs the right glasses to see the problem clearly. And the best part? These glasses only cost a tiny bit of extra energy to wear.

Technical Summary: AOT-POT

Problem Statement

The pre-training of neural operators on diverse Partial Differential Equation (PDE) datasets represents a promising paradigm for building general-purpose surrogate models in scientific machine learning. However, this task faces a fundamental challenge: the inherent complexity and structural diversity of PDE solution operators. Existing approaches primarily attempt to address this by scaling model capacity (e.g., wider layers, deeper networks, or sparse Mixture-of-Experts). These methods leave the target solution operators unchanged, forcing a single network to approximate a heterogeneous family of complex operators simultaneously, which inevitably incurs large errors. Furthermore, simply increasing capacity does not exploit PDE-specific characteristics to simplify the learning target itself.

Methodology: AOT-POT

Inspired by classical numerical analysis, where problem-tailored operator transformations (e.g., Fourier basis for Laplacians) map complex PDEs to simpler, more tractable forms, the authors propose AOT-POT (Adaptive Operator Transformation for Pre-training Operator Transformer). The core insight is that instead of enlarging the model, one should transform the diverse solution operators into a unified, simpler form that is easier to model jointly.

Since the optimal transformation varies across PDE types, AOT-POT implements an adaptive, input-dependent mechanism directly within the neural operator architecture. The architecture expands hidden representations into $n$ parallel streams and employs three complementary components at every layer to mimic an adaptive operator transformation:

Adaptive Aggregation ( $a_l$ ): Before each sub-layer, the $n$ parallel streams are aggregated into a single input using input-dependent weights. This selects a specific combination of basis components for the current PDE instance.
Adaptive Redistribution ( $d_l$ ): After the sub-layer processes the aggregated input, the output is redistributed back to all $n$ streams using separately learned, input-dependent weights. This reshapes which components receive the update.
Sinkhorn-Projected Doubly Stochastic Mixing ( $T_l$ ): A transformation kernel mixes information across the streams. To ensure training stability and prevent signal explosion or attenuation across deep layers, $T_l$ is constrained to the Birkhoff polytope (the set of doubly stochastic matrices) via the Sinkhorn-Knopp projection. This ensures the transformation is volume-preserving with bounded spectral norms.

These mechanisms collectively perform a per-input "change of basis" on the latent solution operator. The backbone (a Fourier attention layer) effectively learns to approximate a transformed, simpler equivalent $\tilde{G}$ rather than the original heterogeneous operator $G$ . The entire process is parameterized by lightweight linear projections derived from the hidden state, adding only approximately 3% to the total parameter count.

Key Contributions

New Perspective: The paper introduces adaptive operator transformation as a new axis for improving multi-PDE pre-training, orthogonal to the prevailing paradigm of scaling model capacity. It demonstrates that architectural modules can act as learnable, input-dependent transformations, recasting complex operators into simpler equivalents without altering the underlying backbone.
Novel Architecture (AOT-POT): A novel architecture that generalizes standard residual connections to dynamically computed, doubly stochastic multi-stream transformations. It enables a single model to handle diverse PDEs by emulating an adaptive operator transformation at every layer.
Empirical Validation: Extensive experiments on 12 PDE benchmarks demonstrate that AOT-POT achieves state-of-the-art performance with minimal parameter overhead.

Experimental Results

The authors evaluated AOT-POT on 12 datasets spanning four benchmarks (FNO, PDEBench, PDEArena, CFDBench) across three model scales (Tiny, Small, Medium).

Zero-Shot Performance: AOT-POT achieves state-of-the-art results on 11 out of 12 datasets without fine-tuning. Compared to the strong baseline DPOT-S (30M parameters), AOT-POT-S (31M parameters) reduces the relative $L_2$ error by an average of 40.9%, with a maximum reduction of 77.6% (on FNO- $\nu=10^{-3}$ ).
Fine-Tuning: After fine-tuning, AOT-POT achieves the best results on all 12 datasets. It reduces $L_2$ error by up to 92% on in-domain PDEs and 89% on out-of-domain PDEs (unseen types during pre-training).
Parameter Efficiency: AOT-POT-S (31M) outperforms DPOT-M (122M) on 11 of 12 datasets, demonstrating that adaptive transformations are significantly more parameter-efficient than simply scaling model size.
Stability and Generalization:
- Long-Trajectory Rollouts: AOT-POT exhibits superior stability in autoregressive rollouts, with error growth rates significantly lower than baselines (e.g., 9.8 $\times$ growth over 500 steps vs. 24.1 $\times$ for DPOT-FT).
- Sample Efficiency: In few-shot fine-tuning scenarios, AOT-POT requires significantly fewer samples to match or exceed the performance of larger baselines.
- Interpretability: The learned transformation kernels ( $T_l$ ) encode PDE-specific signatures. A nearest-neighbor classifier on these features achieves 99.7% accuracy in distinguishing PDE families, generalizing to unseen PDE types with 100% accuracy.

Significance and Claims

The paper claims that AOT-POT establishes adaptive operator transformation as an effective and complementary direction for advancing PDE foundation models, distinct from and orthogonal to the strategy of simply scaling model capacity.

The authors emphasize that their approach is grounded in classical numerical analysis but instantiated as a learnable, input-dependent module within deep learning. By reshaping diverse solution operators into a unified form, AOT-POT allows a single architecture to effectively approximate an entire family of operators. The results suggest that the optimal transformation is strongly PDE-specific and that a single static transformation is insufficient; hence, the adaptive, input-dependent nature of the AOT block is crucial. The paper concludes that this architectural innovation offers a path to more robust, generalizable, and parameter-efficient scientific machine learning models.