Multiple Physics Pretraining for Spatiotemporal Surrogate Models
Authors: 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
Venue: NeurIPS 2024 poster (2024)
Keywords: transfer learning
physics
pretraining
finetuning
surrogate models
spatiotemporal
Primary area: machine_learning_for_physical_sciences
TL;DR: We develop approaches to enable autoregressive pretraining on multiple physical systems and show it can improve transfer performance across domain gaps.

Abstract: 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.

Multiple Physics Pretraining for Spatiotemporal Surrogate Models
Michael McCabe sl.2, Bruno Régaldo-Saint Blancard Liam Parker Ruben Ohana Miles Cranmer Alberto Bietti
Michael Eickenberg .Siavash Golkar?, Geraud Krawezik 1 Francois Lanusse 1.4 Mariel Pettee 1.5
Tiberiu Tesileanu 1, Kyunghyun Cho 6,7,8, Shirley Ho Ho 1,6,9
The Polymathic Collaboration
1 Flatiron Institute, 2 University of Colorado Boulder, Uniiversity of Cambridge, Université Sacis-Sacla,, Université Paris Cité, CEA, CNRS, AIM, Physics Division, Lawrence Berkeley National Laboratory, New York University, Prescient Design, Genentech, 8 CIFAR Fellow, 9 Princeton University
Abstract
We introduce multiple physies pretraining (MPP), an autoregressive task task-agnostic pretraining approach for physical surrogate modeling of spatiotemporal systems with transformers. In MPP, rather than training one model on specific physical system, train backbone model to predict the dynamics of multiple heteroge- neous 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 shared embedding and normalization strategy that projects the fields of multiple systems into shared embedding space. We validate the efficacy our approach on both pretraining and downstream tasks over a broad fluid mechanics-oriented benchmark. We show that a single MPP-pretraine, trans- former is able to match or outperform task-specific baselines on all pretraining sub-tasks without the need for finetuning. For downstream tasks, demonstrate that finetuning MPP-trained models results in more accurate predictions across multiple time-steps on systems with previously unseen physical components higher dimensional systems compared to training from scratch or finetuning pre- trained video foundation models. We open-source our code and model weights trained at multiple scales for reproducibility.
1 Introduction
In recent years, the fields of natural language processing and computer vision have been revolutionized by the success of large models pretrained with task- aas-osic objectives on massive, diverse datasets (Chen et al.) 2020) Devlin et al.] 20188: He et al.) 202T). This has. part, been driven by the development of self-supervised pretraining methods which allow models to utilize far more training data than would be accessible with supervised training (Balestriero et al. 2023). These so-called "foundation models" have enabled transfer learning on entirely new scales. Despite their task-agnostic pretraining, the features they extract have been leveraged as a basis for task-specific finetuning, outperforming supervised training alone across numerous problems especially for transfer to settings that are insufficiently data-rich to train large models from scratch (Bommasani et al. 202I).
" Contacc: mmccabefflatiiiniisttttt.t..
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Deep learning for computational science has begun to see first steps in this direction. Large domain- specific pretrained models have emerged in diverse fields such as chemistry (Bran et al. [023] Chithrananda et al. 2020), medicine (Jiang et al. 2023; Tu et al. 2023), astrophysics (Lenn & Bovy 2023; Nguyen et al. 2023b), and climate (Nguyen et al. 2023a and the trend only seems to be
growing as more and more models are developed for new fields both as refined versions of existing
large language models and as new models trained entirely on field-specific data.
In this work, we demonstrate that similar approaches can be extended to the surrogate modeling of spatiotemporal physical systems. Spatiotemporal prediction tasks, like those found in fluids,
solids, or general continuum mechanics, have attracted significant attention from the deep learning
community, From direct prediction methods (Dang al. [022] Li et al. al.] 2202 Lusch al., 2018 Plaff et al. [021]; Stachenfeld et al. (022) to neural PDE solvers (Bruna et al.
2022, Raissi et al.
2019), researchers have sought to develop fast, accurate models for physies either as faster surrogates
for the partial differential equation (PDE) solvers that dominate the field or to simulate systems that cannot be exactly described or resolved by current mechanistic models and available hardware. While directly outperforming PDE solvers is difficult (Grossmann al. [0233). deep learning has already begun to impact fields like atmospheric science (Ben-Bouallegue et al. 2023 Bi et al. 2023 Lam
2023 Pathak et al.] 2022) and cosmology (Cranmer et al. 202 He et al. 2019: Jamieson et al.
et al.
2023). where the systems are too large or too imprecisely described to be simulated exactly.
Unfortunately, outside of few observation-rich outliers, settings where numerical simulation is expensive or unreliable also tend be settings where the difficulty of acquiring training data makes it impractical to train surrogates conventionally. Most deep learning- baaed surrogates thus far have focused on specific physical systems or families of parameterized PDEs where data can easily be acquired. However, for the low-data settings often found in simulation-drivn exploration and design, would be valuable to have large, task-agnostic models with a broad understanding of common physical behavior to act as a foundation for finetuning.
Contributions. To address this need, we introduce Multiple Physics Pretraining (MPP), new approach for task-agnostic pretraining of physical surrogate models. Our method enables large-scale pretraining for transfer across diverse physics which we study using fluid-oriented benchmarks. Our specific contributions are:
We develop MPP, a pretraining approach in which we embed multiple hetereogeneous physical systems into shared embedding space and learn to autoregressively predict the dynamics of all systems simultaneously.
We show that single transformer models pretrained with MPP are able to match or surpass modern task- specific baselines without applying task-specific finetuning to the MPP models.
We demonstrate the transfer capabilities of models trained with MPP to systems with limited training examples (referred to as low-data systems thereafter) displaying new physics in the form of previously unseen parameter regimes generating notably different qualitative behavior and inflated to higher dimensions.
We open-source our code and provide our pretrained models at a variety of sizes for the community to experiment with on their own tasks.
2 Background
Notation. Let S be an arbitrary physics-driven spatiotemporal dynamical system, either described by parameterized family of PDEs with fixed parameters, or where snapshots are gathered from observation of a unique physical phenomenon. To simplify notation, we discuss systems with single state variable in one spatial dimension. A continuous state variable for system S is represented as u (x, [0, Ls] [0, oo) + R. We discretize the system uniformly in space and time at resolutions Ns: Ts respectively. A snapshot uf TR:Ns represents the value of state variable uS at all Ns spatial discretization points time f. Our pretraining task is then to learn a single model M that can take a uniformly spaced sequence of T's snapshots U.S = =[uŞ_ "ts "t...... u. from system S sampled from some distribution over systems and predict M(US) such that M(U u ui+Ats
Autoregressive Pretraining. In vision and language, the dominant pretraining strategies include autoregressive prediction (Radford et al. 2018), masked reconstruction (Devlin et al.) 2018 He

---

et al. 2021). and contrastive learning (Chen et al., 2020). In language, autoregressive generation emerged as a convenient self-supervised task. In surrogate modeling of dynamical systems, next-step prediction is often a primary goal. This makes autoregressive pretraining natural choice of objective for training time-dependent surrogate models.
We note that it is common to use the simulation parameters to condition the predictions of models operating on PDE-generated data (Gupta & Brandstetter 2022 Subramanian et al. 2023 Takamoto ct al. )023). Our goal is not to develop a new PDE sol ver, but rather to design an approach that is broadly applicable to both observed and simulated dynamics. Thus, we do not assume a known functional form in MPP and the model must instead implicitly infer the impact of these parameters on the dynamics from the history provided in U.S.
Surrogate Modeling for Spatiotemporal Physical Systems. We are primarily concerned with modeling dynamical systems varying in both time and space, where the time evolution of the system is intrinsically tied to spatial relationships amongst the state variables according to physical laws. PDEs are one of the primary modeling tools for this setting. They are often derived from fundamental conservation laws of properties such as mass, momentum, and energy (Farlow 1993). Many PDEs describe variations of the same physical laws, which is why concepts like diffusion, advection, reactivity, and connections between time and spatial gradients appear in many different PDEs. These shared underlying principles suggest we can extract features relevant to multiple physical systems.
3 Related Work
Foundation models. Massive pretrained models dubbed "foundation models" (Bommasani et al. 202T), particularly large transformer-baat based architectures (Vaswani et al. (2017). have recently attracted significant attention. The most prevalent foundation models are pretrained language models like
GPT (Brown et al. |020000 Radford et al.) 2018 2019) and BERT (Devlin et al., 2018). Emergent abilities (Wei et al.] (022) demonstrated by large language models highlight the importance of scale in
manifesting higher-order capabilities absent at smaller scales. Vision has seen similar developments with the growth of masked (He al., 2021; Tong et al., 2022) and contrastive (Chen et al. 2020) pretraining. The data in this work is insufficiently diverse to call the resulting models "foundational" However, we provide the first large-scale implementation of successful multiple nonlinear physics pretraining for spatiotemporal systems.
Scientific machine learning. While wide range of architectures have been employed for physical
surrogate modeling (Bar & Sochen, 2019 Han et al., 2018 SSrignano & Spiliopoulos 2018 Yu
al. 2018 Zang et al.] 2020), we position our work with respect to three three major classes, One www.eetwwokk-------ooooo
approach (Bruna et al. 2022 Raissi et al. 2019) which requires the PDE be known and solves a single system on a single domain. Other methods do not learn the solution directly, but instead augment a PDE-solver as learned corrections
prominent class is the
(Dresdner et al. 2023; Rackauckas et al. 2021 [m al. [202I), learned closures (Duraisamy et al.
2019 Sirignano & MacArt|202), or learned algorithmic components (Bar & Sochen (
al. 2021). A broader, but less physically constrained approach. is learning solution operator
from the data without knowledge of the governing equations (Cao) [222]] Kovachki et al. 2023 Li
ct al. 2020 202T: Lu al.] 2019). While these methods are often evaluated using PDE- PDE-generated benchmarks, these are designed to learn directly from data rather than learning to solve a PDE. Neural operators typically do not reach the accuracy of numerical PDE solvers, but they are applicable for domains without explicitly provided equations. This last family is the most similar to our approach, cspecially Cao [02T] as we use a transformer-basee architecture. However, our pretraining procedure is developed for training across multiple operators.
The high cost of training scientific models from scratch has led significant exploration of transfer learning. Prior work has explored transfer learning in operator networks in such scenarios as conditional shift (Goswami al. 2022 or new domains, boundary conditions, or distributions over parameters (Li et al.) [021]] Subel et al. 2023 Wang et al.) 2022a] Xu et al.. 2023). However, these too need to be retrained from scratch for new differential operators in the PDE. More recently, efforts have been made to explore transfer across operators and benefits from training on multiple physical systems simultaneously. (Subramanian et al.) (2023) explores how transfer scales in this setting. However, their study is limited to steady-state linear systems with periodic boundary conditions.
3

---

