Towards Universal Neural Operators through Multiphysics Pretraining

AI summary

Overview

Research area: Scientific machine learning, specifically neural operators (Fourier Neural Operators and transformer-based operator blocks) applied to partial differential equations (PDEs), with a focus on transfer learning and pretraining across multiple physical systems.

Technical level: Advanced. The paper assumes familiarity with operator learning, kernel integral operators, attention mechanisms, and PDE terminology.

Scope: The paper evaluates whether transformer-based and state-space neural operator architectures can be pretrained on multiple PDE problems and then efficiently fine-tuned on new dynamics using lightweight, problem-specific adapters.

What This Paper Is About

Training neural operators for physical simulation is computationally expensive, and prior pretraining work has been largely case-specific rather than general. The authors investigate whether advanced neural operator architectures, previously applied only to specific problems, can transfer knowledge across diverse PDE problems: extrapolation to unseen equation parameters, incorporation of new variables, and transfer from multi-equation datasets. Their goal is a standardized recipe for pretraining an operator on simpler physics and fine-tuning it cheaply on harder or different physics.

Key Contributions

1. An adapter-based multiphysics training approach. The paper introduces a scheme in which the lifting and projection blocks of a neural operator serve as problem-specific adapters, letting the model train simultaneously on PDE problems that require different sets of input functions.
2. A general transfer learning evaluation for transformer-based operators. Rather than confining transformer-style operators to a single problem, the authors test them across several transfer scenarios without modifying the modeling approach.
3. Demonstrated gains from pretraining plus fine-tuning. Results indicate the transfer learning approach can significantly enhance model quality and reduce fine-tuning costs relative to training from scratch, measured in both error and average epoch time.
4. An open experimental release. Code and experiments are provided at https://anonymous.4open.science/r/multiphysics_neurop-F385/.

Main Findings

• Pretraining improves accuracy and speed. For the out-of-sample parameter scenario, Mamba FNO pretrained reached MSE 1.009×10⁻⁷ and NMAE 0.0120 versus 1.193×10⁻⁷ and 0.0213 from scratch, with average epoch time 21.91 s versus 40.14 s.
• The largest speedup came from Perceiver IO. Perceiver (pretrained) recorded an average epoch time of 3.21 s versus 204.73 s from scratch in Table 1, and 20.4 s versus 804.0 s in Table 2, while also lowering error (NMAE 0.0169 vs 0.0219 in Table 1; 0.0051 vs 0.0074 in Table 2).
• Swin-v2 achieved the best error in Table 1 but at high cost. Swin-v2 (labeled "p.+s.") reached MSE 4.391×10⁻⁸ and NMAE 0.0092, the lowest error in that table, but with an average epoch time of 101.3 s and approximately 10⁹ parameters.
• Pretrained CoDA-NO consistently beat its scratch counterpart. CoDA-NO (pretrained) gave MSE 2.881×10⁻⁷ / NMAE 0.0343 versus 4.912×10⁻⁷ / 0.0712 from scratch in Table 1, and 1.043×10⁻⁵ / 0.013 versus 1.239×10⁻⁵ / 0.018 in Table 2.
• The default FNO baseline is smaller but weaker. FNO (scratch) used roughly 10⁶ parameters and recorded MSE 1.774×10⁻⁷ / NMAE 0.0204 (Table 1) and MSE 7.286×10⁻⁶ / NMAE 0.0121 (Table 2).
• Transfer held across three distinct scenarios. The approach was validated on out-of-sample parameter values, input function set extension (adding convection to the heat equation, adding advection to reaction-diffusion), and multi-physics pretraining with fine-tuning on different dynamics.
• Pretrained models operate on a reduced parameter search. The reported speedups are attributed to optimizing only a subset of parameters during fine-tuning, whereas the "from scratch" runs involved a full parameter search.

Methodology in Plain English

