PaperlayerView exact revision on arXiv

Source-linked AI summary

Multiple Physics Pretraining for Physical Surrogate Models

Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, Mariel Pettee, Tiberiu Tesileanu, Kyunghyun Cho, Shirley Ho

arXiv:2310.02994v2cs.LGcs.AIstat.ML

TL;DR

Low-data settings make it difficult to train physical surrogates conventionally, motivating models that can transfer across systems. MPP pretrains a transformer to predict heterogeneous physical dynamics in a shared representation; it matches or surpasses task-specific baselines without finetuning and shows transfer to low-data systems with unseen physics or higher dimensions.

  • Problem

    In simulation-driven exploration and design, training data can be difficult to acquire, while conventional surrogates often focus on specific systems where data is readily available.

    Source: §1 · p. 2

  • Method

    MPP uses autoregressive pretraining to predict multiple heterogeneous physical systems simultaneously in a shared embedding space.

    Source: §1 · p. 2

  • Results

    MPP models match or surpass modern task-specific baselines without task-specific finetuning, and finetuned models show stronger performance on the evaluated benchmarks.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    MPP models transfer to low-data systems with previously unseen parameter regimes and higher dimensions.

    Source: §1 · p. 2

    The architecture is limited to uniform grids, and transfer beyond the largely fluid-like PDEBench data distribution remains uncertain.

    Source: §6 · p. 10

Abstract

from arXiv · show

We introduce multiple physics pretraining (MPP), an autoregressive task-agnostic pretraining approach for physical surrogate modeling of spatiotemporal systems with transformers. In MPP, rather than training one model on a specific physical system, we train a backbone model to predict the dynamics of multiple heterogeneous physical systems simultaneously in order to learn features that are broadly useful across systems and facilitate transfer. In order to learn effectively in this setting, we introduce a shared embedding and normalization strategy that projects the fields of multiple systems into a shared embedding space. We validate the efficacy of our approach on both pretraining and downstream tasks over a broad fluid mechanics-oriented benchmark. We show that a single MPP-pretrained transformer is able to match or outperform task-specific baselines on all pretraining sub-tasks without the need for finetuning. For downstream tasks, we demonstrate that finetuning MPP-trained models results in more accurate predictions across multiple time-steps on systems with previously unseen physical components or higher dimensional systems compared to training from scratch or finetuning pretrained video foundation models. We open-source our code and model weights trained at multiple scales for reproducibility.

1 Introduction

Low-data physical simulations motivate task-agnostic models that can transfer across systems, while conventional surrogates often focus on systems where training data is readily available. MPP trains one model across heterogeneous physics, matching or surpassing task-specific baselines without task-specific finetuning.

1 Introduction

  • Low-data simulation and design settings motivate task-agnostic models with broad physical understanding for finetuning.Numerical simulation may be expensive or unreliable where acquiring training data is also difficult.Source: §1 · p. 2
  • MPP embeds heterogeneous physical systems in a shared space and autoregressively predicts their dynamics simultaneously.The approach is presented as task-agnostic pretraining for physical surrogate models.Source: §1 · p. 2
  • MPP models match or surpass modern task-specific baselines without task-specific finetuning.The paper studies transfer using fluid-oriented benchmarks.Source: §1 · p. 2
  • MPP models transfer to low-data systems with previously unseen parameter regimes and higher dimensions.The new parameter regimes generate notably different qualitative behavior.Source: §1 · p. 2
  • The authors release code and pretrained models at multiple sizes for community experimentation.Source: §1 · p. 2

2 Background

The paper frames physical surrogate modeling as predicting spatiotemporal dynamics from snapshot histories. It uses autoregressive next-step prediction without assuming known PDE forms or parameters.

2 Background

  • The pretraining task learns one model to predict system dynamics from a uniformly spaced history of snapshots.The notation describes a system with a single state variable in one spatial dimension for simplicity.Source: §2 · p. 2
  • Autoregressive next-step prediction is a natural pretraining objective because next-step prediction is often a primary goal in dynamical-system surrogates.Source: §2 · p. 3
  • MPP does not assume a known PDE form; the model must infer parameter impacts from the provided history.The stated goal is applicability to both observed and simulated dynamics.Source: §2 · p. 3
  • Shared principles such as diffusion, advection, and reactivity across PDEs suggest features may apply to multiple physical systems.Source: §2 · p. 3

