PIBERT
Physics-Informed BERT-style Transformer for Multiscale PDE Modeling
Physics-Informed BERT-style Transformer for Multiscale PDE Modeling
A novel framework for solving multiscale partial differential equations (PDEs) that integrates hybrid spectral embeddings, physics-biased attention mechanisms, and self-supervised pretraining.
Three major innovations that enable PIBERT to outperform existing approaches
Combines Fourier and Wavelet transforms to capture both global patterns and localized features, enabling accurate multiscale representation.
Incorporates PDE residuals directly into attention calculation for physically consistent predictions and improved convergence.
Includes Masked Physics Prediction (MPP) and Equation Consistency Prediction (ECP) tasks for robust generalization.
Designed specifically for PDEs with rich multiscale behavior, capturing dynamic structures in a stable latent space.
Works across different hardware configurations, from consumer GPUs to high-end accelerators.
Provides interpretable latent space representations and attention patterns for physical insight.
PIBERT demonstrates state-of-the-art performance across multiple benchmarks
| Model | Relative L1 | Relative L2 | MAE |
|---|---|---|---|
| PINN | 0.0651 | 0.0803 | 0.0581 |
| FNO | 0.0123 | 0.0150 | 0.0100 |
| Transformer | 0.0225 | 0.0243 | 0.0200 |
| PINNsFormer | 0.0065 | 0.0078 | 0.0060 |
| PIBERT | 0.0061 | 0.0074 | 0.0056 |
| Model | MSE(u) | MSE(v) | MSE(p) |
|---|---|---|---|
| PINNs | 0.0500 | 0.0300 | 0.01500 |
| Spectral PINN | 0.0200 | 0.0045 | 0.00085 |
| FNO | 0.0113 | 0.0012 | 0.00021 |
| PINNsFormer | 0.0065 | 0.0007 | 0.00003 |
| PIBERT (Lite) | 0.0103 | 0.0011 | 0.000046 |
| Model Variant | MSE (Test) | NMSE (Test) |
|---|---|---|
| PIBERT (Full) | 0.4975 | 1.3409 |
| Fourier-only | 1.6520 | 12.4010 |
| Wavelet-only | 0.4123 | 1.1021 |
| Standard-attention | 1.3201 | 9.8760 |
| FNO | 1.8099 | 13.5830 |
| UNet | 3.7006 | 29.2627 |
| Task | Minimum (GTX 3060) | Recommended (A100) | Notes |
|---|---|---|---|
| Model Inference (2D) | ✓ | ✓ | 64×64 grids work on both |
| Model Training (2D) | ✓ (small batches) | ✓ | GTX 3060 requires gradient checkpointing |
| 3D Problem Inference | ✗ | ✓ | Requires 40+ GB VRAM |
| Pretraining | ✗ | ✓ | Not feasible on consumer GPUs |
PIBERT's performance on various benchmark problems
Get started with PIBERT in minutes
# 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]
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}")
All results in the paper can be reproduced using the provided code:
jupyter notebook examples/ablation_study_gpu.ipynb
If you find PIBERT useful in your research, please cite our paper
@article{chakraborty2024pibert,
title={PIBERT: A Physics-Informed Transformer with Hybrid Spectral Embeddings for Multiscale PDE Modeling},
author={Chakraborty, Somyajit, Pan, Ming and Xizhong, Chen},
year={2025}
}
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