Source-linked AI summary
Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior
Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael Mahoney, Amir Gholami
TL;DR
The paper asks whether foundation-model-style pre-training can support transfer learning for SciML problems, where downstream data and physics may differ from pre-training. It evaluates neural operators across model and data scales, physics shifts, and multiple PDE operators, finding that appropriate fine-tuning often reaches target accuracy with far fewer examples than training from scratch. The results support pre-train-and-fine-tune as a potential route toward SciML foundation models, while the analysis is limited to FNO.
Problem
The paper investigates whether foundation-model-style pre-training and transfer learning can generalize to SciML problems across varied PDE tasks, scales, and physics distributions.
Method
The study pre-trains Fourier Neural Operators on diverse datasets spanning PDE coefficients, source functions, and multiple PDE systems, then evaluates zero-shot and fine-tuned transfer.
Results
Transfer learning outperforms training from scratch with limited downstream data, including moderately out-of-distribution physics, while fine-tuning gains more from model scaling.
Takeaways & Limitations
Pre-train-and-fine-tune has demonstrated potential for building SciML models that solve multiple downstream PDE tasks with shared weights.
Takeaways & Limitations
The analysis focuses only on the FNO model and needs to be extended to other architectures.
Abstract
from arXiv · showhide
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 pre-trained 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 behavior 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. We open-source our code for reproducibility.
1 INTRODUCTION
The paper investigates whether foundation-model-style pre-training and transfer learning can extend to SciML neural operators. It evaluates scaling, downstream data efficiency, physics shifts, and transfer across multiple PDE operators.
- Motivation: SciML models often require retraining for new PDE coefficients, sources, or initial and boundary conditions, motivating reusable neural operators.Neural operators instead learn mappings between function spaces from input-output examples.
- Study scope: The study analyzes FNO scaling and transfer across diverse PDE datasets, model sizes, downstream data sizes, training recipes, and OOD physics.The evaluated systems include Poisson’s, Advection-Diffusion, and Helmholtz PDEs.
- Method: Pre-training samples PDE coefficients and source functions across operators, predicts solver-generated solutions, and supports zero-shot or few-shot downstream adaptation.The same model can target in-domain or out-of-domain PDE systems after minimal fine-tuning.
- Downstream scaling: Orders of magnitude fewer downstream examples can achieve desired accuracy with transfer learning than with training from scratch.The comparison concerns zero-shot and few-shot adaptation under limited downstream data.
- Model scaling: Fine-tuned models gain more from parameter scaling than models trained from scratch as model size increases from 64K to 256M parameters.The study reports error saturation at small model sizes and monotonically decreasing error with larger models.
- Transfer behavior: Pre-trained models outperform scratch-trained models for in-distribution and moderately OOD physics, while a mixed-operator model retains gains across downstream PDE tasks.The mixed pre-training setup includes qualitatively different operators such as Poisson’s and Helmholtz.
2 RELATED WORK
Prior work established transfer-learning results for neural operators but generally examined limited dimensions in isolation. The paper distinguishes its analysis by jointly varying diverse PDE inputs, operators, model scale, and dataset scale.
- Existing transfer-learning studies: Prior SciML transfer-learning studies largely focused on PINNs, neural operators, or individual changes in target domains, PDE coefficients, or source functions.Related work also examined changing geometries, discretizations, OOD detection, and physics-informed fine-tuning.
- Research gap: Earlier studies did not systematically pre-train models on diverse combinations of PDE coefficients and source functions across different operators.Some prior work focused on geometric domains or reported accuracy saturation as fine-tuning examples increased.
- Research gap: The paper positions its contribution as simultaneous analysis of interpolation, extrapolation, model scaling, dataset scaling, and diverse pre-training corpora.This scope addresses dimensions that prior related work investigated only partially or separately.
3 METHODS
The study evaluates FNO-based pre-training and fine-tuning for PDE learning across diverse physics systems, source functions, coefficients, model sizes, and downstream distributions. It compares zero-shot and fine-tuned adaptation, including in-distribution and out-of-distribution tasks, with models trained from scratch.
- The method uses foundation-model transfer learning, pre-training a base model on diverse data before adapting it to downstream applications.
- PDE systems: The evaluation covers 2D Poisson, Advection-Diffusion, and Helmholtz systems with periodic boundary conditions.
- PDE systems: The Helmholtz system is included as a challenging case because its solutions can exhibit high-frequency oscillatory spatial patterns.
- Data generation: Training data vary source functions and PDE coefficients, including diffusion tensors, velocity vectors, and wavenumbers, to produce diverse and heterogeneous examples.
- Transfer experiments: Downstream adaptation includes zero-shot testing, few-shot fine-tuning with varying dataset sizes, and comparisons against training from scratch.
- Model scaling: Model scaling varies FNO embedding dimensions and Fourier modes, increasing parameter counts approximately 16× from 64K to 256M.
4 RESULTS
Across downstream-data, model-size, physics-shift, and multi-operator experiments, pre-training followed by fine-tuning consistently outperforms training from scratch, especially when downstream data are limited. The benefits persist under moderate and large distribution shifts and across mixed PDE systems, although they diminish with abundant downstream data and insufficient model capacity.
- Downstream dataset scaling: About 100× fewer downstream examples reach a desired 1e-2 error with fine-tuning than training from scratch in the moderate-OOD SYS-1 test.Fine-tuning requires about 64 examples versus 8K from scratch, and generally saves O(1K−10K) examples in few-shot settings.
- Model scaling: Fine-tuning gains increase with model size from 64K to 256M parameters, while 64K-parameter models show insufficient capacity.Both training recipes improve with scale, but the 256M fine-tuned model has the lowest errors across downstream example counts.
- Scaling limits: Transfer-learning benefits diminish in the large-data regime, while zero-shot transfer is optimal for downstream datasets within the pre-training distribution.The large-data regime contains O(2^15) examples, approximately matching the pre-training data volume.
- Underlying physics: As physics shifts from in-distribution to moderately and strongly OOD, zero-shot performance degrades and fine-tuning recovers performance more slowly.Even under large shifts, transfer learning improves over training from scratch, although the system may require more data because of greater anisotropy.
- Multiple solution operators: A single model pre-trained on Poisson, Advection-Diffusion, and Helmholtz data outperforms training from scratch across all three downstream PDE tasks.The mixed pre-trained model performs well with both zero-shot evaluation and task-specific fine-tuning, using coefficient channels to select the correct operator.
5 CONCLUSIONS
The paper finds that pre-train-and-fine-tune neural operators can support general SciML models across multiple PDE tasks, including small-to-moderate distribution shifts. It identifies remaining questions around architectures, pre-training, fine-tuning, and compute constraints.
- Pre-train-and-fine-tune neural operators can solve multiple PDE tasks with the same weights, including small-to-moderate distribution shifts.
- The study characterizes scaling and transfer behavior across model size, downstream dataset size, physics shifts, multiple PDEs, and relevant problem parameters.
- Building SciML foundation models remains an open direction requiring further work on architectures, pre-training protocols, fine-tuning strategies, and compute constraints.
A.1 Pre-train and downstream data creation
The datasets vary source functions and PDE coefficients to create diverse input-output pairs for pre-training and downstream evaluation. Source functions use Gaussian basis expansions, while coefficients are jointly varied across PDE systems.
- Dataset construction: 128 × 128 numerical discretization is used for all experiments, with source functions and PDE coefficients sampled as dataset inputs.
- Source function sampling: Source functions are linear combinations of Gaussian basis functions centered at grid points.The Gaussian width is fixed at 1/32, centers are spaced by 2σ, and coefficients are sampled from U(0, 1).
- Source function sampling: 20% to 80% sparsity levels are sampled uniformly to vary how many Gaussian coefficients are set to zero.A sparsity of 0.6 leaves 40% of coefficients randomly sampled and sets the remaining 60% to zero.
- PDE coefficient sampling: For pre-training, all PDE inputs are simultaneously varied to create a diverse training corpus.For SYS-2, diffusion anisotropy and direction, advection direction and scale, advection-diffusion ratios, and source-function parameters are varied.
- Solution generation: The PDE solutions are computed with spectral methods, including Fourier-domain expressions for the diffusion system.The diffusion solution uses Fourier transforms of the source and diffusion-tensor coefficients, with uniqueness given for zero-mean source functions.
A.2 Model architecture
The Fourier Neural Operator approximates complex PDE solution operators by composing Fourier-space kernel integral blocks. Its complexity is controlled by Fourier mode cutoffs and embedding dimension.
- FNO architecture: The FNO composes Fourier-space kernel integral operator blocks to approximate solution operators for complex PDEs.The implementation uses Fourier transforms for systems discretized at uniform resolution.
- FNO architecture: Each block combines a pointwise linear transformation with a learned Fourier-space transformation and a nonlinear activation.The Fourier coefficients are produced by the FFT, while complex weights parameterize the Fourier transformation.
- Model complexity: Fourier mode cutoffs determine which frequencies the learned kernel retains, with modes beyond the cutoff ignored.The theoretical maximum is the Nyquist limit, but smaller cutoffs can reduce overfitting and accelerate training.
- Input-output variability: Figure A.1 illustrates how varying velocity, diffusion-tensor, and source inputs produces diverse SYS-2 solution outputs.
- Model complexity: For square problems, symmetric mode truncation uses m = m_h = m_w, and per-layer parameter count is quadratic in mode cutoff and embedding dimension.
A.3 Training details and code open-source
Training uses Adam with cosine learning-rate decay, validation-based model selection, and tuned hyperparameters. Dataset-size and model-size experiments use learning-rate searches to mitigate overfitting on small datasets.
- Training procedure: Models are trained for 500 epochs with Adam, cosine learning-rate decay, and selection by lowest validation loss.Batch size and initial learning rates are tuned with grid search using four NVIDIA A100 GPUs.
- Compute and implementation: A full pre-training run on 32K examples takes around 2.5 hours on four NVIDIA A100 GPUs.
- Hyperparameter tuning: Five learning rates are searched for each model and dataset-size combination because smaller datasets require different hyperparameter values.Very small downstream datasets use significantly smaller learning rates, especially for zero- and few-shot transfer, to mitigate overfitting.
A.5 Input normalization
Input normalization scales source functions and PDE coefficients relative to a reference source norm, making proportionally scaled inputs equivalent while supporting transfer across shifted tasks.
- Normalization procedure: The reference source value is the median source norm over training examples, and source functions and coefficients are normalized by each example’s relative norm.This keeps source norms within a reasonable range and makes jointly scaled inputs equivalent.
- Transfer across shifts: The normalization supports transfer learning across downstream tasks whose PDE parameters are systematically shifted out of distribution.The study evaluates such shifts using different samplings of PDE coefficients and related task distributions.
- Scale invariance: For Poisson’s equation, f and K are treated equivalently to 10f and 10K because both input pairs normalize by the same factor.The normalization preserves their shared solution function under proportional scaling.
B.1 TL behavior over underlying physics
Transfer learning remains effective as downstream physics moves out of distribution, generally requiring far fewer examples than training from scratch, while performance varies for extreme shifts and across PDE systems.
- Physics shifts: TL for SYS-2 retains good performance gains under increasingly out-of-distribution physics, including zero-shot and few-shot transfer until shifts become substantial.In-distribution performance is strong and largely independent of downstream dataset size.
- Data savings: Orders of magnitude fewer downstream examples are needed with TL than from scratch to reach a given accuracy in SYS-1 and SYS-2.The comparison estimates equivalent from-scratch data requirements by interpolating the number of examples needed to match TL error.
- Data savings: Across shifts from None to Large, TL reaches performance that from-scratch models would require thousands to tens of thousands of training examples.The largest SYS-1 OOD shift yields smaller, though still substantial, zero-shot and few-shot advantages.
- Multiple operators: A mixed pre-trained model retains the transfer gains of system-specific pre-training across three PDE systems and adapts to each through downstream fine-tuning.For Helmholtz, moderate OOD shifts affect both approaches, while mixed and Helmholtz-specific pre-training perform similarly on OOD tasks.
B.3 Sensitivity to random seeds
Repeated experiments show small sensitivity to random seeds for both transfer learning and training from scratch, with greater variability generally occurring in the from-scratch models.
- Experimental design: Five random-seed trials measure testing-error distributions for SYS-1 and SYS-2 under the Med OOD shift.The experiments record mean testing error and the 1st-to-3rd quartile spread across downstream dataset sizes.
- Variability: Both TL and training from scratch show small sensitivity to random data shuffles across downstream dataset sizes.The comparison uses the mean testing error and quartile ranges from repeated trials.
- Variability: Variability across random seeds is generally larger for training from scratch than for TL from the pre-trained model.The low TL variability is difficult to see because the testing-error axis is logarithmically spaced.