Strategies for Pretraining Neural Operators

TL;DR

This work provides a model-agnostic evaluation of pretraining strategies for neural operators solving PDEs, highlighting that transfer learning and physics-informed objectives generally yield the strongest gains, especially with data augmentations. It demonstrates that pretraining effectiveness depends on both model architecture and dataset, with transformer- and CNN-based backbones benefiting more than vanilla operators. Data augmentations consistently improve performance, particularly shift-based augmentations that preserve key physics, and the combination of transfer learning with shift augmentation performs best in many cases. The findings offer practical guidance for data-efficient PDE modeling and motivate future work on principled, PDE-specific pretraining frameworks across architectures and PDE families.

Abstract

Pretraining for partial differential equation (PDE) modeling has recently shown promise in scaling neural operators across datasets to improve generalizability and performance. Despite these advances, our understanding of how pretraining affects neural operators is still limited; studies generally propose tailored architectures and datasets that make it challenging to compare or examine different pretraining frameworks. To address this, we compare various pretraining methods without optimizing architecture choices to characterize pretraining dynamics on different models and datasets as well as to understand its scaling and generalization behavior. We find that pretraining is highly dependent on model and dataset choices, but in general transfer learning or physics-based pretraining strategies work best. In addition, pretraining performance can be further improved by using data augmentations. Lastly, pretraining can be additionally beneficial when fine-tuning in scarce data regimes or when generalizing to downstream data similar to the pretraining distribution. Through providing insights into pretraining neural operators for physics prediction, we hope to motivate future work in developing and evaluating pretraining methods for PDEs.

Figures (5)

• Figure 1: An illustration of pretraining strategies adapted from computer vision (CV) and predicting PDE characteristics. Left: CV methods can be described by different shuffling mechanisms and losses. Binary pretraining only classifies if a sequence is shuffled or not, while TimeSort, SpaceSort and Jigsaw sort sequences shuffled in various ways, either along the spatial, temporal, or combined dimensions. Right: PDE data has inherent structure that can be leveraged to predict underlying characteristics. Coefficients of the PDE can be regressed, as well as its spatial and temporal derivatives. Additionally, inputs can be masked to regress the solution field $u$ and learn underlying dynamics.
• Figure 2: Experimental Setup. During pretraining we consider different data augmentations, model choices, and pretraining tasks and evaluate their downstream performance through fine-tuning on physics prediction tasks. During fine-tuning, we leverage the same pretrained model to improve fixed-future or autoregressive prediction on the same pretraining data distribution, unseen coefficients, or new PDEs. Through this setup we can explore a wide variety of pretraining strategies and augmentations and quantify their effects on different models, PDEs, datasets, and tasks.
• Figure 3: Fixed Future Scaling Behavior: For each model, a specific PDE/distribution is displayed. Within each graph, the error of various pretraining strategies at different sample sizes is displayed. Validation errors are averaged over 5 seeds, and error bars denote 1 std-dev. Derivative errors are omitted as outliers.
• Figure 4: Auto-regressive Scaling Behavior: For each model, a specific PDE/distribution is displayed. Within each graph, the performance of various pretraining strategies at different sample sizes is displayed. Validation errors are averaged over 5 seeds, and error bars denote 1 std-dev.
• Figure 5: t-SNE Embeddings of CV Pretrained Models: We display latent embeddings after pretraining models on Binary, TimeSort. or Jigsaw objectives. We see that models learn to sort/classify shuffled samples well and can visualize the relative difficulties of the proposed pretraining strategies.