OmniArch: Building Foundation Model for Scientific Computing
Authors: Tianyu Chen
Haoyi Zhou
Ying Li
Hao Wang
Chonghan Gao
Rongye Shi
Shanghang Zhang
Jianxin Li
Venue: ICML 2025 poster (2025)
Keywords: AI for Science
Partial Differential Equations(PDEs)
Foundation Model
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: We build the foundation model for PDE solvers with 1D-2D-3D united pre-training and exhibit zero-shot capability.

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.

OmniArch: Building Foundation Model For Scientific Computing
Tianyu Chen
Haoyi Zhou'
Ying Li Hao Wang² Chonghan Gao
Rongye Shi
Shanghang Zhang?
Jianxin Li
Abstract
①
Multi-Physscs
Foundation models have revolutionized language modeling, while whether this success repli- cated in scientific computing remains unexplored. We present OmniArch, the first prototype aim- ing at solving multi-scale and multi-physics sci- entific computing problems with physical align- ment. 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 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 per- formance benchmarks for 1D, 2D, and PDEs but also demonstrates exceptional adaptability physics via in-context and zero-shot learning approaches, which supports realistic engineering applications and foresight physics discovery.
U-Nitt
PNO
te-2
Possidonç0 Only
DPOT/2D10))
Arbitrary Combination?
Multi-Scale
OmniArch
G$ Physics Alignment
Predicted State
162
Equations? ?
2D PDEs
3D PDEs
Figure 1: OmniArch achieves state-of-the-- performance (nRMSE Loss) on ID-2D-3D PDE tasks with single foun- dation model. The baselines include the task-specific expert models and the pre-trained models.
on state-of-the-ar clusters. To expedite PDE solving, pioneers have explored the construction of neural operators that learn mappings between function spaces, offering the potential generalize across different discretizations. For the requisite precision, neural operators are often enhanced with physics-informed normalization techniques, such as customized loss functions derived from the governing physical equations (Raissi et al., 2019).
Computing (HPC)
1. Introduction
Developing robust neural surrogate models for temporal partial differential equations (PDEs) is crucial for various scientific and engineering applications, including aircraft design, weather forecasting, and semiconductor manufac- turing (Allen et al., 2022; Pathak al., 2022). These PDEs describe spatial-temporal dynamic systems that are founda- tional to these industries. Traditional scientific computing methods, such as Finite Element Methods (FEMs) and Finite Volume Methods (FVMs) (Oden, 1989), require extensive handcrafted coding and are computationally intensive, even
The primary limitation of neural operator methods lies in their case-specific design, restricting their application scope and hindering broad transferability across diverse physical systems. Recent efforts aim to enhance the transferabil- ity neural operators by developing foundational models that leverage advancements in learning strategies, archi- tectural design, and data curation (Alkin et al., 2024; Sun et al., 2024: Shen al.). terms of learning. the pre-train and fine fine-tune tune paradigm, proven effective for Fourier Neural Operator (FNO) models (Subramanian et al., 2023), has been adapted to PDE contexts. Additionally, Lic group- based self-supervised learning (Lie-SSL) (Mialon et al., 2023) introduces physics-cone constrained transformations for PDEs, primarily addressing inverse problems. Architec-
'KLCCSE, School of Computer Science and Engineering, Beihang University, Beijing, China "hina'SKLMM SKLMIP, School of Computer Science, Peking University, Beijing, China School of Artificial Intelligence, Beihang University, Beijing, China. Correspondence to: Jianxin Li <lijx @ bbu...cdu.cn>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

