PIBERT: Physics-Informed BERT-style Transformer

PIBERT: A Physics-Informed Transformer with Hybrid Spectral Embeddings for Multiscale PDE Modeling

The latest final figures, supplementary figures, and lightweight RealPDEBench result summaries are tracked in this branch. See ARTIFACTS.md for the published bundle layout and for the list of manuscript files and oversized raw artifacts that are intentionally excluded from git.

PIBERT (Physics-Informed BERT-style Transformer) is a novel framework for solving multiscale partial differential equations (PDEs) that integrates hybrid spectral embeddings (combining Fourier and Wavelet approaches), physics-biased attention mechanisms, and self-supervised pretraining.

Unlike existing approaches that only partially address the multiscale challenge, PIBERT unifies three major innovations:

• A hybrid Fourier-Wavelet embedding that captures both global structures and localized phenomena
• A physics-informed attention bias derived from PDE residuals
• A dual-task self-supervised pretraining strategy (Masked Physics Prediction & Equation Consistency Prediction)

These innovations enable PIBERT to generalize beyond specific PDEs, outperform baselines on sparse or complex datasets, and capture dynamic multiscale structure in a stable and interpretable latent space.

• Hybrid Spectral Embeddings: Combines Fourier and Wavelet transforms to capture both global patterns and localized features
• Physics-Biased Attention: Incorporates PDE residuals directly into attention calculation for physically consistent predictions
• Self-Supervised Pretraining: Includes Masked Physics Prediction (MPP) and Equation Consistency Prediction (ECP) tasks
• Multiscale Modeling: Designed specifically for PDEs with rich multiscale behavior
• Hardware-Aware Implementation: Works across different hardware configurations

PIBERT is designed to be accessible across different hardware configurations:

# Basic installation
pip install pibert

# For development with testing and documentation tools
pip install pibert[dev]

# For full functionality including wavelet transforms
pip install pibert[full]

Verify installation with CPU (runs in <60s on any system):

from pibert import PIBERT
from pibert.utils import load_dataset

# Load a small sample dataset
dataset = load_dataset("reaction_diffusion")

# Initialize a small model
model = PIBERT(
    input_dim=1,
    hidden_dim=64,
    num_layers=2,
    num_heads=4
)

# Perform prediction
pred = model.predict(dataset["test"]["x"][:1], dataset["test"]["coords"][:1])

print(f"Prediction shape: {pred.shape}")

For more examples, see the examples directory.

PIBERT demonstrates state-of-the-art performance across multiple benchmarks:

The ablation study confirms the importance of each component:

Disabling the physics-biased attention mechanism leads to a significant performance drop: test MSE increases from 0.4975 to 1.3201, and NMSE jumps from 1.34 to 9.88.

PIBERT was evaluated on the RealPDEBench benchmark suite covering two challenging real-world fluid dynamics datasets: Cylinder Wake and Fluid-Structure Interaction (FSI).

All-component metrics (u + v combined):

All-component metrics (u + v combined):

All results in the paper can be reproduced using the provided code. The ablation studies were verified on a GTX 3060 (12GB VRAM), while the full-scale experiments used A100 GPUs. We provide configuration files for both hardware setups.

To reproduce the ablation study:

jupyter notebook examples/ablation_study_gpu.ipynb

If you find PIBERT useful in your research, please cite our paper: !Link

@article{chakraborty6203487pibert,
  title={PIBERT: A Physics-Informed Bi-directional Hybrid Spectral Transformer for Multiscale CFD Surrogate Modeling},
  author={Chakraborty, Somyajit and Pan, Ming and Chen, Xizhong},
  journal={Available at SSRN 6203487}
}

Distributed under the Apache 2.0 License. See LICENSE for more information.

For support and questions, please open an issue on GitHub or contact the authors:

• Somyajit Chakraborty: chksomyajit@sjtu.edu.cn
• Pan Ming: panming@sjtu.edu.cn
• Chen Xizhong: chenxizh@sjtu.edu.cn

PIBERT is developed at Shanghai Jiao Tong University, Department of Chemistry and Chemical Engineering