The authors build on the standard neural operator pattern: a lifting network raises the input functions into a higher-dimensional hidden representation, a stack of integral-operator blocks transforms that representation, and a projection network maps it back to the output functions. Because the lifting and projection blocks are small and tied to the specific problem (for example, how many input functions the equation has), the authors treat them as adapters, similar in spirit to adapter modules used in large language models.

During pretraining, all parameters are optimized at once across several PDE problems that may have different input function sets. During fine-tuning, the core operator blocks are frozen and only the new lifting and projection adapters are trained. This decoupling keeps storage simple and makes it easy to extend the model to new problems.

For the operator blocks themselves, the authors test two families. The first inserts a Mamba state-space module after the lifting map, which they describe as a latent preconditioner that aligns embeddings with common dynamical motifs such as transport, diffusion, and oscillation before the Fourier integral layers. The second uses attention, specifically Perceiver IO blocks, where latent variables are the queries and keys and values come from FNO-based mappings of the inputs, followed by self-attention between latents and a symmetric cross-attention for outputs. Codomain attention is discussed as an alternative that measures similarity between features rather than samples.

Quality is measured with range-normalized mean absolute error (NMAE), normalized by the range of the ground-truth solution plus a small epsilon, alongside MSE. The experiments use Burgers' equation, the Gray-Scott reaction-diffusion model, the incompressible Navier-Stokes equations, heat equation variants extended with convection, reaction-diffusion extended with advection, and the PDEBench dataset.

Why This Matters

Impact on research: The work pushes neural operators toward the foundation-model paradigm, where one pretrained operator is reused across many physical systems instead of training a separate model per problem. It also provides a controlled comparison of state-space, attention-based, Swin-v2, and codomain-attention operators under identical transfer conditions, which helps clarify which architectural choices actually help generalization.

Real-world applications:

• Meteorological forecasting, where numerical solution of differential equations is computationally costly.
• Multi-physics simulation in engineering, cited as another setting where numerical solution is expensive.
• Fluid and gas dynamics modeled with variants of the Euler/Navier-Stokes equations.
• Reaction-diffusion systems exhibiting pattern formation, and advective transport processes.

Industry relevance: The adapter design keeps model storage small and lets organizations add a new physical regime by training only a small adapter on top of a frozen pretrained operator, which reduces the cost of obtaining a usable model for novel PDE-based problems.

Future Directions

• Scaling from this first stage toward a foundational model pretrained on vast and heterogeneous multiphysics dataset collections.
• Improving model performance and training procedures.
• Implementing data augmentation tools, both generic and PDE-specific, with the paper naming Lie symmetries as an example.
• Further examining neural operator generalizations, including the mesh-agnostic limitations the authors cite from prior work.
• Open question: the relative accuracy-versus-cost trade-off remains unresolved, since the most accurate model in Table 1 (Swin-v2, MSE 4.391×10⁻⁸, NMAE 0.0092) is also the largest (≈10⁹ parameters) and slowest per epoch among the pretrained-style entries (101.3 s), while Perceiver variants are far cheaper per epoch but less accurate.

Target Audience

Researchers and practitioners in scientific machine learning, operator learning, and computational physics who are interested in transfer learning, foundation models for PDEs, or designing cheaper surrogate solvers. It is also relevant to engineers who need to adapt a surrogate model to new physical regimes without paying the full cost of training from scratch. Readers unfamiliar with neural operators, kernel integral formulations, or attention mechanisms will need background reading first.

Authors’ abstract

Although neural operators are widely used in data-driven physical simulations, their training remains computationally expensive. Recent advances address this issue via downstream learning, where a model pretrained on simpler problems is fine-tuned on more complex ones. In this research, we investigate transformer-based neural operators, which have previously been applied only to specific problems, in a more general transfer learning setting. We evaluate their performance across diverse PDE problems, including extrapolation to unseen parameters, incorporation of new variables, and transfer from multi-equation datasets. Our results demonstrate that advanced neural operator architectures can effectively transfer knowledge across PDE problems.