Paper: Multiple Physics Pretraining for Physical Surrogate Models (ICLR-2024)

Reviewer comment (Experiments):
The authors frame their question on the low-data regime (Section 5.2) as “Does MPP provide a fine-tuning advantage over existing spatiotemporal foundation models for new autoregressive prediction tasks?”. Their experimental results indeed support the claim that MMP is superior. However I believe that in order to develop useful foundation models, such models should be compared to the “best one can do” task-specific model in the low-data settings. Monitoring the progress on tasks that are too hard for task-specific baselines (like those implemented in the PDEBench paper) will provide rich insights to the community. The first line of experiments (Section 5.1) does not enable this, as all tasks are used at training time, including the test task. A possible experiment to add to section 5.2 would consist in training one of the PDEBench baseline models on the two low-data tasks.

Author response:
Thanks for the suggestion! We agree that this would be valuable. Since PDEBench does include results for these systems (albeit with access to the full training data), we plan on adding those results in 5.2 as dashed horizontal lines in Fig 5 for comparison so that reviewers can see at which data level our approach surpasses these widely used baselines.

Reviewer comment (Experiments):
The model architecture is completely novel and more experimental validation should be done to maximize the value for the community. Specifically, to separate the added-value of the transformer architecture from the added value of the multi-systems training, it could have been interesting to train the MMP architecture on each of the PDEBench tasks individually, and include it in Table 1 as a new single-task baseline. It is also quite common to include ablation studies to validate intuitions or design choices, and the claim that positional encodings help to learn boundary conditions (Section 4.2, last paragraph) should be validated either by such experiment, or by a proof in the supplementary material.

Author response:
This is also a very good suggestion. We have now broken out the "B" family of models into "train from scratch", "pretraining only", and "finetuned" models. Previously we were only showing "pretraining only" since our goal was to show that, since multi-task pretraining is new for physical dynamics, our approach was able to at least match modern baselines. On the whole though, the pretraining performance can be improved on most systems through finetuning. We note this in the paper, but did not include numerical evidence before.

Reviewer comment (Presentation):
The conclusion reminds the motivation behind multi-systems pretraining, but does not clearly recapitulate which questions have been answered and what is left to future studies. The two paragraphs are not well articulated and I feel would benefit from a rework. Ideally the reader should finish the paper with a clear view of the unanswered questions or limitations to be addressed.

Author response:
Thanks. We agree and appreciate the feedback here. We do discuss limitations, but these could certainly be less about problems faced by the entire field and be more specific to our approach. Improving this within the space constraints is something we need to think a bit more about, but we will update this when we post a revised copy prior to the discussion deadline.

Reviewer comment (Experiments):
Put the low-data results (Section 5.2) in perspective with task-specific baseline performances: are these tasks way too hard for the PDEBench baselines? Ideally, report the performance of one task-specific baseline in Figure 5.

Author response:
This is also a very good suggestion. We have now broken out the "B" family of models into "train from scratch", "pretraining only", and "finetuned" models. Previously we were only showing "pretraining only" since our goal was to show that, since multi-task pretraining is new for physical dynamics, our approach was able to at least match modern baselines. On the whole though, the pretraining performance can be improved on most systems through finetuning. We note this in the paper, but did not include numerical evidence before.

Reviewer comment (Presentation):
Rework the conclusion to simply recapitulate the questions answered and which questions should be explored in future studies.

Author response:
Thanks. We agree and appreciate the feedback here. We do discuss limitations, but these could certainly be less about problems faced by the entire field and be more specific to our approach. Improving this within the space constraints is something we need to think a bit more about, but we will update this when we post a revised copy prior to the discussion deadline.

Reviewer comment (Experiments):
In section 4.2 last paragraph, clarify whether the claim on boundary conditions is justified theoretically, validated by an ablation study, or both. Ideally, add relevant information in the supplementary material.

Author response:
This is a good point. The language is likely a bit strong in the current version and we should tone this down to suggest that it improves multi-task learning across BCs rather than enables zero-shot learning as we only have experimental support for the former. We have performed a small supplementary study on a pair of 1D advection system that differ only in boundary conditions (absorbing vs periodic) and show that our approach has a significant advantages when trained on both systems compared to traditional encodings. This will be referenced in the main text and added to the supplement.