3 Related Work

Related work spans scientific surrogate modeling, pretrained models, and transfer across physical systems. MPP is positioned as pretraining across multiple operators, with the authors noting that its fluid-like data is not diverse enough to establish a foundation model.

3 Related Work

  • The authors describe their work as the first large-scale implementation of successful multiple nonlinear physics pretraining for spatiotemporal systems, while stating its data is insufficiently diverse to call the models foundational.Source: §3 · p. 3
  • Scientific machine-learning approaches include PDE-solution networks, learned solver augmentations, and data-driven solution operators; MPP focuses on pretraining across multiple operators.The solution-operator family is described as most similar to MPP, especially Cao (2021), which also uses a transformer-based architecture.Source: §3 · p. 3
  • Prior transfer work across operators includes Subramanian et al. (2023), whose study is limited to steady-state linear systems with periodic boundary conditions.Source: §3 · p. 3
  • Other multiple-system studies have examined restricted classes or low-dimensional, low-resolution systems.Source: §3 · p. 4

4 Scalable Multiple Physics Pretraining

The section tests whether pretraining on partially shared physics transfers to a combined system, then develops an architecture and training strategy for learning across heterogeneous physical systems.

4.1 Compositionality and Pretraining

  • To isolate transfer from shared physics, the authors test constant-coefficient linear advection-diffusion on a periodic 1D domain with one shared scalar field.The motivating hypotheses are that one model can learn multiple physical behaviors and that partially overlapping physics can benefit transfer.Source: §4.1 · p. 4
  • Pretraining on advection and diffusion separately achieved lower error with fewer advection-diffusion samples, without seeing both physics together during pretraining.The comparison spans 16–100K available examples.Source: Figure 1 · p. 4, §4.1 · p. 4

4.2 Architecture

  • AViT uses fully axial attention to improve scalability in higher dimensions, favoring scalability over maximizing accuracy.Attention proceeds across axes in turn, reducing the stated cost from O((HWT)^2) to O(H^2 + W^2 + T^2).Source: §4.2 · pp. 4–5
  • RevIN normalizes each channel over space and time, while system-specific field projections map inputs into a shared embedding space.Saved input statistics are used to denormalize predictions; new field types can be added by initializing new embedding channels.Source: §4.2 · pp. 5–6
  • For periodic boundaries, relative position distances account for neighbors across the boundary, improving generalization when learning periodic and non-periodic conditions together.The paper otherwise expects the model to infer boundary conditions from the input history.Source: §4.2 · p. 6

4.3 Balancing Objectives During Training

  • Uniformly sampling one system per micro-batch and accumulating gradients reduces workload variance across GPUs without requiring differently shaped systems in one batch.The authors report that this random scheme was sufficient and incurred no additional cost with gradient accumulation.Source: §4.3 · p. 6
  • Normalized MSE partially limits large-scale systems from dominating updates, and the authors report strong results across multiple systems despite remaining differences in task difficulty.The objective does not account for all variation in difficulty, including systems that converge at different rates.Source: §4.3 · p. 6

5 Experiments

The experiments test whether MPP can learn multiple physical systems simultaneously, transfer to low-data systems with new physics, and extend to 3D. Results show strong pretraining performance, advantages in low-data transfer, and improvements when 2D-pretrained models are inflated to 3D.

5 Experiments

  • The experiments use PDEBench’s 2D simulations of four nonlinear PDE systems and evaluate simultaneous prediction, unseen-physics transfer, and transfer relative to spatiotemporal foundation models.Training uses one-step-ahead autoregressive prediction; the VideoMAE comparison constrains the time dimension to 16.Source: §5 · pp. 6–7, §5.1 · p. 7

5.1 Pretraining Representations

  • MPP-pretrained models achieve high-end performance across datasets and improve with scale, while the smallest model outperforms PDEBench baselines except on SWE.The pretrained AViT-B catches up to the 115M-parameter FNO-B, and the authors conclude that large transformers can learn multiple dynamics simultaneously, though multitask training can hurt individual-task performance.Source: §5.1 · pp. 7–8, Table 1 · p. 7