Other works have explored similarly restricted classes or low dimensional, low resolution systems
(Desa al., 2022 Yang et al. )
4
Scalable Multiple Physics Pretraining
4.1
Compositionality and Pretraining
Many specialized PDEs demonstrate form of compositionality, as a range of physical phenomena can be described by core components like nonlinear advection or diffusion, but then are augmented restricted by specialized terms representing concepts like buoyancy or system constraints. To motivate a useful pretraining procedure from this compositionality. we examine two hypotheses:
1. Single models can learn the dynamics for multiple classes of physical behavior.
2. Learning partially overlapping physics beneficial for transfer learning.
Since many real-world systems share core components, under these hypotheses, training single models on many distinct systems is natural approach for developing foundation models for physical dynamics. We therefore start by eliminating the complexity related to hypothesis (1) in order to isolate hypothesis (2). We this by choosing a simple problem setting with one shared scalar field along with consistent scales and geometry: constant-coefficit
linear advection -ifcfionniiffsiion on periodic ID domain. Let (a, t) be a scalar-valued function defined on a periodic spatial domain, v constant one-dimensional velocity coefficient and 8 constant diffusion coefficient, then:
Advection:
+ V
(v6) = 0
01
0e
Diffusion:
+ (-oVe) v) =
Of
Ob + (cv-oVe) = 0.
Advection-Difftssior
Ot
hypothesis (2) holds, we would expect pretraining on advection and diffusion systems individually could be beneficial for transfer to advection-diffusee systems.
find that this is indeed the case. We pretrain spatiotemporal transformer model on a large amount of trajectories (100,000 each) with uni- formly sampled coefficients (v |-3,--_1
Zero Shot
Pretrained
[11,3] 5 E [10 1.I) generated from the ad- vection and diffusion equations while finetuning on restricted samples from
Test Test RMSE
10
Scratch
simulations. The pretrained model is able to achieve much lower error with far fewer sam- ples (Figure without observing advection and diffusion together in the same trajectory during pretraining.
10-2
10a
10a
104
10
Number samples
However, to validate hypothesis (1), we must handle much larger spatial resolutions, varying scales, and heterogeneous relationships between fields. Over the rest of this section, develop an approach for handling these challenges.
Figure 1: Finetuning model pretrained large amounts of advection and diffusion data outper- forms models trained from scratch advection- diffusion data across a wide range of data availabil- ity 100K examples).
4.2 Architecture
Axial Attention. Given the success of large transformer models in other domains, we employ scalable axial attention (Dong et al. 2022 Ho al. 2019; Huang et al. 2019) transformer backbone. For (2+1)ddmenssome system with T x H x W tokens, conventional dense attention attends over all tokens simultaneously and cost O ((HWT))") Axial attention instead performs a series of attention operations over each axis in turn, limiting the cost to O(H² + W2 + T In Figure 2 it can be seen that while we perform attention on each axis independently, the spatial K, Q. V projections shared between the height (y) and width (x) axes.

---

Field Embedding
Multiple Physics Pretraining applied Axial Vision Transformer (AViT) for training Physical Surrogates
De-RevIN
Prediction
Field Reconstruction
convolve each filter
Decoder
Data
M, 5
Nx
P P
MLP
spatiotemporal transformer block
(fixed
Tokens]
time)
o(x,
alx, (ee
plx, Oee
Patch
sip Iju, 1j..
Spatial self-attention
embedding
Field Embedding
filter assembly
RevIN
Temporal
field
A-A
self-attention
1xl conv filters
(Field Indices)
Figure 2: (Left) MPP works by individually normalizing each example using Reversible Instance Normalization (RevIN) then embedding each field individually into a shared, normalized space. A single transformer backbone can then predict the next step for multiple sets of physics. We use AVIT backbone which attends over space and time axis sequentially. Spatial attention further split by axis, though these share linear projection weights. (Right) The embedding and reconstruction matrices are formed by subsampling a larger x convolutional filter based on input fields.
Axial attention has been used in video transformers (Arnab al.) 202 Bertasius al., 2021 due to the improved scalability in higher dimensions. While the tools used in our transformer backbone were introduced in prior work, our choice of using fully axial attention differs from ViVIT which opted to only separate space and time attention. We favor scalability over maximizing accuracy and so chose fully axial attention. In the following,
refer to this architecture as Axial ViT (AVIT).
Field Embedding and Normalization. Embedding multiple physical systems into a single shared representation is complicated by the fact that fields from different systems may operate on entirely different scales terms of both magnitude and resolution. This is one of the primary challenges that must be addressed for multiple-physiee pretraining.
To unify magnitudes, we use Reversible Instance Normalization (Kim al. 2022
RevIN). We compute the mean and standard deviation of each channel over space-time dimensions and use them normalize input fields. These statistics are saved and used to denormalize model outputs. While this approach was initially developed for time-series forecasting, in practice the effect is similar to that of the norm scaling utilized Subramanian al. (2023).
After rescaling, the data is projected into a shared embedding space. This is the only com- ponent with unique weights for each source system. Given system S with state variables u(x, f), vaz t), p(x, t) E R, we project each point or "pixel" into a space of dimension Demb,
e(x, 1) = u(x l)en + v(x, l)e, + p(æ, t)e,
where e are embedding vectors in IR Drmb This can be seen as a convolution with x 1 filters where the input channels of the filter are sub-selected to correspond to the fields present within given dataset. On the right side Figure the filter is assembled by sub-selected columns of the larger filter corresponding to the provided fields. It is important to note that this initial projection setup amenable to fine-tuning to unseen field types. This can achieved by adding new channels to the initial embeddings, and training them from random initialization. In our models, the shared full resolution space is converted into patched tokens by sequence of strided convolutions separated by pointwise nonlincarities in Touvron et al. (022))
5

---