Reviewer comment (Reproducibility):
Code clarity: Papers that propose pretrained models should be as easy as possible to finetune, so that the community can make quick progress towards more efficient pretraining approaches on challenging downstream tasks. I would recommend rewriting the train.py file to remove a lot of boilerplate code and make it easier to reuse and tweak.

Author response:
Thanks. We are making several changes to the format of the repo (especially the README). These should come online in the anonymous repo by the end of discussion. We are still working some API changes to facilitate finetuning as some of the tensor shapes currently make that frustrating, but that is not something that will be completed by the deadlines.

Reviewer comment (Presentation):
Clarity of Figure 2 and its legend can be improved: ReVIN abbreviation for the normalization layer is not defined. “physics metadata” is not defined neither in text nor in figure legend.

Author response:
Thanks for this. We're still looking for a better word internally, but we will update this to mention "field indices" with the caption discussing what these are. On ReVIN - yes, that is also something we need to define. We mention the full word in the main text but didn't actually connect it to the acronym. We will do so both in the main text and in the caption of the figure.

Reviewer comment (Writing):
There is a typo in section 5.1 (word “all” written twice): our models (denoted by MPP-AViT-*) must handle all all systems and regimes without finetuning.

Author response:
Thanks. We will fix this.

Reviewer comment (Experiments):
Firstly, the paper forcibly combines different types of datasets for training without considering the potential conflicts in PDE solutions. For instance, one can construct two PDEs that have identical or minimally differing data within a certain frame count, yet due to the non-linearity of PDEs or inherent differences in the equations, they exhibit substantial differences in subsequent evolution. Unlike Ref [1] which unifies the form of PDEs, this paper's approach to mixed dataset training could lead to the model learning meaningless representations, especially in the presence of conflicting data. From the experiments and provided open-source code, it seems the paper selected equations with vastly different properties for pre-training, which doesn't address this challenge faced in PDE pre-training.

Author response:
There is an important distinction that must be made here: while the benchmark systems are derived from PDEs, our goal is to develop methods that are not dependent on prior knowledge of the equations. In this setting, it is not always possible to restrict the functional form to avoid chaotic or bifurcative behavior because the physical parameters are often outside of our control. This is a challenge faced by all methods in this space that are not simply trying to learn to solve a single PDE.

Knowing we are operating in spaces that can potentially be difficult to distinguish, we made three design choices to mitigate risk:
1. We use a short trajectory of snapshots (16 in our experiments) rather than a single snapshot because of the need to differentiate between systems.
2. Different sources of data are embedded into the shared space independently.
3. We intended for the model to be finetuned for downstream tasks.

Pretraining is intended to learn to identify shared features that the model can refine for specific applications. For example, trajectories from two systems on either side of a bifurcation may follow a similar trajectory prior to diverging, but similar is not identical and the trajectories do evolve differently. (1) ensures that this gap is visible and the pretraining objective encourages the model to learn to differentiate between these small variations. It is not necessary for the model to be able to perform perfect system identification, as it is expected that it will be finetuned to the downstream task where system identification is no longer necessary. 

The only case where the trajectories could realistically be identical is a system described by the weak solution of a PDE. In this case, due to discontinuity, certain terms may only be active in certain regions of phase space. If this region is never seen, like for any data-driven method, our approach will not be able to predict here. However, if this region is seen, since each data source is embedded separately, the model is given a priori knowledge to differentiate between systems even if the underlying equations are quite similar. 

If the reviewer's concern is something that we have not discussed here, please let us know.

Reviewer comment (Experiments):
Although the experimental section is logically well-structured, the descriptions of the experimental results and settings remain unclear. For example, PDEBench provides multiple datasets for CNS M1.0 and CNS M0.1, but the paper doesn't clearly state which dataset was used for training or testing. Besides, the source of the results for other baselines is not clarified, and the comparisons have too few baselines.

Author response:
With regards to the comparisons, our goals in the experiments are two-fold. First, we want to show that pretraining is learning useful behavior by showing that it can perform on par with dedicated modern baselines. The goal is not to show that pretraining alone results in state-of-the-art results. We have made this clearer by breaking out the B family of models into pretrained and finetuned variants to indicate that the pretrained models are not the upper ceiling on performance with these tasks.

