Pretraining a Neural Operator in Lower Dimensions

TL;DR

This paper proposes a pretraining strategy for neural PDE solvers using lower-dimensional data, enabling effective transfer to higher-dimensional PDEs with reduced data costs.

Contribution

It introduces PreLowD, a pretraining approach on lower-dimensional PDEs, and demonstrates its effectiveness for higher-dimensional PDE solving using the FFNO model.

Findings

Pretraining on lower dimensions improves higher-dimensional PDE performance.

The FFNO model effectively transfers learned features across dimensions.

Fine-tuning strategies significantly impact transfer success.

Abstract

There has recently been increasing attention towards developing foundational neural Partial Differential Equation (PDE) solvers and neural operators through large-scale pretraining. However, unlike vision and language models that make use of abundant and inexpensive (unlabeled) data for pretraining, these neural solvers usually rely on simulated PDE data, which can be costly to obtain, especially for high-dimensional PDEs. In this work, we aim to Pretrain neural PDE solvers on Lower Dimensional PDEs (PreLowD) where data collection is the least expensive. We evaluated the effectiveness of this pretraining strategy in similar PDEs in higher dimensions. We use the Factorized Fourier Neural Operator (FFNO) due to having the necessary flexibility to be applied to PDE data of arbitrary spatial dimensions and reuse trained parameters in lower dimensions. In addition, our work sheds light on the effect of the fine-tuning configuration to make the most of this pretraining strategy. Code is available at https://github.com/BaratiLab/PreLowD.