ICML 2025 | Beihang University Introduces OmniArch, a Scientific Foundation Model for PDE Solving with 11x Performance Improvement}

Author: Chen Tianyu

Editor: ScienceAI

"The books of nature are written in the language of mathematics." — Galileo

Partial Differential Equations (PDEs) are among the most fundamental chapters in this book, describing universal physical laws from quantum fluctuations to galaxy evolution, fluid impacts to electromagnetic propagation. However, solving PDEs is like deciphering nature’s code—traditional numerical methods are costly, and neural network models are often problem-specific, limiting scalability and knowledge transfer across complex multi-scale, multi-physics systems.

Can we build a universal model that understands and solves the 'language of physics' as ChatGPT does with language?

Recently, Beihang University’s Professor Li Jianxin’s team proposed OmniArch, a scientific foundation model that achieves a major breakthrough by solving 1D, 2D, and 3D PDEs with a single unified approach, surpassing existing methods in key metrics. The results are published as “OmniArch: Building Foundation Model For Scientific Computing” and the code is available at https://openi.pcl.ac.cn/cty315/OmniArch. The project is part of the CNAI4S platform: https://cnai4s.com/.

Challenges in Building a Universal PDE Solver

Constructing a general PDE foundation model faces three core obstacles:

• Multi-dimensional representation: Physical data varies from 1D sequences to 3D volumes, requiring models to handle different dimensions and share knowledge.
• Multi-physical quantities: Different PDEs involve varying numbers of physical variables (e.g., velocity, pressure, density, electromagnetic components), demanding flexible handling of complex interactions.
• Physical constraints alignment: Each PDE embodies specific physical laws (conservation, boundary conditions, symmetry). How can a flexible model adaptively satisfy these constraints to ensure physically valid solutions?

Innovations in OmniArch

OmniArch creatively integrates frequency domain transforms, attention mechanisms, and contrastive learning to overcome these challenges:

1. Fourier Encoder for Unified Representation

Despite different physical data forms, they share multi-scale structures in the frequency domain. OmniArch applies FFT to input coordinates and fields, retaining top-K frequency components to convert diverse data into a unified spectral representation.

Advantages include: (1) computational complexity reduces from O(N²) to O(N log N); (2) captures global patterns at low frequencies and local details at high frequencies; (3) different resolutions align naturally in the low-frequency spectrum.

2. Temporal Mask for Multi-Physics Coupling

Standard causal attention breaks the interaction between physical variables at the same time step. OmniArch introduces Temporal Mask:

• Intra-time step: All variables at the same time can interact, modeling coupling (e.g., velocity-pressure in Navier-Stokes).
• Inter-time step: Enforces causality, future states do not influence past.

This encodes the core temporal and instantaneous coupling features of physical systems.

3. PDE-Aligner for Physical Consistency

During fine-tuning, PDE-Aligner aligns symbolic PDEs with numerical evolution via contrastive learning:

• Equation encoding: Pretrained BERT processes LaTeX PDE texts.
• Physical features: phase difference Δφ and amplitude ratio R encode evolution patterns.
• Energy constraint: Ensures sum of |R| approximates 1, satisfying Parseval’s theorem.

This soft alignment avoids hard physical constraints, allowing adaptive handling of various equations and boundary conditions.

Experimental Validation

OmniArch was rigorously tested on 11 classic PDEs across fluid dynamics, electromagnetism, and reaction-diffusion (datasets: PDEBench, PDEArena). It outperformed specialized models (PINNs, U-Net, FNO) and state-of-the-art unified models (MPP, Poseidon, DPOT, PDEformer-1).

Performance highlights include:

• Accuracy improvements up to 98.7%.
• In 1D CFD, error reduced by 98.7% compared to the best baseline.
• Physics alignment during fine-tuning increased accuracy by over 20%.

OmniArch demonstrates excellent zero-shot generalization, handling complex PDEs without retraining, and adapts to different resolutions with minimal loss in accuracy, solving a major bottleneck in scientific computing.

Future Directions

The team plans to extend OmniArch to more PDE types, complex boundary conditions, and applications in climate modeling, aerospace design, and other scientific fields, aiming to revolutionize computational physics.

Author: Chen Tianyu, Beihang University Postdoctoral Fellow, expert in scientific AI, large models, and safety, with multiple publications and awards, contributing to the CNAI4S platform and major competitions.