5.2 Transfer to Low-data Domains

  • The low-data transfer study holds out all compressible-fluid data during pretraining, then fine-tunes on compressible Navier-Stokes datasets labeled “near” and “far.”“Near” has low Mach number and viscosity similar to incompressible pretraining; “far” exhibits markedly different turbulent behavior and unseen small-scale structure, alongside other simulation differences.Source: §5.2 · p. 8, Figure 4 · p. 8
  • MPP models outperform VideoMAE and training from scratch by a large margin in the low-data regime.Physical pretraining appears to provide an even larger short-term advantage on the turbulent “far” dataset than on the smoother “near” data, although both approaches are less accurate over long rollouts on “far.”Source: Figure 3 · p. 8, §5.2 · p. 9

5.3 Inflation to 3D

  • For 3D inflation, the study uses full pretrained models at the Ti scale and compares inflated models with scratch-trained models and PDEBench baselines across a substantial 2D-to-3D physics gap.The gap includes known behavioral differences between 2D and 3D turbulence; the table reports NRMSE with initial-condition type as its subheadings.Source: §5.3 · p. 9, Table 2 · p. 9
  • 11.7% improvement on the smaller randomly initialized dataset and 4.1% on the larger turbulently initialized dataset demonstrate gains from 2D-to-3D inflation.The comparison uses the PDEBench compressible Navier-Stokes datasets; 3D high-resolution training is described as enormously expensive.Source: §5.3 · p. 9, Table 2 · p. 9

6 Conclusion

MPP is an autoregressive, multi-system pretraining strategy for physical surrogates that shows transfer to unseen physics and 3D simulations. The authors identify limits in geometry, data cleanliness, and transfer scope, and note a possible weapons-research risk.

6 Conclusion

  • MPP models matched dedicated baselines on fluid and reaction benchmarks from pretraining alone, with stronger performance after task-specific finetuning.The authors also report positive transfer to unseen physics and improved 3D results over training from scratch.Source: §6 · p. 10
  • The architecture is limited to uniform grids, while extending it to non-uniform geometries requires further work balancing efficiency and flexibility.The authors also describe 3D extension as enormously expensive for many model families and identify conservation-law constraints as an open challenge across closed and open systems.Source: §6 · p. 10
  • The evaluated data is clean and fully observed, so generalization to incomplete or noisy observations remains unknown.Some PDEBench tasks omit relevant physical fields, such as velocity fields for shallow water, while available data is observed at every grid point without noise.Source: §6 · p. 10
  • Transfer limits remain poorly understood because PDEBench is largely fluid-like, and the distance between physics tasks is not yet defined.The authors call for more diverse data to investigate how far transfer benefits extend from the training distribution.Source: §6 · p. 10

A Impact Statement

  • The authors identify potential weapons-research use as a risk, but consider it fairly small because the approach targets lower-precision modeling of poorly understood physics.They hope the approach will help researchers and engineers develop more accurate dynamics models from limited observations.Source: A Impact Statement · p. 17

B Data Details

The appendix describes the PDEBench and PDEArena simulation data, model configurations, and pretraining setup. The study pretrains on multiple 2D PDEBench systems using uniformly sampled tasks and trajectories.

B.1 PDEBench

  • PDEBench supplies time-dependent and time-independent PDE simulations, and the study focuses on 2D fluid dynamics with velocity, pressure, and density fields.The appendix describes compressible and incompressible Navier–Stokes, shallow water, and diffusion-reaction data, with varying parameters, initial conditions, and boundary conditions.Source: §B.1 · pp. 17–18
  • PDEBench includes compressible Navier–Stokes datasets with 1,000 trajectories at 512 × 512 resolution and 10,000 at 128 × 128; incompressible simulations also include an immersed tracer.The 512 × 512 set has 100 time steps, while the 128 × 128 set uses the same time-step and initialization description; shallow-water data uses 1,000 trajectories at 128 × 128 over 100 steps.Source: §B.1 · pp. 17–18

B.2 PDEArena

  • PDEArena adds 2D incompressible Navier–Stokes simulations with spatially constant buoyancy and 5,200 training trajectories, plus 1,300 each for validation and test.The data uses a 128 × 128 grid and 16 prediction timesteps, with viscosity 0.01 and Dirichlet boundary conditions.Source: §B.2 · p. 18

