SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For training, we use the Adam optimizer with a cosine learning rate decay schedule. All models are trained for 500 epochs with the best model saved at lowest validation loss. We tune batch size and initial learning rates using a grid hyperparameter search, and train every model using 4 NVIDIA A100 GPUs (using standard data-parallelism) on the Perlmutter supercomputer. A full training run (on all  $32K$  examples used for pre-training) completes in around 2.5 hours. To evaluate model performance, we use the mean relative error between predicted and target solution, defined as:  $\mu_{\ell_2} \coloneqq 1 / N\|u - u_0\|_2 / \|u_0\|_2$ , where  $u$  is the predicted solution,  $u_0$  is the target solution, and  $N$  is total (fixed) number of testing (around  $4K$ ) examples. We open source our model training and data generation code at [1].
A.4 Hyperparameter tuning
For studying the relationship between model and dataset size, we perform a simple grid-search over 5 different learning rates for each combination of model and dataset size. This is important because different dataset sizes may demand different hyperparameter values—for instance, we observe that for very small dataset sizes (small number of downstream examples), the tuned learning rates are significantly smaller (especially for zero- and few-shot transfer learning) to help mitigate over-fitting and optimization difficulties at that scale, whereas larger learning rates perform better as model and dataset sizes increase. Hence, for any result (including the model scaling), these learning rates are tuned with the best values picked from the validation set metrics.
A.5 Input normalization
As mentioned in §3, a key aspect of our training is the input normalization (for any PDE system). We first normalize every training example source and coefficient inputs with respect to a reference source value defined as the median source norm over the training dataset, i.e.,  $f_{\mathrm{ref}} \coloneqq \mathrm{median}_i(\| f_i \|_2)$  where  $i$  are the training examples. Hence for any example  $f_i$ , we first compute  $\| f_i \|_2$  and normalize  $f_i$  and all coefficients with the relative norm  $\| f_i \|_2 / f_{\mathrm{ref}}$ . First, this ensures the source norms are within a reasonable scale range. Additionally, it implies that scaling both coefficients as well as source
(a) SYS-2(0.2,0.4)
(c) SYS-2(1,2)
(b) SYS-2(0.4,1.6)
(d) SYS-2(2,5)
Figure A.3: Addressing (Q3). Testing error as a function of number of downstream dataset examples for different downstream tasks used in SYS-2. We show the extent of overlap between the pretrained and downstream dataset at the top using the range of advection-diffusion ratios of the two datasets. Similar to SYS-1, we observe good zero-shot performance that gradually decreases with distributional shifts but can be recovered through few-shot learning. With large shifts, the number of examples needed to reach desired error levels also increases.
functions of the inputs by a constant yields the same inputs to the neural operator. For example: the inputs  $f$  and  $K$  to the Poisson's equation (SYS-1) are equivalent to the inputs  $10f$  and  $10K$  both have the same solution function. The above normalization makes sure that these two input pairs are equivalent, since  $10f$  and  $10K$  simply get normalized by 10 before passing to the network. Then, to ensure the coefficient inputs are comparable scales and within a reasonable range, the coefficient values (that typically have very different scales depending on the physical process, up to  $100\times$  differences) are normalized by constant reference values pre-computed across the training dataset as the median values of the coefficients. We observe that without normalization the performance of FNO can be quite poor, especially in cases with multiple channels representing different physical scales (such as in advection-diffusion SYS-2).
B.1 TL behavior over underlying physics and model scaling
We show the TL performance under physics shift (Q3) for SYS-2 in Fig. A.3. Similar to SYS-1, we observe excellent TL performance for in-distribution tasks (see Fig. A.3a) that is independent of the downstream dataset size. As we systematically go OOD, we continue to observe good performance gains with TL (both zero-shot and few-shot) until we are significantly OOD as in Fig. A.3d.
In both SYS-1 and SYS-2, we find that TL from pre-trained models outperforms training from scratch with the number of downstream examples required to reach a given accuracy being orders of magnitude smaller with TL. To quantify this in detail, we estimate the amount of "data savings" for TL by interpolating the equivalent number of examples needed to reach a given model error when training from scratch. We plot these "from-scratch" data requirements as a function of number of
Figure B.1: Addressing (Q3). Equivalent data needed to reach the accuracy of TL when training from scratch, for SYS-1 and SYS-2 distribution shifts defined in Table 1. Data from experiments where TL outperforms even the best (i.e., largest dataset) from-scratch models is not plotted.
Figure B.2: Some additional results: (left) for extreme shifts and challenging systems, the number of examples to reach a desired accuracy increases. For example, for TL to SYS-1(10,20), the training from scratch and TL curves intersect at 128K examples. (right) model parameter scaling for TL to SYS-3(1,5) showing similar trends as other systems.
TL fine-tuning examples in Fig. B.1. In these plots, any points above the dashed black line indicate an improvement over training from scratch; points for TL experiments which outperform even the best from-scratch models (which use a maximum of  $32K$  training examples) are not plotted. We observe that even as the downstream dataset distribution shift moves from "None" to "Large" (as defined in Table 1), the TL models consistently reach the performance of from-scratch models which would otherwise require thousands to tens of thousands of training examples. One exception is the largest OOD shift for SYS-1, where we observe smaller (but still substantial) advantages for the TL approach in zero- and few-shot learning. We note that characterizing distribution shifts consistently for different systems is challenging; this result suggests perhaps the "Large" distribution shift for SYS-1 is a more significant shift than that of SYS-2.
Finally, we include two additional results: For extreme shifts such as SYS-1(10,20), we note that the errors are higher as its is a more challenging system (high anisotropy in diffusion, in this case) and hence needs more data to achieve lower errors. In Fig. B.2 (left) we show that the training from scratch and TL from pre-trained model curves converge at a much larger dataset of 128K examples. We also include the model scaling experiments for a Helmholtz system SYS-3(1,5) in Fig. B.2 (right).
B.2 TL behavior underlying multiple operators
We show additional results for TL behavior over multiple operators in Fig. B.3. Each row represents a PDE system with in-distribution (left) and OOD (right) downstream tasks. We observe that the mixed pre-trained model is able to retain TL gains as the system-specific pre-trained model over all the systems. Hence, a single model is now adapted to three different PDE systems in downstream fine-tuning. We also note SYS-3 OOD results—we observe that for the Helmholtz, going even moderately OOD, affects both TL and training from scratch performance. This PDE system is
Downstream Pretrain SYS-1(1,2.5) SYS-1(1,5) 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 5.0 Diffusion tensor eigenvalue
(c) SYS-2(0.2,0.4): pre-training using SYS-2(0.2,1) and mixed dataset
(e) SYS-3(1,5): pre-training using SYS-3(1,10) and mixed dataset
(b) SYS-1(2.5,7.5): pre-training using SYS-1(1,5) and mixed dataset
(d) SYS-2(0.4,1.6): pre-training using SYS-2(0.2,1) and mixed dataset
(f) SYS-3(2,12): pre-training using SYS-3(1,10) and mixed dataset
Figure B.3: Addressing (Q4). Testing error as a function of downstream examples for SYS-1 (top), SYS-2 (middle), and SYS-3 (bottom) with fine-tuning from their respective PDE systems and from the mixed dataset (combination of SYS-1, SYS-2, and SYS-3). In each row: on the left, the downstream task is in-distribution and, on the right, it is OOD. 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 also note that SYS-3 is a particularly challenging system that shows larger performance drops as we go OOD—the mixed pre-training still retains the same gains as the task-specific pre-trained model.
particularly challenging and is very sensitive to the input wavenumber. Slight changes introduce more oscillatory features in the solution, pushing the downstream task OOD easily. However, we note that both the mixed pre-training and pre-training from a Helmholtz dataset show similar performance even on the OOD task.
B.3 Sensitivity to random seeds
To quantify the variability of both training from scratch and performing TL from the pre-trained model, we repeat the "Med" OOD shift experiment for SYS-1 and SYS-2 (see Table 1) 5 times with different random seeds and record the testing error for each trial. The resulting distributions indicate how sensitive each approach is to the random data shuffle used when training on the downstream dataset. We plot the mean testing error along with  $1^{st}$  and  $3^{rd}$  quartiles for each downstream dataset size in Fig. B.4. We observe small sensitivity for both TL and training from scratch. Further, not surprisingly, we observe that the variability across random seeds is generally larger when training from scratch. We note that the small variability in TL is difficult to see due to the log-spaced y-axis.
(b) SYS-2(1,2) pre-trained from SYS-2(0.2,1)
Figure B.4: Testing error as a function of downstream examples for SYS-1 and SYS-2, aggregated over 5 trials with different random seeds for each experiment. The dots indicate the mean testing error at each downstream dataset size, and the shaded region represents the spread between the  $1^{st}$  and  $3^{rd}$  quartiles. We observe small variance that is slightly larger when training from scratch compared to TL from the pre-trained model.