OmniArch: Building Foundation Model For Seientific Computing
turally. innovations like ICON_LM (Yang et al., 2023b), and PITT (Lorsung et al., 2023) incorporate language model principles to enhance neural operator learning, en- abling generalization through equation captions. The Fact- former (Li et al., 2023) introduces a scalable transformer for multi-dimensional PDE data, with the Multi-Physics Pre-training (MPP) (McCabe et al., 2023), Poseidon (Herde et al.. 2024) and DPOT (Hao et al., 2024) further extend- ing this approach to 2D data pre-training- From a data- centric viewpoint, resources such as PDEBench (Takamoto et al., 2022), PDEArena (Gupta& & Brandstetter, 2022) and The-Well (Ohana et al., 2024) offer well-structured datasets that facilitate pre-training and the establishment of rigorous benchmarks.
contributions and findings include:
We introduce OmniArch, the first foundation model successfully conduct ID-2D-3D united pre-training. Using a Fourier Encoder-decoder, OmniArch allows for flexible grid inputs, enabling unified multi-scale training. The Temporal Mask effectively addresses inconsistencies in multi-physscss systems, allowing dif- ferent physical quantities and time steps to be learned simultaneously within a shared Transformer backbone.
We develop the PDE-Aligner for physics-informed fine- tuning, which leverages hidden representations of equa- tions and other physical priors to align with observed physical field dynamics.
While attempting unified learning of multiple PDE solvers in single model, multi-scale and multi-physic challenges persist. The above surrogate models, often constrained by the fixed mapping grid (MPP, Lie-SSI. ICON_LM, and single-time step observation window (MPP. Factformer, PITT. Poseidon), struggle with flexible spatial grid input and long-sequence roll-out predictions.
After fine-tuning, OmniArch achieves state-of-the-an performance on 11 types of PDEs from the PDEBench and PDEArena benchmarks. The model exhibits in- context learning capabilities and demonstrates promis- ing zero-shot performance.
2. Related Works
In this work, we study how to frame the foundation model learning paradigms for Scientific Computing tasks w.r.t PDEs, namely OmniArch. For the pre-training stage, we de- fine a flexible pipeline to deal with muultiple-pyssiee spatial- temporal data and convert the forward problem learning into popular auto-regressive tasks that can be scaled up easily. For the pre-training stage, we devise a flexible pipeline to handle multi- multi-ppyysies spatio-temporal data and reformulate the forward problem as scalable autoregressive tasks. Specif- ically, we employ Fourier encoder to convert coordinate and observation data into frequency components (modes). We use truncated modes to form PDE token embeddings, sequenced for processing by transformer blocks, and we design the PDE-Aligner during fine-tuning to align predic- tions with known physical laws and principles, improving the model concordance to conventional physical constraints.
Learned PDE Solvers. Deep Learning for solving PDEs has been a recent focal point of research (Lu et al., 2021 bb; Karniadakis et al., 2021), including physics-informed meth- ods (Raissi et al., 2019), GNN-based techniques (Veličković et al., 2017: Pfaff et al., 2020), and neural operator models like DeepONet (Lu et al., 2021a) and FNO (Li et al., 2020). While effective, these models often require task-specific training and struggle with generalization. ICON LM (Yang et al., 2023a). MPP (MeCabe et al., 2023), PDEformer- (Ye al., 2024) aim to generalize across diverse physical systems but limit to single dimension.
Foundation Models for Science. The Foundation Mod- els (Devlin et al., 2018; Brown et al., 2020; Radford et al., 2019; 2018; Touvron et al., 2023; Radford et al., 2021) have emerged as pivotal elements in the field of natural language processing, computer vision, and cross-modal tasks. After large-scale pre-trained with the transformer backbone, they serve the bedrock for multitude of downstream tasks by fine-tuning (Zhang et al., 2023) or in-context learning (Li, 2023). Recently, they have shown promise in scientific fields, exemplified by FourcastNet (Pathak et al., 2022) for weather forecasting, OpenLAM (Zhang et al., 2022) for chemistry, and HyenaDNA (Nguyen et al., 2023) for biomedical tasks. However, applying foundation models to scientific computing, particularly PDE solving, remains an emerging and pioneering area.
We release our models' base and large variants', concur- rently addressing ID. 2D, and 3D PDEs. Evaluating perfor- mance across 11 PDE types from PDEBench and PDEArreaa our OmniArch achieves state-of-the of-the--rt results, as illustrated in Figure For the Computational fluid dynamics (CFD) related tasks, we observe one two orders of magnitude re- ductions in normalized root mean squared error. Moreover, our models exhibit emergent capabilities, such as zero-shot generalization to novel PDE systems and in in-context learn- ing of neural operators. The representations learned by OmniArch demonstrate versatility, readily adaptable to in- verse problems. Notably, OmniArch facilitates multi-scale inference, accommodating a range of input grid resolutions with moderate precision trade-offs. In summary, our key http://pppn..polaaaaccccaaaa/a/e
3. Method
The foundation models (Devlin et al., 2018; Brown et al., 2020; Radford et al., 2019; 2018; Touvron et al., 2023;
2

---

OmniArch: Building Foundation Model For Scientific Computing
Coordinates Observations (V,U,..)
Fourier Encoder
Fourier Decoder
Predict Physics
Predict Embeddings
(Shifted Right)
1D
02
02
F3
D
P4
0,
Blocks
Multi-Head
Transformer (Shared Layers)
Attention
2D
Temporal
Vi
U1
V22
U2
V3
U1
Taken Embeddings
3D
(by Time Steps)
3D
wm VMMVMM
Projection
Projection
IRk - C
www
IFFT
Frequency
C -
Frequency
Top-K
Zero
Reserved Modes
Predicted Modes
Padding
Fourier Encoder
Fourier Decoder
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 Temporal Mask for prediction. Predicted modes are decoded using IFFT with zero padding to recover the physical field.
Figure 2:
Radford et al., 2021) have shown significant success with broad generation to various inputs and downstream tasks. Building similar model for scientific computing should require addressing dynamic and complex physical systems and learning intrinsic laws from wild physical phenomena. We highlight the major challenges three-fold:
The proposed OmniArch Model follows the predominant trrainingg-henffin---ee paradigm. In subsection 3.1, we utilize Fourier Encoders and Decoders address the multi-scale challenge and employ the Temporal Attention mechanism to handle multi-physics generalization problems. In subsection 3.2, we leverage the PDE-Aligner in the fine- tuning stage, allowing the incorporation of physical priors in textual form into the model's learning and adaptation process.
Multi-Scale The ability to handle inputs of different dimen- sions (ID, 2D. 3D), varying grid resolutions, and diverse grid shapes. For example, fluid dynamics simulations can range from simple one-dimensional pipe flow complex three-dimensional turbulent flow, and the model must main- tain accuracy and consistency across these different scales,
3.1. Pre-training OmniArch: Flexibly Learning from Different Dynamic Systems
Multi-Physics The capability to handle dynamic systems involving different physical quantities. For instance, me- teorology, multiple physical quantities such as wind speed, temperature, and humidity interact, requiring the model process these different physical fields simultancously.
The overall pre-training framework of OmniArch is illus- trated in Figure 2. For physical data different dimensions (ID, 2D, 3D), form their coordinates and observed physical field values into the frequency domain. High and low frequencies are truncated in the frequency domain so that data from differ- ent grids have the same length of embedded representations. Then, these representations are processed through shared Transformer modules model the integral operators along the time axis. We leverage the Temporal Mask ensure that each physical quantity can simultaneously attend to all physical quantities and previous time steps. Finally, the predicted embedding representations are used to recover
separate Fourier Encoders trans-
Physical Alignment Allowing flexible incorporation of physical priors such governing equations, symmetries, conservation laws, and boundary conditions into the solu- tion process. For example, in heat conduction problems, the law of conservation energy and boundary conditions is crucial for predicting temperature distributions.
3

---

OmniArch: Building Foundation Model For Scientific Computing
the predicted frequency domain signals. We involve zero- padding to keep these signals with the target physical field shape and perform individual inverse Fourier transforms to output the corresponding physical field predictions.
dure of k-thh largest components üg (k: () as:
üx (k: t) = TopK [FT V V[u(x (), u(x 201
))
(I)
where (.p denotes the selection operator over F com- ponents, (-) denotes the linear projection for the dimen- sion alignment and the FFT(.) operator is performed the individual time step.
3.1.1. ENCODER/DECODER IN FOURIER DOMAIN
The multi-scale challenge needs proper representation of inputs from different dimensions, varying grid resolutions, and shapes. Inspired by the Fourier transforms (Brigham, 1988) convert the sequential signals into frequency compo- nents, we re-organize the multi-scale inputs in the spatial domain into the multi-component ones in the frequency domain. The traditional pipeline includes convolutional encoders (Raonic et al., 2023), which capture the local features in separated dimensions while the global infor- mation exchange happens at the channels' explicit mixing. The results of Fourier transforms are complex coefficients that measure the magnitude and phase of decomposed pe- riodic components and the global information is naturally weighted, which also applies to the complex boundary con- ditions and heterogeneous grids. Based on that, we further introduce the filter-like components selecting mechanism that distinguishes the high-frrquency (detailed variations) and low-frequency (overall trends) ones in physical inputs, which may maintain different patterns and distribution ratios among the local and global representation. Thus, we can build a universal representation with different resolutions and grid shapes in one flexible network architecture.
In the decoding stage, the predicted frequency domain fea- preed (k:, () are adapted to the target shape using zero tures Vkx"
padding. Then, the inverse Fourier transform (IFFT) is applied to revert the frequency domain features (Pro t) back to the spatial domain, ultimately obtaining the pre- dicted physical field u (z((() tt + 1) as:
r mod (a (1) ,t 1)
V (IFFT( Pero-Pddding [reeed (1, (1,4)...., u (K, [K,)]] ()])))
This encoding and decoding process in the frequency do- main is maintained throughout the whole OmniArch net- work. Since the encoding and decoding operations are al- ways conducted along specific dimensions, thus we omit the d-th index indicator in the following context.
3.1.2. TRANSFORMER AS AN INTEGRAL NEURAL OPERATOR
perspective, the forward proce- dure of Fourier Encoders can be implemented through the Fast Fourier Transform (FFT) with the O(N log N) com- plexity while the convolution operation ends in O The sparsity and separability of frequency domain features fa- cilitate the subsequent Transformer modules in efficiently processing temporal information, reducing the model's 's pa- rameters and computational overhead for better training and inference efficiency.
From a computing-efficien
To achieve multi-physics versatility, we leverage the Trans- former backbone to simulate integral neural operators. In physics, multi-physics systems often exhibit complex spatio- temporal dependencies, requiring effective long-range de- pendency modeling. The multi-head self-attention mecha- nism of the Transformer, with the introduction of the Tem- poral Mask, allows each time step to attend to all physical quantities at the same and previous time steps, enabling efficient temporal information integration. This design en- sures the robustness and adaptability of the model in multi- physics systems. Additionally, by padding variable-length sequences, systems with different numbers of physical quan- tities can use the model for temporal regression predictions in batches, ensuring accuracy and stability.
Let le E Txxx stand for the physical field inputs. If we have real-valued input u(z(d)) t) E IR from d-th index and t-th time step, the Fourier Encoder firstly applies FFT to convert it from the spatial domain to the frequency domain. Note that D is the total dimension and d denotes the sequen- tial index (1, (1.,,3... -), for example, D = D2 + D2 + D3 for ID. 2D and 3D inputs. Then we have the frequency do- main representation L/ E CTxFxx x1 after traversing through all time steps and dimensions. As previously discussed, we design fikter-like mechanism by applying the TopK selec- tion on all F components (modes) in the frequency domain. For the t-th time step, all the K significant (K < F) compo- nents 0 (t) are retained and form the truncated frequency domain. To be clarifying, we can present the forward proce-
Moreover, the autoregressive mechanism of the Transformer bears a strong mathematical resemblance to traditional multi- step methods for solving equations. Traditional multi-step methods approximate solutions iteratively, capturing the dynamic changes of the system. Similarly, the multi-head self-attention mechanism of the Transformer models the global dependencies at each time step, achieving precise capture of dynamic changes in the system.
Specifically, traditional multi-step methods for solving equa-
4

---

OmniArch: Building Foundation Model For Scientific Computing
Seapshot
Fourler Encoder
POE Snapshot
PDE-Aligner
OmniArch
PDE Snapshot
1D
Fourier Encader
Fourier Decoder
Coreent State
lebaal State
20
[aph Aayg (K: fo- (A44,
MM
Lugq
Liston
Transformer
3D
Mac(K,4a 4a)
PDE Caps Eacodar
Fourier Encoder
Advection equation
Burgers equation
Novier Stokes equation
PDE Encoder
Navier-Stokes Stokes equation Ap+ sdd 19 tp al a0Re
Navier
Initial State
R+
Add+ w-the 5--000
E (P)
Similarity(? Aug))
Energy Conservation
PDE-Aligner Training
Fine-tuning OmniArch
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 downstream PDEs like Navier- Stokes equations for physics-informed learning.
tions can be expressed iteratively as:
where and represent the i-th and j-th tokens in the se- quence, and [ represents the time step. Unlike standard causal masking that enforces strict sequential dependencies, our Temporal Mask enables all physical quantities within the same timestep to attend each other, addressing the fundamental coupling inherent in multi-physics systems. Specifically, for system with C physical quantities at each timestep, tokens {i, i + 1,..., i + C 1} corresponding to timestep have full visibility of each other (intra-timestep attention), while maintaining causal relationships across timesteps (inter-timestep attention). This hierarchical atten- tion pattern ensures that coupled physical quantities- velocity and pressure in fluid dynamics evolve while respecting temporal causality. The design is particularly crucial for systems where physical variables must satisfy simultaneous constraints (e.g., continuity equa- tions in Navier-Stokey, that cannot be properly modeled through sequential token processing.
pred (x, + 1) u(x, () + AL f(u(x,
(3)
In contrast, the autoregressive mechanism of the Trans- former updates the current state by a weighted sum of previ-
ous time steps through attention weights:
prod x 1)
at, i(x, = uuxx -
(4)
where ai,t refer to the attention weights. Both approaches update based on previous time steps, with the attention mechanism acting neural surrogate (Sun al., 2020) for the integral operator f.
jointly
Assume that we have physical system with two physical quantities u(I, t) and v(zz t), where the total number of quantities is recorded by = In OmniArch's 's computa- tion, the frequency domain features å (k, t) and 0x (k: t) obtained from the Fourier Encoder are further transformed into real-valued embeddings through R(-), resulting in the input embeddings for the Transformer U, and V.. These embeddings are grouped by time steps to form the input sequence Z, For each time step 1,
The input sequence then passes through multiple shared Transformer blocks, outputting the shifted right predicted feature sequence for each time step (2,)T+] T+1.
T+1 1=
TransformerBlocks/ ( M) (7)
Due numerical differences between dynamic systems, use nRMSE to calculate the loss function Lsim for batch during training:
Z. = (U,, V.) = (R(ūx (k, t)), ((ox(k,t))
(5)
The Temporal Mask M ensures that each time step & can access all physical quantities the current and previous time steps, which is defined as:
werodd (x, - u(x 1)
I =
sim
B Ea
(8)
if
M(i.))
(6)
if
Lsim
C
5

---

OmniArch: Building Foundation Model For Scientific Computing
This design can effectively capture the temporal evolution of physical fields, achieving high-precision dynamic sys- tem predictions and ensuring that systems with different numbers of physical quantities can adapt to this model for temporal regression predictions.
PDE-Aligner serves as physics-aware discriminator, helping OmniArch distinguish between different physical regimes encountered during pre-training. The fine-tuning loss La = Lsi Leg encourages predictions that are both accurate (via mim) and physically consistent with the spec- ified PDE system (via Leq), effectively steering the model toward the correct physical behavior among many learned dynamics.
3.2. Fine-tuning OmniArch: Enabling Physics-Informed Learning via Equation Supervision
The PDE equations are natural and intuitive 'supervision' methods for real-world physical phenomena. To perform the physical alignment, we incorporate the PDE-Aligner to achieve physics-informed learning. Unlike the pre-training stage, the OmniArch is designed to comply with specific physical laws during fine-tuning. As illustrated in Figure 3 (left), the PDE-Alignes employs a contrastive learning paradigm in the frequency domain.
4. Experiments
4.1. Dataset and Baselines
Dataset. We collect ID, 2D, and 3D datasets from the pub- lic PDEBench and PDEArena. The ID datasets include: (1) CFD, generated by the compressible Navier-Stokes equa- tion with velocity (Vz), density, and pressure. (2) Bur. the Burgers' equation with velocity. (3) Dift., the diffusion- sorption equation with concentration (p). (4) Adv., the ad- vection equation with velocity (V²). (5) Reac. the reaction- diffusion equation with concentration (p). The 2D datasets include: (6) CFD, generated by the compressible Navier- Stokes equation with velocities (Vz, Vy),, density, and pres- sure. (7) Reac., the craction-diffusion equation with activa- tor (u) and inhibitor (v). (8) SWE, the shallow- water equa- tion with velocities (h). (9) Incom generated by 2D Inho- mogeneous, Incompressible Navier-Stokes equations, with velocities (Vz, V2y) and particles. The 3D datasets include: (10) CFD, generated by the compressible Navier-Stokes equation with velocities (Vz, Vgg Vz), density, and pressure. (11) Maxw., the Maxwell equation with electric displace- ment (D,, Dy, Dz) and magnetic field (H,, Hy. H2). More details can be found in Appendix C.
The key insight is that physical evolution manifests dis- tinctively in frequency srace-conservationn laws constrain energy distribution across modes, while different PDEs ex- hibit characteristic spectral signatures. By operating in this domain, PDE-Aligner captures these fundamental pat- terns more effectively than spatial approaches. It compares the dynamic system's semantics with statistical characters of the frequency domain, where the dynamical system de- scriptions, namely equations, boundaries, initial conditions, and other physical priors, are encoded into representation E
To characterize physical evolution, we acquire the ini- tial state u(I, ta) and the current state u(x, tz) of the physical field, applying the Fourier Encoder to obtain their k-tth frequency domain representations 0g (k. t.) and üx (k, to). The phase difference AΦ = (áx (k: t.) iix (k, to))/((aa (k: ( (k, (n))) captures wave propa- gation and dispersion characteristics, while the magnitude ratio R (k; tt)l/(xx (k, to)l quantifies energy transfer across scales-bote serving as physics-awand fingerprints of the underlying PDE. Thus, we have the alignment loss function as:
Baselines. The baselines are divided into two categories: (I) Task-specific expert models, which include Physics- Informed Neural Networks (PINNs) (Raissi et al., 2019), U-Net (Ronneberger et al., 2015), and Fourier Neural Operator (FNO) (Li et al., 2020), all of which require training from scratch for each specific case (cach equa- tion/coefficient, etc.). (2) Unified pre-training models, which include PDEformer-1 (Ye et al., 2024), Multiple Physics Pre-training (MPP) (McCabe et al., 2023), SWIN- transformer (Liu et al., 2021) used for the ORCA task, the large size pretrained checkpoint of Poseidon (Herde et al., 2024) and DPOT (Hao et al., 2024). More details on the baselines are provided in Appendix D.
LAtign = -Ley + ALE,
Log = =S( Erexe (P), 4 ( T)
(9)
LE = Σ K R- 1-1
where A a hyperparameter balancing the energy conserva- tion loss. The energy term LE enforces Parseval's theorem, ensuring physical consistency in the frequency domain. By minimizing the alignment loss function LAlign. the PDE- Aligner aligns the changes in the physical field with the textual descriptions within the constraints of energy conser- vation.
Training Details. The OmniArch model uses single-layer encoders and decoders for data of various dimensions, with the LLaMA model (trained from scratch) as the shared Transformer architecture. The PDE-Aligner employs the pre-trained Fourier encoder from OmniArch to encode phys- ical fields and the pre-trained BERT model to encode PDE captions. Additional training details are in Appendix E.
In the fine-tuning stage (Figure: 3 Right), the pre-trained
6

---

OmniArch: Building Foundation Model For Scientific Computing
Table 1: The nRMSE on various PDEs. We evaluate base-size(-B) and large-size(-L). The previous state-of-the-art performance underlined and our best performance is bolded.
ID Bur.
2D Reac.
3D Maxw.
Methods
CFD Adv.
Diff.
Reac.
CFD
SWE
Incom
CFD
Baselines -Task specific Expert Models
PINNs
0.8130 0.9450 0.2200 0.2140
1.6000 0.0170
U-Net
2.6700 0.7760 0.3201 0.1507 0.0026
1.0700
0.8401
0.0830
1.1200
0.7989
0.2999 0.3052 0.1906
FNO
1.4100 0.0091 0.0174 0.0017 0.0005
0.2060
0.1203 0.0044 0.2574
Baselines- Unified Pre-training and Fine-tuning
PDEformer-1
0.0043
0.0095
0.0009
ORCA-SWIN-B
0.8201 0.0062
MPP-AVIT-B
0.0227
0.0106
0.0024
MPP-AVIT-L
0.0178
0.0098 0.0022
Poseidon-L
0.1079
0.0949 0.0243
DPOT-L
0.0112
0.0263 0.0451
0.4321
Full Pre- Tre-Trannn ID.DD.3D Data
AmniArch-B(Ours)
0.0340 0.0238 0.0089 0.0020
0.0006
0.0196
0.0158 0.0016 0. 0.1726
0.5209
0.2834
OmniArch-L(OOrrs)
0.0250 0.0182 0.0063
0.0015
0.0004
0.0148
0.0105 0.0014 0.1494
0.4531
0.2268
+ PDE-Aligner Fine-tuning
OmniArch-B(Oury)
0.0302 0.0201
0.0071 0.0017
0.0003
0.0153
0.0102 0.0015 0.0955
0.4032
0.1813
OOniArch-I.OOurs)
0.0200 0.0041 0.0032
0.0006
0.0002
0.0125
0.0084
0.0012 0.0827
0.3723
0.1671
std.
0.0031 0.0012 0.0004 0.0001 0.0001
0.0017
0.0004 0.0003 0.0023
0.0443
0.0197
Improvement
98.70%
4.65%
66.32% 64.75% 60.00%
14.28% 45.45% 67.87%
12.32%
Notes:
Symbol Y means model did not converge while '-" means model applicable this dataset.
4.2. Results and Analysis
niArch Performance. OmniArch, pre-trained on ID, 2D, and 3D data, demonstrates superior performance across all evaluated datasets. Both the base (B) and large (L) ver- sions of ArniArch outperform existing models, validating its robustness in multi-physics contexts. To validate our architectural design choices, we conduct ablation studies on the Temporal Mask mechanism (Table 2). The results confirm that our Temporal Mask, which enables full atten- tion among physical quantities within each timestep, sig- nificantly outperforms standard causal masking across ious multi-physies systems. (4) PDE-Aligner Fine-tuning. Fine-tunee tuning with PDE-Aligner significantly enhances Om- ni Arch's 's accuracy, particularly for complex datasets. This step utilizes a pre-trained Fourier encoder and BERT-base- cased model, ensuring precise alignment between physical fields and PDE descriptions. Table 3 quantifies the impact of PDE-Aligner across different dimensions, showing con- sistent improvements of over 20% compared pre-training alone. OmniArch demonstrates substantial performance gains over baselines, with up to 98.70% improvement on CFD-ID and notable enhancements across other PDEs.
OmniArch is designed to support multi-scale, physii,hysics, and flexible physies alignment. Table presents the normal- ized root mean square error (nRMSE) across various PDEs for different methods.
Multi-Physics Results. (1) Compared with Task-specific Expert Models. PINNs, U-Net, and FNO require training from scratch for each specific equation or coefficient. While FNO shows strong performance, PINNs and U-Net strug- gle with convergence and accuracy in some cases (Like the CFD-1D, and CFD-2D). (2) Compared with Unified Pre-training Models. PDEformer-1 exhibits proficiency in specific ID equations but fails to generalize beyond its for- mulation structure. MPP and ORCA-SWIN leverage 2D pre-training and fine-tuning, improving generalization, yet their effectiveness remains constrained by the diversity of the pre-training data. Poseidon enables single-step infer- ence arbitrary timesteps, though its accuracy still leaves room for improvement. DPOT successfully transfers knowl- edge from 2D to 3D CFD through weight sharing, but it lacks support for ID CFD and its performance on non-CFD physics systems requires further cahancement.
Ablation Study on Masking Strategies. As illustrated in Table the superiority of Temporal Mask (18-20% improve-
(3) Om-

---

OmniArch: Building Foundation Model For Scientific Computing
Table 2: Ablation study on masking strategies
Methods
Shock
KH
OTVortex
Dataset
Causal Mask
No Mask
Temporal Mask
FNO
0.7484
1.0891
0.5946
2D Incom.
0.0277
0.0285
0.0227
U-Net
1.6667
0.1677
0.4217
2D CFD
0.0198
0.0205
0.0148
MPP-L
0.3243
1.3261
0.3025
3D CFD
0.1842
0.1923
0.1494
OmniArch-L 0.2126
0.2763
0.1718
ment) reveals fundamental insight: multi multi-physics systems require simultaneous rather than sequential processing of coupled variables. This advantage is most pronounced in 3D CFD, where the complex interplay between five physical quantities (velocities, density, pressure) demands holistic attention patterns.
Table 4: The Performance on Zero-shot PDEs.
nRMSE -Incom
10
Table 3: Impact of PDE-Aligner model performance (OmniArch--.)
DA-B-Single
OA-B-Rodee
Grid (12-513))))1))
OA-L-Single
Grid Siac](32-122)
Configuration
ID PDEs 2D PDEs 3D PDEs
32
64
128 Grid Size
256
512
Pre-training only
0.0103
0.0440
0.3399
Fine-tuning w/o Aligner
0.0073
0.0345
0.3432
Figure The multi-scale capability.
Fine-tuning w/ Aligner
0.0056
0.0262
0.2697
Improvement
23.3%
24.1%
21.4%
T+1
T-2
T-3
T-4
T+5
T+6
Ground fruth
Impact of PDE-Aligner. We report the impact of PDE- Aligner in Table 3, where the consistent 22% improvement across dimensions suggests that PDE-Aligner serves as more than a physics constraint-it helps OmniArch dis- ambiguate between different physical regimes learned dur- ing pre- pre-training. Notably, the similar improvement ratios across ID-3D indicate that physical alignment dimension- agnostic, validating our unified architecture design.
OmmärchL MPP-
Multi-scale Results. In Figure 4, we present the multi-scale inference performance of OmniArch Base and OmniArch- Large on the 2D Incom. Dataset. Due to the frequency truncation capability the Fourier Encoder, OmniArch can handle inputs of varying grid sizes without requiring re-training. In the red-shaded area, the nRMSE decreases the grid size becomes smaller. Conversely, in the blue- shaded area, the nRMSE slightly increases. However, even with a grid size of 512, the maximum nRMSE remains below 0.2. In the rollout settings, grid size of 256 some- times leads to better or comparable performance to a grid size of 128. The non-monotonic relationship between grid resolution and error (red vs. blue regions Figure re- veals an intriguing property frequency-domain learning: OmniArch naturally identifies the intrinsic resolution of physical phenomena. The optimal performance at interme- diate resolutions (128-256) suggests the model has learned to distinguish between meaningful physical scales and nu- merical artifacts. Additional visualizations are provided in Appendix H.5.
Figure 5: Zero-shot prediction results (Rollout) of OmniArch-L and MPP-L on KH dataset. Displaying time steps T+1 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.
0.150
nAMSE
Tine (6)
0.327
0.05
0.100
0.067
0.075
0.02
0.025
0.035
0.310
0.050
0.014
0.025
0.00
0.000
Length of Prompt
Figure 6: In-context learning SWE with OmniArch-B.
8

---

OmniArch: Building Foundation Model For Scientific Computing
Table 5: nRMSE for Inverse Problems.
confusion maatrix
7.37%
1.0%
1.0%
a.0%%
odess
ades
cales
Methods
Forcing
Buoyancy
AJON
duees
ocets
cains
MPP 0.2 = 0.008
0.78 0.006
4.4%
1.6
avies
ades
OmniArch
0.16 + 0.005
0.73 0.012
Jurger
4.5%%
100%
odes
Scratch
0.39 0.012
0.83 = 0.027
Referencee
6.00%
1.0%%
adds
96
Flexible Physies Alignment. To verify the Plllneg's''s ability to perceive physical information, we equipped it with a classification head to classify physical fields. In Figure 7, the PDE-Aligner can perceive physical field cate- gories based equation text information and physical field features, and the classification accuracy rate exceeds 0.94 on all ten categories. More details are in Appendix F.3.
adeu anew avi
nam
19. a.oms sokee adru adese adew ade
dlees aaess adew aden axen
adds
suew
30
10 37. 10 32 FaedD bageniD OF 20.33 20 feac 20 Dweca froor 32 Cro Moos
Predict
Zero-shot Performance. Our examination of 2D PDE pre- dictions, as illustrated in Figure 5, reveals that OmniArch effectively captures both low- and fieguerryeencyy patterns even in zero scenarios, surpassing former 2D models like MPP. MPP often misses key features, leading to er- roneous representations of the primary physics and failed rollouts. Details of zero-shot dataset are in Appendix C.2.
Figure 7: The confusion matrix of the PDE-Aligner classifi- cation results.
OmniArch outperforms MPP in parameter estimation tasks, with lower RMSE values indicating more accurate predic- tions. Models trained from scratch yield the highest errors, underscoring the effectiveness of our fine-tuning approach. This evidence supports the notion that OmniArch is not only proficient in forward simulations but also exhibits superior performance in deducing hidden dynamics within complex systems. More details are in Appendix H.3.
As shown in Table 4, nRMSE scores indicate that all mod- els, except OmniArch, tend to underperform in zero-shot transfer. This suggests that OmniArch's use of Fourier En- coders and unified training approach enhances its ability to generalize across different PDFs. By leveraging flexi- ble grid inputs and dynamic observation windows during pre-training, OmniArch effectively captures the underlying physics of the observed field states, which may not be ade- quately addressed by methods adhering strictly to explicit grid and temporal dependencies. The 4-7x error reduction compared to MPP in zero-shot scenarios (Table 4) indicates that OmniArch has learned transferable physical operators rather than sataset-specific patterns. The success on shock- dominated flows (Shock, notoiolyyy difficult for neural setthodd-ddmoostrattte that requeney-domait rep- resentations capture discontinuities more effectively than spatial approaches.
Other Results. In addition to the primary experiments, we include more rollout case studies in Appendix H.4 and report the inference-tine GPU Memory usage compared with baselines in Appendix H.7. We also include ablation studies for training settings in Appendix G and the detailed performance for CFD PDEs in Appendix H.6. These ad- ditional evaluations highlight OmniArch's robustness and accuracy in complex physical simulations, surpassing other state-of-the-art models.
5. Conclusion
In this study, we introduced a pioneering foundation model for scientific computing, specifically tailored for the resolu- tion of partial differential equations (PDEs). By integrating this model with a novel PDE-Aligner for fine-tuning, we have established new state-of-the-an benchmarks across a comprehensive suite of tasks within the PDEBench. Addi- tionally, we investigated the zero-shot learning capabilitiess of our pre-trained model, uncovering degree of transfer- ability that mirrors the emergent properties found in large- scale language models. Despite the successes, we recognize the challenges posed by 3D PDE systems to our Omni/Arhh model, which may leave for future research. We envisage that OmniArch will serve as cornerstone for developing foundation models in the domain of PDE learning, foster- ing a significant convergence between scientific machine learning (SciML) and broader deep learning disciplines.
In- Learning, After autoregressively pre-trained various dynamic systems, we observe that OmniArch could learn neural operators within the observations of several time steps, which is similar to the in-context learning in Large Language Models. Here, we define the given time- series of observations as PDE Prompt. Our approach varies the prompt length from 2 tokens (derived from a 50 time step interval) to 100 tokens (from 1 time step interval). More details are in Appendix H.2.
Fine-tuning for Inverse Problems. Demonstrating a model's 'ss capability to infer hidden physical parameters from known equations is a critical test of its ability to learn un- derlying physics. The results in Table demonstrate that

---

OmniArch: Building Foundation Model For Scientific Computing
Acknowledgement
networks for solving parameterized partial differential
equations. J. Comput. Phys., 455:1 10996, 2022. doi: 10. 10160... URL https: htt/:/oo.oooooo/ 10. 10.1016/). 2022. 110996.
This work was supported by the National Science and Technology Major Sotect(N..222201112330 and Young Elite Scientists Sponsorship Program by 2SSTNNo.202QQRRC0001
This work was also sponsored by CAAl-Huawei MindSpore Open Fund (CAAIXS)LU22233MddSSqSu
Cho, W., Lee, K., Rim, D., and Park, N. Hypernetwork- based meta-learning for low-rank physics-informed neu- ral networks. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, (eds.), Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023. NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
and developed on openl community. Thanks for the computing infrastructure pro- vided by Beijing Advanced Innovation Center for Big Data and Brain Computing.
Impact Statement
OmniArch represents a significant advancement in scien- tific computing. unifies multi-scale and multi-physies PDE-solving capabilities within a single foundation model framework. This unified approach has profound implica- tions for accelerating scientific discovery and engineering applications across domains such as fluid dynamics, weather forecasting, and materials science. The model's demon- strated ability to handle diverse physical systems and grid resolutions while maintaining physical consistency could dramatically reduce the computational resources required for complex simulations in industrial and research settings. Additionally, there many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Cuomo, S., Di Cola, V. S., Giampaolo, F., Rozza, G., Raissi, M., and Piccialli, F. Scientific machine learning through physics-informed informed neural networks: Where we are and what's next. Journal of Scientific Computing, 92(3):88, 2022.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding, arXiv preprint arXiv: 181.04485. 2018.
Du, G., Cao, X., Liang, J., Chen, X., and Zhan, Y. Medical image segmentation based on u-net: A review. Journal of Imaging Science and Technology, 2020.
Gupta, J. K. and Brandstetter, J. Towards multi- spatiotemporal--sat generalized PDE modeling. CoRR, abs/2209. 15616, 2022. doi: 10.48550/ARXIV.2200. 15616. URL https: /tt::///ooiii http//iiioogg/////////// orr....0g/10.885555 arXiv. 2209.15616 15616.
References
Alkin, B., Fürst, A., Schmid, S., Gruber, L., Holzleitner, M., and Brandstetter, J. Universal physics transformers. CoRR, abs/2402. 12365, 2024. doi: 10.48550/ARXIV. 2402.12365. URL https: http:::/ddo..oggg/// arXiv. 2402 12365.
48550/
Hao, Z., Su, C., Liu, S., Berner, J., Ying, C., Su, H., Anandkumar, A., Song, J., and Zhu, J. DPOT: auto- regressive denoising operator transformer for large-scale PDE pre-training. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: / //openreview net forum? d=w.ddDeeHUHOO.X.
Allen, K. R., Sanchez-Gonzalez A., Battaglia, P. W., Hamrick, J. B., and Plaff, T. Physical design using differentiable learned simulators. CoRR, abs/2202.00722 2022. URL https: // //arxiv. org/ abs///s//// /t////bb22222..2 2202 00728.
T., Stachenfeld, K. L.,
Herde, M., Raonic, B., Rohner, T., Käppeli, R., Molinaro, R., de Bézenac, E., and Mishra, S. Poseidon: Efficient foundation models for pdes. CoRR, abs/2405. 19101, 2024. doi: 10.48550/ARXIV.2405..0 19101. URL https: / //doi. orr. 10. (ww.85550//r/// 2405 .99011.
Brigham, E. O. The fast Fourier transform and its applications. Prentice-Hall, Inc., 1988.
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 in neural information processing systems, 33: 1877-1901, 2020.
Huang, X., Ye, Z., Liu, H., Shi, B., Wang, Z., Yang, K., Li, Y., Wang, M., Chu, H. Yu. F., Hua, B., Chen, L., and Dong, B. Meta-auto-decoder for solving parametric partial differential equations. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A sim- ple framework for contrastive learning of visual represen- tations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Chen, Y., Dong, B., and Xu, Meta-mgnet: Meta multigrid
10

---

OmniArch: Building Foundation Model For Scientific Computing
Karniadakis, G. E., Kevrekidis, I. G., Lu, L., Perdikaris, P., Wang, S., and Yang, L.. inysice-informed machine learning. Nature Reviews Physics, 3(6):422-440, 2021.
Mialon, G., Garrido, Q.. Lawrence, H., Rehman, D., Le- Cun, Y., and Kiani, B. T. Self-supervised learning with lie symmetries for partial differential equations. CoRR, 0ass/2077.55433 2023. doi: 10.48550/ARXIV. 2307.05432. URL https::/p///s.... /ooi..i.../// wtt.ddiorgggooooooooooooooooooo. orgg 48850/ arXiv. 2307. 230..05437... 05432.
Langley, P. Crafting papers on machine learning- In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207-1216, Stanford, CA, 2000. Morgan Kaufmann.
Nguyen, E., Poli, M., Faizi, M., Thomas, A. W., Wornow, M., Birch-Sykes, C., Massaroli, S., Patel, A., Rabideau, C. M., Bengio. Y., Ermon, S., Ré, C., and Baccus, S. Hye- nadna: Long-range genomic sequence modeling at single nucleotide resolution. In Oh, A., Naumann, T., Glober- son, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Li, Y. A practical survey zero-shot prompt design for in-context learning. In Mitkov. R. and Angelova, G. (eds.), Proceedings of the 14th International Conference on Recent Advances in Natural Language Processing, RANLP 2023, Varna, Bulgaria, 4-6 September 2023, pp. 641-647. INCOMA Ltd., Shoumen, Bulgaria, 2023. URL https::///// httt::://acllnthooooyy... org/2023. ranl 1anlp-1.... 19.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.000 0,0885,, 2020.
Oden, J. T. An introduction to the finite element method with applications to nonlinear problems (R. e. white). SIAM Rev., 31(3):112, 1989. doi: 10.1 10.1137/1031144 URL https: /tt.//ooiioo/// wtt.://.///...../1...... 113.
Li, Z., Shu, D., and Farimani, A. B. Scalable transformer for PDE surrogate modeling, CoRR, abs/2305. 2///2335..755 17560,, 2023. doi: 10.48550ARRXI.2205.155555
Ohana, R., McCabe, M., Meyer, L., Morel, R., Agocs, F. J.. Beneitez, M., Berger, M., Burkhart, B., Dalziel, S. B., Fielding, B., Fortunato, D., Goldberg, A., Hirashima, K. Jiang, Y., Kerswell, R. R., Maddu, S.. Miller, J., Mukhopadhyay, P., Nixon, S. S., Shen, J., Wattcaux, R., Blancard, B. R., Rozet, F., Parker, L. H., Cranmer, M. D., and Ho, S. The well: large-scale collection of diverse physics simulations for machine learning. In Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomozak, M., and Zhang, C. (eds.). Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
URL ht http //doi. www.orog////. 4855088550// (tt555///xxivv/ ......15656 17560.
Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B. Swin transformer: Hierarchical vision transformer using shifted windows. 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021, pp. 10002. IEEE, 2021. doi: 10. 10.1499/1CC448922.. 2021.00986. URL https::////. TCCV48922.2021.002 00986.
Lorsung, C., Li, Z., and Farimani, A. B. Physics informed token transformer. CoRR, abs/2305.08757, 2023. doi: 1044550/ARXXV222500055555 URL https: //doi. org/10. 000.448855052113 23w.xxiii.33.... 2305. 08757.
Pathak, J., Subramanian, S., Harrington, P., Raja, S., Chat- topadhyay, A., Mardani, M., Kurth, T., Hall, D., Li, Z., Azizzadenesheli, K., Hassanzadeh, P., Kashinath, K., and Anandkumar, A. Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators. CoRR, 1214 2022. URL Ttttgs///errii.orr/a///22222..
Lu, L., Jin, P., Pang, G., Zhang, Z., and Karniadakis, G. E. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218-22) 2021a.
Lu, L., Meng, X., Mao, Z., and Karniadakis, E. Deepxde: A deep learning library for solving differential equations. SIAM review, 63(1):208 2228 2021b.
Pfaff, T., Fortunato, M., Sanchez-Gonzalez, A., and Battaglia, P. W. Learning mesh-based simulation with graph networks. arXiv preprint arXiv:2010.03409, 2020.
McCabe, M., Blancard, B. R., Parker, L. H. Ohana, R., Cranmer, M. D., Bietti, A., Eickenberg, M., Golkar, S., Krawezik, G., Lanusse, F., Pettee, Tesilenu,, T., Cho, K., and Ho, S. Multiple physics pretraining for physical surrogate models. CoRR, 10.22300,02944, 2023. doi: 10.4550/ARX1V2310002994 URL https: //doi. org/10. 48550/ /tt.ss///xxvvvv. 2310. 02994.
Radford, A., Narasimhan, K., Salimans, T., Sutskever, I., et al. Improving language understanding by generative pre-training. OpenAl blog, 2018.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.

---

OmniArch: Building Foundation Model For Scientific Computing
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, Learning transferable visual models from natural language supervision. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 8748- 8763. PMLR, 2021. URL httt:://proceete mlr A.l.ppees/vv13 http://////aadddd1ll.a...a.m..aa html.
Siddique, N., Paheding, S., Elkin, C. P., and Devabhaktuni,
V. U-net and its variants for medical image segmentation: A review of theory and applications. leee Access, 9: 2021.
Subramanian, S., Harrington, P., Keutzer, K., Bhimji,
W., Morozov, D., Mahoney, M. W., and Gholami, A. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. CoRR, abs/2306.00258. 2023. doi: 10,48550/ARXIV, 2306.00258. URL https: /ww./ii.ooooo /t/d/iiorg/....e. 488550/ /855 arXiv. 2306 00258.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics- informed neural networks: A deep learning frame- work for solving forward and inverse problems involy- ing nonlinear partial differential equations. Comput. Phys., 378:686-707, 2019. doi: 10. 10.1016/0.JCP2208..II 045. URL https: / /wi. http://.......11111111110 jep. 2018. 10.010..55 045.
Sun, J., Liu, Y., Zhang, Z., and Schaeffer, H. To- wards a foundation model for partial differential equations: Multi-operator learning and extrapolation. CoRR,
2024, doi: 10.48550/ARXIV, 2404. 12355. URL https: www.i..oggo...... 48550/ a arXiv 2404. 12355.
Raonic, B., Molinaro, R., Ryck, T. D., Rohner, T., Bar- tolucei, F. Alaifari, R., Mishra, S., and de Bézenac, E. Convolutional neural operators for robust and accurate learning of pdes. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M. and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023.
Sun, L., Gao, H., Pan, S., and Wang, J.-X. Surrogate mod- cling for fluid flows based on physics-constrainee deep learning without simulation data. Computer Methods in Applied Mechanics and Engineering, 361:1 361:112732, 2020.
Takamoto, M., Praditia, T. Leiteritz, R., MacKin- lay, D., Alesiani, F. Pflüger, D., and Niepert, M. Pdebench: extensive benchmark for scientific ma- chine learning. In Koyejo, S., Mohamed, S., Agar- wal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Reusch, A.. Thicle, M., and Lehner, W. Transformer- encoder and decoder models for questions on math. In Faggioli, G., Ferro, N., Hanbury, A., and Potthast, M. (eds.), Proceedings of the Working Notes of CLEF 2022 Conference and Labs of the Evaluation Forum, Bologna, Italy, September 5th to 8th, 2022, vol- ume 3180 of CEUR Workshop Proceedings, pp. 119- 137. CEUR-WS.org. 2022. https:/ pttttts://ceurrwtt org/Vol- 3180/1100 /etpp/ape---eee 07. 07.pdf.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lample, G. Llama: Open and efficient foundation language models. CoRR, abs/2302. 13971, 2023. doi: 10-48550/ARXIV22302.19999
Ronneberger, O., Fischer, P., and Brox, T. U- net: Convolutional networks for biomedical image segmentation. Medical Image Computing and Computer-Assised Intervention-MICCA. 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings. Part III 18, pp. 234-241. Springer, 2015.
URL https: //doi. org/10. 48540 00850//50Xiiii http://axvvvvvv 2302 13971.
Veličković, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., and Bengio, Y. Graph attention networks. arXiv preprint arXiv: 1ttp:::111.100000 10903, 2017.
Yang, L.,, Liu, S., Meng, T., and Osher, S. J. In-context operator learning with data prompts for differential equa- tion problems. Proceedings of the National Academy of Sciences, 120(39):e23100 1212220, 2023a.
Shen, J., Marwah, T., and Talwalkar, A. Ups: Efficiently building foundation models for pde solving via cross- modal adaptation. Transactions on Machine Learning Research.
Yang, L., Meng, T., Liu, S., and Osher, S. J. Prompting in-context operator learning with sensor data, equations, and natural language. CoRR, abs/2308 000001, 2023b. doi: 10.48550/ARXIV22008 2.000.00000000066.... URL attps:/ / doi doi.org/10 48500.48555000 rxiv. 2308, 05061.
Shen, J., Li, L., Dery, L. M., Staten, C., Khodak, M., Neubig, G., and Talwalkar, A. Cross-modal fine-tuning tuning: Align then refine. In International Conference on Machine Learning, pp. 11000-10066 PMLR, 2023.
12

---

OmniArch: Building Foundation Model For Scientific Computing
Ye, Z. Huang, X., Chen, L., Liu, H., Wang. Z., and Dong, B. Pdeformer: Towards foundation model for one-dimensional partial differential equations. CoRR, 0e222402.12652, 2024. doi: 10.48550/ARXIV. 2402. 12652. URL. htt.t:t/t/s//ete .tt.//ooooooooooooooooo/oo./ 10 488550 arXiv. 2402 12652.
Zhang, D., Bi, H., Dai, F., Jiang. W. Zhang, L., and Wang, H. DPA-1: pretraining of attention- based deep potential model for molecular simulation. CoRR, abs/2208. 08236, 2022. doi: 10.45500AARXI122200002225 URL https: org.i.oog//a http:/...4445.50/
2208. 08236.
Zhang, S., Dong, L., Li, X., Zhang, S., Sun, X., Wang, S., Li, J., Hu, R., Zhang, T., Wu, F, and Wang, G. In- struction tuning for large language models: A survey. CoRR, 10792, 2023. doi: 10.48550/ARXIV. 2308.10792. URL https::////... htRp::/t::///ao http:::///oo..g//....../ 48550/ arXiv. 1079.
13

---

OmniArch: Building Foundation Model For Scientific Computing
Supplementary Material for:
OMNIARCH: Building Foundation Model For Scientific Computing
Contents
A. Table of notations
15
B. Limitations
15
C. Dataset details
16
C.1. OmniArch Pre-training Dataset
16
C.2. Dataset For Zero-shot Learning
16
D. Baseline implementation details
17
E. OmniArch implementation details
18
E.1. Pre-training OmniArch
18
E.2. Fine-tuning OmniArch
19
E.3. Parameter Efficiency Analysis
19
F. PDE-Aligner implementation details
19
F.1. PDE-Aligner Pre-training Dataset
19
F22 Examples of Generated PDEs
20
F.3. Pre-training process of PDE-Aligner
21
G. Further ablation study
23
G.I. Dynamic Prompt Length for Efficient Inference
24
G.2. Fine-tumed for Inverse Problems
24
H. More results
23
H.1. Zero-shot Learning Capability
23
H.2. Rollout Predictions
24
H.3. Multi-scale Inference Results
24
H.4. More results in different problem settings
.44
H.5. GPU Memory Usage and Inference Time
27
H.6. Comparison with Traditional Solvers
27
I. More Discussions
28
LI. Meta-Learning vs. Scaling Laws PDE Solving..
28
14

---

OmniArch: Building Foundation Model For Scientific Computing
Table 6: Table of notations
Basic Notations
2,6
Spatial and Temporal coordinate (time)
At
Time interval
T
Total time steps
D
Total dimensions of physical fields
C
Number of physical fields
B
Batch size
U
Physical field inputs
"( u(x(d)00
Physical field of dimension d at spatial coordinate x and time
(-)
Linear projection
OmniArch Related Notations
F,
F-F-1
Fourier transformation and inverse
F
Components (modes) in the frequency domain
ú
Physical field frequency domain representation
k
Frequency Variable
K
Number of retained Fourier modes (cut-off' frequency) Fourier transform of u(x, t) at frequency and time t
K (ke,
Truncated Fourier modes (TopK modes) at frequency & and time &
pred (k,t)
Predicted Fourier modes at frequency and time Predicted physical field spatial coordinate x and time t
med (x.t)
f()
Integral operator
au
Attention weights at spatial coordinate and time
R()
Real-valued embedding function of the frequency domain features
U, Va
Physical field embedding token from R at time
Z.
Input sequence consist of grouped embeddings from Ug, Ve
Z.
Shifted right predicted feature sequence
M
Temporal mask used in Transformer Blocks
ou
Normalization factor ||w||2 for nRMSE calculation
L n
Normalized RMSE loss: V [[(rred /
sim
L Laim
Mean nRMSE across all physical fields
PDE-Aligner Related Notations
PDE text description (captions)
Evx (-)
Text encoder used for PDE captions
Phase difference between the initial state and current state Amplitude ratio between two states
R
X
The hyperparameter balancing the the energy conservation loss Similarity between text embedding and physical embedding
Loa
LE
Energy conservation loss
I Latign
PDF-Aligner training loss
Li
OmniArch fine-tune loss
A. Table of notations
A table of notations is given in Table 6.
B. Limitations
Despite its advancements, OmniArch remains fundamentally data-driven, and its interpretability requires further improve- ment, even with the PDE-Aligner enhancing physical prior alignment. Constraints computational power and data
15

---

OmniArch: Building Foundation Model For Scientific Computing
availability have limited OmniArch's scalability, affecting its generalization capabilities, particularly in complex and abrupt dynamical systems such 3D tasks and shock wave PDEs. Addressing these limitations is crucial for further development and broader applicability in scientific and engineering contexts.
C. Dataset details
C.1. OmniArch Pre-training Dataset
Pre-training Stage. We structured the PDEBench data into distinct training, validation, and testing subsets. For one- dimensional (1D) PDEs, the training dataset comprises selection from the CFD-ID ID, ReacDiff, Advection, Burgers, and diff-sorp datasets. From these, we reserve random 10% sample of trajectories the in-domain test set for each respective PDE equation. The Shock Tube Equation is designated as the out-of-domain test set. Additionally, the test portions of the reacdiff and diff-sorp datasets are utilized as part of the test set.
In the two-dimensional (2D) PDE case, we allocate 90% of trajectories from the CFD, diff-react, NSincom, and shallow water datasets for training. The remaining 10% form the in-domain test set. The Shock Tube, Gelvin-Helmholie instability (KH), and Oopman-Oppennhemer-- Volkoff (TOV) scenarios are included as out-of-domain test sets.
For three-dimensional (3D) PDEs, 90% of trajectories from the CFD-3D dataset are utilized for training, with the remaining 10% serving the in-domain test set. The complete datasets for blastwave and turbulence simulations used as out-of-domain test sets. The Details of our pre-training dataset can be found in Table 7.
Table
Data Statistics for OmniArch Pre-training
Dataset
#Train
#Validation
#Physical quantities
N
Ng
CFD
45000
5000
velocities Vg, density, pressure
100
1024
Reac.
144000
16000
concentration p velocities Va velocities Vi concentration p
200
1024
1D
Adv. 72000
8000
200
1024
Bur. 108000
12000
200
1024
Diff.
9000
1000
100
1024
CFD
39600
4400
velocities Va, Vy, density, pressure
21
512
Reac.
900
100
activator u, inhibitor v velocities Va, Vy, particle velocities h
100
128
2D
Incom
900
100
1000
256
SWE
900
100
100
128
CFD
630
70
velocities VA, Vy, Vz, density, pressure
21
128
3D
electric displacement Dr, Dy, Dz magnetic field Ha, Hy, H₂
Maxw.
8640
960
8
64
C.2. Dataset For Zero-shot Learning
We choose three test datasets from PDEBench to validate the zero-shot ability of our model. They all belong to two- dimensional compressible Navier-Stokee equations but are different fluid phenomena that exhibit distinct physical mecha- nisms and characteristics. Brief introductions and details of the datasets are as follows:
OTVortex:
The Orszag- Tang Vortex system is compressible flow problem that generates highly complex vortex structures through the careful selection of initial conditions. The dataset includes one example, which is 1024 x 1024 resolution physical field evolved over 101 time steps with time interval of 0.01.
2D Shock: Shock waves are characterized by abrupt changes in flow properties resulting from sudden discontinuities
16

---

OmniArch: Building Foundation Model For Scientific Computing
in fluid flow, such as rapid changes in pressure, temperature, and density. The dataset includes one example, which is also 1024 1024 resolution physical field evolved over 101 time steps with time interval of 0.01.
2D KH: The Kelvin-Helmholtz instability is fluid instability that occurs at the interface between two fluid layers with different velocities or densities. This dataset consists of seven examples generated based on different parameters M, dk, and Re. Each is 1024 x 1024 resolution physical field evolved over 51 time steps with a time interval of 0. We conducted experiments on all samples and averaged the results.
D. Baseline implementation details
In our experiments, adopt the benchmarking framework provided by PDEBench (Takamoto et al., 2022) and select three well-established methods for comparative analysis. Furthermore, we have incorporated the Multiple Physics Pre-training (MPP) model into our comparative analysis to address the need for retraining that is inherent to the aforementioned methods when faced with novel sets of conditions, the detailed training hyperparameters of FNO, U-Net, and PINN is provided in Table 8, following PDEbench (Takamoto et al., 2022). The first hyperparameter of U-Net the unroll steps (denoted as us), and the second the train steps (denoted as ts). The hyperparameters shared by both FNO and U-Net are the initial steps (denoted as is) and batch size (denoted bs). The hyperparameter in PINNs is the hidden size (denoted as hid). The learning rate, shared by FNO, U-Net, and PINNs, is denoted as Ir.
Table 8: Setting details when training FNO, U-Net, and PINN, means shared setting for FNO and U-Net, shared setting for FNO, U-Net, and PINN is denoted with symbol t.
FNO
U-Net
PINNs
is
bs
Irt
modes
width
us
ts
hid
Adv.
12
20
20
200
10
50
40
0.001
Bur.
12
20
20
200
10
50
40
0.001
ID
CFD
12
20
20
100
10
50
40
0.001
Diff.
12
20
20
101
10
50
40
0.001
Reac.
20
101
0.001
CFD
12
20
20
21
10
20
40
0.001
2D
Reac.
12
20
20
101
10
40
0.001
SWE
12
20
20
101
10
40
0.001
Incom
12
20
101
40
0.001
CFD
12
20
20
21
10
5
40
0.001
3D
Maxw.
12
20
7
8
7
40
0.001
Physics-Informe Neural Networks (PINNs) (Raissi et al., 2019). PINNs utilize neural networks to solve differential equations by embedding physical laws into a multi-objective optimization framework, minimizing PDE residuals and boundary/initial condition errors (Cuomo al., 2022).
U-Net (Ronneberger et al., 2015). U-Net, designed for biomedical image segmentation, uses encoder-decoder structure for context capture and precise localization (Siddique et al., 2021; Du et al., 2020). We adapt U-Net into ID and 3D forms to analyze spatio-temporal patterns in physical fields.
Fourier Neural Operator (FNO) (Li et al., 2020). FNO pioneers in learning function-to-soluuie mappings by parameter- izing integral kernels in the Fourier domain, enabling efficient and accurate iollution-invariann neural operators.
PDEformer-1 (Ye et al., 2024). PDEformer-1 is neural solver capable of simultaneously addressing various types of 1D partial differential equations. It uses a graph Transformer and implicit neural representation (INR) to generate mesh-free predicted solutions.
Multiple Physics Pre-training (MPP) (McCabe et al., 2023). MPP extends PDEBench's 2D physics scenarios to learn versatile features for predicting dynamics across various physical systems and comprises pre-training and fine-tuning phases, warranting its inclusion in our comparative analysis.
17

---

OmniArch: Building Foundation Model For Scientific Computing
ORCA-SWIN (Shen al., 2023; Liu al., 2021). ORCA fine-tunes the SWIN Transformer for different PDES by first aligning the embedded feature distribution of the target PDE data with the pre-training modality, and then refining the model on this aligned data to effectively leverage shared knowledge across various PDEs.
E. OmniArch implementation details
E.1. Pre-training OmniArch
In our training process, the following strategies or decisions were made:
Pre/Post Norm: Pre-norm
Norm Type: RMS Norm Type
Architecture: Decoder-Only
Attention-Type Multi-scaled Attention
Position Embedding: RoPE
Casual Masking: True- We only evaluate the loss on the T physical fileds prediction.
Hidden Size: 1024
initializer_rann range: 0.02
intermediate_sion 4096
num attention heads: 16
Table 9: Detailed setting of hyperparameters in pre-training the base and large models. The batch sizes, modes, and widths are provided as lists, with values corresponding to ID, 2D, and 3D data respectively.
Hyperparameters
Base
Large
#Layers
12
24
Hidden Size
768
1024
#Heads
12
16
Intermediate Size
3072
4096
Batch Sizes
[42,3,1]
[32,2,1]
Modes
[12,12,12]
[12,12,12]
Widths
[8,8,8]
[8,8,8]
Learning Rate
0.0001
0.0001
Scheduling Method Cosine Annealing Cosine Annealing
We trained two different sizes of model: base and large, which primarily differ in the number of layers, hidden sizes, number of heads, and intermediate sizes, as detailed in Figure For the base model, we selected batch sizes of [42, 3, 1] for the ID, 2D, and 3D trajectories, respectively. These batch sizes represent the maximum capacities our acceleration devices could handle while maintaining the ratio of data trajectories. This configuration allows for optimal training efficiency by minimizing idle time and maximizing device utilization. For the large model, due to its significantly increased size, we adjusted the batch sizes to [32, 2, 1] to ensure that the GPU memory is fully utilized. This reduction in batch sizes accommodates the larger model's 's memory requirements while still enabling effective training across the different dimensions of data trajectories.
18

---

OmniArch: Building Foundation Model For Scientific Computing
E.2. Fine-tuning OmniArch
Fine-tuning is performed on an A40 GPU cluster, which has 40GiB of memory per device. The fine-tuning settings for each dataset are shown in Table 10. We set the learning rate 1e-5, which results in fast convergence. Using 2 GPUs in Distributed Data-Parallel mode, we fine-tune each dataset for maximum of 30 epochs and apply early stopping.
Table 10: Detailed Fine-tuning Settings: The table provides the learning rate, width, modes, and batch size for ID, 2D, and 3D data.
Dims
learning rate width modes batch size Scheduling Method
ID
le-5
Cosine Annealing
2D
le-5
12
Cosine Annealing
3D
le-5
12
2
Cosine Annealing
E.3. Parameter Efficiency Analysis
Table 11: Static Parameter Distribution (Millions)
Model Component OmniArch-B (316M)
OmniArch-L. (672M)
Shared Backhone
138 (43.7%)
435 (64.7%)
ID Encoder/Deccder
0.3
0.4
2D Encoder/Decoder
7.0
9.0
3D Encoder/Deccder
171
227
Table 12: Active Parameters During Task Execution (Millions)
PDE Type
Model
ID PDEs
2D PDEs 3D PDEs
OmniArch-B OmniArch-L
138
144
308
435
445
663
As illustrated in Table and Table 12, the parameter distribution reveals OmniArch's hierarchical design philosophy. Three key observations emerge: (1) The shared backbone dominates the parameter count (43.7-64.7%) facilitating cross- dimensional knowledge transfer while requiring only modest modality additions (0.3-227M). (2) For 2D tasks (MPP's primary domain), OmniArch-B activates merely 144M parameters a 24. 1% increase over MPP-B's 116M that brings three key advantages: (a) unified architecture reduces system complexity, (b) enables latent cross-modal learning, and (c) provides future-proof extensibility, (3) The scaling pattern shows intelligent allocation-3D processing requires 2. 1-2.7 more dedicated parameters than 2D, reflecting its inherent higher dimensionality while maintaining efficient reuse of the shared backbone.
F. PDE-Aligner implementation details
F.1. PDE-Aligner Pre-training Dataset
PDE-Aligner equation augmentation. Given the significant imbalance between equation caption data and physical field data, single equation can yield multitude of physical field simulations. To augment equation captions effectively, it is crucial to preserve the equation's s solutions and boundaries while adhering to physical laws and exploring wide array of possible substitutions. To achieve this, we have developed five-step augmentation pipeline: Equation Rewriting, Form Transformation, Linear Combination, Symbol Substitution, and Physical Checking:
Equation Rewriting, We apply mathematical identities to modify the equation, ensuring the core propertics remain intact.
19

---

OmniArch: Building Foundation Model For Scientific Computing
Form Transformation. We transform equations between differential and integral forms and employ techniques such Green's functions to broaden the equation's representations.
Linear Combination. For systems of equations, we derive new variants through linear combinations, enriching the dataset without altering the system's 's nature.
Symbol Substitution. We systematically swap variables with alternative symbols, such as replacing x with &, to maintain consistency and avoid ambiguity.
Physical Checking, A panel of GPT-4-based experts evaluates the augmented equations, filtering out those that do not align with physical principles.
Leveraging the first four steps, we generate 200 augmented instances per equation type. Subsequently, during the Physical Checking phase, we select the top 50% of these examples based on quality for pre-training, Representative samples of the augmented examples are available Appendix F.2.
Additionally, we randomly sample the numerical distributions of different physical quantities at two distinct time steps within the physical field to represent the field's temporal variations. Each set of two-step physical field data is paired with corresponding enhanced equation text form single data instance. This approach is used to compile a comprehensive pre-training dataset for the PDE-Aligner.
F.2. Examples of Generated PDEs
F.2.1. BURGERS 1D
Original form:
0el(txz) (0,())))))(( = v v/T a 1), e (0,2),
u(0,z) = uo(#), I E (0,1),
After augmented:
2 u(1, (2
8 v² (1,xx)
0.77v f dh u(t x) dt
0.77
dl
Ox
0.73lv(O, x) = 0.73/vn
Explanation: We replace with and 0 with a We integrate and multiply factors both sides of the equation at the same time.
F.2.2. ADVECTION
Original form:
0,((t,x) + 30,u(t, = =0, x (0. 1),t (0,2),
= =(0,))=v((z)) x x (0,1),
After augmented:
1.45 /
Ô (c- A(t,
" A(t,2) dt -0
A(0,z) = Ao(a)
Explanation: We replace u with A, d, a with de a and with c. We integrate and multiply some factors on both sides of the equation the same time.
20

---

OmniArch: Building Foundation Model For Scientific Computing
F.2.3. CFD-1D
Original form:
p(d.v + vv.V) -vp+ n+v+ + v),
c+ po2
+ V. (c pod 2
d1
2 1
) v.-v.0.
=0,
After augmented:
e(t, (tt,))a w(t, ()t + e(t,x) =0
(10)
dx
0.61 x(t,x) θ w(t,
ot w. e(t,x)
+
0x
0.61v) w (1, x) + 0.61 (x+ ?) 8² 3 dr w(1, 0.61
02
On ((,,
(11)
dr2
Explanation: We replaced many symbols, such replacing V with 0 and A with J We integrate and multiply some factors both sides of the equation at the same time. We also swapped the order of some items, such as Ç + n/3.
Table 13: Detailed Data Information: The total amounts of training data, sampled training data, total validation data, and sampled validation data are presented as lists. These lists correspond to ID, 2D, and 3D data respectively.
Dims
Total training Sampled training Total validation Sampled Validation
ID
218T
378K
269M
42K
2D
3.13T
42K
38M
5K
3D
748K
0.63K
9K
0.07K
F.3. Pre-training process of PDE-Aligner
In our architecture, the PDE-Aligner is divided into two components: a text encoder and a physics encoder. The text encoder utilizes the pre-trained albert-math model (Reusch et al., 2022), which is highly capable of processing Lenoe-encoded PDE captions due to its extensive training a large corpus of LaTeX data. For the physics encoder, employ the pre-trained Fourier encoder from OmniArch, known for its strong ability capture physical field features. We adopt large-batch contrastive learning approach similar to SimCLR (Chen et al., 2020). The training involves stochastic sampling strategy with an equal probability (50%) of selecting either canonical PDE captions sourced directly from textbooks or augmented PDE captions. The latter is assumed to enhance the text encoder's generalization capabilities while retaining critical PDE information textual form. The weights of the text encoder and physics encoder are fixed during the PDE-Aligner training process. The training data details for PDE-Aligner are shown Table 13, and the hyperparameter settings are provided in Table 14.
During the fine-tuning phase, the PDE-Aligner evaluates the alignment of gold-standard PDE captions with the state of physical fields at each step of generator G's decoding process. The resulting rewards are averaged over the temporal dimension and finalized upon the completion of inference. The intuition behind the PDE-Aligner fine-tuning is to help OmniArch distinguish the patterns behind different PDE systems. To verify the PDE-Aligner's ability to perceive physical information, we equipped it with a classification head to classify physical fields. The results, shown in Figure 8, indicate that the PDE-Aligner effectively aligns with physical laws.
21

---

OmniArch: Building Foundation Model For Scientific Computing
Table 14: Detailed Settings:
The init learning rate, optimizer, scheduler, hidden size, trainable params, total params, steps, and GPU hrs are presented lists.
Value
Init Learning Rate
le-4
Optimizer
Adam
Scheduler
Cosine Annealing
Hidden Size
768
Trainable params
1.2M
Total Params
195M
Steps
37k
GPU hrs
75
confusion matrix
1.0
4906
14
13 0.26%
40 0.80%
2
3 0.06%
7
0 0.00%
1D
98.28%
0.00%
0.14%
23
7953
2 0.03%
4 0.05%
18 0.22%
0 0.00%
0 0.00%
1D Adv.
0.29%
99.41%
0.00%
0.00%
0.00%
0.00%
0.8
7 0.04%
15733
252
2 0.01%
0 0.00%
1D Reac.
0.04%
98.33%
1.57%
0.00%
0.00%
0.00%
0.00%
0.00%
142
253
11504
1D Burgers
0.24%
1.18%
2.11%
95.87%
0.60%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
1D Diff. 0.00%
3.12%
0.00%
0.00%
96.88%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.6
Reference
0 0.00%
4164
170
28 0.64%
2D
0.18%
0.02%
0.00%
0.00%
94 6.%%
0.66%
0.00%
3.86%
0.00%
0 0.00%
0 0.00%
0 0.00%
96
0 0.00%
0 0.00%
0.00%
0.00%
100.00%
0.000
0 0.00%
0 0.00%
96
0 0.00%
0 0.00%
0 0.00%
2D SWE
0.00%
0.00%
0.00%
0.00%
0.00%
100.00%
0 0.00%
96 100.00%
0 0.00%
2D Incom
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.2
3D CFD
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
100.00%
0.00%
3D Maxw.
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
0.00%
5.52%
94.48%
0.0
1D CFD
1D Adv. 1D ReacID BurgersID Diff. 2D CFD 2D Reac. 2D SWE 2D Incom 3D CFD 3D Maxw.
Predict
Figure 8: The confusion matrix of the PDE-Aligner classification results. PDE-Aligner can perceive physical field categories based on equation text information and physical field features, and the classification accuracy rate exceeds 0.94 on all categories.
22

---

(KH), and OmniArch: Building Foundation Model For Scientific Computing
20. Shock
20 KH OTVortex
Comparison between different training metrics
0.5
10a
0.45
0.3
0.25
3640
10260 13570
16880
20190
23500
26810 30120 33430 36740
10
STEPS
30
50
PDE Prompt Length
Figure 10: Zero-shot learning nRMSE for T+1 timesteps with varying context lengths.
Figure 9: The training loss curve using different metrics.
G. Further ablation study
We have conducted an ablation study batch-wise nRMSE. We nRMSE, RMSE, and MSE respectively loss functions in the training process. We found that nRMSE leads OmniArch's convergence. Table 15 shows nRMSE yielded lower training losses compared to MSE and RMSE (up to 9.3% improvement). While we did encounter gradient calculation issues in extreme cases, these were mitigated by adding a small e to the squared norm of the true labels (averaged over the spatial dimension). Utilizing nRMSE as a training loss function aims to simultaneously reduce all channels, irrespective of their relative numerical values. The loss curve is shown in Figure
unified loss scale for different PDEs, benefiting
Table 15: Training loss metrics ablation study.
Steps
MSE
RMSE
nRMSE
10K
0.3624
0.3458
0.3386 (-2.08%)
20K
0.3371
0.3289
0.3175 (-3.47%)
30K
0.3240
0.3225
0.3005 (-6.22%)
40K
0.3181
0.3183
0.2887 (-9.83%)
H. More results
H.1. Zero-shot Learning Capability
Our examination of 2D PDE predictions reveals that, in contrast task-tuned models, the OmniArch model adeptly captures both low- and high-frequency patterns in in-domain PDEs such as Reaction Diffusion, CFD, Shallow Water, and Incompressible NS. Task-tuned models often miss key features, occasionally leading to erroneous representations of the primary physics. For oft-of-domain PDEs, delineated by red-dotted box in the figure, we evaluated the models' ability to predict unseen PDEs without fine-tuning or parameter adjustment. While task-tuned models consistently failed at this zero-shot learning task, OmniArch successfully predicted essential low-frequency background patterns, though struggled with high- frequencyycy details. Details on the zero-shot dataset, including shock wave, Orszag- Tang Vortex (OTVortte) phenomena, are provided in Appendix C.2.
23

---

OmniArch: Building Foundation Model For Scientific Computing
In zero-shot learning evaluation, we explore the minimum number of time steps necessary formulate accurate neural operators. We also probe the OmniArch model's ability to generalize new physics scenarios without parameter adjustments. As indicated in Table 4 and Figure 10, longer temporal context typically enhances model performance, resulting in lower nRMSE scores across tasks. Notably, our model exhibits impressive zero-shot learning capabilities, maintaining robustness against mesh and temporal interpolation variations, even with fewer than 20 time steps of context.
H.2. Dynamic Prompt Length for Efficient Inference
We examine the trade-off between inference speed and accuracy using dynamic prompt lengths in our model. The goal is to determine whether shorter prompts can accelerate inference times on the CPU without significantly sacrificing precision.
Our approach varies the prompt length from 2 tokens (derived from a 50 time step interval) to 100 tokens (from a I time step interval) to predict physical fields U101- As shown in Figure 6, longer prompts yield higher precision with less variance, while shorter prompts can expedite inference by up 10 times compared to full-length prompts. In particular, model demonstrates an inherent ability to learn temporal differences from the input sequence, negating the need for explicit time-step inputs.
H.3. Fine-tuned for Inverse Problems
Demonstrating a model's capability to infer hidden physical parameters from known equations is critical test of its ability to learn underlying physics. Following the methodology of MPP (McCabe et al., 2023), we evaluate our model on two inverse problems for incompressible Navier-Stokes equations: 1) Forcing Identification, and 2) Buoyancy Determination.
Table 16: RMSE for Parameter Estimation in Inverse Problems.
Methods
Forcing
Buoyancy
MPP
0.2 + 0.008
0.78 + 0.006
OmniArch
0.16 + 0.005
0.73 + 0.012
Scratch 0.39 + 0.012 0.83 + 0.027
The results in Table 16 demonstrate that OmniArch outperforms MPP in parameter estimation tasks, with lower RMSE values indicating more accurate predictions. Models trained from scratch yield the highest errors, underscoring the effectiveness of fine-tuning approach. This evidence supports the notion that OmniArch is not only proficient in forward simulations but also exhibits superior performance in deducing hidden dynamics within complex systems.
H.4. Rollout Predictions
We perform rollout experiments to compare the performance of the Fourier Neural Operator (FNO) model and our proposed OmniArch model, depieted in Figure 11, 12, 13, 14. Our findings indicate that OmniArch demonstrates superior adherence to the underlying physics laws in the initial timesteps, opposed to merely replicating patterns from other trajectories. This improved fidelity is likely result of fine-tuning with PDE-Aligner, which isolates the model from the influences of alternate PDE systems, thereby enhancing the model's 's ability to generalize physical dynamics.
H.S. Multi-scale Inference Results
To thoroughly evaluate the multi-scale forecasting capabilities of OmniArch, extensive experiments were conducted across four different grid resolutions: 32 x 32, 64 x 64, 128 x 128, and 256 x 256. Figure 15 presents the visualization results at + 50 time step on the Incom datasct. These results demonstrate OmniArch's 's robust ability to accurately capture local patterns across varying grid sizes, confirming its effectiveness in handling
data without losing detail or accuracy.
H.6. More results in different problem settings
We tested our model on CFD-2D problems under various settings of the Navier-Stokes equations to evaluate its performance across different scenarios. The goal was to determine the robustness and adaptability of our model, OmniArch, compared to
24

---

OmniArch: Building Foundation Model For Scientific Computing
T+1
T+2
T+3
T+4
T+5
T+6
Ground Truth
Predicted
Figure 11: Prediction results of OmniArch CFD-2D dataset. Displaying time steps T+1 to T+6, the top shows ground truth data, and the bottom row illustrates OmniArch's predictions.
T+1
T+2
T+3
T+4
T+5
T+6
Ground Truth
Predicted
Figure 12: Prediction results of OmniArch on CFD-2D dataset. Displaying time steps T+1 to T+6, the top row shows ground truth data, and the bottom row illustrates Orni'sch's predictions.
T+5
T+10
T+15
T+20
T+25
T+30
Ground Truth
Predicted
Figure 13: Prediction results of OmniArch on SWE dataset. Displaying time steps T+5 to T+30, the top row shows ground truth data, and the bottom row illustrates OmniArch's predictions.
25

---

OmniArch: Building Foundation Model For Scientific Computing
T+1 T+2
T+3
T+4
T+5
T+6
Ground Truth
Predicted
Figure 14: Prediction results of OmniArch on Incom dataset. Displaying time steps T+1 T+6, the top row shows ground truth data, and the bottom illustrates OmniArch's predictions.
Grid Size 32
Grid Size 64
Grid =128
Grid Size =256
Label
Pred
Figure 15: Multi-scale results of OmniArch-Large with different grid sizes.
26

---

OmniArch: Building Foundation Model For Scientific Computing
other state-of-the-art models like MPP, FNO, and U-Net.
Table 17 summarizes the performance results of our model, OmniArch (FT), against MPP (FT), FNO, and U-Net across multiple problem settings. These settings include variations in Mach number (M), viscosity (n), and diffusivity (ς), both in inviscid and turbulent conditions with random periodic boundary conditions.
Table 17: The different problem settings in 2D Navier Stokes equation performance.
Problem Settings
OmniArch(FT) MPP(FT)
FNO
U-Net
M = 0.1, inviscid Rand periodic
0.1600
0.5866
0.38
0.66
M = 0.1, 1) = _ 0.01 Rand periodic
0.1215
0.5286
0.17
0.71
M 0.1, 77 { 0.1 Rand periodic
0.0273
0.5761
0.36
5.1
M 1.0,
0.01 Rand periodic
0.1301
0.5096
0.196
0.36
M 1.0, inviscid Rand periodic
0.1387
0.5391
0.35
0.47
M 1.0. 7) < 0.1 Rand periodic
0.0308
0,5033
0.098
0.92
M = 0.1, inviscid Turb periodic
0.2219
0.3949
0.16
0.19
M = 1.0, inviscid Turb periodic
0.1624
0.5412
0.43
0.14
These results consistently show that OmniArch performs better across various settings, demonstrating its robustness and effectiveness. The performance advantage of OmniArch evident across different Mach numbers, viscosity, and diffusivity settings, both in inviscid and turbulent conditions. These findings highlight the model's capability generalize and maintain high accuracy in diverse and challenging CFD scenarios.
H.7. GPU Memory Usage and Inference Time
We also report the runtime and memory usage in Table 18. OmniArch consistently uses less GPU memory than MPP across all model sizes, demonstrating its efficiency in resource utilization. While FNO and U-Net have lower GPU memory usage and faster inference times, OmniArch 's performance remains competitive, particularly considering its ability to handle wider range of PDE tasks across ID, 2D, and 3D domains.
Table 18: The runtime and memory usage between different models.
Model
Size
GPU Memory Inference Time
Tiny
671MB
0.0125s
Small
866MB
0.0129s
OmniArch
Base
1591MB
0.0136s
Large
3109MB
0.0248s
Tiny
1378MB
0.0387s
Small
1532MB
0.0390s
MPP
Base
1620MB
0.0391s
Large
3270MB
0.0831s
FNO
690MB
0.0018s
U-Net
830MB
0.0027s
H.8. Comparison with Traditional Solvers
Our benchmarks reveal three key advantages of OmniArch over traditional solvers:
Resolution Invariance: While FDM computation time scales quadratically (O(n²)) with grid resolution, OmniArch maintains nearly constant inference time (23-26ms) due to its fxed-frequeet processing in the spectral domain. This yields exponential speedup (155x at 512x512) for high-resolution simulations.
27

---

OmniArch: Building Foundation Model For Scientific Computing
Table 19; Computational Efficiency Comparison (2D Advection)
Resolution
FDM Time/Step (ms)
OmniArch Time (ms)
Speedup Relative Error
64x64
1.123
23.567
0.048
1.24x
128 128
15.264
23.820
0.641
1.18x
x92x 192
75.360
24.098
3.128x
1.15x
256x256
254.027
24.083
10.55x
1.12x
x20x320
583.218
23.866
24.44x
1.09x
384 x384
1130.561
23.453
48.20x
1.07x
448 448
2272.206
23.677
95.96x
1.05x
512
4073.472
26.212
155.4x
1.03x
Accuracy Preservation: Despite dramatic speed improvements, OmniArch maintains comparable accuracy with relative error consistently below 1.25x of FDM results. The enror margin decreases higher resolutions (1.03x x at x12), suggesting better performance in practical high-fidelity scenarios.
Generalization Capability: Unlike traditional methods requiring re-discretization for new PDEs, OmniArch's unified architecture achieves this performance across multiple physics domains (Navier-Stokes, Advection- Diffuusion, etc.) without algorithmic modifications, demonstrated in Section 4.2.
The results validate our design choice of spectral-doman processing while sacrificing some interpretability inherent to mesh-based methods, OmniArch gains orders-of-magnitude efficiency improvements crucial for large-scale multi-physies simulations. This trade-off aligns with emerging trends in scientific ML where learned simulators complement (rather than replace) traditional methods for specific high-throughput applications.
I. More Discussions
L1. Meta-Learning vs. Scaling Laws in PDE Solving
While meta-learning methods (Chen et al., 2022: Huang et al., 2022; Cho et al., 2023) address generalization through gradient-based adaptation, OmniArch explores an orthogonal axis: scaling laws for in-context learning. The distinction mirrors "learning optimize" versus "learning from data" peradigms-meta-Pine new PDEs, whereas foundation models leverage scale discover physics-award aware primitives. These approaches need not compete; future work might hybridize them. We may imagine meta-learning the hypernetworks of foundation model.
refine their optimization trajectory for
28


Reviews:
[{"review_id": "olMkigjDVI", "reviewer": "Reviewer_6qGA", "summary": "OmniArch is a foundation model for solving partial differential equations (PDEs) across 1D, 2D, and 3D domains. It addresses three key challenges: multi-scale modeling (handling different grid dimensions and resolutions), multi-physics capability (processing multiple physical quantities simultaneously), and physical alignment (incorporating physics constraints). The architecture combines a Fourier encoder-decoder for unified multi-dimensional training with a transformer backbone using temporal masking for multi-physics systems. A novel PDE-Aligner module enables physics-informed fine-tuning. Pre-trained on PDEBench datasets and fine-tuned with the PDE-Aligner, OmniArch achieves state-of-the-art performance on 11 PDE types while demonstrating emergent capabilities like zero-shot generalization and in-context learning for unseen PDEs.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank Reviewer  6qGA  for the constructive feedback and recognition of our work. We deeply appreciate  Reviewer 6qGA's support and hope this work can contribute to the research community. Below, we address the reviewers' questions and provide additional clarifications:\n\n**Q1: The zero-shot studies are limited in nature.**\n\n**A1**:  We acknowledge the scope limitation in zero-shot tests(currently on 3 PDE families in Table 2). However, emergent behaviors like **in-context learning**(in Figure 6) and **cross-resolution generalization** (in Figure 4) are observed and we will add a dedicated section, which includes more tests on new PDE datasets (such as  **the-well dataset** mentioned by Reviewer urwm).\n\n----\n\n**Q2: The focus on accuracy neglects computational efficiency, scaling, and deeper physical consistency analysis.**\n\n**A2**: We agree this is critical. While accuracy was our primary focus, we included inference time in Appendx(H.2&H.7).  More related details:\n+  **Efficiency**:  We add  the modal params and inference time below(2DCFD, 128*128, Single H800):\n| Model        | Params | Inference Time|\n|---|---|---|\n|Poseidon-L  | 629 M  |77ms        |\n|MPP-AVIT-L |409 M  | 83ms        |\n|DPOT-L        |509 M  | 32ms        |\n|OmniArch-L |445 M  | 21ms        |\n\nWhere the parameters of pre-trained models are similar and the inference time is In the same order of magnitude.\n\n+ **Scaling** : We'll add a discussion on how scaling in numerical foundation models differs from LMs or VLMs, addressing not only model/data size but also PDE diversity and input resolution dimensions. Here we can only provide midterm results for reference(Sorry for not having enough time & resources for full training currently):\n| Model Size |  Params |  nRMSE(2D)  |\n|---| ---|---|\n|Tiny| 26M |0.0847|\n|Small| 38M|0.0362|\n|Base | 144M|0.0153|\n|Large| 445M|0.0125|\n\n+ **Physical consistency analysis**:  We will include a section analyzing conservation properties, boundary condition satisfaction, and physical law adherence across different PDE types. For instance, we'll quantify momentum/energy conservation in CFD predictions and demonstrate how the PDE-Aligner specifically improves physical consistency. \nHere we give an example analysis on 2DCFD, which we may extend to other PDEs in the final revision.\n\n| Evaluation Metric | OmniArch-L | MPP-L | FNO | U-Net |\n|---------|------------|-------|-----|-------|\n| **Mass Conservation** |\n| Average Relative Error (%) | **0.32** | 0.65 | 2.17 | 3.24 |\n| Maximum Relative Error (%) | **0.78** | 1.21 | 4.85 | 7.36 |\n| **Energy Conservation** |\n| Average Relative Error (%) | **0.58** | 0.92 | 3.25 | 4.89 |\n| Maximum Relative Error (%) | **1.25** | 4.52 | 5.94 | 9.17 |\n| Continuity Error (∇·(ρu)) | **5.2e-4** | 2.1e-3 | 3.7e-3 | 5.4e-3 |\n\n----\n\n**Q3:  The insight into the possibilities of emergent behavior in such numerical foundation models.**\n\n**A3**:  Yes,  we are very happy to share our findings:\n\n(1) **Numerical foundation models can transfer Cross dimension** : We find pre-train on 1D/2D PDEs improves 3D performance(**in Table 1**), suggesting the natural dynamic pattern may be dimension-agnostic and has a much smaller latent dimension. The potential of a unified neural solver is still under-explored.\n\n(2) **Numerical foundation models can learn in-context**: Different from previous small neural solvers (which feed one step and predict the next step), we find OmniArch could learn from the temporal trajectory (the previous k steps) and adjust its prediction based on its observations (**see in Figure 6**).\n\n(3) **Numerical foundation models can be Resolution-agnostic**:  Due to trained directly on the frequency domain, OmniArch could learn in infinite resolutions (thanks to low-frequency truncating), and the low-resolution patterns could help understand high-resolution inputs(**see in Figure 4**).", "claims_and_evidence": "The paper's claims are largely supported by evidence, including performance analysis on 11 PDE tasks and ablation studies.  However, the zero-shot studies are limited in nature.", "methods_and_evaluation_criteria": "Methods address stated challenges, with Fourier transforms for multi-scale data and temporal masking for multi-physics systems. nRMSE is an appropriate evaluation metric, tested across PDEBench and PDEArena. The focus on accuracy neglects computational efficiency, scaling, and deeper physical consistency analysis, which limits the evaluation's breadth.", "theoretical_claims": "No significant theoretical claims are made in this paper.", "experimental_designs_or_analyses": "The experimental design is generally sound, comparing against both task-specific expert models and other unified pre-training approaches.", "relation_to_broader_scientific_literature": "The relationship to the broader scientific literature in numerical foundation models is appropriate for the paper.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "None.", "other_comments_or_suggestions": "None.", "questions_for_authors": "Did this work give you any insight into the possibilities of emergent behavior in such numerical foundation models?", "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": "CABMlszSEs", "reviewer": "Reviewer_u5RW", "summary": "This paper introduces OmniArch, a foundation model designed for solving multi-scale and multi-physics PDEs. Inspired by foundation models in NLP, OmniArch aims to generalize across different PDEs using a Fourier-based encoder-decoder, a Transformer backbone, and a physics-informed fine-tuning method. The Fourier encoder-decoder enables learning across varying spatial scales and the transformer captures complex temporal dependencies. Furthermore, PDE-Aligner ensures physics consistency by aligning model predictions with governing equations. OmniArch is trained across different PDEs and is showing superior performance compared to task-specific methods and pre-trained PDE solvers.\n\n\n## update after rebuttal\nAuthors have included a discussion and additional experiments to compare their method against traditional solvers and will include a discussion about relevant meta-learning papers in the final version which will address my initial concerns about this paper.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer  u5RW for the thoughtful feedback and valuable support of our work. Below, we address each of their questions in detail:\n\n**Q1**: How does OmniArch compare to traditional solvers (e.g., FEM, spectral methods) in accuracy/interpretability?**\n\n**A1**: We thank u5Rw for this critical point. While traditional solvers excel in interpretability, OmniArch targets fully leverages the efficiency in GPU settings, especially for input size scaling(similar to Q1 in 6qGA):\n+ **Accuracy**: on 1D Advection Rollout, OmniArch(RMSE: 0.0321) matches FDM methods(0.0258) but avoids costly re-discretization for new PDEs.\n+ **Speed**: Traditional solvers(FDM, FEM, Spectral methods):  for an input grid with s,s nodes to iterate forward the computation of FDM is $O(s^4)$ while the computation of Spectral methods is $O(s^2log s)$, however, the computation of OmniArch is  $O(tk^2)$, k in a fixed frequency, which allows OmniArch achieves 155x faster inference(512x512,2D, 200 steps ->0.026s) than FDM(~4.5s) at comparable error. We provide a table below, you can also find in [Anonymous Visualization](https://various-easy-trillium.glitch.me/).\n\n| Resolution | FDM Time/Step (s) | OmniArch Time (s) | Speedup (FDM/OmniArch) |\n|------------|-------------------|-------------------|------------------------|\n| 64×64      | 0.001123          | 0.023567          | 0.048x                 |\n| 128×128    | 0.015264          | 0.023820          | 0.641x                 |\n| 192×192    | 0.075360          | 0.024098          | 3.128x                 |\n| 256×256    | 0.254027          | 0.024083          | 10.55x                 |\n| 320×320    | 0.583218          | 0.023866          | 24.44x                 |\n| 384×384    | 1.130561          | 0.023453          | 48.20x                 |\n| 448×448    | 2.272206          | 0.023677          | 95.96x                 |\n| 512×512    | 4.073472          | 0.026212          | 155.4x                 |\n\n\n\n*Revision*: We will add a more dedicated table comparing accuracy/runtime against traditional solvers in the final version.\n\n----\n\n**Q2: Why omit meta-learning approaches (e.g., Meta-PINNs) despite their relevance to generalization?**\n\n**A2**:  We appreciate this insight. Meta-learning is indeed complementary:\n+ **key difference**: Meta-PINNs adapt via gradient updates, while OmniArch enables **token-based in-context learning**(Fig 10)  without fine-tuning.\n\nWe will cite suggested references and clarify this distinction in the final version, emphasizing OmniArch's **architecture-driven generalization** vs. optimization-based meta-learning.\n\n----\n\n**Q3: Can PDE-Aligner guarantee physically valid solutions for all PDEs?**\n\n**A3**: The aligner ensures soft constraints via textual equation supervision, it only help the omniarch better distinguish between different PDE systems. There is no theoretical guarantee for all PDEs(an open challenge even for traditional solvers). But we believe this maybe solved by better augmentation ways(equation-search, lie symmetry, etc.) and more efficient contrastive learning techniques, which we hope to see in the near future. This discussion will be added in the camera-ready version.\n\n----\n\n**Q4:Typographical errors in the abstract?**\n\n**A4**: Thanks for the kindly reminder, we will correct them in the final version.", "claims_and_evidence": "I believe, the paper presents strong experimental results and clear methodological justifications for their claims. \n1) Generalization across multiple PDEs: This claim is supported by training on 11 PDE types from PDEBench and PDEArena datasets and the results show superior performance over task-specific and pre-trained baselines.\n2) OmniArch is multi-scale: This claim is justified through frequency-space transformation, allowing the model to handle different spatial resolutions and experiments include multi scale PDEs.\n3) PDE-Aligner improves physics alignment: Supported by physics-based losses.\n4) zero-shot generalization: Supported by experiments on solving unseen PDEs with no additional training.\n\nMissing:\n1) Comparison to traditional PDE Solvers: The paper does not benchmark runtime or accuracy against finite element or spectral methods, which remain the gold standard in PDE solving. Understanding the trade-offs between accuracy and interpratability compared to these methods would strengthen the paper.", "methods_and_evaluation_criteria": "I believe the methods and evaluation criteria are well-aligned with the goal of designing a foundation model for PDE solving. \n1) Use of Fourier encoder-decoder for multi-scale PDEs seems like a reasonable approach.\n2) The self-attention mechanism enables long-range temporal dependencies.\n3) The design of PDE-Aligner ensures physics consistency of the solutions. \n4) The datasets cover a diverse range of PDEs. \n\nThe only missing evaluation in my opinion, is benchmarking accuracy and interpretability against traditional solvers.", "theoretical_claims": "The paper does not include formal theoretical proofs, as its focus is primarily on empirical validation and model design.", "experimental_designs_or_analyses": "I believe the paper has a strong and well-structured experimental setup and extensive experiments. \n1) The method is evaluated across 11 PDEs. \n2) The method is evaluated against multiple task-specific and pre-trained benchmarks.\n3) The experiments setup includes both zero-shot and in-context learning tests.", "relation_to_broader_scientific_literature": "This paper builds on prior work in neural PDE solvers, foundation models in NLP, and physics-informed learning. It integrates multiple established ideas into a unified framework. I believe this paper has the potential to shape the future in scientific machine learning and it can initiate significant research in foundation models for scientific computing, both for further technical advancements and also different applications.", "essential_references_not_discussed": "While the authors discuss neural PDE solvers and physics-informed learning, they do not address meta-learning approaches for PDE solving, such as Meta-PINNs, which focus on improving generalization that is a key motivation of this paper. Additionally, the authors mention the need for retraining PINNs as a limitation, yet meta-learning techniques have been developed specifically to resolve this issue by enabling faster adaptation to new PDEs. Relevant references on meta-learning for PINNs are missing, despite their direct relevance to OmniArch’s goals. Here are a few examples:\n1) \"Hypernetwork-based Meta-Learning for Low-Rank Physics-Informed Neural Networks\",NeurIPS, 2023\n2) \"Meta-MGNet: Meta Multigrid Networks for Solving Parameterized Partial Differential Equations\", Journal of computational physics, 2022\n3) \"Meta-Auto-Decoder for Solving Parametric Partial Differential Equations.\", NeurIPS, 2022", "other_strengths_and_weaknesses": "To the best of my knowledge, the method is very novel and contributions are significant. The only weakness, is the lack of evaluation against traditional and numerical solvers and also missing references and discussions regarding meta-learning approaches for PDE solving such as meta-PINNs.", "other_comments_or_suggestions": "There are typos in the abstract: 1) \"while whether\" in the second line, and 2) \"we first conduct\" in line 15.", "questions_for_authors": "While PDE-Aligner enforces physics constraints through contrastive learning, is there any theoretical guarantee that it produces physically valid solutions for all PDE types?", "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": "4fZmNGeUBZ", "reviewer": "Reviewer_urwm", "summary": "This paper proposes OmniArch, a foundation model for numerical simulations. They pretrain on 1D/2D/3D data and comapre to other models of the literature. It relies on spatial Fourier encoders/decoders and a causal temporal attention. It relies on PDE-Aligner for fine-tuning.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer urwm for the thoughtful feedback. \n\n**Q1: MPP/Poseidon already did united pretraining.**\n\n**A1**:   We **disagree** with highly respect. There is **factual error**, MPP/Poseidon are designed *only* for 2D (no 1D/3D experiments in their papers). DPOT uses a convolution trick for 3D but lacks unified weights. OmniArch is the first to jointly learn 1D-2D-3D PDEs with shared weights in one architecture. We will clarify this distinction.\n\n-----\n\n**Q2: Table 1 mixes nRMSE/VRMSE.**\n\n**A2**:  We believe there is **misinterpretation** here.   In Equation 8, We implement nRMSE following the common practice in previous works(MPP/PDEBench). The notation $\\sigma_u$ here is equivalent to the $||\\cdot||^2 + \\epsilon$, which measures the norm of GT. We evaluate the predictions from all baselines and OmniArch with the same codebase. For fairness, we compare our nRMSE implementation and MPP, the results are the same and reproducible. Here is the  code:\n\n```python\ndef nRMSE_MPP(output, tar, spatial_dims=None):\n# code from https://github.com/PolymathicAI/multiple_physics_pretraining\n    if spatial_dims is None:\n        spatial_dims = tuple(range(output.ndim))[2:]  # Assume 0, 1, 2 are T, B, C\n    residuals = output - tar\n    tar_norm = (1e-7 + tar.pow(2).mean(spatial_dims, keepdim=True))\n    raw_loss = (residuals.pow(2).mean(spatial_dims, keepdim=True) / tar_norm)\n    return raw_loss.sqrt().mean()\n\n# Ours：\ndef nRMSE(pred, label, mask=None, dim=2):\n    reduce_dims = list(range(-dim, 0))\n    res = pred - label\n    label_norm = label.pow(2).mean(dim=reduce_dims, keepdims=True) + 1e-8\n    norm_loss = res.pow(2).mean(dim=reduce_dims, keepdims=True) / label_norm + 1e-8 # in extreme case, label_norm may be 0\n    if mask is not None:\n        norm_loss = norm_loss[mask.bool()]\n    return norm_loss.sqrt().mean()\n```\n----\n\n**Q3: OmniArch-B (316M params) vs. MPP-B (116M) is unfair.**\n\n**A3**:  There is **Misunderstanding** here, Note that OmniArch uses different encoders/decoder pairs for 1D,2D,3D PDEs while the weight of transformer backbone is shared.(works like Mixture-of-Experts Models). We report the Total Parameters(316M, including transformer backbone +  1D, 2D, 3D encoder/decoder weights, the 3D encoder/decoder weights take a heavy part~170M) in Table 7. But for 2D task(MPP mainly experiment on), only 2D encoder/decoders are used for training/testing , thus the real runtime parameters of **OmniArch-B is 144M** which is comparable to **MPP-B(116M)**  but outperforms in all 2D tasks **(see in Table 1)** , The same for large size model, **OmniArch-L(445M) is also comparable to MPP-L(409M)**.  We will clarify this detailed difference further in the next version, but  we believe the experiments are indeed fair.\n\n----\n\n**Q4: No evidence of emergent behaviors; term is undefined.**\n\n**A4**:  *Emergent behaviors* is the concept borrowed from foundation models, which could be explained as \"unseen capabilities beyond training\" , it can be the capability that large model exhibits but small models not. In our settings, \n\n**(1) Zero-shot PDE solving**(Figure 5, beyond the training scope of PDE systems, not exhibit by smaller models like PINNs/FNO) \n\n**(2)  In-context learning**(Figure 6), which OmniArch could learn from input trajectories. We are glad to include more experiments to comprehensively evaluate the emergent behaviors together with the research community. \n\n----\n\n**Q5: How is OmniArch different from a Fourier-encoded transformer?**\n\n**A5**: OmniArch differs fundamentally from a Fourier-encoded transformer in :\n\n1. **Temporal Mask mechanism**: OmniArch implements a specialized Temporal Mask that enables each physical quantity to attend to all quantities at current and previous time steps, facilitating complex cross-physics interactions that standard transformers cannot model effectively.\n\n2. **Physics-informed alignment**: The PDE-Aligner component provides a dedicated mechanism to incorporate physical constraints and prior knowledge during fine-tuning, which goes far beyond basic Fourier encoding.\n\n---\n**Q6: Show ablation with/without PDE-Aligner.**\n\n**A6**:  Please find the ablation in Q2(Reviewer VHRd), where PDE-Aligner continually improves the performance across 1D-2D-3D PDEs.\n\n----\n\n\n**Q7: How are frequency-space tokens handled**\n\n**A7**: OmniArch converts physical fields into frequency space via FFT, then processes the data by first selecting only the most important frequency components (Top-K by magnitude) for efficiency. These complex-valued coefficients are transformed into real-number embeddings through a learnable projection, making them compatible with standard transformer architectures.\n\n----\n\n**Q8: Why use BERT for tiny text vocab?**\n\n**A8**: Because (1) Vocabulary is not tiny (augmented via symbol replacement, Appx. F.1). (2) Pretrained embeddings ensure generalization (vs. retraining from scratch).\n\n----\n\n**Q9: Cite the-well paper?**\n\n**A9**: We will include this in the next version.", "claims_and_evidence": "I don't think it is fair to say that this is the first foundation model on PDE data. The cited models like MPP, Poseidon and DPOT were already pretrained on multiple physics.\n\nThere are no experiments on the emerging capabilities? I am not sure it is defined somewhere? \n\nAs you can see below, I also don't trust the reported metrics in the Table.", "methods_and_evaluation_criteria": "Table 1 makes sense, even thought it would have been better to add the size of the models. For example, OmniArch-B is 316.13M parameters whereas e.g. MPP-B is 116M parameters. Same for Large models. I think the comparison between models is not totally fair which makes the experimental results difficult to interpret.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "The comparison between models is not 100% fair due to their different size. \n\nOne important point: The authors gave the definition of the VRMSE (Variance Root Mean Square error) but name is nRMSE (normalized root mean square error). This is confusing and makes me wonder if the reported number in Table 1 are for the same metrics (MPP and Poseidon used the original nRSME but the authors compute the VRMSE for their models)", "relation_to_broader_scientific_literature": "I think the authors wrongly name themselves the \"first foundation models for 1D-2D-3D united pre-training\". This has been done before in MPP and Poseidon. However, the fine-tuning part of OmniArch seems novel and interesting to me. \n\nThe authors could also have cited [1] and maybe benchmark their model on it as it seems to be a very relevant point of the literature? \n\n[1] Ohana, R., McCabe, M., Meyer, L., Morel, R., Agocs, F., Beneitez, M., Berger, M., Burkhart, B., Dalziel, S., Fielding, D. and Fortunato, D., 2024. The well: a large-scale collection of diverse physics simulations for machine learning. Advances in Neural Information Processing Systems, 37, pp.44989-45037.", "essential_references_not_discussed": "Except from [1], I don't see any reference not discussed.\n\n\n[1] Ohana, R., McCabe, M., Meyer, L., Morel, R., Agocs, F., Beneitez, M., Berger, M., Burkhart, B., Dalziel, S., Fielding, D. and Fortunato, D., 2024. The well: a large-scale collection of diverse physics simulations for machine learning. Advances in Neural Information Processing Systems, 37, pp.44989-45037.", "other_strengths_and_weaknesses": "- Abstract: \"as far as we know, we first conduct 1D-2D-3D united pre-training on the PDEBench,[...] should be rephrased.\n- In the introduction line 55, I don't think MPP and Poseidon are an extension of the Factformer? This is weirdly phrased.\n- I don't think I have noticed some emergent capabilities of the model? I am not even sure of what this mean or if this is well defined for this type of data?\n- the paragraph line 199 is not clear. What is the d-th index? What's the sequential index? What's the total? It should be rephrased clearly?\n- I don't see the justification of equation 4, as this is not for a transformer as stated, but only with one attention layer. I don't see the point of this remark since OmniArch is a whole transformer?\n- Line 240, I don't think there is a need for 2 paragraphs to explain what causal masking is? This is standard in the literature, and the written maths are more confusing than anything else. \n- As discussed above, Equation 8 is the VRMSE and not the NRSME, which makes me suspicious about the results shown in table 1.\n- Could the authors explain how the PDE-Aligner is really useful in that case? I may have missed it, but an ablation study with/without would have been useful.\n- How is the spatial tokenization done? I know we are in frequency space, so is there any tokenization done, or the authors are just doing a ravel in frequency space? What's the dimensionality then? \n- Why would the authors need a whole pre-trained BERT model for text encoding when the vocabulary is so small? \n- I struggle to understand the difference between OmniArch and a transformer with Fourier layers as encoder/decoders. Could the authors elaborate on that?\n- Can the model work on non-uniform grids? In theory it should?", "other_comments_or_suggestions": "See weaknesses above.", "questions_for_authors": "See weaknesses above.", "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": "FOZVnh3L7c", "reviewer": "Reviewer_VHRd", "summary": "The paper introduces *OmniArch*, a foundation model for scientific computing designed to solve multi-scale and multi-physics Partial Differential Equations (PDEs) using a unified architecture. It employs a Fourier Encoder-Decoder to transform spatial-temporal PDE data into the frequency domain and a Transformer Backbone to capture temporal dependencies, enabling it to handle 1D, 2D, and 3D PDEs within a single framework. A key innovation is the *PDE-Aligner*, which fine-tunes the model with physics-informed constraints, ensuring alignment with governing physical laws. OmniArch achieves state-of-the-art performance across 11 PDE types, demonstrating strong generalization capabilities, including zero-shot learning, in-context learning, and multi-scale inference. Compared to existing models, it significantly improves accuracy, with up to 98.7% enhancement in some cases, making it a versatile tool for applications in computational fluid dynamics, weather prediction, and engineering simulations.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate reviewer VHRd's constructive feedback. Below, we address each concern with clarifications and additional evidence:\n\n**Q1: [Temporal Mask & Multi-Physics Consistency] How does the temporal mask address inconsistencies in multi-physics, and where is the ablation study?**\n\n**A1**: The \"inconsistencies\" arise when handling PDE systems with **varying numbers of physical variables per timestep** (e.g., 3D Navier-Stokes vs. 1D Advection). Traditional causal masks (e.g., GPT-style) fail here because:\n\n+ **Token misalignment**: For systems with multiple variables, causal masks process tokens sequentially, forcing the first token (e.g., velocity) to ignore dependencies on other variables (e.g., pressure, temperature) within the same timestep. \n\nOur  **Attention with Temporal Mask**: Instead of masking tokens individually, we group all variables within a timestep and forms a **hierarchical attention**:\n+ *Intra-timestep*: Full attention among variables (e.g., velocity, pressure) to capture couplings.\n+ *Inter-timestep*: Causal attention across time (future masking).\n\n**Benefits**: Physics-aware modeling: Ensures variables at timestep jointly condition on each other (e.g., enforcing continuity equations).\n\n**Ablations**:\n| Method          | RMSE |\n|-----------------|--------|\n|Causal Mask     |0.0277 |\n|No Mask           |0.0285|   \n|**Temporal Mask** |**0.0227**|\n\n----\n\n**Q2: [PDE-Aligner's Hidden Representations]  How does PDE-Aligner leverage hidden representations, and where is the evidence?**\n\n**A2**: We have report the fine-tuning results with PDE-Aligner in **Table 1(Line 298-302)**, We reorganize the results below:\n\n|Settings          | 1D      | 2D      |3D       |\n|--------------- |--------|--------|--------|\n|Pretrain           |0.0103|0.0440 |0.3399|\n|ft(w/o Aligner)|0.0073|0.0345 |0.3432|\n|ft + Aligner     |0.0056|0.0262 |0.2697|\n|**improvement** | **51.33%**|**23.62%**|**22.07%**|\n\nwhere **PDE-Aligner continually improves the performance across 1D-2D-3D PDEs**.\n\nWe also doing the probing experiment in **Figure 7**, which explains why PDE-Aligner helps the performance, mainly because it helps distinguish the difference between different PDE Systems.\n\n----\n\n**Q3: [Top-K Frequency Selection] Why select top-K frequencies, and does it improve efficiency/generalization?**\n\n**A3**: We adopt **Top-K frequency truncation** for two key reasons:\n\n **(a) Input-Size Agnosticism:** OmniArch jointly trains on 1D-2D-3D PDEs with vastly different resolutions (e.g., 1D-1024, 2D-128×128, 3D-64×64×64). Without truncation, we would either:\n+ *(Option 1)* Require separate encoders/decoders per resolution (losing unified modeling), or\n+ *(Option 2)* Downsample spatial dimensions, forcing 2D/3D tasks to compromise for 1D (hurting performance).\n\nTop-K ensures consistent spectral representation across dimensions, similar to FNO’s infinite-dimensional operator learning [Li et al.].\n\n**(b)Fluid Data Prior:** The PDEBench and PDEArena Datasets exhibit **long-tailed high-frequency noise[1,2,3]**. Retaining only dominant low-frequency modes (Top-K) preserves physically meaningful features while suppressing numerical artifacts. \n\n>[1] Takamoto, Makoto, et al. \"Pdebench: An extensive benchmark for scientific machine learning.\" Advances in Neural Information Processing Systems 35 (2022): 1596-1611.\n\n>[2] Lippe, Phillip, et al. \"Pde-refiner: Achieving accurate long rollouts with neural pde solvers.\" Advances in Neural Information Processing Systems 36 (2023): 67398-67433.\n\n>[3] Zakharov, Vladimir E., Victor S. L'vov, and Gregory Falkovich. Kolmogorov spectra of turbulence I: Wave turbulence. Springer Science & Business Media, 2012.\n\n----\n\n**Q4: [Model Size & Inference Time] How does OmniArch compare in size/speed to baselines?**\n\n**A4**: We include the Memory usage and Inference Time(single A800) in **Appendix H.7**, We reorganize the main results（Pretrained Models baselines，2D，Large size） below for reference. OmniArch is **competitive in efficiency** with other pretrained model baselines.\n| Model        | Params | Inference Time|\n|-------------|----------|----------------|\n|Poseidon-L  | 629 M|0.07712s           |\n|MPP-AVIT-L |409 M| 0.08343s        |\n|DPOT-L        |509 M| 0.03154s          |\n| **OmniArch-L** |**445 M**  | **0.02075s**  |\n\n\n----\n\n**Q5:  [Missing citaions.]Including PINNs/FNO/BSDE works.**\n\n**A5**: We appreciate the reviewer’s suggestion. **In our paper, PINN and FNOs are explicitly discussed as baselines (Section 5.1) and compared in Related Works (Lines 74–82)**. However, we omitted the Deep BSDE Method because its primary applications (e.g., financial derivatives pricing) are orthogonal to our focus on multi-physics PDEs in scientific computing (e.g., fluid dynamics, material science). While BSDEs excel in high-dimensional finance problems , their scope differs fundamentally from our work's goals. That said, we will include a brief comparison in the final version to clarify this distinction.", "claims_and_evidence": "This paper offers three claims in the introduction. However, some claims are not convincing.\n\n(1) The author claimed that \"The temporal mask effectively addresses inconsistencies in multi-physics\". However, I could not find the explanation. Also, I could not find the ablation study to prove it.\n\n\n(2) The PDE-Aligner is claimed to leverage the hidden representations for equations. However, these claims are not supported by the ablation studies.", "methods_and_evaluation_criteria": "The benchmark datasets and evaluation criteria is reasonable.\n\nHowever,\n\nIn Eq.(1), the author select top K significant parts. However, what is the reason? Can it improve the efficiency? Can it improve the generalization ability?\n\nAlso, I can not find how temporal mask be used in the proposed method. Also, there is no explicit ablation to verify its effectiveness.", "theoretical_claims": "There is no theoretical analysis.", "experimental_designs_or_analyses": "The experimental results are fully convincing.\n\n[1] The variation is not clarified in the experimental results.\n\n[2] The model size, training time, and inference time is not compared in the table. I seems that the proposed method is larger and slower than the competing methods.", "relation_to_broader_scientific_literature": "The key contributions of the *OmniArch* paper align with and extend several existing concepts in the broader scientific and machine learning literature. The unified architecture for multi-physics PDEs builds upon traditional neural operators like the Fourier Neural Operator (FNO), which are typically tailored to specific PDEs, by enabling simultaneous learning across diverse equations and conditions. The integration of a Fourier Encoder-Decoder with a Transformer Backbone leverages the efficiency of frequency domain representations and the capability of transformers to model long-range dependencies, enhancing generalization across different scales and physical phenomena. The introduction of the PDE-Aligner for physics-informed fine-tuning extends the principles of Physics-Informed Neural Networks (PINNs) by aligning model predictions with governing physical laws through contrastive learning in the frequency domain. Moreover, *OmniArch*'s emergent generalization capabilities, such as zero-shot and few-shot learning, parallel advancements in foundation models in other domains, demonstrating adaptability to novel PDE systems without retraining.", "essential_references_not_discussed": "Upon reviewing the paper *OmniArch: Building Foundation Model for Scientific Computing*, I have identified several pertinent works related to partial differential equations (PDEs) that are not currently cited but are essential for understanding the context of the paper's key contributions:\n\n1. **Physics-Informed Neural Networks (PINNs)**: PINNs integrate physical laws described by PDEs into the learning process of neural networks. They offer a mesh-free alternative to traditional numerical methods for solving PDEs and have been applied to various problems in computational science. Notable works include:\n\n   - Raissi, M., Perdikaris, P., & Karniadakis, G. E. (2019). 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.\n\n   - Mao, Z., Jagtap, A. D., & Karniadakis, G. E. (2020). Physics-informed neural networks for high-speed flows. *Computer Methods in Applied Mechanics and Engineering, 360*, 112789.\n\n2. **Neural Operators**: Neural operators are deep learning architectures designed to learn mappings between infinite-dimensional function spaces, making them effective for approximating solution operators of PDEs. Key publications include:\n\n   - Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). Fourier neural operator for parametric partial differential equations. *arXiv preprint arXiv:2010.08895*.\n\n   - Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhattacharya, K., Stuart, A., & Anandkumar, A. (2021). Neural operator: Learning maps between function spaces. *Journal of Machine Learning Research, 22*(1), 139–166.\n\n3. **Deep Backward Stochastic Differential Equation (BSDE) Method**: This method combines deep learning with backward stochastic differential equations to solve high-dimensional PDEs, particularly in financial mathematics. A seminal paper in this area is:\n\n   - Han, J., Jentzen, A., & E, W. (2018). Solving high-dimensional partial differential equations using deep learning. *Proceedings of the National Academy of Sciences, 115*(34), 8505–8510.\n\nIncorporating discussions of these methodologies into the paper would provide a more comprehensive understanding of existing advancements in PDE-solving techniques and AI applications in scientific computing, thereby contextualizing OmniArch's contributions within the broader research landscape.", "other_strengths_and_weaknesses": "NO.", "other_comments_or_suggestions": "NO.", "questions_for_authors": "[1] The authors should add variations, model size and inference time of the proposed method and the competing method.\n\n[2] The authors should add ablation studies to verify the claims in the introduction.\n\n[3] The author should explain how temporal mask  benefits solving PDE.", "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: This paper presents OmniArch, a foundation model for solving PDEs across 1D, 2D, and 3D domains. Reviewers highlighted its novel unified architecture, which addresses limitations of prior work (e.g., MPP, Poseidon), and its broad evaluation across 11 PDEs.

Initial opinions were split, with two reviewers in favor and two leaning against. Concerns on emergent capabilities, dataset diversity, and presentation were raised. However, following rebuttals and reviewer-AC discussions, there was general agreement that:
- OmniArch shows zero-shot generalization and in-context learning for unseen PDEs.
- While the evaluation dataset is standard, it covers three entirely distinct classes, i.e., parabolic, hyperbolic, and elliptic PDEs, and can serve as a starting point for further advancements in the field.

We recommend weak acceptance to reflect the paper’s valuable contributions to the SciML community. To strengthen the final version, the authors should:
- Clarify architectural details, especially temporal masking.
- Streamline equations and improve presentation quality.
- Add parameter count tables by dimension.
- Include ablation studies on key components.

With these revisions, this paper can make contributions to developing foundation models beyond traditional domains.
Decision: Accept (poster)