C.1 Model Configurations

  • AViT model sizes scale multiple architecture dimensions together, and the appendix lists model configurations and inference times on an A6000 GPU.The models use pre-norm, instance normalization, GeLU activations, QK normalization, hMLP patching, and a transposed hMLP decoder; their inference is reported as slightly slower on average than similarly sized baselines.Source: §C.1 · pp. 18–19, Table 3 · p. 19, Table 4 · p. 19

C.2 Exp 1: Pretraining Performance

  • MPP pretraining uses all 2D time-dependent PDEBench datasets and trains for 200K steps with uniform sampling of tasks and trajectories.The experiment includes compressible and incompressible Navier–Stokes, 2D diffusion-reaction, and shallow-water data; each dataset is split 0.8/0.1/0.1 by trajectory.Source: §C.2 · pp. 19–20

C.3 Experiment 2: Transfer to Low-Data Domains

This appendix documents transfer experiments to low-data compressible Navier–Stokes tasks and the adaptations used for MPP and VideoMAE. It specifies data splits, model adaptation, and VideoMAE finetuning choices, but these passages do not report comparative results.

C.3 Experiment 2: Transfer to Low-Data Domains

  • The experiment compares MPP-pretrained models with general-purpose pretrained VideoMAE models for frame prediction on PDEBench data.MPP pretraining uses incompressible Navier–Stokes, shallow water, and 2D diffusion-reaction; transfer targets the “Near” and “Far” compressible Navier–Stokes datasets with periodic boundaries.Source: §C.3 · p. 20

C.3.1 MPP Finetuning Procedure

  • MPP finetuning expands field projection and reconstruction layers when the target adds previously unseen state variables, while loading other weights normally.For compressible Navier–Stokes, the pretraining data contains vx, vy, and ρ, while the target also includes pressure; the full model is finetuned because partially frozen training performed noticeably worse.Source: §C.3.1 · p. 21

C.3.2 VideoMAE Settings

  • The VideoMAE comparison uses ViT-base models pretrained on Kinetics-400 or Something-Something V2, with input and output layers adapted to PDE data.The appendix describes adapting channel counts by repeating pretrained weights and positional encodings by trilinear interpolation.Source: §C.3.2 · p. 21

C.3.3 Video MAE Finetuning Procedure

  • VideoMAE is finetuned to predict the next 2 frames (Tp = 2) of a video from a 16-frame context (Tc = 16), masking the last two frames and minimizing NMSE loss.Targets are channel-wise normalized by subtracting the global means of corresponding context frames and dividing by their global standard deviations; Table 5 reports effective learning rates for finetuning.Source: §C.3.3 · pp. 21–22, Table 5 · p. 22

C.4 Experiment 3: Inflation to 3D

The appendix specifies training settings and datasets for the 3D inflation experiments. It identifies two PDEBench compressible Navier–Stokes datasets, but the supplied passages do not state their comparative results.

C.4 Experiment 3: Inflation to 3D

  • For MPP and training from scratch, the setup uses 100 true epochs, a .8/.1/.1 train/validation/test split, accumulation steps of 1 (no accumulation), Adan, and DAdaptation for base learning rate.Source: §C.4 · p. 22
  • The 3D experiments use two PDEBench compressible Navier–Stokes datasets, each simulated for 21 steps with shear and bulk viscosities of 10^-8.One dataset has 600 trajectories with turbulent initialization; the other has 100 trajectories with random initialization.Source: §C.4 · p. 22
  • Table 9 is identified as a validation NRMSE comparison of position bias across training data that differs only in boundary conditions.The supplied passages do not provide the table values or its outcome.Source: Table 6 · p. 23

C.5 Appendix Experiment: Broader Usage of Pretrained Representations

This appendix experiment reports the training settings used for MPP and training from scratch.

C.5 Appendix Experiment: Broader Usage of Pretrained Representations

  • Training ran for 20K optimization steps.The listed settings also specify a batch size of 24 and one accumulation step.Source: §C.5 · p. 23
  • The train/validation/test split used 1000/100/1000 examples, drawn from the original validation set or randomly depending on whether data was used for training.Source: §C.5 · p. 23
  • The optimizer was Adan, with weight decay 1E-3, drop path 0.1, a DAdaptation base learning rate, cosine decay, and gradient clipping at 1.0.Source: §C.5 · p. 23

