Pretraining a Neural Operator in Lower Dimensions

AmirPouya Hemmasian, Amir Barati Farimani

Abstract: There has recently been increasing attention towards developing foundational neural Partial Differential Equation (PDE) solvers and neural operators through large-scale pertaining. 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.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

Addressed reviewer comments:
- Split the large introduction into an introduction section and a related works section
- added further clarification and explanation at the end of section 2 (related works) in response to the reviewers' feedback
- added a new paragraph and a pseudocode at the beginning of the methodology section to explain the methodology and its limitation more clearly according to the reviewers' feedback
- added a short explanation on how table 1 was designed in the methodology section.
- improved the writing of the dataset setup to fix the confusion about the pertaining datasets and their correspondence to the downstream tasks
- added runtime of pertaining and fine-tuning for our experiments for the sake of comparison and an example of the cost difference between 1D and 2D data.
- Added a sentence to the conclusion to provide a general recommendation for the fine-tuning configuration.
- fixed the issue with citation style.
- fixed typos.

Camera ready revision:

- deanonymized
- added link to code

Code: https://github.com/BaratiLab/PreLowD

Assigned Action Editor: ~Xingyou_Song1

Submission Number: 3192