OmniArch: Building Foundation Model For Scientific Computing

TL;DR

OmniArch presents a foundation-model approach to scientific computing that unifies multi-scale and multi-physics PDE solving via a Fourier-domain encoder-decoder and a Transformer backbone, augmented with a PDE-Aligner for physics-informed fine-tuning. The method enables 1D-2D-3D united pre-training and demonstrates in-context and zero-shot generalization across 11 PDE types, achieving state-of-the-art nRMSE. Key innovations include a frequency-domain representation with TopK mode selection, a Temporal Mask that allows full intra-timestep coupling among multiple physical quantities, and a PDE-Aligner that enforces conservation-like constraints in the frequency domain. The results show substantial accuracy gains, multi-scale inference capabilities, and effective inverse-problem and zero-shot performance, suggesting significant practical impact for fast, scalable PDE solving, while acknowledging remaining challenges in 3D tasks and interpretability.

Abstract

Foundation models have revolutionized language modeling, while whether this success is replicated in scientific computing remains unexplored. We present OmniArch, the first prototype aiming at solving multi-scale and multi-physics scientific computing problems with physical alignment. We addressed all three challenges with one unified architecture. Its pre-training stage contains a Fourier Encoder-decoder fading out the disharmony across separated dimensions and a Transformer backbone integrating quantities through temporal dynamics, and the novel PDE-Aligner performs physics-informed fine-tuning under flexible conditions. As far as we know, we first conduct 1D-2D-3D united pre-training on the PDEBench, and it sets not only new performance benchmarks for 1D, 2D, and 3D PDEs but also demonstrates exceptional adaptability to new physics via in-context and zero-shot learning approaches, which supports realistic engineering applications and foresight physics discovery.

Figures (15)

• Figure 1: OmniArch achieves state-of-the-art performance (nRMSE Loss) on 1D-2D-3D PDE tasks with single foundation model. The baselines include the task-specific expert models and the pre-trained models.
• Figure 2: The overview of OmniArch. The Fourier Encoder converts coordinates and physical fields into frequency domains, enabling unified training for 1D, 2D, and 3D data. Reserved frequency modes form PDE token embeddings for Shared Transformer Blocks. Tokens are grouped by timestep to create a Temporal Mask for prediction. Predicted modes are decoded using IFFT with zero padding to recover the physical field.
• Figure 3: (Left) PDE-Aligner architecture with Fourier Encoders for initial/current state, and PDE Caps Encoder enforcing consistency via PDE constraints. (Right) Fine-tuning OmniArch with PDE-Aligner on downstream PDEs like Navier-Stokes equations for physics-informed learning.
• Figure 4: The multi-scale capability.
• Figure 5: Zero-shot prediction results (Rollout) of OmniArch-L and MPP-L on KH dataset. Displaying time steps T+1 to T+6, the top row shows ground truth data, while the middle and the bottom row illustrate MPP-L's and OmniArch-L's predictions respectively.