1. Pre-training dataset generation. We develop a large suite of datasets, and we train our models on data where all the variables (inputs) of any PDE operator are sampled. This is an important step towards developing NNs that can generalize across a variety of downstream tasks, and it extends several previous works, including the original FNO [30], where certain inputs are kept fixed. Not sampling can trivially push the neural operator OOD if, e.g., the source function was changed. We study transfer learning to both in-domain and out-of-domain distributions, characterized by different samplings of PDE coefficients (diffusion, advection, wavenumbers, etc.) and inhomogeneous source functions. We emphasize the construction of pre-trained datasets with sufficient diversity as well as normalization strategies, without which we observe significant degradation of performance.
2. Downstream (target) data scaling. We study the effect of the scale of downstream (target) data in TL performance from the pre-trained model. Here, we assume that a large amount of data is available for pre-training, and data are limited for the downstream tasks (as in many scientific examples); and we are interested in reaching desired accuracy levels with the least amount of additional downstream data. We consider both zero-shot and few-shot TL: zero-shot is the direct evaluation of the pre-trained model on the downstream dataset; and few-shot involves using O(10) downstream data to fine-tune the pre-trained model. While few-shot typically refers to prompting with in-context examples during inference in NLP, here we define it as fine-tuning the model with a few examples. We observe that TL from the pre-trained model can lead to significant performance gains over training the model from scratch on the downstream (target) data, with orders of magnitude less data needed to reach a desired accuracy level (see Fig. 2). We observe this gain over a wide range of data scales, until we enter the "large target data" regime (as much data as pre-training), where we observe similar accuracies for TL as training from scratch.
3. Model (parameter) size scaling. We study the parameter scaling of the model by scaling our model sizes from 64K to 256M parameters (a multiplicative factor of 4K). We observe an error saturation at small model sizes (due to insufficient model expressivity) that monotonically drops as we increase the model size. While both fine-tuned models and models trained from scratch exhibit gains with increased model size, we observe that fine-tuning achieves greater performance gains with parameter scaling (see Fig. 3).
4. Transfer learning behavior over underlying physics. We study the effect of varying the underlying physics in the target domain. In SciML (unlike traditional non-scientific ML), there are typically fundamental constraints such as conservation laws that govern the behaviour of the solution. In some cases, we may even have access to parameters of the underlying physical constraints, and thus a physical understanding of the distribution of the data. It is natural, and necessary, to systematically quantify the effect of these parameters as our downstream tasks go OOD, as this provides a good way to test the (OOD) generalization capability of pre-trained models for SciML applications. We find that for in-distribution TL, the pre-trained model can significantly outperform a model trained from scratch, irrespective of how many new data examples were used for fine-tuning, until the large target data regime (e.g., see Fig. 4a), showing orders of magnitude better accuracies than training from scratch. We also observe these gains for downstream tasks that are moderately OOD, with few-shot fine-tuning providing again orders of magnitude better