The predictions are reconstructed from the output tokens by reversing this process. The tokens are decoded by a sequence of transposed convolution blocks and projected onto the output fields by taking coordinate-wist inner products with reconstruction vectors r:
u(x, /. + At) (e(x, + Al),
(2)
This can similarly be implemented as a 1 x convolution with the outpul channels of the convolution filter sub-selected. The mean and standard deviation computed from the inputs are then applicd to these normalized outputs to produce the final de-normalized predictions as in Kim et al. (2022).
Position Biases and Boundaries. While most cases, we would like the model to infer boundary conditions from the provided history. we make an exception to this policy for periodic boundaries as they change the continuity of the domain. Transformers are inherently permutation equivariant, and is essential to include position biases so that the model can learn locality.
With a slight modification, can use our position biases capture the change in locality imposed by periodic boundaries. TS-style (Raffel et al. (020) relative position encodings (RPE) utilize a lookup table to access learned embeddings corresponding ranges of "relative distance". For periodic boundary conditions, we modify the relative distance computation to account for neighbors across the periodic boundary. In Appendix D.I] we examine simple systems that differ only in boundary conditions and find that this minor change improves generalization in the case where we must learn both periodic and non- pen-piiooii conditions.
4.3 Balancing Objectives During Training
Task Sampling. Our pretraining procedure operates on multiple levels of sampling. The task distribution varies in system S, spatial resolution Ns, and time resolution I's and we want diverse batches that accurately capture the signal this provides. However, sampling a full batch from multiple systems at different resolutions simultaneously would be inefficient on modern hardware as it would require batch processing of differently shaped tensors. Multi- UPU training adds an additional complication as the variance in execution time due to unbalanced workloads can lead to inefficient hardware usage.
We mitigate both of these concerns with a simple randomization scheme involving gradient accu- mulation. Gradient accumulation utilizes multiple backward passes per synchronization step. We therefore sample a single system S uniformly from S for each micro-batch. With m micro-batches per synchronization step, we reduce the work-per-GPU variance o to m L. significantly reducing the average lost cycles due to work discrepancies. This could likely be further reduced by an approximate packing problem solution (Cormen al., 2022). but we found the random approach was sufficient for our needs. As we employ gradient accumulation in order to increase our batch sizes, this sampling procedure incurs no additional cost.
Scaled Training Objective.
The simplest approach to obtaining updates from the different tasks is to add their gradients. However, the magnitudes of the state variables can vary significantly between systems, unweighted losses will result in the gradients from the problems with the largest scales drowning out losses on smaller scales (Yu et al. 2020). To partially control this behavior, we train using the normalized MSE (NMSE) defined as:
1 |
|M((P)) U 2
LNSEE
11 +c
SES
where B C S denotes the micro-batch and is a small number added for numerical stability. This does not account for the full variation in difficulty. Even if sub-task losses have similar magnitudes at the start of training, it is possible for some systems to converge quickly while other losses remain high. Nonetheless, we found that this allows our training process to produce strong results on multiple systems simultaneously.
5 Experiments
We design our experiments probe three vital questions about the utility of MPP:
Can large transformer models learn the dynamics of multiple physical systems simultane- ously?

---

Table 1: NRMSE comparison between MPP- aPPP-prerraine models and dedicated baselines on shallow water equations (SWE), 2D Diffusion-Reaction (DiffRe2D), and compressible Kavier-Stokes (CNS) at Mach numbers M .1 and M 1. Complex parameters counted as two real. Top performing within size range and overall are bolded. Dashes indicate precision not provided by source.
MODEL
#PARAM
SWE
DIFFRE2D
CNS M1.0
CNS MO.1
MPP-AVIL IVPT-VIT-TI
7.6M
0.0066
0.0168
0.0442
0.0312
UNET
7.7M
0.083-
0.84
0,4725
1.6650
FNO
927K
0.0044
0.12-
0.1685
0.2425
FNO-B
115M
0.00246
0.0599
0.1451
0.1978
ORCA-SWIN-B
88M
0.00600
0.82-
AVIT-B
TASK-SPECIFIC
116M
0.00047
0.0110
0.0316
0.0261
MPP
116
0.00240
0.0106
0.0281
0.0172
MPP + FINETUNED
116M
0.00043
0.0087
0.0187
0.0079
MPP-AVIT-S
29M
0.0039
0.0112
0.0319
0.0213
MPP-AVIT-L
409M
0.0022
0.0098
0.0208
0.0147
2. Does pretraining lead to improved accuracy previously unseen physics?
3. Does MPP provide finetuning advantage over existing spatiotemporal foundation models?
Data. We use the full collection of two-dimensional time-dependent simulations from PDEBench (Takamoto et al., 2022) as our primary source for diverse pretraining data. This includes systems governed by four unique nonlinear PDEs defined over a variety of state variables, resolutions, initial conditions, boundary conditions, and simulation parameters. The specific PDEs are the compressible and incompressible Navier Stokes equations (CNS/INS), the shallow-watee equations (SWE), and 2D Diffusion- inffusionnReecttinn equation (DiffRe2D). Full details on the data used can be found in Appendix B.
Training settings. Tis fixed at 16 for all experiments as our VideoMAE comparison Section 5.2 was unable to scale to larger sizes without gradient checkpointing. Autoregressive training performed only one step ahead longer rollouts, noise corruption, or post-processing are included for stability. Training from scratch and MPP pretraining always performed the AVIT architecture described in section 4.2 Full training details including data splits, optimization details, and hardware are documented in Appendix
5.1 Pretraining Representations
First, we evaluate whether pretraining on multiple task actually leads to effective representations by comparing MPP-pretrained models to dedicated basclines from prior work across all available systems. The models are pretrained scaling our approach. Precise model sizes can be found in Appendix C.T Unlike the baselines which are trained on only one system and so must only learn one parameter regime, our models (denoted by MPP-AViT-*) must handle all systems and regimes without finetuning. The effect of physical parameters, forcing, and simulation parameters must be inferred from context U.S. The UNet (Ronneberger et al., 2015) and FNO (Li et al. 2020) results are sourced from Takamoto et al. 2022 while the results from Shen et al. (023) with finetuned SWIN (Liu et al.. (021) are used for ORCA. As the lightweight FNO proved to be the most competitive comparison, we train additional FNO beyond the PDEBench results that has been scaled to 115M parameters (labeled "FNO-B") for fairness. Results are reported in terms of Normalized RMSE (NRMSE, the square root of Equation 3) averaged over fields and examples, as in Takamoto et al. (2023). Our Compressible Navier-Stoker results aggregated based on the mach number here due space limitations. Fully granular results can be found in Appendix D.4 Ablation results demonstrating the importance of balanced losses and normalization are shown in Appendix [.2]
variety of sizes so we can begin to explore benefits of
Our pretrained models are able achieve high-end performance on all datasets (Table I despite the difficulty of multi-task training Yu et al. 2020) while showing improved performance with scale. Our smallest pretrained model, the MPP-AViT-T outperforms the PDEBench baselines on all problems except for SWE. However, though both models improve in absolute performance with scale, the
7

---

pretrained AVTT-B catches up to the FNO-B. It is important to clarify that we are not claiming the pretrained models are optimal-with a series of comparisons on the AViT-B models, we show that at times, the multi-task training does hurt performance on individual tasks and that we can improve upon the pretrained model performance by finetuning our own models on specific tasks. What this experiment answers affirmatively is that large transformers can learn multiple sets of dynamics simultaneously. Trajectories from pretrained models are displayed in Appendix D.6
5.2 Transfer to Low-data Domains
MPP AVIT B
Seratch AVIT B
V-MAE (K4M)-VIT-B
V-MAE (S8V2)-VIT-I
NIMSS
200
100
500
800
200
400
600
900
200
400
900
800
200
400
600
900
Number samples
Number somples
Number of samples
Number of somples
"Near" Transfer
(b) "Far"
Transfer
Figure 3: NRMSE for transfer learning tasks. Solid lines are one-step error. Dashed lines are averaged error over five step rollouts. The MPP model shows clear performance benefits in both cases. The more turbulent behavior of "far" seems to be difficult to learn from scratch or from video data, but pretraining on physical data leads to much stronger results.
Harkening back to Section 4.1 since we have shown that dynamics with a single model, we return to the question of whether these multiple physics mod- els are well suited to transfer learning. For more realistic exploration of transfer, we con- struct a setting where the model must learn new physical behavior by removing all compressible fluid data from the pretraining corpus and pre- training on the three remaining spatiotemporal systems. We then evaluate transfer to two spe- cific compressible Navier-Stokes datasets:
learn multiple sets of
Training
"Near"
"Far"
Figure 4: Kinetic energy for incompressible pre- training and compressible finetuning examples. The "near" compressible snapshot resembles the pretraining snapshot while "far" displays new tur- bulent small scales.
"Near": M 0.1, viscosity= 10-2, Random Periodic Initial Conditions
"Far"; M 1.0, viscosity= 10-8
Turbulent Initial Conditions
Snapshots of the kinetic energy for the finetuning systems and incompressible pretraining data are shown in Figure While quantitatively evaluating the physics gap is unsolved problem, the names reflect both prior physical knowledge and qualitative evaluation. "Near" features low Mach number, the dimensionless quantity that correlates with compressible behavior, and viscosity similar that of the incompressible simulation. "Far" has wildly different turbulent behavior that induces small scale structure never seen during training. However, despite the similarity in physical behavior, the simulations are still quite different: the compressible and incompressible simulations in PDEBench differ in spatial and temporal resolution, initial condition distribution, boundary conditions, viscosity, and velocity range in addition to the difference in compressibility. We use these sets to compare finetuning performance MPP, training from scratch, and an existing pretrained spatiotemporal transformer, VideoMAE (Tong et al. 2022) pretrained on both K400 (Kay et al. (017) and SSV2 (Goyal et al. 2017) datasets. Details on the finetuning procedure followed can be found in Appendix
C.3.I

---

Figure 3 shows that the MPP models outperform VideoMAE and training from scratch by large margin in the low-data regime. Numerical results are listed in Appendix [ VideoMAE displays surprisingly strong finetuning performance given that the pretraining data is conventional video, but it is unable to match the much lower memory (VideoMAF at 79.3 GB vs. AVīT-B at 24.7 GB peak VRAM for batch size 1) MPP-AViT-B in either setting. Predictably, both pretraining approaches are less accurate in the long-run on the turbulent "far" dataset. However, in the short-term the physical pretraining seems to provide an even larger advantage in this regime compared to the far smoother "near" data. Rollout visualizations are included in Appendix (D.7
One possible explanation for the strong performance of finetuning for "far" is that this experiment can be viewed as more realistic example of the compositionality exploration in Section |4. from the perspective of classification of second-ordes PDEs. The solution to Wavier-Stokes in the vanishing viscosity limit represents one possible weak solution to the Fuler equations which are classically hyperbolic (L.e Veeue & Leveque, 1992). While the viscous incompressible flow from pretraining governed by the same transport equations as "far", those solutions are dominated by smoothing locally parabolic behavior. However, the inviscid shallow water simulations in the training set are hyperbolic. The model has therefore seen two major components of "far", but has never seen them within one system until finetuning.
5.3 Inflation to 3D
While 2D problems offer compelling middle ground between complexity and cost for experimen- tation, most physical phenomena of real-world interest are fundamentally three-dimensional. We therefore examine the usefulness of our pretrained models when "inflated" to 3D. Inflation techniques were first demonstrated in Carreira& & Zisserman (2017) and have seen use for extending 2D visual
(image) classifiers to 3D spatiotemporal (video) settings (Nguyen et al. 2022 Xie et al. 2018). Here we employ the technique add an additional spatial dimension.
The factored architecture of the AVIT is well-suited to inflation. We initialize the projection weights discussed in Section 4.2 those from the 2D compressible Navier-Stokes data seen during training with the new velocity direction initialized as the average of the previous two velocity projections. Since the transformer backbone acts on each spatial axis independently. the only dimension-deependent operations are the learned downsampling or "patching" convolutions. These convolutional layers are modified by following the inflation procedure of Carreira & Zisserman (2017): 2D kernel of size P x P is inflated into 3D kernel by repeating the 2D kernel P times along the new axis and rescaling by p. This gives us a 3D convolutional operation that is equivalent to applying previous 2D filters then average pooling in the new direction. In practice, we find performance can be slightly improved by adding low magnitude Gaussian noise to the resulting filter. Due to the low resolution of the "Turbulent" dataset, we additionally reduced the stride of the first convolution in the hMLP by factor of 2 for a total downsampling factor of 8 rather than the 16 used elsewhere. Standard training details are found in Appendix C.4
We compare these inflated 2D to 3D models to both training an identical architecture from scratch and PDEBench baselines. Unlike in Sec- tion 50m25.2eeeer where we held out the 2D Compressible Navier-Stokes data, we use the full pretrained models from section 5.1 here. Nonetheless, this remains a significant physical gap as 2D and 3D turbulence are well-understood to have major differences in behavior (Ouellette. 2012). Due to the difficulty of scaling 3D training, results in this section are reported at the "Ti" scale.
Table 2: NRMSE for 2D to 3D inflation. Sub- headings are initial condition type.
TURBULENT (600, 21, 64
RANDOM (100, 21, 128*)
SIZE (B, T, N)
FNO
0.240
0.370
UNET AVIT-SCRATCH
0.230
1.000
0.098
0.299
AVIT-MPP
0.094
0.264
Table demonstrates significant improvements
on both the turbulently and randomly initialized Compressible Navier-Stokes datasets from PDEBench with a 11.7% improvement for the smaller dataset of randomly initialized simulations and even a 4.1% improvement for the larger turbulently initialized dataset where both the dimensionality and sampling are adjusted. Training on high resolution 3D data is an enormously expensive procedure. Our results suggest that pretraining on 2D data and inflating to 3D a promising strategy for developing models that can be used in this space.
9

---

Conclusion
We introduced an autoregressive pretraining strategy, Multiple Physies Pretraining, for the devel- opment of multi-use physical surrogates. Through per-sample normalization, field embeddings, appropriately scaled losses, and efficient task sampling, we are able train scalable transformer models capable of predicting multiple sets of independent dynamics simultaneously. We evaluated several sizes of model and observed that the approach benefits from scale. MPP models were able to match dedicated modern baselines on benchmarks containing fluid and reaction simulations derived from multiple equations, simulation parameters, and boundary conditions from pretraining alone with even stronger performance after undergoing task task-specific finetuning. Our pretrained models also showed positive transfer by outperforming both training from scratch and existing video models on previously unseen physics. Furthermore, through kernel inflation approaches, were able to demonstrate improved results on 3D simulation compared to training from scratch.
Limitations and Future Work. The focus of our work is on transfer and learning from multiple data sources, however, many interesting questions remain before we can develop true foundation models for spatiotemporal physics. For example, our choice of architecture enabled us to scale to some of the higher resolution benchmarks available today, but also limits our models to uniform grids. Extensions to the non-uniform geometries frequently encountered in physical simulation will require further work balancing efficiency and flexibility. Extending to 3D, for example, is an enormously expensive task for many families of models. Our inflation approach is promising and future work can help identify better approaches for performing this inflation. terms of more physically motivated concerns, we must also explore the role of constraints and conservation laws in models developed for multiple physical systems. Hard constraints that are necessary for a closed system may ensure error in open system, thus adaptable approaches may need to be developed.
Similarly, many physically interesting tasks require handling incomplete or noisy data. While certain datasets in PDEBench do not provide all relevant physical fields (the Shallow Water equations, for instance, only contain rather than full velocity fields needed simulate the data) all the data that is available is observed at every grid point without noise. It remains to be seen if these approaches generalize to settings without this type of clean data.
Finally, as these are early days in the development of foundation models for physics, the limits of transfer in these spaces are poorly understood. PDEBench, which we used here, constructed from largely fluid-like data. While we demonstrated strong transfer benefits from pretraining, it remains to be seen how far away from the training distribution or training tasks these benefits persist. It even an open question how to define distance in this space. Future work will need to expand data diversity to push these questions and others forward. MPP opens up many new research directions and paves the way for this development in the future.
10

---

Acknowledgments
The computations in this work were, in part, run facilities supported by the Scientific Computing Core at the Flatiron Institute, division of the Simons Foundation. M.P. is supported by the Department Energy, Office of Science under contract number DE-AC02-05CH11231 M.M. would like to thank Jed Brown for his valuable insight on the simulation data. Polymathic Al acknowledges support provided by the Simons Foundation and Schmidt Sciences, LLC.
References
Arnb, A., Dehghani, M., Heigold, G., Sun, C., Lučić, M., and Schmid, Vivit: A video vision transformer, 2021.
Balestriero, R., Ibrahim, M., Sobal, V., Morcos, A., Shekhar, S., Goldstein, T., Bordes, F., Bardes, A.,
Mialon, G., Tian, Y., Schwarzschild, A., Wilson, A. G., Geiping, J., Garrido, Q.. Fernandez, P., Bar, A., Pirsiavash, H., LeCun, Y., and Goldblum, M. A cookbook of self-supervised learning, 2023.
Bar, L. and Sochen, N. Unsupervised deep learning algorithm for pde-based forward and inverse problems, arXiv preprint orXiv:1904.054/1 2019.
Ben- Non-Bouallegue, Z., Clare, M. C. A., Magnusson, L., Gascon, E., Maier-Gerber, M.. Janousek, M., Rodwell, M., Pinault, F., Dramsch, S., Lang, S. K., Raoult, B., Rabier, F., Chevallier, M., Sandu, I., Dueben, P., Chantry, M., and Pappenberger, F. The rise of data-driven weather forecasting, 2023.
Bertasius, G., Wang. H., and Torresani. L. Is space-time attention all you need for video under- standing? In Proceedings the International Conference Machine Learning (ICML), July 2021.
Bi, K., Xie, L., Zhang, H., Chen, X., Gu, X., and Tian, Q. Accurate medium-range global weather forecasting with 3d neural networks. Nature, 619(7970):533-533 2023.
Biewald, L. Experiment tracking with weights and biases, 2020. URL httpttt::/// h/t::://wwwa wandb. com/ Software available from wandb.com.
Bommasani, R., Hudson, D. A., Adeli, E., Altman, R., Arora, S., Arx, S., Bernstein, M. S., Bohg,
J., Bosselut, A., Brunskill, E., et al. On the opportunities and risks foundation models. arXiv preprint ar 0rXiv:2108.07258 2021.
Bran, M., Cox, S., White, A. D., and Schwaller, Chemcrow: Augmenting large-language models with chemistry tools, 2023.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances neural information processing systems, 33: 33:1877-1901, 2020.
Bruna, J., Peherstorfer, B., and Vanden-Fijnden, E. Neural galerkin scheme with active learning for high-dimensional evolution equations, 2022.
Cao, S. Choose a transformer: Fourier or galerkin, 2021.
Carreira, J. and Zisserman, A. Quo vadis, action recognition? new model and the kinetics dataset. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), July 2017.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations, 2020.
Chithrananda, S., Grand, G., and Ramsundar, B. Chemberta: Large-scale self-supervised pretraining for molecular property prediction, 2020.
Cormen, H., Leiserson, C. E., Rivest, R. L.,, and Stein, C. Introduction to algorithms. MIT press, 2022.

---

Cranmer, M., Tamayo, D., Rein, H., Battaglia, P., Hadden, S., Armitage, P. J., Ho, S., and Spergel, D. N. A bayesian neural network predicts the dissolution of compact planetary systems. Pro- ceedings of the National Academy of Sciences, 118440):60026663111 2021. doi: 10. 2026053118, BtR.https://// //ww://wwwww... pnas wtwpppp......o.aaa/////////////.
2026053118
Dang, Y., Hu, Z., Cranmer, M., Eickenberg, M., and Ho, S. Tnt: Vision transformer for turbulence simulations, 2022.
Defazio, A. and Mishchenko, K. Learning-rate-free learning by d-adaptation, 2023.
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, Heek, J., Gilmer. J., Steiner, A., Caron, M., Geirhos, R., Alabdulmohsin, I., Jenatton, R., Beyer, L., Tschannen, M., Arnab, A., Wang, X., Riquelme, C., Minderer, M., Puigcerver, J., Evei, U., Kumar, M., van Steenkiste, S., Elsayed, G. F., Mahendran, A., Yu, F., Oliver, A., Huot, F., Bastings, J., Collier, M. P., Gritsenko, A., Birodkar, V., Vasconcelos, C., Tay, Y., Mensink, T., Kolesnikov, A., Pavetić, F., Tran, D., Kipf. T., Lučić, M., Zhai, X., Keysers, D., Harmsen, J., and Houlsby, N. Scaling vision transformers to 22 billion parameters, 2023.
Desai, S., Mattheakis, M., Joy, H., Protopapas, P., and Roberts, S. One-shot transfer learning of physics-infomed neural networks, 2022.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0488 2018.
Dong, X., Bao. J., Chen, D., Zhang, W., Yu, N., Yuan, L., Chen, D., and Guo, B. Cswin transformer: A general vision transformer backbone with cross-shaped windows, 2022.
Dresdner, G., Kochkov, D., Norgaard, P. C.,, Zepeda-Nuneds L.,, Smith, J., Brenner, M., and Hoyer, S. Learning to correct spectral methods for simulating turbulent flows. Transactions Machine Learning Research, 2023. ISSN 2835-8856. URL https: https:// http::////oppeeeeeeeeeeeeeroooa
JNBARGxoJnn
Duraisamy, K., laccarino, G., and Xiao, H. Turbulence modeling in the age of data. Annual Review of Fluid Mechanics, 51(1):357-377, 2019. doi: 10. Saltaannnurr--fud-------- 1 URL https: //doi. wrg/o0..... www.:1///aaaarreeelluleeeeeeeeee 040547
Farlow, S. J. Partial differential equations for scientists and engineers. Courier Corporation, 1993.
Goswami, S., Kontolati, K., Shields, M. D., and Karniadakis, Deep transfer operator learning for partial differential equations under conditional shift. Nature Machine Intelligence, 4(12); 1155-1164, 2022.
Goyal, R., Kahou, S. E., Michalski, V., Materzyúska, J., Westphal, S., Kim, H., Haenel, V., Fruend, I,, Yianilos, P., Mueller-Freitag. M., Hoppe, F., Thurau, C., Bax, I., and Memisevie, R. The "something something" video database for learning and evaluating visual common sense, 2017.
Grossmann, T. G., Komorowska, J. Latz, J., and Schönlieb, C.-B. Can physics-informet neural networks beat the finite element method?, 2023.
Gupta, K. and Brandstetter, J. Towards multi colliispatiotemppral--sse generalized pde modeling. arXiv preprint arXiv:2209. 1t..222999.661. 15616,7 2022.
Han, J., Jentzen, A., and E, W. Solving high-dimensional partial differential equations using deep learning. Proceedings of the National Academy of Sciences, 115(34)88505-855 8510. 2018.
He, K., Chen, X., Xie, S., Li, Y., Dollár, P., and Girshick, R. Masked autoencoders are scalable vision learners, 2021.
He. S.. Li. Y., Feng. Y., Ho. S., Ravanbakhsh, S., Chen, W., and Póczos, B. Learning predict the cosmological structure formation. Proceedings of the National Academy of Sciences, 116(28): 18825-13832. 2019. doi: 10. 10.1100//sss.1100 1821458116. 18215//1 10. 10.1073/pnas /pa/.///a...... 1821458116
URL UtLhsttps: h/wp:::///wwww pnas org/doi/ /bs/
Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus), 2016.
12

---

Hernandez, D., Kaplan, J., Henighan, T., and McCandlish, S. Scaling laws for transfer. arXiv preprint arXiv:2102.01293, 2021.
Ho, Kalehbrenner, N., Weissenborn, D., and Salimans, Axial attention in multidimensional transformers, 2019.
Hoffmann, J., Borgeaud. S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L.,
Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models.
arXiv preprint arXiv: 15556, 2022.
Huang, Z., Wang, X., Huang, L., Huang, C., Wei, Y., and Liu, W. Ccnet: Criss-cross attention for semantic segmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 603-612, 2019.
Jamieson, D., Li, Y., de Oliveira, A., Villaescusa-Navaan F., Ho, S., and Spergel, D. N. Field-level neural network emulator for cosmological n-body simulations. The Astrophysical Journal, 952(2): 145, 2023.
Jiang, I.. Y., Liu, X. C., Nejatian, N. P., Nasir-Moin, M., Wang, D., Abidin, A., Faton, K., Riina, H. A., Laufer, I., Punjabi, P., et al. Health system-scale language models are all-purpose prediction engines. Nature, pp. 1-6, 2023.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint orXiv:2001.0836/ 2020.
Kay, W., Carreira, J., Simonyan, K., Zhang, B., Hillier, C., Vijayanarasinhan, S., Viola, F,, Green, T., Back, T., Natsev, P., Suleyman, M., and Zisserman, A. The kinetics human action video dataset, 2017.
Kim, T., Kim, J., Tae, Y., Park, C., Choi, J.-H., and Choo, J. Reversible instance normalization for accurate time-series forecasting against distribution shift. In International Conference on Learning Representations, 2022. URL https://openrreiiieeee
Klaasen, G. A. and Troy, W. C. Stationary wave solutions of system of reaction-diffusion equations derived from the fiazuoo 96-110, 1984.
nagumo equations. SIAM Journal on Applied Mathematics, 44(1):
Kochkov, D., Smith, J. A., Alieva, A., Wang. Q., Brenner, M. P., and Hoyer, S. Machine learn- ing-accelerated computational fluid dynamics. Proceedings of the National Academy of Sciences, 118(21)e210118441 8, 2021. doi: 10. 10.1073/pnas:2101788 18. URL/https: //www pnas org/ doi/abs/10. 1073/ pnas 2101784118
Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhattacharya, K., Stuart, A., and Anandkumar, A. Neural operator: Learning maps between function spaces, 2023.
Lam, R., Gonzale, A., Willson, M., Wirnsherger, P., Fortunato, M., Alet, F., Ravuri, S., Ewalds, T., Raton-Rosen, Z.,, Hu, W., Merose, A., Hoyer, S., Holland, G., Vinyals, O., Stott, J., Pritzel, A., Mohamed, S., and Battaglia, P. Learning skillful medium-range global weather forecasting. Science, 382(6677) 2023. doi: 10. hetl///ienc..ai22...... URL https
science. wwwo/ogg////ooooooooooo htt:///o//a//s/e/e wwwp://ieeccceee adieeii.eeeeeeee........ adi2336
Leung. H. W. and Bovy, Towards an astronomical foundation model for stars with transformer- based model, 2023.
Veque, R. J. and Leveque, R. J. Numerical methods for conservation laws, volume 214. Springer, 1992.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandku- mar, A. Fourier neural operator for parametric partial differential equations. arXiv preprint axXiv:2010.08895. 2020.
Li, 7., Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Arizzadenesheli, K., and Anandkumar, A. Physics-informed neural operator for learning partial differential equations. arXiv preprint arXiv:2111.03794, 2021.
13

---

Liu, Z., Lin. Y., Cao, Y. Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierar- chical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In International Conference on Learning Reprexentations, 2019. URL
Lu, L.. Jin, P., and Karniadakis, G. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv:1910.03199 2019.
Lusch, B., Kutz, N., and Brunton, S. L. Deep learning for universal linear emheddings nonlinear dynamics. Nature Communicaations (1):4950, 2018. doi: 10. 0010088/441467-018--7770000 URL 0ttp:/10..0033/11117700-----2120
Mialon, G., Garrido, Q., Lawrence, H., Rehman, D., LeCun, Y., and Kiani, B. T. Self-supervised learning with lie symmetries for partial differential equations, 2023.
Nguyen, Goel, K., Gu, A., Downs, G. W., Shah, P., Dao, T., Baccus, S. A., and Ré, C. S4nd: Modeling images and videos as multidimensional signals using state spaces, 2022.
Nguyen, T., Brandstetter, J., Kapoor, A., Gupta, J. K., and Grover, A. Climax: A foundation model for weather and climate. arXiv preprint arXiv:2301.1034/ 2023a.
Nguyen, T. D., Ting, Y.-S., Ciucä, I., O'Neill, C., Sun, Z.-C., Jablońska, M., Kruk, S., Perkowski, E., Miller, J., Li, J., Peek, J., Iyer, K., Rózański, T., Khetarpal, P., Zaman, S., Brodrick, D., Méndez, S. J. R., Bui, T., Goodman, A., Accomazzi, A., Naiman, J., Cranney, J., Schawinski, K., and Universe Astrollama: Towards specialized foundation models in astronomy, 2023b.
Ouellette, N. T. Turbulence in two dimensions. Physics Today, 65(5):68 69, 05 2012. ISSN 0031-9228. doi: 10.1063/PT.331570
httlhttps://d//: /ttpttpp://aoii.og/1..00000 111.:/....1003P....11..
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, I.., Desmaison, A., Kopf, A., Yang. E., Vito, Z.,, Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, high-performance deep learning library. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran As- sociates, Inc., 2019. URL http:/ttttt///t/eeeee h111///ccc2288.e7722222888811111
www.neuuiipp.. wtwp:ippp.ccppppee.iilleeeeeeeee
Pathak, J., Subramanian, S., Harrington, P., Raja, S., Chattopadhyay, A., Mardani, M., Kurth, T., Hall, D., Li, Z., Arizzadenesheli, K., Hassanzadeh, P., Kashinath, K., and Anandkumar, Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202. 1214, 2022.
Pfaff, T., Fortunato, M., MonchezzGonzzlzz.. A., and Battaglia, P. Learning mesh-based simulation with graph networks. In International Conference Learning Representations, 2021. URL
Rackauckas, C., Ma, Y., Martensen, J., Warner, C., Zubov, K., Supekar, R., Skinner, D., Ramadhan,
A., and Edelman, A. Universal differential equations for scientific machine learning. 2021.
Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training. 2018.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 188)9, 2019.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer, 2020.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
14

---

Ronneberger. O., Fischer, P., and Brox, T. U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assiset Assisteed Intervention-MICCC 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp- Springer, 2015.
Shen, J., Li, L., Dery, L. M., Staten, C., Khodak, M., Neubig, G., and Talwalkar, A. Cross-modal fine-tuning: Align then refine, 2023. httpss//:://rr wwwp:/wwww..rrga/////ba/0/ 050000003332320
Sirignano. J. and MacArt, J. F. Deep learning closure models for large-eddy simulation of flows around bluff bodies. Journal of Fluid Mechanics, 966, jul 2023. doi: 211017/ffm.22233444 URL
https: http:::://oo..rg/1o...o 1017 2063.446
Sirignano, and Spiliopoulos, K. Dgm: deep learning algorithm for solving partial differential equations. Journal of computational physics, 375:1339-1364 2018.
Stachenfeld, K., Fielding, D. B., Kochkov, D., Cranmer, M., Pfaff, T., Godwin, J., Cui, C., Ho, S., Battaglia, P., and Sanzhez-Gonzalez, A. Learned coarse models for efficient turbulence simulation, 2022.
Subel, A., Guan, Y., Chattopadhyay, A., and Hassanzadeh, P. Explaining the physics of transfer learning in data-driven turbulence modeling. PNAS nexus, 2(3)pgad0115, 2023.
Subramanian, S., Harrington, P., Keutzer, K., Bhimji, W., Morozov, D.., Mahoney, M., and Gholami, A. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior, 2023.
Takamoto, M., Praditia, T., Leiteritz, R., MacKinlay, D., Alesiani, F, Pflüger, D., and Niepert, M. PDEBench: An Extensive Benchmark for Scientific Machine Learning. In 36th Conference on Neural Information Processing Systems (NeurIPS 2022) Track on Datasets and Benchmarks, 2022.
URL httph:tt/axxiiii //arxiv.. org/ abs/2210. 07182
Takamoto, M., Alesiani, F., and Niepert, M. Learning neural pde solvers with parameter-guides channel attention, 2023.
Tong, Z., Song, Y., Wang, J., and Wang, L. VideoMAE: Masked autoencoders are data effticeffi learners for supe-supe video pre-training. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (cds.), Advances in Neural Information Processing Systems, 2022. URL httpts://ppeeeeiiiee net/forumiid-Ahcon
Touvron, H., Cord, M., El-Nouby, A., Verbeek, J., and Jegou, H. Three things everyone should know about vision transformers. arXiv preprint N2X.v:2203309799 2022.
Tu, T., Azizi, S., Driess, D., Schaekermann, M., Amin, M., Chang, P.-C., Carroll, A., Lau, C., Tanno, R., Ktena, I., Mustafa, B., Chowdhery, A., Liu, Y., Kornblith, S., Fleet, D., Mansfield, P., Prakash, S., Wong, R., Virmani, S., Semturs, C., Mahdavi, S. S., Green, B., Dominowska, E., Arcas, B. A.,
Barral, J., Webster, D., Corrado, G. S., Matias, Y., Singhal, K., Florence, P., Karthikesalingam, A., and Natarajan, V. Towards generalist biomedical ai, 2023.
Ulyanov, D., Vedaldi, A., and Lempitsky, V. Instance normalization: The missing ingredient for fast stylization, 2017.
Um, K., Brand, R., Yun, Fei, Holl, P., and Thuerey, N. Solver-in-the-lopp Learning from differen- tiable physics to interact with iterative pde-solvers, 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, I. Attention is all you need, Advances in neural information processing systems, 30, 2017.
Wang, H., Planas, R., Chandramowlishwaran A., and Bostanabad, R. Mosaic flows: A transferable deep learning framework for solving pdes on unseen domains. Computer Methods in Applied Mechanics and Engineering, 389: 14424, 2022a.
Wang, T., Roberts, A., Hesslow, D., Scao, T. L., Chung, H. W., Beltagy, I., Launay, J., and Raffel, C. What language model architecture and pretraining objective work best for zero-shot generalization?, 2022b.
15

---

Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M.. Zhou, D., Metzler, D., et al. Emergent abilities of large language models. arXiv preprint arXiv:2206.07682, 2022.
Xie, S., Sun, C., Huang. J., Tu, Z.,, and Murphy, K. Rethinking spatiotemporal feature learning: Speed-accuracy trade-offs in video classification, 2018.
Xie, X., Zhou, P., Li, H., Lin, Z., and Yan, S. Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models, 2023.
Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., and Liu, T.Y. On layer normalization in the transformer architecture, 2020.
Xu, W. Lu, Y., and Wang, L. Transfer learning enhanced deeponet for long-time prediction of evolution equations. In Pruceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 10629-10636, 2023.
Yang, L., Liu, S., Meng, T., and Osher, S. In-context operator learning for differential equation problems. arXiv preprint arXiv:2304.07993, 2023.
Yu, B. et al. The deep ritz method: a deep learning- based numerical algorithm for solving variational problems. Communications Mathematics and Statistics, 6(1):1-12, 2018.
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-task learning, 2020.
Zang. Y., Bao, G., Ye, X., and Zhou, H. Weak adversarial networks for nigh-dimensional partial differential equations. Journal of Computational Physics, 411::09409, 2020.
Zhai, X.. Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. In Proceedings of
the JEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 12104 1144-1211,, 2022.
16

---

A Impact Statement
This paper contributes the field of machine learning for computational physics. In terms of positive impact, we would hope that our approach enables researchers and engineers to develop more accurate dynamics models from limited observations. While we see mostly positive implications from this, the danger of any computational physics research is potential use in weapons research. In this case, we feel that risk is fairly small as research in such spaces tends to focus on well-understood physies where numerical methods already obtain extremely high precision while our approach is more suited to the more scientifically oriented space of modeling poorly understood physics at lower precision.
B
Data Details
B.1 PDEBench
To train and evaluate our models, we use the publicly available PDEBench datase Takamoto et al. 2022). We summarize the data included in this section. This dataset comprises a suite of time dependent and time independent simulations based on common PDE systems, generated with varying parameters, initial conditions, and boundary conditions. Specifically, PDEBench uses a discretized ground-truth solver with high precision to evolve the vector-valued solution to a given PDE one time step to the solution at one time step later. When compiled across time steps, the vector-valued solutions take the form I E Tx H where I denotes the total number of times steps, H and W denote the spatial height and width of the simulation grid and C denotes the parameter space representing the velocity (Ug and By). pressure (p) and density (p) fields, such that C = 4. For our study, we focus on the 2D fluid dynamics simulations in PDEBench. These are outlined loosely below; for more details, we refer the reader to Takamoto et al. (2022):
Compressible Navier-Stokes:
These equations are used to model the pressure and velocity of both laminar and turbulent Newtonian fluids, and applied to many real-world problems, from aerodynamics to interstellar gas dynamics. In the regime in which the density of the fluid can change due to pressure variation, the equations can be expressed:
0.p + V. (pv) 0,
p +v. Vv) = =VVp / v (s+
(5)
ac(c + po²/2) + V [(+ (pp+ce 22
v o'] = =0,
where P is the fluid density, v is the fluid velocity, is the fluid pressure, e is the internal energy, o' the viscous stress tensor, 7) is the shear viscosity, and is the bulk viscosity. For our transfer experiments, we use the following two sets of data in particular:
1. A set of 1. 1,000 trajectories on a H x W 512 x 512 regular grid over T 100 time steps (where the separation between steps is At = 0.005). Additionally, (M, C) (1.0, 10-* 10-*), where M, 1), C denote the Mach number, the shear viscosity, and the bulk viscosity, respectively. The velocity field is initialized with turbulent field, while the inital pressure and density fields are taken to be uniform.
2. A set of 10,000 trajectories 10.1, 0.01, 0.01). The time steps and initializations are as above.
H x W 128 128 regular grid with (M, () =
Incompressible NS: In the incompressible regime, which typically occurs in fluids with low Mach numbers (as it rules out density and pressure waves like sound or shock waves), the Navier-Stokes equations simplify to:
= 0,
p (a,v + Vv) = Vp p VV
(8)
where v is the velocity, p is the density, p is the pressure, 7) is the viscosity, and f is a spatially varying external force. The simulation in PDE bench is augmented by immersed tracer that is transported by the velocity field:
Oepsmoke -V V Pamake
(9)
17

---

The system uses Dirichlet boundary conditions on the velocity field v = 0 and Neumann on the density Opamaie = 0, x {0, 1), Ppamoke = 0, y E {0, These equations are typically used By
Oz
to model variety of hydrodynamics systems such as weather. This data is produced at resolution 512 x 512 with time step of .0005. The dataset contains a total of 1000 trajectories with 1000 time steps each.
Shallow water: In the event that the horizontal length scale of the fluid significantly greater than the vertical length scale, the incompressible Navier-Stokes equations can be depth-integrate to derive the shallow water equations. These describe flow below pressure surface in a fluid, and are given by
0,h + V (hv) = 0, 0,(1v) V. V. ( 1 2 29.12 9
(10)
(11)
where h is the water depth, v is the velocity, b is the bathymetry, and gr is the reduced gravity. For our data, we use 1,000 trajectories on a H W 128 x 128 regular grid over 7 100 time steps. The specific simulation used is 2D radial dam break scenario, where the water height is initialized as circular bump in the center of the domain with uniformly randomly sampled radius.
Diffusion-Reactionn The Diffusion-Reaction equations arise in systems with many interacting components and can be represented in the general form
0,u DV R(U).
(12)
where u is a vector of concentration variables, D is a diagonal matrix of diffusion coefficients, and R describes all local reaction kinetics. The most common application of biffusion-reaction equations is in chemical reactions, however they can also used to describe a variety dynamical processes. For
our data, we use 1 1,000 trajectories on a H x W = 128 x 128 regular grid over T = 100 time steps. The reaction functions for the activator and inhibitor are defined by the Fitzhugh-Nagune (Klaasen & Troy. 1984), and their diffusion coefficients are D., 1 X 10-3 and D,, = 5 x 10 respectively. The initial conditions are generated as standard Gaussian random noise.
equation
B.2 PDEArena
In addition to the 2D Incompressible Navier-Stoker data incorporated from PDEBench, we also include 2D Incompressible Navier-Stoker data from PDEArena (Gupta & Brandstetter 2022).
These follow roughly Equation] 7
with a minor variation:
0,
(13)
p Vv) -
(14)
where b E [0.2, 0.5) represents buoyancy in the y direction. Unlike PDEBench, this is spatially constant term. This includes set of 5,200 training trajectories (and 1,300 validation and test trajectories each) a H W 128 x 128 regular grid from which take T 16 timesteps for prediction. As with the PDEBench simulations, the PDEArena simulations include a viscosity parameters of 0.01 and Dirichlet boundary conditions.
C Experiment Details
C.1 Model Configurations
The following architectural decisions were used across all AVTT models trained in this paper:
Pre/Post Norm: Pre-norm (Xiong et al. 2020
Normalization Type: Instance Normalization Ulyanov et al. 2017
Activations: GeLU (Hendrycks Gimpel 2016)
QK Norm: Yes (Dehghani et al.
2023
Patching: hMLP Touvron et al.
2022
18

---

Table 3: Details of the various model architectures and scales explored.
MODEL
EMBED DIM.
MLP DIM.
# HEADS
# BLOCKS
PATCH SIZE
# PARAMS
AVIT-TI
192
768
3
12
[16. 16]
7.6M
AVIT-S
384
1536
12
[16. 16]
29M
AVIT-B
768
3072
12
12
[16. 16]
116M
AVIT-L
1024
4096
16
24
[16, 16]
409M
Table 4: Inference time for various models on A6000 GPU.
MODEL
TIME (MS)
UNET
67.7
FNO AVIT-TI
7.2
19.2
ORCA-SWIN-B
98.5
AVIT-B
105.6
Decoder: Transposed hMLP (this is equivalent to the transposed convolutions mentioned in the main text).
Causal Masking:
False We only evaluate the loss on the T + prediction.
Furthermore, we examine the performance of our models on the aforementioned PDE systems when the size of the model is scaled. Vision transformers have variety of parameters that control the model's size, including the number of processor blocks, the dimensionality of patch embeddings and self-attention, the dimensionality of Multi-Layer Perceptron (MLP) blocks, the number of attention heads, and the patch size applied on the input tensors. In previous studies on language (Hernandez al. 1021 Hoffmann et al. 2022 Kaplan al. 2020) and vision (Zhai et al. (022). it has generally been noted that model performance is typically only weakly dependent on shape parameters, and instead depends largely non non-embedding parameter count given fixed compute budget and dataset size. As such, we follow the general scaled architectures set forth by Zhai et al. (2022) for vision, and scale all aspects of the model shapes simultaneously select variety of model sizes for testing.
These are detailed in 3 In practice, the AVTT models are on average slightly slower than the similarly sized baseline models from Table
Software. All model development and training in this paper is performed using PyTorch 2.0 (Paszke al.
2019).
Hardware. All training for both pretraining and finetuning is done using Distributed Data Parallel (DDP) across Nvidia H100-8008 GPUs.
C.2
Exp 1: Pretraining Performance
For MPP, we train using the following settings:
Training Duration: 200K steps
Train/Val/Test. 1//.1.1 split per dataset on the trajectory level.
Task sampling: Uniformly sample task, then uniformly sample trajectory from task without replacement. We treat every 400 model updates (1 model update=5 micro-balee batches) as an "epoch" and reset the task pool.
Micro-batch size: 8
Accumulation Steps: 5
Optimizer: Adan (Xie et al. 2023)
Weight Decay: IE-3
Drop Path: 0.1
19

---

Base LR:
DAdaptation (Defazio & Mishchenko 2023
LR Schedule: Cosine decay
Gradient clipping: 1.0
For training from scratch, no task sampling occurs and sampling without replacement continues until the dataset is exhausted as in conventional epoch. Note, we use the automated learning selection strategy DAdaptation during pretraining runs in large part avoid excessive hyperparameter tuning of our own models. In finetuning experiments, comparison models are tuned manually following the recommended settings from the model publishers avoid differences being due to compatibility with the parameter-fre method.
FNO-B For the scaled-up FNO, we modify the parameters used in PDEBench to 6 layers, 24 modes, and width 100. Training configurations were taken from PDEBench (500 full passes through the dataset with the Adam optimizer) with learning rate selected by DAdaptation as in the training of our models for consistency.
Data For pretraining, we use all 2D time-dependent PDEBench datasets. These are described in Section B.I. In particular, use the compressible and incompressible Navier-Stokes, Diffusion- Reaction 2D, and Shallow Water data.
C.3 Experiment 2: Transfer to Low-Data Domains
In this experiment, we compare the transfcrability our MPP-Pretrained models to general purposes pretrained video masked autoencoders (VideoMAE; Tong al. 2022) for frame prediction video-like PDEBench data (Takamoto et al. 2022)
Data We study transferability of VideoMAE models for spatiotemporal prediction on video-like scientific data.
MPP-laheled models are pretrained on datasets generated from three PDEs: Incompressible Navier- Stokes, Shallow Water, and Diffusion Reaction 2D. This is performed using the same training settings
as in Section C
We focus on transfer to the two datasets "Near" and "Far" (see Sect. 5.2) of fluid dynamics simulations taken from the PDEBench dataset (Takamoto al. 2022). These simulations solve the compressible Navier-Stokes equations in 2D geometry with periodic boundary conditions (see Appendix B.] for additional details).
C.3.1
MPP Finetuning Procedure
For MPP and training from scatch, we use the following settings:
Training Duration: 500 epochs (true epochs due to the restricted dataset)
Train/Val/Test: 1/11/1 split per dataset on the trajectory level. Note that X is due to the fact that test varying amounts of training data. These subsampled from the training split of 80%.
Batch size: 8
Accumulation Steps: (No accumulation)
Optimizer: Adan (Xie et al. 2023)
Weight Decay: 1E-3
Drop Path: 0.1
Base LR: DAdaptation (Defazio & Mishchenko 2023
LR Schedule: Cosine decay
Gradient clipping: 1.0
20

---

Channel expansion. The only component of the architecture used that is aware of the particular fields being ingested is the field embedding and debedding projection layers. As the finetuning data adds previously unseen state variables, this matrix needs to be expanded. To be concrete, out of the four state variables present in the CNS data - Uz. Vy. P. p - the training data only contains Uaa By- If the original field embedding layer was therefore 1x convolution with 3 input channels, we follow the following procedure to add the new projection.
1. Instantiate new fheeddpprjecctiiee and raald_rrconstuuctiioo layer with a suffi- ciently large number of channels to process both previously seen fields and new fields.
2. If W is a (..., Cin, Cout) field_projection weight matrix, set W [....: Outd :] Wold, Perform the corresponding replacement within the reeedddrccostrucction layer.
All other weights are loaded as normal without modification.
Weights trained. We found that partially frozen training resulted in noticeably worse performance and therefore finetuned the full model.
C.3.2 VideoMAE Settings
While VideoMAE does utilize spatiotemporal information, it was developed for different setting, so we fully document all details of our adaptation of it here both for reproducibility and fairness in our comparison.
VideoMAE models are video transformers that were proven to be efficient data-learners for self- supervised video pretraining (Tong et al., 2022). They rely on asymmetric encoder-decoder architecture building on a vanilla VIT backbone with joint space-tima time attention. VideoMAE models pretrained by learning to reconstruct masked videos using a random tube-masking strategy with extremely high masking ratio (~~90% %).
We make use of two publicly available models, hereafter called CitcoMAE-K4000 and VideoMAE. SSV2, that were pretrained on Kinetics-400 dataset (K400; Kay et al. 2017) and Something- Something V2 dataset (SSV2; Goyal et al. 2017), respectively. Both datasets are made of short videos (typically S 10 long) human-object or
interactions. VideoMAF-K400 (respectively, VideoMAE-SSV2: was pretrained ~ 240k (~ 170k) vidcos. We focus the models that build on a ViT-base backbone, so that their size (in terms of number trainable parameters) remains comparable to that of MPP-AViT-B. After adaptation of the input and output linear layers as described below, the number of trainable parameters of these models reaches 95 M.
Number of channels. Same as the original pretraining procedure, the input data x E RCXTXXXW divided into non-overlapping joint space-time cubes of size 2 x 16 x 16. These are embedded through Conv3d layer, resulting in I' x H W' tokens. Since our PDEBench data has C 4 channels instead of 3 for the RGB videos from the pretraining set, we had to adapt the number of input channels of this Conv3d layer accordingly. The weights of this new layer were defined using (rescaled) repetition of the pretrained weights from the original layer. Similarly, the output number of features of the final linear projection layer of the model had to be adapted C' 4 channels. The weights and biases of this layer were extended by consistently repeating the original pretrained weights and biases.
Positional encoding. The number of tokens resulting from our PDEBench data did not match the number of tokens resulting from the pretraining datasets. Consequently, we also had to adapt the pretraining positional encoding. We chose to interpolate accordingly the original ID sine/cosine positional encoding (Vaswani et al. 2017) using a trilinear interpolation after having reshaped the token index axis onto 3D grid.
C.3.3 Video MAE Finetuning Procedure
We describe the finetuning procedure of the pretrained VideoMAE models for frame prediction. Frame prediction consists in predicting the next T, frames of video given a context of T. frames. Since the pretrained models manipulates space-time cubes of size 2 in time, we naturally choose Tp = 2. The context size is taken to be T. 16 for consistency with AVP-AVIT models. We finetune
21

---

Table 5: Effective learning rate for the finetuning of VideoMAE.
"NEAR"
"FAR"
VIDEOMAE (K400)
0,00039
0.00198
VIDEOMAE (SSV2)
0.00186
0.00150
the pretrained models for frame prediction by adapting the self-supervised training strategy order to reconstruct the last Ip frames of masked video of T Te + Tp frames.
Masking strategy. For frame prediction, instead of the random tube-masking strategy, we simply mask the last Tp frames of the input data.
Loss. We finetune our models by minimizing NMSE loss. In this context, denoting by I, E xxxxT,***X the output of our model and the target (masked frames), respectively, the NMSE loss is defined by ((x,y) )) = EeE Σ |Ze, - ye,t /e.r11//llee. 1ll,l,,
Normalization of the data. Each set of PDEBench simulations is globally and channel-wise rescaled so that pixel values all in [0, 1]. Additionally, we normalize channel-wise the targets y X by subtracting the global mean of the corresponding context frames and then dividing by their global standard deviation.
Optimization. We finetune the pretrained models over 500 epochs (full epochs due to restricted data size) and (total) batch size using Adam optimizer (Loshchilov Hutter 2019). Except for the learning rate, the remaining optimization hyperparameters those used in the finetuning experiments of Tong et al. (022) (Table 10). In particular, we choose weight decay À 0.05. (31, 32) (0.9, 0.999), cosine learning rate decay scheduler with warmup epochs, a drop path rate of 0.1, and a layer- learning rate decay parametrized by 0.75. In this setting, the learning rate is adjusted by performing hyperparameter search monitored with WandB (Biewald 2020). We report the resulting optimal values per pretrained model and dataset in Table
chosen to be consistent with
C.4 Experiment 3: Inflation to 3D
For MPP and training from scatch, we use the following settings:
Training Duration: 100 True Epochs
Train/Val/Test: .8/11.11
Batch size: 8
Accumulation Steps: (No accumulation)
Optimizer: Adan (Xie et al. 2023)
Weight Decay: 1E-3
Drop Path: 0.1
Base LR: DAdaptation (Defazio & Mishchenko 2023
LR Schedule: Cosine decay
Gradient clipping: 1.0
Data We use both 3D Compressible Navier-Stokes datasets provided by PDEBench for these experiments. The first has 600 total trajectories of 21 steps simulated with sheer and bulk viscosities of 10-8 with explicitly turbulent initialization. The second contains 100 total trajectories of steps with viscosities of 10 and random initialization.
22

---

Table 6: Validation NRMSE for position bias comparison. Compares training performance on data that differs only in boundary conditions.
TRAINING
PERIODIC
ABSORBING
PERIODIC BASELINE
0.032
ABSORBING BASELINE
0.295
COMBINED
STANDARD RPE
0.188 0.081
0.189
PERIODIC-ADIUSTOO RPE
0.143
C.5 Appendix Experiment: Broader Usage of Pretrained Representations
For MPP and training from scatch, the following settings:
Training Duration: 20K Optimization Steps
Train/Val/Test: 1000/100/1000 taken from original validation set randomly depending whether data was used for training.
Batch size: 24
Accumulation Steps: (No accumulation)
Optimizer: Adan (Xie et al. 2023)
Weight Decay: 1E-3
Drop Path: 0.1
Base LR: DAdaptation (Defazio & Mishchenko 2023
LR Schedule: Cosine decay
Gradient clipping: 1.0
D Additional and Extended Results
D.1 Position Bias Evaluation
We isolate the impact of position biases on our multi-task training objectives by constructing an experiment that isolates their influence. Recall the advection equation from Equation
Oyb
+ V (vwb = 0
(15)
Öt
We will define two sets of physics. In both cases, the function is defined on the ID domain x E [0, 1]. We sample Unif(-1,1) and use initial conditions sampled from the of circular Gaussians with variances sampled from Unni((1/166 1/5) and means sampled from f systems vary only the choice of boundary conditions. The first uses periodic boundary conditions, implying <(0) é(1). The second uses absorbing boundary conditions in which waves are not reflected back into the solution space. The restricted functional form allows us implement this exactly by extending the domain and solving the periodic equations such that the constant velocity implies the waves exiting the solution space never return.
The two
In this experiment, we first train models (AViT-Ti with 1D patches) on each system individually using 10,000 examples each for 100 epochs to get sense of the baseline performance. We then train models with and without our modified position biases on the two systems jointly (20,000 examples) evaluate the impact of our change.
Table shows that our modified position biases are more effective at training in the joint setting, Both RPE schemes able to improve absorbing boundary with the additional data. Standard RPE on the other hand struggles to learn the periodic baseline. Our Periodic-adjusted variant is much more effective learning the periodic data, though does not outperform the baseline.
interesting to note how large the effect of boundary conditions is on this problem. The model trained on only periodic condition reaches nearly an order of magnitude higher precision. While
23

---

Table 7: Ablation table showing changes in NRMSE proposed modifications to the architecture and model are removed.
SETTING
SWE
DIFFRE2D
CNS M1.0
CNS M0.1
MPP-AVITE
0,00240
0.0106
0.0281
0.0172
REMOVE NRMSE TRAINING
0.01353
0.1502
0.1245
0.1213
REMOVE REVIN
0.02651
0.0661
0.2601
0.3266
REMOVE BOTH
0.03619
0.2952
0.4049
1.5892
Table 8: RMSE comparison between MPP-pretrained models and dedicated baselines on shallow water equations (SWE), 2D Diffusion-Reaid Reaction (DiffRe2D). Complex parameters counted as two real.
MODEL
PARAM
SWE
DIFFRE2D
MPP-AVIT-TI
7.6M
6.9E-3
1.IE-3
UNET
7.7M
8.6E-2
6.1B-2
FNO
927K
4.5E-3
8.1E-3
AVIT-B
TASK-SPECIFIC
116M
4.9E-4
8.8E-4
MPP
1 116M
3.6E-3
6.6E-4
MPP FINETUNED
1 116M
4.2E-4
6.2E-4
MPP-AVIT-S
29M
4.0E-3
7.8E-4
MPP-AVIT-L
409M
2.3E-3
S.0B-4
absorbing boundaries are complicated for numerical solvers, it seems as though attention should be able to simply not attend to waves passing out of the domain. The interaction boundary conditions with attention therefore seems to be an important direction for future study.
D.2 Normalization Ablations
We perform several ablations to record the impact our proposed normalization and loss balancing approaches to the task of multiple physics pretraining. These results are shown in Table 7 using AViT-B the baseline. As we can both modifications greatly improve the ability of base model to learn in this setting. Note that all of these experiments are performed in the multiple-physics settings. Results for dedicated models are listed in the main text.
Note that while training on a normalized loss appears to be important for MPP, the ordinal value of the final results is agnostic to the metric chosen. For example, include Table version of Table [ which reports loss in RMSE instead of NRMSE on the smaller datasets to demonstrate that while not all reported models trained using the same metric, this does significantly change the results. Note that not all sources report RMSE so this table contains fewer results compared to the main text.
D.3 Broader Usage of Pretrained Representations
One of the fascinating aspects of large pretrained models is the utility of their learned features for entirely new types of prediction problems. We explore this behavior by comparing the ability of pretrained MPP-AViT-B model to one trained from scratch to solve the inverse problem of parameter estimation for two parameters.
Forcing Identification for Incompressible Navier-Stokes The two sources variation in the Incompressible Navier-Stokes simulations (Equation 7) are the initial conditions and the spatially varying forcing f applied to the velocity evolution at each step. We compare the performance between the pretrained the constant forcing term used in the incompressible Navier-Stokes simulation from an input trajectory U.S. We divide the validation set from pretraining, taking 1,000 trajectories as the new training and using the rest for validation. Results reported the original test set.
Buoyancy for Incompressible Navier-Stokes For this, we turn to an additional fluid mechanics benchmark, PDEArena (Gupta & Brandstetter 2022). This benchmark includes incompressible

---

Table 9: RMSE for inverse problem tasks. Error from constant prediction included for context.
TRAINING
FORCING
BUOYANCY
MPP
0.20
0.078
SCRATCH
0.43 0.43 2.12
0.0772.005
MIALON ET AL. 2023)
0.062
PREDICT MEAN
1.00
0.088=.000
Table 10: Per dataset NRMSE comparison for M = 0.1 Compressible Navier-Stokes data, R/T denote "random" and "turbulent" initial conditions from PDEBench. 77 = Ç are the bulk and sheer viscosity.
MODEL R-y 10-6
R-n 10-2
R-7 10-1
T- 10-8
MPP-AVIT
0.0493
0.0274
0.0116
0.0339
UNET
0.66-
0,71
5.1
0.19-
FNO
0.28-
0.17-
0.36-
0.16-
MPP-AVIT-S
0.0335
0.0176
0.0071
0.0217
FNO-B
0.1810
0.1129
0,3800
0.1175
MPP-AVIT-B
0,0286
0.0162
0.0078
0.0169
MPP-AVIT-L 0.0234
0.0145
0.0099
0.0136
Navier-Stokes simulation with variable buoyancy (b from Equation 13). Since this set was used during training, we take 1,000 randomly sampled trajectories for train, 100 for validation, and further 1,000 for testing. Since we are now predicting scalar, we train a linear probe on top of the final hidden representation consisting of global average pooling and linear head.
We observe mixed results (Table D. Pretraining reduces the error in the forcing task nearly half, but shows no improvement over training from scratch the scalar prediction. Prior work (Mialon al. 2023) was able to achieve better performance buoyancy through Lie-transformation based contrastive pretraining using a convolutional architecture. MPP does not seem to hurt performance on this task, as the AVIT trained from scratch also barely outperforms a mean prediction. However, we would expect the scalar prediction task to be easier. It is plausible that the dense prediction pretraining task is not well-suited for scalar inference or that the pooled representation frequently used for finetuning classification models in vision is not well suited to parameter inference, but the comparison of performance this non-generative task also echoes prior work in NLP (Wang et 2022b) where autoregressive training has underperformed on non non-gennrratvve tasks.
D.4 Expl: Expanded Results
Here we break out the Compressible Navier-Stokes (CNS) results from Table 1 Table shows the comparison between our pretrained models and task-specific baselines; however, due to space limitations the CNS was aggregated by mach number in the main text, so we share the full CNS results here. M0.1 can be seen in Table IO MI.0 can be seen in Table Note that while is conventional to describe these simulations in terms of dimensionless numbers like the Reynolds number, these simulations are performed at relatively low resolution, so it is likely they incur significant numerical diffusion. Thus we report the results in terms of the nominal diffusion coefficients without making claims about the Reynolds numbers the simulation.
In examining the full CNS data, one interesting result jumps out the most viscous systems 1) = .1 seem to perform relatively worse with scale. For both subsets, was the top performing model the highest viscosity. All other viscosities seem to benefit from scale. This does seem to have a limit, however, as Ti again loses performance. during multi-tast training, they cannot be directly interpreted in the single-task setting.
also important to remember that these results oceur
25

---

5.2 for Table 11: Per dataset NRMSE comparison for M = 1.0 Compressible Navier-Stokes simulations. R/T denote "random" and "turbulent" initial conditions from PDEBench. "I are the bulk and sheer viscocity.
MODEL R-1
10
R-n
10
R-7
10
T-1 10
MPP-AVIT-TI
0.0615
0.0327
0.0171
0.0594
UNET
0.47-
0.36-
0.92-
0.14-
FNO
0.35
0.096-
0.098
0.13
MPP-AVIT-S
0.0451
0.0223
0.0108
0.0425
FNO-B
0,3216
0.0573
0.0914
0.1100
MPP-AVE AVIT-B
0.0386
0,0195
0.0119
0.0365
MPP-AVIT-L
0.0314
0.0171
0.0132
0.0282
Table 12: Test NRMSE for "Near" Compressible Navier-Stokes M0.1, 7) .01.
# TRAINING SAMPLES (NRMSE x10-1)
MODEL
100
200
400
600
800
T+1
T+5
T+1
T+5
T+1
T+5
T+1
T+5
T+1
T+5
VIDEOMAE (K400)
1.26
1.98
0.78
1.25
0.49
0.83
0.39
0.62
0.33
0.50
VIDEOMAE (SSV2)
0.95
1.61
0.63
1.04
0.42
0.66
0.33
0.52
0.25
0.39
MPP-AVIT-B
0.66
1.13
0.42
0.81
0.27
0.55
0.22
0.35
0.19
0.30
D.5
Exp2: Numerical Results
We provide numerical results corresponding to Figure in Tables 12 and
13
We refer to Sect.
discussion.
D.6 Pretraining Trajectories
Here we show example trajectories from pretrained models. Videos are included in attached supplementary material. After pretraining, find that the model initially produces strong predictions, but patch artifacts creep in over time.
t+1 1
t+3
t+5
t+7
t+
Predicted
Ground Truth
Dynamics: swe, Field
26

---

Table 13: Test NRMSE for "Far" Compressible Wavier-Stokes
# TRAINING SAMPLES (NRMSE x10-
MODEL
100
200
400
600
800
T+1
T+5
T+1
T+5
T+1
T+5
T+1
T+5
T+1
T+5
VIDEOMAE (K400)
1.16
1.60
0.79
1.10
0.73
0.96
0.53
0.70
0.49
0.65
VIDEOMAE (SSV2)
0.98
1.42
0.75
0.62
0.84
0.55
0.74
0.51
0.67
MPP-AVIT-B
0.60
1.15
0.37
0.77
0.27
0.66
.32
0.63
0.24
0.48
t+1
t+3
t+5
t+7
t+
Predicted
Ground Truth
Dynamics: diffre2d, Field
D.7 Finetuning Trajectories
After finetuning, we find that the patch-based instability mostly disappears. Again, videos displaying longer trajectories are available the supplementary material.
t+1 1
t+3 3
t+5
t+7 7
t+9
Predicted
Ground Truth
Dynamics: diffre2d, Field v
27

---

t+1 1
t+2
t+3
t+4
t+5
Predicted
Ground Truth
Dynamics: CNS, Field Vx
t+1
t+ 2
t+3 3
t+4
t+5
Predicted
Ground Truth
Dynamics: CNS, Field Vy
t+1
t+2
t+3
t+4
t+5
Predicted
Ground Truth
Dynamics: CNS, Field density
28

---

t+1
t+2
t+3
t+4
t+5
Predicted
Ground Truth
Dynamics: CNS, Field pressure
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The paper claims the following. 1. We develop an approach that allows a single transformer model to learn from hetereogeneous physical systems. This developed in 4.2 and demonstrated in 5.1. 2. Training on hetereogeneous systems can be beneficial in transfer settings on new physical behavior. This is shown on toy example 4.1 and more complex examples in 5.2 and 5.3. 2.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: The conclusion extensively discusses the limits in terms of geometry and compute as well as more physically based limitations in terms of conservation laws. also discusses the limited setting in which this has been applied thus far.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
29

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
There are no theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: Code is provided and the settings for every experiment are recorded in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often
30

---

one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model ge.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code is directly provided and the data is gathered from open well-established datasets.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (tttps://nips.co)
public/guides//odeSuubae wtw.eesuubissiioooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https
/
gttggguiidee/eesee wwt.submiisiiooooooocccoooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Experimental settings are listed in the appendix.
Guidelines:
31

---

The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: As these are explorations of large-scale models with multiple terabytes of training data, experiments are too expensive to repeat sufficient number of times gather useful estimates of statistical significance.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: These are listed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
32

---

The submission obeys copyright law, does not involve human subjects, and has no immediate negative impacts on society.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Impact statement included appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
There are major risk associated with model use.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
33

---

12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The only existing asset in use is the training data which is appropriately cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, pawpess..pppppiipbaaa. toodee com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The only new asset included is the model which is covered by the repository license.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Submission does not involve crowdsourcing.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
34

---

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No crowdsourcing or human subject research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "3qqjXJZTjT", "reviewer": "Reviewer_1VF5", "summary": "The work proposes the idea of pre-training models on multiple PDE problems and demonstrates that such a pretrained model can be effectively fine-tuned on a target PDE when the pre-training and fine-tuning PDEs are similar. The authors carefully constructed the scalable transformer architecture by employing an axial attention mechanism of high-resolution PDE datasets. They validated their hypothesis through empirical analysis of the 2D Shallow Water equation, Diffusion-Reaction equation, and Navier-Stokes equation in both 2D and 3D domains.", "questions": "1. The model does not consider any PDE-specific input. If the pre-training dataset includes data from very similar PDEs (such as Navier-Stokes equations with nearly identical viscosities), would the model be able to distinguish between them effectively?\n2. In Table 1, multiple bold elements appear in the same column. What do they signify?\n3. How well does the proposed architecture handle changes in spatial and temporal resolution during testing?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The work tackles a contemporary yet unresolved issue of designing a foundation model for solving PDEs. \n\n2. The proposed model showed impressive performance on the selected PDE problems.", "weaknesses": "1. **Only transferable to very similar PDEs**. \nThe argument for pre-taining on multiple PDEs was to generalize to unknown PDEs in low data regimes. However, the experiment design lacks practicality. The demonstrated example of finetuning on CNS while pre-training on a large amount of INS data lacks the similarity of any practical situation. For example, a demonstration of the effectiveness of the model on the introduction of new variables, change of domain or geometry, or a moderately different PDE (such as keeping out SWE from the pretraining dataset and finetuning on a small amount of SWE data) would have been a more practical choice. \n\n2. The domain expansion from 2D to 3D adaptation experiment was intriguing. However, the performance of AVIT-SCRATCH and AVIT-MPP is quite similar. This raises the question of whether the benefit comes from MPP or the designed architecture. Nonetheless, this does not align with the primary message of the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "First off, we’d like to thank the reviewer for their time and effort in reviewing our paper. Your feedback will make the paper stronger. There are a few points raised in the review that we hope to clarify. \n\n---\n## W1 - Only transferable to very similar PDEs.\n\nThe hypothesis we explore in this work is that “learning partially overlapping physics is beneficial for transfer learning”. This is what we believe our work has shown, though the level of heterogeneity is quite a bit larger than what may be obvious from the equations themselves.\n\nCNS and INS are separated by the Mach number. Low Mach CNS simulations, as in the dataset we call \"Near\" are qualitatively very similar to INS simulations, but this is no longer true at high Mach number. In PDEBench specifically, there are even more significant differences. __The CNS data in PDEBench possesses both previously unseen geometries and new variables when compared to the pretraining data in Section 5.2.__ We’ll break these differences into several sections: geometry, physics and state variables, and the significance of the experiment.  \n\n. \n### Geometry\nThe INS data uses no-slip boundary conditions - these approximate a domain with solid walls. The CNS data uses periodic boundary conditions which imply no walls. Somewhat unusually since periodic boundaries are typically less expensive to simulate, in our transfer experiments, the model has actually never seen periodic geometry. Without walls, the small-scale behavior in the flow is a product of colliding shocks rather than the forcing and wall dynamics from the INS simulations. \n\n. \n### Physics and Variable Comparison\nThe difference in settings also results in different represented fields. The INS setting in PDEBench is a closed, inhomogeneous, highly viscous (coefficient=$10^{-2}$) flow driven by a forcing term. It is a domain without a pressure gradient and where any density variation is driven by the presence of immersed particles in the flow rather than density differences due to compression or expansion. It is fully characterized by *velocity*, *forcing*, and the *density of the immersed particles*. The behavior can locally be classified as parabolic everywhere - viscous forces damp any emerging shocks such that even discontinuous initial conditions will become smooth. \n\nThe CNS physics, particularly for “Far” which has viscosity coefficients of $10^{-8}$, are very different. They are unforced, the density of the fluid itself varies freely due to compressible behavior, and pressure is not constant. Thus these simulations add fluid density and pressure fields that were also never seen during pretraining. In terms of the physics, near-inviscid flow with Mach number of 1 allows for local behavior resembling hyperbolic systems - shock formation and transport. \n\n. \n### Significance\nWhat’s really exciting to us, and we hope for readers as well, is that this experiment reinforces the compositionality hypothesis: the model has seen hyperbolic and momentum-driven transport during training, but never in the same system. The inviscid shallow water equations used to generate the data are hyperbolic and shock forming while the model had seen momentum-driven transport in INS. Thus this transfer experiment can be interpreted as a more complex version of our motivating example in Section 4. \n\nThis is also one of the most practically useful examples of transfer. Incompressible flow solvers tend to be significantly cheaper to run than compressible flow, so learning from incompressible data would provide an enormous advantage for data acquisition. \n\n---\n## W2 - 3D transfer\n\nWe’re glad that you found the adaptation experiment interesting. 2D-3D transfer is an area that we're excited about as well due to the scaling benefits. Most of the 2D data in PDEBench can be generated in less than 24 hours on current hardware. Equivalent resolution 3D data can take weeks or months on similar hardware. From a ML training perspective, 512 x 512 input is straightforward. Handling 512 x 512 x 512 requires specialized distribution strategies. The ability to use 2D data can be an enormous cost savings for both training data generation and pretraining. \n\nWhile “quite similar” is a matter of opinion, we’d note that these are 4.1% and 11.7% improvements. It looks small because the AViT itself is a significant improvement over the PDEBench baselines. We felt it was important to contextualize the results with the stronger AViT comparison to more accurately demonstrate the impact which is greater than 10% for the smaller dataset where transfer is more vital.\n\n---\n## Q1 - ... would the model be able to distinguish between them effectively\n\nThis depends on the level of similarity, but in general it is possible and demonstrated in our experiments. The PDEBench CNS data uses a range of viscosity parameters. In Tables 9/10 in the supplementary material, we can see that the pretrained models trained on all datasets/parameters are still able to outperform the PDEBench baselines on each dataset. \n\n---\n## Q2 - Bolded tables\n\nBolding indicates top performance within a parameter count bracket. “L” is the top performer overall showing that we see significant benefits from scale. \n\n---\n## Q3 - Varying resolution\n\nIn general, the model needs to have seen problems in a similar range - we would not expect zero shot transfer to unseen discretizations. Using 1D example from section 4, we can intuit that if the model has always been trained at velocities such that density moves exactly one cell during training, the model will learn to only explore neighboring cells, and decreasing the temporal resolution will result in nonsensical zero-shot answers. This is actually one of the perks of MPP - the data coverage from using MPP helps avoid overfitting due to data idiosyncrasies.\n\n---\n## Conclusion\n\nAgain, we appreciate the reviewer's time and effort. If we were able to alleviate your concerns, we'd ask that you adjust your evaluation to reflect it.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZwZLXiEN1W", "reviewer": "Reviewer_1X1Y", "summary": "This article introduces a transformer model which turns across multiple special temporal physics. Using a clever choice of standardization and scaled training, they managed to train a model that can predict the next step given context of snapshots. In the article, they show that a single model can learn dynamics from multiple classes of physical behavior, and a model can learn partially overlapping physics for transfer learning. In the experiment, they show competitive results across multiple physics with a strong emphasis on the challenging fluid dynamics. They also show capacity to transfer efficiently to low data regime. They also show promising results of using their 2-D model for 3-D solves.", "questions": "Could you please address the following minor comments?\n- Please provide details of the fine tuning with MPP, I could only find details for videoMAE.\n- Please provide details on how to predict the first T_S first snapshots since they have incomplete context. \n- Could you please give an intuition why MPP performs so much better than the rest in figure 3b?\n- Please discuss whether the learned tokens from multiple physics could be interpretable or how could this MPP be leveraged to make the predictive model more trustworthy?\n- L. 146. The word “function” is missing after “scalar”\n- L. 253 it should be T_S instead of T^S\n- L. 350 what does “Ti” scale mean? Is it before De-RevIn?", "limitations": "Limitations are adequately addressed.", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is an outstanding article and I have only a few minor comments.", "weaknesses": "Please refer to the minor comments in Questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We'd like to thank the reviewer for their thorough reading of our paper. The reviewer raises some great questions and discussion topics which we aim to address below. \n\n#### Q1 - \n\nThe finetuning settings are provided in C.3 for 2D and C.4 for 3D. However, while the instructions for initializing new field embeddings for new fields are contained in the repository, those details are not currently in the paper. In addition to what is stated in C.3/4, for new fields, one must also expand the field embedding weights and randomly initialize the new embeddings. \n\nWe will add this and expand the finetuning details in a manner more similar to what is described for VideoMAE to centralize the information.\n\n#### Q2 - \n\nThe initial snapshots must come from the data source. It is possible that the context length can be reduced or use variable context lengths with some minor data augmentation during training - 16 was largely chosen for an equitable comparison to VideoMAE - but ultimately some history is needed. This is because the model needs to use a form of in-context learning to differentiate between similar dynamics. For instance, if we’re looking at the PDEBench compressible Navier-Stokes data, we often have similar initial conditions in systems with slightly different viscosities. In a single snapshot, the two cannot be disentangled, but over a moderate length trajectory, one can observe the dynamics evolving differently. \n\nIn designing models for this space, the model either needs to know the exact system coefficients or be able to implicitly infer dynamics from a provided histor. We felt the history approach was more in line with potential workflows. Given observational data, for instance, it is unlikely a user could provide the Reynolds number, but they can almost certainly take a small number of measurements separated in time. Or in a design optimization setting where one might want to approximate complex non-differentiable legacy code with the model with a differentiable surrogate, it might be trivial to generate a short history, but several months of effort to identify all of the conditional modeling assumptions in the code base. \n\n#### Q3 - \n\nThanks for raising this point! This is something we shied away from discussing explicitly in the paper as it requires some level of PDE background, but since it has come up with several reviews, we will add this information to the paper itself. \n\nIn short, there is a degree to which 3B can be seen as a more realistic instance of what we see in our motivating example of section 4. In Section 4, we explore whether learning partially overlapping physics can be beneficial in the extremely simplified setting of 1D linear PDEs on periodic domains. Our pretraining data contains two families of fluid flow - viscous incompressible Navier-Stokes (INS) and the inviscid shallow water equations (SWE). The INS data can be seen as globally parabolic. Even locally, there is no point where viscous forces will allow for shock formation. The SWE simulations have locally hyperbolic behavior meaning we can see discontinuities (shocks) form naturally in the flow.\n \n3A incorporates new physics, but with high viscosity and low mach number, qualitatively it should be similar to the INS pretraining data. 3B is nearly inviscid and occurs at high mach number. This allows for local behavior resembling hyperbolic systems to emerge. The model never saw this in momentum-driven transport problems (INS), but it has seen SWE which has hyperbolic behavior but does not provide the velocity fields for computing transport. So like in Section 4, we have two sources of training data that contain sub-components of the new physical process, but finetuning is required since the model has never seen those components together. \n\n#### Q4 - \n\nThis is a very interesting question. It's not clear that this approach improves interpretability, though analyzing MPP with mechanistic interpretability tools to identify what type of algorithm the model is employing to make the forecast is a direction we're excited about for future work. However, there is a strong case to be made for reliability. Imagine we train a transformer on a single dataset. In this dataset, fluids move only a single cell over each timestep. The model would predictably learn to move based on a sign function on velocity - this isn't the correct function, but its a function that is simple to learn and fits the data perfectly. MPP enables the model to learn from multiple datasets simultaneously ensuring coverage for a wide variety of data. This makes it difficult for models to learn trivial idiosyncrasies of single datasets as the model must learn compressed representations relevant to many datasets simultaneously.\n\n#### Quick Qs\n- Q5 - Thanks! We will fix this.\n- Q6 - Thank you! Fixed.\n- Q7 - This is the model size \"tiny\". It corresponds to the smallest AViT in Table 1\n\n#### Conclusion\n\nAgain, we'd like to thank the reviewer for their efforts. We hope this discussion is interesting and provides useful information!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ENppGnTIYl", "reviewer": "Reviewer_mqdr", "summary": "This paper introduces a multiple physics pretraining approach for surrogate modeling, which learns general useful features across diverse physical tasks with a shared embedding and normalization strategy. The experiment results show the proposed MPP-pretrained model outperforms task-specific baselines on all pretraining sub-tasks and also show superior finetuning results on new physics tasks.", "questions": "1. Does MPP need to be trained in a context of very similar physical backgrounds (e.g., SWE, DiffRe2D and CNS)? How can we determine the similarity of multiple physics fields and whether they can be learned simultaneously?\n2. Single MPP can learn the dynamics for multiple classes of physical behavior. If the physical equations are vastly different and include different forms of dynamic processes, can one model still perform well? For example, can a single model handle Newtonian fluids and non-Newtonian fluids, rigid body dynamics or elastodynamics?\n3. Do baseline methods use the same normalized training loss as MPP? Why do you evaluate using Normalized Mean Squared Error (NMSE) instead of Mean Squared Error (MSE) during evaluation?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea of constructing a large pre-trained base model for physical simulations is promising. The experiments are performed for diverse physical systems. The results show that large surrogate models outperform strong baselines, and even models with a relatively few parameters can learn such diverse physics evolutions and perform competitively.", "weaknesses": "- The validation of the approach is primarily conducted on fluid mechanics-oriented benchmarks. While this is a solid start, the applicability of the approach to other domains of physics remains to be demonstrated extensively.\n- The proposed model can handle simulations on structured meshes. However, simulations on unstructured mesh are under exploration.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "To begin, we'd like to thank the reviewer for the effort they put forth in reviewing our work. Their feedback will help us strengthen the submission and we are grateful for their expertise. We hope that we can address their concerns through our rebuttal. If we are able to address any of your concerns, we would ask that you consider raising your evaluation accordingly. \n\n---\n### W1 - ...primarily conducted on fluid mechanics-oriented benchmarks...\n\nThis is a great point which we mention in our limitations section. However, we'd emphasize that limited data scope can also be useful. MPP is a large change from established approaches for training surrogate models. It is therefore valuable to ground the work in the existing literature rather than creating an entirely new set of tests and comparisons - the latter invites the question of tuning effort and specialization which detracts from the goal of the paper. \n\nPDEBench, despite the restriction to fluids, is one of the more diverse spatiotemporal physics benchmarks published to date incorporating multiple equations, parameters, boundary conditions, and dimensionalities; and with ~125 citations has seen relatively wide adoption. As we state in the limitations, we believe training true \"foundation models\" for this space will require more data maturity, but PDEBench serves as ideal testbed to explore the question of partially overlapping physics while sticking with battle-tested benchmarks.\n\n--- \n### W2 - Unstructured Meshes\n\nWe agree that processing unstructured meshes is vital for many engineering applications, but as an active research area in its own right, irregular mesh handling is orthogonal in many senses to the questions we wanted to answer in this work. In this paper, we wanted to answer two fundamental questions - first whether it is possible for individual models to learn multiple sets of dynamics. Secondly - whether this is beneficial for transfer to new systems whose physics partially overlap with the training data. Uniform grids were sufficient for answering both of these emphatically in the affirmative. Extending our approach to non-uniform grids opens up some really exciting research directions related to transfer across geometry, but we felt the initial exploration would benefit from tighter scope and reduced degrees of freedom.\n\n---\n### Q2/Q1 - Major differences in physical process and the limits of transfer\n\nThese are very interesting questions that highlight the type of research questions opened up by MPP. Large scale pretraining in vision and language have opened up new avenues for interpretability research and we believe MPP could do the same for data-driven spatiotemporal physics. In this paper, we're focused on partially overlapping physics on a fluid-dominated dataset. PDEBench is clearly not the whole of physics, but as one can see comparing INS and CNS, the restricted class covered is still quite diverse. Demonstrating transfer over this range is a significant advancement over prior work and an important step forward in the development of foundation models for this space. \n\nUltimately exploring the limits of transfer will require more mature datasets and significant amounts of compute. Identifying a priori approaches for quantitatively evaluating whether dynamics are similar or finding where model capacity begins to breakdown are natural and interesting research questions that emerge once MPP is established as a possibility. While there is more to explore, we feel opening these new exciting directions is a strength of our work rather than a weakness. \n\n---\n### Q3 - \n\nNMSE was chosen primarily for consistency with the existing literature. It is the only metric reported across all papers we drew baselines from. However, we also agree with this convention. Normalized metrics are viewed as more interpretable. In a vacuum, it can be difficult to tell if an RMSE of $10^{-6}$ is “accurate” since the underlying field may have values on order $10^{100}$ or $10^{-14}$. NRMSE avoids that issue, though the learning task itself can still vary in difficulty. \n\nSince all baselines are drawn from independent studies, they are also trained using the procedures chosen by the original authors. The PDEBench baselines are trained using RMSE while ORCA is trained with NRMSE. For single task models, there is very little difference here as PDEBench fields have similar norms across samples within a given dataset. The exception to this are the NS where there is variation within fields so the relative weighting of fields in the aggregate loss is different. However, this has little effect on aggregate metrics in practice - the ordinal rankings do not change in RMSE. We understand the concern and will add an RMSE table to the appendix to address it.\n\n---\n### Conclusion\n\nOnce again, we'd like to thank the reviewer for the energy they put into helping us refine our paper. If we were able to address any of your concerns, we'd ask that you raise your score proportionally.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ipfUBLj4GN", "reviewer": "Reviewer_cf8J", "summary": "This paper introduces the Multiple Physics Pretraining (MPP) model, a pretraining approach for physical surrogate modeling of spatiotemporal systems using transformers. MPP uses a backbone model to predict the dynamics of several heterogeneous physical systems simultaneously. The authors include a shared embedding and normalization strategy to facilitate effective learning and transfer across different scales and magnitudes. The accuracy of MPP is validated on a fluid mechanics-oriented benchmark, showing that a MPP can perform better than the baselines without finetuning. For downstream tasks, MPP-trained models show more accurate predictions compared to training from scratch or finetuning pretrained video foundation models.", "questions": "* Line 4-5: The paper claims to \"predict the dynamics of multiple heterogeneous physical systems\". However, the method is only focused on 2D fluid mechanics examples. How generalizable is this approach to fully heterogeneous physics data (say, mixing electromagnetism + solid mechanics + fluid mechanics problems)?. In that case, the AViT might not a viable option anymore.\n* Section 4.1: Related to the previous question. This paper relies heavily on the compositonality assumption. Indeed, all the examples are particular cases of the full compressible NS equations together with a transport equation. What if the compositionality assumption does not hold? Often in physics there are certain kinds of phenomena which might not compose trivially from the equations (for example, fluid-structure interaction, contact, etc).\n* Line 167: Apart from the multihead attention from the AViT, have the authors considered a mixture of experts architecture? Would it increase the performance with respect to a single common backbone network?\n\nFinal comment: Despite the limited applicability to fluid mechanics examples, this paper contributes significantly towards the direction of foundation models for science. The methodology is novel, the results outperform baseline models and the paper is well written and structured. Based on the comments above I suggest an accept.", "limitations": "Limitations are addressed appropriately.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The model is robust and outperforms task-specific baselines without finetuning.\n* The authors provide open-source code and pretrained models for reproducibility.\n* The method has good transfer capabilities to systems with limited training data.", "weaknesses": "* The network architecture is heavily oriented to 2D fluid mechanics examples, but the applicability to other domains and physics systems is not clear.\n* The model is validated only with good quality data, which might not be the case for real-world measured datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and expertise. The reviewer raised a number of interesting points that we'd like to address. Some of our responses are more discussion than rebuttal, but we hope we can ease the reviewer's concerns in certain areas and allow them to feel confident in their review.\n\n___\n### W1 - 2D Fluids\n\nThis is a point where we believe we can offer clarification. The AViT architecture itself is designed to easily generalize to varying dimensions - for instance in our 3D inflation experiments, the axial architecture simply repeats the spatial attention operation in the added experiments. No new weights need to be added to the “inner” blocks of the network. Similarly, one could run on 1D data by collapsing a dimension of the encoder/decoder and applying spatial attention along the one dimension. However, we agree that it is currently designed for uniformly gridded meshes and that architecture changes may be required to address non-uniform grids in future work.\n\n___\n### W2 - Data quality\n\nThis is an excellent point. We will update our limitations section to reflect this. While this paper uses data entirely from an established benchmark in order to ground our results with a novel approach firmly in the literature, we believe exploring transfer to experimentally obtained data would be an interesting and valuable study. \n\n---\n### Q1/2 - Generalization\n\nWe feel these points are both fundamentally about generalization, so we'll split them up slightly differently and discuss generalization in terms of architecture and in terms of model capacity.\n\nFirst, architecture: as we mentioned in W1, we’d clarify that the AViT’s limitation is non-uniform grids rather than being strictly designed for 2D. Our inflation experiments for instance show multiple physics pretraining, even in 2D, provides value for 3D downstream tasks as well. The axis-aligned architecture, however, does require either uniformly gridded data or data that can be projected onto a uniform grid. While we believe the use of data-dependent spatial mixing (as found in attention or recent selective state space models) is important for multiple physics training, MPP does not require an AViT specifically. \n\nThe question of generalizability in the sense of model capacity is more complex and likely requires downstream study. If we imagine that we were generating training data by solving random PDEs, it is intuitive that we would eventually reach a capacity limit - we’re not going to learn an infinite set of mappings from a finite set of parameters. However, we’re not generating random PDEs. Most systems of interest to researchers are either derived from conservation laws or are aggregated from particle/density level relations which leads to the very interesting question of where exactly these capacity limits are. As the data ecosystem for scientific machine learning matures, we believe these could become rich, active research areas. While we don’t answer these questions in this work, we feel that enabling the community to explore these questions is a positive rather than a negative of our work.\n\n---\n### Q3 - MoE options\n\nThere are some appealing conceptual arguments - under the compositionality hypothesis, fields could hypothetically be routed towards modules that implement previously learned physical components. However, while MoE seems like a high potential area of exploration as models scale larger, as the dense models were working quite well, we left such explorations for future work. \n\n---\n### Conclusion\n\nWe greatly appreciate your time and effort both in your initial review and in reading our responses. We enjoyed considering your questions and hope our responses are interesting and informative.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The rebuttal has addressed the reviewers' concerns. Reviewer 1VF5 brought up a good point regarding the exact definition of "new physics". The main contention is whether "new physics" is limited to only the "difference" in the physical parameters or also the difference governing equations. The authors have promised to clarify and properly define this term.
Decision: Accept (poster)