Multiple Physics Pretraining for Physical Surrogate Models (2310.02994v2)

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.

• The paper introduces a novel MPP framework that leverages a unified embedding and scalable transformer for physical surrogate modeling.
• It employs autoregressive next-step prediction and normalized MSE to achieve robust generalization and excel on fluid mechanics benchmarks.
• Demonstrated transferability allows MPP models to outperform traditional methods in low-data regimes and diverse physical systems.

Overview of "Multiple Physics Pretraining for Physical Surrogate Models"

The paper presents a novel approach to the pretraining of large-scale surrogate models for predicting the dynamics of physical systems through a method called Multiple Physics Pretraining (MPP). At its core, MPP is an autoregressive, task-agnostic approach developed to handle diverse spatiotemporal data, particularly focused on fluid mechanics. The researchers aim to leverage the benefits of foundation models—prevalent in fields like natural language processing and computer vision—by applying similar pretraining strategies to the domain of physical systems governed by Partial Differential Equations (PDEs).

Key Contributions

The authors introduce several innovative strategies and techniques as part of their MPP framework:

1. Unified Embedding for Heterogeneous Systems:
   • MPP proposes a shared embedding space where multiple physical systems project their state variables, regardless of the underlying dynamics or discrepancy in spatial and temporal resolutions. This shared space is facilitated through reversible instance normalization and parameterized field embeddings, allowing cross-domain applicability without task-specific architectures.
2. Scalable Transformer Architecture:
   • Central to MPP is the Axial Vision Transformer (AViT) architecture, which emphasizes scalability by employing axial attention mechanisms. This design efficiently handles spatiotemporal data by decoupling attention operations across spatial and temporal axes, enabling larger input sizes and higher-resolution data without a prohibitive computational cost.
3. Resilient Pretraining Objective:
   • The paper utilizes autoregressive next-step prediction as a primary objective. The normalization of the mean squared error (NMSE) across different scales in multiple systems ensures that learning signals remain balanced, allowing the model to generalize across disparate tasks.
4. Exemplary Performance on Multiphysics Surrogate Modeling:
   • MPP demonstrates exemplary performance, where a single pre-trained model matches or even surpasses specialized baseline models across several fluid mechanics benchmarks. It excels without task-specific finetuning, underscoring the robustness of its learned representations.
5. Transferability Beyond Original Domains:
   • When fine-tuning on new and data-limited systems, MPP outperformed models trained from scratch and conventional video foundation models. This transferability indicates MPP’s potential utility in low-data regimes often characterizing complex physical systems.

Numerical Results and Claims

The paper provides robust numerical validation of the MPP framework. MPP models surpass existing models, such as UNet and FNO, in terms of accuracy and efficiency across tasks like the compressible and incompressible Navier-Stokes simulations, among others. Notably, the pretrained models achieve high performance on specialized tasks, thereby challenging the need for finetuning procedures traditionally necessary in scientific machine learning.

Implications and Future Directions

The implications of this work are significant, as it suggests the feasibility of developing generalized foundation models within the field of physical science, models capable of being fine-tuned with minimal additional data. This has potential ramifications for vastly enhancing computational efficiency and model generalization across numerous scientific and engineering applications. The pre-trained models promote a better understanding of spatiotemporal dynamics, crucial for domains that are traditionally data-scarce or computationally expensive.

Future work may extend into enhancing the resolution and complexity of the input data, further developing the transformer architectures to accommodate diverse grid types beyond uniform discretization, and evaluating the integration of MPP models with existing mechanistic models for predictive augmentation.

In conclusion, this paper effectively pioneers a methodological bridge from foundational domain-agnostic models to the field of physical sciences, underscoring a promising horizon for the incorporation of learned surrogate models within the broader landscape of physics-driven AI research.