D.1 Position Bias Evaluation

The position-bias evaluation compares joint training on advection systems with periodic and absorbing boundaries. The modified position biases improve joint learning, especially for the periodic data.

D.1 Position Bias Evaluation

  • The two systems share a 1D domain and differ only in boundary conditions: periodic boundaries equate the endpoints, while absorbing boundaries let waves exit without reflection.The absorbing-boundary implementation extends the domain and solves periodic equations so exiting waves do not return.Source: §D.1 · p. 23
  • The joint-training experiment compares models with and without modified position biases after training individual-system baselines on 10,000 examples each for 100 epochs.The joint models train on 20,000 examples.Source: §D.1 · p. 23
  • Standard RPE struggles to learn the periodic baseline, while the periodic-adjusted variant learns periodic data more effectively but does not outperform the baseline.Both RPE schemes improve absorbing-boundary performance with the additional joint-training data.Source: §D.1 · p. 23
  • The periodic-only model reaches nearly an order of magnitude higher precision.Source: §D.1 · p. 23
  • The interaction between attention and boundary conditions is identified as an important direction for future study.The authors note that attention might avoid waves passing out of the domain, but describe this as an expectation rather than a demonstrated result.Source: §D.1 · p. 24

D.2 Normalization Ablations

The appendix reports ablations of normalization and loss balancing, then examines MPP representations on inverse parameter-estimation tasks. Results vary by task: pretraining helps forcing estimation but not scalar prediction.

D.2 Normalization Ablations

  • The proposed normalization and loss-balancing modifications greatly improve AViT-B’s ability to learn in the multiple-physics setting.The ablations use AViT-B as the baseline.Source: §D.2 · p. 24
  • Although normalized loss appears important for MPP, reporting RMSE instead of NRMSE does not significantly change the ordinal results on the smaller datasets.The RMSE comparison includes fewer results because not all sources report RMSE.Source: §D.2 · p. 24

D.3 Broader Usage of Pretrained Representations

  • The experiment compares pretrained MPP-AViT-B with a model trained from scratch on inverse estimation of two parameters, then evaluates representations on an additional incompressible Navier–Stokes variable-buoyancy benchmark.The buoyancy set was not used during training; it uses 1,000 randomly sampled trajectories for training, 100 for validation, and 1,000 for testing.Source: §D.3 · pp. 24–25
  • Pretraining nearly halves error on the forcing task but does not improve scalar prediction over training from scratch.On buoyancy, the scratch-trained AViT also barely outperforms a mean prediction.Source: §D.3 · p. 25

D.4 Exp1: Expanded Results

Expanded compressible Navier–Stokes results show that scaling effects vary with viscosity, while example trajectories develop patch artifacts before finetuning that mostly disappear afterward. The checklist also records reproducibility details and the absence of statistical-significance estimates.

D.4 Exp1: Expanded Results

  • The expanded CNS results report nominal diffusion coefficients rather than Reynolds numbers because the simulations’ relatively low resolution likely incurs significant numerical diffusion.Source: §D.4 · p. 25
  • The most viscous systems (η = .1) perform relatively worse with scale, while other viscosities generally benefit; the smallest model again loses performance.These are multi-task training results and cannot be directly interpreted in the single-task setting.Source: §D.4 · p. 25

D.5 Exp2: Numerical Results

  • Tables 12 and 13 provide numerical results corresponding to Figure 3 for near M0.1 CNS and far CNS, respectively.The supplied captions identify Table 12 as test NRMSE for near M0.1 with η = .01 and Table 13 as test NRMSE for far CNS.Source: §D.5 · p. 26, Table 12 · p. 26, Table 13 · p. 27

D.6 Pretraining Trajectories

  • Pretrained trajectories initially make strong predictions, but patch artifacts creep in over time.Source: §D.6 · p. 26

D.7 Finetuning Trajectories

  • After finetuning, the patch-based instability mostly disappears.Source: §D.7 · p. 27

NeurIPS Paper Checklist

  • The checklist reports that code is provided, data come from open established datasets, and experimental settings are listed in the appendix.It also states that experiments were too expensive to repeat enough times to gather useful statistical-significance estimates.Source: NeurIPS Paper Checklist · pp. 31–32

Related papers

Loading 2310.02994v2…