Second, we want to show that given the same architecture, pretraining does improve performance on the downstream task. Since our method is the first to demonstrate the ability to learn from multiple physics, we also need to show that learning in this way outperforms learning from arbitrary spatiotemporal data (VideoMAE) as well. 

Thank you for mentioning the CNS labeling. These are obtained by averaging results across the subsets with these parameters. We will add the results per file to the appendix. There is limited space in the main text, and the important distinction was the inclusion of compressible behavior which was not seen in the training corpus, so we aggregated by mach number. The pretrained models are still trained on the training partitions of all subsets and the baselines are still trained per-dataset. The split is performed per file as described in Appendix B.2.

Reviewer comment (Experiments):
It seems that the current model mainly focus on predicting the solution one time step further. The video results (diffre, incompNS, mpp_swe) shown in the supplemental materials exhibits strong checkboard artifacts as the timesteps grow.

Author response:
> one time step further...

> ...checkboard artifacts... numerical unstable...

We're grateful that you examined our supplementary material! We will improve the labeling to make this clear, but the listed videos are actually all pretraining only rollouts. The CNS video is of a fine-tuned model. We've found these checkboard artifacts are addressed for many systems through finetuning as can be observed, though not all.

We suspect these are linked to well-studied stability issues [1, 2, 3, 4] that have been connected to aliasing behavior, though we would note that as you observed in the videos, the model does not seem to completely diverge while the linked works have established that most of the baselines do in similar settings. Long-run stability of these models is an important research direction, but is out of scope here. 

References:

