OmniArch: Building Foundation Model for Scientific Computing

We sincerely appreciate reviewer VHRd's constructive feedback. Below, we address each concern with clarifications and additional evidence:

Q1: [Temporal Mask & Multi-Physics Consistency] How does the temporal mask address inconsistencies in multi-physics, and where is the ablation study?

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:

• 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.

Our Attention with Temporal Mask: Instead of masking tokens individually, we group all variables within a timestep and forms a hierarchical attention:

• Intra-timestep: Full attention among variables (e.g., velocity, pressure) to capture couplings.
• Inter-timestep: Causal attention across time (future masking).

Benefits: Physics-aware modeling: Ensures variables at timestep jointly condition on each other (e.g., enforcing continuity equations).

Ablations:

---

Q2: [PDE-Aligner's Hidden Representations] How does PDE-Aligner leverage hidden representations, and where is the evidence?

A2: We have report the fine-tuning results with PDE-Aligner in Table 1(Line 298-302), We reorganize the results below:

where PDE-Aligner continually improves the performance across 1D-2D-3D PDEs.

We 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.

---

Q3: [Top-K Frequency Selection] Why select top-K frequencies, and does it improve efficiency/generalization?

A3: We adopt Top-K frequency truncation for two key reasons:

(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:

• (Option 1) Require separate encoders/decoders per resolution (losing unified modeling), or
• (Option 2) Downsample spatial dimensions, forcing 2D/3D tasks to compromise for 1D (hurting performance).

Top-K ensures consistent spectral representation across dimensions, similar to FNO’s infinite-dimensional operator learning [Li et al.].

(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.

[1] Takamoto, Makoto, et al. "Pdebench: An extensive benchmark for scientific machine learning." Advances in Neural Information Processing Systems 35 (2022): 1596-1611.

[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.

[3] Zakharov, Vladimir E., Victor S. L'vov, and Gregory Falkovich. Kolmogorov spectra of turbulence I: Wave turbulence. Springer Science & Business Media, 2012.

---

Q4: [Model Size & Inference Time] How does OmniArch compare in size/speed to baselines?

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.

---

Q5: [Missing citaions.]Including PINNs/FNO/BSDE works.

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.

We sincerely thank the reviewer urwm for the thoughtful feedback.

Q1: MPP/Poseidon already did united pretraining.

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.

---

Q2: Table 1 mixes nRMSE/VRMSE.

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:

def nRMSE_MPP(output, tar, spatial_dims=None):
# code from https://github.com/PolymathicAI/multiple_physics_pretraining
    if spatial_dims is None:
        spatial_dims = tuple(range(output.ndim))[2:]  # Assume 0, 1, 2 are T, B, C
    residuals = output - tar
    tar_norm = (1e-7 + tar.pow(2).mean(spatial_dims, keepdim=True))
    raw_loss = (residuals.pow(2).mean(spatial_dims, keepdim=True) / tar_norm)
    return raw_loss.sqrt().mean()

# Ours：
def nRMSE(pred, label, mask=None, dim=2):
    reduce_dims = list(range(-dim, 0))
    res = pred - label
    label_norm = label.pow(2).mean(dim=reduce_dims, keepdims=True) + 1e-8
    norm_loss = res.pow(2).mean(dim=reduce_dims, keepdims=True) / label_norm + 1e-8 # in extreme case, label_norm may be 0
    if mask is not None:
        norm_loss = norm_loss[mask.bool()]
    return norm_loss.sqrt().mean()

---

Q3: OmniArch-B (316M params) vs. MPP-B (116M) is unfair.

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.

---

Q4: No evidence of emergent behaviors; term is undefined.

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,

(1) Zero-shot PDE solving(Figure 5, beyond the training scope of PDE systems, not exhibit by smaller models like PINNs/FNO)

(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.

---

Q5: How is OmniArch different from a Fourier-encoded transformer?

A5: OmniArch differs fundamentally from a Fourier-encoded transformer in :

1. 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.

2. 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.

---

Q6: Show ablation with/without PDE-Aligner.

A6: Please find the ablation in Q2(Reviewer VHRd), where PDE-Aligner continually improves the performance across 1D-2D-3D PDEs.

---

Q7: How are frequency-space tokens handled

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.

---

Q8: Why use BERT for tiny text vocab?

A8: Because (1) Vocabulary is not tiny (augmented via symbol replacement, Appx. F.1). (2) Pretrained embeddings ensure generalization (vs. retraining from scratch).

---

Q9: Cite the-well paper?

A9: We will include this in the next version.

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:

Q1: How does OmniArch compare to traditional solvers (e.g., FEM, spectral methods) in accuracy/interpretability?**

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):

• Accuracy: on 1D Advection Rollout, OmniArch(RMSE: 0.0321) matches FDM methods(0.0258) but avoids costly re-discretization for new PDEs.
• 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.

Revision: We will add a more dedicated table comparing accuracy/runtime against traditional solvers in the final version.

---

Q2: Why omit meta-learning approaches (e.g., Meta-PINNs) despite their relevance to generalization?

A2: We appreciate this insight. Meta-learning is indeed complementary:

• key difference: Meta-PINNs adapt via gradient updates, while OmniArch enables token-based in-context learning(Fig 10) without fine-tuning.

We will cite suggested references and clarify this distinction in the final version, emphasizing OmniArch's architecture-driven generalization vs. optimization-based meta-learning.

---

Q3: Can PDE-Aligner guarantee physically valid solutions for all PDEs?

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.

---

Q4:Typographical errors in the abstract?

A4: Thanks for the kindly reminder, we will correct them in the final version.

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:

Q1: The zero-shot studies are limited in nature.

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).

---

Q2: The focus on accuracy neglects computational efficiency, scaling, and deeper physical consistency analysis.

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:

• Efficiency: We add the modal params and inference time below(2DCFD, 128*128, Single H800): | Model | Params | Inference Time| |---|---|---| |Poseidon-L | 629 M |77ms | |MPP-AVIT-L |409 M | 83ms | |DPOT-L |509 M | 32ms | |OmniArch-L |445 M | 21ms |

Where the parameters of pre-trained models are similar and the inference time is In the same order of magnitude.

• 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): | Model Size | Params | nRMSE(2D) | |---| ---|---| |Tiny| 26M |0.0847| |Small| 38M|0.0362| |Base | 144M|0.0153| |Large| 445M|0.0125|

• 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. Here we give an example analysis on 2DCFD, which we may extend to other PDEs in the final revision.

---

Q3: The insight into the possibilities of emergent behavior in such numerical foundation models.

A3: Yes, we are very happy to share our findings:

(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.

(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).

(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).