SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shashank Subramanian  
Lawrence Berkeley National Lab  
shashanksubramanian@lbl.gov
Kurt Keutzer  
UC Berkeley  
keutzer@eecs.berkeley.edu
Dmitriy Morozov  
Lawrence Berkeley National Lab  
dmorozov@lbl.gov
Peter Harrington  
Lawrence Berkeley National Lab  
pharrington@lbl.gov
Wahid Bhimji  
Lawrence Berkeley National Lab  
wbhimji@1bl.gov
Michael W. Mahoney  
ICSI, LBNL, and Department of Statistics  
mmahoney@stat.berkeley.edu
Amir Gholami  
ICSI, UC Berkeley  
amirgh@berkeley.edu
Pre-trained machine learning (ML) models have shown great performance for a wide range of applications, in particular in natural language processing (NLP) and computer vision (CV). Here, we study how pre-training could be used for scientific machine learning (SciML) applications, specifically in the context of transfer learning. We study the transfer behavior of these models as (i) the pretrained model size is scaled, (ii) the downstream training dataset size is scaled, (iii) the physics parameters are systematically pushed out of distribution, and (iv) how a single model pre-trained on a mixture of different physics problems can be adapted to various downstream applications. We find that—when fine-tuned appropriately—transfer learning can help reach desired accuracy levels with orders of magnitude fewer downstream examples (across different tasks that can even be out-of-distribution) than training from scratch, with consistent behaviour across a wide range of downstream examples. We also find that fine-tuning these models yields more performance gains as model size increases, compared to training from scratch on new downstream tasks. These results hold for a broad range of PDE learning tasks. All in all, our results demonstrate the potential of the "pre-train and fine-tune" paradigm for SciML problems, demonstrating a path towards building SciML foundation models. Our code is available as open-source at [1].
Foundation models have received considerable interest recently [3]. This terminology refers to certain models that are trained on extremely large and diverse quantities of data and applied to a wide range of tasks. Rather than being designed for any single task, a foundation model serves as a "prior"
Figure 1: Our setup consists of creating diverse training datasets, sampling both PDE coefficients and source functions simultaneously with different PDE operators and input data (coefficients, sources) distributions for pre-training. A neural operator is then pre-trained to predict the PDE solutions given these inputs and the ground truth solutions (computed through PDE solvers). The pre-trained model is then adapted with minimal fine-tuning (zero-shot or few-shot), and it is used in various downstream tasks (PDE systems) that can be in-domain or out-of-domain from the pre-training datasets. The pre-training with multiple solution operators allows the same model to transfer to several very different systems. For instance, PDE 2 (Helmholtz) manifests highly oscillatory solutions compared to PDE 1 (Advection-Diffusion) or PDE 3 (Poisson's). We further characterize the scaling and transfer properties of this model as a function of downstream data scale and model size scale.
or "foundation" upon which other models can be built. It does so by using transfer learning (TL) methods to fine-tune or adapt the foundation model to a wide range of downstream tasks, using minimal additional data for each additional task. Perhaps the most well-known foundation models are pre-trained large-language models (LLMs) such as BERT [11] and the GPT models [38, 39, 5]. The scaling with respect to the amount of data, the size of the model, and the amount of compute [23, 20, 21] is key to the training of these models. An important aspect of a trained foundation model is the ability to perform tasks seemingly different than those for which it was trained by leveraging shared features across the training tasks. This approach to model development is quite different than the traditional approach of training a one-off model from scratch for each specific problem and each specific dataset. Naturally, it is of interest how broadly this methodological approach can be applied.
Scientific machine learning (SciML) [45] is an area that combines tools from ML and scientific computing to address domain-specific scientific and engineering challenges. It holds promise to drive the next wave of data-driven discovery in the physical and engineering sciences. Recent work has highlighted the promise [42, 32, 30, 24] as well as some of the many challenges [28, 12] of developing SciML models—in general as well as with the traditional one-off learning approach. Many SciML models emulate physical systems described by Partial Differential Equations (PDEs). For example, Physics-Informed Neural Networks [42, 28] impose the PDE as a soft penalty in the loss function. However, they are restricted to solving a single instance of the PDE. The Neural Network (NN) needs to be retrained for each new set of PDE physics coefficients, sources, and/or initial/boundary conditions (IC/BCs). Subsequent models have been developed to learn the full solution operator [30, 32] by training across different coefficients (and/or initial and boundary conditions). These neural operators learn mappings between two function spaces from a finite collection of input-output pairs (that represent the coefficients/initial or boundary conditions as the input and the PDE solution function as the output). This makes them more general and versatile in emulating any PDE system. However, with new coefficients/sources or new differential operators, they too need to be retrained from scratch.
In this paper, we adopt and evaluate the methodology that has been applied successfully in CV and NLP to develop foundation models, with the goal of determining whether such a model is even possible for SciML problems. In particular, we provide an extensive analysis of the scaling and TL behavior of neural operators trained on diverse training datasets from multiple PDE systems. An important aspect of this approach is to explore several dimensions that include the model (architecture and scale), data (diversity and scale), training recipes (pre-training and fine-tuning), and out-of
distribution (OOD) transfer behaviour. For transfer behavior, a major difference from CV/NLP is the diversity of scales, features, and solution behaviors across PDEs. For example, while OOD shifts in CV/NLP typically involve things like different color distributions or semantic contexts, in SciML, it is possible to get significantly different behavior of a PDE solution as physics parameters and conditions change that manifest as novel features (for example, increased emphasis on very different frequency scales of the solution causing significant shifts in detail, sharpness, and other image attributes). For LLMs, given the maturity of the NLP community within ML, these dimensions are well-explored. For SciML problems, in contrast, all these dimensions are open questions. Here, we explore several of these questions. We do so in the context of a specific model architecture, namely, the Fourier Neural Operator (FNO), as a prototypical SciML model that has demonstrated promising results modeling PDEs across a wide range of scientific applications [30, 37, 49, 53, 29]. We focus on the scaling and TL behavior of the FNO on common PDE systems that include Poisson's, Advection-Diffusion, and Helmholtz PDE systems. These systems underpin a wide range of physical systems: fluid flow systems; biological simulations; wave propagation systems; and many others.
See Fig. 1 for a schematic summary of our methodological approach. Our main results demonstrate the potential of the "pre-train and fine-tune" paradigm for SciML problems, demonstrating a path towards building SciML foundation models. In more detail, our main contributions are the following.
accuracies (see Fig. 4b, Fig. 4c). As we systematically go further OOD (see the quantification in Tab. 1), we observe the performance gains expectedly reduce, with more significant drop in the low data regimes (e.g., see Fig. 4d).
5. Transfer learning behavior over multiple operators. We study the effect of simultaneous pre-training on multiple PDE systems that exhibit qualitatively different solution behaviors (e.g., Poisson's and Helmholtz operator solutions show very dissimilar patterns, given a source function). We include the coefficient/source functions for all the operators as inputs to the model, with zero values if those terms do not exist in a given PDE instance. During inference, the zero inputs restrict the neural network to make predictions for the correct operator (see §4 for details). Among other things, we show that the same model pre-trained on different operators retains its gains across different downstream tasks (see Fig. 5), paving the way for it to be used in the foundational sense.
Pre-training a foundation model requires that we first collect a large amount of diverse training data and then train a base model that could subsequently be used with TL for a downstream application. There are several possible methods for TL of the foundation model. One approach is in-context learning, where the foundation model is prompted with few-shot input-output examples for the downstream problem followed by the target input. The model then sees these examples and learns how to compute the target output. This is the approach used by modern NLP models such as
GPT [38, 39, 5] as well as the work of [52]. While this approach is very useful for cases with very few training datapoints available, it is often better to fine-tune a foundation model for a downstream task, when one has access to more downstream training data. This is also supported by NLP models and it often results in better performance—if enough training data is available. We focus on the latter setup, and we study how the TL performance behaves for different problem setups. Our goal is to understand the different moving parts associated with training a foundation model for SciML applications and, specifically, the impact of model scaling, dataset size, and different physics involved in the problem. Below, we discuss (i) the different physics operators considered and our pre-training setup that includes the training dataset generation, (ii) NN model architecture setup for training and inference, and (iii) performance metrics.
PDE/physics system setup. We consider three PDE systems that are common building blocks for many scientific application: 2D Poisson's; Advection-Diffusion; and Helmholtz equations. These PDE systems can be formulated as follows:
1. Poisson's (SYS-1): We consider a prototypical elliptic system with periodic boundary conditions in domain  $\Omega = [0,1]^2$ :
$$
- \operatorname {d i v} K \nabla u = f \quad \text {i n} \Omega , \tag {1}
$$
where  $u(\pmb{x})$  is the solution (state) function,  $f(\pmb{x})$  is a source (forcing) function, and  $K$  is the diffusion coefficient tensor. We use  $K$  to quantify the physics of this system.
2. Advection-Diffusion (SYS-2): We also consider a steady-state advection-diffusion equation that illustrates competing physical processes (adjective and diffusive processes) through two differential operators. We use periodic boundary conditions in domain  $\Omega = [0,1]^2$ :
$$
- \operatorname {d i v} K \nabla u + \boldsymbol {v} \cdot \nabla u = f \quad \text {i n} \Omega , \tag {2}
$$
where  $u(\pmb{x})$  is the solution (state) function,  $f(\pmb{x})$  is a source (forcing) function,  $\pmb{K}$  is the diffusion coefficient tensor, and  $\pmb{v}$  is the velocity vector. To quantify the competing advective/diffusive scales of this system, we define the ratio of advection to diffusion as  $\Psi = \| \pmb{v} \cdot \nabla u \| / \| \mathrm{div} \pmb{K} \nabla u \|$ .
3. Helmholtz (SYS-3): Finally, we also consider the inhomogeneous Helmholtz equation with periodic boundary conditions in domain  $\Omega = [0,1]^2$ . We take this as an example challenging system that can exhibit high-frequency oscillatory spatial patterns that can be difficult to generalize. This system is formulated as:
$$
- \Delta u + \omega u = f \quad \text {i n} \Omega , \tag {3}
$$
where  $u(\pmb{x})$  is the solution (state) function,  $f(\pmb{x})$  is a source (forcing) function,  $\omega > 0$  is the wavenumber used to quantify the underlying physics of this system.
Data setup. For the above PDE systems, we are interested in (i) a large and diverse training dataset for pre-training and (ii) several downstream datasets (tasks) to quantify the TL performance. Given that we can solve these PDEs numerically, we can generate diverse set of training and testing datasets in a controllable fashion by varying the different parameters in these PDEs. In particular, we vary the following parameters for dataset generation:
(i) Source function sampling: We sample different source functions  $f \sim S(\sigma, s)$ , where  $S$  is a distribution that generates diverse and heterogeneous functions. (see Fig. 1 and Fig. A.1 for examples). Here,  $S$  represents a parameterization of the source function as a linear combination of  $n_g$  radial (Gaussian) basis functions  $\{\phi_i(\pmb{x})\}_{i=1}^{n_g}$ , where  $\phi_i(\pmb{x}) = \phi(\pmb{x} - \pmb{x}_i)$  is a Gaussian function centered at grid point  $\pmb{x}_i$ . Specifically:  $f(\pmb{x}) = \sum_{i=1}^{n_g} \phi_i(\pmb{x}) p_i$ , with  $\pmb{p} = \{p_i\}_{i=1}^{n_g}$  as the parameterization vector. The spatial profile controlled by  $\sigma$ , the standard deviation of the Gaussian function, is preset to a small value to encourage high variability. Examples are sampled by uniformly randomly sampling  $p_i \sim \mathcal{U}(0,1)$ . We further introduce heterogeneity by controlling the sparsity  $s$  of  $\pmb{p}$  ( $s$  defined as the number of non-zero components; see Appendix §A.1 for full details).
(i) PDE coefficient sampling: In SYS-1, we sample diffusion coefficient tensors  $\pmb{K} \sim \mathcal{K}(\lambda)$ , where  $\mathcal{K}$  is a distribution that generates varying scales of anisotropy and spread in the diffusion process:  $\pmb{K} = \pmb{R}^{-1}\pmb{D}\pmb{R}$  with  $D = \mathrm{diag}(1,e)$  and  $\pmb{R} = \mathrm{rot}(\theta)$ , where  $e$  is an eigenvalue of the tensor that controls the anisotropy and extent of diffusion and  $\mathrm{rot}(\theta)$  is a rotation matrix with angle  $\theta \sim \mathcal{U}(0,2\pi)$  that controls the general diffusion direction. In SYS-2, we additionally also sample the velocity vector  $v$  direction from  $\mathcal{U}(0,2\pi)$ . The ratio of advection to diffusion  $\Psi$  is changed by scaling the velocity. In SYS-3, we sample the wavenumber  $\omega$  as uniform
integers. We visualize the sampling process and resulting solutions with different ranges of underlying physics in Fig. A.1.
For each of the problems that we consider in the results section, we generate  $2^{15}$  input-output samples (pairs) of data, where the inputs include the source  $f$  as well as any PDE coefficients  $(K, v, \omega)$ , along with  $2^{12}$  validation and testing samples each. The validation dataset is used for hyperparameter optimization, and the testing dataset is used for quantifying model performance. The pre-trained model is trained on the  $2^{15}$  training examples.
We then perform different experiments to evaluate how this model can adapt/TL to different downstream tasks whose data could come from the following distributions: (i) same distribution as in the pre-training dataset (i.e., different input/output pairs but drawn from the same distribution of PDE coefficients/sources); and (ii) the harder task of adapting/TL to a downstream problem that can have slight/large deviation from the dataset used to pre-train the model. For the latter, we create the OOD data by keeping the PDE operator the same as the pre-training task, but sample the coefficients from a different range as in the pre-training. Given this dataset, we then study the TL behaviour for each case (both within distribution and OOD) by scaling both the downstream dataset size, as well the model architecture size which is discussed next.
Pre-training method for training and inference. The inputs to our model are 2D spatial functions discretized at  $h \times w$  and represent the sources and PDE coefficients. These input discretized functions are batched together to form an input tensor in  $\mathbb{R}^{h \times w \times c}$ . The output of the model is the numerical solution of the PDE in  $\mathbb{R}^{h \times w}$ . For the model architecture, we consider the FNO (details in Appendix §A.2). This model bears similarities to both vision transformer-like architectures (fixed image/feature resolution across the depth of the network) and convolutional architectures (successive global convolutions facilitated via FFTs). FNO is also a good choice for our setup as the problems we consider all have periodic boundary conditions (the FNO can also be adapted for non-periodic boundaries). The main modification that we make to the FNO is to incorporate a per-instance normalization layer in the model. We found that this is a critical component as the norm of the input data spans a wide range of values (up  $100 \times$  for our dataset). Please see Appendix §A.5 for details. Furthermore, we consider an  $\ell_2$  loss function for pre-training. That is, the model is trained to predict the output solution given the input by minimizing a mean-squared error loss between the prediction and the ground truth.
To test how this pre-trained model can adapt to different downstream applications we consider the following cases. First, we consider zero-shot adaptation where the pre-trained model is tested on a downstream application without any fine-tuning. Second, we consider few-shot fine-tuning, where the pre-trained model can be fine-tuned on the downstream training dataset. We consider varying sizes for this downstream dataset size. Ideally, we prefer the pre-trained model to achieve good performance with as few examples as needed from the downstream application. We perform an ablation study by training a model from scratch on the downstream dataset alone to see how much gain can be achieved by using a pre-trained model. We also test how the pre-training adaptation performance changes as we scale the model size. This is motivated by observations in NLP [23] where we expect larger models to show better adaptation performance. To do the model scaling, we focus on two model hyperparameters: the embedding dimension  $d$  and the number of Fourier modes used in the FNO  $m$ . (See Appendix §A.2 for the details.) We first fix  $d = 16$  and scale  $m \in \{4, 16\}$ , then fix  $m = 32$  and scale  $d \in \{32, 128\}$  to approximately increase the parameter count  $16 \times$  in each scaling experiment from  $64K$  to  $256M$  parameters.
The main limitation of our work is that we focus on only the FNO model, whereas several other architectures (e.g. ViT models, DeepONets) exist. As discussed in §1, the FNO has been applied to a wide range of applications in SciML and is representative of a performant SciML architecture—restricting our analysis to one model architecture makes the scope of this study more feasible. However, this analysis needs to be done across other models architectures, as well.
Our main results are the following. We demonstrate that pre-training a model on a diverse corpus of data and then fine-tuning it on downstream tasks leads to significantly better performance than training
<sup>1</sup>Constant coefficients are simply replicated across the  $h \times w$  dimensions.
(a) SYS-1(5,10) pre-trained from SYS-1(1,5)
(b) SYS-2(1,2) pre-trained from SYS-2(0.2,1)
Figure 2: Addressing (Q1). Testing error as a function of downstream examples for SYS-1 and SYS-2. We visualize the distribution of pre-training and downstream dataset physics at the top to illustrate (and quantify) the extent of distributional shifts. We observe excellent zero-shot and few-shot TL performance of the pre-trained model despite the modest OOD shifts and in medium-data regimes about  $100 \times$  increase in data efficiency. We observe diminishing returns from pre-training at the large-data regime ( $O(2^{15})$  examples), which has as many examples as used in pre-training.
a model from scratch. This holds even when the downstream data falls outside of the pre-training distribution, including when different physics models are combined. The advantage of pre-training is especially pronounced when the downstream data is limited, which is the most significant setting in practice, motivating the creation of foundation models for scientific machine learning.
To justify these conclusions, we focus on four key questions. What is the effect of (Q1) downstream dataset size and (Q2) neural operator model parameter size on TL? What is the TL behavior of the neural operator (Q3) over the underlying physics and (Q4) over multiple solution operators?
(Q1): Downstream dataset scaling. For SYS-1, we consider the pre-training system SYS-1(1,5), with diffusion constructed by sampling eigenvalue  $e \sim \mathcal{U}(1,5)$  of the diffusion tensor  $K$ . See Fig. A.1 for visualizations. This represents near isotropic to  $5 \times$  anisotropic diffusion. We use  $e \sim \mathcal{U}(5,10)$  as the downstream dataset SYS-1(5,10). This represents  $5 \times -10 \times$  anisotropic diffusion and is moderately out-of-distribution (OOD) from the pre-training dataset. While we systematically quantify the OOD effects in (Q2), we use this specific test-case to illustrate the effect of the size of the downstream dataset. We plot the behaviour of testing error as a function of downstream examples in Fig. 2a. We train an FNO model for each number of downstream examples (x-axis on the plot) starting from "scratch" (random initialization) as well as from the pre-trained model parameters, with tuned hyperparameters for each experiment (see details in Appendix §A.4).
We illustrate the extent of distributional shift between the pre-training and downstream datasets through the range of diffusion tensor eigenvalue  $e$  — in this test case, a modest shift with no overlap (but relatively close). The testing error monotonically decreases as more downstream data is used for training, as we expect. The zero-shot TL shows excellent performance despite the moderate OOD shift of downstream examples. When training from scratch, we define "zero-shot" predictions as the output of the model with random initialization. With "few-shot" learning ( $O(10)$  downstream examples), we observe a consistent performance increase over training from scratch. Given a desired error for downstream performance, TL from the pre-trained model can require orders of magnitude less data—for example, a desired error of 1E-2 needs only about 64 downstream data examples for fine-tuning, whereas training from scratch requires  $8K$  (about  $100 \times$  more) examples to reach the same accuracy level. We further explore this in Appendix §B.1 and find the pre-trained model generally saves  $O(1K - 10K)$  data compared to training from scratch in the few-shot learning setting, and outperforms training from scratch at all scales. Finally, with greater amounts of downstream data, the pre-training provides consistent performance gains until we enter the "large target data" regime, where the number of fine-tuning examples approaches the size of the entire pre-training dataset, and we see diminishing returns from pre-training. We repeat this experiment for SYS-2, using the following test-case: the pre-training system SYS-2(0.2,1) consists of advection-to-diffusion rates  $\Psi \in \sim (0.2,1)$ , representing about  $1 \times -5 \times$  diffusion (relative to advection). For the downstream test, we use a modest TL with SYS-2(1,2) with  $\Psi \in \sim (1,2)$  representing  $1 \times -2 \times$  advection