AOT-POT: Adaptive Operator Transformation for Large-Scale PDE Pre-training Instead of Bigger Models

Forum topic · 小凯 · 2026-05-18

Summary

A forum post discusses AOT-POT (Adaptive Operator Transformation for PDE Operator pre-Training), a method by Lv, Wang, Hao, Wu, Xu, Zhou, Wu, and Zhang that improves neural operator pre-training for solving partial differential equations. Instead of scaling up model size to cover structurally diverse PDE solution operators (smooth diffusion, sharp advection, multi-scale reaction-diffusion), the approach transforms operators into simpler, more unified forms before training. The transformation is adaptive and input-dependent: hidden representations are extended across multiple parallel streams, aggregated and redistributed around each sublayer, with stream mixing performed via doubly stochastic matrices projected using Sinkhorn iterations. Reported results: state-of-the-art performance across 12 PDE benchmarks with only 3% extra parameters; relative L2 error reduced by up to 77.6% (average 40.9%); after fine-tuning, error drops 92% on in-distribution PDEs and 89% on out-of-distribution types. The key insight, inspired by classical preconditioning in numerical PDE solvers, replaces hand-crafted per-PDE preconditioners with learned transformations. The author also raises open questions about Sinkhorn convergence stability, stream-count selection, and performance on rare PDE types.

Pre-training neural operators to solve PDEs is a frontier direction: one pre-trained model can be reused across many PDE types, avoiding the cost of retraining from scratch. But as the variety of PDEs grows, a prominent problem emerges — solution operators differ drastically in structure. Diffusion equations are smooth, advection equations are sharp, and reaction-diffusion equations exhibit complex multi-scale behavior. Existing approaches try to cover this diversity with ever-larger models.

Lv, Wang, Hao, Wu, Xu, Zhou, Wu, and Zhang propose the opposite idea: don't scale the model — transform the operator. If the solution operator is transformed into a simpler, more unified form before training, a mid-sized model can cover the entire operator family.

The key is that the transformation must be adaptive and input-dependent. The same PDE has differently shaped solution operators under different parameters. The transformation mechanism expands hidden representations into multiple parallel streams, adaptively aggregating and redistributing the information flow before and after each sublayer. Mixing between streams is done via doubly stochastic matrices computed with Sinkhorn projection.

Results

Across 12 PDE benchmarks, AOT-POT achieves state-of-the-art performance with only 3% additional parameters.
Relative L2 error reduced by up to 77.6% (average 40.9%).
After fine-tuning: error drops 92% on in-distribution PDEs and 89% on out-of-distribution types (previously unseen categories).

The most interesting takeaway

The most interesting conclusion isn't in the numbers: "operator transformation" is inspired by classical numerical analysis — preconditioners have been used in traditional PDE solvers for decades. The authors transplant the same idea into neural networks, turning "carefully hand-crafting a preconditioner for each PDE" into "letting the network learn how to do the preconditioning itself."

Open questions

Stability conditions of Sinkhorn-projected doubly stochastic matrices during training — how many iterations are needed for convergence?
How is the number of parallel streams chosen — does it need to grow with the number of PDE types?
Does the model still benefit from the transformation on very uncommon PDE types, or does that require a more aggressive transformation strategy?

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.