SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 3: Addressing (Q2). Model size scaling for SYS-1 and SYS-2 from  $64K$  to  $256M$  parameters for medium OOD test-cases. While finetuning consistently improves the model performance and data efficiency, we observe higher errors for small parameter regimes at  $64K$  due to insufficient model capacity. The performance gains are significantly boosted through finetuning with a larger model set sizes monotonically up to  $256M$  parameters.
(relative to diffusion). We visualize this shift in Fig. 2b (top) and also show the testing error as a function of downstream examples. Both experiments reveal the same trend: TL delivers much higher performance, which improves with fine-tuning up to a point of diminishing returns.
(Q2): Model (parameter count) scaling. As described in our model setup, we vary the embedding  $d$  and maximum Fourier modes  $m$  to approximately increase the parameter count  $16 \times$  in each scaling experiment from  $64K$  to  $256M$  parameters. For models trained from scratch, we repeat the data scaling experiments for each parameter count. For the pre-trained model, we first identify the ideal hyperparameters (through grid-search hyperparameter tuning) for each model scale and repeat the above training experiments. We visualize the testing errors as a function of downstream examples used for SYS-1(5,10) (pre-training dataset used: SYS-1(1,5) signifying moderate OOD) for the different model scales in Fig. 3 (left). At the  $64K$  parameter regime, the model capacity is insufficient, with large errors (greater than 1E-2) for either training recipe across the whole range of downstream example counts. As we move to larger models, both training from scratch and fine-tuning show higher performance that increase with more examples. Fine-tuning the pre-trained model boosts its performance, compared to training from scratch, as we increase the model scale and particularly across a wide range of downstream example counts (with  $256M$  parameter model showing the least errors). We repeat the model scaling process for Advection-Diffusion SYS-2 (pre-training dataset SYS-2(0.2,1) with moderately OOD SYS-2(1,2) for downstream) and observe similar trends in Fig. 3 (right).
Table 1: Different downstream datasets and extents of overlap with the pre-training dataset for SYS-1 and SYS-2, controlled by extent of anistropy (eigenvalue e) in diffusion tensor for SYS-1 and amount of advection relative to diffusion (ratio  $\Psi$ ) for SYS-2.
Pre-training	Downstream	Shift
SYS-1(1,5): e ~ U(1,5)	SYS-1(1,2.5): e ~ U(1,2.5)	None
SYS-1(2.5,7.5): e ~ U(2.5,7.5)	Mild
SYS-1(5,10): e ~ U(5,10)	Med
SYS-1(10,20): e ~ U(10,20)	Large
SYS-2(0.2,1): Ψ ∈ ~(0.2,1)	SYS-2(0.2,0.4): Ψ ∈ ~(0.2,0.4)	None
SYS-2(0.2,0.4): Ψ ∈ ~(0.2,0.4)	Mild
SYS-2(1,2): Ψ ∈ ~(1,2)	Med
SYS-2(2,5): Ψ ∈ ~(2,5)	Large
(Q3): TL behavior over underlying physics. We test both in-domain and out-of-domain physics effects by constructing downstream datasets that systematically deviate from the pre-training dataset. For SYS-1, we sample different ranges for  $e$  with varying overlap with the pre-training dataset. Similarly, for SYS-2, we use different ranges of advection-to-diffusion ratio  $\Psi$  showing different overlap. We highlight these systems (downstream and pre-training) in Tab. 1 for the two PDE systems.
(a) SYS-1(1,2.5)
(c) SYS-1(5,10)
(b) SYS-1(2.5,7.5)
(d) SYS-1(10,20)
Figure 4: Addressing (Q3). Testing error as a function of downstream examples for different downstream tasks used in SYS-1. We show the extent of overlap (signifying distributional shifts) between the pre-trained and downstream dataset at the top using the range of sampled diffusion tensor eigenvalue. For datasets within distribution, zero-shot TL is optimal. As the downstream dataset shifts moderately OOD, the zero-shot learning suffers gradually and is recovered through fine-tuning. This recovery is slower as the distributional shifts increase.
We repeat our downstream dataset scaling experiments on the different downstream tasks and show the trends for SYS-1 in Fig. 4. In particular, in Fig. 4a, we consider the downstream dataset within distribution of the pre-training dataset (as visualized by the  $e$  distribution at the top). We observe excellent zero-shot performance that is unaffected by further fine-tuning. In Fig. 4b, the downstream dataset is shifted mildly OOD. Although the zero-shot performance drops, it still shows low errors, significantly smaller than training from scratch. Further, the performance is improved with few-shot TL up to the point of diminishing returns with large numbers of downstream data examples. With further distributional shift (no overlap) in Fig. 4c, the zero-shot performance suffers more, but with a larger amount of fine-tuning recovers good performance. Finally, for large distributional shifts in Fig. 4d, the zero-shot and few-shot performance is poor, with TL showing relatively high errors, but even here TL improves over training from scratch. In this case, due to larger anisotropy, the system is also harder to emulate and might require more data in general. We repeat this analysis for SYS-2 and observe similar trends across different OOD downstream tasks (see Appendix §B.1).
(Q4): TL behavior over multiple operators. We further diversify the pre-training by including examples with different solution operators. We combine the datasets from three PDEs—Poisson's SYS-1(1,5), Advection-Diffusion SYS-2(0.2,1), and Helmholtz SYS-3(1,10) (where the wavenumber  $\omega \sim \mathcal{U}(1,10)$ ). Here, we have additionally included the Helmholtz PDE, a challenging system due to the highly oscillatory behavior of the solutions (see PDE 2 in Fig. 1 and Fig. A.1, for examples), very sensitive to the range of wavenumbers. When pre-training a single model on this "mixed" dataset, we simply use zero channels for those coefficients that do not exist when using examples from a specific operator. For example, the Helmholtz equation has a diffusion tensor input (identity matrix) with an additional input for the wavenumber but no advection (zero channel), while the Poisson's equation only has a diffusion tensor input and hence we append zero channels to signify no wavenumbers and advection; similarly for Advection-Diffusion. This, effectively, serves
(a) SYS-1(1,2.5): pre-training using SYS-1(1,5) and mixed dataset
(b) SYS-3(1,5): pre-training from SYS-3(1,10) and mixed dataset
Figure 5: Addressing (Q4). Testing error as a function of downstream examples for SYS-1, SYS-2, and SYS-3 with fine-tuning from their respective PDE systems and from the mixed dataset (combination of SYS-1, SYS-2, and SYS-3). The model pre-trained on the mixed dataset performs better than training from scratch. More importantly, the same pre-trained model yields low errors on all the downstream PDEs with both zero-shot and task-specific fine-tuning. We show the results for SYS-2 and OOD performance in §B.2.
as selection of the solution operator during the forward pass to predict the solution to the right operator. While more advance techniques such as in-context prompting (from LLMs) exist, here we are interested in understanding if this simple and minimal selection/prompting is sufficient for the model to transfer effectively to downstream tasks. For the downstream tasks, we consider three within-distribution tasks of Poisson's SYS-1(1,2.5), Advection-Diffusion SYS-2(0.2,0.4), and Helmholtz SYS-3(1,5) and show our dataset scaling results in Fig. 5.
The results support our most compelling conclusion: fine-tuning from the mixed dataset retains the substantial performance gains over training from scratch for all downstream tasks. The same model (pre-trained on three different tasks) is useful in all downstream tasks, in both the zero-shot and the fine-tuning settings. This indicates the input coefficient channels are sufficient to prompt the model to predict the correct downstream solution. We show the OOD downstream performance of this model in Appendix §B.2 and observe similar behavior.
We have provided an extensive analysis of the scaling and transfer behavior of neural operator models on multiple PDE systems. This involved characterizing behavior as a function of model size, downstream dataset size, underlying physics of the downstream tasks in relation to pre-training, the adaptation of these models to multiple downstream PDEs, and how all these behaviors scale with relevant problem and model parameters. Among other things, we have shown that it is possible and beneficial to develop more general SciML models capable of solving multiple tasks with the same set of weights, even when downstream tasks involve small-to-moderate distribution shifts relative to the pre-training data. All in all, this demonstrates the potential of the "pre-train and fine-tune" paradigm for SciML problems, paving a path towards building SciML foundation models. Moving forward, many questions remain. These include further exploration of model architectures (balancing expressivity and flexibility), pre-training protocols (including self-supervision components at scale), fine-tuning strategies, and the integration of these within a specific compute (and memory) envelope. There are a number of future directions raised by our work. Our pre-train and fine-tune recipe may need more sophisticated prompting during inference, especially if only the operator form changed between two different PDE systems. Also, we do not look at self-supervision with the PDE loss penalty as a means of large-scale pre-training, and we limit our analysis to 2D spatial systems. Moving to other architectures, larger scales, and more complex PDEs in space-time is a focus of future work.
This work was partially supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) program, under Contract Number DE-AC02-05CH11231 at Lawrence Berkeley National Laboratory. This research used the Perlmutter supercomputing resources of the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory, operated under Contract No. DE-AC02-05CH11231. We acknowledge gracious support from Google Cloud, Google TRC team, and specifically Jonathan Caton, and Prof. David Patterson. Prof. Keutzer's lab is sponsored by Intel corporation, Intel VLAB team, Intel One-API center of excellence, as well as funding through BDD and BAIR. Amir Gholami acknowledges support from Samsung SAIT and Intel corporations. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
[1] https://github.com/ShashankSubramanian/neuraloperators-TL-scaling, 2023.
[2] Kaushik Bhattacharya, Bamdad Hosseini, Nikola B Kovachki, and Andrew M Stuart. Model reduction and neural networks for parametric pdes. The SMAI journal of computational mathematics, 7:121-157, 2021.
[3] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258, 2021.
[4] Nicolas Boullé and Alex Townsend. Learning elliptic partial differential equations with randomized linear algebra. Foundations of Computational Mathematics, 23(2):709-739, 2023.
[5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[6] Ayan Chakraborty, Cosmin Anitescu, Xiaoying Zhuang, and Timon Rabczuk. Domain adaptation based transfer learning approach for solving pdes on complex geometries. Engineering with Computers, 38(5):4569-4588, 2022.
[7] Souvik Chakraborty. Transfer learning based multi-fidelity physics informed deep neural network. Journal of Computational Physics, 426:109942, 2021.
[8] Xinhai Chen, Chunye Gong, Qian Wan, Liang Deng, Yunbo Wan, Yang Liu, Bo Chen, and Jie Liu. Transfer learning for deep neural network-based partial differential equations solving. Advances in Aerodynamics, 3(1):1-14, 2021.
[9] Maarten De Hoop, Daniel Zhengyu Huang, Elizabeth Qian, and Andrew M Stuart. The cost-accuracy trade-off in operator learning with neural networks. arXiv preprint arXiv:2203.13181, 2022.
[10] Shaan Desai, Marios Mattheakis, Hayden Joy, Pavlos Protopapas, and Stephen Roberts. One-shot transfer learning of physics-informed neural networks. arXiv preprint arXiv:2110.11286, 2021.
[11] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[12] C. Edwards. Neural networks learn to speed up simulations. Communications of the ACM, 65(5):27-29, 2022.
[13] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. Domain-adversarial training of neural networks. The journal of machine learning research, 17(1):2096-2030, 2016.
[14] Yihang Gao, Ka Chun Cheung, and Michael K Ng. Svd-pinns: Transfer learning of physics-informed neural networks via singular value decomposition. arXiv preprint arXiv:2211.08760, 2022.
[15] Somdatta Goswami, Cosmin Anitescu, Souvik Chakraborty, and Timon Rabczuk. Transfer learning enhanced physics informed neural network for phase-field modeling of fracture. Theoretical and Applied Fracture Mechanics, 106:102447, 2020.
[53] Yan Yang, Angela F. Gao, Jorge C. Castellanos, Zachary E. Ross, Kamyar Azizzadeneshli, and Robert W. Clayton. Seismic wave propagation and inversion with neural operators, 2021.
[54] Jason Yosinski, Jeff Clune, Yoshua Bengio, and Hod Lipson. How transferable are features in deep neural networks? Advances in neural information processing systems, 27, 2014.
[55] Min Zhu, Handi Zhang, Anran Jiao, George Em Karniadakis, and Lu Lu. Reliable extrapolation of deep neural operators informed by physics or sparse observations. arXiv preprint arXiv:2212.06347, 2022.
Figure A.1: Visualization of the source function sampling (left) and the effect of certain PDE coefficients (right) on the solutions for the different systems. On the left, as we go down, the sparsity of Gaussians is increasing, leading to more sparse and spread out source functions encouraging heterogeneity in the dataset. For one of these source functions, we apply the different PDE operators with varying ranges of certaino PDE coefficients to illustrate their effect on the solutions. On the top row, for SYS-1 (Poisson's), we show that by increasing the diffusion tensor eigenvalue  $e$  (but keeping the direction  $\theta$  fixed), we increasing anisotropy and diffusion as we move towards the right. In the middle, we increase the velocity scales for SYS-2 (Advection-Diffusion), but keep the diffusion tensor and velocity direction the same, to demonstrate the increasing competing advection and diffusion processes as we go right. Finally, at the bottom, we show the highly oscillatory behavior in SYS-3 (Helmholtz) as we increase the wavenumber  $\omega$ . Note the significant differences between the solutions of the different systems.
A Appendix: Additional Details
A.1 Pre-train and downstream data creation
As described in §3, our sampling strategy involves (i) source functions and (ii) PDE coefficients. We use a numerical discretization of  $128 \times 128$  for all our experiments. Some additional details are as follows:
(i) Source function sampling: As described in §3, our source functions are 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 1/32 (hence 4 times the spatial resolution) to encourage high variability. The spacing between Gaussians is fixed at 2σ. 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}$  (we define  $s$  as the fraction of number of zero components of  $p_i$ ). Hence,  $s = 0.6$ , implies only 40% of  $\pmb{p}$  is randomly sampled in  $\mathcal{U}(0,1)$  and the rest (60%) are set to zero. Visualizations of this are in Fig. A.1 on the left. As we go down, the sparsity is increased (and hence Gaussians are more sparse and spread apart). In our pre-training and downstream datasets, we sample sparsity levels from 20% to 80% uniformly.
(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. We visualize the effect of  $e$  in Fig. A.1 (top row). With a fixed  $\theta$  (direction of diffusion), we see that with larger  $e$ , the solution is more anisotropic and diffuse. The sampling of  $e$  and  $\theta$  in a systematic fashion
Figure A.2: We illustrate the variability in source function inputs (left) and solution outputs (right) for SYS-2, where the velocity direction, scales, and the diffusion tensor direction and anisotropy scales (eigenvalue) are all changed along with the source sampling to produce the input-output pairs for the training dataset.
quantifies the underlying diffusion physics. In SYS-2, we additionally also sample the velocity vector  $\mathbf{v}$  direction from  $\mathcal{U}(0,2\pi)$ . We define  $\Psi = \| \pmb {v}\cdot \nabla u\| /\| \mathrm{div}K\nabla u\|$ , the ratio of advection to diffusion to quantify the different processes and  $\Psi$  is changed by scaling the velocity. In Fig. A.1 (middle row), we visualize increasing advection (from left to right) when the diffusion tensor is kept the same—observe that the solution changes significantly as  $\Psi$  increases and the two processes (advection and diffusion) compete more strongly. In SYS-3, we sample the wavenumber  $\omega$  as uniform integers. We visualize the effect of increasing  $\omega$  in Fig. A.1 (bottom row) and observe that increasing frequency leads to highly oscillatory behavior. We also underscore the significant differences in the output of the three systems here. For more details, see Appendix §A.1. We visualize the sampling process and resulting solutions with different ranges of underlying physics in Fig. A.1. In Fig. A.2, we show a small sample of input sources and output solutions for SYS-2 to illustrate the variability in our datasets—note here all the variables (advection directions and scales, diffusion directions and scales, source functions) are changed to create the full dataset.
For each PDE system, we use spectral methods to compute the ground truth solution of the PDE. For example, for SYS-1:
$$
- \operatorname {d i v} K \nabla u = f, \tag {4}
$$
the solution is:
$$
u = \mathcal {F} ^ {- 1} \left(\frac {- 1}{- \left(k _ {x} ^ {2} k _ {1 1} + k _ {y} ^ {2} k _ {2 2} + 2 k _ {x} k _ {y} k _ {1 2}\right)} \mathcal {F} (f (\boldsymbol {x}))\right), \tag {5}
$$
where  $\mathcal{F}$  is the Fourier transform,  $k_x, k_y$  are the frequencies in the Fourier domain, and  $k_{11}, k_{22}, k_{12}$  are the diagonal and off-diagonal coefficients of the diffusion tensor  $\pmb{K}$ . The solution is unique given zero-mean source functions. Similar solutions can be derived for the other systems.
A.2 Model architecture
We describe the architecture details of the FNO here. The central building block of FNO learns a kernel integral operator parameterized by a function operating in Fourier space. By composing several of these blocks together in sequence, the FNO can learn to approximate the solution operator for complex PDEs. For Euclidean systems discretized at uniform resolution (e.g., in 2D with inputs and solutions having dimension  $\mathbb{R}^{h\times w}$ ), the implementation can be made efficient by using the
Fast Fourier Transform (FFT), denoted by  $\mathcal{F}$  with inverse FFT  $\mathcal{F}^{-1}$ . We refer the reader to the original FNO paper [30] for details, but briefly summarize the basic mechanism, to highlight how the model complexity depends on key hyperparameters. The inputs to the 2D FNO are spatial functions discretized at this resolution  $h\times w$  that represent sources, PDE coefficients, ICs/BCs of the PDE system--each of these is represented as a separate channel, leading to input tensors in  $\mathbb{R}^{h\times w\times c}$ , with  $c$  input channels. As mentioned before, we use the resolution  $h\times w = 128\times 128$ . Given this input tensor  $A\in \mathbb{R}^{h\times w\times c}$  ( $c$  input channels representing the PDE inputs), the FNO first projects  $A$  into a tensor  $X\in \mathbb{R}^{h\times w\times d}$  with embedding dimension  $d$ , which is passed through a series of FNO blocks. For a given block  $l\in \{1,\dots,L\}$  with input  $X^l\in \mathbb{R}^{h\times w\times d}$  the output at spatial index  $(i,j)$  is computed as
$$
X _ {(i, j)} ^ {l + 1} = \sigma \left(W _ {l} X _ {(i, j)} ^ {l} + \mathcal {F} ^ {- 1} \left[ \mathcal {K} _ {l} \left(\hat {X} ^ {l}\right) \right] _ {(i, j)}\right), \tag {6}
$$
where  $\sigma$  is a pointwise nonlinear activation function,  $W_{l}\in \mathbb{R}^{d\times d}$  is a learnable weight matrix, which performs pointwise linear transformation, and  $\hat{X}^l = \mathcal{F}(X^l)\in \mathbb{C}^{h\times w\times d}$  are complex-valued Fourier coefficients output by the FFT. The transformation  $\kappa_{l}$  in Fourier space is parameterized by complex weights  $\Phi^l\in \mathbb{C}^{d\times d\times m_1\times m_2}$  according to
$$
\mathcal {K} _ {l} \left(X ^ {l}\right) _ {\left(k _ {1}, k _ {2}\right)} = \Phi_ {\left(k _ {1}, k _ {2}\right)} ^ {l} X _ {\left(k _ {1}, k _ {2}\right)} ^ {l}, \tag {7}
$$
for pairs of Fourier frequencies/wavenumbers  $k_{1} \in \{1, \dots, m_{h}\}$  and  $k_{2} \in \{1, \dots, m_{w}\}$ . The hyperparameters  $m_{h}$  and  $m_{w}$  control the "mode cutoff", beyond which Fourier modes are ignored by  $\mathcal{K}$ , and they have a theoretical maximum of  $m_{h} = h / 2$  and  $m_{w} = w / 2$ , i.e., the Nyquist limit along each spatial dimension. In practice, the mode cutoffs are a key hyperparameter controlling model complexity along with the embedding dimension  $d$ , and they are often tuned to be less than the Nyquist limit to prevent overfitting and accelerate training. For square problems  $(h = w)$ , a symmetric truncation is adopted such that  $m = m_{h} = m_{w}$ . Thus, the per-layer parameter count is quadratic in both  $m$  and  $d$ , dominated by the complex weights  $\Phi$ . The above ( $d$  and  $m$ ) hyperparameters are the focus of our exploration in the model scaling experiments.
A.3 Training details and code open-source