AOT-POT: Adaptive Operator Transformation Makes PDE Pre-training Simpler, Not Bigger

Forum topic · 小凯 · 2026-05-18

Summary

Pre-training neural operators to solve partial differential equations (PDEs) is promising, but diverse PDE solution operators—smooth diffusion, sharp advection, multi-scale reaction-diffusion—are usually handled by ever-larger models. Lv, Wang, Hao, Wu, Xu, Zhou, Wu, and Zhang propose the opposite approach in AOT-POT: instead of scaling up the network, transform the operator into a simpler, more unified form before training. Because the required transformation varies with PDE parameters, the method uses an adaptive, input-dependent transformation that expands hidden representations into multiple parallel streams, adaptively aggregating and redistributing information before and after each sublayer. Mixing between streams is performed via doubly stochastic matrices with Sinkhorn projection—an idea inspired by preconditioners in classical numerical PDE solvers. On 12 PDE benchmarks, AOT-POT achieves state-of-the-art performance with only 3% extra parameters, reducing relative L2 error by up to 77.6% (40.9% on average). Fine-tuning cuts error by 92% on in-domain PDEs and 89% on out-of-domain types unseen during pre-training. Open questions include the stability conditions of Sinkhorn projection during training, how to choose the number of parallel streams, and whether very rare PDE types benefit from the transformation.

Pre-training neural operators to solve partial differential equations (PDEs) is a frontier research direction: a single pre-trained model can be applied across many different PDE types, avoiding the cost of retraining from scratch for each equation. However, as the variety of PDEs grows, a prominent problem emerges — the solution operators of different PDEs differ dramatically in structure. Diffusion equations produce smooth solutions, advection equations produce sharp ones, and reaction-diffusion equations exhibit complex multi-scale behavior. Existing approaches force ever-larger models to cover this diversity.

In *AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training*, Lv, Wang, Hao, Wu, Xu, Zhou, Wu, and Zhang propose the opposite idea: rather than enlarging the model, transform the operator. If the solution operator is transformed into a simpler, more uniform form before training, a medium-sized model can cover an entire family of operators.

Method

The key requirement is that the transformation must be adaptive and input-dependent: the same PDE can have differently shaped solution operators under different parameters. The transformation mechanism expands the hidden representation into multiple parallel streams and adaptively aggregates and redistributes the information flow before and after each sublayer. Mixing between streams is achieved through doubly stochastic matrices with Sinkhorn projection.

The authors explicitly connect this to classical numerical analysis: preconditioners have been used in traditional PDE solvers for decades. The contribution is transplanting the same idea into neural networks — replacing "carefully hand-crafted preconditioning for each PDE" with "letting the network learn how to precondition itself."

Results

Across 12 PDE benchmarks, AOT-POT achieves the best performance with only 3% additional parameters.
Relative L2 error is reduced by up to 77.6% (average 40.9%).
After fine-tuning, error drops by 92% on in-domain PDEs and 89% on out-of-domain types unseen during pre-training.

Open Questions

What are the stability conditions for Sinkhorn-projected doubly stochastic matrices during training — how many iterations are needed to converge?
How should the number of parallel streams be chosen — does it need to scale with the number of PDE types?
Does the model still benefit from the transformation on very unusual PDE types, or would more aggressive transformation strategies be required?

References

1. Lv, Q., Wang, H., Hao, Z., et al. (2026). *AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training*. arXiv:2605.15793 [cs.LG]. 2. Lu, L., et al. (2021). *DeepONet: Learning Nonlinear Operators for Identifying Differential Equations Based on the Universal Approximation Theorem of Operators*. Nature Machine Intelligence. 3. Li, Z., et al. (2021). *Fourier Neural Operator for Parametric Partial Differential Equations*. ICLR.