[1] Learned coarse models for efficient turbulence simulation. (https://arxiv.org/abs/2112.15275).

[2] Towards Stability of Autoregressive Neural Operators. (https://arxiv.org/abs/2306.10619).

[3] Pde-refiner: Achieving accurate long rollouts with neural pde solvers. (https://arxiv.org/abs/2308.05732). 

[4] Are neural operators really neural operators? Frame theory meets operator learning. (https://arxiv.org/abs/2305.19913).

Reviewer comment (Experiments):
The proposed model can only handle simulations on structured mesh. However unstructured mesh (e.g., triangular, tetrahedra) is a more common choice in real world simulations. In addition, the multi-resolution i.e., mesh with adaptive resolution, which is also a common technique in large scale simulations, has not been taken into consideration.

Author response:
> ...unstructured mesh (e.g., triangular, tetrahedra) ... adaptive resolution...

Thank you for pointing this out! Natively handling arbitrary meshes is a very exciting problem both in machine learning and numerics and certainly is not solved in either space. Outside of GNNs, most deep learning methods, including all of the baselines discussed, assume uniform grids either natively or via interpolation. We fully agree that handling arbitrary meshes would be a valuable contribution, but our work is focused specifically on the pretraining task.

Reviewer comment (Reproducibility):
It is mentioned in the appendix that the model is trained for 500 epochs for one task. How long does the training process take?

Author response:
> How long does the training process take?

We are planning on updating the language to switch from "epoch", which we arbitrarily define as 2000 optimization steps as we are sampling with replacement, to listing the number of steps for clarity. 

The base model (AViT-B) is pretrained using ~960 GPU hours on H100s, though finetuning is comparable to the UNet and ORCA baselines. The inference time of each model on an A6000 GPU is listed below (we plan on adding this to an appendix since it is important information to convey)

| Model | Time (ms) |
| -------- | -------- |
| UNet     | 67.7      | 
| FNO     | 7.2    | 
| AViT-Ti     | 83.5    | 
| AViT-B     | 105.6    | 
| ORCA | 98.5 |

Reviewer comment (Evaluation):
Are there any intuitive explanation for the checkboard artifacts as the timesteps grow shown in the video? Does it mean the proposed method is numerical unstable to some extent?

Author response:
> one time step further...

> ...checkboard artifacts... numerical unstable...

We're grateful that you examined our supplementary material! We will improve the labeling to make this clear, but the listed videos are actually all pretraining only rollouts. The CNS video is of a fine-tuned model. We've found these checkboard artifacts are addressed for many systems through finetuning as can be observed, though not all.

We suspect these are linked to well-studied stability issues [1, 2, 3, 4] that have been connected to aliasing behavior, though we would note that as you observed in the videos, the model does not seem to completely diverge while the linked works have established that most of the baselines do in similar settings. Long-run stability of these models is an important research direction, but is out of scope here. 

References:

[1] Learned coarse models for efficient turbulence simulation. (https://arxiv.org/abs/2112.15275).

[2] Towards Stability of Autoregressive Neural Operators. (https://arxiv.org/abs/2306.10619).

[3] Pde-refiner: Achieving accurate long rollouts with neural pde solvers. (https://arxiv.org/abs/2308.05732). 

[4] Are neural operators really neural operators? Frame theory meets operator learning. (https://arxiv.org/abs/2305.19913).

Reviewer comment (Experiments):
Does the proposed method have the potential to be applied to 3D simulations while keeping the efficiency and accuracy? (as the max memory usage of current model is ~60GB)

Author response:
> limited to 2-D cases... unclear whether the proposed model... more realistic 3D physical simulations.

> ...applied to 3D simulations while keeping the efficiency and accuracy? 

For introducing this method, we felt it was best to evaluate on a pre-existing benchmark within the community. 2D simulations are conventionally used for this as they elicit complex behavior without the extreme cost of 3D methods. We do choose our base architecture with scaling in mind, but since we're introducing fundamentally new capabilities to deep learning for physical systems (our approach is the first to demonstrate the ability to learn multiple nonlinear time-dependent dynamics simultaneously), it was important to place it in context.

We do agree that 3D fluids are a more realistic setting and did account for this future goal in our architectural choices. Using 3D patching, we are able to train on the (5x128x128x128) Compressible Navier-Stokes (CNS) AViT-B in mixed-precision at a memory cost of 43.90 GB per sample. Training from scatch is still enough to surpass the PDEBench baselines: 

| Model | NRMSE CNS (Turb) |  NRMSE CNS (Rand)|
| -------- | -------- | -------- |
| UNet     | 1.00      |  .23|
| FNO     | .37    |  .24 | 
| AViT-B     | .32   | .18 | 

However, our goal here is exploring transfer and there are not yet accepted benchmarks that cover a diverse set of 3D problems. Transfer from 2D to 3D is non-trivial due to the downsampling assuming a fixed patch dimension, but dimension-agnostic "patching" is a promising future research direction for further expanding multiple physics pretraining and our design choices do enable this type of exploration.

Reviewer comment (Evaluation):
Although the generalizing capability is significant, the computational cost of the proposed models is still concerning.

Author response:
These are good points. While the purpose of pretraining methods is to develop multi-use models where the pretraining cost is borne by the pretrainers, this is important information and we have added an appendix section showing this. For reference, the pretraining for AViT-B currently takes ~960 GPU hours in H100s in single precision training (we plan on release the weights to amortize cost for users). During finetuning, the cost is more comparable to baselines. The time per forward pass on an A6000 GPU is listed below:

| Model | Time (ms) |
| -------- | -------- |
| UNet     | 67.7      | 
| FNO     | 7.2    | 
| AViT-Ti     | 83.5    | 
| AViT-B     | 105.6    | 
| ORCA | 98.5 |

Our AViT code is currently unoptimized - it is for instance possible to fuse the axial attentions into a single cuda kernel launch. Despite that, the runtime of B is comparable to ORCA while Ti is comparable to the UNet. The FNO from the benchmark is quite a bit faster, but the performance is competitive only on SWE.

Reviewer comment (Evaluation):
How do the authors expect the computational cost of the proposed models and how does it compare against other baselines in Table 1?

Author response:
These are good points. While the purpose of pretraining methods is to develop multi-use models where the pretraining cost is borne by the pretrainers, this is important information and we have added an appendix section showing this. For reference, the pretraining for AViT-B currently takes ~960 GPU hours in H100s in single precision training (we plan on release the weights to amortize cost for users). During finetuning, the cost is more comparable to baselines. The time per forward pass on an A6000 GPU is listed below:

| Model | Time (ms) |
| -------- | -------- |
| UNet     | 67.7      | 
| FNO     | 7.2    | 
| AViT-Ti     | 83.5    | 
| AViT-B     | 105.6    | 
| ORCA | 98.5 |

Our AViT code is currently unoptimized - it is for instance possible to fuse the axial attentions into a single cuda kernel launch. Despite that, the runtime of B is comparable to ORCA while Ti is comparable to the UNet. The FNO from the benchmark is quite a bit faster, but the performance is competitive only on SWE.

Reviewer comment (Experiments):
Is the model applicable to 3D simulations?

Author response:
This is something we are excited to explore in future work. The architecture as a whole is selected with scalability in mind and can be run on 3D data with minor changes, but established benchmarks for evaluating cross-system transfer in 3D  do not currently exist. Using mixed precision, we can fit an AViT-B on the largest 3D system in PDEBench (5x128x128x128) with ~43 GB of VRAM. Transfer from 2D to 3D has additional complications, like the dimension-specific downsampling (patching) process.

Reviewer comment (Writing):
How do you define the inverse problem and what model was used for the tasks?

Author response:
Thanks, we agree with these criticisms and will update the manuscript with a more detailed description.

There are two parameter inference problems discussed here. Both are tackled using the pretrained AViT-B model. The first is the forcing identification in INS. In equation 9 in the appendix, the INS equations are described with a forcing **f**. This **f** is randomly sampled per trajectory and constant through the simulation. In this task, we try to predict what this forcing is by finetuning (or training from scratch) the AViT-B model from Experiment 1. 

The second is identifying the buoyancy parameter which is a single scalar serving a similar role in an outside simulation entirely separate from pretraining. We will update the appendix to include this precise equation as well for reference in the main text.

Reviewer comment (Writing):
Typo in section 5.1: must handle all all systems and regimes without finetuning.

Author response:
This has been fixed, thanks!

Reviewer comment (Experiments):
More attempts can be made to address the transferability problems between different types of physics equations. For instance, when discussing the Navier-Stokes (NS) equations, whether at high or low Reynolds numbers, the equation forms are quite similar, which allows for the investigation into whether a single model still possesses robust merged learning capabilities for more different classes of equations, such as diffusion equations and wave equations.

Author response:
These are interesting suggestions! While we agree that more diverse exploration would be valuable, multi-task training is new to this space, so in this initial paper, we wanted use an established benchmark that has been previously used for analysis in this space. We agree that exploring the limits of transfer is an exciting area going forward.

About the Reynolds remark, the CNS data in PDEBench does actually include a range of Reynolds numbers - we do not report them as we believe the simulations have significant numerical diffusion due to insufficient resolution which we cannot account for in the computation of the effective Reynolds numbers. We do want to make this range clearer and we updated the appendix to reflect the results per nominal viscocity coefficient with the numerical diffusion caveat included in the text.

Reviewer comment (Novelty):
What are the advantages and disadvantages of this new approach compared to the Finite Element Method (FEM) for systems with explicit control equations or empirical formulas? One of my biggest problem is whether the physics task should be treated as a purely data-driven problem, or if it ought to incorporate certain explicit priors or equation-based guidelines.

Author response:
This is an important distinction and one we will try to address in the discussion of neural PDE solvers vs operators in the related work. We do consider this to be a data-driven method. In this setting, PDEs are simply convenient testbeds rather than the true target. Numerical methods like FEM are clearly superior in both accuracy and predictability for tasks where the physics can be fully resolved by numerical methods. The advantage of data-driven methods is applicability to faster surrogates for difficult-to-resolve system or for learning directly from observational data in less controlled settings. As we note in our intro, these tend to also be areas where we have limited training data, which is why one of our focuses is transfer in the low-data regime.

Explicit physical priors are another valuable approach in this regime, but recently vision and language have been finding success with less constrained models that are pretrained on large amounts of data, but it has not been clear how to do the same in a dynamics setting. Our work is the first that has seen some success for nonlinear dynamics. We believe both paths to be interesting research directions.

Reviewer comment (Experiments):
In the appendix, experimental data from 1-step to 9-step show little change in the field, whether looking at ground truth or predicted solutions. If the selected time steps were longer, would the model still be able to accurately predict future changes?

Author response:
The answer to this is likely task-specific. For our method and most methods in this space there is likely going to be a trade-off between small steps that rapidly accumulate autoregressive error and large steps which render the problem too difficult. Our goal in experiment 1 was to show that our pretrained models learn the system at least as well as modern baselines, so we used the benchmark steps which vary considerably system-to-system. CNS, for example, has highly visible changes step-to-step and is one of the systems where we see some of the largest improvements by percentage. This does indicate that our method performs no worse in this regard than standard approaches.